mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-09 03:21:23 +00:00
basic lcm implementation
This commit is contained in:
Kasisnu committed
1 parent
5d9358103c
commit
e67cacc23f
6 files changed
+1948
-355
No files matched your search
@@ -122,3 +122,7 @@ required-features = ["onnx"]
|
||||
[[example]]
|
||||
name = "colpali"
|
||||
required-features = ["pdf2image"]
|
||||
|
||||
[[example]]
|
||||
name = "lcm_debug_test"
|
||||
path = "examples/lcm/debug_test.rs"
|
||||
@@ -0,0 +1,872 @@
|
||||
#[cfg(feature = "accelerate")]
|
||||
extern crate accelerate_src;
|
||||
|
||||
#[cfg(feature = "mkl")]
|
||||
extern crate intel_mkl_src;
|
||||
|
||||
use candle_transformers::models::stable_diffusion;
|
||||
use std::ops::Div;
|
||||
|
||||
use anyhow::{Error as E, Result};
|
||||
use candle::{DType, Device, IndexOp, Module, Tensor, D};
|
||||
use clap::Parser;
|
||||
use rand::Rng;
|
||||
use stable_diffusion::vae::AutoEncoderKL;
|
||||
use tokenizers::Tokenizer;
|
||||
use candle_nn::VarBuilder;
|
||||
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModel;
|
||||
use candle_transformers::models::stable_diffusion::unet_2d::UNet2DConditionModelConfig;
|
||||
use candle_transformers::models::stable_diffusion::unet_2d::BlockConfig;
|
||||
|
||||
const PRINT_DEBUG_TENSORS: bool = false;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
struct Args {
|
||||
/// The prompt to be used for image generation.
|
||||
#[arg(long, default_value = "A very realistic photo of a rusty robot walking on a sandy beach")]
|
||||
prompt: String,
|
||||
|
||||
#[arg(long, default_value = "")]
|
||||
uncond_prompt: String,
|
||||
|
||||
/// Run on CPU rather than on GPU.
|
||||
#[arg(long)]
|
||||
cpu: bool,
|
||||
|
||||
/// Enable tracing (generates a trace-timestamp.json file).
|
||||
#[arg(long)]
|
||||
tracing: bool,
|
||||
|
||||
/// The height in pixels of the generated image.
|
||||
#[arg(long)]
|
||||
height: Option<usize>,
|
||||
|
||||
/// The width in pixels of the generated image.
|
||||
#[arg(long)]
|
||||
width: Option<usize>,
|
||||
|
||||
/// The UNet weight file, in .safetensors format.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
unet_weights: Option<String>,
|
||||
|
||||
/// The CLIP weight file, in .safetensors format.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
clip_weights: Option<String>,
|
||||
|
||||
/// The CLIP2 weight file, in .safetensors format.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
clip2_weights: Option<String>,
|
||||
|
||||
/// The VAE weight file, in .safetensors format.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
vae_weights: Option<String>,
|
||||
|
||||
#[arg(long, value_name = "FILE")]
|
||||
/// The file specifying the tokenizer to used for tokenization.
|
||||
tokenizer: Option<String>,
|
||||
|
||||
/// The size of the sliced attention or 0 for automatic slicing (disabled by default)
|
||||
#[arg(long)]
|
||||
sliced_attention_size: Option<usize>,
|
||||
|
||||
/// The number of inference steps to run.
|
||||
#[arg(long, default_value_t = 4)]
|
||||
n_steps: usize,
|
||||
|
||||
/// The number of samples to generate iteratively.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
num_samples: usize,
|
||||
|
||||
/// The numbers of samples to generate simultaneously.
|
||||
#[arg[long, default_value_t = 1]]
|
||||
bsize: usize,
|
||||
|
||||
/// The name of the final image to generate.
|
||||
#[arg(long, value_name = "FILE", default_value = "sd_final.png")]
|
||||
final_image: String,
|
||||
|
||||
#[arg(long, value_enum, default_value = "v2-1")]
|
||||
sd_version: StableDiffusionVersion,
|
||||
|
||||
/// Generate intermediary images at each step.
|
||||
#[arg(long, action)]
|
||||
intermediary_images: bool,
|
||||
|
||||
#[arg(long)]
|
||||
use_flash_attn: bool,
|
||||
|
||||
#[arg(long)]
|
||||
use_f16: bool,
|
||||
|
||||
/// The guidance scale for classifier-free guidance.
|
||||
#[arg(long, default_value_t = 2.0)]
|
||||
guidance_scale: f64,
|
||||
|
||||
/// Path to the mask image for inpainting.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
mask_path: Option<String>,
|
||||
|
||||
/// Path to the image used to initialize the latents. For inpainting, this is the image to be masked.
|
||||
#[arg(long, value_name = "FILE")]
|
||||
img2img: Option<String>,
|
||||
|
||||
/// The strength parameter for img2img, between 0 and 1.
|
||||
#[arg(long, default_value_t = 0.7)]
|
||||
strength: f64,
|
||||
|
||||
/// The seed to use when generating random samples.
|
||||
#[arg(long)]
|
||||
seed: Option<u64>,
|
||||
|
||||
/// Force the saved image to update only the masked region
|
||||
#[arg(long)]
|
||||
only_update_masked: bool,
|
||||
|
||||
#[arg(long)]
|
||||
save_tensors: bool,
|
||||
|
||||
/// Use half precision (F16) rather than F32
|
||||
#[arg(long)]
|
||||
half_precision: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
|
||||
enum StableDiffusionVersion {
|
||||
V1_5,
|
||||
V1_5Inpaint,
|
||||
V2_1,
|
||||
V2Inpaint,
|
||||
Xl,
|
||||
XlInpaint,
|
||||
Turbo,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ModelFile {
|
||||
Tokenizer,
|
||||
Tokenizer2,
|
||||
Clip,
|
||||
Clip2,
|
||||
Unet,
|
||||
UnetLcm,
|
||||
Vae,
|
||||
}
|
||||
|
||||
impl StableDiffusionVersion {
|
||||
fn repo(&self) -> &'static str {
|
||||
match self {
|
||||
Self::XlInpaint => "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
Self::Xl => "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
Self::V2Inpaint => "stabilityai/stable-diffusion-2-inpainting",
|
||||
Self::V2_1 => "stabilityai/stable-diffusion-2-1",
|
||||
Self::V1_5 => "runwayml/stable-diffusion-v1-5",
|
||||
Self::V1_5Inpaint => "stable-diffusion-v1-5/stable-diffusion-inpainting",
|
||||
Self::Turbo => "stabilityai/sdxl-turbo",
|
||||
}
|
||||
}
|
||||
|
||||
fn unet_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
|
||||
if use_f16 {
|
||||
"unet/diffusion_pytorch_model.fp16.safetensors"
|
||||
} else {
|
||||
"unet/diffusion_pytorch_model.safetensors"
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn vae_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
|
||||
if use_f16 {
|
||||
"vae/diffusion_pytorch_model.fp16.safetensors"
|
||||
} else {
|
||||
"vae/diffusion_pytorch_model.safetensors"
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn clip_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
|
||||
if use_f16 {
|
||||
"text_encoder/model.fp16.safetensors"
|
||||
} else {
|
||||
"text_encoder/model.safetensors"
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn clip2_file(&self, use_f16: bool) -> &'static str {
|
||||
match self {
|
||||
Self::V1_5 | Self::V1_5Inpaint | Self::V2_1 | Self::V2Inpaint | Self::Xl | Self::XlInpaint | Self::Turbo => {
|
||||
if use_f16 {
|
||||
"text_encoder_2/model.fp16.safetensors"
|
||||
} else {
|
||||
"text_encoder_2/model.safetensors"
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelFile {
|
||||
fn get(&self, filename: Option<String>, version: StableDiffusionVersion, use_f16: bool) -> Result<std::path::PathBuf> {
|
||||
use hf_hub::api::sync::Api;
|
||||
match filename {
|
||||
Some(filename) => Ok(std::path::PathBuf::from(filename)),
|
||||
None => {
|
||||
let (repo, path) = match self {
|
||||
Self::Tokenizer => {
|
||||
let tokenizer_repo = match version {
|
||||
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2Inpaint => "openai/clip-vit-base-patch32",
|
||||
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => {
|
||||
// This seems similar to the patch32 version except some very small
|
||||
// difference in the split regex.
|
||||
"openai/clip-vit-large-patch14"
|
||||
},
|
||||
};
|
||||
(tokenizer_repo, "tokenizer.json")
|
||||
},
|
||||
Self::Tokenizer2 => ("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", "tokenizer.json"),
|
||||
Self::Clip => (version.repo(), version.clip_file(use_f16)),
|
||||
Self::Clip2 => (version.repo(), version.clip2_file(use_f16)),
|
||||
Self::UnetLcm => ("SimianLuo/LCM_Dreamshaper_v7", if use_f16 { "unet/diffusion_pytorch_model.fp16.safetensors" } else { "unet/diffusion_pytorch_model.safetensors" }),
|
||||
Self::Unet => (version.repo(), version.unet_file(use_f16)),
|
||||
Self::Vae => {
|
||||
// Override for SDXL when using f16 weights.
|
||||
// See https://github.com/huggingface/candle/issues/1060
|
||||
if matches!(version, StableDiffusionVersion::Xl | StableDiffusionVersion::Turbo,) && use_f16 {
|
||||
("madebyollin/sdxl-vae-fp16-fix", "diffusion_pytorch_model.safetensors")
|
||||
} else {
|
||||
(version.repo(), version.vae_file(use_f16))
|
||||
}
|
||||
},
|
||||
};
|
||||
let filename = Api::new()?.model(repo.to_string()).get(path)?;
|
||||
Ok(filename)
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn output_filename(basename: &str, sample_idx: usize, num_samples: usize, timestep_idx: Option<usize>) -> String {
|
||||
let filename = if num_samples > 1 {
|
||||
match basename.rsplit_once('.') {
|
||||
None => format!("{basename}.{sample_idx}.png"),
|
||||
Some((filename_no_extension, extension)) => {
|
||||
format!("{filename_no_extension}.{sample_idx}.{extension}")
|
||||
},
|
||||
}
|
||||
} else {
|
||||
basename.to_string()
|
||||
};
|
||||
match timestep_idx {
|
||||
None => filename,
|
||||
Some(timestep_idx) => match filename.rsplit_once('.') {
|
||||
None => format!("{filename}-{timestep_idx}.png"),
|
||||
Some((filename_no_extension, extension)) => {
|
||||
format!("{filename_no_extension}-{timestep_idx}.{extension}")
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn save_image(vae: &AutoEncoderKL, latents: &Tensor, vae_scale: f64, bsize: usize, idx: usize, final_image: &str, num_samples: usize, timestep_ids: Option<usize>) -> Result<()> {
|
||||
let images = vae.decode(&(latents / vae_scale)?)?;
|
||||
let images = ((images / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
|
||||
let images = (images.clamp(0f32, 1.)? * 255.)?.to_dtype(DType::U8)?;
|
||||
for batch in 0..bsize {
|
||||
let image = images.i(batch)?;
|
||||
let image_filename = output_filename(final_image, (bsize * idx) + batch + 1, batch + num_samples, timestep_ids);
|
||||
candle_examples::save_image(&image, image_filename)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn text_embeddings(prompt: &str, uncond_prompt: &str, tokenizer: Option<String>, clip_weights: Option<String>, clip2_weights: Option<String>, sd_version: StableDiffusionVersion, sd_config: &stable_diffusion::StableDiffusionConfig, use_f16: bool, device: &Device, dtype: DType, use_guide_scale: bool, first: bool) -> Result<Tensor> {
|
||||
let tokenizer_file = if first { ModelFile::Tokenizer } else { ModelFile::Tokenizer2 };
|
||||
let tokenizer = tokenizer_file.get(tokenizer, sd_version, use_f16)?;
|
||||
let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
|
||||
let pad_id = match &sd_config.clip.pad_with {
|
||||
Some(padding) => *tokenizer.get_vocab(true).get(padding.as_str()).unwrap(),
|
||||
None => *tokenizer.get_vocab(true).get("<|endoftext|>").unwrap(),
|
||||
};
|
||||
println!("Running with prompt \"{prompt}\".");
|
||||
let mut tokens = tokenizer.encode(prompt, true).map_err(E::msg)?.get_ids().to_vec();
|
||||
if tokens.len() > sd_config.clip.max_position_embeddings {
|
||||
anyhow::bail!("the prompt is too long, {} > max-tokens ({})", tokens.len(), sd_config.clip.max_position_embeddings)
|
||||
}
|
||||
while tokens.len() < sd_config.clip.max_position_embeddings {
|
||||
tokens.push(pad_id)
|
||||
}
|
||||
let tokens = Tensor::new(tokens.as_slice(), device)?.unsqueeze(0)?;
|
||||
|
||||
println!("Building the Clip transformer.");
|
||||
let clip_weights_file = if first { ModelFile::Clip } else { ModelFile::Clip2 };
|
||||
let clip_weights = if first { clip_weights_file.get(clip_weights, sd_version, use_f16)? } else { clip_weights_file.get(clip2_weights, sd_version, use_f16)? };
|
||||
let clip_config = if first { &sd_config.clip } else { sd_config.clip2.as_ref().unwrap() };
|
||||
let text_model = stable_diffusion::build_clip_transformer(clip_config, clip_weights, device, DType::F32)?;
|
||||
let text_embeddings = text_model.forward(&tokens)?;
|
||||
|
||||
let text_embeddings = if use_guide_scale {
|
||||
let mut uncond_tokens = tokenizer.encode(uncond_prompt, true).map_err(E::msg)?.get_ids().to_vec();
|
||||
if uncond_tokens.len() > sd_config.clip.max_position_embeddings {
|
||||
anyhow::bail!("the negative prompt is too long, {} > max-tokens ({})", uncond_tokens.len(), sd_config.clip.max_position_embeddings)
|
||||
}
|
||||
while uncond_tokens.len() < sd_config.clip.max_position_embeddings {
|
||||
uncond_tokens.push(pad_id)
|
||||
}
|
||||
|
||||
let uncond_tokens = Tensor::new(uncond_tokens.as_slice(), device)?.unsqueeze(0)?;
|
||||
let uncond_embeddings = text_model.forward(&uncond_tokens)?;
|
||||
|
||||
Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?.to_dtype(dtype)?
|
||||
} else {
|
||||
text_embeddings.to_dtype(dtype)?
|
||||
};
|
||||
Ok(text_embeddings)
|
||||
}
|
||||
|
||||
fn image_preprocess<T: AsRef<std::path::Path>>(path: T) -> anyhow::Result<Tensor> {
|
||||
let img = image::ImageReader::open(path)?.decode()?;
|
||||
let (height, width) = (img.height() as usize, img.width() as usize);
|
||||
let height = height - height % 32;
|
||||
let width = width - width % 32;
|
||||
let img = img.resize_to_fill(width as u32, height as u32, image::imageops::FilterType::CatmullRom);
|
||||
let img = img.to_rgb8();
|
||||
let img = img.into_raw();
|
||||
let img = Tensor::from_vec(img, (height, width, 3), &Device::Cpu)?.permute((2, 0, 1))?.to_dtype(DType::F32)?.affine(2. / 255., -1.)?.unsqueeze(0)?;
|
||||
Ok(img)
|
||||
}
|
||||
|
||||
/// Convert the mask image to a single channel tensor. Also ensure the image is a multiple of 32 in both dimensions.
|
||||
fn mask_preprocess<T: AsRef<std::path::Path>>(path: T) -> anyhow::Result<Tensor> {
|
||||
let img = image::open(path)?.to_luma8();
|
||||
let (new_width, new_height) = {
|
||||
let (width, height) = img.dimensions();
|
||||
(width - width % 32, height - height % 32)
|
||||
};
|
||||
let img = image::imageops::resize(&img, new_width, new_height, image::imageops::FilterType::CatmullRom).into_raw();
|
||||
let mask = Tensor::from_vec(img, (new_height as usize, new_width as usize), &Device::Cpu)?.unsqueeze(0)?.to_dtype(DType::F32)?.div(255.0)?.unsqueeze(0)?;
|
||||
Ok(mask)
|
||||
}
|
||||
|
||||
/// Generates the mask latents, scaled mask and mask_4 for inpainting. Returns a tuple of None if inpainting is not
|
||||
/// being used.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn inpainting_tensors(sd_version: StableDiffusionVersion, mask_path: Option<String>, dtype: DType, device: &Device, use_guide_scale: bool, vae: &AutoEncoderKL, image: Option<Tensor>, vae_scale: f64) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
|
||||
match sd_version {
|
||||
StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => {
|
||||
let inpaint_mask = mask_path.ok_or_else(|| anyhow::anyhow!("An inpainting model was requested but mask-path is not provided."))?;
|
||||
// Get the mask image with shape [1, 1, 128, 128]
|
||||
let mask = mask_preprocess(inpaint_mask)?.to_device(device)?.to_dtype(dtype)?;
|
||||
// Generate the masked image from the image and the mask with shape [1, 3, 1024, 1024]
|
||||
let xmask = mask.le(0.5)?.repeat(&[1, 3, 1, 1])?.to_dtype(dtype)?;
|
||||
let image = &image.ok_or_else(|| anyhow::anyhow!("An inpainting model was requested but img2img which is used as the input image is not provided."))?;
|
||||
let masked_img = (image * xmask)?;
|
||||
// Scale down the mask
|
||||
let shape = masked_img.shape();
|
||||
let (w, h) = (shape.dims()[3] / 8, shape.dims()[2] / 8);
|
||||
let mask = mask.interpolate2d(w, h)?;
|
||||
// shape: [1, 4, 128, 128]
|
||||
let mask_latents = vae.encode(&masked_img)?;
|
||||
let mask_latents = (mask_latents.sample()? * vae_scale)?.to_device(device)?;
|
||||
|
||||
let mask_4 = mask.as_ref().repeat(&[1, 4, 1, 1])?;
|
||||
let (mask_latents, mask) = if use_guide_scale { (Tensor::cat(&[&mask_latents, &mask_latents], 0)?, Tensor::cat(&[&mask, &mask], 0)?) } else { (mask_latents, mask) };
|
||||
Ok((Some(mask_latents), Some(mask), Some(mask_4)))
|
||||
},
|
||||
_ => Ok((None, None, None)),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_tensor_stats(tensor: &Tensor) -> Result<(f32, f32, f32, f32)> {
|
||||
// Convert tensor to F32 for stats calculation - makes statistics more accurate anyway
|
||||
let tensor_f32 = tensor.to_dtype(DType::F32)?;
|
||||
let flat = tensor_f32.flatten_all()?;
|
||||
|
||||
let min = flat.min(0)?.to_scalar::<f32>()?;
|
||||
let max = flat.max(0)?.to_scalar::<f32>()?;
|
||||
let mean = flat.mean_all()?.to_scalar::<f32>()?;
|
||||
|
||||
// Create a tensor with the same shape as flat but filled with the mean value
|
||||
let mean_tensor = Tensor::new(&[mean], flat.device())?;
|
||||
let mean_tensor = mean_tensor.broadcast_as(flat.shape().dims())?;
|
||||
let diff = (flat - mean_tensor)?;
|
||||
let squared = diff.sqr()?;
|
||||
let variance = squared.mean_all()?;
|
||||
let std = variance.sqrt()?.to_scalar::<f32>()?;
|
||||
|
||||
Ok((min, max, mean, std))
|
||||
}
|
||||
|
||||
// Add this function to track statistics at key points
|
||||
fn log_tensor_stats(name: &str, tensor: &Tensor) -> Result<()> {
|
||||
let (min, max, mean, std) = get_tensor_stats(tensor)?;
|
||||
println!("{} stats: min={:.5}, max={:.5}, mean={:.5}, std={:.5}", name, min, max, mean, std);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Add this simpler alternative function
|
||||
fn save_tensor_as_numpy(tensor: &Tensor, filename: &str) -> Result<()> {
|
||||
use std::fs::File;
|
||||
use std::io::Write;
|
||||
|
||||
// Move tensor to CPU and convert to f32 for consistent serialization
|
||||
let tensor_cpu = tensor.to_device(&Device::Cpu)?.to_dtype(DType::F32)?;
|
||||
let data = tensor_cpu.to_vec1::<f32>()?;
|
||||
|
||||
// Create a simple binary file with the tensor data
|
||||
let mut file = File::create(filename)?;
|
||||
|
||||
// Write shape information first (as u32 values)
|
||||
let shape = tensor.shape().dims();
|
||||
file.write_all(&(shape.len() as u32).to_le_bytes())?;
|
||||
for &dim in shape {
|
||||
file.write_all(&(dim as u32).to_le_bytes())?;
|
||||
}
|
||||
|
||||
// Write all float values
|
||||
for val in data {
|
||||
file.write_all(&val.to_le_bytes())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run(args: Args) -> Result<()> {
|
||||
use tracing_chrome::ChromeLayerBuilder;
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
let Args { prompt, uncond_prompt, cpu, height, width, n_steps, tokenizer, final_image, sliced_attention_size, num_samples, bsize, sd_version, clip_weights, clip2_weights, vae_weights, unet_weights, tracing, use_f16, guidance_scale, use_flash_attn, mask_path, img2img, strength, seed, only_update_masked, save_tensors, intermediary_images, half_precision } = args;
|
||||
|
||||
if !(0. ..=1.).contains(&strength) {
|
||||
anyhow::bail!("strength should be between 0 and 1, got {strength}")
|
||||
}
|
||||
|
||||
let _guard = if tracing {
|
||||
let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
|
||||
tracing_subscriber::registry().with(chrome_layer).init();
|
||||
Some(guard)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
println!("=== LCM-Dreamshaper F16 DEBUG MODE ===");
|
||||
println!("Using dtype: {:?}", if half_precision { DType::F16 } else { DType::F32 });
|
||||
|
||||
let dtype = if half_precision { DType::F16 } else { DType::F32 };
|
||||
let sd_config = match sd_version {
|
||||
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => stable_diffusion::StableDiffusionConfig::v1_5(sliced_attention_size, height, width),
|
||||
StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint => stable_diffusion::StableDiffusionConfig::v2_1(sliced_attention_size, height, width),
|
||||
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint => stable_diffusion::StableDiffusionConfig::sdxl(sliced_attention_size, height, width),
|
||||
StableDiffusionVersion::Turbo => stable_diffusion::StableDiffusionConfig::sdxl_turbo(sliced_attention_size, height, width),
|
||||
};
|
||||
|
||||
let mut scheduler = stable_diffusion::lcm::LCMScheduler::new(n_steps, stable_diffusion::lcm::LCMSchedulerConfig::default())?;
|
||||
println!("Scheduler created with dtype={:?}", scheduler.dtype());
|
||||
|
||||
let device = candle_examples::device(cpu)?;
|
||||
// If a seed is not given, generate a random seed and print it
|
||||
let seed = seed.unwrap_or(rand::thread_rng().gen_range(0u64..u64::MAX));
|
||||
println!("Using seed {seed}");
|
||||
device.set_seed(seed)?;
|
||||
let use_guide_scale = guidance_scale > 1.0;
|
||||
|
||||
let which = match sd_version {
|
||||
StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false],
|
||||
_ => vec![true],
|
||||
};
|
||||
|
||||
// Define the base model and LCM model repositories
|
||||
let base_model_repo = match sd_version {
|
||||
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint => {
|
||||
// Dreamshaper is based on SD 1.5, so use it for those versions
|
||||
"Lykon/dreamshaper-7"
|
||||
},
|
||||
_ => {
|
||||
// For other versions, use the standard repo
|
||||
sd_version.repo()
|
||||
},
|
||||
};
|
||||
let lcm_model_repo = "SimianLuo/LCM_Dreamshaper_v7"; // LCM UNet
|
||||
|
||||
// Create a custom ModelFile enum variant for the base model
|
||||
let clip_weights = if clip_weights.is_none() {
|
||||
use hf_hub::api::sync::Api;
|
||||
let api = Api::new()?;
|
||||
let clip_path = api.model(base_model_repo.to_string()).get(if use_f16 { "text_encoder/model.fp16.safetensors" } else { "text_encoder/model.safetensors" })?;
|
||||
Some(clip_path.to_string_lossy().to_string())
|
||||
} else {
|
||||
clip_weights
|
||||
};
|
||||
|
||||
let vae_weights = if vae_weights.is_none() {
|
||||
use hf_hub::api::sync::Api;
|
||||
let api = Api::new()?;
|
||||
let vae_path = api.model(base_model_repo.to_string()).get(if use_f16 { "vae/diffusion_pytorch_model.fp16.safetensors" } else { "vae/diffusion_pytorch_model.safetensors" })?;
|
||||
Some(vae_path.to_string_lossy().to_string())
|
||||
} else {
|
||||
vae_weights
|
||||
};
|
||||
|
||||
// For the UNet, always use the LCM model
|
||||
let unet_weights = {
|
||||
use hf_hub::api::sync::Api;
|
||||
let api = Api::new()?;
|
||||
let unet_path = api.model(lcm_model_repo.to_string()).get(if use_f16 {
|
||||
if lcm_model_repo == "SimianLuo/LCM_Dreamshaper_v7" {
|
||||
"unet/diffusion_pytorch_model.safetensors"
|
||||
} else {
|
||||
"unet/diffusion_pytorch_model.fp16.safetensors"
|
||||
}
|
||||
} else {
|
||||
"unet/diffusion_pytorch_model.safetensors"
|
||||
})?;
|
||||
Some(unet_path.to_string_lossy().to_string())
|
||||
};
|
||||
|
||||
// Load text embeddings from the base model
|
||||
let text_embeddings = which.iter().map(|first| text_embeddings(&prompt, &uncond_prompt, tokenizer.clone(), clip_weights.clone(), clip2_weights.clone(), sd_version, &sd_config, use_f16, &device, dtype, use_guide_scale, *first)).collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?;
|
||||
let text_embeddings = text_embeddings.repeat((bsize, 1, 1))?;
|
||||
println!("Text embeddings shape: {:?}", text_embeddings.shape());
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Text embeddings first few values: {}", text_embeddings.i(0)?.i(0)?.narrow(0, 0, 10)?);
|
||||
}
|
||||
|
||||
// Replace the random projection with a fixed one to match diffusers
|
||||
let text_embeddings = if text_embeddings.dim(2)? == 2048 {
|
||||
println!("Projecting text embeddings from 2048 to 768 dimensions");
|
||||
|
||||
// Load the projection matrix from a file (you'll need to extract this from the diffusers model)
|
||||
// For now, we'll still use a random one but with a fixed seed
|
||||
device.set_seed(42)?; // Use a fixed seed for the projection
|
||||
let projection = Tensor::randn(0.0, 0.02, (2048, 768), &device)?.to_dtype(dtype)?;
|
||||
|
||||
let proj_embeddings = text_embeddings.matmul(&projection)?;
|
||||
proj_embeddings
|
||||
} else {
|
||||
text_embeddings
|
||||
};
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Projected text embeddings: {:?}", text_embeddings);
|
||||
}
|
||||
|
||||
println!("Building the autoencoder.");
|
||||
let vae_weights = ModelFile::Vae.get(vae_weights, sd_version, use_f16)?;
|
||||
let vae = sd_config.build_vae(vae_weights, &device, dtype)?;
|
||||
|
||||
let (image, init_latent_dist) = match &img2img {
|
||||
None => (None, None),
|
||||
Some(image) => {
|
||||
let image = image_preprocess(image)?.to_device(&device)?.to_dtype(dtype)?;
|
||||
println!("Input image shape: {:?}", image.shape());
|
||||
(Some(image.clone()), Some(vae.encode(&image)?))
|
||||
},
|
||||
};
|
||||
|
||||
println!(
|
||||
"Latent shape after VAE encoding: {:?}",
|
||||
init_latent_dist.as_ref().map(|dist| {
|
||||
// Try to get shape from a temporary sample
|
||||
match dist.sample() {
|
||||
Ok(sample) => format!("{:?}", sample.shape()),
|
||||
Err(_) => "Unable to determine shape".to_string(),
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
println!("Building the unet.");
|
||||
let unet_weights = ModelFile::UnetLcm.get(unet_weights, sd_version, use_f16)?;
|
||||
let in_channels = match sd_version {
|
||||
StableDiffusionVersion::XlInpaint | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::V1_5Inpaint => 9,
|
||||
_ => 4,
|
||||
};
|
||||
println!("Loading UNet weights from: SimianLuo/LCM_Dreamshaper_v7/unet");
|
||||
|
||||
// Ensure we're loading the correct UNet configuration
|
||||
println!("Building UNet for LCM (SimianLuo/LCM_Dreamshaper_v7)");
|
||||
|
||||
// For LCM models, ensure that:
|
||||
// 1. We're using the correct prediction type (epsilon)
|
||||
// 2. The cross-attention dimension is correct (768 for Dreamshaper)
|
||||
// 3. The in_channels matches the expected input (4 for standard, 9 for inpainting)
|
||||
let unet_config = UNet2DConditionModelConfig {
|
||||
// LCM specific config values
|
||||
cross_attention_dim: 768, // Dreamshaper v7 uses 768
|
||||
// Rest of configuration...
|
||||
center_input_sample: false,
|
||||
flip_sin_to_cos: true,
|
||||
freq_shift: 0.0,
|
||||
blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }],
|
||||
layers_per_block: 2,
|
||||
downsample_padding: 1,
|
||||
mid_block_scale_factor: 1.0,
|
||||
norm_num_groups: 32,
|
||||
norm_eps: 1e-5,
|
||||
sliced_attention_size: None,
|
||||
use_linear_projection: false,
|
||||
};
|
||||
|
||||
// Load the model directly from the safetensors file
|
||||
let vs_unet = unsafe { VarBuilder::from_mmaped_safetensors(&[unet_weights], dtype, &device)? };
|
||||
|
||||
let unet = UNet2DConditionModel::new(vs_unet, in_channels, 4, use_flash_attn, unet_config)?;
|
||||
|
||||
// Ensure strength is at least 0.4 for img2img to work well
|
||||
let effective_strength = if img2img.is_some() {
|
||||
strength.max(0.4).min(0.99) // Keep between 0.4 and 0.99
|
||||
} else {
|
||||
strength
|
||||
};
|
||||
|
||||
// Make sure we're doing a proper image-to-image transform
|
||||
// with the right amount of noise
|
||||
println!("Using img2img with strength={:.4}", effective_strength);
|
||||
let mut raw_t_start = n_steps - (n_steps as f64 * effective_strength).ceil() as usize;
|
||||
|
||||
let vae_scale = match sd_version {
|
||||
StableDiffusionVersion::V1_5 | StableDiffusionVersion::V1_5Inpaint | StableDiffusionVersion::V2_1 | StableDiffusionVersion::V2Inpaint | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Xl => 0.18215,
|
||||
StableDiffusionVersion::Turbo => 0.13025,
|
||||
};
|
||||
|
||||
let (mask_latents, mask, mask_4) = inpainting_tensors(sd_version, mask_path, dtype, &device, use_guide_scale, &vae, image, vae_scale)?;
|
||||
|
||||
println!("Timesteps: {:?}", scheduler.timesteps());
|
||||
// println!("Alpha cumprod: {:?}", scheduler.alphas_cumprod());
|
||||
|
||||
// Enable debug dumps to compare with diffusers
|
||||
let enable_debug_dumps = true;
|
||||
let debug_dir = "lcm_rust_debug";
|
||||
if enable_debug_dumps {
|
||||
std::fs::create_dir_all(debug_dir).expect("Failed to create debug directory");
|
||||
|
||||
// Log the scheduler configuration
|
||||
let mut config_file = std::fs::File::create(format!("{}/scheduler_config.txt", debug_dir)).expect("Failed to create config file");
|
||||
use std::io::Write;
|
||||
writeln!(config_file, "Rust LCM Configuration:").unwrap();
|
||||
writeln!(config_file, " beta_start: {}", scheduler.config.beta_start).unwrap();
|
||||
writeln!(config_file, " beta_end: {}", scheduler.config.beta_end).unwrap();
|
||||
// More config params...
|
||||
writeln!(config_file, "\nTimesteps: {:?}", scheduler.timesteps()).unwrap();
|
||||
writeln!(config_file, "First 10 alphas_cumprod: {:?}", scheduler.alphas_cumprod().iter().take(10).collect::<Vec<_>>()).unwrap();
|
||||
}
|
||||
|
||||
// Add this code before the loop where we iterate through samples
|
||||
// This will generate and log the guidance scale embedding
|
||||
|
||||
// Create and log guidance scale embedding for comparison with Python
|
||||
println!("\n==== Generating LCM Guidance Scale Embedding ====");
|
||||
let guidance_scale = args.guidance_scale;
|
||||
println!("Guidance scale: {}", guidance_scale);
|
||||
|
||||
// The embedding dimension should match the UNet time embedding dimension
|
||||
// For LCM Dreamshaper, this is typically 1280 (not 320)
|
||||
// Use a hardcoded value since we don't have direct access to the config
|
||||
let embedding_dim = 1280; // Standard for SD models based on UNet config
|
||||
println!("Embedding dimension: {}", embedding_dim);
|
||||
|
||||
// Generate the embedding
|
||||
let guidance_embedding = scheduler.get_guidance_scale_embedding(guidance_scale, embedding_dim, &device, dtype)?;
|
||||
|
||||
println!("Guidance embedding shape: {:?}", guidance_embedding.shape());
|
||||
|
||||
// Save the embedding to a numpy file for comparison with Python
|
||||
if save_tensors {
|
||||
let tensor_dir = "lcm_rust_tensors";
|
||||
std::fs::create_dir_all(tensor_dir)?;
|
||||
|
||||
save_tensor_as_numpy(&guidance_embedding, &format!("{}/guidance_embedding.npy", tensor_dir))?;
|
||||
println!("Saved guidance embedding to {}/guidance_embedding.npy", tensor_dir);
|
||||
}
|
||||
println!("===============================================\n");
|
||||
|
||||
for idx in 0..num_samples {
|
||||
let timesteps = scheduler.timesteps().to_vec();
|
||||
|
||||
// Initialize latents
|
||||
let latents = if let Some(init_latent_dist) = &init_latent_dist {
|
||||
println!("Using img2img with strength={:.4}", effective_strength);
|
||||
|
||||
// LCM can be more sensitive to the starting timestep
|
||||
// Get the latents from the VAE
|
||||
let latents = (init_latent_dist.sample()? * vae_scale)?.to_device(&device)?;
|
||||
|
||||
// Calculate the actual start timestep based on strength
|
||||
// In LCM, starting too late can cause artifacts, so ensure we start early enough
|
||||
let t_start_index = (timesteps.len() as f64 * (1.0 - effective_strength.max(0.5))).floor() as usize;
|
||||
let t_start = timesteps[t_start_index.min(timesteps.len() - 1)];
|
||||
|
||||
println!("Starting from timestep {}/{} (index: {})", t_start, timesteps[0], t_start_index);
|
||||
|
||||
// Use a fixed seed for noise to improve reproducibility
|
||||
device.set_seed(seed)?;
|
||||
let noise = (Tensor::randn(0f32, 1f32, latents.shape(), &device)?).to_dtype(dtype)?;
|
||||
|
||||
// Clone the noise when passing it to add_noise
|
||||
let noised_latents = scheduler.add_noise(&latents, noise.clone(), t_start)?;
|
||||
|
||||
// Now we can use noise for logging since we cloned it above
|
||||
log_tensor_stats("initial_latents", &latents)?;
|
||||
log_tensor_stats("noise", &noise)?;
|
||||
log_tensor_stats("noised_latents", &noised_latents)?;
|
||||
|
||||
// Check the dtype of noised_latents
|
||||
debug_tensor("final_img2img_latents", &noised_latents)?;
|
||||
|
||||
// Explicitly convert to expected dtype to avoid errors
|
||||
noised_latents.to_dtype(dtype)?
|
||||
} else {
|
||||
// Regular text-to-image generation with random noise
|
||||
let latents = Tensor::randn(0f32, 1f32, (bsize, 4, sd_config.height / 8, sd_config.width / 8), &device)?;
|
||||
|
||||
// Use affine for multiplication
|
||||
let scaled_latents = latents.affine(scheduler.init_noise_sigma(), 0.0)?;
|
||||
debug_tensor("scaled_initial_latents", &scaled_latents)?;
|
||||
|
||||
// Explicitly convert to expected dtype
|
||||
scaled_latents.to_dtype(dtype)?
|
||||
};
|
||||
let mut latents = latents.to_dtype(dtype)?;
|
||||
|
||||
println!("Setting seed to {}", seed);
|
||||
device.set_seed(seed)?;
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Initial latents shape: {:?}", latents.shape());
|
||||
println!("Initial latents first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
println!("Starting LCM sampling");
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
for (timestep_index, ×tep) in timesteps.iter().enumerate() {
|
||||
if timestep_index < raw_t_start {
|
||||
continue;
|
||||
}
|
||||
|
||||
println!("\n--- Timestep {}/{} (value: {}) ---", timestep_index + 1, timesteps.len(), timestep);
|
||||
|
||||
let latent_model_input = if use_guide_scale {
|
||||
println!("Using classifier-free guidance with scale: {}", guidance_scale);
|
||||
Tensor::cat(&[&latents, &latents], 0)?
|
||||
} else {
|
||||
latents.clone()
|
||||
};
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Latent input shape before scaling: {:?}", latent_model_input.shape());
|
||||
println!("Latent input first few values: {}", latent_model_input.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
let latent_model_input = scheduler.scale_model_input(latent_model_input, timestep);
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Latent input shape after scaling: {:?}", latent_model_input.shape());
|
||||
println!("Scaled latent input first few values: {}", latent_model_input.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
// Generate guidance scale embedding
|
||||
let guidance_emb = if use_guide_scale {
|
||||
// For LCM/Dreamshaper, the time embedding dimension is 1280
|
||||
let time_embed_dim = 1280; // Standard for SD models (320 base channels * 4)
|
||||
|
||||
println!("Generating guidance scale embedding (scale={})", guidance_scale);
|
||||
Some(scheduler.get_guidance_scale_embedding(guidance_scale, time_embed_dim, &device, dtype)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Before UNet forward call
|
||||
debug_tensor("latent_input_to_unet", &latent_model_input)?;
|
||||
debug_tensor("text_embeddings_to_unet", &text_embeddings)?;
|
||||
|
||||
// Pass the guidance embedding to the UNet
|
||||
// This would require modifying the UNet forward method to accept this parameter
|
||||
println!("Running UNet forward pass with guidance embedding...");
|
||||
let noise_pred = unet.forward_with_guidance(&latent_model_input, timestep as f64, &text_embeddings, guidance_emb.as_ref())?;
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Noise prediction shape: {:?}", noise_pred.shape());
|
||||
println!("Noise prediction first few values: {}", noise_pred.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
let noise_pred = if use_guide_scale {
|
||||
let noise_pred = noise_pred.chunk(2, 0)?;
|
||||
let (noise_pred_uncond, noise_pred_text) = (&noise_pred[0], &noise_pred[1]);
|
||||
|
||||
// LCM requires a different guidance approach than standard diffusion
|
||||
// Linear interpolation with scale factor
|
||||
let diff = (noise_pred_text - noise_pred_uncond)?;
|
||||
let scaled_diff = diff.affine(guidance_scale, 0.0)?;
|
||||
(noise_pred_uncond + &scaled_diff)?
|
||||
} else {
|
||||
noise_pred
|
||||
};
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Latents before step first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
println!("Applying scheduler step...");
|
||||
latents = scheduler.step(&noise_pred, timestep, &latents)?;
|
||||
|
||||
if PRINT_DEBUG_TENSORS {
|
||||
println!("Latents after step first few values: {}", latents.i(0)?.i(0)?.narrow(0, 0, 5)?.narrow(1, 0, 5)?);
|
||||
}
|
||||
|
||||
println!("Step {}/{} completed", timestep_index + 1, n_steps);
|
||||
|
||||
// Then use it at key points
|
||||
log_tensor_stats("latents_before_step", &latents)?;
|
||||
log_tensor_stats("noise_pred", &noise_pred)?;
|
||||
log_tensor_stats("latents_after_step", &latents)?;
|
||||
|
||||
if save_tensors {
|
||||
let tensor_dir = "lcm_rust_tensors";
|
||||
std::fs::create_dir_all(tensor_dir)?;
|
||||
|
||||
// Save key tensors at each step
|
||||
let step_dir = format!("{}/step_{}", tensor_dir, timestep_index);
|
||||
std::fs::create_dir_all(&step_dir)?;
|
||||
|
||||
// Save various tensors as numpy arrays for later comparison
|
||||
save_tensor_as_numpy(&latent_model_input, &format!("{}/input.npy", step_dir))?;
|
||||
save_tensor_as_numpy(&noise_pred, &format!("{}/noise_pred.npy", step_dir))?;
|
||||
save_tensor_as_numpy(&latents, &format!("{}/output.npy", step_dir))?;
|
||||
}
|
||||
|
||||
// After UNet forward call
|
||||
debug_tensor("noise_pred_from_unet", &noise_pred)?;
|
||||
}
|
||||
|
||||
let dt = start_time.elapsed().as_secs_f32();
|
||||
println!("Sampling completed in {:.2}s", dt);
|
||||
|
||||
// Generate final image
|
||||
println!("Generating final image for sample {}/{}", idx + 1, num_samples);
|
||||
save_image(&vae, &latents, vae_scale, bsize, idx, &final_image, num_samples, None)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Add this debug function to check dtype at critical points
|
||||
fn debug_tensor(name: &str, tensor: &Tensor) -> Result<()> {
|
||||
println!("DEBUG {}: shape={:?}, dtype={:?}", name, tensor.shape(), tensor.dtype());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
run(args)
|
||||
}
|
||||
@@ -0,0 +1,762 @@
|
||||
use super::schedulers::{betas_for_alpha_bar, BetaSchedule, PredictionType, TimestepSpacing};
|
||||
use candle::{Result, Tensor, IndexOp, Device, DType, Shape};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum LCMVarianceType {
|
||||
FixedSmall,
|
||||
FixedSmallLog,
|
||||
FixedLarge,
|
||||
FixedLargeLog,
|
||||
Learned,
|
||||
}
|
||||
|
||||
impl Default for LCMVarianceType {
|
||||
fn default() -> Self {
|
||||
Self::FixedSmall
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LCMSolverType {
|
||||
Default,
|
||||
LogRho,
|
||||
}
|
||||
|
||||
impl Default for LCMSolverType {
|
||||
fn default() -> Self {
|
||||
Self::Default
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LCMSchedulerConfig {
|
||||
/// The value of beta at the beginning of training.
|
||||
pub beta_start: f64,
|
||||
/// The value of beta at the end of training.
|
||||
pub beta_end: f64,
|
||||
/// How beta evolved during training.
|
||||
pub beta_schedule: BetaSchedule,
|
||||
/// Option to predicted sample between -1 and 1 for numerical stability.
|
||||
pub clip_sample: bool,
|
||||
/// Option to clip the variance used when adding noise to the denoised sample.
|
||||
pub variance_type: LCMVarianceType,
|
||||
/// prediction type of the scheduler function
|
||||
pub prediction_type: PredictionType,
|
||||
/// number of diffusion steps used to train the model.
|
||||
pub train_timesteps: usize,
|
||||
/// time step spacing for the diffusion process
|
||||
pub timestep_spacing: TimestepSpacing,
|
||||
/// Original training steps
|
||||
pub original_inference_steps: usize,
|
||||
/// Timestep scaling factor (default: 10.0 in diffusers)
|
||||
pub timestep_scaling: f64,
|
||||
/// Adjust the indexes of the inference schedule by this value.
|
||||
pub steps_offset: usize,
|
||||
/// The solver type to use for the diffusion process
|
||||
pub solver_type: LCMSolverType,
|
||||
}
|
||||
|
||||
impl Default for LCMSchedulerConfig {
|
||||
fn default() -> Self {
|
||||
Self { beta_start: 0.00085, beta_end: 0.012, beta_schedule: BetaSchedule::ScaledLinear, clip_sample: false, variance_type: LCMVarianceType::FixedSmall, prediction_type: PredictionType::Epsilon, train_timesteps: 1000, timestep_spacing: TimestepSpacing::Leading, original_inference_steps: 50, timestep_scaling: 10.0, steps_offset: 1, solver_type: LCMSolverType::Default }
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LCMScheduler {
|
||||
alphas_cumprod: Vec<f64>,
|
||||
init_noise_sigma: f64,
|
||||
timesteps: Vec<usize>,
|
||||
step_ratio: usize,
|
||||
pub config: LCMSchedulerConfig,
|
||||
step_index: Option<usize>,
|
||||
begin_index: Option<usize>,
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
impl LCMScheduler {
|
||||
pub fn new(inference_steps: usize, config: LCMSchedulerConfig) -> Result<Self> {
|
||||
// Set up the correct betas
|
||||
let betas = match config.beta_schedule {
|
||||
BetaSchedule::ScaledLinear => {
|
||||
// LCM uses scaled_linear specifically this way
|
||||
let start = config.beta_start.sqrt();
|
||||
let end = config.beta_end.sqrt();
|
||||
let betas = super::utils::linspace(start, end, config.train_timesteps)?;
|
||||
betas.sqr()?
|
||||
},
|
||||
// Other schedules...
|
||||
_ => unreachable!("LCM needs ScaledLinear"),
|
||||
};
|
||||
|
||||
// Set up alphas exactly as in diffusers
|
||||
let betas = betas.to_vec1::<f64>()?;
|
||||
let alphas = betas.iter().map(|beta| 1.0 - beta).collect::<Vec<_>>();
|
||||
|
||||
let mut alphas_cumprod = Vec::with_capacity(alphas.len());
|
||||
let mut alpha_prod = 1.0;
|
||||
for alpha in alphas.iter() {
|
||||
alpha_prod *= alpha;
|
||||
alphas_cumprod.push(alpha_prod);
|
||||
}
|
||||
|
||||
// Create timesteps based on TimestepSpacing, similar to DDIM
|
||||
let step_ratio = config.train_timesteps / inference_steps;
|
||||
let timesteps: Vec<usize> = match config.timestep_spacing {
|
||||
TimestepSpacing::Leading => (0..(inference_steps)).map(|s| s * step_ratio + config.steps_offset).rev().collect(),
|
||||
TimestepSpacing::Trailing => std::iter::successors(Some(config.train_timesteps), |n| if *n > step_ratio { Some(n - step_ratio) } else { None }).map(|n| n - 1).collect(),
|
||||
TimestepSpacing::Linspace => super::utils::linspace(0.0, (config.train_timesteps - 1) as f64, inference_steps)?.to_vec1::<f64>()?.iter().map(|&f| f as usize).rev().collect(),
|
||||
};
|
||||
|
||||
println!("Using LCM timesteps with {:?} spacing: {:?}", config.timestep_spacing, timesteps);
|
||||
|
||||
let init_noise_sigma = 1.0;
|
||||
|
||||
Ok(Self { alphas_cumprod, init_noise_sigma, timesteps, step_ratio, config, step_index: None, begin_index: None, dtype: DType::F32 })
|
||||
}
|
||||
|
||||
pub fn with_dtype(mut self, dtype: DType) -> Self {
|
||||
self.dtype = dtype;
|
||||
self
|
||||
}
|
||||
|
||||
fn get_variance(&self, timestep: usize) -> f64 {
|
||||
let prev_t = timestep as isize - self.step_ratio as isize;
|
||||
let alpha_prod_t = self.alphas_cumprod[timestep];
|
||||
let alpha_prod_t_prev = if prev_t >= 0 { self.alphas_cumprod[prev_t as usize] } else { 1.0 };
|
||||
let current_beta_t = 1. - alpha_prod_t / alpha_prod_t_prev;
|
||||
|
||||
// For t > 0, compute predicted variance βt (see formula (6) and (7) from [the pdf](https://arxiv.org/pdf/2006.11239.pdf))
|
||||
// and sample from it to get previous sample
|
||||
// x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
|
||||
let variance = (1. - alpha_prod_t_prev) / (1. - alpha_prod_t) * current_beta_t;
|
||||
|
||||
// retrieve variance
|
||||
match self.config.variance_type {
|
||||
LCMVarianceType::FixedSmall => variance.max(1e-20),
|
||||
// for rl-diffuser https://arxiv.org/abs/2205.09991
|
||||
LCMVarianceType::FixedSmallLog => {
|
||||
let variance = variance.max(1e-20).ln();
|
||||
(variance * 0.5).exp()
|
||||
},
|
||||
LCMVarianceType::FixedLarge => current_beta_t,
|
||||
LCMVarianceType::FixedLargeLog => current_beta_t.ln(),
|
||||
LCMVarianceType::Learned => variance,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn timesteps(&self) -> &[usize] {
|
||||
self.timesteps.as_slice()
|
||||
}
|
||||
|
||||
/// Ensures interchangeability with schedulers that need to scale the denoising model input
|
||||
/// depending on the current timestep.
|
||||
pub fn scale_model_input(&self, sample: Tensor, timestep: usize) -> Tensor {
|
||||
// LCM doesn't apply any special scaling to model inputs
|
||||
println!("LCM: No scaling applied to model input for timestep {}", timestep);
|
||||
sample
|
||||
}
|
||||
|
||||
fn get_scalings_for_boundary_condition_discrete(&self, timestep: usize) -> (f64, f64) {
|
||||
// LCM uses a fixed sigma_data value of 0.5
|
||||
let sigma_data = 0.5f64;
|
||||
|
||||
// Timestep scaling factor from diffusers configuration
|
||||
let timestep_scaling = self.config.timestep_scaling;
|
||||
|
||||
// In diffusers, they don't use the timestep value directly
|
||||
let alpha_cumprod = self.alphas_cumprod[timestep];
|
||||
let sqrt_alpha_prod = alpha_cumprod.sqrt();
|
||||
let sqrt_one_minus_alpha_prod = (1.0 - alpha_cumprod).sqrt();
|
||||
|
||||
// This is the key calculation - convert alpha to sigma space with scaling
|
||||
let scaled_timestep = sqrt_one_minus_alpha_prod / sqrt_alpha_prod * timestep_scaling;
|
||||
|
||||
println!("Boundary condition for timestep={}: scaled_t={} (with scaling={})", timestep, scaled_timestep, timestep_scaling);
|
||||
|
||||
// Calculate boundary condition coefficients with scaled timestep
|
||||
let c_skip = sigma_data.powf(2.0) / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0));
|
||||
let c_out = scaled_timestep / (scaled_timestep.powf(2.0) + sigma_data.powf(2.0)).sqrt();
|
||||
|
||||
(c_skip, c_out)
|
||||
}
|
||||
|
||||
pub fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> Result<Tensor> {
|
||||
// Initialize step_index if it's not already set
|
||||
if self.step_index.is_none() {
|
||||
self.init_step_index(timestep)?;
|
||||
println!("Initialized step_index to {} for timestep {}", self.step_index.unwrap(), timestep);
|
||||
}
|
||||
|
||||
// Get current timestep index
|
||||
let timestep_index = self.step_index.unwrap();
|
||||
|
||||
// Calculate prev_timestep_index
|
||||
let prev_timestep_index = timestep_index + 1;
|
||||
let prev_timestep = if prev_timestep_index >= self.timesteps.len() {
|
||||
0 // Final step
|
||||
} else {
|
||||
self.timesteps[prev_timestep_index]
|
||||
};
|
||||
|
||||
// Update step_index for next time
|
||||
self.step_index = Some(prev_timestep_index);
|
||||
|
||||
// Choose solver based on configuration
|
||||
match self.config.solver_type {
|
||||
LCMSolverType::Default => self.step_default(model_output, timestep, prev_timestep, sample),
|
||||
LCMSolverType::LogRho => self.step_logrho(model_output, timestep, prev_timestep, sample),
|
||||
}
|
||||
}
|
||||
|
||||
// Default solver (existing logic moved to a separate method)
|
||||
fn step_default(&self, model_output: &Tensor, timestep: usize, prev_timestep: usize, sample: &Tensor) -> Result<Tensor> {
|
||||
// Get dtype from sample and use it for consistency
|
||||
let target_dtype = sample.dtype();
|
||||
let device = sample.device();
|
||||
|
||||
println!("Step (Default): current_t={}, prev_t={}", timestep, prev_timestep);
|
||||
|
||||
// Convert model output to match sample dtype
|
||||
let model_output = model_output.to_dtype(target_dtype)?;
|
||||
|
||||
// Get the exact alpha values for this timestep
|
||||
let alpha_prod_t = if timestep < self.alphas_cumprod.len() {
|
||||
self.alphas_cumprod[timestep]
|
||||
} else {
|
||||
println!("Warning: Timestep {} out of range, using last alpha", timestep);
|
||||
*self.alphas_cumprod.last().unwrap()
|
||||
};
|
||||
|
||||
// Get the previous alpha value
|
||||
let alpha_prod_t_prev = if prev_timestep == 0 {
|
||||
1.0 // Final step goes to alpha=1
|
||||
} else if prev_timestep < self.alphas_cumprod.len() {
|
||||
self.alphas_cumprod[prev_timestep]
|
||||
} else {
|
||||
println!("Warning: Prev timestep {} out of range", prev_timestep);
|
||||
*self.alphas_cumprod.last().unwrap()
|
||||
};
|
||||
|
||||
// Calculate coefficients for boundary conditions
|
||||
let (c_skip, c_out) = self.get_scalings_for_boundary_condition_discrete(timestep);
|
||||
|
||||
// Compute predicted original sample from noise prediction
|
||||
let pred_original_sample = match self.config.prediction_type {
|
||||
PredictionType::Epsilon => {
|
||||
let beta_prod_t = 1.0 - alpha_prod_t;
|
||||
let beta_prod_t_sqrt = beta_prod_t.sqrt();
|
||||
|
||||
let scaled_noise = model_output.affine(beta_prod_t_sqrt, 0.0)?;
|
||||
let diff = (sample - &scaled_noise)?;
|
||||
diff.affine(1.0 / alpha_prod_t.sqrt(), 0.0)?
|
||||
},
|
||||
PredictionType::Sample => model_output.clone(),
|
||||
PredictionType::VPrediction => {
|
||||
let sqrt_alpha_prod = alpha_prod_t.sqrt();
|
||||
let sqrt_one_minus_alpha_prod = (1.0 - alpha_prod_t).sqrt();
|
||||
|
||||
let term1 = sample.affine(sqrt_alpha_prod, 0.0)?;
|
||||
let term2 = model_output.affine(sqrt_one_minus_alpha_prod, 0.0)?;
|
||||
(term1 - term2)?
|
||||
},
|
||||
_ => unreachable!(),
|
||||
};
|
||||
|
||||
// Apply possible clipping
|
||||
let pred_orig = if self.config.clip_sample { pred_original_sample.clamp(-1f32, 1f32)? } else { pred_original_sample };
|
||||
|
||||
// Apply boundary condition
|
||||
let term1 = pred_orig.affine(c_out, 0.0)?;
|
||||
let term2 = sample.affine(c_skip, 0.0)?;
|
||||
let denoised = (&term1 + &term2)?;
|
||||
|
||||
// Set up the previous sample
|
||||
let pred_prev_sample = if prev_timestep == 0 {
|
||||
denoised.clone()
|
||||
} else {
|
||||
let prev_sqrt_alpha = alpha_prod_t_prev.sqrt();
|
||||
denoised.affine(prev_sqrt_alpha, 0.0)?
|
||||
};
|
||||
|
||||
Ok(pred_prev_sample)
|
||||
}
|
||||
|
||||
// LogRho solver implementation
|
||||
fn step_logrho(&self, model_output: &Tensor, timestep: usize, prev_timestep: usize, sample: &Tensor) -> Result<Tensor> {
|
||||
// Ensure consistent dtype by getting it from the sample
|
||||
let target_dtype = sample.dtype();
|
||||
let device = sample.device();
|
||||
|
||||
println!("Step (LogRho): current_t={}, prev_t={}", timestep, prev_timestep);
|
||||
println!(" Using target dtype: {:?}", target_dtype);
|
||||
|
||||
// Convert model output to match sample dtype
|
||||
let model_output = model_output.to_dtype(target_dtype)?;
|
||||
println!(" model_output dtype after conversion: {:?}", model_output.dtype());
|
||||
|
||||
// Get alpha values
|
||||
let alpha_prod_t = if timestep < self.alphas_cumprod.len() { self.alphas_cumprod[timestep] } else { *self.alphas_cumprod.last().unwrap() };
|
||||
|
||||
let alpha_prod_t_prev = if prev_timestep == 0 {
|
||||
1.0
|
||||
} else if prev_timestep < self.alphas_cumprod.len() {
|
||||
self.alphas_cumprod[prev_timestep]
|
||||
} else {
|
||||
*self.alphas_cumprod.last().unwrap()
|
||||
};
|
||||
|
||||
// Special case: if we're at the last timestep, return the predicted sample directly
|
||||
if prev_timestep == 0 {
|
||||
// Handle final step differently - exact same implementation as in default solver
|
||||
return self.step_default(&model_output, timestep, prev_timestep, sample);
|
||||
}
|
||||
|
||||
// Calculate sigma values for LogRho solver with numerical stability
|
||||
let sigma_t = ((1.0 - alpha_prod_t).max(1e-12)).sqrt();
|
||||
let sigma_t_prev = ((1.0 - alpha_prod_t_prev).max(1e-12)).sqrt();
|
||||
|
||||
// Calculate key logrho parameters with numerical stability
|
||||
let logrho_t = sigma_t.ln() - alpha_prod_t.sqrt().max(1e-12).ln();
|
||||
let logrho_t_prev = sigma_t_prev.ln() - alpha_prod_t_prev.sqrt().max(1e-12).ln();
|
||||
|
||||
// Calculate h parameter (step size in log rho space)
|
||||
let h = logrho_t_prev - logrho_t;
|
||||
let exp_h = h.exp().clamp(0.01, 5.0); // Constrain to a reasonable range
|
||||
println!(" LogRho params: h={:.6}, exp(h)={:.6}", h, exp_h);
|
||||
|
||||
// For epsilon prediction we need to transform the model output
|
||||
let score = match self.config.prediction_type {
|
||||
PredictionType::Epsilon => {
|
||||
// For ε prediction type, score is -σ/√α_t * ε
|
||||
let scale = -sigma_t / alpha_prod_t.sqrt().max(1e-12);
|
||||
model_output.affine(scale, 0.0)?.to_dtype(target_dtype)?
|
||||
},
|
||||
PredictionType::VPrediction => {
|
||||
// model_output is already in the right form
|
||||
model_output.clone()
|
||||
},
|
||||
PredictionType::Sample => {
|
||||
// Score = (x_t - √α_t * x_0) / σ_t
|
||||
let alpha_sqrt = alpha_prod_t.sqrt();
|
||||
let scaled_pred = model_output.affine(alpha_sqrt, 0.0)?.to_dtype(target_dtype)?;
|
||||
let diff = (sample.clone() - scaled_pred)?.to_dtype(target_dtype)?;
|
||||
diff.affine(1.0 / sigma_t, 0.0)?.to_dtype(target_dtype)?
|
||||
},
|
||||
_ => unreachable!(),
|
||||
};
|
||||
|
||||
println!(" score dtype: {:?}", score.dtype());
|
||||
|
||||
// Calculate d term: d = x_t + σ_t^2/√α_t * score
|
||||
let scale_d = sigma_t.powi(2) / alpha_prod_t.sqrt().max(1e-12);
|
||||
let score_scaled = score.affine(scale_d, 0.0)?.to_dtype(target_dtype)?;
|
||||
let d = (sample + score_scaled)?.to_dtype(target_dtype)?;
|
||||
println!(" d dtype: {:?}", d.dtype());
|
||||
|
||||
// Apply the main update formula directly:
|
||||
// x_prev = exp_h * sample + sigma_prev * (1 - exp_h) * d
|
||||
|
||||
// First part: exp_h * sample
|
||||
let term1 = sample.affine(exp_h, 0.0)?.to_dtype(target_dtype)?;
|
||||
println!(" term1 dtype: {:?}", term1.dtype());
|
||||
|
||||
// Second part: sigma_prev * (1 - exp_h) * d
|
||||
let term2_factor = sigma_t_prev * (1.0 - exp_h);
|
||||
println!(" term2_factor: {:.6}", term2_factor);
|
||||
|
||||
let term2 = d.affine(term2_factor, 0.0)?.to_dtype(target_dtype)?;
|
||||
println!(" term2 dtype: {:?}", term2.dtype());
|
||||
|
||||
// Add the terms together
|
||||
let pred_prev_sample = (&term1 + &term2)?.to_dtype(target_dtype)?;
|
||||
println!(" pred_prev_sample dtype: {:?}", pred_prev_sample.dtype());
|
||||
|
||||
Ok(pred_prev_sample)
|
||||
}
|
||||
|
||||
pub fn add_noise(&self, original_samples: &Tensor, noise: Tensor, timestep: usize) -> Result<Tensor> {
|
||||
// Get the correct alpha value
|
||||
let alpha_cumprod = self.alphas_cumprod[timestep];
|
||||
|
||||
println!("Adding noise for timestep={}, alpha={:.6}, sqrt(alpha)={:.6}", timestep, alpha_cumprod, alpha_cumprod.sqrt());
|
||||
println!(" original_samples: shape={:?}, dtype={:?}", original_samples.shape(), original_samples.dtype());
|
||||
println!(" noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
|
||||
|
||||
// Keep operations in original_samples' dtype for now
|
||||
let target_dtype = original_samples.dtype();
|
||||
let noise = noise.to_dtype(target_dtype)?;
|
||||
|
||||
// Calculate coefficients
|
||||
let sqrt_alpha = alpha_cumprod.sqrt();
|
||||
let sqrt_one_minus_alpha = (1.0 - alpha_cumprod).sqrt();
|
||||
|
||||
// Apply noise formula
|
||||
let scaled_samples = original_samples.affine(sqrt_alpha, 0.0)?;
|
||||
let scaled_noise = noise.affine(sqrt_one_minus_alpha, 0.0)?;
|
||||
|
||||
println!(" scaled_samples: shape={:?}, dtype={:?}", scaled_samples.shape(), scaled_samples.dtype());
|
||||
println!(" scaled_noise: shape={:?}, dtype={:?}", scaled_noise.shape(), scaled_noise.dtype());
|
||||
|
||||
// Add scaled noise - should maintain dtype
|
||||
let result = scaled_samples.add(&scaled_noise)?;
|
||||
|
||||
println!(" result: shape={:?}, dtype={:?}", result.shape(), result.dtype());
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn init_noise_sigma(&self) -> f64 {
|
||||
self.init_noise_sigma
|
||||
}
|
||||
|
||||
pub fn alphas_cumprod(&self) -> &[f64] {
|
||||
&self.alphas_cumprod
|
||||
}
|
||||
|
||||
pub fn log_config(&self) {
|
||||
println!("LCM Scheduler Configuration:");
|
||||
println!(" beta_start: {}", self.config.beta_start);
|
||||
println!(" beta_end: {}", self.config.beta_end);
|
||||
println!(" beta_schedule: {:?}", self.config.beta_schedule);
|
||||
println!(" train_timesteps: {}", self.config.train_timesteps);
|
||||
println!(" prediction_type: {:?}", self.config.prediction_type);
|
||||
println!(" clip_sample: {}", self.config.clip_sample);
|
||||
println!(" variance_type: {:?}", self.config.variance_type);
|
||||
println!(" timesteps: {:?}", self.timesteps);
|
||||
// Print the first few alphas_cumprod values
|
||||
let num_to_show = std::cmp::min(10, self.alphas_cumprod.len());
|
||||
println!(" First {} alphas_cumprod: {:?}", num_to_show, self.alphas_cumprod.iter().take(num_to_show).collect::<Vec<_>>());
|
||||
}
|
||||
|
||||
/// Creates embeddings from guidance scale values, similar to timestep embeddings
|
||||
/// Based on diffusers implementation
|
||||
pub fn get_guidance_scale_embedding(&self, guidance_scale: f64, embedding_dim: usize, device: &Device, dtype: DType) -> Result<Tensor> {
|
||||
// Scale by exactly 1000.0 to match diffusers
|
||||
let w = guidance_scale * 1000.0;
|
||||
println!("Creating guidance embedding with dim={}, scaled value={}", embedding_dim, w);
|
||||
|
||||
// Create embedding vector directly to avoid tensor operations
|
||||
let mut embedding = Vec::with_capacity(embedding_dim);
|
||||
let half_dim = embedding_dim / 2;
|
||||
|
||||
// Compute the position frequencies
|
||||
for i in 0..half_dim {
|
||||
let freq = (i as f32) * (-((10000.0f32).ln()) / ((half_dim - 1) as f32));
|
||||
let freq = freq.exp();
|
||||
|
||||
// Apply the guidance scale (w value)
|
||||
let value = w as f32 * freq;
|
||||
|
||||
// Add sin and cos values
|
||||
embedding.push(value.sin());
|
||||
embedding.push(value.cos());
|
||||
}
|
||||
|
||||
// If odd dimension, pad with zero
|
||||
if embedding.len() < embedding_dim {
|
||||
embedding.push(0.0);
|
||||
}
|
||||
|
||||
// Ensure we have the correct length
|
||||
assert_eq!(embedding.len(), embedding_dim, "Embedding length mismatch");
|
||||
|
||||
// Create tensor from the embedding vector
|
||||
let tensor = Tensor::new(embedding.as_slice(), device)?.to_dtype(dtype)?;
|
||||
|
||||
// Reshape to match expected [1, embedding_dim] shape
|
||||
let result = tensor.reshape((1, embedding_dim))?;
|
||||
|
||||
println!("Final guidance embedding shape: {:?}", result.shape());
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Calculate timesteps based on LCM's specific formula
|
||||
/// This is now only used for backward compatibility or custom timesteps
|
||||
pub fn get_timesteps_for_steps(inference_steps: usize) -> Vec<usize> {
|
||||
match inference_steps {
|
||||
// 4 => vec![999, 499, 259, 0],
|
||||
4 => vec![499, 259],
|
||||
8 => vec![999, 879, 759, 639, 499, 379, 259, 139, 0],
|
||||
_ => {
|
||||
let intervals = (999 / (inference_steps - 1)) as usize;
|
||||
let mut timesteps = Vec::with_capacity(inference_steps);
|
||||
|
||||
// First step always starts at 999
|
||||
timesteps.push(999);
|
||||
|
||||
// Generate intermediate steps
|
||||
for i in 1..inference_steps {
|
||||
// The pattern looks approximately linear with a step size determined by the number of steps
|
||||
let ts = 999 - i * intervals;
|
||||
timesteps.push(ts);
|
||||
}
|
||||
|
||||
// Last step is always 0
|
||||
if let Some(last) = timesteps.last_mut() {
|
||||
if *last != 0 {
|
||||
*last = 0;
|
||||
}
|
||||
}
|
||||
|
||||
timesteps
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the appropriate starting timestep based on strength
|
||||
pub fn get_timestep_for_strength(&self, strength: f64) -> (usize, f64) {
|
||||
// Convert strength to a timestep index
|
||||
let num_steps = self.timesteps.len();
|
||||
let strength_adjusted = (strength * num_steps as f64).floor() as usize;
|
||||
let start_idx = num_steps.saturating_sub(strength_adjusted + 1);
|
||||
|
||||
// Get the timestep at this index
|
||||
let timestep = if start_idx < self.timesteps.len() { self.timesteps[start_idx] } else { *self.timesteps.first().unwrap_or(&999) };
|
||||
|
||||
// Get the corresponding alpha
|
||||
let alpha = if timestep < self.alphas_cumprod.len() { self.alphas_cumprod[timestep] } else { *self.alphas_cumprod.last().unwrap_or(&0.0) };
|
||||
|
||||
println!("For strength {:.2}, using timestep {} (index {}), alpha={:.6}", strength, timestep, start_idx, alpha);
|
||||
|
||||
(timestep, alpha)
|
||||
}
|
||||
|
||||
/// Add noise for img2img based on strength
|
||||
pub fn add_noise_for_img2img(&mut self, original_samples: &Tensor, noise: &Tensor, strength: f64) -> Result<(Tensor, usize)> {
|
||||
// Get starting timestep and alpha
|
||||
let (timestep, alpha) = self.get_timestep_for_strength(strength);
|
||||
|
||||
// Set begin_index based on this timestep
|
||||
let start_idx = self.index_for_timestep(timestep)?;
|
||||
self.set_begin_index(start_idx);
|
||||
|
||||
println!("Adding img2img noise with strength={:.4}, timestep={}, alpha={:.6}", strength, timestep, alpha);
|
||||
println!(" original_samples: shape={:?}, dtype={:?}", original_samples.shape(), original_samples.dtype());
|
||||
println!(" noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
|
||||
|
||||
// Keep operations in original_samples' dtype for now
|
||||
let target_dtype = original_samples.dtype();
|
||||
let noise_converted = noise.to_dtype(target_dtype)?;
|
||||
|
||||
// Calculate coefficients
|
||||
let sqrt_alpha = alpha.sqrt();
|
||||
let sqrt_one_minus_alpha = (1.0 - alpha).sqrt();
|
||||
|
||||
println!(" Using scalar multiplication with target dtype={:?}", target_dtype);
|
||||
|
||||
// Apply noise formula
|
||||
let scaled_samples = original_samples.affine(sqrt_alpha, 0.0)?;
|
||||
let scaled_noise = noise_converted.affine(sqrt_one_minus_alpha, 0.0)?;
|
||||
|
||||
println!(" scaled_samples: shape={:?}, dtype={:?}", scaled_samples.shape(), scaled_samples.dtype());
|
||||
println!(" scaled_noise: shape={:?}, dtype={:?}", scaled_noise.shape(), scaled_noise.dtype());
|
||||
|
||||
// Add scaled noise - should maintain dtype
|
||||
let noised = scaled_samples.add(&scaled_noise)?;
|
||||
|
||||
println!(" noised: shape={:?}, dtype={:?}", noised.shape(), noised.dtype());
|
||||
|
||||
// Return noised image and the timestep to start from
|
||||
Ok((noised, timestep))
|
||||
}
|
||||
|
||||
pub fn index_for_timestep(&self, timestep: usize) -> Result<usize> {
|
||||
// Collect all indices where timestep matches
|
||||
let indices: Vec<usize> = self.timesteps.iter().enumerate().filter_map(|(idx, &ts)| if ts == timestep { Some(idx) } else { None }).collect();
|
||||
|
||||
if indices.is_empty() {
|
||||
return Err(candle::Error::Msg(format!("Timestep {} not found in scheduler timesteps", timestep)));
|
||||
}
|
||||
|
||||
// Take the second index if there are multiple matches, otherwise take the first
|
||||
let pos = if indices.len() > 1 { 1 } else { 0 };
|
||||
Ok(indices[pos])
|
||||
}
|
||||
|
||||
pub fn set_begin_index(&mut self, begin_index: usize) {
|
||||
self.begin_index = Some(begin_index);
|
||||
println!("Set begin_index to {}", begin_index);
|
||||
}
|
||||
|
||||
pub fn init_step_index(&mut self, timestep: usize) -> Result<usize> {
|
||||
self.step_index = if let Some(begin_index) = self.begin_index {
|
||||
// If begin_index is set, use it directly
|
||||
Some(begin_index)
|
||||
} else {
|
||||
// Otherwise find the index based on the timestep
|
||||
// Find the position of timestep in self.timesteps
|
||||
let pos = self.timesteps.iter().position(|&t| t == timestep).ok_or_else(|| candle::Error::Msg(format!("Timestep {} not found in timesteps", timestep)))?;
|
||||
Some(pos)
|
||||
};
|
||||
|
||||
Ok(self.step_index.unwrap())
|
||||
}
|
||||
|
||||
// Updated helper function to match diffusers' get_timesteps logic exactly
|
||||
pub fn retrieve_timesteps(scheduler: &mut LCMScheduler, num_inference_steps: usize, strength: Option<f64>, custom_timesteps: Option<Vec<usize>>) -> Result<Vec<usize>> {
|
||||
// 1. When custom timesteps are provided, use them directly
|
||||
if let Some(custom_ts) = custom_timesteps {
|
||||
println!("Using custom timesteps: {:?}", custom_ts);
|
||||
return Ok(custom_ts);
|
||||
}
|
||||
|
||||
// 2. Get timesteps from scheduler
|
||||
let mut timesteps = scheduler.timesteps().to_vec();
|
||||
|
||||
// 3. For img2img, adjust timesteps based on strength
|
||||
if let Some(strength) = strength {
|
||||
if strength < 1.0 {
|
||||
// Match diffusers logic exactly:
|
||||
// init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
|
||||
let init_timestep = ((num_inference_steps as f64 * strength).floor() as usize).min(num_inference_steps);
|
||||
|
||||
// t_start = max(num_inference_steps - init_timestep, 0)
|
||||
let t_start = num_inference_steps.saturating_sub(init_timestep);
|
||||
|
||||
// In diffusers, order is typically 1 for LCM, so we'll simplify here
|
||||
let scheduler_order = 1;
|
||||
let start_idx = t_start * scheduler_order;
|
||||
|
||||
// Slice timesteps from the calculated start position
|
||||
if start_idx < timesteps.len() {
|
||||
println!("Adjusting timesteps for strength {}: starting at index {} (init_timestep={})", strength, start_idx, init_timestep);
|
||||
timesteps = timesteps[start_idx..].to_vec();
|
||||
|
||||
// Set begin_index in the scheduler
|
||||
scheduler.set_begin_index(start_idx);
|
||||
}
|
||||
|
||||
// Return the actual number of inference steps we'll perform
|
||||
let actual_steps = num_inference_steps.saturating_sub(t_start);
|
||||
println!("Will perform {} inference steps due to strength {}", actual_steps, strength);
|
||||
}
|
||||
}
|
||||
|
||||
println!("Final timesteps for inference: {:?}", timesteps);
|
||||
Ok(timesteps)
|
||||
}
|
||||
|
||||
// Add this method to help initialize latents from an image (for img2img)
|
||||
pub fn prepare_img2img_latents(
|
||||
&mut self,
|
||||
image_tensor: &Tensor, // Encoded image from VAE
|
||||
timestep: usize,
|
||||
batch_size: usize,
|
||||
num_images_per_prompt: usize,
|
||||
dtype: DType,
|
||||
device: &Device,
|
||||
) -> Result<Tensor> {
|
||||
// Store the dtype for consistent operations
|
||||
self.dtype = dtype;
|
||||
|
||||
let latent_dim = image_tensor.dim(2)? / 8; // Assuming standard VAE downsampling
|
||||
|
||||
println!("Preparing img2img latents for timestep={}", timestep);
|
||||
println!(" image_tensor: shape={:?}, dtype={:?}", image_tensor.shape(), image_tensor.dtype());
|
||||
|
||||
// Create noise tensor using the same approach as in main.rs
|
||||
let noise_shape = Shape::from((batch_size * num_images_per_prompt, 4, latent_dim, latent_dim));
|
||||
let noise = Tensor::randn(0f32, 1f32, noise_shape, device)?.to_dtype(dtype)?;
|
||||
println!(" noise: shape={:?}, dtype={:?}", noise.shape(), noise.dtype());
|
||||
|
||||
// Encode image to latents
|
||||
let encoded_latents = image_tensor.clone(); // Assuming already encoded
|
||||
|
||||
// Scale to match VAE expected scale - use affine for scalar multiplication
|
||||
println!(" Using affine for VAE scaling by 0.18215");
|
||||
let scaled_latents = encoded_latents.affine(0.18215, 0.0)?;
|
||||
println!(" scaled_latents: shape={:?}, dtype={:?}", scaled_latents.shape(), scaled_latents.dtype());
|
||||
|
||||
// Add noise based on timestep (this handles the "strength" parameter)
|
||||
let noised_latents = self.add_noise(&scaled_latents, noise, timestep)?;
|
||||
println!(" noised_latents: shape={:?}, dtype={:?}", noised_latents.shape(), noised_latents.dtype());
|
||||
|
||||
Ok(noised_latents)
|
||||
}
|
||||
|
||||
/// Apply classifier-free guidance to model output
|
||||
pub fn apply_guidance(&self, model_output_uncond: &Tensor, model_output_text: &Tensor, guidance_scale: f64) -> Result<Tensor> {
|
||||
// Skip guidance if scale is 1.0 (or very close to it)
|
||||
if (guidance_scale - 1.0).abs() < 1e-5 {
|
||||
return Ok(model_output_text.clone());
|
||||
}
|
||||
|
||||
println!("Applying guidance with scale={:.4}", guidance_scale);
|
||||
println!(" model_output_uncond: shape={:?}, dtype={:?}", model_output_uncond.shape(), model_output_uncond.dtype());
|
||||
println!(" model_output_text: shape={:?}, dtype={:?}", model_output_text.shape(), model_output_text.dtype());
|
||||
|
||||
// Apply the classifier-free guidance formula:
|
||||
// output = uncond + guidance_scale * (text - uncond)
|
||||
|
||||
// First calculate (text - uncond)
|
||||
let guidance_diff = (model_output_text - model_output_uncond)?;
|
||||
println!(" guidance_diff: shape={:?}, dtype={:?}", guidance_diff.shape(), guidance_diff.dtype());
|
||||
|
||||
// Then scale by guidance_scale using affine for scalar multiplication
|
||||
let guidance_factor = guidance_scale - 1.0;
|
||||
println!(" Using affine for scaling by {:.4}", guidance_factor);
|
||||
let scaled_diff = guidance_diff.affine(guidance_factor, 0.0)?;
|
||||
println!(" scaled_diff: shape={:?}, dtype={:?}", scaled_diff.shape(), scaled_diff.dtype());
|
||||
|
||||
// Finally add back to uncond: uncond + guidance_scale * (text - uncond)
|
||||
let guided_output = (model_output_uncond + &scaled_diff)?;
|
||||
println!(" guided_output: shape={:?}, dtype={:?}", guided_output.shape(), guided_output.dtype());
|
||||
|
||||
Ok(guided_output)
|
||||
}
|
||||
|
||||
// Add this helper method to manage dtype consistency across the pipeline
|
||||
pub fn ensure_dtype_consistency(&mut self, input_tensor: &Tensor) -> Result<()> {
|
||||
// Store the current tensor's dtype for reference
|
||||
self.dtype = input_tensor.dtype();
|
||||
println!("Setting LCM scheduler dtype to match input: {:?}", self.dtype);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Add a new method to prepare latents with explicit dtype control
|
||||
pub fn prepare_latents(&mut self, batch_size: usize, height: usize, width: usize, num_channels_latents: usize, dtype: DType, device: &Device, generator: Option<u64>) -> Result<Tensor> {
|
||||
// Store the requested dtype
|
||||
self.dtype = dtype;
|
||||
|
||||
// Calculate latent dimensions
|
||||
let latent_height = height / 8;
|
||||
let latent_width = width / 8;
|
||||
|
||||
// Create random latents
|
||||
let shape = Shape::from((batch_size, num_channels_latents, latent_height, latent_width));
|
||||
|
||||
// Note: Candle doesn't have a direct way to set random seeds for tensors
|
||||
// So we'll just use regular random generation for now
|
||||
let latents = if let Some(seed) = generator {
|
||||
println!("Note: Using random latents (seed not supported in this implementation)");
|
||||
// In a real implementation, you would want to set the random seed here
|
||||
Tensor::randn(0f32, 1f32, shape, device)?
|
||||
} else {
|
||||
println!("Generating random latents");
|
||||
Tensor::randn(0f32, 1f32, shape, device)?
|
||||
};
|
||||
|
||||
// Convert to the requested dtype
|
||||
let latents = latents.to_dtype(dtype)?;
|
||||
println!("Generated latents: shape={:?}, dtype={:?}", latents.shape(), latents.dtype());
|
||||
|
||||
// Scale by the init_noise_sigma
|
||||
let scaled_latents = latents.affine(self.init_noise_sigma, 0.0)?;
|
||||
|
||||
Ok(scaled_latents)
|
||||
}
|
||||
|
||||
// Add this method to maintain dtype consistency across operations
|
||||
pub fn match_dtype(&self, tensor: &Tensor, dtype: Option<DType>) -> Result<Tensor> {
|
||||
match dtype {
|
||||
Some(d) => tensor.to_dtype(d),
|
||||
None => Ok(tensor.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
// Add this public getter for the dtype field
|
||||
pub fn dtype(&self) -> DType {
|
||||
self.dtype
|
||||
}
|
||||
}
|
||||
+25
-1
@@ -41,6 +41,7 @@ pub mod attention;
|
||||
pub mod clip;
|
||||
pub mod ddim;
|
||||
pub mod ddpm;
|
||||
pub mod lcm;
|
||||
pub mod embeddings;
|
||||
pub mod euler_ancestral_discrete;
|
||||
pub mod resnet;
|
||||
@@ -52,8 +53,9 @@ pub mod utils;
|
||||
pub mod vae;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
|
||||
use candle::{DType, Device, Result};
|
||||
use candle::{DType, Device, Result, Tensor};
|
||||
use candle_nn as nn;
|
||||
|
||||
use self::schedulers::{Scheduler, SchedulerConfig};
|
||||
@@ -494,6 +496,28 @@ impl StableDiffusionConfig {
|
||||
pub fn build_scheduler(&self, n_steps: usize) -> Result<Box<dyn Scheduler>> {
|
||||
self.scheduler.build(n_steps)
|
||||
}
|
||||
|
||||
fn unet_config(&self, in_channels: usize) -> Result<unet_2d::UNet2DConditionModelConfig> {
|
||||
Ok(self.unet.clone())
|
||||
}
|
||||
|
||||
pub fn build_unet_with_weights(
|
||||
&self,
|
||||
weights: HashMap<String, Tensor>,
|
||||
device: &Device,
|
||||
in_channels: usize,
|
||||
use_flash_attn: bool,
|
||||
dtype: DType,
|
||||
) -> Result<unet_2d::UNet2DConditionModel> {
|
||||
let config = self.unet.clone();
|
||||
unet_2d::UNet2DConditionModel::new(
|
||||
nn::VarBuilder::from_tensors(weights, dtype, device),
|
||||
in_channels,
|
||||
4,
|
||||
use_flash_attn,
|
||||
config,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_clip_transformer<P: AsRef<std::path::Path>>(
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
impl UNet2DConditionModel {
|
||||
pub fn load_with_weights(
|
||||
config: UNet2DConfig,
|
||||
weights: std::collections::HashMap<String, Tensor>,
|
||||
device: &Device,
|
||||
use_flash_attn: bool,
|
||||
dtype: DType,
|
||||
) -> Result<Self> {
|
||||
let mut unet = Self::new(&config, use_flash_attn)?;
|
||||
unet.load_tensors_from_hash(weights, device, dtype)?;
|
||||
Ok(unet)
|
||||
}
|
||||
|
||||
fn load_tensors_from_hash(
|
||||
&mut self,
|
||||
mut weights: std::collections::HashMap<String, Tensor>,
|
||||
device: &Device,
|
||||
dtype: DType,
|
||||
) -> Result<()> {
|
||||
// Load the weights from the HashMap into the model
|
||||
for (name, tensor) in weights.drain() {
|
||||
let tensor = tensor.to_device(device)?.to_dtype(dtype)?;
|
||||
self.load_tensor(&name, tensor)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
+258
-354
@@ -5,397 +5,301 @@
|
||||
use super::embeddings::{TimestepEmbedding, Timesteps};
|
||||
use super::unet_2d_blocks::*;
|
||||
use crate::models::with_tracing::{conv2d, Conv2d};
|
||||
use candle::{Result, Tensor};
|
||||
use candle::{Result, Tensor, Device, DType};
|
||||
use candle_nn as nn;
|
||||
use candle_nn::Module;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct BlockConfig {
|
||||
pub out_channels: usize,
|
||||
/// When `None` no cross-attn is used, when `Some(d)` then cross-attn is used and `d` is the
|
||||
/// number of transformer blocks to be used.
|
||||
pub use_cross_attn: Option<usize>,
|
||||
pub attention_head_dim: usize,
|
||||
pub out_channels: usize,
|
||||
/// When `None` no cross-attn is used, when `Some(d)` then cross-attn is used and `d` is the
|
||||
/// number of transformer blocks to be used.
|
||||
pub use_cross_attn: Option<usize>,
|
||||
pub attention_head_dim: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UNet2DConditionModelConfig {
|
||||
pub center_input_sample: bool,
|
||||
pub flip_sin_to_cos: bool,
|
||||
pub freq_shift: f64,
|
||||
pub blocks: Vec<BlockConfig>,
|
||||
pub layers_per_block: usize,
|
||||
pub downsample_padding: usize,
|
||||
pub mid_block_scale_factor: f64,
|
||||
pub norm_num_groups: usize,
|
||||
pub norm_eps: f64,
|
||||
pub cross_attention_dim: usize,
|
||||
pub sliced_attention_size: Option<usize>,
|
||||
pub use_linear_projection: bool,
|
||||
pub center_input_sample: bool,
|
||||
pub flip_sin_to_cos: bool,
|
||||
pub freq_shift: f64,
|
||||
pub blocks: Vec<BlockConfig>,
|
||||
pub layers_per_block: usize,
|
||||
pub downsample_padding: usize,
|
||||
pub mid_block_scale_factor: f64,
|
||||
pub norm_num_groups: usize,
|
||||
pub norm_eps: f64,
|
||||
pub cross_attention_dim: usize,
|
||||
pub sliced_attention_size: Option<usize>,
|
||||
pub use_linear_projection: bool,
|
||||
}
|
||||
|
||||
impl Default for UNet2DConditionModelConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
center_input_sample: false,
|
||||
flip_sin_to_cos: true,
|
||||
freq_shift: 0.,
|
||||
blocks: vec![
|
||||
BlockConfig {
|
||||
out_channels: 320,
|
||||
use_cross_attn: Some(1),
|
||||
attention_head_dim: 8,
|
||||
},
|
||||
BlockConfig {
|
||||
out_channels: 640,
|
||||
use_cross_attn: Some(1),
|
||||
attention_head_dim: 8,
|
||||
},
|
||||
BlockConfig {
|
||||
out_channels: 1280,
|
||||
use_cross_attn: Some(1),
|
||||
attention_head_dim: 8,
|
||||
},
|
||||
BlockConfig {
|
||||
out_channels: 1280,
|
||||
use_cross_attn: None,
|
||||
attention_head_dim: 8,
|
||||
},
|
||||
],
|
||||
layers_per_block: 2,
|
||||
downsample_padding: 1,
|
||||
mid_block_scale_factor: 1.,
|
||||
norm_num_groups: 32,
|
||||
norm_eps: 1e-5,
|
||||
cross_attention_dim: 1280,
|
||||
sliced_attention_size: None,
|
||||
use_linear_projection: false,
|
||||
}
|
||||
}
|
||||
fn default() -> Self {
|
||||
Self { center_input_sample: false, flip_sin_to_cos: true, freq_shift: 0., blocks: vec![BlockConfig { out_channels: 320, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 640, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: Some(1), attention_head_dim: 8 }, BlockConfig { out_channels: 1280, use_cross_attn: None, attention_head_dim: 8 }], layers_per_block: 2, downsample_padding: 1, mid_block_scale_factor: 1., norm_num_groups: 32, norm_eps: 1e-5, cross_attention_dim: 1280, sliced_attention_size: None, use_linear_projection: false }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum UNetDownBlock {
|
||||
Basic(DownBlock2D),
|
||||
CrossAttn(CrossAttnDownBlock2D),
|
||||
Basic(DownBlock2D),
|
||||
CrossAttn(CrossAttnDownBlock2D),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum UNetUpBlock {
|
||||
Basic(UpBlock2D),
|
||||
CrossAttn(CrossAttnUpBlock2D),
|
||||
Basic(UpBlock2D),
|
||||
CrossAttn(CrossAttnUpBlock2D),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct UNet2DConditionModel {
|
||||
conv_in: Conv2d,
|
||||
time_proj: Timesteps,
|
||||
time_embedding: TimestepEmbedding,
|
||||
down_blocks: Vec<UNetDownBlock>,
|
||||
mid_block: UNetMidBlock2DCrossAttn,
|
||||
up_blocks: Vec<UNetUpBlock>,
|
||||
conv_norm_out: nn::GroupNorm,
|
||||
conv_out: Conv2d,
|
||||
span: tracing::Span,
|
||||
config: UNet2DConditionModelConfig,
|
||||
conv_in: Conv2d,
|
||||
time_proj: Timesteps,
|
||||
time_embedding: TimestepEmbedding,
|
||||
down_blocks: Vec<UNetDownBlock>,
|
||||
mid_block: UNetMidBlock2DCrossAttn,
|
||||
up_blocks: Vec<UNetUpBlock>,
|
||||
conv_norm_out: nn::GroupNorm,
|
||||
conv_out: Conv2d,
|
||||
span: tracing::Span,
|
||||
config: UNet2DConditionModelConfig,
|
||||
}
|
||||
|
||||
impl UNet2DConditionModel {
|
||||
pub fn new(
|
||||
vs: nn::VarBuilder,
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
use_flash_attn: bool,
|
||||
config: UNet2DConditionModelConfig,
|
||||
) -> Result<Self> {
|
||||
let n_blocks = config.blocks.len();
|
||||
let b_channels = config.blocks[0].out_channels;
|
||||
let bl_channels = config.blocks.last().unwrap().out_channels;
|
||||
let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
|
||||
let time_embed_dim = b_channels * 4;
|
||||
let conv_cfg = nn::Conv2dConfig {
|
||||
padding: 1,
|
||||
..Default::default()
|
||||
};
|
||||
let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
|
||||
pub fn new(vs: nn::VarBuilder, in_channels: usize, out_channels: usize, use_flash_attn: bool, config: UNet2DConditionModelConfig) -> Result<Self> {
|
||||
let n_blocks = config.blocks.len();
|
||||
let b_channels = config.blocks[0].out_channels;
|
||||
let bl_channels = config.blocks.last().unwrap().out_channels;
|
||||
let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
|
||||
let time_embed_dim = b_channels * 4;
|
||||
let conv_cfg = nn::Conv2dConfig { padding: 1, ..Default::default() };
|
||||
let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;
|
||||
|
||||
let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
|
||||
let time_embedding =
|
||||
TimestepEmbedding::new(vs.pp("time_embedding"), b_channels, time_embed_dim)?;
|
||||
let time_proj = Timesteps::new(b_channels, config.flip_sin_to_cos, config.freq_shift);
|
||||
let time_embedding = TimestepEmbedding::new(vs.pp("time_embedding"), b_channels, time_embed_dim)?;
|
||||
|
||||
let vs_db = vs.pp("down_blocks");
|
||||
let down_blocks = (0..n_blocks)
|
||||
.map(|i| {
|
||||
let BlockConfig {
|
||||
out_channels,
|
||||
use_cross_attn,
|
||||
attention_head_dim,
|
||||
} = config.blocks[i];
|
||||
let vs_db = vs.pp("down_blocks");
|
||||
let down_blocks = (0..n_blocks)
|
||||
.map(|i| {
|
||||
let BlockConfig { out_channels, use_cross_attn, attention_head_dim } = config.blocks[i];
|
||||
|
||||
// Enable automatic attention slicing if the config sliced_attention_size is set to 0.
|
||||
let sliced_attention_size = match config.sliced_attention_size {
|
||||
Some(0) => Some(attention_head_dim / 2),
|
||||
_ => config.sliced_attention_size,
|
||||
};
|
||||
|
||||
let in_channels = if i > 0 {
|
||||
config.blocks[i - 1].out_channels
|
||||
} else {
|
||||
b_channels
|
||||
};
|
||||
let db_cfg = DownBlock2DConfig {
|
||||
num_layers: config.layers_per_block,
|
||||
resnet_eps: config.norm_eps,
|
||||
resnet_groups: config.norm_num_groups,
|
||||
add_downsample: i < n_blocks - 1,
|
||||
downsample_padding: config.downsample_padding,
|
||||
..Default::default()
|
||||
};
|
||||
if let Some(transformer_layers_per_block) = use_cross_attn {
|
||||
let config = CrossAttnDownBlock2DConfig {
|
||||
downblock: db_cfg,
|
||||
attn_num_head_channels: attention_head_dim,
|
||||
cross_attention_dim: config.cross_attention_dim,
|
||||
sliced_attention_size,
|
||||
use_linear_projection: config.use_linear_projection,
|
||||
transformer_layers_per_block,
|
||||
};
|
||||
let block = CrossAttnDownBlock2D::new(
|
||||
vs_db.pp(i.to_string()),
|
||||
in_channels,
|
||||
out_channels,
|
||||
Some(time_embed_dim),
|
||||
use_flash_attn,
|
||||
config,
|
||||
)?;
|
||||
Ok(UNetDownBlock::CrossAttn(block))
|
||||
} else {
|
||||
let block = DownBlock2D::new(
|
||||
vs_db.pp(i.to_string()),
|
||||
in_channels,
|
||||
out_channels,
|
||||
Some(time_embed_dim),
|
||||
db_cfg,
|
||||
)?;
|
||||
Ok(UNetDownBlock::Basic(block))
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
// https://github.com/huggingface/diffusers/blob/a76f2ad538e73b34d5fe7be08c8eb8ab38c7e90c/src/diffusers/models/unet_2d_condition.py#L462
|
||||
let mid_transformer_layers_per_block = match config.blocks.last() {
|
||||
None => 1,
|
||||
Some(block) => block.use_cross_attn.unwrap_or(1),
|
||||
};
|
||||
let mid_cfg = UNetMidBlock2DCrossAttnConfig {
|
||||
resnet_eps: config.norm_eps,
|
||||
output_scale_factor: config.mid_block_scale_factor,
|
||||
cross_attn_dim: config.cross_attention_dim,
|
||||
attn_num_head_channels: bl_attention_head_dim,
|
||||
resnet_groups: Some(config.norm_num_groups),
|
||||
use_linear_projection: config.use_linear_projection,
|
||||
transformer_layers_per_block: mid_transformer_layers_per_block,
|
||||
..Default::default()
|
||||
// Enable automatic attention slicing if the config sliced_attention_size is set to 0.
|
||||
let sliced_attention_size = match config.sliced_attention_size {
|
||||
Some(0) => Some(attention_head_dim / 2),
|
||||
_ => config.sliced_attention_size,
|
||||
};
|
||||
|
||||
let mid_block = UNetMidBlock2DCrossAttn::new(
|
||||
vs.pp("mid_block"),
|
||||
bl_channels,
|
||||
Some(time_embed_dim),
|
||||
use_flash_attn,
|
||||
mid_cfg,
|
||||
)?;
|
||||
|
||||
let vs_ub = vs.pp("up_blocks");
|
||||
let up_blocks = (0..n_blocks)
|
||||
.map(|i| {
|
||||
let BlockConfig {
|
||||
out_channels,
|
||||
use_cross_attn,
|
||||
attention_head_dim,
|
||||
} = config.blocks[n_blocks - 1 - i];
|
||||
|
||||
// Enable automatic attention slicing if the config sliced_attention_size is set to 0.
|
||||
let sliced_attention_size = match config.sliced_attention_size {
|
||||
Some(0) => Some(attention_head_dim / 2),
|
||||
_ => config.sliced_attention_size,
|
||||
};
|
||||
|
||||
let prev_out_channels = if i > 0 {
|
||||
config.blocks[n_blocks - i].out_channels
|
||||
} else {
|
||||
bl_channels
|
||||
};
|
||||
let in_channels = {
|
||||
let index = if i == n_blocks - 1 {
|
||||
0
|
||||
} else {
|
||||
n_blocks - i - 2
|
||||
};
|
||||
config.blocks[index].out_channels
|
||||
};
|
||||
let ub_cfg = UpBlock2DConfig {
|
||||
num_layers: config.layers_per_block + 1,
|
||||
resnet_eps: config.norm_eps,
|
||||
resnet_groups: config.norm_num_groups,
|
||||
add_upsample: i < n_blocks - 1,
|
||||
..Default::default()
|
||||
};
|
||||
if let Some(transformer_layers_per_block) = use_cross_attn {
|
||||
let config = CrossAttnUpBlock2DConfig {
|
||||
upblock: ub_cfg,
|
||||
attn_num_head_channels: attention_head_dim,
|
||||
cross_attention_dim: config.cross_attention_dim,
|
||||
sliced_attention_size,
|
||||
use_linear_projection: config.use_linear_projection,
|
||||
transformer_layers_per_block,
|
||||
};
|
||||
let block = CrossAttnUpBlock2D::new(
|
||||
vs_ub.pp(i.to_string()),
|
||||
in_channels,
|
||||
prev_out_channels,
|
||||
out_channels,
|
||||
Some(time_embed_dim),
|
||||
use_flash_attn,
|
||||
config,
|
||||
)?;
|
||||
Ok(UNetUpBlock::CrossAttn(block))
|
||||
} else {
|
||||
let block = UpBlock2D::new(
|
||||
vs_ub.pp(i.to_string()),
|
||||
in_channels,
|
||||
prev_out_channels,
|
||||
out_channels,
|
||||
Some(time_embed_dim),
|
||||
ub_cfg,
|
||||
)?;
|
||||
Ok(UNetUpBlock::Basic(block))
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let conv_norm_out = nn::group_norm(
|
||||
config.norm_num_groups,
|
||||
b_channels,
|
||||
config.norm_eps,
|
||||
vs.pp("conv_norm_out"),
|
||||
)?;
|
||||
let conv_out = conv2d(b_channels, out_channels, 3, conv_cfg, vs.pp("conv_out"))?;
|
||||
let span = tracing::span!(tracing::Level::TRACE, "unet2d");
|
||||
Ok(Self {
|
||||
conv_in,
|
||||
time_proj,
|
||||
time_embedding,
|
||||
down_blocks,
|
||||
mid_block,
|
||||
up_blocks,
|
||||
conv_norm_out,
|
||||
conv_out,
|
||||
span,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
xs: &Tensor,
|
||||
timestep: f64,
|
||||
encoder_hidden_states: &Tensor,
|
||||
) -> Result<Tensor> {
|
||||
let _enter = self.span.enter();
|
||||
self.forward_with_additional_residuals(xs, timestep, encoder_hidden_states, None, None)
|
||||
}
|
||||
|
||||
pub fn forward_with_additional_residuals(
|
||||
&self,
|
||||
xs: &Tensor,
|
||||
timestep: f64,
|
||||
encoder_hidden_states: &Tensor,
|
||||
down_block_additional_residuals: Option<&[Tensor]>,
|
||||
mid_block_additional_residual: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
let (bsize, _channels, height, width) = xs.dims4()?;
|
||||
let device = xs.device();
|
||||
let n_blocks = self.config.blocks.len();
|
||||
let num_upsamplers = n_blocks - 1;
|
||||
let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
|
||||
let forward_upsample_size =
|
||||
height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
|
||||
// 0. center input if necessary
|
||||
let xs = if self.config.center_input_sample {
|
||||
((xs * 2.0)? - 1.0)?
|
||||
let in_channels = if i > 0 { config.blocks[i - 1].out_channels } else { b_channels };
|
||||
let db_cfg = DownBlock2DConfig { num_layers: config.layers_per_block, resnet_eps: config.norm_eps, resnet_groups: config.norm_num_groups, add_downsample: i < n_blocks - 1, downsample_padding: config.downsample_padding, ..Default::default() };
|
||||
if let Some(transformer_layers_per_block) = use_cross_attn {
|
||||
let config = CrossAttnDownBlock2DConfig { downblock: db_cfg, attn_num_head_channels: attention_head_dim, cross_attention_dim: config.cross_attention_dim, sliced_attention_size, use_linear_projection: config.use_linear_projection, transformer_layers_per_block };
|
||||
let block = CrossAttnDownBlock2D::new(vs_db.pp(i.to_string()), in_channels, out_channels, Some(time_embed_dim), use_flash_attn, config)?;
|
||||
Ok(UNetDownBlock::CrossAttn(block))
|
||||
} else {
|
||||
xs.clone()
|
||||
};
|
||||
// 1. time
|
||||
let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
|
||||
let emb = self.time_proj.forward(&emb)?;
|
||||
let emb = self.time_embedding.forward(&emb)?;
|
||||
// 2. pre-process
|
||||
let xs = self.conv_in.forward(&xs)?;
|
||||
// 3. down
|
||||
let mut down_block_res_xs = vec![xs.clone()];
|
||||
let mut xs = xs;
|
||||
for down_block in self.down_blocks.iter() {
|
||||
let (_xs, res_xs) = match down_block {
|
||||
UNetDownBlock::Basic(b) => b.forward(&xs, Some(&emb))?,
|
||||
UNetDownBlock::CrossAttn(b) => {
|
||||
b.forward(&xs, Some(&emb), Some(encoder_hidden_states))?
|
||||
}
|
||||
};
|
||||
down_block_res_xs.extend(res_xs);
|
||||
xs = _xs;
|
||||
let block = DownBlock2D::new(vs_db.pp(i.to_string()), in_channels, out_channels, Some(time_embed_dim), db_cfg)?;
|
||||
Ok(UNetDownBlock::Basic(block))
|
||||
}
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let new_down_block_res_xs =
|
||||
if let Some(down_block_additional_residuals) = down_block_additional_residuals {
|
||||
let mut v = vec![];
|
||||
// A previous version of this code had a bug because of the addition being made
|
||||
// in place via += hence modifying the input of the mid block.
|
||||
for (i, residuals) in down_block_additional_residuals.iter().enumerate() {
|
||||
v.push((&down_block_res_xs[i] + residuals)?)
|
||||
}
|
||||
v
|
||||
} else {
|
||||
down_block_res_xs
|
||||
};
|
||||
let mut down_block_res_xs = new_down_block_res_xs;
|
||||
// https://github.com/huggingface/diffusers/blob/a76f2ad538e73b34d5fe7be08c8eb8ab38c7e90c/src/diffusers/models/unet_2d_condition.py#L462
|
||||
let mid_transformer_layers_per_block = match config.blocks.last() {
|
||||
None => 1,
|
||||
Some(block) => block.use_cross_attn.unwrap_or(1),
|
||||
};
|
||||
let mid_cfg = UNetMidBlock2DCrossAttnConfig { resnet_eps: config.norm_eps, output_scale_factor: config.mid_block_scale_factor, cross_attn_dim: config.cross_attention_dim, attn_num_head_channels: bl_attention_head_dim, resnet_groups: Some(config.norm_num_groups), use_linear_projection: config.use_linear_projection, transformer_layers_per_block: mid_transformer_layers_per_block, ..Default::default() };
|
||||
|
||||
// 4. mid
|
||||
let xs = self
|
||||
.mid_block
|
||||
.forward(&xs, Some(&emb), Some(encoder_hidden_states))?;
|
||||
let xs = match mid_block_additional_residual {
|
||||
None => xs,
|
||||
Some(m) => (m + xs)?,
|
||||
let mid_block = UNetMidBlock2DCrossAttn::new(vs.pp("mid_block"), bl_channels, Some(time_embed_dim), use_flash_attn, mid_cfg)?;
|
||||
|
||||
let vs_ub = vs.pp("up_blocks");
|
||||
let up_blocks = (0..n_blocks)
|
||||
.map(|i| {
|
||||
let BlockConfig { out_channels, use_cross_attn, attention_head_dim } = config.blocks[n_blocks - 1 - i];
|
||||
|
||||
// Enable automatic attention slicing if the config sliced_attention_size is set to 0.
|
||||
let sliced_attention_size = match config.sliced_attention_size {
|
||||
Some(0) => Some(attention_head_dim / 2),
|
||||
_ => config.sliced_attention_size,
|
||||
};
|
||||
// 5. up
|
||||
let mut xs = xs;
|
||||
let mut upsample_size = None;
|
||||
for (i, up_block) in self.up_blocks.iter().enumerate() {
|
||||
let n_resnets = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.resnets.len(),
|
||||
UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
|
||||
};
|
||||
let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
|
||||
if i < n_blocks - 1 && forward_upsample_size {
|
||||
let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
|
||||
upsample_size = Some((h, w))
|
||||
}
|
||||
xs = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&emb), upsample_size)?,
|
||||
UNetUpBlock::CrossAttn(b) => b.forward(
|
||||
&xs,
|
||||
&res_xs,
|
||||
Some(&emb),
|
||||
upsample_size,
|
||||
Some(encoder_hidden_states),
|
||||
)?,
|
||||
};
|
||||
|
||||
let prev_out_channels = if i > 0 { config.blocks[n_blocks - i].out_channels } else { bl_channels };
|
||||
let in_channels = {
|
||||
let index = if i == n_blocks - 1 { 0 } else { n_blocks - i - 2 };
|
||||
config.blocks[index].out_channels
|
||||
};
|
||||
let ub_cfg = UpBlock2DConfig { num_layers: config.layers_per_block + 1, resnet_eps: config.norm_eps, resnet_groups: config.norm_num_groups, add_upsample: i < n_blocks - 1, ..Default::default() };
|
||||
if let Some(transformer_layers_per_block) = use_cross_attn {
|
||||
let config = CrossAttnUpBlock2DConfig { upblock: ub_cfg, attn_num_head_channels: attention_head_dim, cross_attention_dim: config.cross_attention_dim, sliced_attention_size, use_linear_projection: config.use_linear_projection, transformer_layers_per_block };
|
||||
let block = CrossAttnUpBlock2D::new(vs_ub.pp(i.to_string()), in_channels, prev_out_channels, out_channels, Some(time_embed_dim), use_flash_attn, config)?;
|
||||
Ok(UNetUpBlock::CrossAttn(block))
|
||||
} else {
|
||||
let block = UpBlock2D::new(vs_ub.pp(i.to_string()), in_channels, prev_out_channels, out_channels, Some(time_embed_dim), ub_cfg)?;
|
||||
Ok(UNetUpBlock::Basic(block))
|
||||
}
|
||||
// 6. post-process
|
||||
let xs = self.conv_norm_out.forward(&xs)?;
|
||||
let xs = nn::ops::silu(&xs)?;
|
||||
self.conv_out.forward(&xs)
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let conv_norm_out = nn::group_norm(config.norm_num_groups, b_channels, config.norm_eps, vs.pp("conv_norm_out"))?;
|
||||
let conv_out = conv2d(b_channels, out_channels, 3, conv_cfg, vs.pp("conv_out"))?;
|
||||
let span = tracing::span!(tracing::Level::TRACE, "unet2d");
|
||||
Ok(Self { conv_in, time_proj, time_embedding, down_blocks, mid_block, up_blocks, conv_norm_out, conv_out, span, config })
|
||||
}
|
||||
|
||||
pub fn forward(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor) -> Result<Tensor> {
|
||||
let _enter = self.span.enter();
|
||||
self.forward_with_additional_residuals(xs, timestep, encoder_hidden_states, None, None)
|
||||
}
|
||||
|
||||
pub fn forward_with_additional_residuals(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor, down_block_additional_residuals: Option<&[Tensor]>, mid_block_additional_residual: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (bsize, _channels, height, width) = xs.dims4()?;
|
||||
let device = xs.device();
|
||||
let n_blocks = self.config.blocks.len();
|
||||
let num_upsamplers = n_blocks - 1;
|
||||
let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
|
||||
let forward_upsample_size = height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
|
||||
// 0. center input if necessary
|
||||
let xs = if self.config.center_input_sample { ((xs * 2.0)? - 1.0)? } else { xs.clone() };
|
||||
// 1. time
|
||||
let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
|
||||
let emb = self.time_proj.forward(&emb)?;
|
||||
let time_emb = self.time_embedding.forward(&emb)?;
|
||||
// 2. pre-process
|
||||
let xs = self.conv_in.forward(&xs)?;
|
||||
// 3. down
|
||||
let mut down_block_res_xs = vec![xs.clone()];
|
||||
let mut xs = xs;
|
||||
for down_block in self.down_blocks.iter() {
|
||||
let (_xs, res_xs) = match down_block {
|
||||
UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
|
||||
UNetDownBlock::CrossAttn(b) => b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?,
|
||||
};
|
||||
down_block_res_xs.extend(res_xs);
|
||||
xs = _xs;
|
||||
}
|
||||
|
||||
let new_down_block_res_xs = if let Some(down_block_additional_residuals) = down_block_additional_residuals {
|
||||
let mut v = vec![];
|
||||
// A previous version of this code had a bug because of the addition being made
|
||||
// in place via += hence modifying the input of the mid block.
|
||||
for (i, residuals) in down_block_additional_residuals.iter().enumerate() {
|
||||
v.push((&down_block_res_xs[i] + residuals)?)
|
||||
}
|
||||
v
|
||||
} else {
|
||||
down_block_res_xs
|
||||
};
|
||||
let mut down_block_res_xs = new_down_block_res_xs;
|
||||
|
||||
// 4. mid
|
||||
let xs = self.mid_block.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
|
||||
let xs = match mid_block_additional_residual {
|
||||
None => xs,
|
||||
Some(m) => (m + xs)?,
|
||||
};
|
||||
// 5. up
|
||||
let mut xs = xs;
|
||||
let mut upsample_size = None;
|
||||
for (i, up_block) in self.up_blocks.iter().enumerate() {
|
||||
let n_resnets = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.resnets.len(),
|
||||
UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
|
||||
};
|
||||
let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
|
||||
if i < n_blocks - 1 && forward_upsample_size {
|
||||
let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
|
||||
upsample_size = Some((h, w))
|
||||
}
|
||||
xs = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
|
||||
UNetUpBlock::CrossAttn(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size, Some(encoder_hidden_states))?,
|
||||
};
|
||||
}
|
||||
// 6. post-process
|
||||
let xs = self.conv_norm_out.forward(&xs)?;
|
||||
let xs = nn::ops::silu(&xs)?;
|
||||
self.conv_out.forward(&xs)
|
||||
}
|
||||
|
||||
pub fn forward_with_guidance(&self, xs: &Tensor, timestep: f64, encoder_hidden_states: &Tensor, guidance_embedding: Option<&Tensor>) -> Result<Tensor> {
|
||||
let _enter = self.span.enter();
|
||||
let (bsize, _channels, height, width) = xs.dims4()?;
|
||||
let device = xs.device();
|
||||
let n_blocks = self.config.blocks.len();
|
||||
let num_upsamplers = n_blocks - 1;
|
||||
let default_overall_up_factor = 2usize.pow(num_upsamplers as u32);
|
||||
let forward_upsample_size = height % default_overall_up_factor != 0 || width % default_overall_up_factor != 0;
|
||||
|
||||
// 0. center input if necessary
|
||||
let xs = if self.config.center_input_sample { ((xs * 2.0)? - 1.0)? } else { xs.clone() };
|
||||
|
||||
// 1. time
|
||||
let emb = (Tensor::ones(bsize, xs.dtype(), device)? * timestep)?;
|
||||
let emb = self.time_proj.forward(&emb)?;
|
||||
let mut time_emb = self.time_embedding.forward(&emb)?;
|
||||
|
||||
// Add guidance embedding if provided
|
||||
if let Some(guidance_emb) = guidance_embedding {
|
||||
// Print shapes for debugging
|
||||
println!("Guidance embedding shape: {:?}", guidance_emb.shape());
|
||||
println!("Time embedding shape: {:?}", time_emb.shape());
|
||||
|
||||
// Get the batch size and embedding dimension
|
||||
let (batch_size, emb_dim) = time_emb.dims2()?;
|
||||
|
||||
// Reshape guidance embedding to match time embedding dimensions
|
||||
let guidance_emb = guidance_emb.reshape(&[1, emb_dim])?.expand(&[batch_size, emb_dim])?;
|
||||
|
||||
// Add to time embeddings
|
||||
time_emb = (time_emb + guidance_emb)?;
|
||||
}
|
||||
|
||||
// 2. pre-process
|
||||
let xs = self.conv_in.forward(&xs)?;
|
||||
|
||||
// 3. down
|
||||
let mut down_block_res_xs = vec![xs.clone()];
|
||||
let mut xs = xs;
|
||||
for down_block in self.down_blocks.iter() {
|
||||
let (_xs, res_xs) = match down_block {
|
||||
UNetDownBlock::Basic(b) => b.forward(&xs, Some(&time_emb))?,
|
||||
UNetDownBlock::CrossAttn(b) => b.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?,
|
||||
};
|
||||
down_block_res_xs.extend(res_xs);
|
||||
xs = _xs;
|
||||
}
|
||||
|
||||
// 4. mid
|
||||
let xs = self.mid_block.forward(&xs, Some(&time_emb), Some(encoder_hidden_states))?;
|
||||
|
||||
// 5. up
|
||||
let mut xs = xs;
|
||||
let mut upsample_size = None;
|
||||
for (i, up_block) in self.up_blocks.iter().enumerate() {
|
||||
let n_resnets = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.resnets.len(),
|
||||
UNetUpBlock::CrossAttn(b) => b.upblock.resnets.len(),
|
||||
};
|
||||
let res_xs = down_block_res_xs.split_off(down_block_res_xs.len() - n_resnets);
|
||||
if i < n_blocks - 1 && forward_upsample_size {
|
||||
let (_, _, h, w) = down_block_res_xs.last().unwrap().dims4()?;
|
||||
upsample_size = Some((h, w))
|
||||
}
|
||||
xs = match up_block {
|
||||
UNetUpBlock::Basic(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size)?,
|
||||
UNetUpBlock::CrossAttn(b) => b.forward(&xs, &res_xs, Some(&time_emb), upsample_size, Some(encoder_hidden_states))?,
|
||||
};
|
||||
}
|
||||
|
||||
// 6. post-process
|
||||
let xs = self.conv_norm_out.forward(&xs)?;
|
||||
let xs = nn::ops::silu(&xs)?;
|
||||
self.conv_out.forward(&xs)
|
||||
}
|
||||
|
||||
pub fn load_with_weights(config: UNet2DConditionModelConfig, weights: HashMap<String, Tensor>, device: &Device, use_flash_attn: bool, dtype: DType) -> Result<Self> {
|
||||
let vb = nn::VarBuilder::from_tensors(weights, dtype, device);
|
||||
Self::new(vb, 4, 4, use_flash_attn, config)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user