Update app_lora.py

app_lora.py

@@ -20,16 +20,17 @@ 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 ---
+# --- Model Loading at Startup (Your Correct Method) ---
+# This loads the entire model into GPU VRAM when the Space starts.
+# This is correct for your H200 hardware to ensure fast inference.
 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
+    MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
 )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
-pipe.enable_model_cpu_offload()
+pipe.to("cuda")
 
-# LoRA Loading
 try:
     causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
     print("✅ LoRA downloaded to:", causvid_path)
@@ -46,74 +47,65 @@ 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, MIN_FRAMES, MAX_FRAMES = 24, 8, 240
+FIXED_FPS, MIN_FRAMES, MAX_FRAMES = 24, 8, 81
 default_prompt = "make this image come alive, cinematic motion, smooth animation"
-default_neg_prompt = "static, blurry, watermark, text, signature, ugly, deformed"
+default_neg_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"
 
-# --- 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)
+# This function correctly provides a static duration to the decorator at startup.
+def get_duration(steps, duration_seconds):
+    if steps > 4 and duration_seconds > 2: return 90
+    if steps > 4 or duration_seconds > 2: return 75
+    return 60
+
+@spaces.GPU(duration=60) # Default duration, the get_duration logic inside the function is not effective for the decorator itself
 def generate_video(input_image, prompt, height, width,
                    negative_prompt, duration_seconds,
                    guidance_scale, steps, seed, randomize_seed,
                    progress=gr.Progress(track_tqdm=True)):
-
+    
     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)
+    
+    # Using a robust frame calculation to prevent potential model errors
     raw_frames = int(round(duration_seconds * FIXED_FPS))
     num_frames = ((raw_frames - 1) // 4) * 4 + 1
     num_frames = np.clip(num_frames, MIN_FRAMES, MAX_FRAMES)
-
-    if num_frames > 120 and max(target_h, target_w) > 768:
-        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)
 
-    try:
-        torch.cuda.empty_cache()
-        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
-            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 as e:
-        raise gr.Error("Out of GPU memory. Try reducing duration or resolution.")
-    except Exception as e:
-        raise gr.Error(f"Generation failed: {e}")
-    finally:
-        torch.cuda.empty_cache()
+    with torch.inference_mode():
+        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)
+        ).frames[0]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
+        # Using a more robust video exporter for better quality and compression
         import imageio
         writer = imageio.get_writer(video_path, fps=FIXED_FPS, codec='libx264',
                                    pixelformat='yuv420p', quality=8)
         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("# Wan 2.1 I2V FusionX-LoRA")
+    gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) fusionx-lora")
+    gr.Markdown("Note: The Space will restart after a period of inactivity, causing a one-time long load.")
     
     with gr.Row():
         with gr.Column():
             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)")
+            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):
                 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)
@@ -126,20 +118,16 @@ with gr.Blocks() as demo:
             gen_button = gr.Button("Generate Video", variant="primary")
         with gr.Column():
             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_upload(img):
         if img is None: return gr.update(value=DEFAULT_H), gr.update(value=DEFAULT_W)
         try:
-            w, h = img.size
-            a = h / w
-            h_new = int(np.sqrt(MAX_AREA * a))
-            w_new = int(np.sqrt(MAX_AREA / a))
+            w, h = img.size; 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)
+        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])
     
@@ -148,4 +136,4 @@ with gr.Blocks() as demo:
     gen_button.click(fn=generate_video, inputs=inputs, outputs=outputs)
 
 if __name__ == "__main__":
-    demo.queue(max_size=3).launch()
+    demo.queue().launch()