Update app_lora.py

app_lora.py

@@ -7,40 +7,65 @@ import tempfile
 import os
 import subprocess
 
+# The spaces library IS required for ZeroGPU.
+import spaces
 from huggingface_hub import hf_hub_download
 import numpy as np
 from PIL import Image
 import random
 
 import warnings
-warnings.filterwarnings("ignore", message=".*Attempting to use legacy OpenCV backend.*")
-warnings.filterwarnings("ignore", message=".*num_frames - 1.*")
+warnings.filterwarnings("ignore")
 
 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"
 
-# Global variable to hold the pipeline. It's initialized to None.
+# --- Global variable for the pipeline ---
+# We use a global variable to cache the model between calls.
 pipe = None
 
-def initialize_pipeline():
+# --- Constants and Helper Functions ---
+MOD_VALUE = 32
+DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE = 640, 1024
+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, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL = 24, 8, 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 get_duration(duration_seconds):
     """
-    Initializes the model pipeline on the first request.
-    This function is designed for serverless GPU environments like ZeroGPU.
+    Dynamically set the timeout for the @spaces.GPU decorator based on video length.
     """
+    if duration_seconds > 7: return 180
+    if duration_seconds > 5: return 120
+    if duration_seconds > 3: return 90
+    return 60
+
+# --- The Main GPU Function ---
+# The @spaces.GPU decorator is ESSENTIAL for ZeroGPU.
+# It tells the platform that this function needs a GPU.
+@spaces.GPU(duration=60) # Default duration, can be updated dynamically
+def generate_video(input_image, prompt, height, width,
+                   negative_prompt, duration_seconds,
+                   guidance_scale, steps, seed, randomize_seed,
+                   progress=gr.Progress(track_tqdm=True)):
+    
     global pipe
-    # The 'pipe' global variable acts as a flag. If it's not None, we've already initialized.
+    
+    # --- LAZY LOADING of the model ---
+    # This block will only run on the very first generation request.
     if pipe is None:
-        print("First time setup: Initializing model pipeline...")
-        gr.Info("Cold start: The first generation will take longer as the model is loaded.")
-
-        if not torch.cuda.is_available():
-            raise gr.Error("GPU not available. This application requires a GPU to run.")
-
+        progress(0, desc="Cold start: Initializing model...")
+        print("Cold start: Initializing model pipeline...")
+        
         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)
         
-        # All model loading happens here, when a GPU is guaranteed to be active.
         pipe = WanImageToVideoPipeline.from_pretrained(
             MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.float16
         )
@@ -49,74 +74,17 @@ def initialize_pipeline():
 
         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()
+            print("✅ LoRA loaded successfully.")
         except Exception as e:
             raise gr.Error(f"Error loading LoRA: {e}")
         
         print("✅ Pipeline initialized successfully.")
 
-# --- Constants and Helper Functions ---
-# (These are unchanged)
-MOD_VALUE = 32
-DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE = 640, 1024
-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, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL = 24, 8, 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))
-    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):
-    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 calculating new dimensions.")
-        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
-
-def export_video_with_ffmpeg(frames, output_path, fps=24):
-    try:
-        import imageio
-        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()
-    except ImportError:
-        export_to_video(frames, output_path, fps=fps)
-
-def generate_video(input_image, prompt, height, width,
-                   negative_prompt, duration_seconds,
-                   guidance_scale, steps, seed, randomize_seed,
-                   progress=gr.Progress(track_tqdm=True)):
-
-    # --- LAZY LOADING TRIGGER ---
-    # This will load the model on the first run, and do nothing on subsequent runs.
-    initialize_pipeline()
+    # Update the GPU duration based on user input for longer videos.
+    spaces.set_timeout(get_duration(duration_seconds))
     
     if input_image is None:
         raise gr.Error("Please upload an input image.")
@@ -143,7 +111,8 @@ def generate_video(input_image, prompt, height, width,
                 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)
+                generator=torch.Generator(device="cuda").manual_seed(current_seed),
+                callback_on_step_end=lambda p, s, t: progress(s/int(steps))
             ).frames[0]
     except torch.cuda.OutOfMemoryError:
         raise gr.Error("Out of GPU memory. Try reducing duration or resolution.")
@@ -152,11 +121,31 @@ def generate_video(input_image, prompt, height, width,
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
-        export_video_with_ffmpeg(output_frames_list, video_path, fps=FIXED_FPS)
+        # (Video export logic is unchanged)
+        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 UI ---
+# (Helper functions for UI are unchanged)
+def handle_image_upload_for_dims_wan(uploaded_pil_image):
+    if uploaded_pil_image is None: return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+    try:
+        orig_w, orig_h = uploaded_pil_image.size
+        aspect_ratio = orig_h / orig_w
+        calc_h = round(np.sqrt(NEW_FORMULA_MAX_AREA * aspect_ratio))
+        calc_w = round(np.sqrt(NEW_FORMULA_MAX_AREA / aspect_ratio))
+        calc_h = max(MOD_VALUE, (calc_h // MOD_VALUE) * MOD_VALUE)
+        calc_w = max(MOD_VALUE, (calc_w // MOD_VALUE) * MOD_VALUE)
+        new_h = int(np.clip(calc_h, SLIDER_MIN_H, SLIDER_MAX_H))
+        new_w = int(np.clip(calc_w, 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)
+
 with gr.Blocks() as demo:
     gr.Markdown("# Wan 2.1 I2V FusionX-LoRA (ZeroGPU Ready)")
     gr.Markdown("The first generation will be slow due to a 'cold start'. Subsequent generations will be much faster.")
@@ -186,5 +175,4 @@ with gr.Blocks() as demo:
     generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
 if __name__ == "__main__":
-    # We launch the demo unconditionally now. The GPU check is deferred until the first click.
-    demo.queue(max_size=3).launch()
+    demo.queue().launch()