From 49c010db59d206823475f09dbec07728673b0487 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Tue, 25 Mar 2025 17:34:33 -0400
Subject: [PATCH] temporarily remove references to flux (temp fix)

---
 crates/desktop/ml_models/src/ml/mod.rs        |  3 +-
 .../src/endpoints/text_to_image_endpoint.rs   | 29 ++++++++++++-------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/crates/desktop/ml_models/src/ml/mod.rs b/crates/desktop/ml_models/src/ml/mod.rs
index 67908b8f65..c889cdbae0 100644
--- a/crates/desktop/ml_models/src/ml/mod.rs
+++ b/crates/desktop/ml_models/src/ml/mod.rs
@@ -4,5 +4,4 @@ pub mod model_cache;
 pub mod model_config;
 pub mod model_file;
 pub mod prompt_cache;
-pub mod stable_diffusion;
-pub mod flux;
\ No newline at end of file
+pub mod stable_diffusion;
\ No newline at end of file
diff --git a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
index fe798720b0..5ce4e18f24 100644
--- a/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
+++ b/crates/desktop/src-tauri/src/endpoints/text_to_image_endpoint.rs
@@ -1,5 +1,5 @@
 #[cfg(not(target_os = "macos"))]
-use ml_models::ml::flux::flux_pipeline::{flux_pipeline, FluxPipelineArgs, FluxModel};
+use ml_models::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
 
 use crate::state::app_config::AppConfig;
 use crate::state::app_dir::AppDataRoot;
@@ -84,24 +84,31 @@ async fn text_to_image_impl(
   app_data_root: &AppDataRoot,
 ) -> Result<RgbImage, String> {
   
-  let args = FluxPipelineArgs {
+  let weights_dir = app_data_root.weights_dir();
+  let vae_path = weights_dir.weight_path(&LYKON_DEAMSHAPER_7_VAE);
+  let unet_path = weights_dir.weight_path(&SIMIANLUO_LCM_DREAMSHAPER_V7_UNET);
+  let clip_json_path = weights_dir.weight_path(&CLIP_JSON);
+  let clip_weights_path= weights_dir.weight_path(&SDXL_TURBO_CLIP_TEXT_ENCODER);
+
+  let args = Args {
     image: &image,
     prompt: prompt.to_string(),
+    uncond_prompt: "".to_string(),
     model_cache,
     prompt_cache: &prompt_cache,
-    app: &app,
-    flux_model: FluxModel::Schnell,
-    active_blocks: None,
-    use_cuda_stream: Some(true),
-    prefetch_next_batch: Some(false),
-    use_full_gpu: Some(true),
     img2img_strength: strength,
-    seed: None,
-    use_quantized_model: Some(true),
+    cfg_scale: config.cfg_scale,
+    use_flash_attn: true,
     model_config: &config.model_config,
+    maybe_seed: None,
+    scheduler_steps: config.scheduler_steps,
+    vae_path: &vae_path,
+    unet_path: &unet_path,
+    clip_json_path: &clip_json_path,
+    clip_weights_path: &clip_weights_path,
   };
 
-  let image = flux_pipeline(args)
+  let image = lcm_pipeline(args)
     .map_err(|err| format!("failure to encode image: {:?}", err))?;
 
   Ok(image)