diff --git a/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_sglang.ipynb b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_sglang.ipynb new file mode 100644 index 00000000..ce7079e9 --- /dev/null +++ b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_sglang.ipynb @@ -0,0 +1,533 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Cosmos3 Nano Transfer with SGLang\n", + "\n", + "This notebook calls an already-running SGLang Cosmos3 server through the OpenAI-compatible video API. It reuses the checked-in transfer control assets and specs from this cookbook, then sends edge, blur, depth, segmentation, and world-scenario-map transfer requests to `POST /v1/videos`.\n", + "\n", + "The runnable examples use one control at a time. SGLang also accepts unweighted multi-control requests by passing a list of control videos in the API request. Per-hint `weight` is supported only by Cosmos Framework.\n", + "\n", + "The notebook does not modify the SGLang source tree.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Prerequisites\n", + "\n", + "Start a SGLang Cosmos3 server before running the request cells. The examples assume `${COSMOS3_WORKDIR}` points to cosmos repo.\n", + "\n", + "Transfer controls are available from SGLang `main` and the released `lmsysorg/sglang:dev` container. \n", + "\n", + "```bash\n", + "export SGLANG_HF_CACHE=\"${SGLANG_HF_CACHE:-$HOME/.cache/sglang-huggingface}\"\n", + "mkdir -p \"$SGLANG_HF_CACHE\"\n", + "chmod 777 \"$SGLANG_HF_CACHE\"\n", + "\n", + "export COSMOS3_WORKDIR=\"$PWD\"\n", + "\n", + "docker run --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"${COSMOS3_WORKDIR}:${COSMOS3_WORKDIR}\" \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Nano \\\n", + " --host 0.0.0.0\n", + "```\n", + "\n", + "Generator guardrails are off by default. To turn it on, set `COSMOS3_SGLANG_GUARDRAILS=true` for these sample requests, and it requires access to the gated `nvidia/Cosmos-1.0-Guardrail` repository. When guardrails are enabled, also set and pass `HF_TOKEN` to the container; it must be a Hugging Face token authorized for that repository. To disable guardrails server-wide, use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in sglang server.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Configure Paths and Endpoints\n", + "\n", + "Run this cell from anywhere inside the `cosmos` checkout. It resolves local assets, output paths, the SGLang endpoint, and repo-local `control_path` values.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import base64\n", + "import html\n", + "import json\n", + "import os\n", + "import shutil\n", + "import subprocess\n", + "import time\n", + "from IPython.display import HTML, display\n", + "\n", + "def find_repo_root(start: Path) -> Path:\n", + " for path in [start, *start.parents]:\n", + " if (path / \"README.md\").exists() and (path / \"cookbooks\").exists():\n", + " return path\n", + " return start\n", + "\n", + "\n", + "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", + "TRANSFER_ROOT = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"generator\" / \"transfer\"\n", + "SPECS_DIR = TRANSFER_ROOT / \"specs\"\n", + "ASSETS_DIR = TRANSFER_ROOT / \"assets\"\n", + "OUTPUT_ROOT = Path(\n", + " os.environ.get(\"COSMOS3_TRANSFER_SGLANG_OUTPUT_ROOT\", TRANSFER_ROOT / \"outputs\" / \"sglang\")\n", + ").resolve()\n", + "BASE_URL = os.environ.get(\"COSMOS3_SGLANG_BASE_URL\", \"http://localhost:30000\").rstrip(\"/\")\n", + "GUARDRAILS = os.environ.get(\"COSMOS3_SGLANG_GUARDRAILS\", \"false\").strip().lower() not in {\"0\", \"false\", \"no\", \"off\"}\n", + "\n", + "OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)\n", + "\n", + "print(\"COSMOS_ROOT:\", COSMOS_ROOT)\n", + "print(\"TRANSFER_ROOT:\", TRANSFER_ROOT)\n", + "print(\"OUTPUT_ROOT:\", OUTPUT_ROOT)\n", + "print(\"BASE_URL:\", BASE_URL)\n", + "print(\"GUARDRAILS:\", GUARDRAILS)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Verify Endpoint Configuration\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from urllib.parse import urlparse\n", + "\n", + "try:\n", + " import requests\n", + "except ImportError as exc:\n", + " raise RuntimeError(\"Install requests in this notebook kernel: pip install requests\") from exc\n", + "\n", + "\n", + "def api_root_url(base_url: str) -> str:\n", + " normalized = base_url.rstrip(\"/\")\n", + " if not normalized.endswith(\"/v1\"):\n", + " normalized = f\"{normalized}/v1\"\n", + " return normalized\n", + "\n", + "\n", + "API_ROOT = api_root_url(BASE_URL)\n", + "VIDEOS_GENERATION_URL = f\"{API_ROOT}/videos\"\n", + "MODELS_URL = f\"{API_ROOT}/models\"\n", + "parsed = urlparse(API_ROOT)\n", + "print(\"api root:\", API_ROOT)\n", + "print(\"videos generation:\", VIDEOS_GENERATION_URL)\n", + "print(\"models:\", MODELS_URL)\n", + "print(\"scheme:\", parsed.scheme)\n", + "print(\"host:\", parsed.netloc)\n", + "\n", + "response = requests.get(MODELS_URL, timeout=30)\n", + "response.raise_for_status()\n", + "print(response.json())\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Define Transfer Request Helpers\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "TRANSFER_CONTROLS = (\"edge\", \"blur\", \"depth\", \"seg\", \"wsm\")\n", + "\n", + "\n", + "def compact_json_file(path: Path) -> str:\n", + " return json.dumps(json.loads(path.read_text()), ensure_ascii=True, separators=(\",\", \":\"))\n", + "\n", + "\n", + "def resolve_spec_path(spec_path: Path, value: str) -> Path:\n", + " path = Path(value)\n", + " if path.is_absolute():\n", + " return path\n", + " return (spec_path.parent / path).resolve()\n", + "\n", + "\n", + "def repo_relative_path(local_path: Path) -> str:\n", + " return local_path.resolve().relative_to(COSMOS_ROOT).as_posix()\n", + "\n", + "\n", + "def spec_size(spec: dict) -> str:\n", + " height = int(spec[\"resolution\"])\n", + " width_ratio, height_ratio = (int(part) for part in spec[\"aspect_ratio\"].split(\",\", 1))\n", + " width = round(height * width_ratio / height_ratio)\n", + " return f\"{width}x{height}\"\n", + "\n", + "\n", + "def load_transfer_spec(control: str) -> tuple[Path, dict]:\n", + " if control not in TRANSFER_CONTROLS:\n", + " raise ValueError(f\"control must be one of {TRANSFER_CONTROLS}, got {control!r}\")\n", + " spec_path = SPECS_DIR / f\"{control}.json\"\n", + " if not spec_path.is_file():\n", + " raise FileNotFoundError(spec_path)\n", + " return spec_path, json.loads(spec_path.read_text())\n", + "\n", + "\n", + "def build_transfer_request(control: str) -> tuple[dict[str, str], Path, Path]:\n", + " spec_path, spec = load_transfer_spec(control)\n", + " hint = dict(spec[control])\n", + " local_control_path = resolve_spec_path(spec_path, hint[\"control_path\"])\n", + " hint[\"control_path\"] = local_control_path.resolve().as_posix()\n", + "\n", + " extra_params = {\n", + " \"use_resolution_template\": False,\n", + " \"use_duration_template\": False,\n", + " \"guardrails\": GUARDRAILS,\n", + " **hint,\n", + " \"control_hint\": control,\n", + " \"resolution\": spec[\"resolution\"],\n", + " \"control_guidance\": spec[\"control_guidance\"],\n", + " \"num_video_frames_per_chunk\": spec[\"num_video_frames_per_chunk\"],\n", + " \"num_conditional_frames\": spec.get(\"num_conditional_frames\", 1),\n", + " \"num_first_chunk_conditional_frames\": spec.get(\"num_first_chunk_conditional_frames\", 0),\n", + " \"share_vision_temporal_positions\": spec.get(\"share_vision_temporal_positions\", True),\n", + " \"max_frames\": spec[\"num_frames\"],\n", + " }\n", + " form = {\n", + " \"prompt\": compact_json_file(resolve_spec_path(spec_path, spec[\"prompt_path\"])),\n", + " \"negative_prompt\": compact_json_file(resolve_spec_path(spec_path, spec[\"negative_prompt_file\"])),\n", + " \"size\": spec_size(spec),\n", + " \"num_frames\": str(spec[\"num_frames\"]),\n", + " \"fps\": str(spec[\"fps\"]),\n", + " \"num_inference_steps\": \"50\",\n", + " \"guidance_scale\": str(spec[\"guidance\"]),\n", + " \"flow_shift\": \"10.0\",\n", + " \"seed\": \"2026\",\n", + " \"extra_params\": json.dumps(extra_params, separators=(\",\", \":\")),\n", + " }\n", + " output_path = OUTPUT_ROOT / control / f\"{spec['name']}.mp4\"\n", + " return form, local_control_path, output_path\n", + "\n", + "\n", + "def run_transfer(control: str) -> Path:\n", + " form, local_control_path, output_path = build_transfer_request(control)\n", + " output_path.parent.mkdir(parents=True, exist_ok=True)\n", + " error_path = output_path.with_suffix(\".error.txt\")\n", + " tmp_path = output_path.with_suffix(\".tmp\")\n", + " print(\"control:\", control)\n", + " print(\"local control:\", local_control_path)\n", + " print(\"size:\", form[\"size\"], \"frames:\", form[\"num_frames\"], \"fps:\", form[\"fps\"])\n", + " print(\"guardrails:\", json.loads(form[\"extra_params\"])[\"guardrails\"])\n", + " print(\"output:\", output_path)\n", + "\n", + " headers = {\"Accept\": \"video/mp4\"}\n", + " api_key = os.environ.get(\"COSMOS3_SGLANG_API_KEY\")\n", + " if api_key:\n", + " headers[\"Authorization\"] = f\"Bearer {api_key}\"\n", + "\n", + " t0 = time.time()\n", + " response = requests.post(VIDEOS_GENERATION_URL, json=form, headers=headers, timeout=3600)\n", + " if not response.ok:\n", + " error_path.write_text(response.text)\n", + " print(\"request failed:\", response.status_code)\n", + " print(response.text[:2000])\n", + " response.raise_for_status()\n", + "\n", + " initial = response.json()\n", + " (output_path.parent / \"response.json\").write_text(json.dumps(initial, indent=2))\n", + "\n", + " while True:\n", + " response = requests.get(f\"{API_ROOT}/videos/{initial['id']}\", timeout=30)\n", + " response.raise_for_status()\n", + " final = response.json()\n", + " (output_path.parent / \"final.json\").write_text(json.dumps(final, indent=2))\n", + " print(initial[\"id\"], final.get(\"status\"), f\"{final.get('progress', 0)}%\")\n", + " if final.get(\"status\") == \"completed\":\n", + " break\n", + " if final.get(\"status\") in {\"failed\", \"cancelled\"}:\n", + " raise RuntimeError(json.dumps(final, indent=2))\n", + " time.sleep(5)\n", + "\n", + " content_response = requests.get(f\"{API_ROOT}/videos/{initial['id']}/content\", timeout=300)\n", + " content_response.raise_for_status()\n", + " if content_response.content:\n", + " tmp_path.write_bytes(content_response.content)\n", + " tmp_path.rename(output_path)\n", + " else:\n", + " raise RuntimeError(\"video content endpoint returned an empty body\")\n", + "\n", + " print(f\"wrote {output_path} in {time.time() - t0:.1f}s\")\n", + " return output_path\n", + "\n", + "\n", + "def _ffmpeg_exe() -> str:\n", + " try:\n", + " import imageio_ffmpeg\n", + "\n", + " return imageio_ffmpeg.get_ffmpeg_exe()\n", + " except ImportError:\n", + " pass\n", + " exe = shutil.which(\"ffmpeg\")\n", + " if exe:\n", + " return exe\n", + " raise RuntimeError(\"Install imageio-ffmpeg or put ffmpeg on PATH to create compact previews.\")\n", + "\n", + "\n", + "def make_preview(src: Path, *, crf: int = 28) -> Path:\n", + " preview = src.with_name(f\"{src.stem}_preview.mp4\")\n", + " if not preview.exists() or preview.stat().st_mtime < src.stat().st_mtime:\n", + " subprocess.run(\n", + " [\n", + " _ffmpeg_exe(),\n", + " \"-y\",\n", + " \"-loglevel\",\n", + " \"error\",\n", + " \"-i\",\n", + " str(src),\n", + " \"-c:v\",\n", + " \"libx264\",\n", + " \"-crf\",\n", + " str(crf),\n", + " \"-preset\",\n", + " \"veryfast\",\n", + " \"-an\",\n", + " \"-pix_fmt\",\n", + " \"yuv420p\",\n", + " str(preview),\n", + " ],\n", + " check=True,\n", + " )\n", + " return preview\n", + "\n", + "\n", + "def display_video(path: Path, *, width: int = 720) -> None:\n", + " data = base64.b64encode(path.read_bytes()).decode(\"ascii\")\n", + " label = html.escape(str(path))\n", + " markup = f'''\n", + "\n", + "
{label}
\n", + "'''\n", + " display(HTML(markup))\n", + "\n", + "\n", + "def view_transfer(control: str, output_path: Path | None = None) -> None:\n", + " form, local_control_path, expected_output = build_transfer_request(control)\n", + " output_path = Path(output_path or expected_output)\n", + " if not output_path.is_file():\n", + " raise FileNotFoundError(f\"missing output: {output_path} (run {control} transfer first)\")\n", + " for label, src in [(\"control\", local_control_path), (\"generated\", output_path)]:\n", + " preview = make_preview(src)\n", + " print(f\"{control} {label}: {src.name} ({src.stat().st_size // 1024} KB -> {preview.stat().st_size // 1024} KB preview)\")\n", + " display_video(preview)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Preview Available Inputs\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for control in TRANSFER_CONTROLS:\n", + " form, local_control_path, _ = build_transfer_request(control)\n", + " prompt = json.loads(form[\"prompt\"])\n", + " caption = prompt.get(\"temporal_caption\") or prompt.get(\"comprehensive_t2i_caption\") or prompt.get(\"extra\", {}).get(\"prompt\", \"\")\n", + " print(f\"{control}: {local_control_path.relative_to(COSMOS_ROOT)}\")\n", + " print(f\" size={form['size']} frames={form['num_frames']} fps={form['fps']}\")\n", + " print(f\" prompt={caption[:180]}{'...' if len(caption) > 180 else ''}\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Edge (Canny) Transfer\n", + "\n", + "Run the `edge` transfer request through SGLang, then display the input control video and generated output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "edge_output = run_transfer(\"edge\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "view_transfer(\"edge\", edge_output)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Blur Transfer\n", + "\n", + "Run the `blur` transfer request through SGLang, then display the input control video and generated output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "blur_output = run_transfer(\"blur\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "view_transfer(\"blur\", blur_output)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Depth Transfer\n", + "\n", + "Run the `depth` transfer request through SGLang, then display the input control video and generated output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "depth_output = run_transfer(\"depth\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "view_transfer(\"depth\", depth_output)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 9. Segmentation Transfer\n", + "\n", + "Run the `seg` transfer request through SGLang, then display the input control video and generated output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "seg_output = run_transfer(\"seg\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "view_transfer(\"seg\", seg_output)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 10. World Scenario Map Transfer\n", + "\n", + "Run the `wsm` transfer request through SGLang, then display the input control video and generated output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "wsm_output = run_transfer(\"wsm\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "view_transfer(\"wsm\", wsm_output)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8501a859", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}