DmitryBaltin/VCExperiment

1"""Minimal conditioned video generation Space."""2 3import gc4import os5import random6import tempfile7from pathlib import Path8 9os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")10os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")11 12import gradio as gr13import imageio.v2 as imageio14import numpy as np15import spaces16import torch17from PIL import Image, ImageOps18from diffusers import AutoencoderKLWan, WanVACEPipeline19from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler20from diffusers.utils import export_to_video21 22MODEL_ID = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"23TARGET_WIDTH = 83224TARGET_HEIGHT = 48025OUTPUT_FPS = 1626MAX_FRAMES = 8127MAX_SEED = 2**31 - 128 29DEFAULT_NEGATIVE_PROMPT = (30    "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "31    "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "32    "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "33    "misshapen limbs, fused fingers, still picture, messy background, three legs, many people "34    "in the background, walking backwards"35)36 37print(f"[vc-experiment] Loading {MODEL_ID}...", flush=True)38vae = AutoencoderKLWan.from_pretrained(39    MODEL_ID,40    subfolder="vae",41    torch_dtype=torch.float32,42)43pipe = WanVACEPipeline.from_pretrained(44    MODEL_ID,45    vae=vae,46    torch_dtype=torch.bfloat16,47)48pipe.scheduler = UniPCMultistepScheduler.from_config(49    pipe.scheduler.config,50    flow_shift=3.0,51)52pipe.vae.enable_tiling()53pipe.to("cuda")54print("[vc-experiment] Model ready.", flush=True)55 56 57def _to_rgb_uint8(frame: np.ndarray) -> np.ndarray:58    if frame.ndim == 2:59        frame = np.repeat(frame[..., None], 3, axis=2)60    elif frame.ndim == 3 and frame.shape[2] == 1:61        frame = np.repeat(frame, 3, axis=2)62    elif frame.ndim == 3 and frame.shape[2] >= 4:63        frame = frame[..., :3]64    if frame.dtype != np.uint8:65        frame = np.clip(frame, 0, 255).astype(np.uint8)66    return frame67 68 69def _resize_control_frame(frame: Image.Image) -> Image.Image:70    return ImageOps.fit(71        frame.convert("RGB"),72        (TARGET_WIDTH, TARGET_HEIGHT),73        method=Image.Resampling.BICUBIC,74        centering=(0.5, 0.5),75    )76 77 78def load_control_video(path: str) -> tuple[list[Image.Image], float]:79    if not path:80        raise gr.Error("Upload a control video first.")81 82    source = Path(path)83    if not source.exists():84        raise gr.Error("The uploaded video is no longer available. Please upload it again.")85 86    reader = None87    try:88        reader = imageio.get_reader(str(source))89        metadata = reader.get_meta_data() or {}90        fps = float(metadata.get("fps") or OUTPUT_FPS)91        frames: list[Image.Image] = []92 93        for index, frame in enumerate(reader):94            if index >= MAX_FRAMES + 1:95                break96            frame = _to_rgb_uint8(np.asarray(frame))97            frames.append(_resize_control_frame(Image.fromarray(frame)))98    except Exception as exc:99        raise gr.Error(f"Could not read the control video: {exc}") from exc100    finally:101        if reader is not None:102            reader.close()103 104    if not frames:105        raise gr.Error("The control video contains no readable frames.")106    if len(frames) > MAX_FRAMES:107        raise gr.Error(108            f"This Space accepts at most {MAX_FRAMES} frames. "109            "Prepare a shorter input clip."110        )111    if len(frames) < 5:112        raise gr.Error("The control video is too short. Use at least 5 frames.")113    if (len(frames) - 1) % 4 != 0:114        raise gr.Error(115            f"The model requires 4n+1 frames; received {len(frames)}. "116            "Use 5, 9, 13, ..., 81 frames."117        )118 119    return frames, fps120 121 122def _gpu_duration(control_video, prompt, steps, guidance, control_scale, seed, negative_prompt) -> int:123    del control_video, prompt, guidance, control_scale, seed, negative_prompt124    steps = int(steps)125    return int(min(300, max(90, 30 + steps * 4)))126 127 128@spaces.GPU(duration=_gpu_duration)129def generate(130    control_video: str,131    prompt: str,132    steps: int,133    guidance: float,134    control_scale: float,135    seed: int,136    negative_prompt: str,137):138    if not prompt or not prompt.strip():139        raise gr.Error("Prompt cannot be empty.")140 141    frames, input_fps = load_control_video(control_video)142    num_frames = len(frames)143    seed = int(seed)144    if seed < 0:145        seed = random.randint(0, MAX_SEED)146 147    generator = torch.Generator(device="cuda").manual_seed(seed)148 149    try:150        result = pipe(151            video=frames,152            mask=None,153            prompt=prompt.strip(),154            negative_prompt=(negative_prompt or "").strip() or None,155            height=TARGET_HEIGHT,156            width=TARGET_WIDTH,157            num_frames=num_frames,158            num_inference_steps=int(steps),159            guidance_scale=float(guidance),160            conditioning_scale=float(control_scale),161            generator=generator,162        ).frames[0]163 164        with tempfile.NamedTemporaryFile(prefix="vc_experiment_", suffix=".mp4", delete=False) as tmp:165            output_path = tmp.name166        export_to_video(result, output_path, fps=OUTPUT_FPS)167 168        status = (169            f"Generated {num_frames} frames at {TARGET_WIDTH}×{TARGET_HEIGHT}, {OUTPUT_FPS} fps. "170            f"Input fps: {input_fps:.2f}. Seed: {seed}."171        )172        return output_path, status173    except gr.Error:174        raise175    except torch.cuda.OutOfMemoryError as exc:176        raise gr.Error(177            "GPU memory was exhausted. Try fewer frames or fewer inference steps."178        ) from exc179    except Exception as exc:180        raise gr.Error(f"Generation failed: {type(exc).__name__}: {exc}") from exc181    finally:182        gc.collect()183        if torch.cuda.is_available():184            torch.cuda.empty_cache()185 186 187with gr.Blocks(title="VC Experiment") as demo:188    gr.Markdown(189        "# VC Experiment\n"190        "Upload a control video, enter a text prompt, and generate a video."191    )192 193    with gr.Row():194        with gr.Column(scale=1):195            control_video = gr.Video(196                label="Control video",197                sources=["upload"],198                format="mp4",199            )200            prompt = gr.Textbox(201                label="Prompt",202                lines=4,203                placeholder="Describe the desired output...",204            )205 206            with gr.Accordion("Inference settings", open=False):207                steps = gr.Slider(10, 40, value=20, step=1, label="Inference steps")208                guidance = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Guidance scale")209                control_scale = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Control scale")210                seed = gr.Number(value=42, precision=0, label="Seed (-1 = random)")211                negative_prompt = gr.Textbox(212                    value=DEFAULT_NEGATIVE_PROMPT,213                    label="Negative prompt",214                    lines=4,215                )216 217            generate_button = gr.Button("Generate", variant="primary")218 219        with gr.Column(scale=1):220            result_video = gr.Video(label="Result", autoplay=True)221            status = gr.Markdown()222 223    generate_button.click(224        fn=generate,225        inputs=[control_video, prompt, steps, guidance, control_scale, seed, negative_prompt],226        outputs=[result_video, status],227    )228 229    gr.Markdown(230        "**Input:** 5–81 frames; frame count must be `4n+1`. "231        "Frames are normalized to 832×480."232    )233 234if __name__ == "__main__":235    demo.queue(default_concurrency_limit=1).launch()236