Update app.py

app.py

@@ -43,20 +43,19 @@ default_negative_prompt = "static, still, no motion, frozen"
 
 # Initialize once on startup
 video_pipe = None
+video_pipeline_ready = False
 
 def initialize_video_pipeline():
-    global video_pipe
-    if video_pipe is None:
+    global video_pipe, video_pipeline_ready
+    if video_pipe is None and not video_pipeline_ready:
         try:
+            print("Starting video pipeline initialization...")
+            
             # Install PyTorch 2.8 (if needed)
             os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
             
-            # Import optimization module
-            try:
-                from optimization import optimize_pipeline_
-            except ImportError:
-                print("Warning: optimization module not found, skipping optimization")
-                optimize_pipeline_ = None
+            # Import LoRA loading utilities
+            from peft import LoraConfig, get_peft_model, TaskType
             
             video_pipe = WanImageToVideoPipeline.from_pretrained(VIDEO_MODEL_ID,
                 transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
@@ -72,14 +71,30 @@ def initialize_video_pipeline():
                 torch_dtype=torch.bfloat16,
             ).to('cuda')
             
-            # Clear memory
-            for i in range(3):
-                gc.collect()
-                torch.cuda.synchronize()
-                torch.cuda.empty_cache()
+            # Clear memory after loading
+            gc.collect()
+            torch.cuda.empty_cache()
+            
+            # Load Lightning LoRA
+            try:
+                print("Loading Lightning LoRA adapter...")
+                video_pipe.transformer.load_adapter("Lightx2v/lightx2v_I2V_14B_480p_cfg_step_4", adapter_name="lightx2v")
+                video_pipe.transformer_2.load_adapter("Lightx2v/lightx2v_I2V_14B_480p_cfg_step_4", adapter_name="lightx2v_2")
+                video_pipe.transformer.set_adapters(["lightx2v"], adapter_weights=[1.0])
+                video_pipe.transformer_2.set_adapters(["lightx2v_2"], adapter_weights=[1.0])
+                print("Lightning LoRA loaded successfully")
+            except Exception as e:
+                print(f"Warning: Could not load Lightning LoRA: {e}")
+                # Continue without LoRA
+            
+            # Clear memory again
+            gc.collect()
+            torch.cuda.empty_cache()
             
-            # Optimize pipeline if module available
-            if optimize_pipeline_ is not None:
+            # Try to optimize if module available
+            try:
+                from optimization import optimize_pipeline_
+                print("Optimizing pipeline...")
                 optimize_pipeline_(video_pipe,
                     image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
                     prompt='prompt',
@@ -87,11 +102,19 @@ def initialize_video_pipeline():
                     width=LANDSCAPE_WIDTH,
                     num_frames=MAX_FRAMES_MODEL,
                 )
+                print("Pipeline optimization complete")
+            except ImportError:
+                print("Optimization module not found, running without optimization")
+            except Exception as e:
+                print(f"Warning: Optimization failed: {e}")
                 
+            video_pipeline_ready = True
             print("Video pipeline initialized successfully!")
+            
         except Exception as e:
             print(f"Error initializing video pipeline: {e}")
             video_pipe = None
+            video_pipeline_ready = False
 
 # ===========================
 # Image Processing Functions
@@ -647,14 +670,13 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
 
 # Launch
 if __name__ == "__main__":
-    # Try to initialize video pipeline on startup
-    try:
-        initialize_video_pipeline()
-    except:
-        print("Video pipeline initialization deferred to first use")
+    # Don't initialize video pipeline on startup to avoid blocking
+    print("Starting application...")
+    print("Note: Video pipeline will initialize on first use")
     
     demo.launch(
         share=True,
         server_name="0.0.0.0",
-        server_port=7860
+        server_port=7860,
+        show_error=True
     )