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, list_repo_files
import numpy as np
from PIL import Image
import random
import os

# ----------------------
#
# CONFIGURATION
#
# ----------------------

MODEL_ID = "gaalos/Wan2.1-I2V-14B-720P-Diffusers-scaled"
LORA_REPO_ID = "hotdogs/wan_nsfw_lora"
MOD_VALUE = 32
FIXED_FPS = 24
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 30 * FIXED_FPS  # 30s max
MAX_SEED = np.iinfo(np.int32).max

DEFAULT_H_SLIDER_VALUE = 640
DEFAULT_W_SLIDER_VALUE = 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

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"

# ----------------------
# LOAD BASE MODEL
# ----------------------

print("Loading models...")
image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
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")
print("Models loaded successfully!")

# ----------------------
# LOAD ALL LORAS FROM "hotdogs/wan_nsfw_lora"
# ----------------------

print("Loading LoRAs...")
lora_files = [f for f in list_repo_files(LORA_REPO_ID) if f.endswith(".safetensors")]
lora_paths = {}
for file in lora_files:
    local_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=file)
    adapter_name = os.path.splitext(os.path.basename(file))[0]
    lora_paths[adapter_name] = local_path
    pipe.load_lora_weights(local_path, adapter_name=adapter_name)
print(f"Loaded {len(lora_paths)} LoRA adapters: {list(lora_paths.keys())}")

# ----------------------
# DIMENSION HELPERS
# ----------------------

def _calculate_new_dimensions(pil_image):
    """Calculate optimal dimensions based on image aspect ratio"""
    orig_w, orig_h = pil_image.size
    if orig_w <= 0 or orig_h <= 0:
        return DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE

    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 // MOD_VALUE) * MOD_VALUE))
    new_w = int(np.clip(calc_w, SLIDER_MIN_W, (SLIDER_MAX_W // MOD_VALUE) * MOD_VALUE))

    return new_h, new_w

def handle_image_upload(image, current_h, current_w):
    """Update height/width sliders when image is uploaded"""
    if 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(image)
        return gr.update(value=new_h), gr.update(value=new_w)
    except Exception as e:
        print(f"Error calculating dimensions: {e}")
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)

# ----------------------
# DURATION CONFIGURATION
# ----------------------

def get_duration(input_image, prompt, height, width, negative_prompt, duration_seconds,
                 guidance_scale, steps, seed, randomize_seed, lora_list, progress):
    """Calculate GPU duration for spaces.GPU decorator"""
    base = 60
    if duration_seconds > 10: base += 30
    if duration_seconds > 20: base += 30
    if steps > 8: base += 20
    return base

# ----------------------
# MAIN GENERATION FUNCTION
# ----------------------

@spaces.GPU(duration=get_duration)
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,
                   lora_list=None, progress=gr.Progress(track_tqdm=True)):
    """
    Generate video from image using Wan2.1 I2V pipeline with optional LoRAs
    """
    if input_image is None:
        raise gr.Error("Please upload an image.")

    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    resized_image = input_image.resize((target_w, target_h))

    selected_adapters = lora_list or []
    pipe.set_adapters(selected_adapters, adapter_weights=[0.95] * len(selected_adapters))
    pipe.fuse_lora()

    with torch.inference_mode():
        output_frames_list = 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
    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
    return video_path, current_seed

# ----------------------
# GRADIO UI - GRADIO 6 SYNTAX
# ----------------------

with gr.Blocks() as demo:
    # Header with anycoder link
    gr.Markdown(
        """
        # Wan2.1 I2V + Multi-LoRA Generator (NSFW LoRA repo)
        Generate videos using LoRAs from `hotdogs/wan_nsfw_lora` with Wan2.1. Up to **30s** per video.
        [Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input & Settings")
            image_input = gr.Image(
                type="pil", 
                label="Input Image",
                sources=["upload", "webcam", "clipboard"],
                height=300
            )
            prompt = gr.Textbox(
                label="Prompt", 
                value=default_prompt_i2v,
                placeholder="Describe the motion you want to see...",
                lines=3
            )
            duration_slider = gr.Slider(
                minimum=round(MIN_FRAMES_MODEL / FIXED_FPS, 1), 
                maximum=30.0, 
                step=0.1, 
                value=2, 
                label="Video Duration (seconds)"
            )

            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
                )
                random_seed_checkbox = gr.Checkbox(
                    label="Randomize seed", 
                    value=True
                )
                
                with gr.Row():
                    height_input = gr.Slider(
                        SLIDER_MIN_H, 
                        SLIDER_MAX_H, 
                        MOD_VALUE, 
                        value=DEFAULT_H_SLIDER_VALUE, 
                        label="Output Height"
                    )
                    width_input = gr.Slider(
                        SLIDER_MIN_W, 
                        SLIDER_MAX_W, 
                        MOD_VALUE, 
                        value=DEFAULT_W_SLIDER_VALUE, 
                        label="Output Width"
                    )
                
                steps_slider = gr.Slider(
                    1, 
                    30, 
                    step=1, 
                    value=4, 
                    label="Inference Steps"
                )
                guidance_slider = gr.Slider(
                    0.0, 
                    20.0, 
                    step=0.5, 
                    value=1.0, 
                    label="Guidance Scale", 
                    visible=False
                )
                lora_selector = gr.CheckboxGroup(
                    choices=list(lora_paths.keys()), 
                    label="Activate LoRA(s)",
                    info="Select one or more LoRAs to apply"
                )

            generate_button = gr.Button(
                "Generate Video", 
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### Output")
            video_output = gr.Video(
                label="Generated Video", 
                autoplay=True,
                height=400
            )
            seed_output = gr.Textbox(
                label="Seed used", 
                interactive=False,
                info="Use this seed to reproduce the same video"
            )

    # Event handlers
    image_input.upload(
        fn=handle_image_upload, 
        inputs=[image_input, height_input, width_input], 
        outputs=[height_input, width_input],
        api_visibility="private"
    )
    image_input.clear(
        fn=handle_image_upload, 
        inputs=[image_input, height_input, width_input], 
        outputs=[height_input, width_input],
        api_visibility="private"
    )

    inputs = [
        image_input, prompt, height_input, width_input,
        negative_prompt_input, duration_slider, guidance_slider,
        steps_slider, seed_input, random_seed_checkbox, lora_selector
    ]

    generate_button.click(
        fn=generate_video, 
        inputs=inputs, 
        outputs=[video_output, seed_output],
        api_visibility="public",
        concurrency_limit=1
    )

# ----------------------
# LAUNCH - GRADIO 6 SYNTAX
# ----------------------

if __name__ == "__main__":
    demo.queue().launch(
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="indigo",
            neutral_hue="slate",
            font=gr.themes.GoogleFont("Inter"),
            text_size="lg",
            spacing_size="lg",
            radius_size="md"
        ).set(
            button_primary_background_fill="*primary_600",
            button_primary_background_fill_hover="*primary_700",
            block_title_text_weight="600",
            block_background_fill="*neutral_50"
        ),
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
            {"label": "Model", "url": f"https://huggingface.co/{MODEL_ID}"},
            {"label": "LoRAs", "url": f"https://huggingface.co/{LORA_REPO_ID}"}
        ],
        show_error=True,
        max_threads=40,
        share=False
    )