feat: config-driven sampling step count for cosmos3 worker

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