app.py

# PyTorch 2.8 (temporary hack)
import os
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')

# Actual demo code
import spaces
import torch
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.utils.export_utils import export_to_video
import gradio as gr
import tempfile
import numpy as np
from PIL import Image
import random
import gc
from optimization import optimize_pipeline_
from deep_translator import GoogleTranslator

def get_translation(context_custom, language = 'en'):
    translated = GoogleTranslator(source = 'auto', target = language).translate(text = context_custom)
    return translated

MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"

LANDSCAPE_WIDTH = 832
LANDSCAPE_HEIGHT = 480
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81

MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)


pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID,
    transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
        subfolder='transformer',
        torch_dtype=torch.bfloat16,
        device_map='cuda',
    ),
    transformer_2=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
        subfolder='transformer_2',
        torch_dtype=torch.bfloat16,
        device_map='cuda',
    ),
    torch_dtype=torch.bfloat16,
).to('cuda')

for i in range(3): 
    gc.collect()
    torch.cuda.synchronize() 
    torch.cuda.empty_cache()

optimize_pipeline_(pipe,
    image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
    prompt='prompt',
    height=LANDSCAPE_HEIGHT,
    width=LANDSCAPE_WIDTH,
    num_frames=MAX_FRAMES_MODEL,
)


default_prompt_i2v = "Энэ зургийг амьд, киноны хөдөлгөөн, жигд хөдөлгөөнт болгох"
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"


def resize_image(image: Image.Image) -> Image.Image:
    if image.height > image.width:
        transposed = image.transpose(Image.Transpose.ROTATE_90)
        resized = resize_image_landscape(transposed)
        return resized.transpose(Image.Transpose.ROTATE_270)
    return resize_image_landscape(image)


def resize_image_landscape(image: Image.Image) -> Image.Image:
    target_aspect = LANDSCAPE_WIDTH / LANDSCAPE_HEIGHT
    width, height = image.size
    in_aspect = width / height
    if in_aspect > target_aspect:
        new_width = round(height * target_aspect)
        left = (width - new_width) // 2
        image = image.crop((left, 0, left + new_width, height))
    else:
        new_height = round(width / target_aspect)
        top = (height - new_height) // 2
        image = image.crop((0, top, width, top + new_height))
    return image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)

def get_duration(
    input_image,
    prompt,
    steps,
    negative_prompt,
    duration_seconds,
    guidance_scale,
    guidance_scale_2,
    seed,
    randomize_seed,
    progress,
):
    return int(steps) * 15

@spaces.GPU(duration=get_duration)
def generate_video(
    input_image,
    prompt,
    steps = 4,
    negative_prompt=default_negative_prompt,
    duration_seconds = MAX_DURATION,
    guidance_scale = 1,
    guidance_scale_2 = 1,    
    seed = 42,
    randomize_seed = False,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate a video from an input image using the Wan 2.2 14B I2V model with Lightning LoRA.
    
    This function takes an input image and generates a video animation based on the provided
    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Image-to-Video model in with Lightning LoRA
    for fast generation in 4-8 steps.
    
    Args:
        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
        prompt (str): Text prompt describing the desired animation or motion.
        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
            Defaults to 4. Range: 1-30.
        negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
            Defaults to default_negative_prompt (contains unwanted visual artifacts).
        duration_seconds (float, optional): Duration of the generated video in seconds.
            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
            Defaults to 1.0. Range: 0.0-20.0.
        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
            Defaults to 1.0. Range: 0.0-20.0.
        seed (int, optional): Random seed for reproducible results. Defaults to 42.
            Range: 0 to MAX_SEED (2147483647).
        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
            Defaults to False.
        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
    
    Returns:
        tuple: A tuple containing:
            - video_path (str): Path to the generated video file (.mp4)
            - current_seed (int): The seed used for generation (useful when randomize_seed=True)
    
    Raises:
        gr.Error: If input_image is None (no image uploaded).
    
    Note:
        - The function automatically resizes the input image to the target dimensions
        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
        - The function uses GPU acceleration via the @spaces.GPU decorator
        - Generation time varies based on steps and duration (see get_duration function)
    """
    if input_image is None:
        raise gr.Error("Please upload an input image.")
    
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    resized_image = resize_image(input_image)
    prompt = get_translation(prompt)
    output_frames_list = pipe(
        image=resized_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=resized_image.height,
        width=resized_image.width,
        num_frames=num_frames,
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        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

with gr.Blocks() as demo:
    gr.Markdown("# Fast Wan 2.2 I2V (14B) with Lightning LoRA")
    gr.Markdown("Run Wan 2.2, with [Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 хэмжигдэхүүн ба AoT эмхэтгэл - 🧨 диффузор болон ZeroGPU⚡️-тай нийцдэг")
    with gr.Row():
        with gr.Column():
            input_image_component = gr.Image(type="pil", label="Зураг оруулах (auto-resized to target H/W)")
            prompt_input = gr.Textbox(label="Текст дүрслэл", value=default_prompt_i2v)
            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
            
            with gr.Accordion("Нарийвчилсан тохиргоо", open=False):
                negative_prompt_input = gr.Textbox(label="Сөрөг сануулга", 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)
                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Тооцооллын шатууд") 
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Удирдамжийн масштаб - дуу чимээ ихтэй үе шат")
                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Удирдамжийн хуваарь 2 - дуу чимээ багатай үе шат")

            generate_button = gr.Button("Видео үүсгэх", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="Үүсгэсэн видео", autoplay=True, interactive=False)
    
    ui_inputs = [
        input_image_component, prompt_input, steps_slider,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox
    ]
    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])

    gr.Examples(
        examples=[ 
            [
                "wan_i2v_input.JPG",
                "POV селфи видео, далайн эрэг дээр зогсож буй нарны шилтэй цагаан муур, тайван инээмсэглэл, ард нь халуун орны далайн эрэг (цэвэр ус, ногоон толгод, үүлтэй хөх тэнгэр). Далайд унасан муур, бөмбөлөг, нарны туяа бүхий камер усан доор живж байна. Усан доорх муурны царайг товчхон харж, дараа нь муур дахин гарч, сельфи зураг авалтаа үргэлжлүүлж, зуны амралтын хөгжилтэй уур амьсгал.",
                4,
            ],
            [
                "wan22_input_2.jpg",
                "Цагаан сансрын хувцас өмссөн сансрын нисэгчид сарыг харайх өвөрмөц хөдөлгөөнтэй хөлөг онгоцон дээр гарч ирэхэд сарны гялалзсан машин зүүнээс баруун тийш гулсаж, сарны тоосыг өшиглөж байна. Алс холын цаана VTOL хөлөг шууд доошоо бууж, гадаргуу дээр чимээгүйхэн бууна. Бүхэл бүтэн үзэгдлийн туршид эфирийн аврора бореалисын туузууд ододоор дүүрэн тэнгэрт бүжиглэж, сарны ландшафтыг өөр ертөнцийн ид шидийн туяанд шингээсэн ногоон, хөх, нил ягаан гэрлийн гялалзсан хөшгийг цутгаж байна.",
                4,
            ],
            [
                "kill_bill.jpeg",
                "Ума Турманы дүр болох Беатрикс Киддо сахлын хурц ирээ киноны гэрэлтүүлэгт тогтвортой байлгадаг. Гэнэт өнгөлсөн ган нь халсан металл бүтцийн бүрэн бүтэн байдлаа алдаж эхэлдэг шиг зөөлөрч, гажиж эхэлдэг. Хутганы төгс ирмэг нь аажмаар мурийж, унждаг бөгөөд хайлсан ган нь метал гялбаагаа хадгалж, мөнгөн устай урсаж доошоо урсаж эхэлдэг. Өөрчлөлт эхлээд маш нарийн эхэлдэг - ир нь бага зэрэг нугалж, дараа нь метал улам шингэн болж хувирах тусам хурдасдаг. Цоолбор нүд нь аажмаар нарийсч, үхлийн үр дагаварт хүргэхгүй, харин төөрөгдөл, түгшүүртэйгээр зэвсгээ нүднийх нь өмнө уусч байгааг хараад камер нүүрэн дээр нь тогтвортой байдаг. Энэ боломжгүй өөрчлөлтийг харахад түүний амьсгал бага зэрэг түргэсдэг. Хайлалтыг эрчимжүүлж, катанагийн төгс хэлбэр улам хийсвэр болж, атгаснаас нь шингэн мөнгөн ус шиг дуслах болно. Хайлсан дуслууд нь зөөлөн металлын нөлөөгөөр газарт унадаг. Өс хонзонгийн домогт хөгжмийн зэмсэг нь гарт нь шингэрч, хамгаалалтгүй, баримжаагаа алдахад түүний нүүр царай тайван бэлэн байдлаас эргэлзэх, санаа зоволт руу шилждэг.",
                6,
            ],
            [
                "character.jpg",
                "Хөөрхөн мэдээний сурвалжлагч эрч хүчтэйгээр ярьна.",
                4,
            ],
        ],
        inputs=[input_image_component, prompt_input, steps_slider], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
    )

if __name__ == "__main__":
    demo.queue().launch(mcp_server=True)