WAMU V2 - Wan 2.2 I2V (14B) on google colab

wan2_2_i2v_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ccf54947",
   "metadata": {},
   "source": [
    "# Wan 2.2 I2V Lightning sample for Google Colab\n",
    "\n",
    "This notebook runs the image-to-video sample from the Hugging Face Space\n",
    "`r3gm/wan2-2-fp8da-aoti-preview` on a regular Colab GPU.\n",
    "\n",
    "Runtime notes:\n",
    "\n",
    "- Use a GPU runtime. A high-VRAM GPU such as A100 is recommended.\n",
    "- The Space uses ZeroGPU AOTI kernels. Colab usually cannot use those kernels,\n",
    "  so this notebook keeps AOTI disabled by default and uses normal Diffusers\n",
    "  execution with optional TorchAO quantization.\n",
    "- If you hit CUDA OOM, reduce `DURATION_SECONDS`, use fewer frames, or disable\n",
    "  optional quantization only if it is failing on your GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fb09153",
   "metadata": {},
   "source": [
    "## 1. Check GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75055808",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError(\"Please switch Colab runtime to GPU: Runtime > Change runtime type > GPU\")\n",
    "\n",
    "print(torch.cuda.get_device_name(0))\n",
    "print(f\"CUDA capability: {torch.cuda.get_device_capability(0)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d211bbbc",
   "metadata": {},
   "source": [
    "## 2. Install dependencies\n",
    "\n",
    "Run this cell once. Restarting the runtime after installation is usually not required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c11233",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "!apt-get -qq update\n",
    "!apt-get -qq install -y ffmpeg git git-lfs\n",
    "!git lfs install\n",
    "!pip -q install --upgrade pip\n",
    "!pip -q install \\\n",
    "  \"git+https://github.com/linoytsaban/diffusers.git@wan22-loras\" \\\n",
    "  \"transformers<5\" \\\n",
    "  accelerate \\\n",
    "  safetensors \\\n",
    "  sentencepiece \\\n",
    "  peft \\\n",
    "  ftfy \\\n",
    "  imageio \\\n",
    "  imageio-ffmpeg \\\n",
    "  opencv-python \\\n",
    "  \"torchao==0.11.0\" \\\n",
    "  huggingface_hub"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eebdf296",
   "metadata": {},
   "source": [
    "## 3. Clone the sample Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc771c85",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pathlib\n",
    "import subprocess\n",
    "\n",
    "REPO_URL = \"https://huggingface.co/spaces/r3gm/wan2-2-fp8da-aoti-preview\"\n",
    "WORKDIR = pathlib.Path(\"/content/wan2-2-fp8da-aoti-preview\")\n",
    "\n",
    "if not WORKDIR.exists():\n",
    "    subprocess.run([\"git\", \"clone\", REPO_URL, str(WORKDIR)], check=True)\n",
    "\n",
    "os.chdir(WORKDIR)\n",
    "subprocess.run([\"git\", \"lfs\", \"pull\"], check=False)\n",
    "print(f\"Working directory: {pathlib.Path.cwd()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b875c3",
   "metadata": {},
   "source": [
    "## 4. Optional Hugging Face login\n",
    "\n",
    "If model download fails with an authorization or rate-limit error, create a\n",
    "Hugging Face token and run `login()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20dada11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "\n",
    "# login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f77963",
   "metadata": {},
   "source": [
    "## 5. Load the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1302bdda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import gc\n",
    "import random\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from IPython.display import Video, display\n",
    "\n",
    "from diffusers import (\n",
    "    DEISMultistepScheduler,\n",
    "    DPMSolverMultistepInverseScheduler,\n",
    "    DPMSolverMultistepScheduler,\n",
    "    DPMSolverSinglestepScheduler,\n",
    "    FlowMatchEulerDiscreteScheduler,\n",
    "    SASolverScheduler,\n",
    "    UniPCMultistepScheduler,\n",
    ")\n",
    "from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline\n",
    "from diffusers.utils.export_utils import export_to_video\n",
    "\n",
    "try:\n",
    "    from torchao.quantization import (\n",
    "        Float8DynamicActivationFloat8WeightConfig,\n",
    "        Int8WeightOnlyConfig,\n",
    "        quantize_,\n",
    "    )\n",
    "except Exception as exc:\n",
    "    quantize_ = None\n",
    "    print(f\"TorchAO quantization is unavailable: {exc}\")\n",
    "\n",
    "\n",
    "MODEL_ID = \"TestOrganizationPleaseIgnore/WAMU_v2_WAN2.2_I2V_LIGHTNING\"\n",
    "USE_TORCHAO_QUANTIZATION = True\n",
    "LOW_VRAM_MODE = True\n",
    "\n",
    "pipe = WanImageToVideoPipeline.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.bfloat16,\n",
    ")\n",
    "\n",
    "original_scheduler = copy.deepcopy(pipe.scheduler)\n",
    "\n",
    "if USE_TORCHAO_QUANTIZATION and quantize_ is not None:\n",
    "    try:\n",
    "        quantize_(pipe.text_encoder, Int8WeightOnlyConfig())\n",
    "        quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())\n",
    "        quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())\n",
    "        print(\"TorchAO quantization enabled.\")\n",
    "    except Exception as exc:\n",
    "        print(f\"Continuing without TorchAO quantization: {exc}\")\n",
    "\n",
    "if LOW_VRAM_MODE:\n",
    "    pipe.enable_model_cpu_offload()\n",
    "    pipe.vae.enable_slicing()\n",
    "    pipe.vae.enable_tiling()\n",
    "    print(\"Low VRAM mode enabled: model CPU offload + VAE slicing/tiling.\")\n",
    "else:\n",
    "    pipe.to(\"cuda\")\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e3abfe",
   "metadata": {},
   "source": [
    "## 6. Define generation helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1625722",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "LOW_VRAM_RESOLUTION = True\n",
    "\n",
    "if LOW_VRAM_RESOLUTION:\n",
    "    MAX_DIM = 640\n",
    "    MIN_DIM = 384\n",
    "    SQUARE_DIM = 512\n",
    "else:\n",
    "    MAX_DIM = 832\n",
    "    MIN_DIM = 480\n",
    "    SQUARE_DIM = 640\n",
    "MULTIPLE_OF = 16\n",
    "FIXED_FPS = 16\n",
    "MIN_FRAMES_MODEL = 8\n",
    "MAX_FRAMES_MODEL = 160\n",
    "MAX_SEED = np.iinfo(np.int32).max\n",
    "\n",
    "SCHEDULER_MAP = {\n",
    "    \"FlowMatchEulerDiscrete\": FlowMatchEulerDiscreteScheduler,\n",
    "    \"SASolver\": SASolverScheduler,\n",
    "    \"DEISMultistep\": DEISMultistepScheduler,\n",
    "    \"DPMSolverMultistepInverse\": DPMSolverMultistepInverseScheduler,\n",
    "    \"UniPCMultistep\": UniPCMultistepScheduler,\n",
    "    \"DPMSolverMultistep\": DPMSolverMultistepScheduler,\n",
    "    \"DPMSolverSinglestep\": DPMSolverSinglestepScheduler,\n",
    "}\n",
    "\n",
    "DEFAULT_PROMPT = \"make this image come alive, cinematic motion, smooth animation\"\n",
    "DEFAULT_NEGATIVE_PROMPT = (\n",
    "    \"色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, \"\n",
    "    \"整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, \"\n",
    "    \"画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, \"\n",
    "    \"手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走\"\n",
    ")\n",
    "\n",
    "\n",
    "def resize_image(image: Image.Image) -> Image.Image:\n",
    "    width, height = image.size\n",
    "    if width == height:\n",
    "        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)\n",
    "\n",
    "    aspect_ratio = width / height\n",
    "    max_aspect_ratio = MAX_DIM / MIN_DIM\n",
    "    min_aspect_ratio = MIN_DIM / MAX_DIM\n",
    "    image_to_resize = image\n",
    "\n",
    "    if aspect_ratio > max_aspect_ratio:\n",
    "        target_w, target_h = MAX_DIM, MIN_DIM\n",
    "        crop_width = int(round(height * max_aspect_ratio))\n",
    "        left = (width - crop_width) // 2\n",
    "        image_to_resize = image.crop((left, 0, left + crop_width, height))\n",
    "    elif aspect_ratio < min_aspect_ratio:\n",
    "        target_w, target_h = MIN_DIM, MAX_DIM\n",
    "        crop_height = int(round(width / min_aspect_ratio))\n",
    "        top = (height - crop_height) // 2\n",
    "        image_to_resize = image.crop((0, top, width, top + crop_height))\n",
    "    elif width > height:\n",
    "        target_w = MAX_DIM\n",
    "        target_h = int(round(target_w / aspect_ratio))\n",
    "    else:\n",
    "        target_h = MAX_DIM\n",
    "        target_w = int(round(target_h * aspect_ratio))\n",
    "\n",
    "    final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF\n",
    "    final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF\n",
    "    final_w = max(MIN_DIM, min(MAX_DIM, final_w))\n",
    "    final_h = max(MIN_DIM, min(MAX_DIM, final_h))\n",
    "    return image_to_resize.resize((final_w, final_h), Image.LANCZOS)\n",
    "\n",
    "\n",
    "def resize_and_crop_to_match(target_image: Image.Image, reference_image: Image.Image) -> Image.Image:\n",
    "    ref_width, ref_height = reference_image.size\n",
    "    target_width, target_height = target_image.size\n",
    "    scale = max(ref_width / target_width, ref_height / target_height)\n",
    "    new_width, new_height = int(target_width * scale), int(target_height * scale)\n",
    "    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)\n",
    "    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2\n",
    "    return resized.crop((left, top, left + ref_width, top + ref_height))\n",
    "\n",
    "\n",
    "def get_num_frames(duration_seconds: float) -> int:\n",
    "    return 1 + int(\n",
    "        np.clip(\n",
    "            int(round(duration_seconds * FIXED_FPS)),\n",
    "            MIN_FRAMES_MODEL,\n",
    "            MAX_FRAMES_MODEL,\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "def set_scheduler(scheduler_name: str, flow_shift: float) -> None:\n",
    "    scheduler_class = SCHEDULER_MAP[scheduler_name]\n",
    "    current_class_name = pipe.scheduler.config._class_name\n",
    "    current_flow_shift = pipe.scheduler.config.get(\"flow_shift\", pipe.scheduler.config.get(\"shift\"))\n",
    "\n",
    "    if scheduler_class.__name__ == current_class_name and flow_shift == current_flow_shift:\n",
    "        return\n",
    "\n",
    "    config = copy.deepcopy(original_scheduler.config)\n",
    "    if scheduler_class == FlowMatchEulerDiscreteScheduler:\n",
    "        config[\"shift\"] = flow_shift\n",
    "    else:\n",
    "        config[\"flow_shift\"] = flow_shift\n",
    "    pipe.scheduler = scheduler_class.from_config(config)\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def generate_video_colab(\n",
    "    input_image: Image.Image,\n",
    "    output_path: str = \"/content/wan22_i2v_output.mp4\",\n",
    "    prompt: str = DEFAULT_PROMPT,\n",
    "    last_image: Image.Image | None = None,\n",
    "    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,\n",
    "    duration_seconds: float = 0.6,\n",
    "    steps: int = 4,\n",
    "    guidance_scale: float = 1.0,\n",
    "    guidance_scale_2: float = 1.0,\n",
    "    seed: int = 42,\n",
    "    scheduler: str = \"UniPCMultistep\",\n",
    "    flow_shift: float = 6.0,\n",
    "    quality: int = 5,\n",
    ") -> str:\n",
    "    resized_image = resize_image(input_image.convert(\"RGB\"))\n",
    "    processed_last_image = None\n",
    "    if last_image is not None:\n",
    "        processed_last_image = resize_and_crop_to_match(last_image.convert(\"RGB\"), resized_image)\n",
    "\n",
    "    num_frames = get_num_frames(duration_seconds)\n",
    "    seed = int(seed if seed >= 0 else random.randint(0, MAX_SEED))\n",
    "    set_scheduler(scheduler, flow_shift)\n",
    "\n",
    "    print(\n",
    "        f\"Generating {num_frames} frames at {resized_image.size[0]}x{resized_image.size[1]} \"\n",
    "        f\"with {steps} steps, seed={seed}\"\n",
    "    )\n",
    "    start = time.time()\n",
    "    result = pipe(\n",
    "        image=resized_image,\n",
    "        last_image=processed_last_image,\n",
    "        prompt=prompt,\n",
    "        negative_prompt=negative_prompt,\n",
    "        height=resized_image.height,\n",
    "        width=resized_image.width,\n",
    "        num_frames=num_frames,\n",
    "        guidance_scale=float(guidance_scale),\n",
    "        guidance_scale_2=float(guidance_scale_2),\n",
    "        num_inference_steps=int(steps),\n",
    "        generator=torch.Generator(device=\"cuda\").manual_seed(seed),\n",
    "        output_type=\"np\",\n",
    "    )\n",
    "    pipe.scheduler = original_scheduler\n",
    "\n",
    "    frames = result.frames[0]\n",
    "    export_to_video(frames, output_path, fps=FIXED_FPS, quality=quality)\n",
    "    torch.cuda.empty_cache()\n",
    "    print(f\"Saved to {output_path} in {time.time() - start:.1f}s\")\n",
    "    return output_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e05098a",
   "metadata": {},
   "source": [
    "## 7. Choose an input image\n",
    "\n",
    "Run this cell to upload your own image. If upload is skipped or unavailable,\n",
    "the notebook falls back to a sample image from the cloned Space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e7e42c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_IMAGE_PATH = None\n",
    "\n",
    "try:\n",
    "    from google.colab import files\n",
    "\n",
    "    uploaded = files.upload()\n",
    "    if uploaded:\n",
    "        INPUT_IMAGE_PATH = next(iter(uploaded))\n",
    "except Exception as exc:\n",
    "    print(f\"Upload is unavailable in this environment: {exc}\")\n",
    "\n",
    "if INPUT_IMAGE_PATH is None:\n",
    "    INPUT_IMAGE_PATH = str(WORKDIR / \"wan_i2v_input.JPG\")\n",
    "    print(f\"Using sample image: {INPUT_IMAGE_PATH}\")\n",
    "else:\n",
    "    print(f\"Using uploaded image: {INPUT_IMAGE_PATH}\")\n",
    "\n",
    "input_image = Image.open(INPUT_IMAGE_PATH).convert(\"RGB\")\n",
    "display(input_image.resize((384, int(384 * input_image.height / input_image.width))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7437ec0f",
   "metadata": {},
   "source": [
    "## 8. Generate video\n",
    "\n",
    "For a first test, keep the duration short. Increase `DURATION_SECONDS` after the\n",
    "basic run succeeds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92bfcd7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROMPT = \"make this image come alive, cinematic camera motion, smooth animation\"\n",
    "DURATION_SECONDS = 0.6\n",
    "STEPS = 4\n",
    "SEED = 42\n",
    "\n",
    "output_path = generate_video_colab(\n",
    "    input_image=input_image,\n",
    "    prompt=PROMPT,\n",
    "    duration_seconds=DURATION_SECONDS,\n",
    "    steps=STEPS,\n",
    "    seed=SEED,\n",
    "    output_path=\"/content/wan22_i2v_output.mp4\",\n",
    ")\n",
    "\n",
    "display(Video(output_path, embed=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7921ce5a",
   "metadata": {},
   "source": [
    "## 9. Download the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f11c2d13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "\n",
    "files.download(output_path)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}