diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py index 2d9c26edc531..3d5862716666 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): @@ -269,6 +270,30 @@ class Cosmos3VideoResponse(BaseModel): _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 +# 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: """Run one generation collectively on all ranks. @@ -303,6 +328,8 @@ def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: 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") @@ -364,6 +391,13 @@ def _generate_blocking(request: Cosmos3VideoRequest) -> Cosmos3VideoResponse: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio + 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. if request.model_mode is not None: payload["model_mode"] = request.model_mode @@ -544,6 +578,13 @@ def create_video(req: Cosmos3VideoRequest) -> dict: } if aspect_ratio is not None: payload["aspect_ratio"] = aspect_ratio + 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. if req.model_mode is not None: payload["model_mode"] = req.model_mode