From 4ceaac94116a6b934856a121ae30071f5a20a10d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Mar 2025 00:23:03 -0400 Subject: [PATCH] test old clip code --- .../ml_models/src/ml/stable_diffusion/clip.rs | 133 ++++++++++++++++++ .../ml_models/src/ml/stable_diffusion/mod.rs | 1 + 2 files changed, 134 insertions(+) create mode 100644 crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs new file mode 100644 index 0000000000..0b897f36a5 --- /dev/null +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/clip.rs @@ -0,0 +1,133 @@ +#[cfg(feature = "accelerate")] +extern crate accelerate_src; +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +use candle_transformers::models::stable_diffusion; +use std::path::PathBuf; + +use crate::ml::model_file::{ModelFile, StableDiffusionVersion}; +use crate::ml::weights_registry::weights::{CLIP_JSON, SDXL_TURBO_CLIP_TEXT_ENCODER, SDXL_TURBO_CLIP_TEXT_ENCODER_2}; +use anyhow::Error as E; +use candle_core::{DType, Device, Module, Tensor, D}; +use log::info; +use tokenizers::Tokenizer; + +#[allow(clippy::too_many_arguments)] +pub fn old_clip( + prompt: &str, + uncond_prompt: &str, + tokenizer: Option, + clip_weights: Option, + clip2_weights: Option, + sd_version: StableDiffusionVersion, + sd_config: &stable_diffusion::StableDiffusionConfig, + use_f16: bool, + device: &Device, + dtype: DType, + use_guide_scale: bool, + weights_dir: &AppWeightsDir +) -> anyhow::Result { + info!("Preparing text embeddings... for {:?}", sd_version); + + let which = match sd_version { + StableDiffusionVersion::Xl | StableDiffusionVersion::XlInpaint | StableDiffusionVersion::Turbo => vec![true, false], + _ => vec![true], + }; + + let text_embeddings = which + .iter() + .map(|first| { + do_infer_clip_text_embeddings( + &prompt, + &uncond_prompt, + None, // tokenizer + clip_weights.clone(), + None, // clip2_weights + sd_version, + &sd_config, + false, // use_f16 + &device, + dtype, + use_guide_scale, + *first, + weights_dir, + ) + }) + .collect::>>()?; + + let text_embeddings = Tensor::cat(&text_embeddings, D::Minus1)?; + + Ok(text_embeddings) +} + +#[allow(clippy::too_many_arguments)] +fn do_infer_clip_text_embeddings( + prompt: &str, + uncond_prompt: &str, + tokenizer: Option, + clip_weights: Option, + clip2_weights: Option, + sd_version: StableDiffusionVersion, + sd_config: &stable_diffusion::StableDiffusionConfig, + use_f16: bool, + device: &Device, + dtype: DType, + use_guide_scale: bool, + first: bool, + weights_dir: &AppWeightsDir +) -> anyhow::Result { + info!("do_infer_clip_text_embeddings called with args {:?}", (prompt, uncond_prompt, tokenizer.clone().unwrap_or_else(|| "None".to_string()), clip_weights.clone().unwrap_or_else(|| "None".to_string()), clip2_weights.clone().unwrap_or_else(|| "None".to_string()), sd_version, sd_config, use_f16, device, dtype, use_guide_scale, first)); + + let tokenizer = weights_dir.weight_path(&CLIP_JSON); + + info!("Loading Clip Tokenizer path: {:?}", tokenizer); + + 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 = if first { + println!(">>>>>>>>> CLIP 1"); + weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER) + } else { + println!(">>>>>>>>> CLIP 2"); + weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER_2) + }; + 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) +} diff --git a/crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs b/crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs index d605b647d8..3736dda82e 100644 --- a/crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs +++ b/crates/desktop/ml_models/src/ml/stable_diffusion/mod.rs @@ -2,4 +2,5 @@ pub mod get_vae_scale; pub mod infer_clip_text_embeddings; pub mod stable_diffusion_pipeline; pub mod lcm_pipeline; +pub mod clip; mod remap_lcm_strength_range; \ No newline at end of file