Update app_lora.py

app_lora.py

@@ -22,7 +22,9 @@ 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
+# --- Model Loading at Startup ---
+# This is the correct pattern for your environment. The model is loaded once
+# when the Space starts, leading to a longer build but a fast experience for users.
 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(
@@ -30,10 +32,12 @@ pipe = WanImageToVideoPipeline.from_pretrained(
 )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
 
-# Enable memory efficient attention and CPU offloading for large videos
+# Enable memory efficient attention and CPU offloading
 pipe.enable_model_cpu_offload()
-pipe.enable_vae_slicing()
-pipe.enable_vae_tiling()
+
+# THE FIX: These two lines caused the original error and have been removed.
+# pipe.enable_vae_slicing()
+# pipe.enable_vae_tiling()
 
 try:
     causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
@@ -48,261 +52,142 @@ except Exception as e:
     print("❌ Error during LoRA loading:")
     traceback.print_exc()
 
+# --- Constants and Helper Functions ---
 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
 MAX_SEED = np.iinfo(np.int32).max
 
 FIXED_FPS = 24
-MIN_FRAMES_MODEL = 8  # Minimum 8 frames (~0.33s)
-MAX_FRAMES_MODEL = 240  # Maximum 240 frames (10 seconds at 24fps)
+MIN_FRAMES_MODEL = 8
+MAX_FRAMES_MODEL = 240
 
 default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
 default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards, watermark, text, signature"
 
 
-def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
-                                 min_slider_h, max_slider_h,
-                                 min_slider_w, max_slider_w,
-                                 default_h, default_w):
-    orig_w, orig_h = pil_image.size
-    if orig_w <= 0 or orig_h <= 0:
-        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))
+def get_duration(duration_seconds):
+    if duration_seconds > 7: return 180
+    if duration_seconds > 5: return 120
+    if duration_seconds > 3: return 90
+    return 60
 
-    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):
-    if uploaded_pil_image is None:
-        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
-    try:
-        new_h, new_w = _calculate_new_dimensions_wan(
-            uploaded_pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
-            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
-            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
-        )
-        return gr.update(value=new_h), gr.update(value=new_w)
-    except Exception as e:
-        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', 
-                                   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, 
+# --- The Main Generation Function ---
+# The @spaces.GPU decorator is correctly placed here.
+@spaces.GPU(duration=60)
+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)):
-    
+
+    spaces.set_timeout(get_duration(duration_seconds))
+
     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)
 
+    if num_frames > 120 and max(target_h, target_w) > 768:
+        scale_factor = 768 / max(target_h, target_w)
+        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")
+
+    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()
 
     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),
-                    generator=torch.Generator(device="cuda").manual_seed(current_seed),
-                    return_dict=True
-                ).frames[0]
+        with torch.inference_mode(), 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),
+                generator=torch.Generator(device="cuda").manual_seed(current_seed),
+                return_dict=True
+            ).frames[0]
     except torch.cuda.OutOfMemoryError:
-        torch.cuda.empty_cache()
         raise gr.Error("Out of GPU memory. Try reducing the duration or resolution.")
     except Exception as e:
-        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()
+    finally:
+        if torch.cuda.is_available():
+            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
-        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
-                ]
-                
-                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
-                    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")
+        import imageio
+        writer = imageio.get_writer(video_path, fps=FIXED_FPS, codec='libx264',
+                                   pixelformat='yuv420p', quality=8)
+        for frame in output_frames_list:
+            writer.append_data(np.array(frame))
+        writer.close()
 
     return video_path, current_seed
 
-# Gradio Interface
+# --- Gradio UI ---
 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)")
+            input_image_component = gr.Image(type="pil", label="Input Image")
             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."
-            )            
+            duration_seconds_input = gr.Slider(minimum=0.3, maximum=10.0, step=0.1, value=2, label="Duration (seconds)")
             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)
                 randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                 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") 
+                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label="Height")
+                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label="Width")
+                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")
-
-    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( 
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
+            gr.Markdown("### Tips:\n- For videos > 5s, consider lower resolutions.\n- 4-8 steps is often optimal.")
+
+    def handle_image_upload(img):
+        if img is None: return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+        try:
+            w, h = img.size
+            aspect = h / w
+            calc_h = round(np.sqrt(NEW_FORMULA_MAX_AREA * aspect))
+            calc_w = round(np.sqrt(NEW_FORMULA_MAX_AREA / aspect))
+            new_h = int(np.clip((calc_h // MOD_VALUE) * MOD_VALUE, SLIDER_MIN_H, SLIDER_MAX_H))
+            new_w = int(np.clip((calc_w // MOD_VALUE) * MOD_VALUE, SLIDER_MIN_W, SLIDER_MAX_W))
+            return gr.update(value=new_h), gr.update(value=new_w)
+        except: return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+
+    input_image_component.upload(handle_image_upload, inputs=input_image_component, 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
-    ]
+    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 gr.Examples requires the files to be in your repo. Commenting out to prevent errors.
+    # 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"
+    # )
 
 if __name__ == "__main__":
     demo.queue(max_size=3).launch()