Fix FLUX.1 weights (huggingface#2457)

* fix FLUX.1 weights * added flux1-dev.safetensors
getong · Aug 29, 2024 · c02b7c3 · c02b7c3
1 parent 86613c0
commit c02b7c3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/candle-examples/examples/flux/main.rs b/candle-examples/examples/flux/main.rs
@@ -147,8 +147,8 @@ fn run(args: Args) -> Result<()> {
             println!("CLIP\n{clip_emb}");
             let img = {
                 let model_file = match model {
-                    Model::Schnell => bf_repo.get("flux1-schnell.sft")?,
-                    Model::Dev => bf_repo.get("flux1-dev.sft")?,
+                    Model::Schnell => bf_repo.get("flux1-schnell.safetensors")?,
+                    Model::Dev => bf_repo.get("flux1-dev.safetensors")?,
                 };
                 let vb =
                     unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
@@ -189,7 +189,7 @@ fn run(args: Args) -> Result<()> {
     println!("latent img\n{img}");
 
     let img = {
-        let model_file = bf_repo.get("ae.sft")?;
+        let model_file = bf_repo.get("ae.safetensors")?;
         let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], dtype, &device)? };
         let cfg = match model {
             Model::Dev => flux::autoencoder::Config::dev(),