{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wan 2.1 (14B) Video Generation – ComfyUI Colab Pro (A100 / G4)\n",
    "\n",
    "**GPU:** Designed for A100 (40 GB) or L4/G4 (24 GB).  \n",
    "**Model:** Wan 2.1 T2V 14B (default) + optional I2V 14B.  \n",
    "**VRAM safeties:** FP8 / `--lowvram` toggles included.\n",
    "\n",
    "Run each cell top-to-bottom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 1. ENVIRONMENT CONFIGURATION\n",
    "# ============================================================\n",
    "\n",
    "#--- USER CONFIG ---\n",
    "MODEL_SIZE    = '14B'          # '1.3B' or '14B'\n",
    "MODEL_TYPE    = 'T2V'          # 'T2V' or 'I2V'\n",
    "USE_FP8       = False          # FP8 weights (saves ~50 % VRAM, slight quality loss)\n",
    "USE_GGUF      = False          # GGUF quant (for <=12 GB VRAM – not needed on A100)\n",
    "LOW_VRAM_MODE = False          # Enable ComfyUI --lowvram (slower, works on 16 GB)\n",
    "USE_NGROK     = True           # True = ngrok tunnel, False = localtunnel\n",
    "NGROK_TOKEN   = ''             # paste your ngrok authtoken here if USE_NGROK=True\n",
    "#-------------------\n",
    "\n",
    "import os, subprocess, sys, json, urllib.request, textwrap\n",
    "from pathlib import Path\n",
    "\n",
    "# Verify GPU\n",
    "!nvidia-smi\n",
    "\n",
    "# Paths\n",
    "BASE_DIR      = Path('/content/ComfyUI')\n",
    "MODELS_DIR    = BASE_DIR / 'models'\n",
    "DIFF_DIR      = MODELS_DIR / 'diffusion_models'\n",
    "TEXT_ENC_DIR  = MODELS_DIR / 'text_encoders'\n",
    "VAE_DIR       = MODELS_DIR / 'vae'\n",
    "CLIP_DIR      = MODELS_DIR / 'clip'\n",
    "WORKFLOW_DIR  = BASE_DIR / 'workflows'\n",
    "\n",
    "for p in [DIFF_DIR, TEXT_ENC_DIR, VAE_DIR, CLIP_DIR, WORKFLOW_DIR]:\n",
    "    p.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "print('Config:', MODEL_SIZE, MODEL_TYPE, 'FP8=', USE_FP8, 'GGUF=', USE_GGUF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 2. INSTALL COMFYUI & DEPENDENCIES\n",
    "# ============================================================\n",
    "\n",
    "%cd /content\n",
    "!git clone https://github.com/Comfy-Org/ComfyUI.git || true\n",
    "%cd /content/ComfyUI\n",
    "!git pull\n",
    "\n",
    "# Core requirements\n",
    "!pip install -q -r requirements.txt\n",
    "\n",
    "# Wan-specific extras\n",
    "!pip install -q transformers accelerate huggingface-hub safetensors\n",
    "\n",
    "# Tunnel tools\n",
    "if USE_NGROK:\n",
    "    !pip install -q pyngrok\n",
    "else:\n",
    "    !npm install -g localtunnel\n",
    "\n",
    "print('\\nComfyUI ready at', BASE_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 3. DOWNLOAD WAN 2.1 MODELS\n",
    "# ============================================================\n",
    "\n",
    "from huggingface_hub import snapshot_download, hf_hub_download\n",
    "\n",
    "# ---- 3A. TEXT ENCODER (shared across all Wan variants) ----\n",
    "print('Downloading UMT5 text encoder...')\n",
    "snapshot_download(\n",
    "    repo_id='Wan-AI/Wan2.1-T2V-14B',\n",
    "    local_dir=str(TEXT_ENC_DIR / 'umt5'),\n",
    "    allow_patterns=['*umt5*', '*tokenizer*'],\n",
    "    local_dir_use_symlinks=False,\n",
    "    resume_download=True\n",
    ")\n",
    "\n",
    "# ---- 3B. VAE (shared) ----\n",
    "print('Downloading Wan 2.1 VAE...')\n",
    "hf_hub_download(\n",
    "    repo_id='Wan-AI/Wan2.1-T2V-14B',\n",
    "    filename='wan_2.1_vae.safetensors',\n",
    "    local_dir=str(VAE_DIR),\n",
    "    local_dir_use_symlinks=False,\n",
    "    resume_download=True\n",
    ")\n",
    "\n",
    "# ---- 3C. DIFFUSION MODEL ----\n",
    "if USE_GGUF:\n",
    "    gguf_repo = 'city96/Wan2.1-T2V-14B-gguf'\n",
    "    gguf_file = 'Wan2.1-T2V-14B-Q4_K_M.gguf'\n",
    "    print(f'Downloading GGUF: {gguf_file}...')\n",
    "    hf_hub_download(repo_id=gguf_repo, filename=gguf_file,\n",
    "                    local_dir=str(DIFF_DIR), local_dir_use_symlinks=False, resume_download=True)\n",
    "elif USE_FP8:\n",
    "    fp8_repo  = 'Wan-AI/Wan2.1-T2V-14B'\n",
    "    fp8_file  = 'wan_2.1_fp8.safetensors'\n",
    "    try:\n",
    "        hf_hub_download(repo_id=fp8_repo, filename=fp8_file,\n",
    "                        local_dir=str(DIFF_DIR), local_dir_use_symlinks=False, resume_download=True)\n",
    "    except Exception as e:\n",
    "        print('FP8 file not found, falling back to FP16. Error:', e)\n",
    "        USE_FP8 = False\n",
    "\n",
    "if not USE_GGUF and not USE_FP8:\n",
    "    model_repo = f'Wan-AI/Wan2.1-{MODEL_TYPE}-{MODEL_SIZE}'\n",
    "    print(f'Downloading FP16 diffusion model from {model_repo}...')\n",
    "    snapshot_download(\n",
    "        repo_id=model_repo,\n",
    "        local_dir=str(DIFF_DIR / f'wan2.1_{MODEL_TYPE.lower()}_{MODEL_SIZE.lower()}'),\n",
    "        ignore_patterns=['*.md', '*.txt', 'umt5*'],\n",
    "        local_dir_use_symlinks=False,\n",
    "        resume_download=True\n",
    "    )\n",
    "\n",
    "print('\\nModel download complete.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 4. DOWNLOAD A WORKFLOW JSON (T2V + I2V templates)\n",
    "# ============================================================\n",
    "\n",
    "import urllib.request, json\n",
    "\n",
    "workflow_urls = {\n",
    "    'T2V': 'https://raw.githubusercontent.com/comfyanonymous/ComfyUI_examples/master/wan/wan_t2v.json',\n",
    "    'I2V': 'https://raw.githubusercontent.com/comfyanonymous/ComfyUI_examples/master/wan/wan_i2v.json'\n",
    "}\n",
    "\n",
    "wf_url = workflow_urls.get(MODEL_TYPE, workflow_urls['T2V'])\n",
    "wf_path = WORKFLOW_DIR / f'wan_{MODEL_TYPE.lower()}.json'\n",
    "\n",
    "try:\n",
    "    urllib.request.urlretrieve(wf_url, wf_path)\n",
    "    print(f'Workflow saved: {wf_path}')\n",
    "except Exception as e:\n",
    "    print('Could not auto-download workflow, you will load it manually in ComfyUI. Error:', e)\n",
    "\n",
    "# Show a minimal in-cell T2V workflow stub you can paste if needed\n",
    "minimal_wf = {\n",
    "    '1': {'class_type': 'CheckpointLoaderSimple', 'inputs': {'ckpt_name': ''}},\n",
    "    '2': {'class_type': 'CLIPTextEncode', 'inputs': {'text': 'a cat playing piano', 'clip': ['1', 0]}},\n",
    "    '3': {'class_type': 'EmptyLatentVideo', 'inputs': {'width': 832, 'height': 480, 'length': 81, 'batch_size': 1}},\n",
    "    '4': {'class_type': 'KSampler', 'inputs': {'model': ['1', 0], 'positive': ['2', 0], 'negative': ['2', 0],\n",
    "             'latent_image': ['3', 0], 'seed': 42, 'steps': 30, 'cfg': 6.0, 'sampler_name': 'uni_pc',\n",
    "             'scheduler': 'simple', 'denoise': 1.0}},\n",
    "    '5': {'class_type': 'VAEDecode', 'inputs': {'samples': ['4', 0], 'vae': ['1', 2]}}\n",
    "}\n",
    "print('\\nMinimal T2V node structure (for reference):')\n",
    "print(json.dumps(minimal_wf, indent=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 5. CREATE MODEL SYMLINKS / RENAME FOR COMFYUI DISCOVERY\n",
    "# ============================================================\n",
    "\n",
    "import glob\n",
    "\n",
    "# ComfyUI expects models in specific subfolders. Let's symlink the big diffusion file.\n",
    "diff_files = list(DIFF_DIR.rglob('*.safetensors')) + list(DIFF_DIR.rglob('*.gguf'))\n",
    "\n",
    "if diff_files:\n",
    "    for f in diff_files:\n",
    "        alias = DIFF_DIR / f.name\n",
    "        if not alias.exists():\n",
    "            alias.symlink_to(f.resolve())\n",
    "            print('Symlinked', f.name)\n",
    "else:\n",
    "    print('WARNING: No diffusion model found in', DIFF_DIR)\n",
    "\n",
    "# Show what's available\n",
    "print('\\n--- diffusion_models ---')\n",
    "for f in DIFF_DIR.iterdir(): print(f.name)\n",
    "print('--- vae ---')\n",
    "for f in VAE_DIR.iterdir(): print(f.name)\n",
    "print('--- text_encoders ---')\n",
    "for f in TEXT_ENC_DIR.iterdir(): print(f.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# 6. LAUNCH COMFYUI + TUNNEL\n",
    "# ============================================================\n",
    "\n",
    "import subprocess, threading, time, os, sys\n",
    "\n",
    "cmd = [sys.executable, 'main.py', '--listen', '0.0.0.0', '--port', '8188']\n",
    "\n",
    "if LOW_VRAM_MODE:\n",
    "    cmd.append('--lowvram')\n",
    "elif MODEL_SIZE == '14B' and not USE_FP8 and not USE_GGUF:\n",
    "    cmd.append('--highvram')\n",
    "\n",
    "cmd.append('--disable-auto-launch')\n",
    "\n",
    "print('Launching ComfyUI with:')\n",
    "print(' '.join(cmd))\n",
    "\n",
    "process = subprocess.Popen(cmd, cwd=str(BASE_DIR), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n",
    "\n",
    "def log_output(proc):\n",
    "    for line in proc.stdout:\n",
    "        print(line, end='')\n",
    "\n",
    "t = threading.Thread(target=log_output, args=(process,))\n",
    "t.daemon = True\n",
    "t.start()\n",
    "\n",
    "time.sleep(15)\n",
    "\n",
    "# ---- TUNNEL ----\n",
    "if USE_NGROK:\n",
    "    from pyngrok import ngrok\n",
    "    if NGROK_TOKEN:\n",
    "        ngrok.set_auth_token(NGROK_TOKEN)\n",
    "    public_url = ngrok.connect(8188, 'http')\n",
    "    print(f'\\n\\n🚀 NGROK URL: {public_url}\\n')\n",
    "else:\n",
    "    lt_proc = subprocess.Popen(['lt', '--port', '8188'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n",
    "    for line in lt_proc.stdout:\n",
    "        if 'your url is' in line.lower():\n",
    "            print('\\n\\n🚀 LOCALTUNNEL URL:', line.strip())\n",
    "            break\n",
    "\n",
    "print('\\nOpen the URL above in your browser. Then load the workflow from /workflows/.')\n",
    "print('If using T2V: use the wan_t2v.json workflow. If using I2V: use wan_i2v.json.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🎬 Using the Notebook\n",
    "\n",
    "1. **Run cells 1→6** in order.\n",
    "2. **Open the tunnel URL** printed in cell 6.\n",
    "3. **In ComfyUI:** click **Load** → choose `/content/ComfyUI/workflows/wan_t2v.json` (or I2V).\n",
    "4. **Select model** in the `Load Diffusion Model` node: pick the symlinked `.safetensors` or `.gguf`.\n",
    "5. **Queue Prompt** to generate video.\n",
    "\n",
    "## 🛠 Troubleshooting\n",
    "\n",
    "| Symptom | Fix |\n",
    "|---|---|\n",
    "| `CUDA OOM` | Enable `LOW_VRAM_MODE = True` or switch to `MODEL_SIZE = '1.3B'` |\n",
    "| `Model not found` | Check cell 5 symlinks; manually browse to `/content/ComfyUI/models/diffusion_models/` |\n",
    "| `ImportError` | Re-run cell 2 (requirements) then **Runtime → Restart session** |\n",
    "| Slow first run | Normal – model is being cached to GPU. Subsequent runs are faster. |\n",
    "| Black output | Increase `steps` to 30–50 or check VAE node connection. |\n",
    "\n",
    "## 💾 Persisting Models on Google Drive\n",
    "\n",
    "To avoid re-downloading every session, mount Drive and change paths:\n",
    "```python\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "BASE_DIR = Path('/content/drive/MyDrive/ComfyUI')  # change all paths above\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}