Avoid flux gpu<>cpu sync in sampling loop

flue-core/Cargo.toml

@@ -30,7 +30,7 @@ tokio = { workspace = true }
 serde_plain = { workspace = true }
 
 [features]
-cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"]
+cuda = ["candle-core/cuda", "candle-nn/cuda"]
 cudnn = ["candle-core/cudnn"]
 metal = ["candle-core/metal", "candle-nn/metal", "candle-transformers/metal"]
 flash-attn = ["candle-flash-attn"]

flue-core/src/flux/model.rs

@@ -563,7 +563,7 @@ pub struct Flux {
     time_in: MlpEmbedder,
     vector_in: MlpEmbedder,
     guidance_in: Option<MlpEmbedder>,
-    pe_embedder: EmbedNd,
+    pub pe_embedder: EmbedNd,
     double_blocks: Vec<DoubleStreamBlock>,
     single_blocks: Vec<SingleStreamBlock>,
     final_layer: LastLayer,
@@ -614,9 +614,8 @@ impl Flux {
     pub fn forward(
         &self,
         img: &Tensor,
-        img_ids: &Tensor,
         txt: &Tensor,
-        txt_ids: &Tensor,
+        pe: &Tensor,
         timesteps: &Tensor,
         y: &Tensor,
         guidance: Option<&Tensor>,
@@ -628,10 +627,6 @@ impl Flux {
             candle_core::bail!("unexpected shape for img {:?}", img.shape())
         }
         let dtype = img.dtype();
-        let pe = {
-            let ids = Tensor::cat(&[txt_ids, img_ids], 1)?;
-            ids.apply(&self.pe_embedder)?
-        };
         let mut txt = txt.apply(&self.txt_in)?;
         let mut img = img.apply(&self.img_in)?;
         let vec_ = timestep_embedding(timesteps, 256, dtype)?.apply(&self.time_in)?;
@@ -645,12 +640,12 @@ impl Flux {
 
         // Double blocks
         for block in self.double_blocks.iter() {
-            (img, txt) = block.forward(&img, &txt, &vec_, &pe)?
+            (img, txt) = block.forward(&img, &txt, &vec_, pe)?
         }
         // Single blocks
         let mut img = Tensor::cat(&[&txt, &img], 1)?;
         for block in self.single_blocks.iter() {
-            img = block.forward(&img, &vec_, &pe)?;
+            img = block.forward(&img, &vec_, pe)?;
         }
         let img = img.i((.., txt.dim(1)?..))?;
         self.final_layer.forward(&img, &vec_)

flue-core/src/flux/sampling.rs

@@ -107,14 +107,21 @@ pub fn denoise(
     let b_sz = img.dim(0)?;
     let dev = img.device();
     let guidance = Tensor::full(guidance as f32, b_sz, dev)?;
+    let t_vec_one = Tensor::full(1f32, b_sz, dev)?;
     let mut img = img.clone();
+
+    let pe = {
+        let ids = Tensor::cat(&[txt_ids, img_ids], 1)?;
+        ids.apply(&model.pe_embedder)?
+    };
+
     for window in timesteps.windows(2) {
         let (t_curr, t_prev) = match window {
             [a, b] => (a, b),
             _ => continue,
         };
-        let t_vec = Tensor::full(*t_curr as f32, b_sz, dev)?;
-        let pred = model.forward(&img, img_ids, txt, txt_ids, &t_vec, vec_, Some(&guidance))?;
+        let t_vec = (&t_vec_one * { *t_curr })?;
+        let pred = model.forward(&img, txt, &pe, &t_vec, vec_, Some(&guidance))?;
         img = (img + pred * (t_prev - t_curr))?
     }
     Ok(img)