Avoid ZeroGPU startup OOM with precompiled Wan AOTI blocks

app.py

@@ -17,8 +17,10 @@ import random
 import gc
 from gradio_client import Client, handle_file # Import for API call
 
-# Import the optimization function from the separate file
-from optimization import optimize_pipeline_
+# Import the quantization configs used by the precompiled ZeroGPU blocks
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
+from torchao.quantization import Int8WeightOnlyConfig
 
 # --- Constants and Model Loading ---
 MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
@@ -67,13 +69,29 @@ for i in range(3):
     torch.cuda.synchronize()
     torch.cuda.empty_cache()
 
-optimize_pipeline_(pipe,
-    image=Image.new('RGB', (MAX_DIMENSION, MIN_DIMENSION)), 
-    prompt='prompt',
-    height=MIN_DIMENSION,
-    width=MAX_DIMENSION,
-    num_frames=MAX_FRAMES_MODEL,
+pipe.load_lora_weights(
+    "Kijai/WanVideo_comfy",
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+    adapter_name="lightx2v",
 )
+pipe.load_lora_weights(
+    "Kijai/WanVideo_comfy",
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+    adapter_name="lightx2v_2",
+    load_into_transformer_2=True,
+)
+pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1.0, 1.0])
+pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3.0, components=["transformer"])
+pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1.0, components=["transformer_2"])
+pipe.unload_lora_weights()
+
+quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+
+spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Wan2", variant="fp8da")
+spaces.aoti_blocks_load(pipe.transformer_2, "zerogpu-aoti/Wan2", variant="fp8da")
+
 print("All models loaded and optimized. Gradio app is ready.")
 
 
@@ -333,4 +351,4 @@ with gr.Blocks() as app:
     )
 
 if __name__ == "__main__":
-    app.launch(mcp_server=True, theme=gr.themes.Citrus(), css=css)
+    app.launch(mcp_server=True, theme=gr.themes.Citrus(), css=css)