Update app.py

app.py

@@ -8,6 +8,7 @@ import os
 # from diffusers import QwenImageEditInpaintPipeline
 from optimization import optimize_pipeline_
 from diffusers.utils import load_image
+from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_inpaint import QwenImageEditInpaintPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
@@ -21,14 +22,44 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 # Initialize Qwen Image Edit pipeline
-pipe = QwenImageEditInpaintPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=torch.bfloat16).to("cuda")
+# Scheduler configuration for Lightning
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": math.log(3),
+    "invert_sigmas": False,
+    "max_image_seq_len": 8192,
+    "max_shift": math.log(3),
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "shift_terminal": None,
+    "stochastic_sampling": False,
+    "time_shift_type": "exponential",
+    "use_beta_sigmas": False,
+    "use_dynamic_shifting": True,
+    "use_exponential_sigmas": False,
+    "use_karras_sigmas": False,
+}
+
+# Initialize scheduler with Lightning config
+scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+
+
+pipe = QwenImageEditInpaintPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=scheduler, torch_dtype=torch.bfloat16).to("cuda")
+pipe.load_lora_weights(
+        "lightx2v/Qwen-Image-Lightning", 
+        weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
+    )
+pipe.fuse_lora()
+
 pipe.transformer.__class__ = QwenImageTransformer2DModel
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
-dummy_mask  = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/mask_cat.png?raw=true")
 
-# --- Ahead-of-time compilation ---
-optimize_pipeline_(pipe, image=Image.new("RGB", (1328, 1328)), prompt="prompt", mask_image=dummy_mask)
+
+# dummy_mask  = load_image("https://github.com/Trgtuan10/Image_storage/blob/main/mask_cat.png?raw=true")
+
+# # --- Ahead-of-time compilation ---
+# optimize_pipeline_(pipe, image=Image.new("RGB", (1328, 1328)), prompt="prompt", mask_image=dummy_mask)
 
 @spaces.GPU(duration=120)
 def infer(edit_images, prompt, negative_prompt="", seed=42, randomize_seed=False, strength=1.0, num_inference_steps=35, true_cfg_scale=4.0, progress=gr.Progress(track_tqdm=True)):
@@ -138,16 +169,16 @@ with gr.Blocks(css=css) as demo:
                     minimum=1.0,
                     maximum=20.0,
                     step=0.5,
-                    value=4.0,
+                    value=1.0,
                     info="Classifier-free guidance scale"
                 )
   
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=10,
-                    maximum=100,
+                    maximum=50,
                     step=1,
-                    value=35,
+                    value=8,
                 )
 
     gr.on(