Update app.py

app.py

@@ -6,6 +6,10 @@ import spaces
 
 from PIL import Image
 from diffusers import QwenImageEditPipeline
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+import torch 
+import math
+
 
 import os
 
@@ -14,7 +18,25 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the model pipeline
-pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
+# From https://github.com/ModelTC/Qwen-Image-Lightning/blob/342260e8f5468d2f24d084ce04f55e101007118b/generate_with_diffusers.py#L82C9-L97C10
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": math.log(3),  # We use shift=3 in distillation
+    "invert_sigmas": False,
+    "max_image_seq_len": 8192,
+    "max_shift": math.log(3),  # We use shift=3 in distillation
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "shift_terminal": None,  # set shift_terminal to None
+    "stochastic_sampling": False,
+    "time_shift_type": "exponential",
+    "use_beta_sigmas": False,
+    "use_dynamic_shifting": True,
+    "use_exponential_sigmas": False,
+    "use_karras_sigmas": False,
+}
+scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=scheduler, torch_dtype=dtype).to(device)
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
@@ -28,7 +50,7 @@ def infer(
     randomize_seed=False,
     guidance_scale=4.0,
     true_guidance_scale=1.0,
-    num_inference_steps=50,
+    num_inference_steps=8,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
@@ -108,7 +130,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=4.0,
+                    value=1.0,
                 )
 
                 true_guidance_scale = gr.Slider(
@@ -124,7 +146,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=30,
+                    value=8,
                 )
 
         # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)