Update app.py

app.py

@@ -11,6 +11,10 @@ from diffusers import QwenImageEditPipeline, FlowMatchEulerDiscreteScheduler
 from huggingface_hub import InferenceClient
 import math
 
+from optimization import optimize_pipeline_
+from qwenimage.pipeline_qwen_image_edit import QwenImageEditPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+
 # --- Prompt Enhancement using Hugging Face InferenceClient ---
 def polish_prompt_hf(original_prompt, system_prompt):
     """
@@ -158,24 +162,11 @@ scheduler_config = {
 # Initialize scheduler with Lightning config
 scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
 
-# Load the edit pipeline with Lightning scheduler
-pipe = QwenImageEditPipeline.from_pretrained(
-    "Qwen/Qwen-Image-Edit", 
-    scheduler=scheduler,
-    torch_dtype=dtype
-).to(device)
-
-# Load Lightning LoRA weights for acceleration
-try:
-    pipe.load_lora_weights(
-        "lightx2v/Qwen-Image-Lightning", 
-        weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
-    )
-    pipe.fuse_lora()
-    print("Successfully loaded Lightning LoRA weights")
-except Exception as e:
-    print(f"Warning: Could not load Lightning LoRA weights: {e}")
-    print("Continuing with base model...")
+pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=scheduler,torch_dtype=dtype).to(device)
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+
+# --- Ahead-of-time compilation ---
+optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max