app.py

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 spaces
from huggingface_hub import hf_hub_download
import logging
import numpy as np
from PIL import Image

# --- Global Model Loading & LoRA Handling ---
MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
LORA_REPO_ID = "Kijai/WanVideo_comfy"
LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Model Loading ---
logger.info(f"Loading Image Encoder for {MODEL_ID}...")
image_encoder = CLIPVisionModel.from_pretrained(
    MODEL_ID,
    subfolder="image_encoder",
    torch_dtype=torch.float32 # Using float32 for image encoder as sometimes bfloat16/float16 can be problematic
)

logger.info(f"Loading VAE for {MODEL_ID}...")
vae = AutoencoderKLWan.from_pretrained(
    MODEL_ID,
    subfolder="vae",
    torch_dtype=torch.float32 # Using float32 for VAE for precision
)
logger.info(f"Loading Pipeline {MODEL_ID}...")
pipe = WanImageToVideoPipeline.from_pretrained(
    MODEL_ID,
    vae=vae,
    image_encoder=image_encoder,
    torch_dtype=torch.bfloat16 # Main pipeline can use bfloat16 for speed/memory
)
flow_shift = 8.0
pipe.scheduler = UniPCMultistepScheduler.from_config(
    pipe.scheduler.config, flow_shift=flow_shift
)
logger.info("Moving pipeline to CUDA...")
pipe.to("cuda")

# --- LoRA Loading ---
logger.info(f"Downloading LoRA {LORA_FILENAME} from {LORA_REPO_ID}...")
causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)

logger.info("Loading LoRA weights...")
pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
logger.info("Setting LoRA adapter...")
pipe.set_adapters(["causvid_lora"], adapter_weights=[1.0])

# --- Constants for Dimension Calculation ---
MOD_VALUE = 32
MOD_VALUE_H = MOD_VALUE_W = MOD_VALUE

DEFAULT_H_SLIDER_VALUE = 512
DEFAULT_W_SLIDER_VALUE = 896

# New fixed max_area for the calculation formula
NEW_FORMULA_MAX_AREA = float(480 * 832)

SLIDER_MIN_H = 128
SLIDER_MAX_H = 896
SLIDER_MIN_W = 128
SLIDER_MAX_W = 896

def _calculate_new_dimensions_wan(pil_image: Image.Image, mod_val: int, calculation_max_area: float,
                                 min_slider_h: int, max_slider_h: int,
                                 min_slider_w: int, max_slider_w: int,
                                 default_h: int, default_w: int) -> tuple[int, int]:
    orig_w, orig_h = pil_image.size

    if orig_w <= 0 or orig_h <= 0: # Changed to <= 0 for robustness
        logger.warning(f"Uploaded image has non-positive width or height ({orig_w}x{orig_h}). Using default slider dimensions.")
        return default_h, default_w

    aspect_ratio = orig_h / orig_w
    
    sqrt_h_term = np.sqrt(calculation_max_area * aspect_ratio)
    sqrt_w_term = np.sqrt(calculation_max_area / aspect_ratio)

    calc_h = round(sqrt_h_term) // mod_val * mod_val
    calc_w = round(sqrt_w_term) // mod_val * mod_val

    calc_h = mod_val if calc_h < mod_val else calc_h
    calc_w = mod_val if calc_w < mod_val else calc_w

    effective_min_h = min_slider_h
    effective_min_w = min_slider_w

    effective_max_h_from_slider = (max_slider_h // mod_val) * mod_val
    effective_max_w_from_slider = (max_slider_w // mod_val) * mod_val
    
    new_h = int(np.clip(calc_h, effective_min_h, effective_max_h_from_slider))
    new_w = int(np.clip(calc_w, effective_min_w, effective_max_w_from_slider))

    logger.info(f"Auto-dim: Original {orig_w}x{orig_h} (AR: {aspect_ratio:.2f}). Max Area for calc: {calculation_max_area}.")
    logger.info(f"Auto-dim: Sqrt terms HxW: {sqrt_h_term:.0f}x{sqrt_w_term:.0f}. Calculated (round(sqrt_term)//{mod_val}*{mod_val}): {calc_h}x{calc_w}.")
    logger.info(f"Auto-dim: Clamped HxW: {new_h}x{new_w} (Effective H_range:[{effective_min_h}-{effective_max_h_from_slider}], Effective W_range:[{effective_min_w}-{effective_max_w_from_slider}]).")

    return new_h, new_w

def handle_image_upload_for_dims_wan(uploaded_pil_image: Image.Image | None, current_h_val: int, current_w_val: int):
    if uploaded_pil_image is None:
        logger.info("Image cleared. Resetting dimensions to default slider values.")
        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, # Use the globally defined max_area for the new formula
            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:
        logger.error(f"Error auto-adjusting H/W from image: {e}", exc_info=True)
        # Fallback to default slider values on error, as in the original code
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)


# --- Gradio Interface Function ---
@spaces.GPU
def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
                   height: int, width: int, duration_seconds: float, # Changed from num_frames
                   guidance_scale: float, steps: int,
                   progress=gr.Progress(track_tqdm=True)): # Removed fps_for_conditioning_and_export
    if input_image is None:
        raise gr.Error("Please upload an input image.")

    # Constants for frame calculation
    FIXED_FPS = 24
    MIN_FRAMES_MODEL = 8  # Based on original num_frames_input slider min
    MAX_FRAMES_MODEL = 81 # Based on original num_frames_input slider max

    logger.info("Starting video generation...")
    logger.info(f"  Input Image: Uploaded (Original size: {input_image.size if input_image else 'N/A'})")
    logger.info(f"  Prompt: {prompt}")
    logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
    logger.info(f"  Target Output Height: {height}, Target Output Width: {width}")

    target_height = int(height)
    target_width = int(width)
    # duration_seconds is already float
    guidance_scale_val = float(guidance_scale)
    steps_val = int(steps)

    # Calculate number of frames based on duration and fixed FPS
    num_frames_for_pipeline = int(round(duration_seconds * FIXED_FPS))
    # Clamp num_frames to be within model's supported range
    num_frames_for_pipeline = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, num_frames_for_pipeline))
    # Ensure at least MIN_FRAMES_MODEL if rounding leads to a very small number (or zero)
    if num_frames_for_pipeline < MIN_FRAMES_MODEL:
        num_frames_for_pipeline = MIN_FRAMES_MODEL

    logger.info(f"  Duration: {duration_seconds:.1f}s, Fixed FPS (conditioning & export): {FIXED_FPS}")
    logger.info(f"  Calculated Num Frames: {num_frames_for_pipeline} (clamped to [{MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL}])")
    logger.info(f"  Guidance Scale: {guidance_scale_val}, Steps: {steps_val}")


    # Ensure dimensions are compatible.
    if target_height % MOD_VALUE_H != 0:
        logger.warning(f"Height {target_height} is not a multiple of {MOD_VALUE_H}. Adjusting...")
        target_height = (target_height // MOD_VALUE_H) * MOD_VALUE_H
    if target_width % MOD_VALUE_W != 0:
        logger.warning(f"Width {target_width} is not a multiple of {MOD_VALUE_W}. Adjusting...")
        target_width = (target_width // MOD_VALUE_W) * MOD_VALUE_W

    target_height = max(MOD_VALUE_H, target_height if target_height > 0 else MOD_VALUE_H)
    target_width = max(MOD_VALUE_W, target_width if target_width > 0 else MOD_VALUE_W)


    resized_image = input_image.resize((target_width, target_height))
    logger.info(f"  Input image resized to: {resized_image.size} for pipeline input.")

    with torch.inference_mode():
        output_frames_list = pipe(
            image=resized_image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=target_height,
            width=target_width,
            num_frames=num_frames_for_pipeline, # Use calculated and clamped num_frames
            guidance_scale=guidance_scale_val,
            num_inference_steps=steps_val,
            generator=torch.Generator(device="cuda").manual_seed(0)
        ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    export_to_video(output_frames_list, video_path, fps=FIXED_FPS) # Use fixed FPS for export
    logger.info(f"Video successfully generated and saved to {video_path}")
    return video_path

# --- Gradio UI Definition ---
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"
penguin_image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"

with gr.Blocks() as demo:
    gr.Markdown(f"""
    # Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA
    """)
    with gr.Row():
        with gr.Column():
            input_image_component = gr.Image(type="pil", label="Input Image (will be resized to target H/W)")
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
            duration_seconds_input = gr.Slider(minimum=0.4, maximum=3.3, step=0.1, value=1.7, label="Duration (seconds)", info="The CausVid LoRA was trained on 24fps, Wan has 81 maximum frames limit, limiting the maximum to 3.3s")
            
            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt_input = gr.Textbox(
                    label="Negative Prompt (Optional)",
                    value=default_negative_prompt,
                    lines=3
                )
                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") 
                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", interactive=False)

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

    inputs_for_click_and_examples = [
        input_image_component,
        prompt_input,
        negative_prompt_input,
        height_input,
        width_input,
        duration_seconds_input,
        guidance_scale_input,
        steps_slider
    ]

    generate_button.click(
        fn=generate_video,
        inputs=inputs_for_click_and_examples,
        outputs=video_output
    )

    gr.Examples(
        examples=[
            [penguin_image_url, "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, 896, 512, 2, 1.0, 4],
            ["https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0001.jpg", "the frog jumps around", default_negative_prompt, 448, 832, 2, 1.0, 4],
        ],
        inputs=inputs_for_click_and_examples, 
        outputs=video_output,
        fn=generate_video,
        cache_examples="lazy"
    )

if __name__ == "__main__":
    demo.queue().launch(share=True, debug=True)