{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro-md",
   "metadata": {},
   "source": [
    "# 🎬 OmniFlows — Wan 2.1 T2V + MMAudio Server\n",
    "\n",
    "**No account or token needed.** Just:\n",
    "1. Set runtime to **T4 GPU** (Runtime → Change runtime type)\n",
    "2. Click **Runtime → Run all**\n",
    "3. Wait ~8 minutes — a big URL box will appear at the bottom\n",
    "4. Copy the URL → paste into **OmniFlows Admin → Colab tab → Save**\n",
    "\n",
    "**What this server can do:**\n",
    "- 🎥 Generate videos from text prompts (Wan 2.1 T2V 1.3B)\n",
    "- 🔊 Add AI-generated audio to any video URL (MMAudio small_16k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-install",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 1: Install dependencies ──────────────────────────────────────────────\n",
    "import subprocess, sys\n",
    "\n",
    "pkgs = [\n",
    "    'diffusers>=0.32.0',\n",
    "    'transformers>=4.46.0',\n",
    "    'accelerate>=1.0.0',\n",
    "    'bitsandbytes>=0.43.0',\n",
    "    'imageio[ffmpeg]>=2.34.0',\n",
    "    'fastapi>=0.111.0',\n",
    "    'uvicorn[standard]>=0.30.0',\n",
    "    'httpx>=0.27.0',\n",
    "    'sentencepiece',\n",
    "    'ftfy',\n",
    "    'torchaudio',\n",
    "    'av',\n",
    "]\n",
    "subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '--upgrade'] + pkgs)\n",
    "\n",
    "# Install MMAudio for video-to-audio generation\n",
    "subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q',\n",
    "    'git+https://github.com/hkchengrex/MMAudio.git'])\n",
    "\n",
    "# Install cloudflared (no account/token needed)\n",
    "subprocess.check_call(['wget', '-q', '-O', '/usr/local/bin/cloudflared',\n",
    "    'https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64'])\n",
    "subprocess.check_call(['chmod', '+x', '/usr/local/bin/cloudflared'])\n",
    "\n",
    "print('✅ All dependencies installed.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-model",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 2: Load models ────────────────────────────────────────────────────────\n",
    "#\n",
    "# VRAM budget on free T4 (15 GB):\n",
    "#   UMT5-XXL text encoder  4-bit NF4  ≈ 2.5 GB\n",
    "#   Wan 2.1 DiT 1.3B       float16    ≈ 2.6 GB\n",
    "#   VAE                    float16    ≈ 0.5 GB\n",
    "#   MMAudio small_16k      float16    ≈ 1.5 GB\n",
    "#   ─────────────────────────────────────────────\n",
    "#   Total loaded:                     ≈ 7.1 GB → 7.9 GB free for inference\n",
    "#\n",
    "import gc, os\n",
    "import torch\n",
    "from diffusers import WanPipeline, UniPCMultistepScheduler\n",
    "from transformers import UMT5EncoderModel, BitsAndBytesConfig\n",
    "\n",
    "os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "MODEL_ID = 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers'\n",
    "print('Step 1/3 -- Loading text encoder in 4-bit NF4 (~2.5 GB VRAM)...')\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_compute_dtype=torch.float16,\n",
    "    bnb_4bit_quant_type='nf4',\n",
    ")\n",
    "text_encoder = UMT5EncoderModel.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    subfolder='text_encoder',\n",
    "    quantization_config=bnb_config,\n",
    "    device_map='cuda:0',\n",
    ")\n",
    "used = torch.cuda.memory_allocated() / 1e9\n",
    "print(f'  Text encoder: {used:.1f} GB used on GPU')\n",
    "\n",
    "print('Step 2/3 -- Loading Wan 2.1 pipeline (transformer + VAE)...')\n",
    "pipe = WanPipeline.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    text_encoder=text_encoder,\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ")\n",
    "pipe.transformer.to('cuda')\n",
    "pipe.vae.to('cuda')\n",
    "pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "\n",
    "for _method in ['enable_vae_slicing', 'enable_vae_tiling']:\n",
    "    try:\n",
    "        getattr(pipe, _method)()\n",
    "        print(f'  {_method}: OK')\n",
    "    except AttributeError:\n",
    "        try:\n",
    "            getattr(pipe.vae, _method.replace('enable_vae_', 'enable_'))()\n",
    "            print(f'  vae.{_method.replace(\"enable_vae_\",\"enable_\")}: OK')\n",
    "        except AttributeError:\n",
    "            pass\n",
    "try:\n",
    "    pipe.enable_attention_slicing(1)\n",
    "    print('  enable_attention_slicing: OK')\n",
    "except AttributeError:\n",
    "    pass\n",
    "\n",
    "total = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "used  = torch.cuda.memory_allocated() / 1e9\n",
    "print(f'\\nWan 2.1 ready -- GPU: {used:.1f} GB used / {total - used:.1f} GB free / {total:.1f} GB total')\n",
    "\n",
    "# ── MMAudio (video-to-audio) ───────────────────────────────────────────────────\n",
    "print('\\nStep 3/3 -- Loading MMAudio small_16k for video-to-audio...')\n",
    "mmaudio_net = None\n",
    "mmaudio_features = None\n",
    "try:\n",
    "    from mmaudio.eval_utils import setup_eval_utils\n",
    "    mmaudio_net, mmaudio_features = setup_eval_utils(\n",
    "        'small_16k', device='cuda', dtype=torch.float16\n",
    "    )\n",
    "    used = torch.cuda.memory_allocated() / 1e9\n",
    "    print(f'  MMAudio ready -- GPU: {used:.1f} GB used / {total - used:.1f} GB free / {total:.1f} GB total')\n",
    "except Exception as _e:\n",
    "    print(f'  ⚠️  MMAudio failed to load (non-fatal): {_e}')\n",
    "    print('  Video generation still works. /add-audio endpoint will return 503.')\n",
    "\n",
    "print('\\n✅ All models ready!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-server",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Cell 3: Start server + cloudflare tunnel ───────────────────────────────────\n",
    "import asyncio, base64, gc, io, re, subprocess, tempfile, threading, time, uuid\n",
    "import urllib.request\n",
    "import imageio\n",
    "import numpy as np\n",
    "import torch\n",
    "from pathlib import Path\n",
    "from typing import List, Optional\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "import uvicorn\n",
    "from fastapi import FastAPI, HTTPException, BackgroundTasks\n",
    "from fastapi.middleware.cors import CORSMiddleware\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "# ── Pydantic models ─────────────────────────────────────────────────────────────\n",
    "class GenerateRequest(BaseModel):\n",
    "    prompts: List[str] = Field(..., min_length=1, max_length=5)\n",
    "    negative_prompt: str = 'ugly, blurry, watermark, low quality, distorted'\n",
    "    width: int = 512\n",
    "    height: int = 288\n",
    "    num_frames: int = 25\n",
    "    num_inference_steps: int = 20\n",
    "    guidance_scale: float = 5.0\n",
    "    fps: int = 8\n",
    "\n",
    "class VideoResult(BaseModel):\n",
    "    index: int\n",
    "    prompt: str\n",
    "    video_b64: str\n",
    "    duration_s: float\n",
    "\n",
    "class GenerateResponse(BaseModel):\n",
    "    job_id: str\n",
    "    videos: List[VideoResult]\n",
    "    total_time_s: float\n",
    "\n",
    "class JobStatus(BaseModel):\n",
    "    job_id: str\n",
    "    status: str\n",
    "    progress: int\n",
    "    total: int\n",
    "    videos: Optional[List[VideoResult]] = None\n",
    "    error: Optional[str] = None\n",
    "\n",
    "class AudioRequest(BaseModel):\n",
    "    video_url: str\n",
    "    prompt: str = ''\n",
    "    negative_prompt: str = 'music, speech, voice, talking'\n",
    "    duration: Optional[float] = None\n",
    "    cfg_strength: float = 4.5\n",
    "    num_steps: int = 25\n",
    "\n",
    "class AudioResponse(BaseModel):\n",
    "    video_b64: str\n",
    "    audio_duration_s: float\n",
    "    duration_s: float\n",
    "\n",
    "# ── Frames → base64 MP4 ─────────────────────────────────────────────────────────\n",
    "def frames_to_b64_mp4(frames, fps: int = 8) -> str:\n",
    "    buf = io.BytesIO()\n",
    "    writer = imageio.get_writer(buf, format='mp4', fps=fps, codec='libx264',\n",
    "                                 quality=7, pixelformat='yuv420p',\n",
    "                                 output_params=['-movflags', 'faststart'])\n",
    "    for frame in frames:\n",
    "        if not isinstance(frame, np.ndarray):\n",
    "            frame = np.array(frame)\n",
    "        writer.append_data(frame)\n",
    "    writer.close()\n",
    "    return base64.b64encode(buf.getvalue()).decode()\n",
    "\n",
    "# ── Job store ───────────────────────────────────────────────────────────────────\n",
    "jobs: dict = {}\n",
    "\n",
    "# ── Video generation worker ──────────────────────────────────────────────────────\n",
    "def run_generation(job_id: str, req: GenerateRequest):\n",
    "    jobs[job_id] = {'status': 'running', 'progress': 0, 'total': len(req.prompts), 'videos': [], 'error': None}\n",
    "    t0 = time.time()\n",
    "    results = []\n",
    "    try:\n",
    "        for i, prompt in enumerate(req.prompts):\n",
    "            jobs[job_id]['progress'] = i\n",
    "            ts = time.time()\n",
    "\n",
    "            gc.collect()\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            output = pipe(\n",
    "                prompt=prompt,\n",
    "                negative_prompt=req.negative_prompt,\n",
    "                width=req.width,\n",
    "                height=req.height,\n",
    "                num_frames=req.num_frames,\n",
    "                num_inference_steps=req.num_inference_steps,\n",
    "                guidance_scale=req.guidance_scale,\n",
    "            )\n",
    "            frames = output.frames[0]\n",
    "            video_b64 = frames_to_b64_mp4(frames, fps=req.fps)\n",
    "\n",
    "            del output, frames\n",
    "            gc.collect()\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            results.append({'index': i, 'prompt': prompt, 'video_b64': video_b64, 'duration_s': round(time.time() - ts, 1)})\n",
    "            jobs[job_id]['videos'] = results[:]\n",
    "            jobs[job_id]['progress'] = i + 1\n",
    "            print(f'[{job_id[:8]}] Video {i+1}/{len(req.prompts)} done in {results[-1][\"duration_s\"]}s')\n",
    "\n",
    "        jobs[job_id]['status'] = 'done'\n",
    "        jobs[job_id]['total_time_s'] = round(time.time() - t0, 1)\n",
    "        print(f'[{job_id[:8]}] All done in {jobs[job_id][\"total_time_s\"]}s')\n",
    "    except Exception as ex:\n",
    "        jobs[job_id]['status'] = 'error'\n",
    "        jobs[job_id]['error'] = str(ex)\n",
    "        print(f'[ERROR] job {job_id}: {ex}')\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "# ── FastAPI app ─────────────────────────────────────────────────────────────────\n",
    "app = FastAPI(title='OmniFlows Colab Video Server', version='3.0.0')\n",
    "app.add_middleware(CORSMiddleware, allow_origins=['*'], allow_methods=['*'], allow_headers=['*'])\n",
    "\n",
    "@app.get('/health')\n",
    "async def health():\n",
    "    return {\n",
    "        'status': 'ok',\n",
    "        'model': 'Wan2.1-T2V-1.3B',\n",
    "        'mmaudio': mmaudio_net is not None,\n",
    "        'device': torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu',\n",
    "    }\n",
    "\n",
    "@app.post('/generate', response_model=GenerateResponse)\n",
    "async def generate(req: GenerateRequest):\n",
    "    if len(req.prompts) > 5:\n",
    "        raise HTTPException(status_code=400, detail='Maximum 5 prompts per batch')\n",
    "    job_id = str(uuid.uuid4())\n",
    "    loop = asyncio.get_event_loop()\n",
    "    await loop.run_in_executor(None, run_generation, job_id, req)\n",
    "    job = jobs.get(job_id, {})\n",
    "    if job.get('status') == 'error':\n",
    "        raise HTTPException(status_code=500, detail=job.get('error', 'Generation failed'))\n",
    "    return GenerateResponse(\n",
    "        job_id=job_id,\n",
    "        videos=[VideoResult(**v) for v in job.get('videos', [])],\n",
    "        total_time_s=job.get('total_time_s', 0),\n",
    "    )\n",
    "\n",
    "@app.post('/generate/async')\n",
    "async def generate_async(req: GenerateRequest, background_tasks: BackgroundTasks):\n",
    "    if len(req.prompts) > 5:\n",
    "        raise HTTPException(status_code=400, detail='Maximum 5 prompts per batch')\n",
    "    job_id = str(uuid.uuid4())\n",
    "    background_tasks.add_task(run_generation, job_id, req)\n",
    "    return {'job_id': job_id, 'status': 'queued', 'total': len(req.prompts)}\n",
    "\n",
    "@app.get('/progress/{job_id}', response_model=JobStatus)\n",
    "async def progress(job_id: str):\n",
    "    job = jobs.get(job_id)\n",
    "    if job is None:\n",
    "        raise HTTPException(status_code=404, detail='Job not found')\n",
    "    videos = [VideoResult(**v) for v in job.get('videos', [])] if job['status'] == 'done' else None\n",
    "    return JobStatus(job_id=job_id, status=job['status'], progress=job['progress'],\n",
    "                     total=job['total'], videos=videos, error=job.get('error'))\n",
    "\n",
    "# ── Add Audio endpoint ─────────────────────────────────────────────────────────\n",
    "@app.post('/add-audio', response_model=AudioResponse)\n",
    "async def add_audio(req: AudioRequest):\n",
    "    \"\"\"Download a video from URL, generate matching audio with MMAudio, merge and return.\"\"\"\n",
    "    if mmaudio_net is None or mmaudio_features is None:\n",
    "        raise HTTPException(status_code=503, detail='MMAudio not loaded. Restart the notebook.')\n",
    "\n",
    "    loop = asyncio.get_event_loop()\n",
    "    result = await loop.run_in_executor(None, _run_add_audio, req)\n",
    "    if 'error' in result:\n",
    "        raise HTTPException(status_code=500, detail=result['error'])\n",
    "    return AudioResponse(**result)\n",
    "\n",
    "def _run_add_audio(req: AudioRequest) -> dict:\n",
    "    import torchaudio\n",
    "    from mmaudio.eval_utils import generate as mmaudio_generate, load_video\n",
    "    from mmaudio.model.flow_matching import FlowMatching\n",
    "\n",
    "    t0 = time.time()\n",
    "    tmpdir = tempfile.mkdtemp()\n",
    "    video_path = Path(tmpdir) / 'input.mp4'\n",
    "    audio_path = Path(tmpdir) / 'audio.wav'\n",
    "    output_path = Path(tmpdir) / 'output.mp4'\n",
    "\n",
    "    try:\n",
    "        # Download the video\n",
    "        print(f'[add-audio] Downloading video from {req.video_url[:80]}...')\n",
    "        headers = {'User-Agent': 'Mozilla/5.0'}\n",
    "        r = urllib.request.Request(req.video_url, headers=headers)\n",
    "        with urllib.request.urlopen(r, timeout=60) as resp:\n",
    "            video_path.write_bytes(resp.read())\n",
    "        print(f'[add-audio] Downloaded {video_path.stat().st_size // 1024} KB')\n",
    "\n",
    "        # Load video frames for MMAudio\n",
    "        clip_frames, sync_frames, duration_sec, fps = load_video(str(video_path), num_frames=8)\n",
    "        if req.duration is not None:\n",
    "            duration_sec = req.duration\n",
    "        print(f'[add-audio] Video duration: {duration_sec:.1f}s, fps: {fps}')\n",
    "\n",
    "        # Generate audio\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "        fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=req.num_steps)\n",
    "        rng = torch.Generator(device='cuda').manual_seed(42)\n",
    "\n",
    "        with torch.inference_mode():\n",
    "            audio_waveform = mmaudio_generate(\n",
    "                clip_frames, sync_frames, duration_sec,\n",
    "                text=[req.prompt],\n",
    "                neg_text=[req.negative_prompt],\n",
    "                feature_utils=mmaudio_features,\n",
    "                net=mmaudio_net,\n",
    "                fm=fm,\n",
    "                rng=rng,\n",
    "                cfg_strength=req.cfg_strength,\n",
    "            )\n",
    "\n",
    "        # audio_waveform shape: [T] or [1, T] at 16000 Hz\n",
    "        if audio_waveform.dim() == 1:\n",
    "            audio_waveform = audio_waveform.unsqueeze(0)\n",
    "        torchaudio.save(str(audio_path), audio_waveform.float().cpu(), 16000)\n",
    "        print(f'[add-audio] Audio generated: {audio_waveform.shape[-1] / 16000:.1f}s at 16kHz')\n",
    "\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        # Merge audio into video with ffmpeg\n",
    "        result = subprocess.run([\n",
    "            'ffmpeg', '-y',\n",
    "            '-i', str(video_path),\n",
    "            '-i', str(audio_path),\n",
    "            '-c:v', 'copy',\n",
    "            '-c:a', 'aac',\n",
    "            '-ar', '44100',\n",
    "            '-b:a', '128k',\n",
    "            '-shortest',\n",
    "            '-movflags', '+faststart',\n",
    "            str(output_path)\n",
    "        ], capture_output=True, text=True)\n",
    "\n",
    "        if result.returncode != 0:\n",
    "            return {'error': f'ffmpeg merge failed: {result.stderr[-500:]}'}\n",
    "\n",
    "        video_b64 = base64.b64encode(output_path.read_bytes()).decode()\n",
    "        elapsed = round(time.time() - t0, 1)\n",
    "        print(f'[add-audio] Done in {elapsed}s, output {output_path.stat().st_size // 1024} KB')\n",
    "\n",
    "        return {\n",
    "            'video_b64': video_b64,\n",
    "            'audio_duration_s': round(audio_waveform.shape[-1] / 16000, 2),\n",
    "            'duration_s': elapsed,\n",
    "        }\n",
    "    except Exception as ex:\n",
    "        import traceback\n",
    "        traceback.print_exc()\n",
    "        return {'error': str(ex)}\n",
    "    finally:\n",
    "        # Cleanup temp files\n",
    "        for p in [video_path, audio_path, output_path]:\n",
    "            try: p.unlink()\n",
    "            except: pass\n",
    "        try: Path(tmpdir).rmdir()\n",
    "        except: pass\n",
    "\n",
    "# ── Start uvicorn in background ─────────────────────────────────────────────────\n",
    "PORT = 8765\n",
    "\n",
    "def run_uvicorn():\n",
    "    uvicorn.run(app, host='0.0.0.0', port=PORT, log_level='warning')\n",
    "\n",
    "server_thread = threading.Thread(target=run_uvicorn, daemon=True)\n",
    "server_thread.start()\n",
    "time.sleep(2)\n",
    "print('✅ Server started on port', PORT)\n",
    "\n",
    "# ── Start cloudflare tunnel (no account or token needed) ────────────────────────\n",
    "print('Starting tunnel...')\n",
    "cf_proc = subprocess.Popen(\n",
    "    ['cloudflared', 'tunnel', '--url', f'http://localhost:{PORT}'],\n",
    "    stdout=subprocess.PIPE, stderr=subprocess.STDOUT\n",
    ")\n",
    "\n",
    "public_url = ''\n",
    "deadline = time.time() + 30\n",
    "while time.time() < deadline and not public_url:\n",
    "    line = cf_proc.stdout.readline().decode('utf-8', errors='replace')\n",
    "    m = re.search(r'https://[a-z0-9\\-]+\\.trycloudflare\\.com', line)\n",
    "    if m:\n",
    "        public_url = m.group(0)\n",
    "\n",
    "if not public_url:\n",
    "    public_url = f'http://localhost:{PORT}  ⚠️ tunnel failed — check cloudflared output'\n",
    "\n",
    "display(HTML(f\"\"\"\n",
    "<div style=\"\n",
    "  font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;\n",
    "  border: 2px solid #7c3aed;\n",
    "  border-radius: 14px;\n",
    "  padding: 24px 28px;\n",
    "  background: linear-gradient(135deg, #1e1b4b, #2e1065);\n",
    "  color: #ede9fe;\n",
    "  margin: 12px 0;\n",
    "  box-shadow: 0 4px 24px rgba(124,58,237,0.35);\n",
    "\">\n",
    "  <div style=\"font-size:13px;color:#a78bfa;font-weight:700;letter-spacing:0.08em;text-transform:uppercase;margin-bottom:8px\">\n",
    "    ✅ OmniFlows Colab Server is Running\n",
    "  </div>\n",
    "  <div style=\"font-size:11px;color:#c4b5fd;margin-bottom:6px\">\n",
    "    Wan 2.1 T2V 1.3B &bull; MMAudio small_16k &bull; T4 GPU &bull; No token needed\n",
    "  </div>\n",
    "  <div style=\"font-size:11px;color:{'#6ee7b7' if mmaudio_net is not None else '#f87171'};margin-bottom:16px\">\n",
    "    {'🔊 MMAudio ready — Add Audio endpoint active' if mmaudio_net is not None else '⚠️ MMAudio not loaded — video generation only'}\n",
    "  </div>\n",
    "  <div style=\"\n",
    "    background:rgba(255,255,255,0.07);\n",
    "    border:1px solid rgba(167,139,250,0.4);\n",
    "    border-radius:10px;\n",
    "    padding:14px 18px;\n",
    "    font-family:monospace;\n",
    "    font-size:16px;\n",
    "    font-weight:700;\n",
    "    color:#f5f3ff;\n",
    "    margin-bottom:16px;\n",
    "    word-break:break-all;\n",
    "  \" id=\"omni-url-box\">{public_url}</div>\n",
    "  <div style=\"display:flex;gap:10px;flex-wrap:wrap;align-items:center\">\n",
    "    <button\n",
    "      onclick=\"navigator.clipboard.writeText('{public_url}').then(()=>{{this.textContent='✓ Copied!';this.style.background='#059669';setTimeout(()=>{{this.textContent='📋 Copy URL';this.style.background='#7c3aed';}},2500);}});\"\n",
    "      style=\"background:#7c3aed;color:white;border:none;border-radius:8px;padding:10px 22px;font-size:14px;font-weight:700;cursor:pointer;\"\n",
    "    >📋 Copy URL</button>\n",
    "    <a href=\"{public_url}/health\" target=\"_blank\"\n",
    "      style=\"color:#a78bfa;font-size:13px;text-decoration:none;border:1px solid rgba(167,139,250,0.35);border-radius:8px;padding:10px 16px;font-weight:600;\"\n",
    "    >🔗 Test Health</a>\n",
    "  </div>\n",
    "  <div style=\"margin-top:18px;font-size:12px;color:#8b5cf6;line-height:1.7\">\n",
    "    <strong style=\"color:#c4b5fd\">Next step:</strong>\n",
    "    Click <strong>Copy URL</strong> &rarr; go to\n",
    "    <strong>OmniFlows Admin &rarr; Colab tab</strong> &rarr; paste &rarr; click <strong>Save URL</strong>.<br>\n",
    "    <span style=\"color:#6d28d9\">Keep this tab open to keep the server alive.</span>\n",
    "  </div>\n",
    "</div>\n",
    "\"\"\"))\n",
    "\n",
    "print(f'\\nPublic URL: {public_url}')\n",
    "print(f'Health check: {public_url}/health')"
   ]
  }
 ]
}