Update app_lora.py

app_lora.py

@@ -14,180 +14,138 @@ 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"
 
 # --- 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(
     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
 pipe.enable_model_cpu_offload()
 
-# THE FIX: These two lines caused the original error and have been removed.
-# pipe.enable_vae_slicing()
-# pipe.enable_vae_tiling()
-
+# LoRA Loading
 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()
+    print(f"❌ Error during LoRA loading: {e}")
 
-# --- Constants and Helper Functions ---
+# --- Constants ---
 MOD_VALUE = 32
-DEFAULT_H_SLIDER_VALUE = 640
-DEFAULT_W_SLIDER_VALUE = 1024
-NEW_FORMULA_MAX_AREA = 640.0 * 1024.0
-
+DEFAULT_H, DEFAULT_W = 640, 1024
+MAX_AREA = DEFAULT_H * DEFAULT_W
 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
-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 get_duration(duration_seconds):
-    if duration_seconds > 7: return 180
-    if duration_seconds > 5: return 120
-    if duration_seconds > 3: return 90
-    return 60
-
-# --- The Main Generation Function ---
-# The @spaces.GPU decorator is correctly placed here.
-@spaces.GPU(duration=60)
+FIXED_FPS, MIN_FRAMES, MAX_FRAMES = 24, 8, 240
+default_prompt = "make this image come alive, cinematic motion, smooth animation"
+default_neg_prompt = "static, blurry, watermark, text, signature, ugly, deformed"
+
+# --- Main Generation Function ---
+# THE FIX: Set a generous, FIXED duration for the decorator. 180 seconds (3 minutes)
+# should be enough for the longest video generation.
+@spaces.GPU(duration=180)
 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,
+                   negative_prompt, duration_seconds,
+                   guidance_scale, steps, seed, randomize_seed,
                    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)
-
     raw_frames = int(round(duration_seconds * FIXED_FPS))
     num_frames = ((raw_frames - 1) // 4) * 4 + 1
-    num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+    num_frames = np.clip(num_frames, MIN_FRAMES, MAX_FRAMES)
 
     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")
+        scale = 768 / max(target_h, target_w)
+        target_h = max(MOD_VALUE, int(target_h * scale) // MOD_VALUE * MOD_VALUE)
+        target_w = max(MOD_VALUE, int(target_w * scale) // MOD_VALUE * MOD_VALUE)
+        gr.Info(f"Reduced resolution to {target_w}x{target_h} for long video.")
 
     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)
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
     try:
+        torch.cuda.empty_cache()
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
-            output_frames_list = pipe(
+            frames = 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:
-        raise gr.Error("Out of GPU memory. Try reducing the duration or resolution.")
+    except torch.cuda.OutOfMemoryError as e:
+        raise gr.Error("Out of GPU memory. Try reducing duration or resolution.")
     except Exception as e:
-        raise gr.Error(f"Generation failed: {str(e)}")
+        raise gr.Error(f"Generation failed: {e}")
     finally:
-        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
         import imageio
         writer = imageio.get_writer(video_path, fps=FIXED_FPS, codec='libx264',
                                    pixelformat='yuv420p', quality=8)
-        for frame in output_frames_list:
+        for frame in frames:
             writer.append_data(np.array(frame))
         writer.close()
-
+    
     return video_path, current_seed
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) FusionX-LoRA")
+    gr.Markdown("# Wan 2.1 I2V FusionX-LoRA")
     
     with gr.Row():
         with gr.Column():
-            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=0.3, maximum=10.0, step=0.1, value=2, label="Duration (seconds)")
+            input_image_comp = gr.Image(type="pil", label="Input Image")
+            prompt_comp = gr.Textbox(label="Prompt", value=default_prompt)
+            duration_comp = gr.Slider(minimum=round(MIN_FRAMES/FIXED_FPS, 1), maximum=round(MAX_FRAMES/FIXED_FPS, 1), step=0.1, value=2, label="Duration (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)
-                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
+                neg_prompt_comp = gr.Textbox(label="Negative Prompt", value=default_neg_prompt, lines=3)
+                seed_comp = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
+                rand_seed_comp = gr.Checkbox(label="Randomize seed", value=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="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")
+                    height_comp = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H, label="Height")
+                    width_comp = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W, label="Width")
+                steps_comp = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Steps")
+                guidance_comp = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="CFG Scale", visible=False)
+            gen_button = gr.Button("Generate Video", variant="primary")
         with gr.Column():
-            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-            gr.Markdown("### Tips:\n- For videos > 5s, consider lower resolutions.\n- 4-8 steps is often optimal.")
+            video_comp = gr.Video(label="Generated Video", autoplay=True, interactive=False)
+            gr.Markdown("### Tips:\n- For long 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)
+    def handle_upload(img):
+        if img is None: return gr.update(value=DEFAULT_H), gr.update(value=DEFAULT_W)
         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])
+            a = h / w
+            h_new = int(np.sqrt(MAX_AREA * a))
+            w_new = int(np.sqrt(MAX_AREA / a))
+            h_final = max(MOD_VALUE, h_new // MOD_VALUE * MOD_VALUE)
+            w_final = max(MOD_VALUE, w_new // MOD_VALUE * MOD_VALUE)
+            return gr.update(value=h_final), gr.update(value=w_final)
+        except Exception:
+            return gr.update(value=DEFAULT_H), gr.update(value=DEFAULT_W)
+
+    input_image_comp.upload(handle_upload, inputs=input_image_comp, outputs=[height_comp, width_comp])
     
-    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])
-
-    # 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"
-    # )
+    inputs = [input_image_comp, prompt_comp, height_comp, width_comp, neg_prompt_comp, duration_comp, guidance_comp, steps_comp, seed_comp, rand_seed_comp]
+    outputs = [video_comp, seed_comp]
+    gen_button.click(fn=generate_video, inputs=inputs, outputs=outputs)
 
 if __name__ == "__main__":
     demo.queue(max_size=3).launch()