fix(inference): reject out-of-range seeds before generation

cosmos_framework/inference/args_test.py

@@ -913,3 +913,46 @@ def test_sample_output_dir_rejects_a_symlink_out_of_the_root(tmp_path: Path) ->
 
     with pytest.raises(ValueError, match="outside the requested output root"):
         sample_output_dir(root, "link")
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+@pytest.mark.parametrize("seed", [0, 2025, 2**32 - 1])
+def test_seed_accepts_the_range_every_sampler_takes(seed: int) -> None:
+    assert OmniSampleOverrides(seed=seed).seed == seed
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+@pytest.mark.parametrize("seed", [-1, 2**32, 2**64])
+def test_seed_rejects_values_a_sampler_cannot_take(seed: int) -> None:
+    # NumPy's RandomState (bidirectional sampler) takes 0..2**32-1, and torch overflows at 2**64,
+    # so these used to fail only after the model had loaded.
+    with pytest.raises(pydantic.ValidationError):
+        OmniSampleOverrides(seed=seed)
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_seed_plus_extra_outputs_must_stay_in_range(tmp_path: Path) -> None:
+    from cosmos_framework.inference.common.args import MAX_SEED, SampleArgs
+
+    fields = dict(output_dir=tmp_path, model="", extra={}, name="s", tensors_file=None, pickle_file=None)
+    assert SampleArgs(**fields, num_outputs=2, seed=MAX_SEED - 1).seed == MAX_SEED - 1
+    with pytest.raises(pydantic.ValidationError, match="per-output seeds up to"):
+        SampleArgs(**fields, num_outputs=2, seed=MAX_SEED)
+
+
+@pytest.mark.L0
+@pytest.mark.CPU
+def test_built_sample_rejects_extra_outputs_past_the_seed_range(tmp_path: Path) -> None:
+    from cosmos_framework.inference.common.args import MAX_SEED
+
+    setup_args = OmniSetupOverrides(
+        checkpoint_path=DEFAULT_CHECKPOINT_NAME, output_dir=tmp_path / "outputs"
+    ).build_setup()
+    model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig)
+    overrides = OmniSampleOverrides(name="s", seed=MAX_SEED, num_outputs=2)
+    overrides.output_dir = tmp_path / "inputs"
+    with pytest.raises(pydantic.ValidationError, match="per-output seeds up to"):
+        overrides.build_sample(model_config=model_dict.config)

cosmos_framework/inference/common/args.py

@@ -1089,6 +1089,12 @@ def build_setup(self) -> SetupArgs:
         """Build setup arguments."""
 
 
+# The bidirectional sampler seeds NumPy's RandomState with the sample seed, which accepts only
+# 0..2**32-1; out-of-range seeds used to fail only after the model had loaded.
+MAX_SEED = 2**32 - 1
+Seed = Annotated[int, pydantic.Field(ge=0, le=MAX_SEED)]
+
+
 class SampleArgs(ArgsBase):
     """Inference sample arguments."""
 
@@ -1098,10 +1104,20 @@ class SampleArgs(ArgsBase):
 
     name: str
     num_outputs: pydantic.PositiveInt
-    seed: int | None
+    seed: Seed | None
     tensors_file: ResolvedFilePath | None
     pickle_file: ResolvedFilePath | None
 
+    @pydantic.model_validator(mode="after")
+    def _check_expanded_seeds(self) -> Self:
+        # Each extra output uses seed + i, so the last one must stay in range too.
+        if self.seed is not None and self.seed + self.num_outputs - 1 > MAX_SEED:
+            raise ValueError(
+                f"seed={self.seed} with num_outputs={self.num_outputs} gives per-output seeds up to "
+                f"{self.seed + self.num_outputs - 1}, above the maximum {MAX_SEED}."
+            )
+        return self
+
     def get_data(self, *, device: str | int = "cpu") -> dict[str, Any]:
         import pickle
 
@@ -1130,8 +1146,8 @@ class SampleOverrides(OverridesBase):
     """Name of the sample."""
     num_outputs: Training[Annotated[pydantic.PositiveInt | None, tyro.conf.arg(aliases=("-n",))]] = None
     """Number of outputs to generate per sample."""
-    seed: int | None = None
-    """Seed for the random number generator."""
+    seed: Seed | None = None
+    """Seed for the random number generator (0 to 2**32-1)."""
 
     tensors_file: ResolvedFilePath | None = None
     """Path to data tensors file."""

cosmos_framework/inference/multiview.py

@@ -38,7 +38,7 @@
 )
 from cosmos_framework.data.generator.sequence_packing import SequencePlan
 from cosmos_framework.inference.args import ConditionVisionMode, OmniSampleArgs
-from cosmos_framework.inference.common.args import IMAGE_EXTENSIONS, VIDEO_EXTENSIONS
+from cosmos_framework.inference.common.args import IMAGE_EXTENSIONS, MAX_SEED, VIDEO_EXTENSIONS
 from cosmos_framework.inference.common.multiview_captions import exported_caption_task
 from cosmos_framework.inference.vision import load_conditioning_image_pixels, load_conditioning_video
 from cosmos_framework.model.generator.mot.replica_partition import ReplicaPartitioner, WorkUnit
@@ -624,7 +624,8 @@ def generate_multiview_caption_chunks(
                 "prompt": chunk.prompt,
                 "num_frames": chunk_frame_count,
                 "condition_frame_indexes_vision": condition_frame_indexes,
-                "seed": base_seed + chunk.index,
+                # Wrap so a seed near the top of the accepted range stays valid for later chunks.
+                "seed": (base_seed + chunk.index) % (MAX_SEED + 1),
             },
         )
         chunk_data_batch = get_multiview_sample_data(

cosmos_framework/inference/transfer.py

@@ -23,6 +23,7 @@
     TransferArgs,
     TransferHintKey,
 )
+from cosmos_framework.inference.common.args import MAX_SEED
 from cosmos_framework.inference.inference import _fallback_seed, _get_prompt_sample_data
 from cosmos_framework.inference.transfer_preprocessing import (
     TransferGeometry,
@@ -1036,7 +1037,8 @@ def generate_transfer_sample(
                 control_guidance=sample_args.control_guidance,
                 control_guidance_interval=sample_args.control_guidance_interval,
             ),
-            seed=[seed + chunk_id],
+            # Wrap so a seed near the top of the accepted range stays valid for later chunks.
+            seed=[(seed + chunk_id) % (MAX_SEED + 1)],
             n_sample=1,
             has_negative_prompt=negative_prompt is not None,
             num_steps=sample_args.num_steps,

docs/inference.md

@@ -227,7 +227,7 @@ For debugging, the full set of sample arguments is saved to `<output_dir>/<sampl
 Common arguments:
 
 - `model_mode`: Generation modality. See [Modes](#modes) above for all options.
-- `seed`: Random seed for reproducibility.
+- `seed`: Random seed for reproducibility, from 0 to 2**32 - 1. Without a seed, one is derived from the sample name and output directory. With `num_outputs` > 1, output `i` uses `seed + i`, so `seed + num_outputs - 1` must also be at most 2**32 - 1. Multiview and Transfer chunks use `seed + chunk_index`, wrapped modulo 2**32, so a seed near the top of the range stays valid for later chunks; a wrapped chunk seed can equal the seed of another sample.
 
 **Note:** Condition file paths are relative to the input file.