From c203e06abbd3c0176645445d7288ca2dd1a0487c Mon Sep 17 00:00:00 2001 From: rwagwani Date: Thu, 29 Oct 2026 12:54:33 +0000 Subject: [PATCH 1/2] fix(inference): reject out-of-range seeds before generation `seed` accepted any integer, but the samplers cannot take all of them: the bidirectional sampler seeds NumPy's RandomState (0..2**32-1) and torch.Generator overflows at 2**64. A request with seed=2**64 therefore loaded the model and then failed with "Overflow when unpacking long long" when the first noise was drawn. Constrain `seed` to 0..2**32-1, the range every sampler accepts. An out-of-range seed is rejected when the request is parsed, before the model loads. num_outputs > 1 uses seed + i per output, so the built sample also checks the last per-output seed (this check runs when the sample is built, still before generation). Chunked bidirectional transfer and chunked multiview add the chunk index to the seed; those sums now wrap within the range, which changes nothing for any seed that runs today. A missing seed still falls back to the per-sample derived seed, which is within 31 bits. docs/inference.md states the range. Behaviour change: negative seeds currently run on the causal (Sim-Transfer) path, because torch.Generator accepts them. They are now rejected on every model, so one range holds for all of them. No seed in the repository's specs, docs, tests or notebooks is outside the range. Signed-off-by: rwagwani Co-Authored-By: Claude Opus 5.5 --- cosmos_framework/inference/args_test.py | 43 +++++++++++++++++++++++ cosmos_framework/inference/common/args.py | 22 ++++++++++-- cosmos_framework/inference/multiview.py | 5 +-- cosmos_framework/inference/transfer.py | 4 ++- docs/inference.md | 2 +- 5 files changed, 69 insertions(+), 7 deletions(-) diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index 10fa68404..318f98d25 100644 --- a/cosmos_framework/inference/args_test.py +++ b/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) diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 4b49a25c2..f3931b1fe 100644 --- a/cosmos_framework/inference/common/args.py +++ b/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.""" diff --git a/cosmos_framework/inference/multiview.py b/cosmos_framework/inference/multiview.py index c5db6a92f..5fd7ab06d 100644 --- a/cosmos_framework/inference/multiview.py +++ b/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( diff --git a/cosmos_framework/inference/transfer.py b/cosmos_framework/inference/transfer.py index 63c615469..6f79a73ef 100644 --- a/cosmos_framework/inference/transfer.py +++ b/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, diff --git a/docs/inference.md b/docs/inference.md index 72d04c291..331fac305 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -227,7 +227,7 @@ For debugging, the full set of sample arguments is saved to `/ Date: Fri, 6 Nov 2026 08:21:07 +0000 Subject: [PATCH 2/2] docs(inference): document per-output seeds and the chunk-seed wrap Co-Authored-By: Claude Opus 5.5 --- docs/inference.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/inference.md b/docs/inference.md index 331fac305..62835faf0 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -227,7 +227,7 @@ For debugging, the full set of sample arguments is saved to `/ 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.