daniel543x/artcraft
1
0 Fork 0
mirror of https://github.com/storytold/artcraft.git synced 2026-11-07 02:21:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
Brandon Thomas committed 2025-02-18 00:30:48 -05:00
1 parent 92f25f884f
commit 851b6f51fe
1 file changed
+31 -147
Download Patch File Download Diff File Expand all files Collapse all files
crates/desktop/src-tauri/src/endpoints/sd.rs
+31 -147
View File
@@ -1,7 +1,6 @@
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use crate::ml::image::dynamic_image_to_tensor::dynamic_image_to_tensor;
use crate::ml::image::save_image_from_tensor::save_image_from_tensor;
use crate::ml::image::tensor_to_image_buffer::{tensor_to_image_buffer, RgbImage};
use crate::ml::infer_clip_text_embeddings::infer_clip_text_embeddings;
use crate::ml::model_cache::ModelCache;
@@ -9,7 +8,7 @@ use crate::ml::prompt_cache::PromptCache;
use crate::ml::stable_diffusion::get_vae_scale::get_vae_scale;
use crate::state::app_config::AppConfig;
use anyhow::{Error as E, Result};
use candle_core::{DType, Device, IndexOp, D};
use candle_core::{DType, IndexOp, D};
use candle_transformers::models::stable_diffusion::vae::DiagonalGaussianDistribution;
use image::DynamicImage;
use log::info;
@@ -25,8 +24,8 @@ pub struct Args<'a> {
pub prompt: String,
pub uncond_prompt: String,
pub guidance_scale: Option<f64>,
pub model_cache: &'a ModelCache,
pub configs: &'a AppConfig,
pub model_cache: &'a ModelCache,
pub prompt_cache: &'a PromptCache,
}
@@ -45,39 +44,10 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
args.configs.scheduler_steps)?;
let seed = args.configs.seed.unwrap_or_else(|| rand::thread_rng().gen());
println!("Using seed: {}", seed);
println!("Using seed: {}", seed);
args.configs.device.set_seed(seed)?;
println!("Initializing Hugging Face API...");
let repo = args.configs.sd_version.repo();
println!("Downloading model files from: {}", repo);
//println!("Downloading VAE ...");
//let vae_file = args.api.model(repo.to_string()).get("vae/diffusion_pytorch_model.safetensors")?;
//println!("VAE Path: {:?}", &vae_file);
//println!("Downloading text encoders...");
//let clip_file = args.api.model(repo.to_string()).get("text_encoder/model.safetensors")?;
//let clip2_file = args.api.model(repo.to_string()).get("text_encoder_2/model.safetensors")?;
//let tokenizer = args.api.model("openai/clip-vit-large-patch14".to_string()).get("tokenizer.json")?;
//let tokenizer2 = args.api.model("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k".to_string()).get("tokenizer.json")?;
//println!("Building VAE model...");
//let vae = sd_config.build_vae(vae_file, &device, dtype)?;
//// Build text encoders
//println!("Building text encoders...");
//let text_model = stable_diffusion::build_clip_transformer(&sd_config.clip, clip_file, &device, dtype)?;
//let text_model2 = stable_diffusion::build_clip_transformer(sd_config.clip2.as_ref().unwrap(), clip2_file, &device, dtype)?;
//
//// Tokenize and encode text
//println!("Processing text prompt...");
//let tokenizer = tokenizers::Tokenizer::from_file(tokenizer).map_err(E::msg)?;
//let tokenizer2 = tokenizers::Tokenizer::from_file(tokenizer2).map_err(E::msg)?;
info!("Checking if prompt is cached");
let maybe_cached = args.prompt_cache.get_copy(&args.prompt)?;
@@ -85,7 +55,7 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
info!("Prompt is cached!");
tensor
} else {
info!("Prompt is NOT cached!");
info!("Prompt is NOT cached! Calculating embedding...");
let tensor = infer_clip_text_embeddings(
&args.prompt,
&args.uncond_prompt,
@@ -99,9 +69,7 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
args.configs.dtype,
false, // use_guide_scale
)?;
args.prompt_cache.store_copy(&args.prompt, &tensor)?;
tensor
};
@@ -114,31 +82,27 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
&args.configs.device,
args.configs.dtype)?;
// REFERENCE IMAGE shape: [1, 3, 1024, 1024]
println!("REFERENCE IMAGE shape: {:?}", input_image.shape());
println!("Calculating start step for diffusion process...");
let img2img_strength = 0.75f64;
let t_start = {
let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;
println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength);
start
};
let use_guide_scale = false; // TODO
println!("Reference image shape: {:?}", input_image.shape());
let vae_scale = get_vae_scale(args.configs.sd_version);
let encoded_image : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
let init_latent_dist = encoded_image;
let init_latent_dist : DiagonalGaussianDistribution = args.model_cache.vae_encode(&input_image)?;
println!("Generating latents from input image...");
let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&args.configs.device)?;
println!("Initial latents shape: {:?}", latents.shape());
println!("Calculating start step for diffusion process...");
let img2img_strength = 0.75f64;
let t_start = {
let start = args.configs.scheduler_steps - (args.configs.scheduler_steps as f64 * img2img_strength) as usize;
println!("Starting from step {} of {} (strength: {})", start, args.configs.scheduler_steps, img2img_strength);
start
};
let latents = if t_start < scheduler.timesteps().len() {
println!("Adding noise to latents...");
let noise = latents.randn_like(0f64, 1f64)?;
@@ -150,29 +114,8 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
println!("Latents initialized successfully");
//let encoded_sample = encoded_image.sample()?;
// Encoded sample shape: [1, 4, 128, 128] ... so close. Just needs to be smaller.
//println!("Encoded sample shape: {:?}", encoded_sample.shape());
//let mask_latents = (encoded_sample * vae_scale)?.to_device(&args.model_configs.device)?;
//println!("Mask latents shape: {:?}", mask_latents.shape());
//let timesteps = scheduler.timesteps().to_vec();
let timesteps: Vec<_> = scheduler.timesteps().iter().copied().collect();
//// TODO: This shape may need to change for inpainting.
//let latents = Tensor::randn(
// 0f32,
// 1f32,
// (1, 4, args.model_configs.sd_config.height / 8, args.model_configs.sd_config.width / 8),
// &args.model_configs.device,
//)?;
//let mut latents = (latents * scheduler.init_noise_sigma())?;
//let mut latents = (mask_latents * scheduler.init_noise_sigma())?;
// Initial noise shape: [1, 4, 64, 64]
println!("Initial noise shape: {:?}", latents.shape());
@@ -184,66 +127,19 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
if timestep_index < t_start {
continue;
}
// println!("Processing step {}/{}", timestep_index + 1, timesteps.len());
// let latent_model_input = latents.clone();
// // Latent input shape: [1, 4, 64, 64]
// println!("Latent input shape: {:?}", latent_model_input.shape());
//
// println!("Scaling model input...");
// let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
// // Scaled input shape: [1, 4, 64, 64]
// println!("Scaled input shape: {:?}", latent_model_input.shape());
// // TODO: THIS IS A HACK
// //let latent_model_input = match HACK_INPAINT_SD_VERSION {
// // StableDiffusionVersion::XlInpaint
// // | StableDiffusionVersion::V2Inpaint
// // | StableDiffusionVersion::V1_5Inpaint => {
// // info!("Concatenating input shape: {:?}", latent_model_input.shape());
// // info!("Mask shape: {:?}", mask.as_ref().unwrap().shape());
// // info!("Mask latents shape: {:?}", mask_latents.as_ref().unwrap().shape());
// // info!("IF THIS FAILS, REVERT THE `HACK_INPAINT_SD_VERSION` 'flag'");
// // Tensor::cat(
// // &[
// // &latent_model_input,
// // mask.as_ref().unwrap(),
// // mask_latents.as_ref().unwrap(),
// // ],
// // 1,
// // )?
// // },
// // _ => latent_model_input,
// //}
// // .to_device(&args.model_configs.device)?;
//
// println!("Running UNet inference with timestep {}...", timestep);
// let noise_pred = args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings)?;
// println!("Noise prediction shape: {:?}", noise_pred.shape());
//
// println!("Applying scheduler step...");
// latents = scheduler.step(&noise_pred, timestep, &latents)?;
// println!("Step {}/{} completed", timestep_index + 1, timesteps.len());
println!("Processing step {}/{}", timestep_index + 1, timesteps.len());
let latent_model_input = latents.clone();
println!("Latent input shape: {:?}", latent_model_input.shape());
println!("Scaling model input...");
let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)?;
println!("Scaled input shape: {:?}", latent_model_input.shape());
println!("Running UNet inference with timestep {}...", timestep);
let noise_pred = match args.model_cache.unet_inference(&latent_model_input, timestep as f64, &text_embeddings) {
Ok(pred) => {
println!("UNet inference successful");
@@ -261,34 +157,22 @@ pub fn run(args: Args<'_>) -> Result<RgbImage> {
println!("Step {}/{} completed", timestep_index + 1, timesteps.len());
}
//println!("Diffusion process completed, decoding image...");
//let image = args.model_cache.vae_decode(&(latents / 0.13025)?)?;
//println!("VAE decode completed");
//let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
//println!("Normalized image values");
//let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
//println!("Converted to 8-bit format");
//
//save_image_from_tensor(&image.i(0)?, "temp.png")?;
//println!("Image generation completed successfully");
println!("Diffusion process completed, decoding image...");
let image = args.model_cache.vae_decode(&(latents / vae_scale)?)?;
//let image = args.model_cache.vae_decode(&(latents.div(vae_scale)?))?;
println!("VAE decode completed");
//let image = ((image.div(2.))? + 0.5)?.to_device(&Device::Cpu)?;
let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
println!("Normalized image values");
let image = (image.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
println!("Converted to 8-bit format");
//save_image_from_tensor(&image.i(0)?, "temp.png")?;
//println!("Image generation completed successfully");
println!("Scaling image back");
let image = ((image / 2.)? + 0.5)?;
// TODO(bt,2025-08-18): This normalization is slow.
println!("Normalized image values");
let image = (image.clamp(0f32, 1.)? * 255.)?;
println!("Convert to int8");
let image = image.to_dtype(DType::U8)?;
println!("Converted to 8-bit format");
let image = tensor_to_image_buffer(&image.i(0)?)?;