Update app.py

app.py

@@ -18,7 +18,8 @@ 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_
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
 
 # --- Constants and Model Loading ---
 MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
@@ -67,13 +68,13 @@ 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,
-)
+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.")