From be21d7e75a38b2d383ae64f028ded2d4ce3c473d Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Wed, 9 Aug 2023 14:23:59 +0200
Subject: [PATCH] Fix the padding used in stable diffusion. (#362)

---
 candle-examples/examples/stable-diffusion/clip.rs | 4 ++--
 candle-examples/examples/stable-diffusion/main.rs | 5 ++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/candle-examples/examples/stable-diffusion/clip.rs b/candle-examples/examples/stable-diffusion/clip.rs
index 12f482fd..ca00b417 100644
--- a/candle-examples/examples/stable-diffusion/clip.rs
+++ b/candle-examples/examples/stable-diffusion/clip.rs
@@ -31,7 +31,7 @@ pub struct Config {
     intermediate_size: usize,
     pub max_position_embeddings: usize,
     // The character to use for padding, use EOS when not set.
-    pad_with: Option<String>,
+    pub pad_with: Option<String>,
     num_hidden_layers: usize,
     num_attention_heads: usize,
     #[allow(dead_code)]
@@ -257,7 +257,7 @@ impl ClipEncoder {
     fn forward(&self, xs: &Tensor, causal_attention_mask: &Tensor) -> Result<Tensor> {
         let mut xs = xs.clone();
         for layer in self.layers.iter() {
-            xs = layer.forward(&xs, causal_attention_mask)?
+            xs = layer.forward(&xs, causal_attention_mask)?;
         }
         Ok(xs)
     }
diff --git a/candle-examples/examples/stable-diffusion/main.rs b/candle-examples/examples/stable-diffusion/main.rs
index 0e0330d7..8bb3c56d 100644
--- a/candle-examples/examples/stable-diffusion/main.rs
+++ b/candle-examples/examples/stable-diffusion/main.rs
@@ -181,8 +181,8 @@ fn run(args: Args) -> Result<()> {
     let device = candle_examples::device(cpu)?;
 
     let tokenizer = Tokenizer::from_file(tokenizer).map_err(E::msg)?;
-    let pad_id = match tokenizer.get_padding() {
-        Some(padding) => padding.pad_id,
+    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}\".");
@@ -212,7 +212,6 @@ fn run(args: Args) -> Result<()> {
     let uncond_embeddings = text_model.forward(&uncond_tokens)?;
     let text_embeddings = Tensor::cat(&[uncond_embeddings, text_embeddings], 0)?;
 
-    println!("text-embeddings: {text_embeddings:?}");
     println!("Building the autoencoder.");
     let vae = sd_config.build_vae(&vae_weights, &device)?;
     println!("Building the unet.");