Update optimization.py (#1)

- Update optimization.py (845a6fa9aafae56c55619862cfd754aa7a5b8d5b)
- Update app.py (5ee50f27f37c63d15584263d7dc04039d4fbc74b)

app.py

@@ -1,3 +1,7 @@
+# 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
@@ -10,8 +14,8 @@ from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
 import torch 
 import math
 
+from optimization import optimize_pipeline_
 
-import os
 
 # --- Model Loading ---
 dtype = torch.bfloat16
@@ -44,9 +48,11 @@ pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", scheduler=s
 pipe.load_lora_weights(
     "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.0.safetensors", adapter_name="lightx2v"
 )
-pipeline.set_adapters(["lightx2v"], adapter_weights=[1.])
-pipeline.fuse_lora(adapter_names=["lightx2v"], lora_scale=1., components=["transformer"])
-pipeline.unload_lora_weights()
+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')
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max

optimization.py

@@ -19,7 +19,7 @@ from optimization_utils import cudagraph
 P = ParamSpec('P')
 
 
-TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=4096, max=8212)
+TRANSFORMER_HIDDEN_DIM = torch.export.Dim('hidden', min=3584, max=8212)
 
 TRANSFORMER_DYNAMIC_SHAPES = {
     'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},