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..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.