Update app_local.py

app_local.py

@@ -4,7 +4,7 @@ import random
 import torch
 import spaces
 from PIL import Image
-from diffusers import QwenImageEditPipeline
+from diffusers import QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler
 from diffusers.utils import is_xformers_available
 import os
 import sys
@@ -211,10 +211,33 @@ def polish_prompt(original_prompt: str) -> str:
             rewritten_prompt = rewritten_prompt.split(': ', 1)[-1].strip()
         
         return rewritten_prompt[:200] if rewritten_prompt else original_prompt
-        
+
+# 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)
+
+
 # Load main image editing pipeline
 pipe = QwenImageEditPipeline.from_pretrained(
     LOC,
+    scheduler=scheduler,
     torch_dtype=dtype
 ).to(device)
 
@@ -249,13 +272,10 @@ def infer(
     num_inference_steps=8,
     rewrite_prompt=True,
     num_images_per_prompt=1,
+    progress=gr.Progress(track_tqdm=True),
 ):
     """Image editing endpoint with optimized prompt handling"""
-    # Clear cache at start
-    if device == "cuda":
-        torch.cuda.empty_cache()
-        gc.collect()
-    
+
     # Resize image to max 1024px on longest side
     def resize_image(pil_image, max_size=1024):
         """Resize image to maximum dimension of 1024px while maintaining aspect ratio"""