Update app.py

app.py

@@ -1,7 +1,3 @@
-# PyTorch 2.8 (temporary hack)
-import os
-os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
-
 import gradio as gr
 import numpy as np
 import random
@@ -20,39 +16,10 @@ from optimization import optimize_pipeline_
 # --- Model Loading ---
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
 
-# Load the model pipeline
-
-# scheduler config needed for the LoRA
-# 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)
-
-# lora loading
-pipe.load_lora_weights(
-    "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.0.safetensors", adapter_name="lightx2v"
-)
-pipe.set_adapters(["lightx2v"], adapter_weights=[1.])
-pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=1., components=["transformer"])
-pipe.unload_lora_weights()
-
-# optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt='prompt')
+# --- 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