Update app_lora.py

app_lora.py

@@ -1,4 +1,3 @@
-import spaces
 import torch
 from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
 from diffusers.utils import export_to_video
@@ -17,41 +16,48 @@ import warnings
 warnings.filterwarnings("ignore", message=".*Attempting to use legacy OpenCV backend.*")
 warnings.filterwarnings("ignore", message=".*num_frames - 1.*")
 
+# This decorator is specific to HuggingFace Spaces and will cause an error in other environments.
+# import spaces
+
 MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 
 LORA_REPO_ID = "vrgamedevgirl84/Wan14BT2VFusioniX"
 LORA_FILENAME = "FusionX_LoRa/Wan2.1_I2V_14B_FusionX_LoRA.safetensors"
 
 # Initialize models with proper dtype handling
-image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float16)
-vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float16)
-pipe = WanImageToVideoPipeline.from_pretrained(
-    MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.float16
-)
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
-
-# Enable memory efficient attention and CPU offloading for large videos
-pipe.enable_model_cpu_offload()
-pipe.enable_vae_slicing()
-pipe.enable_vae_tiling()
-
-try:
-    causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
-    print("✅ LoRA downloaded to:", causvid_path)
-
-    pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
-    pipe.set_adapters(["causvid_lora"], adapter_weights=[0.75])
-    pipe.fuse_lora()
-
-except Exception as e:
-    import traceback
-    print("❌ Error during LoRA loading:")
-    traceback.print_exc()
+# This section requires a GPU and CUDA to be available
+pipe = None
+if torch.cuda.is_available():
+    image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float16)
+    vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float16)
+    pipe = WanImageToVideoPipeline.from_pretrained(
+        MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.float16
+    )
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
+
+    # Enable memory efficient attention and CPU offloading for large videos
+    pipe.enable_model_cpu_offload()
+
+    try:
+        causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
+        print("✅ LoRA downloaded to:", causvid_path)
+
+        pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
+        pipe.set_adapters(["causvid_lora"], adapter_weights=[0.75])
+        pipe.fuse_lora()
+
+    except Exception as e:
+        import traceback
+        print("❌ Error during LoRA loading:")
+        traceback.print_exc()
+else:
+    print("CUDA is not available. This script requires a GPU to run.")
+
 
 MOD_VALUE = 32
 DEFAULT_H_SLIDER_VALUE = 640
 DEFAULT_W_SLIDER_VALUE = 1024
-NEW_FORMULA_MAX_AREA = 640.0 * 1024.0 
+NEW_FORMULA_MAX_AREA = 640.0 * 1024.0
 
 SLIDER_MIN_H, SLIDER_MAX_H = 128, 1024
 SLIDER_MIN_W, SLIDER_MAX_W = 128, 1024
@@ -74,16 +80,16 @@ def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
         return default_h, default_w
 
     aspect_ratio = orig_h / orig_w
-    
+
     calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
     calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
 
     calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
     calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
-    
+
     new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
     new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
-    
+
     return new_h, new_w
 
 def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
@@ -100,89 +106,62 @@ def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_
         gr.Warning("Error attempting to calculate new dimensions")
         return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
 
-def get_duration(input_image, prompt, height, width, 
-               negative_prompt, duration_seconds,
-               guidance_scale, steps,
-               seed, randomize_seed, 
-               progress):
-    # Adjust timeout based on video length and complexity
-    if duration_seconds > 7:
-        return 180  # 3 minutes for very long videos
-    elif duration_seconds > 5:
-        return 120  # 2 minutes for long videos
-    elif duration_seconds > 3:
-        return 90   # 1.5 minutes for medium videos
-    else:
-        return 60   # 1 minute for short videos
 
 def export_video_with_ffmpeg(frames, output_path, fps=24):
     """Export video using imageio if available, otherwise fall back to OpenCV"""
     try:
         import imageio
-        # Use imageio for better quality
-        writer = imageio.get_writer(output_path, fps=fps, codec='libx264', 
+        writer = imageio.get_writer(output_path, fps=fps, codec='libx264',
                                    pixelformat='yuv420p', quality=8)
         for frame in frames:
             writer.append_data(np.array(frame))
         writer.close()
         return True
     except ImportError:
-        # Fall back to OpenCV
         export_to_video(frames, output_path, fps=fps)
         return False
 
-@spaces.GPU(duration=get_duration)
-def generate_video(input_image, prompt, height, width, 
+def generate_video(input_image, prompt, height, width,
                    negative_prompt=default_negative_prompt, duration_seconds=2,
                    guidance_scale=1, steps=4,
-                   seed=42, randomize_seed=False, 
+                   seed=42, randomize_seed=False,
                    progress=gr.Progress(track_tqdm=True)):
-    
+
+    if pipe is None or not torch.cuda.is_available():
+        raise gr.Error("Pipeline not initialized or CUDA not available. Please check the console for errors.")
+
     if input_image is None:
         raise gr.Error("Please upload an input image.")
 
     target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
-    
-    # Calculate frames with proper alignment
+
     raw_frames = int(round(duration_seconds * FIXED_FPS))
-    # Ensure num_frames-1 is divisible by 4 as required by the model
     num_frames = ((raw_frames - 1) // 4) * 4 + 1
     num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
-    
-    # Additional check for very long videos
+
     if num_frames > 120:
-        # For videos longer than 5 seconds, reduce resolution to manage memory
         max_dim = max(target_h, target_w)
         if max_dim > 768:
             scale_factor = 768 / max_dim
             target_h = max(MOD_VALUE, (int(target_h * scale_factor) // MOD_VALUE) * MOD_VALUE)
             target_w = max(MOD_VALUE, (int(target_w * scale_factor) // MOD_VALUE) * MOD_VALUE)
             gr.Info(f"Reduced resolution to {target_w}x{target_h} for long video generation")
-    
+
     print(f"Generating {num_frames} frames (requested {raw_frames}) at {target_w}x{target_h}")
-    
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
+    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     resized_image = input_image.resize((target_w, target_h), Image.Resampling.LANCZOS)
 
-    # Clear GPU cache before generation
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
 
     try:
         with torch.inference_mode():
-            # Generate video with autocast for memory efficiency
             with torch.autocast("cuda", dtype=torch.float16):
                 output_frames_list = pipe(
-                    image=resized_image, 
-                    prompt=prompt, 
-                    negative_prompt=negative_prompt,
-                    height=target_h, 
-                    width=target_w, 
-                    num_frames=num_frames,
-                    guidance_scale=float(guidance_scale), 
-                    num_inference_steps=int(steps),
+                    image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
+                    height=target_h, width=target_w, num_frames=num_frames,
+                    guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
                     generator=torch.Generator(device="cuda").manual_seed(current_seed),
                     return_dict=True
                 ).frames[0]
@@ -193,47 +172,27 @@ def generate_video(input_image, prompt, height, width,
         torch.cuda.empty_cache()
         raise gr.Error(f"Generation failed: {str(e)}")
 
-    # Clear cache after generation
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
-        
-        # Export using imageio if available, otherwise OpenCV
-        used_imageio = export_video_with_ffmpeg(output_frames_list, video_path, fps=FIXED_FPS)
-        
-        # Only try FFmpeg optimization if we have a valid video file
+        export_video_with_ffmpeg(output_frames_list, video_path, fps=FIXED_FPS)
+
         if os.path.exists(video_path) and os.path.getsize(video_path) > 0:
             try:
-                # Check if ffmpeg is available
                 subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
-                
                 optimized_path = video_path + "_opt.mp4"
                 cmd = [
-                    'ffmpeg',
-                    '-y',  # Overwrite without asking
-                    '-i', video_path,  # Input file
-                    '-c:v', 'libx264',  # Codec
-                    '-pix_fmt', 'yuv420p',  # Pixel format
-                    '-profile:v', 'main',  # Compatibility profile
-                    '-level', '4.0',  # Support for higher resolutions
-                    '-movflags', '+faststart',  # Streaming optimized
-                    '-crf', '23',  # Quality level
-                    '-preset', 'medium',  # Balance between speed and compression
-                    '-maxrate', '10M',  # Max bitrate for large videos
-                    '-bufsize', '20M',  # Buffer size
-                    optimized_path
+                    'ffmpeg', '-y', '-i', video_path, '-c:v', 'libx264', '-pix_fmt', 'yuv420p',
+                    '-profile:v', 'main', '-level', '4.0', '-movflags', '+faststart', '-crf', '23',
+                    '-preset', 'medium', '-maxrate', '10M', '-bufsize', '20M', optimized_path
                 ]
-                
                 result = subprocess.run(cmd, capture_output=True, text=True)
-                
                 if result.returncode == 0 and os.path.exists(optimized_path) and os.path.getsize(optimized_path) > 0:
-                    os.unlink(video_path)  # Remove original
+                    os.unlink(video_path)
                     video_path = optimized_path
                 else:
                     print(f"FFmpeg optimization failed: {result.stderr}")
-                    
             except (subprocess.CalledProcessError, FileNotFoundError):
                 print("FFmpeg not available or optimization failed, using original export")
 
@@ -243,19 +202,17 @@ def generate_video(input_image, prompt, height, width,
 with gr.Blocks() as demo:
     gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) FusionX-LoRA")
     gr.Markdown("Generate videos up to 10 seconds long! Longer videos may use reduced resolution for stability.")
-    
+
     with gr.Row():
         with gr.Column():
             input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
             prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
             duration_seconds_input = gr.Slider(
-                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1),  # 0.3s (8 frames)
-                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1),  # 10.0s (240 frames)
-                step=0.1, 
-                value=2,  # Default 2 seconds
-                label="Duration (seconds)", 
-                info=f"Video length: {MIN_FRAMES_MODEL/FIXED_FPS:.1f}-{MAX_FRAMES_MODEL/FIXED_FPS:.1f}s. Longer videos may take more time and use more memory."
-            )            
+                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1),
+                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1),
+                step=0.1, value=2, label="Duration (seconds)",
+                info=f"Video length: {MIN_FRAMES_MODEL/FIXED_FPS:.1f}-{MAX_FRAMES_MODEL/FIXED_FPS:.1f}s."
+            )
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                 seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
@@ -263,46 +220,48 @@ with gr.Blocks() as demo:
                 with gr.Row():
                     height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
                     width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
-                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps") 
+                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
                 guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)
 
             generate_button = gr.Button("Generate Video", variant="primary")
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-            gr.Markdown("### Tips for best results:")
-            gr.Markdown("- For videos longer than 5 seconds, consider using lower resolutions (512-768px)")
-            gr.Markdown("- Clear, simple prompts often work better than complex descriptions")
-            gr.Markdown("- The model works best with 4-8 inference steps")
+            gr.Markdown("### Tips for best results:\n- For videos longer than 5 seconds, consider using lower resolutions (512-768px)\n- Clear, simple prompts often work better than complex descriptions\n- The model works best with 4-8 inference steps")
 
     input_image_component.upload(
         fn=handle_image_upload_for_dims_wan,
         inputs=[input_image_component, height_input, width_input],
         outputs=[height_input, width_input]
     )
-    
-    input_image_component.clear( 
+    input_image_component.clear(
         fn=handle_image_upload_for_dims_wan,
         inputs=[input_image_component, height_input, width_input],
         outputs=[height_input, width_input]
     )
-    
+
     ui_inputs = [
         input_image_component, prompt_input, height_input, width_input,
         negative_prompt_input, duration_seconds_input,
         guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
     ]
     generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
-
-    gr.Examples(
-        examples=[ 
-            ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
-            ["forg.jpg", "the frog jumps around", 448, 832],
-        ],
-        inputs=[input_image_component, prompt_input, height_input, width_input], 
-        outputs=[video_output, seed_input], 
-        fn=generate_video, 
-        cache_examples="lazy"
-    )
+    
+    # The example images 'peng.png' and 'forg.jpg' are not present in this environment,
+    # so the gr.Examples component is commented out to prevent errors.
+    # gr.Examples(
+    #     examples=[
+    #         ["path/to/your/peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
+    #         ["path/to/your/forg.jpg", "the frog jumps around", 448, 832],
+    #     ],
+    #     inputs=[input_image_component, prompt_input, height_input, width_input],
+    #     outputs=[video_output, seed_input],
+    #     fn=generate_video,
+    #     cache_examples="lazy"
+    # )
 
 if __name__ == "__main__":
-    demo.queue(max_size=3).launch()
+    if pipe is not None:
+        demo.queue(max_size=3).launch()
+    else:
+        gr.Blocks()._queue_closed = False # A hack to prevent Gradio from hanging
+        gr.Info("Application not started because a GPU (CUDA) is required but not found.")