diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb index dc93da25..8b9cde2f 100644 --- a/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_fd_with_sglang.ipynb @@ -24,7 +24,17 @@ "\n", "Forward dynamics predicts future visual observations from an initial image and an action trajectory. This notebook contains separate AV and robotics sections that each build their own input spec, run inference, and visualize generated videos.\n", "\n", - "Start the SGLang server:\n", + "### Set up HF Cache:\n", + "\n", + "SGLang container is run under root by default. User can create a writable cache directory on host, or start the container with host user id.\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", + "\n", + "### Start the SGLang server:\n", "\n", "```bash\n", "docker rm -f cosmos3-sglang-notebook 2>/dev/null || true\n", @@ -32,7 +42,7 @@ "docker run -d --init --name cosmos3-sglang-notebook \\\n", " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", " -v \"$PWD:/workspace\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", @@ -1296,7 +1306,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3", "language": "python", "name": "python3" }, diff --git a/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb index c10b4dff..48dfbe3f 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_sglang.ipynb @@ -28,7 +28,17 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Start SGLang Server\n", + "### Set up HF Cache:\n", + "\n", + "SGLang container is run under root by default. User can create a writable cache directory on host, or start the container with host user id.\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", + "\n", + "### Start SGLang Server\n", "\n", "Start the server in a terminal from the `cosmos` repo root.\n", "\n", @@ -38,7 +48,7 @@ "docker run -d --init --name cosmos3-sglang-notebook \\\n", " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", " -v \"$PWD:/workspace\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", @@ -54,6 +64,7 @@ { "cell_type": "code", "execution_count": null, + "id": "4e48ec6b", "metadata": {}, "outputs": [], "source": [ @@ -88,6 +99,7 @@ }, { "cell_type": "markdown", + "id": "54674249", "metadata": {}, "source": [ "## Create the Inverse-Dynamics Input Spec\n", @@ -199,6 +211,7 @@ }, { "cell_type": "markdown", + "id": "8240d23d", "metadata": {}, "source": [ "## Run Inverse-Dynamics Inference\n", @@ -215,6 +228,7 @@ { "cell_type": "code", "execution_count": null, + "id": "d8e1a916", "metadata": {}, "outputs": [], "source": [ @@ -234,6 +248,16 @@ " print(response.json())\n", "\n", "\n", + "def log_sglang_error(response: requests.Response, run_dir: Path, label: str, context: dict | None = None) -> None:\n", + " error_path = run_dir / f\"{label}_error_response.txt\"\n", + " error_path.write_text(response.text)\n", + " print(f\"SGLang {label} request failed:\", response.status_code)\n", + " print(response.text)\n", + " print(\"saved error response:\", error_path)\n", + " if context is not None:\n", + " print(\"context:\", json.dumps(context, indent=2))\n", + "\n", + "\n", "def submit_inverse_dynamics(record: dict) -> dict:\n", " run_dir = id_output_dir / record[\"name\"]\n", " run_dir.mkdir(parents=True, exist_ok=True)\n", @@ -242,7 +266,7 @@ " extra_params = {\n", " \"action_mode\": \"inverse_dynamics\",\n", " \"domain_name\": record[\"domain_name\"],\n", - " \"view_point\": record[\"view_point\"],\n", + " \"action_view_point\": record[\"view_point\"],\n", " \"raw_action_dim\": 9,\n", " \"guardrails\": False,\n", " }\n", @@ -259,38 +283,29 @@ "\n", " with video_path.open(\"rb\") as video_file:\n", " response = requests.post(\n", - " f\"{SGLANG_BASE_URL}/v1/videos\",\n", + " f\"{SGLANG_BASE_URL}/v1/actions/generations\",\n", " data={key: str(value) for key, value in form.items()},\n", " files={\"input_reference\": (video_path.name, video_file, \"video/mp4\")},\n", " timeout=120,\n", " )\n", - " response.raise_for_status()\n", - " initial = response.json()\n", - " (run_dir / \"response.json\").write_text(json.dumps(initial, indent=2))\n", - "\n", - " while True:\n", - " response = requests.get(f\"{SGLANG_BASE_URL}/v1/videos/{initial['id']}\", timeout=30)\n", + " if not response.ok:\n", + " log_sglang_error(\n", + " response,\n", + " run_dir,\n", + " \"submit\",\n", + " {\"form\": form, \"extra_params\": extra_params, \"video_path\": str(video_path)},\n", + " )\n", " response.raise_for_status()\n", - " final = response.json()\n", - " (run_dir / \"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(2)\n", - "\n", - " action = final.get(\"action\")\n", - " if not action or \"data\" not in action:\n", - " raise RuntimeError(f\"SGLang response did not include action data: {json.dumps(final, indent=2)}\")\n", + " response.raise_for_status()\n", + " action = response.json()['data'][0]['action']\n", " (run_dir / \"action.json\").write_text(json.dumps(action, indent=2))\n", "\n", - " sample_outputs = {\"outputs\": [{\"content\": {\"action\": action[\"data\"]}}]}\n", + " sample_outputs = {\"outputs\": [{\"content\": {\"action\": action.get(\"values\")}}]}\n", " (run_dir / \"sample_outputs.json\").write_text(json.dumps(sample_outputs, indent=2))\n", "\n", " print(\"saved\", run_dir / \"sample_outputs.json\")\n", " print(\"action shape:\", action.get(\"shape\"), \"dtype:\", action.get(\"dtype\"))\n", - " return {\"record\": record, \"initial\": initial, \"final\": final, \"run_dir\": run_dir, \"action\": action}\n", + " return {\"record\": record, \"run_dir\": run_dir, \"action\": action}\n", "\n", "\n", "check_sglang_server()\n", @@ -402,7 +417,7 @@ "for record in records:\n", " name = record[\"name\"]\n", " outputs = json.loads((id_output_dir / name / \"sample_outputs.json\").read_text())\n", - " poses_rel = np.array(outputs[\"outputs\"][0][\"content\"][\"action\"][0]) # [T-1, 9] = [translation(3), rot6d(6)]\n", + " poses_rel = np.array(outputs[\"outputs\"][0][\"content\"][\"action\"]) # [T-1, 9] = [translation(3), rot6d(6)]\n", "\n", " # AV action convention (see cosmos_framework/data/vfm/action/av_dataset.py):\n", " # rot6d rotation, backward_framewise, translation_scale = 1.35.\n", @@ -427,7 +442,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3", "language": "python", "name": "python3" }, diff --git a/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb b/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb index c2744f98..dce1897d 100644 --- a/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/action/run_policy_with_sglang.ipynb @@ -24,7 +24,7 @@ "\n", "This notebook runs Cosmos3 Nano **action policy** inference through SGLang using the checked-in DROID LeRobot sample under `assets/droid_lerobot_example`.\n", "\n", - "It sends `POST /v1/videos` requests with a first frame and instruction, then retrieves a rollout video plus top-level `action` metadata." + "It first sends a `POST /v1/actions/generations` request with a first frame and instruction to predict a DROID policy action vector. It can then send that action vector to the forward-dynamics `POST /v1/videos` endpoint to render a rollout video." ] }, { @@ -32,7 +32,17 @@ "id": "policy-server-md", "metadata": {}, "source": [ - "## Start SGLang Policy Server\n", + "### Set up HF Cache:\n", + "\n", + "SGLang container is run under root by default. User can create a writable cache directory on host, or start the container with host user id.\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", + "\n", + "### Start SGLang Server\n", "\n", "Start the server in a terminal from the `cosmos` repo root.\n", "\n", @@ -43,7 +53,7 @@ " --runtime nvidia --gpus '\"device=0\"' \\\n", " -e CUDA_DEVICE_ORDER=PCI_BUS_ID \\\n", " -e PYTHONPATH=/workspace/cosmos-framework \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", " -v \"$PWD:/workspace\" \\\n", " -p 30000:30000 --ipc=host \\\n", " lmsysorg/sglang:dev \\\n", @@ -59,7 +69,9 @@ "\n", "```bash\n", "docker logs -f cosmos3-sglang-policy-notebook\n", - "```" + "```\n", + "\n", + "The notebook uses this same server for both the policy action request and the rollout-video request by default." ] }, { @@ -87,6 +99,7 @@ "COSMOS3_ACTION_ROOT = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"generator\" / \"action\"\n", "DROID_ASSET_ROOT = COSMOS3_ACTION_ROOT / \"assets\" / \"droid_lerobot_example\"\n", "SGLANG_BASE_URL = os.environ.get(\"COSMOS3_SGLANG_BASE_URL\", \"http://localhost:30000\").rstrip(\"/\")\n", + "SGLANG_FD_BASE_URL = os.environ.get(\"COSMOS3_SGLANG_FD_BASE_URL\", SGLANG_BASE_URL).rstrip(\"/\")\n", "SGLANG_MODEL = os.environ.get(\"COSMOS3_SGLANG_MODEL\", \"nvidia/Cosmos3-Nano-Policy-DROID\")\n", "\n", "COSMOS3_INPUT_DIR.mkdir(parents=True, exist_ok=True)\n", @@ -97,6 +110,7 @@ "print(\"COSMOS3_INPUT_DIR:\", COSMOS3_INPUT_DIR)\n", "print(\"COSMOS3_POLICY_OUTPUT_DIR:\", COSMOS3_POLICY_OUTPUT_DIR)\n", "print(\"COSMOS3_SGLANG_BASE_URL:\", SGLANG_BASE_URL)\n", + "print(\"COSMOS3_SGLANG_FD_BASE_URL:\", SGLANG_FD_BASE_URL)\n", "print(\"COSMOS3_SGLANG_MODEL:\", SGLANG_MODEL)" ] }, @@ -182,9 +196,9 @@ "id": "policy-video-md", "metadata": {}, "source": [ - "## Run Policy Inference Through `/v1/videos`\n", + "## Run Policy Inference and Render Rollout\n", "\n", - "This path behaves like the forward/inverse SGLANG action notebooks: it sends a multipart request to the OpenAI-compatible video API, polls the async job, writes the generated rollout video, and saves the predicted action from the response metadata." + "First, this sends a multipart request to `/v1/actions/generations` to get a DROID policy action vector. Then it sends a second request to the forward-dynamics `/v1/videos` endpoint, using that policy vector as the action condition, and saves the generated rollout video." ] }, { @@ -195,9 +209,12 @@ "outputs": [], "source": [ "import json\n", + "import mimetypes\n", "import time\n", "from pathlib import Path\n", "\n", + "from PIL import Image\n", + "\n", "try:\n", " import requests\n", "except ImportError as exc:\n", @@ -240,7 +257,16 @@ " return min(candidates, key=lambda size: abs(input_ratio - size[1] / size[0]))\n", "\n", "\n", - "def submit_policy_video() -> dict:\n", + "def action_values_from_response(action: dict) -> list:\n", + " values = action.get(\"values\") or action.get(\"data\")\n", + " if values is None:\n", + " raise RuntimeError(f\"Policy response did not include action values: {json.dumps(action, indent=2)}\")\n", + " if values and isinstance(values[0], list) and values[0] and isinstance(values[0][0], list):\n", + " values = values[0]\n", + " return values\n", + "\n", + "\n", + "def submit_policy_action() -> dict:\n", " run_dir = COSMOS3_POLICY_OUTPUT_DIR / \"video_api\"\n", " run_dir.mkdir(parents=True, exist_ok=True)\n", "\n", @@ -267,7 +293,7 @@ "\n", " with policy_image_path.open(\"rb\") as image_file:\n", " response = requests.post(\n", - " f\"{SGLANG_BASE_URL}/v1/videos\",\n", + " f\"{SGLANG_BASE_URL}/v1/actions/generations\",\n", " data={key: str(value) for key, value in form.items()},\n", " files={\"input_reference\": (policy_image_path.name, image_file, \"image/png\")},\n", " timeout=120,\n", @@ -279,45 +305,102 @@ " print(\"form:\", json.dumps(form, indent=2))\n", " response.raise_for_status()\n", "\n", - " initial = response.json()\n", - " (run_dir / \"response.json\").write_text(json.dumps(initial, indent=2))\n", + " raw = response.json()\n", + " (run_dir / \"policy_response.json\").write_text(json.dumps(raw, indent=2))\n", + " action = raw[\"data\"][0][\"action\"]\n", + " action_values = action_values_from_response(action)\n", + " (run_dir / \"action.json\").write_text(json.dumps(action, indent=2))\n", + " (run_dir / \"policy_action_for_fd.json\").write_text(json.dumps(action_values, indent=2))\n", + " sample_outputs = {\"outputs\": [{\"content\": {\"action\": action_values}}]}\n", + " (run_dir / \"sample_outputs.json\").write_text(json.dumps(sample_outputs, indent=2))\n", "\n", - " while True:\n", - " response = requests.get(f\"{SGLANG_BASE_URL}/v1/videos/{initial['id']}\", timeout=30)\n", + " print(\"saved\", run_dir / \"action.json\")\n", + " print(\"action shape:\", action.get(\"shape\"), \"dtype:\", action.get(\"dtype\"), \"domain_id:\", action.get(\"domain_id\"))\n", + " return {\"run_dir\": run_dir, \"raw\": raw, \"video_path\": None, \"action\": action, \"action_values\": action_values}\n", + "\n", + "\n", + "def submit_forward_dynamics_rollout(policy_result: dict) -> Path:\n", + " run_dir = policy_result[\"run_dir\"]\n", + " action_values = policy_result[\"action_values\"]\n", + " action_path = run_dir / \"policy_action_for_fd.json\"\n", + "\n", + " input_width, input_height = Image.open(policy_image_path).size\n", + " target_width = (input_width // 16) * 16\n", + " target_height = (input_height // 16) * 16\n", + " mime_type = mimetypes.guess_type(policy_image_path.name)[0] or \"image/png\"\n", + " extra_params = {\n", + " \"action_mode\": \"forward_dynamics\",\n", + " \"domain_name\": \"droid_lerobot\",\n", + " \"action_view_point\": \"concat_view\",\n", + " \"action\": action_values,\n", + " \"guardrails\": False,\n", + " }\n", + " form = {\n", + " \"prompt\": policy_prompt,\n", + " \"num_frames\": len(action_values) + 1,\n", + " \"fps\": 15,\n", + " \"size\": f\"{target_width}x{target_height}\",\n", + " \"num_inference_steps\": 30,\n", + " \"guidance_scale\": 1.0,\n", + " \"flow_shift\": 10.0,\n", + " \"seed\": 0,\n", + " \"extra_params\": json.dumps(extra_params),\n", + " }\n", + "\n", + " print(\"rendering rollout from policy action:\", action_path)\n", + " print(\"forward-dynamics server:\", SGLANG_FD_BASE_URL)\n", + " with policy_image_path.open(\"rb\") as image_file:\n", + " response = requests.post(\n", + " f\"{SGLANG_FD_BASE_URL}/v1/videos\",\n", + " data={key: str(value) for key, value in form.items()},\n", + " files={\"input_reference\": (policy_image_path.name, image_file, mime_type)},\n", + " timeout=120,\n", + " )\n", + " if not response.ok:\n", + " (run_dir / \"rollout_error_response.txt\").write_text(response.text)\n", + " print(\"SGLang rollout request failed:\", response.status_code)\n", + " print(response.text)\n", + " print(\"form:\", json.dumps(form, indent=2))\n", + " print(\"extra_params keys:\", sorted(extra_params))\n", + " print(\"action shape:\", [len(action_values), len(action_values[0]) if action_values else 0])\n", " response.raise_for_status()\n", + "\n", + " initial = response.json()\n", + " (run_dir / \"rollout_response.json\").write_text(json.dumps(initial, indent=2))\n", + " while True:\n", + " response = requests.get(f\"{SGLANG_FD_BASE_URL}/v1/videos/{initial['id']}\", timeout=30)\n", + " if not response.ok:\n", + " (run_dir / \"rollout_poll_error_response.txt\").write_text(response.text)\n", + " print(\"SGLang rollout poll failed:\", response.status_code)\n", + " print(response.text)\n", + " response.raise_for_status()\n", " final = response.json()\n", - " (run_dir / \"final.json\").write_text(json.dumps(final, indent=2))\n", + " (run_dir / \"rollout_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", + " raise RuntimeError(\n", + " f\"SGLang rollout job {final.get('status')} for {initial['id']}. \"\n", + " f\"Full response written to {run_dir / 'rollout_final.json'}:\\n{json.dumps(final, indent=2)}\"\n", + " )\n", " time.sleep(2)\n", "\n", - " action = final.get(\"action\")\n", - " if not action or \"data\" not in action:\n", - " raise RuntimeError(f\"SGLang response did not include action data: {json.dumps(final, indent=2)}\")\n", - " (run_dir / \"action.json\").write_text(json.dumps(action, indent=2))\n", - " sample_outputs = {\"outputs\": [{\"content\": {\"action\": action[\"data\"]}}]}\n", - " (run_dir / \"sample_outputs.json\").write_text(json.dumps(sample_outputs, indent=2))\n", - "\n", - " content_response = requests.get(f\"{SGLANG_BASE_URL}/v1/videos/{initial['id']}/content\", timeout=300)\n", - " content_response.raise_for_status()\n", + " response = requests.get(f\"{SGLANG_FD_BASE_URL}/v1/videos/{initial['id']}/content\", timeout=300)\n", + " if not response.ok:\n", + " (run_dir / \"rollout_content_error_response.txt\").write_text(response.text)\n", + " print(\"SGLang rollout content request failed:\", response.status_code)\n", + " print(response.text)\n", + " response.raise_for_status()\n", " video_path = run_dir / \"policy_rollout.mp4\"\n", - " if content_response.content:\n", - " video_path.write_bytes(content_response.content)\n", - " print(\"saved\", video_path)\n", - " else:\n", - " video_path = None\n", - " print(\"video content endpoint returned an empty body\")\n", - "\n", - " print(\"saved\", run_dir / \"action.json\")\n", - " print(\"action shape:\", action.get(\"shape\"), \"dtype:\", action.get(\"dtype\"), \"domain_id:\", action.get(\"domain_id\"))\n", - " return {\"initial\": initial, \"final\": final, \"run_dir\": run_dir, \"video_path\": video_path, \"action\": action}\n", + " video_path.write_bytes(response.content)\n", + " print(\"saved\", video_path)\n", + " return video_path\n", "\n", "\n", "check_sglang_server()\n", - "policy_video_result = submit_policy_video()" + "policy_video_result = submit_policy_action()\n", + "policy_video_result[\"video_path\"] = submit_forward_dynamics_rollout(policy_video_result)" ] }, { @@ -325,9 +408,9 @@ "id": "policy-preview-md", "metadata": {}, "source": [ - "## Inspect Video API Outputs\n", + "## Inspect Policy and Rollout Outputs\n", "\n", - "Preview the rollout video if the server returned one, and print the first few predicted action rows." + "Print the first few predicted action rows and preview the rollout video generated by the forward-dynamics request." ] }, { @@ -373,12 +456,13 @@ "\n", "\n", "action = policy_video_result[\"action\"]\n", - "action_array = np.asarray(action[\"data\"], dtype=np.float32)\n", + "action_array = np.asarray(policy_video_result[\"action_values\"], dtype=np.float32)\n", "print(\"action array:\", action_array.shape, action_array.dtype)\n", "print(action_array[: min(5, len(action_array))])\n", "\n", "video_path = policy_video_result.get(\"video_path\")\n", "if video_path is not None:\n", + " assert Path(video_path).exists(), f\"missing rollout video: {video_path}\"\n", " preview = make_preview(video_path)\n", " print(f\"preview: {preview}\")\n", " display(Video(str(preview), embed=True))" @@ -395,7 +479,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3", "language": "python", "name": "python3" }, diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_sglang.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_sglang.ipynb index 86f742f8..35051d20 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_sglang.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_sglang.ipynb @@ -46,11 +46,21 @@ "\n", "Run the SGLang server before running the request cells. Use the Docker image for every modality on this page. Mount any directory that contains local media or action files you want the server to read.\n", "\n", + "### Set up HF Cache\n", + "\n", + "SGLang container is run under root by default. User can create a writable cache directory on host, or start the container with host user id.\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", + "\n", "### Docker Image: Cosmos3-Nano\n", "\n", "```bash\n", "docker run --runtime nvidia --gpus all \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", " -v \"$(pwd):/workspace\" \\\n", " -p 30000:30000 \\\n", " --ipc=host \\\n", @@ -68,7 +78,7 @@ "\n", "```bash\n", "docker run --runtime nvidia --gpus all \\\n", - " -v ~/.cache/huggingface:/root/.cache/huggingface \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", " -v \"$(pwd):/workspace\" \\\n", " -p 30000:30000 \\\n", " --ipc=host \\\n", @@ -93,7 +103,35 @@ " --cfg-parallel-size 2\n", "```\n", "\n", - "For Cosmos3, set CFG strength with the request-level `guidance_scale` field. Do not use `true_cfg_scale` for CFG Parallel with these Cosmos3 examples.\n" + "For Cosmos3, set CFG strength with the request-level `guidance_scale` field. Do not use `true_cfg_scale` for CFG Parallel with these Cosmos3 examples.\n", + "\n", + "### Docker Image: Cosmos3-Edge\n", + "\n", + "```bash\n", + "docker run --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"$(pwd):/workspace\" \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Edge \\\n", + " --host 0.0.0.0\n", + "```\n", + "\n", + "### Docker Image: Cosmos3-Distil\n", + "\n", + "```bash\n", + "docker run --runtime nvidia --gpus all \\\n", + " -v $SGLANG_HF_CACHE:/root/.cache/huggingface \\\n", + " -v \"$(pwd):/workspace\" \\\n", + " -p 30000:30000 \\\n", + " --ipc=host \\\n", + " lmsysorg/sglang:dev \\\n", + " sglang serve \\\n", + " --model-path nvidia/Cosmos3-Super-Image2Video-4Step \\\n", + " --host 0.0.0.0\n", + "```" ] }, { @@ -133,6 +171,8 @@ "SGLANG_ENDPOINTS = {\n", " \"Cosmos3-Nano\": os.environ.get(\"COSMOS3_SGLANG_NANO_BASE_URL\", DEFAULT_SGLANG_BASE_URL),\n", " \"Cosmos3-Super\": os.environ.get(\"COSMOS3_SGLANG_SUPER_BASE_URL\", DEFAULT_SGLANG_BASE_URL),\n", + " \"Cosmos3-Edge\": os.environ.get(\"COSMOS3_SGLANG_EDGE_BASE_URL\", DEFAULT_SGLANG_BASE_URL),\n", + " \"Cosmos3-Distil\": os.environ.get(\"COSMOS3_SGLANG_DISTIL_BASE_URL\", DEFAULT_SGLANG_BASE_URL),\n", "}\n", "\n", "os.environ[\"COSMOS3_AUDIOVISUAL_OUTPUT_ROOT\"] = str(COSMOS3_AUDIOVISUAL_OUTPUT_ROOT)\n", @@ -266,12 +306,30 @@ " \"prompt\": \"assets/prompts/text2image/robot_draping.json\",\n", " \"enable_sound\": False,\n", " },\n", + " \"t2i_edge\": {\n", + " \"model\": \"Cosmos3-Edge\",\n", + " \"mode\": \"text2image\",\n", + " \"prompt\": \"assets/prompts/text2image/robot_draping.json\",\n", + " \"enable_sound\": False,\n", + " },\n", + " \"t2i_distil\": {\n", + " \"model\": \"Cosmos3-Distil\",\n", + " \"mode\": \"text2image\",\n", + " \"prompt\": \"assets/prompts/text2image/robot_draping.json\",\n", + " \"enable_sound\": False,\n", + " },\n", " \"t2v_nano_noaudio\": {\n", " \"model\": \"Cosmos3-Nano\",\n", " \"mode\": \"text2video\",\n", " \"prompt\": \"assets/prompts/text2video/robot_kitchen.json\",\n", " \"enable_sound\": False,\n", " },\n", + " \"t2v_edge_noaudio\": {\n", + " \"model\": \"Cosmos3-Edge\",\n", + " \"mode\": \"text2video\",\n", + " \"prompt\": \"assets/prompts/text2video/robot_kitchen.json\",\n", + " \"enable_sound\": False,\n", + " },\n", " \"t2vs\": {\n", " \"model\": \"Cosmos3-Nano\",\n", " \"mode\": \"text2video\",\n", @@ -285,6 +343,20 @@ " \"image\": \"assets/images/image2video/car_driving.jpg\",\n", " \"enable_sound\": False,\n", " },\n", + " \"i2v_edge_noaudio\": {\n", + " \"model\": \"Cosmos3-Edge\",\n", + " \"mode\": \"image2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"image\": \"assets/images/image2video/car_driving.jpg\",\n", + " \"enable_sound\": False,\n", + " },\n", + " \"i2v_distil_noaudio\": {\n", + " \"model\": \"Cosmos3-Distil\",\n", + " \"mode\": \"image2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"image\": \"assets/images/image2video/car_driving.jpg\",\n", + " \"enable_sound\": False,\n", + " },\n", " \"i2vs\": {\n", " \"model\": \"Cosmos3-Nano\",\n", " \"mode\": \"image2video\",\n", @@ -315,6 +387,11 @@ " },\n", "}\n", "\n", + "def _is_distil(model_name: str) -> bool:\n", + " return \"Distil\" in model_name\n", + "\n", + "def _is_edge(model_name: str) -> bool:\n", + " return \"Edge\" in model_name\n", "\n", "def asset_path(relative_path: str) -> Path:\n", " path = COSMOS3_AUDIOVISUAL_ROOT / relative_path\n", @@ -327,27 +404,15 @@ " return json.dumps(json.loads(path.read_text()), ensure_ascii=True, separators=(\",\", \":\"))\n", "\n", "\n", - "def normalize_negative_prompt(value) -> str:\n", - " if value is None:\n", - " return \"\"\n", - " if isinstance(value, str):\n", - " return value\n", - " if isinstance(value, dict):\n", - " parts = []\n", - " for v in value.values():\n", - " if isinstance(v, str):\n", - " parts.append(v)\n", - " elif isinstance(v, list):\n", - " parts.extend(str(x) for x in v)\n", - " return \"\\n\".join(parts)\n", - " return str(value)\n", - "\n", - "\n", "def payload_dimensions(payload: dict) -> tuple[int, int]:\n", " if payload.get(\"resolution\") == \"720\" and payload.get(\"aspect_ratio\") == \"16,9\":\n", " return 720, 1280\n", - " if payload.get(\"resolution\") == \"256\" and payload.get(\"aspect_ratio\") == \"16,9\":\n", - " return 192, 320\n", + " if payload.get(\"resolution\") == \"480\" and payload.get(\"aspect_ratio\") == \"16,9\":\n", + " return 480, 832\n", + " if payload.get(\"resolution\") == \"640\" and payload.get(\"aspect_ratio\") == \"1,1\":\n", + " return 640, 640\n", + " if payload.get(\"resolution\") == \"256\" and payload.get(\"aspect_ratio\") == \"1,1\":\n", + " return 256, 256\n", " raise ValueError(f\"Unsupported payload resolution/aspect ratio: {payload.get('resolution')} {payload.get('aspect_ratio')}\")\n", "\n", "\n", @@ -368,8 +433,8 @@ " prompt_path = asset_path(spec[\"prompt\"])\n", " negative_prompt = \"\"\n", " if spec[\"mode\"] != \"text2image\":\n", - " negative_prompt_path = asset_path(f\"assets/negative_prompts/{spec['mode']}/neg_prompt.json\") if not spec.get(\"negative_prompt\") else spec[\"negative_prompt\"]\n", - " negative_prompt = normalize_negative_prompt(negative_prompt_path)\n", + " negative_prompt_path = asset_path(f\"assets/negative_prompts/{spec['mode']}/neg_prompt.json\") if not spec.get(\"negative_prompt\") else asset_path(spec[\"negative_prompt\"])\n", + " negative_prompt = compact_json_file(negative_prompt_path)\n", " payload_path = payload_dir / f\"{use_case}.json\"\n", " payload = {\n", " \"model_mode\": spec[\"mode\"],\n", @@ -387,6 +452,16 @@ " video_path = asset_path(spec[\"video\"])\n", " payload[\"vision_path\"] = os.path.relpath(video_path, payload_path.parent)\n", "\n", + " if _is_distil(spec[\"model\"]) or _is_edge(spec[\"model\"]):\n", + " payload.pop(\"guidance\", None)\n", + " payload.pop(\"shift\", None)\n", + "\n", + " if _is_distil(spec[\"model\"]):\n", + " payload.pop(\"num_steps\", None)\n", + "\n", + " if _is_edge(spec[\"model\"]):\n", + " payload[\"resolution\"] = \"480\"\n", + " \n", " payload_path.write_text(json.dumps(payload, indent=2) + \"\\n\")\n", "\n", " os.environ[f\"COSMOS3_{backend.upper()}_{use_case.upper()}_INPUT\"] = str(payload_path)\n", @@ -405,7 +480,7 @@ " video_display_path = resolve_payload_path(payload_path, payload[\"vision_path\"])\n", " print(f\"video: {video_display_path.relative_to(COSMOS_ROOT)}\")\n", " display(Video(filename=str(video_display_path), width=420))\n", - " print(json.dumps({k: payload[k] for k in [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]}, indent=2))\n", + " print(json.dumps({k: payload[k] for k in [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"] if k in payload}, indent=2))\n", " return payload_path, output_dir, spec[\"model\"]\n", "\n", "\n", @@ -446,9 +521,9 @@ " \"size\": f\"{width}x{height}\",\n", " \"num_frames\": str(payload[\"num_frames\"]),\n", " \"fps\": str(payload[\"fps\"]),\n", - " \"num_inference_steps\": str(payload[\"num_steps\"]),\n", - " \"guidance_scale\": str(payload[\"guidance\"]),\n", - " \"flow_shift\": str(payload[\"shift\"]),\n", + " \"num_inference_steps\": str(payload[\"num_steps\"]) if \"num_steps\" in payload else None,\n", + " \"guidance_scale\": str(payload[\"guidance\"]) if \"guidance\" in payload else None,\n", + " \"flow_shift\": str(payload[\"shift\"]) if \"shift\" in payload else None,\n", " \"seed\": str(payload[\"seed\"]),\n", " \"extra_params\": json.dumps(extra_params, separators=(\",\", \":\")),\n", " }\n", @@ -465,9 +540,9 @@ " \"negative_prompt\": payload.get(\"negative_prompt\", \"\"),\n", " \"size\": f\"{width}x{height}\",\n", " \"n\": 1,\n", - " \"num_inference_steps\": payload[\"num_steps\"],\n", - " \"guidance_scale\": payload[\"guidance\"],\n", - " \"flow_shift\": payload[\"shift\"],\n", + " \"num_inference_steps\": payload[\"num_steps\"] if \"num_steps\" in payload else None,\n", + " \"guidance_scale\": payload[\"guidance\"] if \"guidance\" in payload else None,\n", + " \"flow_shift\": payload[\"shift\"] if \"shift\" in payload else None,\n", " \"seed\": payload[\"seed\"],\n", " \"response_format\": \"b64_json\",\n", " \"extra_params\": {\n", @@ -501,6 +576,8 @@ " cmd += [\"-H\", f\"Authorization: Bearer {api_key}\"]\n", "\n", " for key, value in build_sglang_video_form(payload).items():\n", + " if value is None:\n", + " continue\n", " cmd += [\"--form-string\", f\"{key}={value}\"]\n", "\n", " if payload[\"model_mode\"] == \"image2video\":\n", @@ -1182,11 +1259,336 @@ "source": [ "view_run(i2v_super_noaudio_output)\n" ] + }, + { + "cell_type": "markdown", + "id": "cd022d1d", + "metadata": {}, + "source": [ + "# Cosmos3-Edge Examples\n", + "\n", + "Use cases for the smaller `Cosmos3-Edge` model. This section is self-contained; you can run it without the other sections above.\n" + ] + }, + { + "cell_type": "markdown", + "id": "0b168305", + "metadata": {}, + "source": [ + "## Edge: Text to Image\n", + "\n", + "Edge text-to-image generation using the same structured JSON prompt.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "af8ddac1", + "metadata": {}, + "outputs": [], + "source": [ + "t2i_edge_payload, t2i_edge_output, t2i_edge_model = create_payload(\"t2i_edge\", backend=\"sglang\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "a42cb79b", + "metadata": {}, + "source": [ + "### Run\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c5515508", + "metadata": {}, + "outputs": [], + "source": [ + "run_sglang_payload(t2i_edge_payload, t2i_edge_output, model=\"Cosmos3-Edge\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "0502a370", + "metadata": {}, + "source": [ + "### View Results\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59b155c8", + "metadata": {}, + "outputs": [], + "source": [ + "view_run(t2i_edge_output)\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "c0592326", + "metadata": {}, + "source": [ + "## Edge: Text to Video Without Audio\n", + "\n", + "Edge text-to-video generation with audio disabled.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ed85b5ff", + "metadata": {}, + "outputs": [], + "source": [ + "t2v_edge_noaudio_payload, t2v_edge_noaudio_output, t2v_edge_noaudio_model = create_payload(\"t2v_edge_noaudio\", backend=\"sglang\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "db0e6065", + "metadata": {}, + "source": [ + "### Run\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c8e28db", + "metadata": {}, + "outputs": [], + "source": [ + "run_sglang_payload(t2v_edge_noaudio_payload, t2v_edge_noaudio_output, model=\"Cosmos3-Edge\")\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "63168603", + "metadata": {}, + "source": [ + "### View Results\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9702913c", + "metadata": {}, + "outputs": [], + "source": [ + "view_run(t2v_edge_noaudio_output)\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "3c3b2d01", + "metadata": {}, + "source": [ + "## Edge: Image to Video Without Audio\n", + "\n", + "Edge image-to-video generation using its paired image asset, with audio disabled.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f4640e93", + "metadata": {}, + "outputs": [], + "source": [ + "i2v_edge_noaudio_payload, i2v_edge_noaudio_output, i2v_edge_noaudio_model = create_payload(\"i2v_edge_noaudio\", backend=\"sglang\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "8d5d0873", + "metadata": {}, + "source": [ + "### Run\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2721deaa", + "metadata": {}, + "outputs": [], + "source": [ + "run_sglang_payload(i2v_edge_noaudio_payload, i2v_edge_noaudio_output, model=\"Cosmos3-Edge\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "13ae89e3", + "metadata": {}, + "source": [ + "### View Results\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10f39756", + "metadata": {}, + "outputs": [], + "source": [ + "view_run(i2v_edge_noaudio_output)\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "d862a1f6", + "metadata": {}, + "source": [ + "# Cosmos3-Distil Examples\n", + "\n", + "Use cases for the smaller `Cosmos3-Distil` model. This section is self-contained; you can run it without the other sections above.\n" + ] + }, + { + "cell_type": "markdown", + "id": "7bb59b25", + "metadata": {}, + "source": [ + "## Distil: Text to Image\n", + "\n", + "Distil text-to-image generation using the same structured JSON prompt.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3eade39c", + "metadata": {}, + "outputs": [], + "source": [ + "t2i_distil_payload, t2i_distil_output, t2i_distil_model = create_payload(\"t2i_distil\", backend=\"sglang\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "a3367e63", + "metadata": {}, + "source": [ + "### Run\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5908dc43", + "metadata": {}, + "outputs": [], + "source": [ + "run_sglang_payload(t2i_distil_payload, t2i_distil_output, model=\"Cosmos3-Distil\")\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "2b23b582", + "metadata": {}, + "source": [ + "### View Results\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e67db0de", + "metadata": {}, + "outputs": [], + "source": [ + "view_run(t2i_distil_output)\n", + " " + ] + }, + { + "cell_type": "markdown", + "id": "35481725", + "metadata": {}, + "source": [ + "## Distil: Image to Video Without Audio\n", + "\n", + "Distil image-to-video generation using its paired image asset, with audio disabled.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d737217e", + "metadata": {}, + "outputs": [], + "source": [ + "i2v_distil_noaudio_payload, i2v_distil_noaudio_output, i2v_distil_noaudio_model = create_payload(\"i2v_distil_noaudio\", backend=\"sglang\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "681f6f85", + "metadata": {}, + "source": [ + "### Run\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3df984a", + "metadata": {}, + "outputs": [], + "source": [ + "run_sglang_payload(i2v_distil_noaudio_payload, i2v_distil_noaudio_output, model=\"Cosmos3-Distil\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "8ae59282", + "metadata": {}, + "source": [ + "### View Results\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06d4095c", + "metadata": {}, + "outputs": [], + "source": [ + "view_run(i2v_distil_noaudio_output)\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64c8ffe5", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3", "language": "python", "name": "python3" },