Cosmos3-edge PR 6: Cosmos3-Edge as a cosmos3 config variant, with streaming rollout and action loop

README.md

@@ -121,6 +121,7 @@ _Note_: The **first request(s) on a fresh environment can be slow** — often te
 | [Orpheus](https://huggingface.co/canopylabs/orpheus-3b-0.1-ft) | Speech LM | text → speech | `/v1/audio/speech` |
 | [Cosmos3 Nano / Super](https://huggingface.co/nvidia/Cosmos3-Nano) | World model | text, image, video → image, video (+ sound), robot actions | `/v1/images/generations`, `/v1/videos/generations` |
 | [Cosmos3 Policy DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID) | Robot policy | text, image, video → robot actions, video | `/generate`, `/v1/images/generations`, `/v1/videos/generations` |
+| [Cosmos3-Edge](https://huggingface.co/nvidia/Cosmos3-Edge) | 480p world model + VLM | text, image, video → image, video (streamed windows), robot actions, text | `/v1/chat/completions`, `/v1/images/generations`, `/v1/videos/generations`, `/generate`, `/generate/ws` |
 | [Pi0.5](https://huggingface.co/lerobot/pi05_base) | Vision-language-action | text, image, state → robot actions | `/generate` |
 | [V-JEPA 2 / 2-AC](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256) | World model | video (+ actions) → latents, rollouts | `/generate` |
 | [Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers) | Video diffusion | text, image → video | `/v1/videos/generations`, `/generate` |

benchmark/cosmos3/bench_action_baselines.py

@@ -0,0 +1,204 @@
+"""Action-policy baseline clients. vLLM-Omni: POST /v1/videos (multipart, extra_params.action_mode=policy) -> poll
+GET /v1/videos/{id} -> top-level `action` {data, shape, dtype, raw_action_dim}. Saves the actions as .npy (for
+notes/served_action_parity.py --ref) and reports per-call latency and actions/s over N rounds.
+SGLang: GET /v1/actions/metadata then one best-effort POST /v1/actions/generations (Cosmos3 may not enable it).
+
+  python benchmark/cosmos3/bench_action_baselines.py vllm-omni --port 8200 --image <jpg> --rounds 5 \
+      --out results/<date>/action_vllm_omni.npy
+  python benchmark/cosmos3/bench_action_baselines.py sglang --port 8400 --image <jpg>
+"""
+
+import argparse
+import base64
+import json
+import statistics
+import sys
+import time
+import urllib.request
+
+import numpy as np
+
+ap = argparse.ArgumentParser()
+ap.add_argument("engine", choices=["vllm-omni", "sglang"])
+ap.add_argument("--port", type=int, required=True)
+ap.add_argument("--model", default="nvidia/Cosmos3-Edge")
+ap.add_argument("--image", required=True)
+ap.add_argument("--prompt", default="Pick up the red cup and place it in the sink.")
+ap.add_argument("--domain", default="droid_lerobot")
+ap.add_argument("--action-dim", type=int, default=10)
+ap.add_argument("--chunk", type=int, default=32)
+ap.add_argument("--size", default="832x480")
+ap.add_argument("--steps", type=int, default=4)
+ap.add_argument("--gs", type=float, default=3.0)
+ap.add_argument(
+    "--flow-shift",
+    type=float,
+    default=5.0,
+    help="the DROID policy recipe (vLLM-Omni ROBOLAB defaults: 4 steps, gs 3.0, shift 5.0)",
+)
+ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--rounds", type=int, default=5)
+ap.add_argument("--warmup", type=int, default=1)
+ap.add_argument("--poll", type=float, default=0.05)
+ap.add_argument("--out", default="")
+a = ap.parse_args()
+BASE = f"http://127.0.0.1:{a.port}"
+
+
+def data_url(path):
+    ext = path.rsplit(".", 1)[-1].lower().replace("jpg", "jpeg")
+    return f"data:image/{ext};base64," + base64.b64encode(open(path, "rb").read()).decode()
+
+
+def _multipart(fields, files=()):
+    """files: (field, filename, bytes, content_type) tuples."""
+    boundary = "----act" + str(int(time.time() * 1e6))
+    body = b""
+    for k, v in fields.items():
+        body += f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode()
+    for k, fn, data, ctype in files:
+        body += (
+            f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"; filename="{fn}"\r\n'
+            f"Content-Type: {ctype}\r\n\r\n".encode()
+            + data
+            + b"\r\n"
+        )
+    return body + f"--{boundary}--\r\n".encode(), f"multipart/form-data; boundary={boundary}"
+
+
+def _body_of(e):
+    try:
+        return e.read().decode(errors="replace")[:600]
+    except Exception:
+        return ""
+
+
+def _get(path, timeout=600):
+    with urllib.request.urlopen(BASE + path, timeout=timeout) as r:
+        return json.load(r)
+
+
+def vllm_omni_once():
+    extra = {
+        "action_mode": "policy",
+        "domain_name": a.domain,
+        "raw_action_dim": a.action_dim,
+        "action_chunk_size": a.chunk,
+    }
+    fields = {
+        "model": a.model,
+        "prompt": a.prompt,
+        "size": a.size,
+        "num_frames": str(a.chunk + 1),
+        "fps": "15",
+        "num_inference_steps": str(a.steps),
+        "guidance_scale": str(a.gs),
+        "flow_shift": str(a.flow_shift),
+        "seed": str(a.seed),
+        "extra_params": json.dumps(extra),
+    }
+    # The recipe ships the conditioning frame as the input_reference upload (video_bench.py does the same for i2v);
+    # a data-URL image_reference is rejected with 400 "did not decode to an image".
+    body, ctype = _multipart(
+        fields, files=[("input_reference", a.image.rsplit("/", 1)[-1], open(a.image, "rb").read(), "image/jpeg")]
+    )
+    t0 = time.perf_counter()
+    req = urllib.request.Request(BASE + "/v1/videos", data=body, headers={"Content-Type": ctype})
+    with urllib.request.urlopen(req, timeout=1800) as r:
+        job = json.load(r)
+    while job.get("status") not in ("completed", "failed"):
+        time.sleep(a.poll)
+        job = _get(f"/v1/videos/{job['id']}")
+    wall = time.perf_counter() - t0
+    if job.get("status") != "completed":
+        raise RuntimeError(f"vllm-omni job failed: {job.get('error')}")
+    act = job.get("action")
+    if not act:
+        raise RuntimeError(f"no action in job response; keys={list(job.keys())}")
+    arr = np.array(act["data"], dtype=np.float32).reshape(act["shape"])
+    while arr.ndim > 2:
+        arr = arr[0]
+    return arr, wall, act
+
+
+def sglang_once():
+    meta = _get("/v1/actions/metadata")
+    print(
+        json.dumps({"sglang_action_metadata": {k: meta.get(k) for k in ("policy_family", "input", "output")}}),
+        flush=True,
+    )
+    keys = (meta.get("input") or {}).get("image_keys") or ["image"]
+    # sglang.multimodal_gen.runtime.entrypoints.action.protocol: the JSON body is
+    # {"input": {"task", "observation"}, "parameters"};
+    # image values are {"b64_json": ...} dicts (data URLs are not decoded).
+    # Cosmos3 policy parameters ride in "parameters".
+    b64 = base64.b64encode(open(a.image, "rb").read()).decode()
+    payload = {
+        "input": {"task": a.prompt, "observation": {"images": {keys[0]: {"b64_json": b64}}}},
+        "parameters": {
+            "seed": a.seed,
+            "num_inference_steps": a.steps,
+            "guidance_scale": a.gs,
+            "flow_shift": a.flow_shift,
+            "action_mode": "policy",
+            "domain_name": a.domain,
+            "raw_action_dim": a.action_dim,
+            "action_chunk_size": a.chunk,
+            "num_frames": a.chunk + 1,
+            "width": int(a.size.split("x")[0]),
+            "height": int(a.size.split("x")[1]),
+        },
+    }
+    t0 = time.perf_counter()
+    req = urllib.request.Request(
+        BASE + "/v1/actions/generations",
+        data=json.dumps(payload).encode(),
+        headers={"Content-Type": "application/json"},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=1800) as r:
+            out = json.load(r)
+    except urllib.error.HTTPError as e:
+        raise RuntimeError(f"sglang {e.code}: {_body_of(e)}") from None
+    wall = time.perf_counter() - t0
+    # action.protocol.action_generation_response:
+    # {"data": [{"action": {"shape": [H, D], "values": [[...]], "raw_action_dim"}}], "usage"}
+    act = out["data"][0]["action"]
+    arr = np.array(act["values"], dtype=np.float32).reshape(act["shape"])
+    raw = int(act.get("raw_action_dim") or a.action_dim)
+    return arr[:, :raw], wall, {"shape": act["shape"], "raw_action_dim": raw, "usage": out.get("usage")}
+
+
+once = vllm_omni_once if a.engine == "vllm-omni" else sglang_once
+print(
+    f"=== {a.engine} action policy port={a.port} domain={a.domain} chunk={a.chunk} "
+    f"steps={a.steps} gs={a.gs} seed={a.seed} ===",
+    flush=True,
+)
+for i in range(a.warmup):
+    try:
+        _, w, _ = once()
+        print(f"  warmup {i}: {w:.3f}s", flush=True)
+    except Exception as exc:  # noqa: BLE001
+        print(f"  warmup {i} failed: {exc!r}", flush=True)
+        sys.exit(2)
+walls, arrs = [], []
+for i in range(a.rounds):
+    arr, w, act = once()
+    walls.append(w)
+    arrs.append(arr)
+    print(
+        f"  round {i}: {w:.3f}s  actions {arr.shape} first={np.round(arr[0, : min(3, arr.shape[1])], 4).tolist()} "
+        f"dtype={act.get('dtype')} raw_dim={act.get('raw_action_dim')}",
+        flush=True,
+    )
+med = statistics.median(walls)
+rep = float(max(np.abs(x - arrs[0]).max() for x in arrs[1:])) if len(arrs) > 1 else 0.0
+print(
+    f"  chunk latency median {med:.3f}s  p95 {sorted(walls)[int(0.95 * (len(walls) - 1))]:.3f}s"
+    f"  -> {a.chunk / med:.1f} actions/s (sequential calls); repeat max-abs-diff {rep:.3e}",
+    flush=True,
+)
+if a.out:
+    np.save(a.out, arrs[0])
+    print("  saved", a.out)

benchmark/cosmos3/bench_chat_oai.py

@@ -0,0 +1,116 @@
+"""Reasoner (understanding tower) latency client for the OpenAI chat endpoint
+both M* (``mstar serve cosmos3_edge``) and vLLM (``vllm serve nvidia/Cosmos3-Edge``)
+expose: time-to-first-token and decode tokens/s, streamed, at a chosen
+concurrency, on the model card's reasoning prompt (image + text) or text only.
+
+Same payload on both engines (greedy, fixed max_tokens, thinking off unless
+asked), client-side timing, warmup excluded, median and p95 reported.
+
+  python bench_chat_oai.py --port 8000 --model nvidia/Cosmos3-Edge --tag vllm --image assets/example_reasoning_input.png
+  python bench_chat_oai.py --port 8100 --model cosmos3_edge --tag ours --image assets/example_reasoning_input.png
+"""
+import argparse
+import base64
+import concurrent.futures as cf
+import json
+import mimetypes
+import statistics
+import time
+import urllib.request
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--port", type=int, required=True)
+ap.add_argument("--model", default="nvidia/Cosmos3-Edge")
+ap.add_argument("--image", default="")  # optional conditioning image path (else text-only)
+ap.add_argument("--prompt", default="The task is to put flower into the red bottle. Generate a plan consisting of "
+                                    "subtasks for accomplish the task.")
+ap.add_argument("--max-tokens", type=int, default=128)
+ap.add_argument("--concurrency", default="1,8,32")
+ap.add_argument("--requests", type=int, default=16)  # per concurrency level (>= concurrency)
+ap.add_argument("--warmup", type=int, default=2)
+ap.add_argument("--thinking", action="store_true")
+ap.add_argument("--tag", default="run")
+ap.add_argument("--out", default="")  # optional JSON results path
+args = ap.parse_args()
+
+URL = f"http://localhost:{args.port}/v1/chat/completions"
+content = []
+if args.image:
+    mime = mimetypes.guess_type(args.image)[0] or "image/png"
+    with open(args.image, "rb") as f:
+        data_url = f"data:{mime};base64," + base64.b64encode(f.read()).decode()
+    content.append({"type": "image_url", "image_url": {"url": data_url}})
+content.append({"type": "text", "text": args.prompt})
+BODY = {
+    "model": args.model,
+    "messages": [{"role": "user", "content": content}],
+    "max_tokens": args.max_tokens,
+    "temperature": 0.0,
+    "stream": True,
+    "chat_template_kwargs": {"enable_thinking": bool(args.thinking)},
+}
+
+
+def one() -> dict:
+    req = urllib.request.Request(URL, data=json.dumps(BODY).encode(), headers={"Content-Type": "application/json"})
+    t0 = time.perf_counter()
+    first = None
+    n_chunks = 0
+    text = []
+    with urllib.request.urlopen(req, timeout=600) as r:
+        for line in r:
+            line = line.decode("utf-8", "replace").strip()
+            if not line.startswith("data:"):
+                continue
+            payload = line[5:].strip()
+            if payload == "[DONE]":
+                break
+            delta = json.loads(payload)["choices"][0].get("delta", {})
+            piece = delta.get("content")
+            if piece:
+                if first is None:
+                    first = time.perf_counter()
+                n_chunks += 1
+                text.append(piece)
+    end = time.perf_counter()
+    ttft = (first or end) - t0
+    total = end - t0
+    decode = max(end - (first or end), 1e-9)
+    return {"ttft": ttft, "total": total, "chunks": n_chunks, "tok_s": (n_chunks - 1) / decode if n_chunks > 1 else 0.0,
+            "text": "".join(text)}
+
+
+def pct(xs, p):
+    xs = sorted(xs)
+    return xs[min(len(xs) - 1, int(round(p * (len(xs) - 1))))]
+
+
+results = {}
+print(f"=== {args.tag} port={args.port} model={args.model} max_tokens={args.max_tokens} "
+      f"image={'yes' if args.image else 'no'} thinking={args.thinking} ===", flush=True)
+for _ in range(args.warmup):
+    one()
+for conc in [int(c) for c in args.concurrency.split(",")]:
+    n = max(args.requests, conc)
+    t0 = time.perf_counter()
+    with cf.ThreadPoolExecutor(max_workers=conc) as ex:
+        outs = list(ex.map(lambda _: one(), range(n)))
+    wall = time.perf_counter() - t0
+    ttfts = [o["ttft"] for o in outs]
+    toks = [o["tok_s"] for o in outs]
+    total_tokens = sum(o["chunks"] for o in outs)
+    rec = {
+        "concurrency": conc, "requests": n, "ttft_p50": statistics.median(ttfts), "ttft_p95": pct(ttfts, 0.95),
+        "decode_tok_s_per_req_p50": statistics.median(toks), "aggregate_tok_s": total_tokens / wall,
+        "wall_s": wall, "sample": outs[0]["text"][:120],
+    }
+    results[conc] = rec
+    print(f"  conc={conc:3d}  TTFT p50 {rec['ttft_p50'] * 1000:7.1f} ms  p95 {rec['ttft_p95'] * 1000:7.1f} ms  "
+          f"decode {rec['decode_tok_s_per_req_p50']:6.1f} tok/s/req  aggregate {rec['aggregate_tok_s']:7.1f} tok/s  "
+          f"(n={n}, wall {wall:.1f}s)", flush=True)
+print("  sample:", repr(results[min(results)]["sample"]))
+if args.out:
+    with open(args.out, "w") as f:
+        json.dump({"tag": args.tag, "model": args.model, "max_tokens": args.max_tokens, "image": bool(args.image),
+                   "results": results}, f, indent=2)
+print("DONE", flush=True)