More generic sampling.

candle-transformers/src/models/flux/mod.rs

@@ -1,3 +1,19 @@
+use candle::{Result, Tensor};
+
+pub trait WithForward {
+    #[allow(clippy::too_many_arguments)]
+    fn forward(
+        &self,
+        img: &Tensor,
+        img_ids: &Tensor,
+        txt: &Tensor,
+        txt_ids: &Tensor,
+        timesteps: &Tensor,
+        y: &Tensor,
+        guidance: Option<&Tensor>,
+    ) -> Result<Tensor>;
+}
+
 pub mod autoencoder;
 pub mod model;
 pub mod quantized_model;

candle-transformers/src/models/flux/model.rs

@@ -575,9 +575,11 @@ impl Flux {
             final_layer,
         })
     }
+}
 
+impl super::WithForward for Flux {
     #[allow(clippy::too_many_arguments)]
-    pub fn forward(
+    fn forward(
         &self,
         img: &Tensor,
         img_ids: &Tensor,

candle-transformers/src/models/flux/quantized_model.rs

@@ -577,9 +577,11 @@ impl Flux {
             final_layer,
         })
     }
+}
 
+impl super::WithForward for Flux {
     #[allow(clippy::too_many_arguments)]
-    pub fn forward(
+    fn forward(
         &self,
         img: &Tensor,
         img_ids: &Tensor,

candle-transformers/src/models/flux/sampling.rs

@@ -92,8 +92,8 @@ pub fn unpack(xs: &Tensor, height: usize, width: usize) -> Result<Tensor> {
 }
 
 #[allow(clippy::too_many_arguments)]
-pub fn denoise(
+pub fn denoise<M: super::WithForward>(
-    model: &super::model::Flux,
+    model: &M,
     img: &Tensor,
     img_ids: &Tensor,
     txt: &Tensor,