app_lora.py

import spaces
import torch
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
from diffusers.utils import export_to_video
from transformers import CLIPVisionModel
import gradio as gr
import tempfile
import os
import subprocess

from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import random

import warnings
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 (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.bfloat16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
pipe.to("cuda")

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:
    print(f"❌ Error during LoRA loading: {e}")

# --- Constants ---
MOD_VALUE = 32
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, MIN_FRAMES, MAX_FRAMES = 24, 8, 81
default_prompt = "make this image come alive, cinematic motion, smooth animation"
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"

# 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)
    
    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)

    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

with gr.Blocks() as demo:
    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)")
            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)
                rand_seed_comp = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    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_comp = gr.Video(label="Generated Video", autoplay=True, interactive=False)

    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))
            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])
    
    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().launch()