From 65a50f2dafcb9b1499d46c02c0c5e56dbd53d7f5 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Wed, 15 Jul 2026 19:08:55 +0000 Subject: [PATCH 1/3] feat: config-driven sampling step count for cosmos3 worker Add COSMOS3_NUM_STEPS env (set via model-config extra_env) as the default denoising-step count, plus a per-request nvext.num_steps override. Lets us tune step count without an image rebuild. Falls back to the framework's per-modality default (35) when unset. Motivation: fewer steps is the one in-house, no-post-training cost lever for Cosmos3 video (sampling loop is ~94% of latency). Step-sweep on a reserved B200 (720p/5s Nano) confirmed clean latency scaling: 35=64s, 28=53s, 20=40s, 16=34s, 12=27s, 8=21s. 28 matches fal's default and looked visually identical to 35. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/cosmos3/worker.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py index 2d9c26edc531..57268cc09c40 100644 --- a/examples/cosmos3/worker.py +++ b/examples/cosmos3/worker.py @@ -177,6 +177,7 @@ class NvExtFields(BaseModel): fps: int | None = None num_frames: int | None = None seed: int | None = None + num_steps: int | None = None class Cosmos3VideoRequest(BaseModel): @@ -267,6 +268,12 @@ class Cosmos3VideoResponse(BaseModel): _served_model_name: str = "" _output_root: Path = Path("/tmp/cosmos3_worker_outputs") _generate_lock = threading.Lock() +# Default denoising steps when a request doesn't specify one. Driven by the +# COSMOS3_NUM_STEPS env (set via model-config extra_env) so step count is tunable +# without an image rebuild; None falls back to the framework's per-modality default. +_default_num_steps: int | None = ( + int(os.environ["COSMOS3_NUM_STEPS"]) if os.environ.get("COSMOS3_NUM_STEPS") else None +) def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: @@ -300,7 +307,7 @@ def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: override_kwargs["model_mode"] = model_mode if payload.get("action_url") is not None: override_kwargs["action_path"] = payload["action_url"] - for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size"): + for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size", "num_steps"): if payload.get(key) is not None: override_kwargs[key] = payload[key] @@ -364,6 +371,9 @@ def _generate_blocking(request: Cosmos3VideoRequest) -> Cosmos3VideoResponse: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio + ns = request.nvext.num_steps if (request.nvext is not None and request.nvext.num_steps is not None) else _default_num_steps + if ns is not None: + payload["num_steps"] = ns # Forward action fields when present. if request.model_mode is not None: payload["model_mode"] = request.model_mode @@ -544,6 +554,9 @@ def create_video(req: Cosmos3VideoRequest) -> dict: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio + ns = req.nvext.num_steps if (req.nvext is not None and req.nvext.num_steps is not None) else _default_num_steps + if ns is not None: + payload["num_steps"] = ns # Forward action fields when present. if req.model_mode is not None: payload["model_mode"] = req.model_mode From 3644ac21d281164a18c903488df71caa67ca283e Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Wed, 15 Jul 2026 21:44:22 +0000 Subject: [PATCH 2/3] fix: scope COSMOS3_NUM_STEPS to video, parse env defensively Address adversarial review of the num_steps change: - Apply the default step count to VIDEO only (num_frames > 1). text2image defaults to 50 steps; the previous version silently forced it to the video value, an unvalidated quality cut on a live t2i path. Image modes now keep their framework default; explicit nvext.num_steps still overrides either. - Parse COSMOS3_NUM_STEPS defensively (non-int / <=0 -> warn + ignore) instead of int() at import, which would crash the worker fleet-wide on a bad value. - Keep all lines within the repo's 88-char limit. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/cosmos3/worker.py | 44 ++++++++++++++++++++++++++++++-------- 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py index 57268cc09c40..c946e8ec800d 100644 --- a/examples/cosmos3/worker.py +++ b/examples/cosmos3/worker.py @@ -268,12 +268,28 @@ class Cosmos3VideoResponse(BaseModel): _served_model_name: str = "" _output_root: Path = Path("/tmp/cosmos3_worker_outputs") _generate_lock = threading.Lock() -# Default denoising steps when a request doesn't specify one. Driven by the -# COSMOS3_NUM_STEPS env (set via model-config extra_env) so step count is tunable -# without an image rebuild; None falls back to the framework's per-modality default. -_default_num_steps: int | None = ( - int(os.environ["COSMOS3_NUM_STEPS"]) if os.environ.get("COSMOS3_NUM_STEPS") else None -) +# Default denoising steps for VIDEO when a request doesn't specify one. Driven by +# the COSMOS3_NUM_STEPS env (set via model-config extra_env) so step count is +# tunable without an image rebuild. Parsed defensively: a bad value degrades to the +# framework default instead of crashing the worker at import. Applied to video only +# (see the payload builders) — image modes keep their own framework default (e.g. +# text2image defaults to 50, unvalidated at a lower count). +def _parse_default_num_steps() -> int | None: + raw = os.environ.get("COSMOS3_NUM_STEPS") + if not raw: + return None + try: + n = int(raw) + except ValueError: + log.warning("ignoring non-integer COSMOS3_NUM_STEPS=%r", raw) + return None + if n <= 0: + log.warning("ignoring non-positive COSMOS3_NUM_STEPS=%r", raw) + return None + return n + + +_default_num_steps: int | None = _parse_default_num_steps() def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: @@ -307,9 +323,11 @@ def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: override_kwargs["model_mode"] = model_mode if payload.get("action_url") is not None: override_kwargs["action_path"] = payload["action_url"] - for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size", "num_steps"): + for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size"): if payload.get(key) is not None: override_kwargs[key] = payload[key] + if payload.get("num_steps") is not None: + override_kwargs["num_steps"] = payload["num_steps"] overrides = OmniSampleOverrides(**override_kwargs) overrides.download(rank_dir / "inputs") @@ -371,7 +389,11 @@ def _generate_blocking(request: Cosmos3VideoRequest) -> Cosmos3VideoResponse: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio - ns = request.nvext.num_steps if (request.nvext is not None and request.nvext.num_steps is not None) else _default_num_steps + ns = None + if request.nvext is not None and request.nvext.num_steps is not None: + ns = request.nvext.num_steps + elif num_frames > 1: # video only; image modes keep their framework default + ns = _default_num_steps if ns is not None: payload["num_steps"] = ns # Forward action fields when present. @@ -554,7 +576,11 @@ def create_video(req: Cosmos3VideoRequest) -> dict: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio - ns = req.nvext.num_steps if (req.nvext is not None and req.nvext.num_steps is not None) else _default_num_steps + ns = None + if req.nvext is not None and req.nvext.num_steps is not None: + ns = req.nvext.num_steps + elif num_frames > 1: # video only; image modes keep their framework default + ns = _default_num_steps if ns is not None: payload["num_steps"] = ns # Forward action fields when present. From 667b827c4175de93047564df34c014d1b14073be Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Wed, 15 Jul 2026 23:18:22 +0000 Subject: [PATCH 3/3] style: black formatting (2 blank lines before top-level fn) Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/cosmos3/worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py index c946e8ec800d..3d5862716666 100644 --- a/examples/cosmos3/worker.py +++ b/examples/cosmos3/worker.py @@ -268,6 +268,8 @@ class Cosmos3VideoResponse(BaseModel): _served_model_name: str = "" _output_root: Path = Path("/tmp/cosmos3_worker_outputs") _generate_lock = threading.Lock() + + # Default denoising steps for VIDEO when a request doesn't specify one. Driven by # the COSMOS3_NUM_STEPS env (set via model-config extra_env) so step count is # tunable without an image rebuild. Parsed defensively: a bad value degrades to the