import os import sys import gc import random import spaces import gradio as gr import torch import numpy as np from PIL import Image, ImageOps from huggingface_hub import snapshot_download sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import wan from wan.configs import WAN_CONFIGS, MAX_AREA_CONFIGS from wan.utils.utils import cache_video # ============================================================ # Wan 2.1 Image-to-Video 14B (720P checkpoint, also serves 480p) # ============================================================ REPO_ID = "Wan-AI/Wan2.1-I2V-14B-720P" TASK_NAME = "i2v-14B" cfg = WAN_CONFIGS[TASK_NAME] print(f"Downloading/loading {REPO_ID}...") ckpt_dir = snapshot_download(REPO_ID) print(f"Model files: {ckpt_dir}") FPS = cfg.sample_fps # Wan 2.1 is trained at 16 fps MAX_FRAMES = 81 # ~5 s, the model's native clip length # Exact native resolutions (Wan uses "width*height" keys) QUALITY_PRESETS = { "720p (HD)": {"landscape": (1280, 720), "portrait": (720, 1280)}, "480p (Fast)": {"landscape": (832, 480), "portrait": (480, 832)}, } DEFAULT_QUALITY = "720p (HD)" print("Initializing WanI2V 14B...") device_id = 0 if torch.cuda.is_available() else -1 pipeline = wan.WanI2V( config=cfg, checkpoint_dir=ckpt_dir, device_id=device_id, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, t5_cpu=True, init_on_cpu=True, ) print("Wan 2.1 I2V 14B initialized.") def pick_resolution(image: Image.Image, quality: str, orientation: str): preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS[DEFAULT_QUALITY]) if orientation == "Auto (match image)": key = "portrait" if image.height > image.width else "landscape" elif orientation.startswith("Portrait"): key = "portrait" else: key = "landscape" return preset[key] def crisp_fit(image: Image.Image, width: int, height: int) -> Image.Image: """Center-crop to the exact target aspect ratio, then resize with high-quality Lanczos so nothing is stretched or blurred.""" return ImageOps.fit(image, (width, height), method=Image.LANCZOS, centering=(0.5, 0.5)) def gpu_duration(image, prompt, quality, orientation, duration, steps, *args, **kwargs): # 720p on 14B is heavy; give ZeroGPU enough time. per_step = 9 if str(quality).startswith("720") else 4 return min(int(60 + per_step * int(steps) * (float(duration) / 5.0 + 0.5)), 600) @spaces.GPU(duration=gpu_duration) def generate_video( image, prompt, quality, orientation, duration, sampling_steps, guide_scale, seed, progress=gr.Progress(track_tqdm=True), ): if image is None: raise gr.Error("Please upload an image.") if not prompt or not prompt.strip(): raise gr.Error("Please enter a motion prompt.") seed = random.randint(0, 2**31 - 1) if int(seed) == -1 else int(seed) if isinstance(image, np.ndarray): image = Image.fromarray(image) image = image.convert("RGB") width, height = pick_resolution(image, quality, orientation) input_image = crisp_fit(image, width, height) # Frame count must be 4n+1 for Wan num_frames = int(round(float(duration) * FPS)) num_frames = ((num_frames - 1) // 4) * 4 + 1 num_frames = int(np.clip(num_frames, 17, MAX_FRAMES)) max_area = MAX_AREA_CONFIGS[f"{width}*{height}"] shift = 5.0 if min(width, height) >= 720 else 3.0 # Wan-recommended per resolution print(f"Generating {width}x{height}, {num_frames} frames @ {FPS}fps, seed={seed}") video_tensor = pipeline.generate( input_prompt=prompt, img=input_image, max_area=max_area, frame_num=num_frames, shift=shift, sample_solver="unipc", sampling_steps=int(sampling_steps), guide_scale=float(guide_scale), seed=seed, offload_model=True, ) video_path = cache_video( tensor=video_tensor[None], save_file=None, fps=FPS, nrow=1, normalize=True, value_range=(-1, 1), ) del video_tensor gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() return video_path, f"{width}x{height} - {num_frames} frames - seed {seed}" # ============================================================ # USER INTERFACE (GRADIO) # ============================================================ with gr.Blocks(title="My Private AI Video Studio") as demo: gr.Markdown("# My Private Wan 2.1 Video Generator\nAnimate any image in crisp 720p HD.") with gr.Row(): with gr.Column(): image_input = gr.Image(type="pil", label="Upload Your Image") prompt_input = gr.Textbox( label="Motion Prompt", placeholder="e.g., slow camera zoom, water flowing smoothly, cinematic motion", lines=3, ) quality_input = gr.Radio( choices=list(QUALITY_PRESETS.keys()), value=DEFAULT_QUALITY, label="Quality", ) orientation_input = gr.Radio( choices=["Auto (match image)", "Landscape 16:9", "Portrait 9:16"], value="Auto (match image)", label="Aspect Ratio", ) duration_input = gr.Slider(1.0, 5.0, step=0.5, value=3.0, label="Duration (sec)") steps_input = gr.Slider(10, 40, step=1, value=25, label="Steps (higher = sharper, slower)") guide_input = gr.Slider(1.0, 10.0, step=0.5, value=5.0, label="Prompt Strength") seed_input = gr.Number(value=-1, precision=0, label="Seed (-1 = Random)") generate_button = gr.Button("Animate Now", variant="primary") with gr.Column(): video_output = gr.Video(label="Your Generated Video", autoplay=True) info_output = gr.Textbox(label="Output Info", interactive=False) generate_button.click( fn=generate_video, inputs=[image_input, prompt_input, quality_input, orientation_input, duration_input, steps_input, guide_input, seed_input], outputs=[video_output, info_output], ) if __name__ == "__main__": demo.queue(max_size=3) demo.launch()