Update app.py

app.py

@@ -9,8 +9,14 @@ import numpy as np
 from PIL import Image
 import random
 import gc
-import torchao.quantization
+import torchao
+from torchao.quantization import quantize, Int8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig
 import aoti
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Model configuration
 MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
@@ -26,6 +32,7 @@ MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
 MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)
 
 # Initialize pipeline
+logger.info("Loading WanImageToVideoPipeline...")
 pipe = WanImageToVideoPipeline.from_pretrained(
     MODEL_ID,
     transformer=WanTransformer3DModel.from_pretrained(
@@ -43,30 +50,31 @@ pipe = WanImageToVideoPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 ).to('cuda')
 
-# Load LoRA weights
-pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy",
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-    adapter_name="lightx2v"
-)
-kwargs_lora = {"load_into_transformer_2": True}
-pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy",
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-    adapter_name="lightx2v_2",
-    **kwargs_lora
-)
-pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
-pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
-pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
-pipe.unload_lora_weights()
+# Load LoRA weights (simplified to single adapter)
+logger.info("Loading LoRA weights...")
+try:
+    pipe.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+        adapter_name="lightx2v"
+    )
+    pipe.set_adapters(["lightx2v"], adapter_weights=[1.0])
+    pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3.0, components=["transformer", "transformer_2"])
+    pipe.unload_lora_weights()
+except Exception as e:
+    logger.error(f"Failed to load LoRA weights: {str(e)}")
+    raise
 
 # Quantization
-quantize_(pipe.text_encoder, torchao.quantization.Int8WeightOnlyConfig())
-quantize_(pipe.transformer, torchao.quantization.Float8DynamicActivationFloat8WeightConfig())
-quantize_(pipe.transformer_2, torchao.quantization.Float8DynamicActivationFloat8WeightConfig())
-aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
-aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
+logger.info(f"Using torchao version: {torchao.__version__}")
+try:
+    quantize(pipe.text_encoder, Int8WeightOnlyConfig())
+    quantize(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+    quantize(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+    aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
+    aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
+except Exception as e:
+    logger.warning(f"Quantization failed: {str(e)}. Proceeding without quantization.")
 
 # Default prompts
 default_prompt_i2v = {
@@ -200,7 +208,7 @@ def generate_video(
     export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
     return video_path, current_seed
 
-# Updated UI with pretranslated texts and enhanced tab logic
+# UI with pretranslated texts and tab logic
 def create_demo():
     with gr.Blocks(css="", title="Fast Image to Video") as demo:
         gr.HTML("""
@@ -547,13 +555,6 @@ def create_demo():
                 return false;
             };
 
-            const observer = new MutationObserver((mutations, obs) => {
-                const tabElements = document.querySelectorAll('[role="tab"], button, div, .gr-tab-item, .tab-item, [data-testid="tab"], [aria-selected], .nav-item, [data-tab]');
-                if (selectTab(tabElements, tab)) {
-                    obs.disconnect();
-                }
-            });
-
             const tryObserveTabs = (attempt = 1, maxAttempts = 3) => {
                 const tabs = document.querySelector('.gr-tabs, .tabs');
                 if (tabs) {
@@ -585,6 +586,13 @@ def create_demo():
                 }
             };
 
+            const observer = new MutationObserver((mutations, obs) => {
+                const tabElements = document.querySelectorAll('[role="tab"], button, div, .gr-tab-item, .tab-item, [data-testid="tab"], [aria-selected], .nav-item, [data-tab]');
+                if (selectTab(tabElements, tab)) {
+                    obs.disconnect();
+                }
+            });
+
             tryObserveTabs();
 
             const toolbarSelectors = [