From 8105762b14cead6ab643f686ab7405d9758af7c1 Mon Sep 17 00:00:00 2001 From: Liang Feng Date: Fri, 9 Oct 2026 09:51:56 +0000 Subject: [PATCH] Release 2026-10-09 from i4 cda81965fb3fa14d7d898c5d19555f39142db344 _source_commit: cda81965fb3fa14d7d898c5d19555f39142db344-dirty _dest_commit: 7c5d20cf330db99f35327166ebe0408318f86a04 --- .file_mapping.json | 26 +- .../callbacks/data_stats_writer.py | 157 ++-- .../callbacks/dataloader_state.py | 118 ++- .../callbacks/sampled_media_recorder.py | 5 +- .../callbacks/sampled_media_recorder_test.py | 21 +- .../callbacks/wall_clock_checkpoint.py | 50 +- .../configs/base/defaults/replay_attention.py | 4 +- .../augmentors/reasoner/prompt_format.py | 6 +- .../generator/lance/lance_dataset_open.py | 6 +- .../generator/multiview/camera_geometry.py | 772 ++++++++++++++++++ .../multiview/camera_geometry_contract.py | 162 ++++ .../multiview/camera_geometry_native.py | 183 +++++ .../generator/reasoner/video_decoder_qwen.py | 5 +- .../sequence_packing/autoregressive.py | 130 ++- .../model/generator/mot/causal_attention.py | 178 +++- .../mot/causal_cosmos3_vfm_network.py | 16 + .../mot/flex_attention_memory_test.py | 174 ++++ .../model/generator/mot/maskless_attention.py | 3 +- .../mot/multiview_action_attention.py | 102 +++ .../model/generator/mot/rigrope.py | 171 ++++ .../model/generator/mot/rigrope_test.py | 146 ++++ .../model/generator/omni_mot_causal_model.py | 655 +++++---------- .../generator/omni_mot_causal_model_test.py | 135 ++- .../generator/tokenizers/lidar/geometry.py | 238 ++++++ .../model/generator/utils/camera_rays.py | 214 +++++ .../model/generator/utils/camera_rays_test.py | 130 +++ .../generator/utils/camera_relative_pose.py | 420 ++++++++++ .../utils/camera_relative_pose_test.py | 165 ++++ .../model/generator/utils/ftheta.py | 114 +++ .../model/generator/utils/multiview_ar.py | 706 +++++++++++++++- .../utils/prope_cross_view_poses_test.py | 146 ++++ .../model/generator/utils/rigrope_geometry.py | 406 +++++++++ .../generator/utils/rigrope_geometry_test.py | 448 ++++++++++ .../utils/rolling_kv/rolling_transfer_ar.py | 18 +- cosmos_framework/utils/generator/multiview.py | 14 + .../utils/generator/multiview_test.py | 19 + .../utils/generator/video_preprocess.py | 73 ++ 37 files changed, 5633 insertions(+), 703 deletions(-) create mode 100644 cosmos_framework/data/generator/multiview/camera_geometry.py create mode 100644 cosmos_framework/data/generator/multiview/camera_geometry_contract.py create mode 100644 cosmos_framework/data/generator/multiview/camera_geometry_native.py create mode 100644 cosmos_framework/model/generator/mot/flex_attention_memory_test.py create mode 100644 cosmos_framework/model/generator/mot/multiview_action_attention.py create mode 100644 cosmos_framework/model/generator/mot/rigrope.py create mode 100644 cosmos_framework/model/generator/mot/rigrope_test.py create mode 100644 cosmos_framework/model/generator/tokenizers/lidar/geometry.py create mode 100644 cosmos_framework/model/generator/utils/camera_rays.py create mode 100644 cosmos_framework/model/generator/utils/camera_rays_test.py create mode 100644 cosmos_framework/model/generator/utils/camera_relative_pose.py create mode 100644 cosmos_framework/model/generator/utils/camera_relative_pose_test.py create mode 100644 cosmos_framework/model/generator/utils/ftheta.py create mode 100644 cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py create mode 100644 cosmos_framework/model/generator/utils/rigrope_geometry.py create mode 100644 cosmos_framework/model/generator/utils/rigrope_geometry_test.py diff --git a/.file_mapping.json b/.file_mapping.json index 60d1258cc..7bd398ba5 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "c956695327e59d1b759868e8e0a36cca5733cdb3-dirty", - "_dest_commit": "8aca062dd4d93111d8afc4c6e7c7c592a81b81bc", - "_generated_at": "2026-10-08T08:05:55Z", + "_source_commit": "cda81965fb3fa14d7d898c5d19555f39142db344-dirty", + "_dest_commit": "7c5d20cf330db99f35327166ebe0408318f86a04", + "_generated_at": "2026-10-09T09:51:37Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -372,6 +372,9 @@ "projects/cosmos3/cosmos3/datasets/lance/conversation_training_export_manifest.py": "cosmos_framework/data/generator/lance/conversation_training_export_manifest.py", "projects/cosmos3/cosmos3/datasets/lance/lance_dataset_open.py": "cosmos_framework/data/generator/lance/lance_dataset_open.py", "projects/cosmos3/cosmos3/datasets/multiview/camera_attributes.py": "cosmos_framework/data/generator/multiview/camera_attributes.py", + "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry.py": "cosmos_framework/data/generator/multiview/camera_geometry.py", + "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry_contract.py": "cosmos_framework/data/generator/multiview/camera_geometry_contract.py", + "projects/cosmos3/cosmos3/datasets/multiview/camera_geometry_native.py": "cosmos_framework/data/generator/multiview/camera_geometry_native.py", "projects/cosmos3/cosmos3/datasets/multiview/caption_chunks.py": "cosmos_framework/data/generator/multiview/caption_chunks.py", "projects/cosmos3/cosmos3/datasets/multiview/caption_format.py": "cosmos_framework/data/generator/multiview/caption_format.py", "projects/cosmos3/cosmos3/datasets/packing_iterable_dataset.py": "cosmos_framework/data/generator/packing_iterable_dataset.py", @@ -414,6 +417,7 @@ "projects/cosmos3/cosmos3/models/mot/dot_product_attention.py": "cosmos_framework/model/generator/mot/dot_product_attention.py", "projects/cosmos3/cosmos3/models/mot/flex_attention.py": "cosmos_framework/model/generator/mot/flex_attention.py", "projects/cosmos3/cosmos3/models/mot/flex_attention_bench.py": "cosmos_framework/model/generator/mot/flex_attention_bench.py", + "projects/cosmos3/cosmos3/models/mot/flex_attention_memory_test.py": "cosmos_framework/model/generator/mot/flex_attention_memory_test.py", "projects/cosmos3/cosmos3/models/mot/flex_attention_test.py": "cosmos_framework/model/generator/mot/flex_attention_test.py", "projects/cosmos3/cosmos3/models/mot/flex_attention_utils.py": "cosmos_framework/model/generator/mot/flex_attention_utils.py", "projects/cosmos3/cosmos3/models/mot/fsdp_cpu_offload_test.py": "cosmos_framework/model/generator/mot/fsdp_cpu_offload_test.py", @@ -434,6 +438,8 @@ "projects/cosmos3/cosmos3/models/mot/parallelize_vfm_network_test.py": "cosmos_framework/model/generator/mot/parallelize_vfm_network_test.py", "projects/cosmos3/cosmos3/models/mot/replica_partition.py": "cosmos_framework/model/generator/mot/replica_partition.py", "projects/cosmos3/cosmos3/models/mot/replica_partition_test.py": "cosmos_framework/model/generator/mot/replica_partition_test.py", + "projects/cosmos3/cosmos3/models/mot/rigrope.py": "cosmos_framework/model/generator/mot/rigrope.py", + "projects/cosmos3/cosmos3/models/mot/rigrope_test.py": "cosmos_framework/model/generator/mot/rigrope_test.py", "projects/cosmos3/cosmos3/models/mot/und_k_norm_example_test.py": "cosmos_framework/model/generator/mot/und_k_norm_example_test.py", "projects/cosmos3/cosmos3/models/mot/unified_mot.py": "cosmos_framework/model/generator/mot/unified_mot.py", "projects/cosmos3/cosmos3/models/mot/unified_mot_test.py": "cosmos_framework/model/generator/mot/unified_mot_test.py", @@ -493,13 +499,21 @@ "projects/cosmos3/cosmos3/models/sensor_encoder.py": "cosmos_framework/model/generator/sensor_encoder.py", "projects/cosmos3/cosmos3/models/utils/__init__.py": "cosmos_framework/model/generator/utils/__init__.py", "projects/cosmos3/cosmos3/models/utils/batch_normalization.py": "cosmos_framework/model/generator/utils/batch_normalization.py", + "projects/cosmos3/cosmos3/models/utils/camera_rays.py": "cosmos_framework/model/generator/utils/camera_rays.py", + "projects/cosmos3/cosmos3/models/utils/camera_rays_test.py": "cosmos_framework/model/generator/utils/camera_rays_test.py", + "projects/cosmos3/cosmos3/models/utils/camera_relative_pose.py": "cosmos_framework/model/generator/utils/camera_relative_pose.py", + "projects/cosmos3/cosmos3/models/utils/camera_relative_pose_test.py": "cosmos_framework/model/generator/utils/camera_relative_pose_test.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition.py": "cosmos_framework/model/generator/utils/data_and_condition.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition_test.py": "cosmos_framework/model/generator/utils/data_and_condition_test.py", + "projects/cosmos3/cosmos3/models/utils/ftheta.py": "cosmos_framework/model/generator/utils/ftheta.py", "projects/cosmos3/cosmos3/models/utils/load_balancing_stats.py": "cosmos_framework/model/generator/utils/load_balancing_stats.py", "projects/cosmos3/cosmos3/models/utils/memory.py": "cosmos_framework/model/generator/utils/memory.py", "projects/cosmos3/cosmos3/models/utils/moe_utils.py": "cosmos_framework/model/generator/utils/moe_utils.py", + "projects/cosmos3/cosmos3/models/utils/prope_cross_view_poses_test.py": "cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py", "projects/cosmos3/cosmos3/models/utils/rig_view_embedding.py": "cosmos_framework/model/generator/utils/rig_view_embedding.py", "projects/cosmos3/cosmos3/models/utils/rig_view_embedding_test.py": "cosmos_framework/model/generator/utils/rig_view_embedding_test.py", + "projects/cosmos3/cosmos3/models/utils/rigrope_geometry.py": "cosmos_framework/model/generator/utils/rigrope_geometry.py", + "projects/cosmos3/cosmos3/models/utils/rigrope_geometry_test.py": "cosmos_framework/model/generator/utils/rigrope_geometry_test.py", "projects/cosmos3/cosmos3/models/utils/safetensors_loader.py": "cosmos_framework/model/generator/utils/safetensors_loader.py", "projects/cosmos3/cosmos3/models/utils/safetensors_loader_indexed_test.py": "cosmos_framework/model/generator/utils/safetensors_loader_indexed_test.py", "projects/cosmos3/cosmos3/models/utils/safetensors_loader_test.py": "cosmos_framework/model/generator/utils/safetensors_loader_test.py", @@ -615,6 +629,7 @@ "projects/cosmos3/interactive/models/mot/maskless_attention.py": "cosmos_framework/model/generator/mot/maskless_attention.py", "projects/cosmos3/interactive/models/mot/maskless_kernels.py": "cosmos_framework/model/generator/mot/maskless_kernels.py", "projects/cosmos3/interactive/models/mot/merge_attention.py": "cosmos_framework/model/generator/mot/merge_attention.py", + "projects/cosmos3/interactive/models/mot/multiview_action_attention.py": "cosmos_framework/model/generator/mot/multiview_action_attention.py", "projects/cosmos3/interactive/models/mot/post_saturation/__init__.py": "cosmos_framework/model/generator/mot/post_saturation/__init__.py", "projects/cosmos3/interactive/models/mot/post_saturation/cuda_graph.py": "cosmos_framework/model/generator/mot/post_saturation/cuda_graph.py", "projects/cosmos3/interactive/models/mot/post_saturation/installer.py": "cosmos_framework/model/generator/mot/post_saturation/installer.py", @@ -624,10 +639,6 @@ "projects/cosmos3/interactive/models/mot/three_way_attention_test.py": "cosmos_framework/model/generator/mot/three_way_attention_test.py", "projects/cosmos3/interactive/models/omni_mot_causal_model.py": "cosmos_framework/model/generator/omni_mot_causal_model.py", "projects/cosmos3/interactive/models/omni_mot_causal_model_test.py": "cosmos_framework/model/generator/omni_mot_causal_model_test.py", - "projects/cosmos3/interactive/models/rolling_prompt.py": "cosmos_framework/model/generator/rolling_prompt.py", - "projects/cosmos3/interactive/models/rolling_replay.py": "cosmos_framework/model/generator/rolling_replay.py", - "projects/cosmos3/interactive/models/rolling_sink_rope.py": "cosmos_framework/model/generator/rolling_sink_rope.py", - "projects/cosmos3/interactive/models/rolling_transfer_ar.py": "cosmos_framework/model/generator/rolling_transfer_ar.py", "projects/cosmos3/interactive/models/teacher_forcing.py": "cosmos_framework/model/generator/teacher_forcing.py", "projects/cosmos3/interactive/models/utils/frustum_retrieval.py": "cosmos_framework/model/generator/utils/frustum_retrieval.py", "projects/cosmos3/interactive/models/utils/kv_cache.py": "cosmos_framework/model/generator/utils/kv_cache.py", @@ -649,6 +660,7 @@ "projects/cosmos3/tokenizer/lidar_tokenizer/checkpoint_test.py": "cosmos_framework/model/generator/tokenizers/lidar/checkpoint_test.py", "projects/cosmos3/tokenizer/lidar_tokenizer/dtypes.py": "cosmos_framework/model/generator/tokenizers/lidar/dtypes.py", "projects/cosmos3/tokenizer/lidar_tokenizer/dtypes_test.py": "cosmos_framework/model/generator/tokenizers/lidar/dtypes_test.py", + "projects/cosmos3/tokenizer/lidar_tokenizer/geometry.py": "cosmos_framework/model/generator/tokenizers/lidar/geometry.py", "projects/cosmos3/tokenizer/lidar_tokenizer/lidar_tokenizer_v1.py": "cosmos_framework/model/generator/tokenizers/lidar/lidar_tokenizer_v1.py", "projects/cosmos3/tokenizer/lidar_tokenizer/network/__init__.py": "cosmos_framework/model/generator/tokenizers/lidar/network/__init__.py", "projects/cosmos3/tokenizer/lidar_tokenizer/network/distributions.py": "cosmos_framework/model/generator/tokenizers/lidar/network/distributions.py", diff --git a/cosmos_framework/callbacks/data_stats_writer.py b/cosmos_framework/callbacks/data_stats_writer.py index a11859c96..7d815f608 100644 --- a/cosmos_framework/callbacks/data_stats_writer.py +++ b/cosmos_framework/callbacks/data_stats_writer.py @@ -51,16 +51,22 @@ def arrow(self) -> pa.Field: _FIELDS = ( - _Field("phase", pa.string(), 'Data split that produced the row: "train" or "val".'), + _Field("phase", pa.string(), 'Data split that produced the row: "train", "val", or "data_only".'), _Field("step", pa.int64(), "Completed optimizer step associated with the sample."), _Field("rank", pa.int32(), "Distributed process rank that consumed the sample."), _Field("sample_rank", pa.int32(), "Rank recorded by the dataset distributor."), _Field("worker_id", pa.int32(), "Dataloader worker that selected the sample."), - _Field("epoch", pa.int32(), "Distributor epoch from which the sample was selected."), - _Field("sample_index", pa.int64(), "Sample index within the worker's epoch shard."), + _Field("epoch", pa.int32(), "WDS distributor epoch or Lance recipe epoch (legacy: per-source pass)."), + _Field("sample_index", pa.int64(), "WDS epoch shard index or zero-based Lance worker-local schedule position."), _Field("dataset_name", pa.string(), "Recipe dataset/category name."), + _Field("source_backend", pa.string(), "Source identity domain: webdataset or lance."), + _Field("dataset_population", pa.int64(), "Lance grouped sample count; null for WebDataset."), + _Field("sample_id", pa.string(), "Stable sample identity independent of GCS or Lustre storage."), + _Field("source_occurrence", pa.int64(), "Worker-local Lance source cursor after this occurrence."), + _Field("group_index", pa.int64(), "Source-local Lance grouped sample index; null for WebDataset or legacy runs."), + _Field("batch_sample_index", pa.int32(), "Logical sample position within the delivered batch."), _Field("file_name", pa.string(), "Sample key with the leading dataset prefix removed."), - _Field("url", pa.string(), "Full source tar URL/path."), + _Field("url", pa.string(), "Source artifact URL/path; a TAR for WebDataset."), _Field("num_trainable_tokens", pa.int32(), "Valid next-token labels for this sample."), _Field("token_ce_sum", pa.float32(), "Cross entropy summed over valid labels in this sample."), _Field("mean_token_ce", pa.float32(), "token_ce_sum divided by num_trainable_tokens."), @@ -110,7 +116,7 @@ def arrow(self) -> pa.Field: DATA_STATS_SCHEMA = pa.schema( [field.arrow() for field in _FIELDS], - metadata={b"schema_version": b"1"}, + metadata={b"schema_version": b"2"}, ) @@ -162,7 +168,7 @@ def merge_save_dir(save_dir: Path) -> None: ) expected_rows = sum(pq.read_metadata(path).num_rows for path in rank_files) temporary_path = save_dir / f"{MERGED_FILE}.tmp" - with pq.ParquetWriter(temporary_path, DATA_STATS_SCHEMA, compression="zstd") as writer: + with pq.ParquetWriter(temporary_path, pq.read_schema(rank_files[0]), compression="zstd") as writer: for path in rank_files: writer.write_table(pq.read_table(path)) # The rank files are deleted next, so the merged file must read back with every row first. @@ -199,51 +205,69 @@ def checkpoint_data_stats_complete(checkpoint_dir: str | Path) -> bool: "sample_worker_id", "sample_epoch", "sample_index", - "dataset_num_urls", ) +_PROVENANCE_KEYS = frozenset({"source_backend", "dataset_population", "sample_id", "source_occurrence", "group_index"}) # Every schema field that _record does not compute must come from the sample's data_stats dict. -_COMPUTED_FIELDS = frozenset( - { - "phase", - "step", - "rank", - "sample_rank", - "worker_id", - "epoch", - "sample_index", - "dataset_name", - "file_name", - "url", - "num_trainable_tokens", - "token_ce_sum", - "mean_token_ce", - "objective_numerator", - "objective_weight", - "objective_contribution", - "context_length", - "dataset_num_urls", - "total_training_steps", - } +_COMPUTED_FIELDS = ( + frozenset( + { + "phase", + "step", + "rank", + "sample_rank", + "worker_id", + "epoch", + "sample_index", + "dataset_name", + "file_name", + "url", + "num_trainable_tokens", + "token_ce_sum", + "mean_token_ce", + "objective_numerator", + "objective_weight", + "objective_contribution", + "context_length", + "dataset_num_urls", + "total_training_steps", + "batch_sample_index", + } + ) + | _PROVENANCE_KEYS ) _METADATA_KEYS = frozenset(DATA_STATS_SCHEMA.names) - _COMPUTED_FIELDS def _check_sample(key: str, metadata: dict[str, Any], batch_values: dict[str, Any]) -> None: """Raise, naming the sample, if its data-stats fields are missing, unexpected, null, or implausible.""" - if set(metadata) != _METADATA_KEYS: + required_metadata = set(metadata) - _PROVENANCE_KEYS + if required_metadata != _METADATA_KEYS: raise KeyError( - f"data_stats of sample {key}: missing {sorted(_METADATA_KEYS - set(metadata))}, " - f"unexpected {sorted(set(metadata) - _METADATA_KEYS)}" + f"data_stats of sample {key}: missing {sorted(_METADATA_KEYS - required_metadata)}, " + f"unexpected {sorted(required_metadata - _METADATA_KEYS)}" ) values = {**batch_values, "dataset_weight": metadata["dataset_weight"]} nulls = [name for name, value in values.items() if value is None] if nulls: raise ValueError(f"Sample {key} has no value for {nulls}") - if not 0 < values["dataset_weight"] <= 1 or values["dataset_num_urls"] <= 0: + population_key = "dataset_population" if metadata.get("source_backend") == "lance" else "dataset_num_urls" + population = metadata.get(population_key, values.get(population_key)) + if not 0 < values["dataset_weight"] <= 1 or population is None or population <= 0: raise ValueError( f"Sample {key} has dataset_weight={values['dataset_weight']} and " - f"dataset_num_urls={values['dataset_num_urls']}; expected a weight in (0, 1] and a positive count" + f"{population_key}={population}; expected a weight in (0, 1] and a positive count" ) + if metadata.get("source_backend") == "lance" and ( + not metadata.get("sample_id") + or not isinstance(metadata.get("source_occurrence"), int) + or metadata["source_occurrence"] <= 0 + ): + raise ValueError(f"Sample {key} has no valid Lance identity/cursor") + group_index = metadata.get("group_index") + if group_index is not None and ( + metadata.get("source_backend") != "lance" or type(group_index) is not int or not 0 <= group_index < population + ): + raise ValueError(f"Sample {key} has invalid group_index={group_index}") def _values(value: Any, size: int) -> list[Any]: @@ -273,11 +297,14 @@ def __init__( enabled: bool = False, merge_in_background: bool = True, merge_timeout_seconds: float = 1800.0, + flush_on_end: bool = False, ) -> None: super().__init__() self.enabled = enabled self.merge_in_background = merge_in_background self.merge_timeout_seconds = merge_timeout_seconds + self.flush_on_end = flush_on_end + self._last_recorded_iteration: int = 0 # Rows recorded since the previous save, one small table per microbatch. self._tables: list[pa.Table] = [] self._rank = 0 @@ -315,7 +342,7 @@ def _record( self, phase: str, data_batch: dict[str, Any], - output_batch: dict[str, torch.Tensor], + output_batch: dict[str, torch.Tensor] | None, iteration: int, ) -> None: if not self.enabled: @@ -330,7 +357,7 @@ def _record( ) if data_batch.get("true_packing") or "cu_seq_lens_q" in data_batch: raise NotImplementedError("Per-sample data-stats logging currently requires padded batches") - missing = [key for key in _LOSS_KEYS if key not in output_batch] + missing = [key for key in _LOSS_KEYS if output_batch is not None and key not in output_batch] if missing: raise KeyError(f"Model output is missing data-stats loss fields: {missing}") missing = [key for key in _BATCH_KEYS if data_batch.get(key) is None] @@ -339,22 +366,33 @@ def _record( size = len(metadata) columns = {key: _values(data_batch[key], size) for key in _BATCH_KEYS} + columns["dataset_num_urls"] = _values(data_batch.get("dataset_num_urls", [None] * size), size) keys = [str(key) for key in columns["__key__"]] for index, sample_metadata in enumerate(metadata): - _check_sample(keys[index], sample_metadata, {name: values[index] for name, values in columns.items()}) - token_sums = _values(output_batch[_LOSS_KEYS[0]], size) - valid_counts = _values(output_batch[_LOSS_KEYS[1]], size) - numerators = _values(output_batch[_LOSS_KEYS[2]], size) - weights = _values(output_batch[_LOSS_KEYS[3]], size) - global_weight = float(output_batch[_LOSS_KEYS[4]].detach().cpu()) + batch_values = {name: values[index] for name, values in columns.items()} + if sample_metadata.get("source_backend") == "lance": + batch_values.pop("dataset_num_urls") + _check_sample(keys[index], sample_metadata, batch_values) + if output_batch is None: + labels = data_batch["labels"] + ignore_indices = _values(data_batch["ignore_index"], size) + valid_counts = [int((labels[index, 1:] != ignore_indices[index]).sum()) for index in range(size)] + token_sums = numerators = weights = [None] * size + global_weight = 0.0 + else: + token_sums = _values(output_batch[_LOSS_KEYS[0]], size) + valid_counts = _values(output_batch[_LOSS_KEYS[1]], size) + numerators = _values(output_batch[_LOSS_KEYS[2]], size) + weights = _values(output_batch[_LOSS_KEYS[3]], size) + global_weight = float(output_batch[_LOSS_KEYS[4]].detach().cpu()) rows: list[dict[str, Any]] = [] context_length = int(self.config.model.config.policy.model_max_length) total_steps = int(self.config.trainer.max_iter) for index, sample_metadata in enumerate(metadata): token_count = int(valid_counts[index]) - token_sum = float(token_sums[index]) - numerator = float(numerators[index]) + token_sum = None if token_sums[index] is None else float(token_sums[index]) + numerator = None if numerators[index] is None else float(numerators[index]) key = keys[index] rows.append( { @@ -366,16 +404,24 @@ def _record( "epoch": int(columns["sample_epoch"][index]), "sample_index": int(columns["sample_index"][index]), "dataset_name": str(columns["dataset_name"][index]), + "batch_sample_index": index, + "source_backend": sample_metadata.get("source_backend", "webdataset"), + "dataset_population": sample_metadata.get("dataset_population"), + "sample_id": sample_metadata.get("sample_id", f"{columns['__url__'][index]}:{key}"), + "source_occurrence": sample_metadata.get("source_occurrence"), + "group_index": sample_metadata.get("group_index"), "file_name": key.split("__", 1)[1] if "__" in key else key, "url": str(columns["__url__"][index]), "num_trainable_tokens": token_count, "token_ce_sum": token_sum, - "mean_token_ce": token_sum / token_count if token_count else None, + "mean_token_ce": token_sum / token_count if token_count and token_sum is not None else None, "objective_numerator": numerator, - "objective_weight": float(weights[index]), - "objective_contribution": numerator / global_weight if global_weight > 0 else None, + "objective_weight": None if weights[index] is None else float(weights[index]), + "objective_contribution": numerator / global_weight + if global_weight > 0 and numerator is not None + else None, "context_length": context_length, - "dataset_num_urls": int(columns["dataset_num_urls"][index]), + "dataset_num_urls": columns["dataset_num_urls"][index], "total_training_steps": total_steps, # Includes dataset_weight: the batch-level tensor is bf16 here, but custom_collate keeps # the exact value per sample. @@ -386,6 +432,11 @@ def _record( self._tables.append(pa.Table.from_pylist(rows, schema=DATA_STATS_SCHEMA)) except (pa.ArrowException, OverflowError) as error: raise ValueError(f"data_stats rows of samples {keys} do not match the schema: {error}") from error + self._last_recorded_iteration = iteration + + def record_data_batch(self, data_batch: dict[str, Any], iteration: int) -> None: + """Record delivered samples without a model forward pass or invented losses.""" + self._record("data_only", data_batch, None, iteration) def on_training_step_batch_end( self, @@ -497,6 +548,16 @@ def on_app_end(self) -> None: if not self.enabled: return self._complete_last_save() + if self.flush_on_end: + iteration = self._last_recorded_iteration + if dist.is_available() and dist.is_initialized(): + iterations = [iteration] * self._world_size + dist.all_gather_object(iterations, iteration) + iteration = max(iterations) + if iteration != self._last_save: + self._flush(iteration) + self._complete_last_save() + return dropped = sum(table.num_rows for table in self._tables) if dropped: log.info(f"[DataStats] Not saving {dropped} rows recorded after the final checkpoint on rank 0.") diff --git a/cosmos_framework/callbacks/dataloader_state.py b/cosmos_framework/callbacks/dataloader_state.py index c3659e236..af1977ba2 100644 --- a/cosmos_framework/callbacks/dataloader_state.py +++ b/cosmos_framework/callbacks/dataloader_state.py @@ -15,13 +15,19 @@ from cosmos_framework.utils import log from cosmos_framework.utils.callback import Callback -LANCE_VLM_RESUME_FORMAT = "lance_vlm_cursor_v1" +LANCE_VLM_RESUME_FORMAT = "lance_vlm_schedule_v2" LANCE_VLM_RESUME_STATE_KEY = "_lance_vlm_resume_state" LANCE_VLM_RESUME_WORKER_ENV_PREFIX = "LANCE_VLM_RESUME_STATE_WORKER_" +LANCE_VLM_VAL_RESUME_WORKER_ENV_PREFIX = "LANCE_VLM_VAL_RESUME_STATE_WORKER_" +LANCE_VLM_VAL_RESUME_BATCH_COUNT_ENV = "LANCE_VLM_VAL_RESUME_BATCH_COUNT" VAL_STATE_KEY = "val" _STATE_LABELS = {"train": "dataloader", "val": "validation dataloader"} +def lance_resume_worker_env_prefix(split: str) -> str: + return LANCE_VLM_RESUME_WORKER_ENV_PREFIX if split == "train" else LANCE_VLM_VAL_RESUME_WORKER_ENV_PREFIX + + def no_replace_resume_env_names(split: str, worker_id: int) -> tuple[str, str]: """Return the env vars that carry one worker's last consumed (epoch, index) for a split.""" prefix = "NSL_STATE_WORKER" if split == "train" else "NSL_STATE_VAL_WORKER" @@ -47,7 +53,10 @@ def __init__( self.state: dict[int, NoReplaceShardlistState] = {} self.val_state: dict[int, NoReplaceShardlistState] = {} self.lance_state: dict[int, dict[str, Any]] = {} + self.val_lance_state: dict[int, dict[str, Any]] = {} + self.val_lance_batches: int = 0 self._pending_lance_state: dict[str, Any] | None = None + self._pending_val_lance_state: dict[str, Any] | None = None self.verbose = True def _tracks_validation(self) -> bool: @@ -69,13 +78,14 @@ def on_training_step_batch_start( raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict, got {type(lance_state).__name__}.") self._pending_lance_state = lance_state - def _update_lance_state(self, update: dict[str, Any]) -> None: + def _update_lance_state(self, update: dict[str, Any], split: str = "train") -> None: if update.get("format") != LANCE_VLM_RESUME_FORMAT: raise ValueError(f"Unsupported Lance VLM resume state format {update.get('format')!r}.") worker_id = update.get("worker_id") draw_count = update.get("draw_count") fingerprint = update.get("fingerprint") source_cursors = update.get("source_cursors") + source_drops = update.get("source_drops") pool = update.get("pool") if not isinstance(worker_id, int) or worker_id < 0: raise ValueError(f"Lance VLM resume worker_id must be non-negative, got {worker_id!r}.") @@ -88,17 +98,22 @@ def _update_lance_state(self, update: dict[str, Any]) -> None: for name, cursor in source_cursors.items() ): raise ValueError("Lance VLM source cursors must map source names to non-negative integers.") + if not isinstance(source_drops, dict) or any( + not isinstance(name, str) or type(count) is not int or count < 0 for name, count in source_drops.items() + ): + raise ValueError("Lance VLM source drops must map names to non-negative counts") if not isinstance(pool, list) or any( not isinstance(entry, list) - or len(entry) != 2 + or len(entry) != 5 or not isinstance(entry[0], str) - or not isinstance(entry[1], int) - or entry[1] < 1 + or any(type(value) is not int or value < 0 for value in entry[1:]) + or entry[4] < 1 for entry in pool ): - raise ValueError("Lance VLM packing pool must be a list of [source name, positive cursor] pairs.") + raise ValueError("Lance VLM pool entries must contain source, group, epoch, position and source cursor") - saved = self.lance_state.get(worker_id) + states = self.lance_state if split == "train" else self.val_lance_state + saved = states.get(worker_id) previous_current = saved["current"] if saved is not None else None if previous_current is not None: if previous_current["fingerprint"] != fingerprint: @@ -111,17 +126,23 @@ def _update_lance_state(self, update: dict[str, Any]) -> None: if previous_cursor is not None and cursor < previous_cursor: raise ValueError(f"Lance VLM source cursor {source_name!r} moved backwards for worker {worker_id}.") merged_cursors = {**previous_cursors, **source_cursors} + previous_drops = previous_current["source_drops"] + if any(count < previous_drops.get(name, 0) for name, count in source_drops.items()): + raise ValueError("Lance VLM source drop count moved backwards") + merged_drops = {**previous_drops, **source_drops} else: merged_cursors = dict(source_cursors) + merged_drops = dict(source_drops) current = { "format": LANCE_VLM_RESUME_FORMAT, "worker_id": worker_id, "draw_count": draw_count, "fingerprint": fingerprint, "source_cursors": merged_cursors, + "source_drops": merged_drops, "pool": pool, } - self.lance_state[worker_id] = {"previous": previous_current, "current": current} + states[worker_id] = {"previous": previous_current, "current": current} @staticmethod def _update_state_from_batch( @@ -172,6 +193,19 @@ def on_training_step_end( msg += f"worker {wid}: epoch={state.epoch}, index={state.index}\n" log.info(msg) + def on_validation_step_start( + self, + model: ImaginaireModel, + data_batch: dict[str, Any], + iteration: int = 0, + ) -> None: + if self.distributor_type != "no_replace": + return + lance_state = data_batch.pop(LANCE_VLM_RESUME_STATE_KEY, None) + if lance_state is not None and not isinstance(lance_state, dict): + raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict, got {type(lance_state).__name__}.") + self._pending_val_lance_state = lance_state if self._tracks_validation() else None + def on_validation_step_end( self, model: ImaginaireModel, @@ -180,8 +214,14 @@ def on_validation_step_end( loss: torch.Tensor, iteration: int = 0, ) -> None: - # Lance batches carry no shardlist cursor. - if "sample_worker_id" in data_batch and self._tracks_validation(): + if not self._tracks_validation(): + return + # Lance batches carry a group-schedule cursor instead of a shardlist cursor. + if self._pending_val_lance_state is not None: + self._update_lance_state(self._pending_val_lance_state, "val") + self.val_lance_batches += 1 + self._pending_val_lance_state = None + elif "sample_worker_id" in data_batch: self._update_state_from_batch(self.val_state, data_batch) def has_checkpoint_state(self) -> bool: @@ -195,7 +235,13 @@ def state_dict(self) -> dict[Any, Any]: state_dict: dict[Any, Any] = {"format": LANCE_VLM_RESUME_FORMAT, "workers": self.lance_state} else: state_dict = self._save_shardlist_state(self.state, "train") - if self.val_state: + if self.val_lance_state: + state_dict[VAL_STATE_KEY] = { + "format": LANCE_VLM_RESUME_FORMAT, + "workers": self.val_lance_state, + "batch_count": self.val_lance_batches, + } + elif self.val_state: state_dict[VAL_STATE_KEY] = self._save_shardlist_state(self.val_state, "val") return state_dict @@ -211,31 +257,47 @@ def load_state_dict(self, state_dict: dict[Any, Any]) -> None: val_state_dict = state_dict.pop(VAL_STATE_KEY, None) if val_state_dict is not None: if self._tracks_validation(): - self.val_state = self._load_shardlist_state(val_state_dict, "val") + if val_state_dict.get("format") == LANCE_VLM_RESUME_FORMAT: + batch_count = val_state_dict.get("batch_count") + if type(batch_count) is not int or batch_count < 0: + raise ValueError("Lance validation checkpoint requires a non-negative batch_count") + self.val_lance_state = self._load_lance_state(val_state_dict, "val") + self.val_lance_batches = batch_count + os.environ[LANCE_VLM_VAL_RESUME_BATCH_COUNT_ENV] = str(batch_count) + elif "format" in val_state_dict: + raise ValueError(f"Unsupported validation sampler state {val_state_dict['format']!r}") + else: + self.val_state = self._load_shardlist_state(val_state_dict, "val") else: log.info("Ignoring validation dataloader state because the validation iterator is not persistent") if state_dict.get("format") == LANCE_VLM_RESUME_FORMAT: - workers = state_dict.get("workers") - if not isinstance(workers, dict): - raise TypeError("Lance VLM checkpoint workers state must be a dictionary.") - self.lance_state = {int(worker_id): state_pair for worker_id, state_pair in workers.items()} - for worker_id, state_pair in self.lance_state.items(): - if not isinstance(state_pair, dict) or not isinstance(state_pair.get("current"), dict): - raise ValueError(f"Invalid Lance VLM resume state for worker {worker_id}.") - os.environ[f"{LANCE_VLM_RESUME_WORKER_ENV_PREFIX}{worker_id}"] = json.dumps( - state_pair, - separators=(",", ":"), - sort_keys=True, - ) - log.info( - f"Loaded Lance VLM dataloader state for worker {worker_id}: " - f"draw_count={state_pair['current']['draw_count']}" - ) + self.lance_state = self._load_lance_state(state_dict, "train") + if val_state_dict is None and self.config is not None and self._tracks_validation(): + log.warning("Checkpoint has no validation dataloader state; validation will start from epoch zero") return + if "format" in state_dict: + raise ValueError(f"Unsupported sampler state {state_dict['format']!r}; finite sampling requires a new run") self.state = self._load_shardlist_state(state_dict, "train") + @staticmethod + def _load_lance_state(state_dict: dict[str, Any], split: str) -> dict[int, dict[str, Any]]: + workers = state_dict.get("workers") + if not isinstance(workers, dict): + raise TypeError("Lance VLM checkpoint workers state must be a dictionary.") + states = {int(worker_id): state_pair for worker_id, state_pair in workers.items()} + prefix = lance_resume_worker_env_prefix(split) + for worker_id, state_pair in states.items(): + if not isinstance(state_pair, dict) or not isinstance(state_pair.get("current"), dict): + raise ValueError(f"Invalid Lance VLM resume state for worker {worker_id}.") + os.environ[f"{prefix}{worker_id}"] = json.dumps(state_pair, separators=(",", ":"), sort_keys=True) + log.info( + f"Loaded Lance VLM {_STATE_LABELS[split]} state for worker {worker_id}: " + f"draw_count={state_pair['current']['draw_count']}" + ) + return states + @staticmethod def _save_shardlist_state(state: dict[int, NoReplaceShardlistState], split: str) -> dict[int, dict[str, int]]: state_dict: dict[int, dict[str, int]] = {} diff --git a/cosmos_framework/callbacks/sampled_media_recorder.py b/cosmos_framework/callbacks/sampled_media_recorder.py index 022cb4400..ab82ad0e5 100644 --- a/cosmos_framework/callbacks/sampled_media_recorder.py +++ b/cosmos_framework/callbacks/sampled_media_recorder.py @@ -181,6 +181,7 @@ def _extract_records(self, data_batch: dict[str, Any], iteration: int, rank: int count = len(sample_ids) media_urls = _as_list(data_batch.get("__url__"), count) dataset_names = _as_list(data_batch.get("dataset_name"), count) + source_streams = _as_list(data_batch.get("transfer_source_stream"), count) source_names = _as_list(data_batch.get("source_dataset_name"), count) captions = _as_optional_string_list(data_batch.get("ai_caption"), count) if self.record_caption else [] source_ids = _as_list(data_batch.get("source_id"), count) @@ -208,6 +209,7 @@ def _extract_records(self, data_batch: dict[str, Any], iteration: int, rank: int "rank": rank, "media_type": media_types[sample_index], "dataset_name": dataset_names[sample_index], + "source_stream": source_streams[sample_index] or dataset_names[sample_index], "source_dataset_name": source_names[sample_index], "source_id": source_ids[sample_index], "sample_id": sample_id, @@ -235,6 +237,7 @@ def _table_schema() -> Any: pa.field("rank", pa.int32(), nullable=False), pa.field("media_type", pa.string(), nullable=False), pa.field("dataset_name", pa.string(), nullable=False), + pa.field("source_stream", pa.string(), nullable=True), pa.field("source_dataset_name", pa.string(), nullable=False), pa.field("source_id", pa.string(), nullable=True), pa.field("sample_id", pa.string(), nullable=False), @@ -254,7 +257,7 @@ def _upgrade_legacy_table_schema(self, existing: Any, storage_options: dict[str, # A tuple in _table_schema() declaration order, not a set: the loop below appends # missing columns in this order, and set iteration is hash-randomized -- two runs # upgrading the same legacy table would otherwise produce different column orders. - optional_field_order = ("source_id", "caption", "caption_mode", "conversation", "media_items") + optional_field_order = ("source_stream", "source_id", "caption", "caption_mode", "conversation", "media_items") optional_fields = frozenset(optional_field_order) expected_by_name = {field.name: field for field in expected_schema} existing_by_name = {field.name: field for field in existing.schema} diff --git a/cosmos_framework/callbacks/sampled_media_recorder_test.py b/cosmos_framework/callbacks/sampled_media_recorder_test.py index 8a012ef74..b13d03f4b 100644 --- a/cosmos_framework/callbacks/sampled_media_recorder_test.py +++ b/cosmos_framework/callbacks/sampled_media_recorder_test.py @@ -25,6 +25,7 @@ def _sample_record(iteration: int, caption: str | None = None) -> dict[str, obje "rank": 0, "media_type": "video", "dataset_name": "video_256", + "source_stream": "video_256", "source_dataset_name": "source", "source_id": "", "sample_id": f"video-{iteration}", @@ -52,11 +53,26 @@ def test_extract_records_from_consumed_image_batch(monkeypatch: pytest.MonkeyPat assert [record["sample_id"] for record in records] == ["image-a", "image-b"] assert [record["media_type"] for record in records] == ["image", "image"] assert [record["source_dataset_name"] for record in records] == ["source-a", "source-b"] + assert [record["source_stream"] for record in records] == ["images", "images"] assert all(record["run_id"] == "12345" for record in records) assert all(record["iteration"] == 7 for record in records) assert all(record["rank"] == 3 for record in records) +def test_extract_records_preserves_mixed_transfer_source_stream() -> None: + callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance") + batch = { + "__key__": ["sample-a"], + "dataset_name": ["transfer_rgb"], + "transfer_source_stream": ["transfer_depth_single_view_720"], + } + + records = callback._extract_records(batch, iteration=7, rank=3) + + assert records[0]["dataset_name"] == "transfer_rgb" + assert records[0]["source_stream"] == "transfer_depth_single_view_720" + + def test_extract_records_prefers_action_fingerprint() -> None: callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance") callback.config = SimpleNamespace(job=SimpleNamespace(name="test_experiment")) @@ -208,8 +224,8 @@ def test_local_lance_append_upgrades_legacy_table_metadata_only(tmp_path: Path) pa = pytest.importorskip("pyarrow") output_uri = str(tmp_path / "samples.lance") callback = SampledMediaRecorder(enabled=True, output_uri=output_uri) - vlm_fields = {"source_id", "conversation", "media_items"} - legacy_schema = pa.schema([field for field in callback._table_schema() if field.name not in vlm_fields]) + missing_fields = {"source_stream", "source_id", "conversation", "media_items"} + legacy_schema = pa.schema([field for field in callback._table_schema() if field.name not in missing_fields]) legacy_table = pa.Table.from_pylist([_sample_record(1)], schema=legacy_schema) lance.write_dataset(legacy_table, output_uri, mode="create") data_files_before = lance.dataset(output_uri).get_fragments()[0].data_files() @@ -221,3 +237,4 @@ def test_local_lance_append_upgrades_legacy_table_metadata_only(tmp_path: Path) assert set(upgraded.schema.names) == set(callback._table_schema().names) assert upgraded.to_table(columns=["caption"])["caption"].to_pylist() == [None, "recorded caption"] assert upgraded.to_table(columns=["source_id"])["source_id"].to_pylist() == [None, ""] + assert upgraded.to_table(columns=["source_stream"])["source_stream"].to_pylist() == [None, "video_256"] diff --git a/cosmos_framework/callbacks/wall_clock_checkpoint.py b/cosmos_framework/callbacks/wall_clock_checkpoint.py index 2a687742d..6b9b3c290 100644 --- a/cosmos_framework/callbacks/wall_clock_checkpoint.py +++ b/cosmos_framework/callbacks/wall_clock_checkpoint.py @@ -158,6 +158,7 @@ import time from collections import deque from concurrent.futures import ThreadPoolExecutor +from typing import Any import torch @@ -178,9 +179,11 @@ #: Written before deletion starts so an interrupted deletion can be resumed. DELETING_MARKER = "wall_clock_deleting.json" -#: The state a resume needs, one DCP component per subdirectory. ``dataloader/`` is +#: The standard state a resume needs, one DCP component per subdirectory. ``dataloader/`` is #: left out: it is optional and written as pickles rather than by DCP, so it has no #: ``.metadata`` and its absence says nothing about whether the save finished. +#: Distillation writes separate optimizer/scheduler components; capture those from +#: the writer's save hook and persist them in each wall-clock marker. DCP_COMPONENTS = ("model", "optim", "scheduler", "trainer") #: Written by DCP once a component's shards are all in, which makes it the closest @@ -255,6 +258,10 @@ def __init__( # Confirmed wall-clock checkpoints, ascending. Only maintained on rank 0, # which is the only rank that prunes. self._wall_clock_iterations: list[int] = [] + self._dcp_components: tuple[str, ...] = DCP_COMPONENTS + # Async completion must use that save's layout, even if another save's + # hook has since described different optimizer/scheduler components. + self._pending_dcp_components: dict[int, tuple[str, ...]] = {} # Captured from on_before_optimizer_step so we can call checkpointer.save(). self._optimizer: torch.optim.Optimizer | None = None @@ -332,14 +339,21 @@ def on_save_checkpoint_start(self, model: ImaginaireModel, iteration: int = 0) - self._last_checkpoint_iteration = iteration self._milestone_skip_logged = False + def on_save_checkpoint(self, model: ImaginaireModel, state_dict: dict[str, Any]) -> None: + """Capture the actual DCP writer layout, excluding optional pickle state.""" + del model + self._dcp_components = tuple(key for key in state_dict if key != "dataloader") + self._pending_dcp_components[self._last_checkpoint_iteration] = self._dcp_components + def on_save_checkpoint_success(self, iteration: int = 0, elapsed_time: float = 0) -> None: self._observe_write_duration(elapsed_time) + components = self._pending_dcp_components.pop(iteration, self._dcp_components) if iteration not in self._triggered_iterations: return self._triggered_iterations.discard(iteration) if not distributed.is_rank0(): return - if self._mark_saved(iteration): + if self._mark_saved(iteration, components): self._wall_clock_iterations.append(iteration) self._prune() @@ -504,11 +518,15 @@ def _checkpoint_dirname(self, iteration: int) -> str: def _backend_key(self) -> str | None: return self.trainer.checkpointer.save_s3_backend_key - def _mark_saved(self, iteration: int) -> bool: + def _mark_saved(self, iteration: int, components: tuple[str, ...]) -> bool: """Tag a checkpoint as ours. Returns False if the tag could not be written.""" path = os.path.join(self._checkpoint_dirname(iteration), SAVED_MARKER) try: - easy_io.dump({"iteration": iteration, "saved_at": time.time()}, path, backend_key=self._backend_key) + easy_io.dump( + {"iteration": iteration, "saved_at": time.time(), "dcp_components": list(components)}, + path, + backend_key=self._backend_key, + ) except Exception as error: # noqa: BLE001 - a missing tag only costs us pruning log.warning(f"[WallClockCheckpoint] Could not write {path}: {error}. This checkpoint will not be pruned.") return False @@ -669,7 +687,10 @@ def _verify_complete(self, iteration: int) -> bool: if not checkpoint_data_stats_complete(dirname): log.error(f"[WallClockCheckpoint] Data stats for {dirname} are not merged yet.") return False - for component in DCP_COMPONENTS: + components = self._saved_dcp_components(dirname) + if components is None: + return False + for component in components: component_dirname = os.path.join(dirname, component) path = os.path.join(component_dirname, DCP_METADATA) try: @@ -684,6 +705,25 @@ def _verify_complete(self, iteration: int) -> bool: return False return True + def _saved_dcp_components(self, dirname: str) -> tuple[str, ...] | None: + """Read the writer's saved layout, retaining support for older markers.""" + path = os.path.join(dirname, SAVED_MARKER) + try: + marker = easy_io.load(path, backend_key=self._backend_key) + # Older markers carry no layout. Use the current writer's known + # components, or the standard layout before its first save hook. + components = marker.get("dcp_components", self._dcp_components) + if ( + not isinstance(components, (list, tuple)) + or not all(isinstance(name, str) and name and os.path.basename(name) == name for name in components) + or not {"model", "trainer"}.issubset(components) + ): + raise ValueError("Checkpoint marker has no valid DCP component layout.") + return tuple(components) + except Exception as error: # noqa: BLE001 - unreadable completion evidence must defer deletion + log.warning(f"[WallClockCheckpoint] Could not read components from {path}: {error}") + return None + def _verify_shard_sizes(self, component_dirname: str) -> bool: """Check every shard is at least as long as this component's metadata says. diff --git a/cosmos_framework/configs/base/defaults/replay_attention.py b/cosmos_framework/configs/base/defaults/replay_attention.py index 3fac24ec1..5335bb93c 100644 --- a/cosmos_framework/configs/base/defaults/replay_attention.py +++ b/cosmos_framework/configs/base/defaults/replay_attention.py @@ -12,7 +12,9 @@ AttentionScope, ) -TeacherForcingKVImplementation = Literal["multiview_flex_kv", "multiview_maskless_kv", "singleview_threeway_kv"] +TeacherForcingKVImplementation = Literal[ + "multiview_flex_kv", "multiview_maskless_kv", "singleview_threeway_kv", "multiview_threeway_kv" +] TEACHER_FORCING_KV_IMPLEMENTATIONS = get_args(TeacherForcingKVImplementation) TeacherForcingControlVisibility = Literal["global", "causal", "current"] diff --git a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py index 631d01d58..963502771 100644 --- a/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py +++ b/cosmos_framework/data/generator/augmentors/reasoner/prompt_format.py @@ -45,7 +45,11 @@ def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: # check if this is list of list of dict or list of dict if isinstance(list_of_conversation[0], list): - selected_conversation = random.sample(list_of_conversation, 1)[0] + selected_index = random.sample(range(len(list_of_conversation)), 1)[0] + selected_conversation = list_of_conversation[selected_index] + conversation_metadata = data_dict.pop("_conversation_metadata", None) + if conversation_metadata is not None: + data_dict.update(conversation_metadata[selected_index]) elif isinstance(list_of_conversation[0], dict): selected_conversation = list_of_conversation else: diff --git a/cosmos_framework/data/generator/lance/lance_dataset_open.py b/cosmos_framework/data/generator/lance/lance_dataset_open.py index ba8ce67e4..1def1635d 100644 --- a/cosmos_framework/data/generator/lance/lance_dataset_open.py +++ b/cosmos_framework/data/generator/lance/lance_dataset_open.py @@ -76,11 +76,15 @@ def open_lance_dataset_with_retry( *, storage_options: dict[str, str], version: int | None = None, + metadata_cache_size_bytes: int | None = None, ) -> lance.LanceDataset: """Open a Lance dataset, dispersing retries for transient object-store errors.""" import lance + cache_options = ( + {} if metadata_cache_size_bytes is None else {"metadata_cache_size_bytes": metadata_cache_size_bytes} + ) return run_lance_io_with_retry( - lambda: lance.dataset(uri, storage_options=storage_options, version=version), + lambda: lance.dataset(uri, storage_options=storage_options, version=version, **cache_options), description=f"dataset-open for {uri!r}", ) diff --git a/cosmos_framework/data/generator/multiview/camera_geometry.py b/cosmos_framework/data/generator/multiview/camera_geometry.py new file mode 100644 index 000000000..c68c74dac --- /dev/null +++ b/cosmos_framework/data/generator/multiview/camera_geometry.py @@ -0,0 +1,772 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Optional camera metadata, selected only after the final RGB training window. + +Pixel coordinates denote pixel centers at integer locations. Calibration stays +in its native lens model; ``image_from_calibration`` records the aspect-preserving +resize and center crop into the returned video. Poses stay float64 in the supplied +OpenCV camera-to-world frame (meters), without per-camera recentering. +""" + +from __future__ import annotations + +import json +import math +import os +import warnings +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, TypedDict, cast + +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor +from cosmos_framework.utils import log +from cosmos_framework.data.generator.multiview.camera_geometry_contract import ( + CAMERA_CALIBRATION_SCHEMA, + CAMERA_CALIBRATION_SCHEMA_V1_1, + POSE_SCHEMA, + TIMESTAMPS_SCHEMA, + TIMESTAMPS_SCHEMA_V1_1, + Matrix4, + QuaternionXyzw, + normalize_quaternion_xyzw, + quaternion_xyzw_to_rotation_matrix, + validate_camera_to_world_matrix, +) +from cosmos_framework.data.generator.multiview.camera_geometry_native import ( + MADS_FULL_FRAME_EXPORT, + mads_media_from_calibration, +) +from cosmos_framework.model.generator.utils.camera_rays import SUPPORTED_RAY_CAMERA_MODELS, validate_token_geometry +from cosmos_framework.utils.generator.video_preprocess import resize_and_crop_affine + +SOURCE_GEOMETRY_KEY = "_camera_geometry_source" +SOURCE_FPS_KEY = "_camera_source_fps" +SOURCE_CALIBRATION_ERROR_KEY = "_camera_calibration_error" +# Per-view ``[V,3]`` float64 rows (zoom, top fraction, left fraction) of the crop frame +# extraction applied; absent means every view was center-cropped at zoom 1. +VIEW_CROP_KEY = "_view_crop" +# Vehicle cameras sit within a few meters of the rear-axle rig origin. +EPISODE_START_RIG_MAX_OFFSET_M = 10.0 +# Rig key marking a pose-world rig, whose camera-to-rig transform is each selected frame's pose. +PER_FRAME_CAMERA_TO_RIG_KEY = "per_frame_camera_to_rig" + +SOURCE_RIG_JSON_KEY = "_rig_json" + + +class CameraCalibrationError(ValueError): + """Malformed sample calibration, distinct from IO and training-layout errors.""" + + +@dataclass +class CameraGeometrySource: + """One physical camera, before view/frame sampling; frame IDs index its video.""" + + calibration: dict[str, Any] | None = None + poses: dict[int, Matrix4] = field(default_factory=dict) + times_ns: dict[int, int] | None = None + native_times_ns: dict[int, int] = field(default_factory=dict) + provenance: dict[str, Any] = field(default_factory=dict) + + +class CameraGeometry(TypedDict): + camera_keys: list[str] + frame_indices: torch.Tensor # [V,T] + camera_to_world: torch.Tensor # [V,T,4,4] + pose_valid: torch.Tensor # [V,T] + calibration: list[dict[str, Any] | None] + calibration_valid: torch.Tensor # [V] + image_from_calibration: torch.Tensor # [V,3,3] + image_hw: torch.Tensor # [V,2] + frame_times_ns: torch.Tensor # [V,T] + time_valid: torch.Tensor # [V,T] + time_source: list[list[str]] + native_timestamps_ns: torch.Tensor # [V,T] + native_time_valid: torch.Tensor # [V,T] + provenance: dict[str, Any] + + +def _format_prompt_number(value: float) -> str: + """Format geometry compactly while preserving useful calibration precision.""" + if value == 0.0: + value = 0.0 + return f"{value:.6g}" + + +def _format_prompt_matrix(matrix: torch.Tensor) -> str: + """Render a small matrix as an unambiguous row-major numeric block.""" + return "[" + "; ".join(", ".join(_format_prompt_number(float(value)) for value in row) for row in matrix) + "]" + + +class AddCameraGeometryToCaption(Augmentor): + """Prefix each selected view caption with its intrinsics and rig-relative pose.""" + + reference_cameras: tuple[str, ...] = ( + "camera_front_wide_120fov", + "camera_front_fisheye_200fov", + ) + + def __init__(self, caption_key: str = "ai_caption", geometry_key: str = "camera_geometry") -> None: + super().__init__([], None, None) + self.caption_key = caption_key + self.geometry_key = geometry_key + + def _reference_view(self, geometry: CameraGeometry) -> int | None: + camera_keys = geometry["camera_keys"] + pose_valid = geometry["pose_valid"] # [V,T] + for camera in self.reference_cameras: + if camera in camera_keys: + preferred = camera_keys.index(camera) + if bool(pose_valid[preferred].any()): + return preferred + return next((view for view in range(len(camera_keys)) if bool(pose_valid[view].any())), None) + + def _intrinsics_text(self, geometry: CameraGeometry, view: int) -> str: + calibration = geometry["calibration"][view] + if calibration is None: + return "intrinsics unavailable" + native_k = torch.tensor( # [3,3] + [ + [calibration["fx_px"], 0.0, calibration["cx_px"]], + [0.0, calibration["fy_px"], calibration["cy_px"]], + [0.0, 0.0, 1.0], + ], + dtype=torch.float64, + ) + distortion = ", ".join(_format_prompt_number(float(value)) for value in calibration["distortion_coefficients"]) + output_height, output_width = (int(value) for value in geometry["image_hw"][view].tolist()) + return ( + f"native {calibration['camera_model']} image {calibration['width_px']}x{calibration['height_px']} pixels, " + f"K={_format_prompt_matrix(native_k)} pixels, distortion={calibration['distortion_model']}[{distortion}], " + f"native-to-model-image affine={_format_prompt_matrix(geometry['image_from_calibration'][view])}, " + f"model image {output_width}x{output_height} pixels" + ) + + def _extrinsics_text(self, geometry: CameraGeometry, view: int, reference_view: int | None) -> str: + if reference_view is None: + return "relative extrinsics unavailable" + jointly_valid = geometry["pose_valid"][reference_view] & geometry["pose_valid"][view] # [T] + valid_frames = torch.nonzero(jointly_valid, as_tuple=False).flatten() # [N_valid] + if valid_frames.numel() == 0: + return "relative extrinsics unavailable" + frame = int(valid_frames[0]) + camera_to_reference = ( # [4,4] + torch.linalg.inv(geometry["camera_to_world"][reference_view, frame]) + @ geometry["camera_to_world"][view, frame] + ) + rotation = camera_to_reference[:3, :3] # [3,3] + position = camera_to_reference[:3, 3] # [3] + viewing_direction = rotation[:, 2] # [3] + horizontal_offset_degrees = math.degrees( # [] + math.atan2(float(viewing_direction[0]), float(viewing_direction[2])) + ) + vertical_offset_degrees = math.degrees( # [] + math.atan2( + -float(viewing_direction[1]), + math.hypot(float(viewing_direction[0]), float(viewing_direction[2])), + ) + ) + source_frame = int(geometry["frame_indices"][view, frame]) + reference_camera = geometry["camera_keys"][reference_view] + return ( + f"relative to reference camera {reference_camera} at source frame {source_frame}, this camera center is " + f"at [{_format_prompt_number(float(position[0]))} right, " + f"{_format_prompt_number(float(position[1]))} down, " + f"{_format_prompt_number(float(position[2]))} forward] meters and its optical axis points along " + f"[{_format_prompt_number(float(viewing_direction[0]))}, " + f"{_format_prompt_number(float(viewing_direction[1]))}, " + f"{_format_prompt_number(float(viewing_direction[2]))}] in the reference-camera coordinates. " + f"Its horizontal viewing offset is {_format_prompt_number(horizontal_offset_degrees)}° " + f"(positive right) and vertical viewing offset is {_format_prompt_number(vertical_offset_degrees)}° " + f"(positive up). The exact current-camera-to-reference transform is " + f"{_format_prompt_matrix(camera_to_reference[:3])}" + ) + + def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any]: + captions = data_dict[self.caption_key] + geometry: CameraGeometry = data_dict[self.geometry_key] + camera_keys = geometry["camera_keys"] + if not isinstance(captions, list) or len(captions) != len(camera_keys): + raise ValueError( + "Camera-geometry caption prefixes require one separate caption per selected camera: " + f"captions={len(captions) if isinstance(captions, list) else None}, cameras={len(camera_keys)}" + ) + reference_view = self._reference_view(geometry) + prefixed: list[str] = [] + for view, caption in enumerate(captions): + caption_text = caption if isinstance(caption, str) else json.dumps(caption) + prefix = ( + "Camera calibration for this view uses OpenCV camera axes (x right, y down, z forward). " + f"Intrinsics: {self._intrinsics_text(geometry, view)}. " + f"Extrinsics: {self._extrinsics_text(geometry, view, reference_view)}." + ) + prefixed.append(f"{prefix}\n\n{caption_text}") + data_dict[self.caption_key] = prefixed + return data_dict + + +def format_camera_geometry_captions(captions: object, geometry: dict[str, Any]) -> list[str]: + """Use the training decorator verbatim in inference and fixed validation.""" + if not isinstance(captions, list) or not all(isinstance(caption, str) for caption in captions): + raise ValueError("Camera calibration captions require one separately formatted caption per view") + return AddCameraGeometryToCaption()({"ai_caption": captions, "camera_geometry": geometry})["ai_caption"] + + +_SIDECAR_SCHEMAS: dict[str, tuple[pa.Schema, ...]] = { + "calibration": (CAMERA_CALIBRATION_SCHEMA, CAMERA_CALIBRATION_SCHEMA_V1_1), + "pose": (POSE_SCHEMA,), + "timestamps": (TIMESTAMPS_SCHEMA, TIMESTAMPS_SCHEMA_V1_1), +} + + +def _without_list_item_names(schema: pa.Schema) -> pa.Schema: + """Parquet writers name list children ``item`` or ``element``; the name has no meaning.""" + return pa.schema( + [ + pa.field(item.name, pa.list_(item.type.value_type), nullable=item.nullable) + if pa.types.is_list(item.type) + else item + for item in schema + ] + ) + + +def _match_sidecar_schema(role: str, table_schema: pa.Schema) -> pa.Schema | None: + metadata = table_schema.metadata or {} + fields = _without_list_item_names(table_schema) + for schema in _SIDECAR_SCHEMAS[role]: + if not fields.equals(_without_list_item_names(schema), check_metadata=False): + continue + expected = (schema.metadata or {}).items() + if schema is CAMERA_CALIBRATION_SCHEMA: + # Static camera calibration has no version-specific interpretation beyond + # its exact Arrow fields. Older Lance exports may omit only the canonical + # schema metadata, so accept that bounded legacy form while still rejecting + # different fields/types or contradictory metadata values. + matches = all(metadata.get(key, value) == value for key, value in expected) + else: + # Versioned layouts and pose/timestamp frame and clock semantics must remain exact. + matches = all(metadata.get(key) == value for key, value in expected) + if matches: + return schema + return None + + +def _calibration_from_v1_1(row: dict[str, Any]) -> dict[str, Any]: + """Pack a v1.1.0 ``nvidia_ftheta_polynomial`` row as media-pixel F-theta coefficients. + + Media and native pixels are related by a per-axis affine resize, so offsets from + the principal point scale exactly by ``(s_x, s_y)``. Measuring radius in + ``s_y``-scaled media pixels turns the native backward polynomial into + ``p_i * s_y**i``; residual x anisotropy is the lens matrix ``[[s_y / s_x, 0], [0, 1]]``. + """ + calibration = dict(row) + scale_x = calibration.pop("pixel_to_native_scale_x") + scale_y = calibration.pop("pixel_to_native_scale_y") + if not all(math.isfinite(scale) and scale > 0 for scale in (scale_x, scale_y)): + raise ValueError(f"Invalid pixel-to-native scale: x={scale_x}, y={scale_y}") + if calibration["distortion_model"] != "nvidia_ftheta_polynomial": + if (scale_x, scale_y) != (1.0, 1.0): + raise ValueError(f"Unsupported pixel-to-native scale for {calibration['distortion_model']} calibration") + return calibration + polynomial = list(calibration["distortion_coefficients"]) + if calibration["camera_model"] != "ftheta" or not 2 <= len(polynomial) <= 6: + raise ValueError(f"Unsupported nvidia_ftheta_polynomial calibration with {len(polynomial)} terms") + polynomial += [0.0] * (6 - len(polynomial)) + lens = [] if scale_x == scale_y else [scale_y / scale_x, 0.0, 0.0] + calibration.update( + distortion_model="ftheta", + distortion_coefficients=[ + calibration["cx_px"], + calibration["cy_px"], + float(calibration["width_px"]), + float(calibration["height_px"]), + *(value * scale_y**power for power, value in enumerate(polynomial)), + 1.0, + *lens, + ], + ) + return calibration + + +def read_camera_sidecars( + stream_id: str, payloads: dict[str, bytes | None], provenance: dict[str, Any] +) -> CameraGeometrySource: + """Decode canonical v1.0/v1.1 Parquet sidecars; absence is distinct from invalid data.""" + # The training-local contract pins the supported canonical sidecar versions. + source = CameraGeometrySource(provenance=dict(provenance)) + for role in ("calibration", "pose", "timestamps"): + payload = payloads.get(role) + if payload is None: + continue + table = pq.read_table(pa.BufferReader(payload)) + metadata = table.schema.metadata or {} + schema = _match_sidecar_schema(role, table.schema) + if schema is None: + raise ValueError(f"Unsupported {role} schema for {stream_id}") + rows = table.to_pylist() + if any( + row["stream_id"] != stream_id + or any(value is None for value in row.values()) + or None in row.get("distortion_coefficients", ()) + for row in rows + ): + raise ValueError(f"Invalid {role} stream identity or null values for {stream_id}") + if role == "calibration": + if len(rows) != 1: + raise ValueError(f"Expected one static calibration for {stream_id}") + source.calibration = ( + _calibration_from_v1_1(rows[0]) if schema is CAMERA_CALIBRATION_SCHEMA_V1_1 else rows[0] + ) + _validate_calibration(source.calibration) + continue + indexed = {row["sample_index"]: row for row in rows} + if len(indexed) != len(rows) or any(index < 0 for index in indexed): + raise ValueError(f"Duplicate or negative {role} frame IDs for {stream_id}") + if role == "timestamps": + origin = metadata.get(b"episode_time_origin", b"episode_start").decode() + native_origin = metadata.get(b"native_timestamp_origin", b"provider_clock").decode() + if origin not in {"episode_start", "sample_start"} or native_origin not in { + "sample_start", + "provider_clock", + }: + raise ValueError(f"Unsupported timestamp origin for {stream_id}") + source.provenance.update(time_origin=origin, native_timestamp_origin=native_origin) + source.times_ns = {index: row["episode_time_ns"] for index, row in indexed.items()} + source.native_times_ns = {index: row["native_timestamp_ns"] for index, row in indexed.items()} + ordered_times = [source.times_ns[index] for index in sorted(indexed)] + # Canonical timestamps are nondecreasing; distinct frames may share a timestamp. + if any(b < a for a, b in zip(ordered_times, ordered_times[1:])): + raise ValueError(f"Decreasing timestamps for {stream_id}") + else: + for index, row in indexed.items(): + quaternion = cast(QuaternionXyzw, tuple(row[f"orientation_q{axis}"] for axis in "xyzw")) + rotation = quaternion_xyzw_to_rotation_matrix(normalize_quaternion_xyzw(quaternion, label=stream_id)) + matrix = cast( + Matrix4, + tuple((*rotation[i], row[f"position_{axis}_m"]) for i, axis in enumerate("xyz")) + + ((0.0, 0.0, 0.0, 1.0),), + ) + validate_camera_to_world_matrix(matrix, label=f"{stream_id} frame {index}") + source.poses[index] = matrix + return source + + +def episode_start_rig_geometry(sources: dict[str, CameraGeometrySource]) -> dict[str, Any]: + """Derive camera-to-rig extrinsics from poses anchored at the episode-start rig. + + Canonical AV pose sidecars (e.g. Alpamayo) express camera-to-world in the FLU ego + rig frame at episode time zero, so each camera's pose at that instant is its static + rig extrinsic. A pose far from the rig origin means a map-anchored world frame. + """ + camera_to_rig: dict[str, Matrix4] = {} + for camera, source in sources.items(): + if source.times_ns is None or source.provenance.get("time_origin") != "episode_start": + raise ValueError(f"{camera}: rig extrinsics require episode-start timestamps") + start = min((index for index, time in source.times_ns.items() if time == 0), default=None) + if start is None or start not in source.poses: + raise ValueError(f"{camera}: missing camera pose at episode start") + pose = source.poses[start] + offset_m = math.hypot(*(row[3] for row in pose[:3])) + if offset_m > EPISODE_START_RIG_MAX_OFFSET_M: + raise ValueError(f"{camera}: episode-start pose is {offset_m:.1f} m from the rig origin") + camera_to_rig[camera] = pose + if not camera_to_rig: + raise ValueError("Sensor rig requires at least one camera") + return {"camera_to_rig": camera_to_rig} + + +def pose_world_rig_geometry(sources: dict[str, CameraGeometrySource]) -> dict[str, Any]: + """Use the shared pose world frame as the rig, for cameras that move within a scene. + + Every camera must carry world poses in one scene-fixed frame. ``camera_to_rig`` + holds each camera's first pose only to identify rig membership; RigRoPE reads the + selected frames' ``camera_to_world`` because these cameras have no static extrinsic. + """ + camera_to_rig: dict[str, Matrix4] = {} + for camera, source in sources.items(): + if not source.poses: + raise ValueError(f"{camera}: a pose-world rig requires camera poses") + camera_to_rig[camera] = source.poses[min(source.poses)] + if not camera_to_rig: + raise ValueError("Sensor rig requires at least one camera") + return {"camera_to_rig": camera_to_rig, PER_FRAME_CAMERA_TO_RIG_KEY: True} + + +def _validate_calibration(calibration: dict[str, Any]) -> None: + coefficients = calibration["distortion_coefficients"] + width, height = calibration["width_px"], calibration["height_px"] + scalars = [calibration[key] for key in ("fx_px", "fy_px", "cx_px", "cy_px")] + if width <= 0 or height <= 0 or not all(math.isfinite(value) for value in (*scalars, *coefficients)): + raise ValueError("Invalid camera calibration dimensions or nonfinite coefficients") + if min(scalars[:2]) <= 0: + raise ValueError(f"Invalid focal length: fx={scalars[0]}, fy={scalars[1]}") + # Validate the principal point in the frozen calibration's own pixel domain, + # before applying the separate calibration-to-media and training resize maps. + if not (0 <= scalars[2] <= width and 0 <= scalars[3] <= height): + raise ValueError( + f"Invalid principal point: cx={scalars[2]}, cy={scalars[3]}, " + f"calibration width={width}, height={height}, camera_model={calibration['camera_model']}" + ) + if calibration["camera_model"] == "ftheta": + if ( + calibration["distortion_model"] != "ftheta" + or len(coefficients) not in (11, 14) + or coefficients[5] == 0 + or min(coefficients[2:4]) <= 0 + ): + raise ValueError("Unsupported ftheta calibration") + if list(coefficients[2:4]) != [width, height]: + warnings.warn( + f"{calibration.get('stream_id', 'ftheta')}: Ftheta coefficient dimensions {coefficients[2:4]} disagree with " + f"calibration dimensions {[width, height]}; retaining supplied calibration and assuming shared pixel coordinates.", + RuntimeWarning, + stacklevel=2, + ) + elif calibration["camera_model"] == "pinhole": + model = calibration["distortion_model"] + if not ( + (model == "none" and not coefficients) or (model == "opencv" and len(coefficients) in {4, 5, 8, 12, 14}) + ): + raise ValueError("Unsupported pinhole distortion model") + elif calibration["camera_model"] == "fisheye": + if calibration["distortion_model"] != "equidistant" or len(coefficients) != 4: + raise ValueError("Unsupported fisheye calibration; expected equidistant k1..k4") + else: + raise ValueError(f"Unsupported camera model: {calibration['camera_model']}") + + +class FinalizeCameraGeometry(Augmentor): + """Align metadata to retained physical views and source frames, without RNG draws. + + Invalid calibration on a selected camera skips the clip (returns ``None``) + so WebDataset workers drop it instead of crashing the job. Other per-clip + geometry ValueErrors (layout, FPS, clock origin) skip the same way. + Canonical sidecar readers still raise on malformed parquet. + """ + + require_camera_poses: bool + require_camera_calibration: bool + filter_invalid_camera_calibration: bool + fallback_invalid_camera_calibration: bool + calibration_token_stride: int + _pose_filter_counts: dict[tuple[int, int], tuple[int, int]] + _calibration_filter_counts: dict[tuple[int, int], tuple[int, int]] + _calibration_fallback_counts: tuple[int, int] + + def __init__( + self, + require_camera_poses: bool = False, + require_camera_calibration: bool = False, + filter_invalid_camera_calibration: bool = False, + calibration_token_stride: int = 32, + fallback_invalid_camera_calibration: bool = False, + ) -> None: + super().__init__([], None, None) + if ( + filter_invalid_camera_calibration or fallback_invalid_camera_calibration + ) and not require_camera_calibration: + raise ValueError("Calibration filtering or fallback requires require_camera_calibration=True") + if filter_invalid_camera_calibration and fallback_invalid_camera_calibration: + raise ValueError("Invalid calibration is either filtered or falls back, not both") + if calibration_token_stride <= 0: + raise ValueError("Calibration token stride must be positive") + self.require_camera_poses = require_camera_poses + self.require_camera_calibration = require_camera_calibration + self.filter_invalid_camera_calibration = filter_invalid_camera_calibration + self.fallback_invalid_camera_calibration = fallback_invalid_camera_calibration + # The caller selects the grid used by its geometry adapter. + self.calibration_token_stride = calibration_token_stride + self._pose_filter_counts = {} + self._calibration_filter_counts = {} + self._calibration_fallback_counts = (0, 0) + + def __call__(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: + if self.fallback_invalid_camera_calibration: + return self._finalize_or_fall_back(data_dict) + if not self.filter_invalid_camera_calibration: + try: + return self._finalize(data_dict) + except ValueError as error: + log.warning( + f"Skipping camera-geometry clip {data_dict.get('__key__')!r}: {error}", + rank0_only=False, + ) + return None + video = data_dict["video"] + resolution = tuple((video[-1] if isinstance(video, list) else video).shape[-2:]) + checked, dropped = self._calibration_filter_counts.get(resolution, (0, 0)) + checked += 1 + worker = torch.utils.data.get_worker_info() + context = ( + f"CameraCalibrationFilter worker={worker.id if worker is not None else 'main'} " + f"pid={os.getpid()} resolution={resolution[0]}x{resolution[1]}" + ) + try: + result = self._finalize(data_dict) + except CameraCalibrationError as error: + dropped += 1 + result = None + # Emit every rejection so the persisted rank logs form a complete rejection ledger. + log.warning( + f"{context} checked={checked} calibration_drops={dropped}: " + f"Dropping sample {data_dict.get('__key__')!r}: {error}", + rank0_only=False, + ) + self._calibration_filter_counts[resolution] = (checked, dropped) + if result is not None: + # Pass cumulative worker counters through accepted rows. The training monitor + # de-duplicates repeated CP-window rows and reports rejections separately from fallback. + result["camera_geometry"]["provenance"]["calibration_filter_stats"] = { + "worker": f"{os.getpid()}:{resolution[0]}x{resolution[1]}", + "checked": checked, + "rejected": dropped, + } + if checked == 1 or checked % 100 == 0: + log.info( + f"{context} checked={checked} calibration_drops={dropped} drop_pct={100.0 * dropped / checked:.2f}", + rank0_only=False, + ) + return result + + def _finalize_or_fall_back(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: + """Keep a sample whose geometry is absent or invalid, with ``camera_geometry=None``. + + A row the source loaded no geometry for (no source, or a ``None`` one) passes through + without geometry. Any other sample goes through the same checks as the filter, and one + failing them keeps its media and captions but carries no geometry, so a geometry- + conditioned model takes its fallback for it instead of losing the sample. The pose + requirement, when set, still drops samples and logs its own drops. + """ + if SOURCE_CALIBRATION_ERROR_KEY not in data_dict and data_dict.get(SOURCE_GEOMETRY_KEY) is None: + data_dict.pop(SOURCE_GEOMETRY_KEY, None) + data_dict.pop(SOURCE_FPS_KEY, None) + data_dict.pop(SOURCE_RIG_JSON_KEY, None) + data_dict["camera_geometry"] = None + return data_dict + checked, fell_back = self._calibration_fallback_counts + checked += 1 + try: + result = self._finalize(data_dict) + except CameraCalibrationError as error: + fell_back += 1 + for key in (SOURCE_CALIBRATION_ERROR_KEY, SOURCE_GEOMETRY_KEY, SOURCE_FPS_KEY, SOURCE_RIG_JSON_KEY): + data_dict.pop(key, None) + data_dict["camera_geometry"] = None + result = data_dict + if fell_back <= 3 or fell_back % 100 == 0: + log.warning( + f"CameraCalibrationFallback pid={os.getpid()} checked={checked} fallbacks={fell_back}: " + f"sample {data_dict.get('__key__')!r} keeps no camera geometry: {error}", + rank0_only=False, + ) + self._calibration_fallback_counts = (checked, fell_back) + if checked == 1 or checked % 100 == 0: + log.info( + f"CameraCalibrationFallback pid={os.getpid()} checked={checked} fallbacks={fell_back} " + f"fallback_pct={100.0 * fell_back / checked:.2f}", + rank0_only=False, + ) + return result + + def _finalize(self, data_dict: dict[str, Any]) -> dict[str, Any] | None: + parse_error = data_dict.pop(SOURCE_CALIBRATION_ERROR_KEY, None) + if parse_error is not None: + raise CameraCalibrationError(parse_error) + sources: dict[str, CameraGeometrySource] = data_dict.pop(SOURCE_GEOMETRY_KEY) + source_fps: list[float] = data_dict.pop(SOURCE_FPS_KEY) + rig_json = data_dict.pop(SOURCE_RIG_JSON_KEY, None) + cameras = list(data_dict["camera_keys_selection"]) + indices = data_dict["frame_indices"].tolist() + views, frames = len(cameras), len(indices) + video = data_dict["video"] + clips = video if isinstance(video, list) else [video] + if ( + not views + or frames != int(data_dict["num_video_frames_per_view"]) + or any( + clip.ndim != 4 or clip.shape[1] != views * frames or clip.shape[-2:] != clips[-1].shape[-2:] + for clip in clips + ) + ): + raise ValueError("Camera geometry does not match the final RGB layout") + height, width = clips[-1].shape[-2:] + original_hw = data_dict["original_hw"].tolist() + if ( + len(original_hw) != views + or any(len(hw) != 2 or min(hw) <= 0 for hw in original_hw) + or len(source_fps) != views + or any(fps <= 0 or not math.isfinite(fps) for fps in source_fps) + ): + raise ValueError("Invalid camera source dimensions or FPS") + view_crops = data_dict.get(VIEW_CROP_KEY) + crops = [(1.0, 0.5, 0.5)] * views if view_crops is None else [tuple(crop) for crop in view_crops.tolist()] + if len(crops) != views or any(len(crop) != 3 for crop in crops): + raise ValueError("View crops do not match the camera views") + poses = torch.eye(4, dtype=torch.float64).repeat(views, frames, 1, 1) # [V,T,4,4] + pose_valid = torch.zeros((views, frames), dtype=torch.bool) # [V,T] + affine = torch.eye(3, dtype=torch.float64).repeat(views, 1, 1) # [V,3,3] + times = torch.zeros((views, frames), dtype=torch.int64) # [V,T] + time_valid = torch.zeros((views, frames), dtype=torch.bool) # [V,T] + native_times = torch.zeros_like(times) # [V,T] + native_valid = torch.zeros_like(time_valid) # [V,T] + calibration: list[dict[str, Any] | None] = [] + time_source: list[list[str]] = [] + provenance: list[dict[str, Any]] = [] + origins = { + sources[camera].provenance.get("time_origin", "sample_start") + for camera in cameras + if camera in sources and sources[camera].times_ns is not None + } + if len(origins) > 1: + raise ValueError("Camera timestamp sidecars declare incompatible clock origins") + clock_origin = next(iter(origins), "sample_start") + for view, camera in enumerate(cameras): + source = sources.get(camera, CameraGeometrySource()) + cal = deepcopy(source.calibration) + if cal is None and rig_json is not None: + raise CameraCalibrationError(f"missing calibration for {camera} with rig_json present") + calibration.append(cal) + provenance.append(dict(source.provenance)) + try: + rig = source.provenance.get("rig_geometry") + if "rig_geometry" in source.provenance and (rig is None or camera not in rig["camera_to_rig"]): + raise ValueError(f"{camera}: missing calibrated rig extrinsic") + if self.require_camera_calibration and cal is None: + raise ValueError(f"{camera}: ray attention requires native calibration") + if cal is not None: + _validate_calibration(cal) + media_from_calibration = source.provenance.get("media_from_calibration") + if ( + self.require_camera_calibration + and media_from_calibration is None + and source.provenance.get("image_export") == MADS_FULL_FRAME_EXPORT + ): + media_from_calibration = mads_media_from_calibration( + cal["distortion_coefficients"], media_hw=tuple(original_hw[view]) + ) + provenance[-1]["media_from_calibration"] = media_from_calibration + if self.require_camera_calibration: + if cal["camera_model"] not in SUPPORTED_RAY_CAMERA_MODELS: + raise ValueError( + f"{camera}: ray attention cannot unproject camera model {cal['camera_model']!r}" + ) + if media_from_calibration is None and [cal["height_px"], cal["width_px"]] != original_hw[view]: + raise ValueError(f"{camera}: unresolved calibration-to-media mapping") + if [cal["height_px"], cal["width_px"]] != original_hw[view] and media_from_calibration is None: + warnings.warn( + f"{camera}: calibration dimensions {[cal['height_px'], cal['width_px']]} differ from " + f"media dimensions {original_hw[view]} (H, W); retaining supplied calibration and assuming " + "shared pixel coordinates for the media resize.", + RuntimeWarning, + stacklevel=2, + ) + zoom, crop_top, crop_left = crops[view] + affine[view] = resize_and_crop_affine( + tuple(original_hw[view]), (height, width), zoom=zoom, crop_center=(crop_top, crop_left) + ) # [3,3] + if self.require_camera_calibration and media_from_calibration is not None: + media_affine = torch.as_tensor(media_from_calibration, dtype=torch.float64) # [3,3] + if media_affine.shape != (3, 3) or not torch.isfinite(media_affine).all(): + raise ValueError(f"{camera}: invalid calibration-to-media mapping") + affine[view] = affine[view] @ media_affine # [3,3] + if self.filter_invalid_camera_calibration or self.fallback_invalid_camera_calibration: + validate_token_geometry( + cal, + affine[view], + image_hw=(height, width), + stride=self.calibration_token_stride, + ) + except ValueError as error: + raise CameraCalibrationError( + f"{camera}: {error}; sample={source.provenance.get('sample_key', data_dict.get('__key__'))}, " + f"calibration_source={source.provenance.get('ftheta_intrinsic')}, " + f"media_hw={original_hw[view]}" + ) from error + labels: list[str] = [] + for frame, index in enumerate(indices): + if index in source.poses: + poses[view, frame] = torch.tensor(source.poses[index], dtype=torch.float64) # [4,4] + pose_valid[view, frame] = True # [] + if source.times_ns is None and clock_origin == "sample_start": + times[view, frame] = round(index * 1e9 / source_fps[view]) # [] + time_valid[view, frame] = True # [] + labels.append("nominal_fps") + elif source.times_ns is not None and index in source.times_ns: + times[view, frame] = source.times_ns[index] # [] + time_valid[view, frame] = True # [] + labels.append("sidecar") + else: + labels.append("missing") + if index in source.native_times_ns: + native_times[view, frame] = source.native_times_ns[index] # [] + native_valid[view, frame] = True # [] + time_source.append(labels) + if self.require_camera_poses: + missing_poses = views > 1 and not bool(pose_valid.all()) + # Worker-local cumulative counts, separated by output resolution. Checked + # includes single-camera samples, which bypass the pose requirement. + checked, dropped = self._pose_filter_counts.get((height, width), (0, 0)) + checked += 1 + dropped += int(missing_poses) + self._pose_filter_counts[(height, width)] = (checked, dropped) + log_summary = checked == 1 or checked % 100 == 0 + log_drop = missing_poses and (dropped <= 3 or dropped % 100 == 0) + if log_summary or log_drop: + worker = torch.utils.data.get_worker_info() + context = ( + f"CameraPoseFilter worker={worker.id if worker is not None else 'main'} " + f"pid={os.getpid()} resolution={height}x{width}" + ) + if log_summary: + log.info( + f"{context} checked={checked} missing_pose_drops={dropped} " + f"drop_pct={100.0 * dropped / checked:.2f}", + rank0_only=False, + ) + if log_drop: + missing = { + camera: [index for frame, index in enumerate(indices) if not pose_valid[view, frame]][:8] + for view, camera in enumerate(cameras) + if not pose_valid[view].all() + } + log.warning( + f"{context}: Dropping sample {data_dict.get('__key__')!r} from {data_dict.get('__url__')!r}: " + f"missing required camera poses (first 8 source frames per camera): {missing}", + rank0_only=False, + ) + if missing_poses: + return None + offset = int(times[time_valid].min()) if time_valid.any() else 0 # [] + times[time_valid] -= offset # [N_valid] + url = data_dict.get("__url__", "") + source_url = f"{url.root}/{url.path}" if hasattr(url, "root") else str(url) + data_dict["camera_geometry"] = CameraGeometry( + camera_keys=cameras, + frame_indices=torch.tensor(indices, dtype=torch.int64).repeat(views, 1), # [V,T] + camera_to_world=poses, + pose_valid=pose_valid, + calibration=calibration, + calibration_valid=torch.tensor([cal is not None for cal in calibration], dtype=torch.bool), # [V] + image_from_calibration=affine, + image_hw=torch.tensor([height, width], dtype=torch.int64).repeat(views, 1), # [V,2] + frame_times_ns=times, + time_valid=time_valid, + time_source=time_source, + native_timestamps_ns=native_times, + native_time_valid=native_valid, + provenance={ + "sample_key": str(data_dict["__key__"]), + "source_url": source_url, + "cameras": provenance, + "time_origin": clock_origin, + "time_offset_ns": offset, + "source_fps": source_fps, + "rig_json": rig_json, + }, + ) + return data_dict diff --git a/cosmos_framework/data/generator/multiview/camera_geometry_contract.py b/cosmos_framework/data/generator/multiview/camera_geometry_contract.py new file mode 100644 index 000000000..a67014bc5 --- /dev/null +++ b/cosmos_framework/data/generator/multiview/camera_geometry_contract.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Training-local reader contract for canonical camera sidecars v1.0.0 and v1.1.0. + +The schema definitions and pose helpers are copied from +``packages/cosmos-data/cosmos_data/multiview/{camera_sidecars,pose}.py`` at +``44167c4f26a``. Keep this bounded copy independent of ingestion dependencies. +Changes to the persisted format need an explicit reader compatibility update; +the canonical Parquet fixtures in ``test_data`` pin the supported v1 format. +""" + +import math +from typing import cast + +import pyarrow as pa + +CAMERA_SIDECAR_SCHEMA_VERSION = b"1.0.0" +CAMERA_CALIBRATION_SCHEMA_NAME = b"cosmos.multiview.camera_calibration" +CAMERA_POSE_SCHEMA_NAME = b"cosmos.multiview.camera_pose" +TIMESTAMPS_SCHEMA_NAME = b"cosmos.multiview.timestamps" + +CAMERA_CALIBRATION_SCHEMA = pa.schema( + [ + pa.field("stream_id", pa.string(), nullable=False), + pa.field("camera_model", pa.string(), nullable=False), + pa.field("width_px", pa.int64(), nullable=False), + pa.field("height_px", pa.int64(), nullable=False), + pa.field("fx_px", pa.float64(), nullable=False), + pa.field("fy_px", pa.float64(), nullable=False), + pa.field("cx_px", pa.float64(), nullable=False), + pa.field("cy_px", pa.float64(), nullable=False), + pa.field("distortion_model", pa.string(), nullable=False), + pa.field("distortion_coefficients", pa.list_(pa.float64()), nullable=False), + ], + metadata={ + b"schema_name": CAMERA_CALIBRATION_SCHEMA_NAME, + b"schema_version": CAMERA_SIDECAR_SCHEMA_VERSION, + b"image_axes": b"x_right_y_down", + b"length_unit": b"pixel", + }, +) + +POSE_SCHEMA = pa.schema( + [ + pa.field("stream_id", pa.string(), nullable=False), + pa.field("sample_index", pa.int64(), nullable=False), + pa.field("episode_time_ns", pa.int64(), nullable=False), + pa.field("position_x_m", pa.float64(), nullable=False), + pa.field("position_y_m", pa.float64(), nullable=False), + pa.field("position_z_m", pa.float64(), nullable=False), + pa.field("orientation_qw", pa.float64(), nullable=False), + pa.field("orientation_qx", pa.float64(), nullable=False), + pa.field("orientation_qy", pa.float64(), nullable=False), + pa.field("orientation_qz", pa.float64(), nullable=False), + ], + metadata={ + b"schema_name": CAMERA_POSE_SCHEMA_NAME, + b"schema_version": CAMERA_SIDECAR_SCHEMA_VERSION, + b"camera_axes": b"x_right_y_down_z_forward", + b"coordinate_system": b"right_handed", + b"quaternion_order": b"qw_qx_qy_qz", + b"transform_direction": b"camera_to_world", + b"translation_unit": b"meter", + }, +) + +TIMESTAMPS_SCHEMA = pa.schema( + [ + pa.field("stream_id", pa.string(), nullable=False), + pa.field("sample_index", pa.int64(), nullable=False), + pa.field("native_timestamp_ns", pa.int64(), nullable=False), + pa.field("episode_time_ns", pa.int64(), nullable=False), + ], + metadata={ + b"schema_name": TIMESTAMPS_SCHEMA_NAME, + b"schema_version": CAMERA_SIDECAR_SCHEMA_VERSION, + b"clock": b"episode_clock", + b"time_unit": b"nanosecond", + }, +) + +# Released v1.1.0 sidecars (e.g. Alpamayo AV). Calibration adds the media-to-native +# pixel scale of ``nvidia_ftheta_polynomial`` rows, whose six backward terms are in +# native sensor pixels while cx/cy/width/height are in media pixels. Timestamps keep +# the v1.0.0 fields and clock semantics. +CAMERA_SIDECAR_SCHEMA_VERSION_V1_1 = b"1.1.0" +CAMERA_CALIBRATION_SCHEMA_V1_1 = pa.schema( + [ + *(CAMERA_CALIBRATION_SCHEMA.field(name) for name in CAMERA_CALIBRATION_SCHEMA.names[:8]), + pa.field("pixel_to_native_scale_x", pa.float64(), nullable=False), + pa.field("pixel_to_native_scale_y", pa.float64(), nullable=False), + *(CAMERA_CALIBRATION_SCHEMA.field(name) for name in CAMERA_CALIBRATION_SCHEMA.names[8:]), + ], + metadata={**(CAMERA_CALIBRATION_SCHEMA.metadata or {}), b"schema_version": CAMERA_SIDECAR_SCHEMA_VERSION_V1_1}, +) +TIMESTAMPS_SCHEMA_V1_1 = TIMESTAMPS_SCHEMA.with_metadata( + {**(TIMESTAMPS_SCHEMA.metadata or {}), b"schema_version": CAMERA_SIDECAR_SCHEMA_VERSION_V1_1} +) + +QuaternionXyzw = tuple[float, float, float, float] +Rotation3 = tuple[tuple[float, float, float], tuple[float, float, float], tuple[float, float, float]] +Matrix4 = tuple[ + tuple[float, float, float, float], + tuple[float, float, float, float], + tuple[float, float, float, float], + tuple[float, float, float, float], +] + + +def normalize_quaternion_xyzw( + quaternion: QuaternionXyzw, + *, + label: str, + tolerance: float = 1e-4, +) -> QuaternionXyzw: + """Return a normalized ``xyzw`` quaternion after validating it is already unit length.""" + norm = math.sqrt(sum(value * value for value in quaternion)) + if not math.isfinite(norm) or norm <= 1e-12: + raise ValueError(f"{label} has an invalid quaternion") + if not math.isclose(norm, 1.0, rel_tol=tolerance, abs_tol=tolerance): + raise ValueError(f"{label} quaternion must be normalized") + return cast(QuaternionXyzw, tuple(value / norm for value in quaternion)) + + +def quaternion_xyzw_to_rotation_matrix(quaternion: QuaternionXyzw) -> Rotation3: + """Convert a unit ``xyzw`` quaternion to a 3x3 rotation matrix.""" + x, y, z, w = quaternion + return ( + (1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)), + (2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)), + (2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)), + ) + + +def validate_camera_to_world_matrix( + matrix: Matrix4, + *, + label: str, + tolerance: float = 1e-4, +) -> None: + """Validate that a homogeneous 4x4 matrix contains a proper camera-to-world rotation.""" + if any(not math.isfinite(value) for row in matrix for value in row): + raise ValueError(f"{label} values must be finite") + if any( + not math.isclose(value, expected, abs_tol=1e-6) for value, expected in zip(matrix[3], (0, 0, 0, 1), strict=True) + ): + raise ValueError(f"{label} must be homogeneous") + rotation = tuple(row[:3] for row in matrix[:3]) + for first_column in range(3): + for second_column in range(3): + product = sum(rotation[row][first_column] * rotation[row][second_column] for row in range(3)) + expected = 1.0 if first_column == second_column else 0.0 + if not math.isclose(product, expected, rel_tol=0.0, abs_tol=tolerance): + raise ValueError(f"{label} rotation must be orthonormal") + determinant = ( + rotation[0][0] * (rotation[1][1] * rotation[2][2] - rotation[1][2] * rotation[2][1]) + - rotation[0][1] * (rotation[1][0] * rotation[2][2] - rotation[1][2] * rotation[2][0]) + + rotation[0][2] * (rotation[1][0] * rotation[2][1] - rotation[1][1] * rotation[2][0]) + ) + if not math.isclose(determinant, 1.0, rel_tol=tolerance, abs_tol=tolerance): + raise ValueError(f"{label} rotation must have determinant +1") diff --git a/cosmos_framework/data/generator/multiview/camera_geometry_native.py b/cosmos_framework/data/generator/multiview/camera_geometry_native.py new file mode 100644 index 000000000..7839a4935 --- /dev/null +++ b/cosmos_framework/data/generator/multiview/camera_geometry_native.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Training-local readers for sparse MADS poses and native F-theta intrinsics. + +Copied from ``pipelines/sila/multiview/ingestion/av/camera_geometry.py`` at +``44167c4f26a``, retaining only the pose and F-theta readers used by WebDataset. +Poses are OpenCV camera-to-world matrices in meters. F-theta coefficients +calibrate the native distorted RGB, without substituting a pinhole projection. +""" + +from __future__ import annotations + +import io +import re +import tarfile +from dataclasses import dataclass +from pathlib import PurePosixPath +from typing import cast + +import numpy as np + +from cosmos_framework.data.generator.multiview.camera_geometry_contract import ( + Matrix4, + validate_camera_to_world_matrix, +) + +_POSE_MEMBER = re.compile(r"\.(?P\d+)\.pose\.camera_(?P.+)\.npy$") +_FTHETA_MEMBER = re.compile(r"\.ftheta_intrinsic\.camera_(?P.+)\.npy$") +_FTHETA_COEFFICIENT_COUNTS = (11, 14) + +# Native sidecars scale calibration uniformly, truncating the reported width. +# Raw rig dimensions were checked against all 1,034 calibrations in the 94-clip +# benchmark. Matched raw/exported frames verify full-frame, half-pixel resizing. +# Keep these profiles explicit: image dimensions alone cannot establish a crop. +MADS_FULL_FRAME_EXPORT = "mads_rdshq_full_frame_v1" +_MADS_RAW_WH_BY_CALIBRATION_WH = { + (1924, 1084): (3848, 2168), + (1920, 1080): (3840, 2160), + (1719, 1080): (1936, 1216), + (1350, 1080): (1920, 1536), +} + + +def mads_media_from_calibration( + coefficients: tuple[float, ...] | list[float], *, media_hw: tuple[int, int] +) -> list[list[float]]: + """Map native MADS calibration pixels into its full-frame exported video. + + This applies to frozen RDS-HQ calibration and full-frame exported media. + Calibration uses origin-zero uniform scaling from the raw rig; video resize + uses pixel centers. In particular, width 1719 represents 1719.473684... raw + pixels after uniform scaling, while the encoder emits 1720 columns. + """ + if len(coefficients) not in _FTHETA_COEFFICIENT_COUNTS: + raise ValueError("MADS image mapping requires native F-theta calibration") + calibration_wh = tuple(coefficients[2:4]) + if min(*calibration_wh, *media_hw) <= 0 or not all(np.isfinite(value) for value in calibration_wh): + raise ValueError("MADS image mapping requires positive finite dimensions") + raw_wh = _MADS_RAW_WH_BY_CALIBRATION_WH.get(calibration_wh) + if raw_wh is None: + # The ablation treats provider calibration as frozen. For other sensor + # windows, honor the requested full-frame resize without fitting a rig. + sx, sy = media_hw[1] / calibration_wh[0], media_hw[0] / calibration_wh[1] + return [[sx, 0.0, (sx - 1.0) / 2], [0.0, sy, (sy - 1.0) / 2], [0.0, 0.0, 1.0]] + raw_w, raw_h = raw_wh + calibration_h = coefficients[3] + scale = calibration_h / raw_h + height, width = media_hw + sx, sy = width / raw_w, height / raw_h + return [[sx / scale, 0.0, (sx - 1.0) / 2], [0.0, sy / scale, (sy - 1.0) / 2], [0.0, 0.0, 1.0]] + + +@dataclass(frozen=True, slots=True) +class CanonicalCameraCalibration: + """One static camera calibration in the Pose QA sidecar contract.""" + + camera_model: str + distortion_model: str + distortion_coefficients: tuple[float, ...] + intrinsics: tuple[float, float, float, float] + + +def parse_pose_entries(payload: bytes) -> dict[str, dict[int, Matrix4]]: + """Read camera-to-world matrices by frame ID, preserving gaps for consumers to mask.""" + arrays = _load_npy_tar(payload, pattern=_POSE_MEMBER, modality="pose") + poses_by_camera: dict[str, dict[int, Matrix4]] = {} + for member_path, match, array in arrays: + camera = match.group("camera") + frame_index = int(match.group("frame")) + matrix = _as_camera_to_world_matrix(array, label=f"pose {member_path}") + camera_frames = poses_by_camera.setdefault(camera, {}) + if frame_index in camera_frames: + raise ValueError(f"duplicate pose frame {frame_index} for camera {camera}") + camera_frames[frame_index] = matrix + return poses_by_camera + + +def parse_ftheta_intrinsic_tar(payload: bytes) -> dict[str, CanonicalCameraCalibration]: + """Return per-camera f-theta calibrations from one nested npy tar.""" + calibrations: dict[str, CanonicalCameraCalibration] = {} + for member_path, match, array in _load_npy_tar(payload, pattern=_FTHETA_MEMBER, modality="ftheta_intrinsic"): + camera = match.group("camera") + if camera in calibrations: + raise ValueError(f"duplicate ftheta intrinsic for camera {camera}") + calibrations[camera] = parse_ftheta_intrinsic(array, label=member_path) + return calibrations + + +def parse_ftheta_intrinsic(array: np.ndarray, *, label: str) -> CanonicalCameraCalibration: + """Interpret a native MADS f-theta vector, retaining optional affine lens terms. + + Provider layout is ``[cx, cy, width, height, a0, a1, a2, a3, a4, a5, flag]``. + The flag distinguishes radius-to-angle from angle-to-radius coefficients. + The legacy equivalent focal length is retained for metadata compatibility; ray attention + uses the complete native polynomial and optional final ``[c,d,e]`` terms. + Pose QA cannot score ``camera_model=ftheta``; those cameras remain attached + because that is the native MADS image-formation model. MADS does not ingest + pinhole as a substitute on the same still-distorted frames. + """ + if array.ndim != 1 or array.size not in _FTHETA_COEFFICIENT_COUNTS: + raise ValueError(f"{label} ftheta intrinsic must have 11 or 14 coefficients") + values = tuple(float(value) for value in array.reshape(-1)) + if not all(np.isfinite(value) for value in values): + raise ValueError(f"{label} ftheta intrinsic must contain finite values") + cx_px, cy_px, _width, _height, _a0, linear_rad_per_px, *_rest = values + if linear_rad_per_px == 0.0: + raise ValueError(f"{label} ftheta linear coefficient must be non-zero") + focal_px = abs(1.0 / linear_rad_per_px) + return CanonicalCameraCalibration( + camera_model="ftheta", + distortion_model="ftheta", + distortion_coefficients=values, + intrinsics=(focal_px, focal_px, cx_px, cy_px), + ) + + +def _load_npy_tar( + payload: bytes, + *, + pattern: re.Pattern[str], + modality: str, +) -> tuple[tuple[str, re.Match[str], np.ndarray], ...]: + members = _npy_tar_members(payload) + if not members: + raise ValueError(f"{modality} sidecar tar contains no .npy members") + matched: list[tuple[str, re.Match[str], np.ndarray]] = [] + for member_path, array in members: + match = pattern.search(PurePosixPath(member_path).name) + if match is None: + continue + matched.append((member_path, match, array)) + if not matched: + raise ValueError(f"{modality} sidecar tar has no camera {modality} npy members") + return tuple(matched) + + +def _npy_tar_members(payload: bytes) -> tuple[tuple[str, np.ndarray], ...]: + members: list[tuple[str, np.ndarray]] = [] + try: + with tarfile.open(fileobj=io.BytesIO(payload), mode="r:*") as archive: + for member in archive.getmembers(): + if not member.isfile() or not PurePosixPath(member.name).name.endswith(".npy"): + continue + extracted = archive.extractfile(member) + if extracted is None: + raise ValueError(f"tar member is not a regular file: {member.name}") + try: + array = np.load(io.BytesIO(extracted.read()), allow_pickle=False) + except (ValueError, OSError) as exc: + raise ValueError(f"could not load npy member {member.name}") from exc + members.append((member.name, np.asarray(array))) + except tarfile.TarError as exc: + raise ValueError("camera geometry sidecar must be a nested tar of .npy members") from exc + return tuple(members) + + +def _as_camera_to_world_matrix(array: np.ndarray, *, label: str) -> Matrix4: + if array.shape != (4, 4): + raise ValueError(f"{label} must have shape (4, 4), got {array.shape}") + matrix = cast(Matrix4, tuple(tuple(float(value) for value in row) for row in array)) + validate_camera_to_world_matrix(matrix, label=label) + return matrix diff --git a/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py index d7598a039..4ab7391b7 100644 --- a/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py +++ b/cosmos_framework/data/generator/reasoner/video_decoder_qwen.py @@ -9,6 +9,7 @@ """ import hashlib +import math import random import re from concurrent.futures import ThreadPoolExecutor @@ -215,9 +216,9 @@ def _video_decoder_qwen_func( if video_fps < 1: raise ValueError("Video fps lower than 1, skipping") - if video_fps < min_fps_thres: + if video_fps < min_fps_thres and not math.isclose(video_fps, min_fps_thres, rel_tol=1e-12): raise ValueError(f"Video fps {video_fps} lower than {min_fps_thres}, skipping") - if video_fps > max_fps_thres: + if video_fps > max_fps_thres and not math.isclose(video_fps, max_fps_thres, rel_tol=1e-12): raise ValueError(f"Video fps {video_fps} higher than {max_fps_thres}, skipping") configured_max_video_token_length = max_video_token_length diff --git a/cosmos_framework/data/generator/sequence_packing/autoregressive.py b/cosmos_framework/data/generator/sequence_packing/autoregressive.py index c5cd367d2..e574f1fc7 100644 --- a/cosmos_framework/data/generator/sequence_packing/autoregressive.py +++ b/cosmos_framework/data/generator/sequence_packing/autoregressive.py @@ -1,13 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Autoregressive sequence packing for framewise and chunkwise AR generation.""" +"""Interactive temporal-causal packing for autoregressive generation and replay training.""" from collections.abc import Sequence +from dataclasses import replace from typing import Any, cast import torch +from cosmos_framework.configs.base.defaults.model_config import MultiviewActionConditioningConfig from cosmos_framework.data.generator.augmentors.text_tokenizer import TEXT_SYSTEM_PROMPT_KEY from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.data.generator.sequence_packing import ( @@ -15,6 +17,132 @@ SequencePlan, pack_input_sequence, ) +from cosmos_framework.data.generator.sequence_packing.packers import uses_single_timestep +from cosmos_framework.data.generator.sequence_packing.sequence import PackedSequenceBuilder +from cosmos_framework.data.generator.sequence_packing.temporal_causal import pack_supertokens_temporal_causal + + +def pack_multiview_action_temporal_causal( + plans: list[SequencePlan], + text_tokens: list[list[int]], + data: GenerationDataClean, + timesteps: torch.Tensor, + *, + config: MultiviewActionConditioningConfig, + special_tokens: dict[str, int], + patch_spatial: int, + temporal_factor: int, + temporal_margin: int, + reset_spatial_ids: bool, + enable_fps_modulation: bool, + base_fps: float, +) -> PackedSequence: + """Pack one synchronized robot clip, reusing the single-view causal supertokens. + + Camera-major views may have different spatial grids: head uses the 480p + bucket, wrists use 256p. Each view has [null, V0], [a1..a4, V1], ...; + all views reset to the same mRoPE clock. Actions remain clean inputs. + """ + if data.batch_size != 1 or len(plans) != 1 or len(text_tokens) != 1: + raise ValueError("Robot multiview temporal-causal packing requires one logical clip per batch.") + plan = plans[0] + if ( + not plan.has_text + or not plan.has_vision + or not plan.has_action + or plan.has_sound + or plan.has_lidar + or plan.has_radar + ): + raise ValueError("Robot multiview temporal-causal packing supports text, RGB, and conditioned actions only.") + if timesteps.is_cuda: + raise ValueError("Temporal-causal packing timesteps must remain on CPU.") + assert data.x0_tokens_vision is not None and data.x0_tokens_action is not None + num_views = (data.num_vision_items_per_sample or [1])[0] + if num_views not in (1, len(config.view_codes)) or len(data.x0_tokens_vision) != num_views: + raise ValueError("Robot temporal-causal packing expects a single head view or the configured synchronized rig.") + latent_t = data.x0_tokens_vision[0].shape[2] + action = data.x0_tokens_action[0] + rows_per_view = (latent_t - 1) * temporal_factor + if ( + any(view.shape[2] != latent_t for view in data.x0_tokens_vision) + or action.ndim != 2 + or action.shape[0] != num_views * rows_per_view + ): + raise ValueError("Every view must retain all frame-rate actions aligned to T+1 observations.") + vision_fps = float(data.fps_vision[0]) if data.fps_vision is not None else base_fps + action_fps = float(data.fps_action[0]) if data.fps_action is not None else vision_fps + timestep_row = timesteps[0] + input_timestep: float | torch.Tensor = ( + float(timestep_row.item()) if timestep_row.numel() == 1 else timestep_row.flatten() + ) + builder = PackedSequenceBuilder(uses_single_timestep=uses_single_timestep(timesteps)) + builder._mrope_reset_spatial = reset_spatial_ids + builder.begin_sample(0) + builder.pack_text_tokens( + text_tokens[0], + special_tokens, + has_generation=True, + use_float_positions=enable_fps_modulation, + ) + builder.advance_mrope_temporal_offset(temporal_margin) + temporal_offset = builder.mrope_temporal_offset + generation_start = builder.current_seq_index + # Never pass OmegaConf containers to Torch: their probing can retain caller + # frames (and training tensors) until cyclic GC. + codes = action.new_tensor([list(code) for code in config.view_codes]) + domains: list[torch.Tensor] | None = [] if data.action_domain_id is not None else None + for view_id, vision in enumerate(data.x0_tokens_vision): + builder.set_mrope_temporal_offset(temporal_offset) + real_actions = action[view_id * rows_per_view : (view_id + 1) * rows_per_view].clone() + start = config.view_code_start + real_actions[:, start : start + codes.shape[1]] = codes[view_id] + _, null_prefix = pack_supertokens_temporal_causal( + builder, + vision, + real_actions, + plan.condition_frame_indexes_vision, + input_timestep, + patch_spatial, + temporal_factor, + action.shape[1], + vision_fps=vision_fps, + action_fps=action_fps, + enable_fps_modulation=enable_fps_modulation, + base_fps=base_fps, + ) + assert null_prefix and builder.vision is not None and builder.action is not None + # The null action's physical channels stay zero, but its fixed code still + # identifies the view at V0. The head code is exactly zero, as in single view. + builder.action.tokens[-1][:temporal_factor, start : start + codes.shape[1]] = codes[view_id] + builder.vision.seconds_per_frame.append(temporal_factor / vision_fps) + builder.action.seconds_per_frame.append(1.0 / action_fps) + if domains is not None: + assert data.action_domain_id is not None + domain = data.action_domain_id[0].reshape(-1) + if domain.numel() not in (1, action.shape[0]): + raise ValueError("Action domains must be scalar or have one ID per real frame-rate action.") + domains.append( + domain if domain.numel() == 1 else domain[view_id * rows_per_view : (view_id + 1) * rows_per_view] + ) + builder.null_action_supertokens = True + builder.num_action_tokens_per_supertoken = temporal_factor + builder.finish_sample(builder.current_seq_index - generation_start, builder.current_seq_index) + packed = builder.finalize( + replace( + data, + num_vision_items_per_sample=[num_views], + num_views_per_vision_item=[1] * num_views, + num_views_per_action_item=[1] * num_views, + action_domain_id=domains, + raw_action_dim=data.raw_action_dim * num_views if data.raw_action_dim is not None else None, + action_valid_mask=data.action_valid_mask * num_views if data.action_valid_mask is not None else None, + ) + ) + # Dynamic replay metadata is copied by make_teacher_forcing_clean_pack and + # keeps the generic interactive multiview helpers opt-in for this layout. + setattr(packed, "multiview_action_conditioning", True) + return packed def resolve_text_system_prompt(data_batch: dict[str, Any]) -> Any: diff --git a/cosmos_framework/model/generator/mot/causal_attention.py b/cosmos_framework/model/generator/mot/causal_attention.py index 56ec7797f..c0b126547 100644 --- a/cosmos_framework/model/generator/mot/causal_attention.py +++ b/cosmos_framework/model/generator/mot/causal_attention.py @@ -42,6 +42,10 @@ replay_maskless_attention, ) from cosmos_framework.model.generator.mot.merge_attention import merge_attentions_ac_safe +from cosmos_framework.model.generator.mot.multiview_action_attention import ( + MultiviewActionAttentionPlan, + same_instant_multiview_rgb_attention, +) from cosmos_framework.model.generator.utils.kv_cache import ( ARMemoryValue, JointARMemoryValue, @@ -49,6 +53,7 @@ MultiviewARMemoryValue, TFNoisyMemoryValue, TFReplayCleanMemoryValue, + zero_null_action_values, ) from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( RollingPromptLayout, @@ -1180,6 +1185,53 @@ def _clean_pass_target_attention_components( raise ValueError(f"Unknown clean-pass causality: {policy.clean_pass_causality!r}") +def _teacher_forcing_view_attention_components( + query: torch.Tensor, # [T*S,H,D] + key: torch.Tensor, # [T*S,H_kv,D] + value: torch.Tensor, # [T*S,H_kv,D] + memory_value: TFReplayCleanMemoryValue | TFNoisyMemoryValue, + *, + token_shape: tuple[int, int, int], + num_action_tokens: int, + cached_clean_gen_k: torch.Tensor | None = None, # [1,T*S,H_kv,D] + cached_clean_gen_v: torch.Tensor | None = None, # [1,T*S,H_kv,D] +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Run the existing three-way TF receptive fields for one camera. + + Single-view replay calls this once. Multiview replay calls the same helper + independently for each camera before adding its one RGB-only varlen pass. + """ + frames, height, width = token_shape + tokens_per_frame = num_action_tokens + height * width + item_len = frames * tokens_per_frame + if query.shape[0] != item_len or key.shape[0] != item_len or value.shape[0] != item_len: + raise ValueError( + "Teacher-forcing view length does not match its token shape: " + f"shape={token_shape}, actions={num_action_tokens}, expected={item_len}, " + f"got Q/K/V={query.shape[0]}/{key.shape[0]}/{value.shape[0]}." + ) + query_2d = query.reshape(1, frames, tokens_per_frame, query.shape[-2], query.shape[-1]) + key_2d = key.reshape(1, frames, tokens_per_frame, key.shape[-2], key.shape[-1]) + value_2d = value.reshape(1, frames, tokens_per_frame, value.shape[-2], value.shape[-1]) + if isinstance(memory_value, TFNoisyMemoryValue): + return teacher_forcing_gen_attention( + query_2d, + key_2d, + value_2d, + memory_value, + memory_value.frames_per_chunk, + cached_clean_gen_k=cached_clean_gen_k, + cached_clean_gen_v=cached_clean_gen_v, + ) + return _clean_pass_target_attention_components( + query_2d, + key_2d, + value_2d, + memory_value.teacher_forcing_replay_policy, + memory_value.frames_per_chunk, + ) + + def teacher_forcing_transfer_attention( control_q: torch.Tensor, # [T*S,H,D] control_k: torch.Tensor, # [T*S,H_kv,D] @@ -1428,6 +1480,7 @@ def three_way_attention_with_memory( memory_value: KVTrainMemoryValue, attention_meta: SplitInfo | None = None, packed_key_states_normalized: SequencePack | None = None, + multiview_action_attention_plan: MultiviewActionAttentionPlan | None = None, ) -> SequencePack: """Branchless three-way attention with cached K/V memory. @@ -1475,15 +1528,24 @@ def three_way_attention_with_memory( video_out, text_out = teacher_forcing_target_only_attention(video_q, video_k, video_v, memory_value) return from_mode_splits(text_out, video_out, packed_query_states) + uses_multiview_action_attention = multiview_action_attention_plan is not None if attention_meta is not None and attention_meta.null_action_supertokens: - video_v = video_v.clone() - starts = video_pack_q_offsets[:-1].long() - null_positions = (starts.unsqueeze(1) + torch.arange(num_action_tokens, device=starts.device)).reshape(-1) - video_v[null_positions] = 0 + if uses_multiview_action_attention: + video_v = zero_null_action_values( + video_v.unsqueeze(0), + memory_value.vision_token_shapes, + num_action_tokens, + True, + ).squeeze(0) + else: + video_v = video_v.clone() + starts = video_pack_q_offsets[:-1].long() + null_positions = (starts.unsqueeze(1) + torch.arange(num_action_tokens, device=starts.device)).reshape(-1) + video_v[null_positions] = 0 # --- video self-attention: temporal-causal via multi_dimensional_attention --- vision_token_shapes = memory_value.vision_token_shapes - is_transfer = len(vision_token_shapes) == 2 + is_transfer = len(vision_token_shapes) == 2 and not uses_multiview_action_attention if is_transfer: if not isinstance(memory_value, (TFReplayCleanMemoryValue, TFNoisyMemoryValue)): raise TypeError("Two-item temporal-causal transfer is supported only by replay teacher forcing.") @@ -1496,7 +1558,11 @@ def three_way_attention_with_memory( T, H_p, W_p = vision_token_shapes[0] S_super = num_action_tokens + H_p * W_p item_len = T * S_super - video_len = item_len * (2 if is_transfer else 1) + video_len = ( + sum(frames * (num_action_tokens + height * width) for frames, height, width in vision_token_shapes) + if uses_multiview_action_attention + else item_len * (2 if is_transfer else 1) + ) num_heads = video_q.shape[1] num_kv_heads = video_k.shape[1] head_dim = video_q.shape[2] @@ -1507,6 +1573,7 @@ def three_way_attention_with_memory( # requires an exact (T, S_super) reshape. # The padding will be added back after merge_attentions. padded_video_len = video_q.shape[0] + padded_video_q = video_q video_q = video_q[:video_len] video_k = video_k[:video_len] video_v = video_v[:video_len] @@ -1514,7 +1581,61 @@ def three_way_attention_with_memory( # Naming: ``_sa`` = self-attention, ``_ca`` = cross-attention, ``_lse`` = log-sum-exp. attn_outputs: list[torch.Tensor] lse_outputs: list[torch.Tensor] - if is_transfer: + if uses_multiview_action_attention: + assert multiview_action_attention_plan is not None + if not isinstance(memory_value, (TFReplayCleanMemoryValue, TFNoisyMemoryValue)): + raise TypeError("Hybrid multiview replay requires teacher-forcing replay memory.") + view_components: list[tuple[torch.Tensor, torch.Tensor]] = [] + offset = 0 + for token_shape in vision_token_shapes: + frames, height, width = token_shape + spatial_tokens = num_action_tokens + height * width + view_len = frames * spatial_tokens + components = _teacher_forcing_view_attention_components( + video_q[offset : offset + view_len], + video_k[offset : offset + view_len], + video_v[offset : offset + view_len], + memory_value, + token_shape=token_shape, + num_action_tokens=num_action_tokens, + cached_clean_gen_k=( + memory_value.cached_clean_gen_k[:, offset : offset + view_len] + if isinstance(memory_value, TFNoisyMemoryValue) + else None + ), + cached_clean_gen_v=( + memory_value.cached_clean_gen_v[:, offset : offset + view_len] + if isinstance(memory_value, TFNoisyMemoryValue) + else None + ), + ) + view_components.append(_merge_transfer_item_components(components, item_len=view_len)) + offset += view_len + same_view_out, same_view_lse = view_components[0] + for view_out, view_lse in view_components[1:]: + same_view_out, same_view_lse = ConcatenateAttentionsBridge.apply( + same_view_out, + same_view_lse, + view_out, + view_lse, + ) + attn_outputs = [same_view_out] + lse_outputs = [same_view_lse] + + # The only additional keys are live RGB at the same instant. Actions + # stay view-local, and V=1 retains precisely the original three-way path. + if multiview_action_attention_plan.rgb_indexes.numel(): + # Normalized K is reserved for GEN-to-UND/text attention below. + # RGB self-attention must use the same raw GEN K as the per-view + # three-way component and Stage-0 decomposed attention. + gen_k, _ = get_full_only_seq(packed_key_states) # [N_gen_padded,H_kv,D] + gen_v, _ = get_full_only_seq(packed_value_states) # [N_gen_padded,H_kv,D] + multiview_rgb_out, multiview_rgb_lse = same_instant_multiview_rgb_attention( + padded_video_q, gen_k, gen_v, multiview_action_attention_plan + ) + attn_outputs.append(multiview_rgb_out) + lse_outputs.append(multiview_rgb_lse) + elif is_transfer: assert isinstance(memory_value, (TFReplayCleanMemoryValue, TFNoisyMemoryValue)) transfer_components = teacher_forcing_transfer_attention( video_q[:item_len], # [T*S,H,D] @@ -1529,33 +1650,23 @@ def three_way_attention_with_memory( attn_outputs = [out for out, _lse in transfer_components] # each [1,2*T*S,H,D] lse_outputs = [lse for _out, lse in transfer_components] # each [1,2*T*S,H] else: - # Reshape to expose temporal dimension for the causal mask. - video_q_2d = video_q.reshape(1, T, S_super, num_heads, head_dim) # [1,T,S_super,H,D] - video_k_2d = video_k.reshape(1, T, S_super, num_kv_heads, head_dim) # [1,T,S_super,H_kv,D] - video_v_2d = video_v.reshape(1, T, S_super, num_kv_heads, head_dim) # [1,T,S_super,H_kv,D] - video_components: list[tuple[torch.Tensor, torch.Tensor]] - if isinstance(memory_value, TFNoisyMemoryValue): - # Teacher forcing: two merge components framewise, four chunkwise - # (frames_per_chunk > 1, chunk partition [1, C, C, ...]). - video_components = teacher_forcing_gen_attention( - video_q_2d, - video_k_2d, - video_v_2d, + if isinstance(memory_value, (TFReplayCleanMemoryValue, TFNoisyMemoryValue)): + # This exact helper is applied once for single-view and once per + # camera for multiview; only the latter adds an RGB varlen component. + video_components = _teacher_forcing_view_attention_components( + video_q, + video_k, + video_v, memory_value, - memory_value.frames_per_chunk, - ) - elif isinstance(memory_value, TFReplayCleanMemoryValue): - # Single-view three-way replay uses only clean-pass causality from - # the unified policy; multiview scope does not apply to this layout. - video_components = _clean_pass_target_attention_components( - video_q_2d, - video_k_2d, - video_v_2d, - memory_value.teacher_forcing_replay_policy, - memory_value.frames_per_chunk, + token_shape=vision_token_shapes[0], + num_action_tokens=num_action_tokens, ) else: + # Reshape to expose temporal dimension for the causal mask. + video_q_2d = video_q.reshape(1, T, S_super, num_heads, head_dim) # [1,T,S_super,H,D] + video_k_2d = video_k.reshape(1, T, S_super, num_kv_heads, head_dim) # [1,T,S_super,H_kv,D] + video_v_2d = video_v.reshape(1, T, S_super, num_kv_heads, head_dim) # [1,T,S_super,H_kv,D] # --- Standard temporal-causal self-attention --- video_sa, video_sa_lse = multi_dimensional_attention( # [1,T,S_super,H,D], [1,T,S_super,H] video_q_2d, @@ -2086,6 +2197,12 @@ def dispatch_attention_with_memory( - ``ARMemoryValue`` with ``frame_idx == 0`` → interactive no-memory dispatch - ``None`` without a replay plan → interactive no-memory dispatch """ + multiview_action_attention_plan = getattr(attention_mask, "multiview_action_attention_plan", None) + if multiview_action_attention_plan is not None: + if not isinstance(multiview_action_attention_plan, MultiviewActionAttentionPlan) or not isinstance( + memory_value, (TFReplayCleanMemoryValue, TFNoisyMemoryValue) + ): + raise TypeError("Multiview action attention requires a varlen plan and teacher-forcing memory.") replay_maskless_plan = getattr(attention_mask, "replay_maskless_plan", None) if replay_maskless_plan is not None and not isinstance(replay_maskless_plan, ReplayMasklessPlan): raise TypeError("Maskless attention requires a ReplayMasklessPlan.") @@ -2132,6 +2249,7 @@ def dispatch_attention_with_memory( memory_value=memory_value, attention_meta=attention_meta, packed_key_states_normalized=packed_key_states_normalized, + multiview_action_attention_plan=multiview_action_attention_plan, ) return output, None # Keep the post-saturation static-compile predicate first so Dynamo can diff --git a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py index e98c40a82..5682b11c4 100644 --- a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py @@ -36,6 +36,9 @@ build_teacher_forcing_multiview_flex_metadata, ) from cosmos_framework.model.generator.mot.maskless_attention import build_replay_maskless_plan +from cosmos_framework.model.generator.mot.multiview_action_attention import ( + build_multiview_action_attention_plan, +) from cosmos_framework.model.generator.utils.rolling_kv.rolling_replay import ( RollingReplayRequest, build_rolling_replay_metadata, @@ -93,6 +96,7 @@ class InteractiveCosmos3VFMNetwork(Cosmos3VFMNetwork): teacher_forcing_replay_policy: TeacherForcingReplayPolicyConfig | None teacher_forcing_frames_per_chunk: int teacher_forcing_maskless: bool = False + teacher_forcing_multiview_threeway: bool = False def __init__(self, language_model: torch.nn.Module, config: Any) -> None: super().__init__(language_model=language_model, config=config) @@ -116,6 +120,16 @@ def _replace_teacher_forcing_mask( packed_seq = self._active_packed_seq if packed_seq is None: return None + if self.teacher_forcing_multiview_threeway: + if not getattr(packed_seq, "multiview_action_conditioning", False): + return None + assert packed_seq.vision is not None + attention_meta = kwargs["attention_mask"] + gen, _ = get_full_only_seq(args[0]) # [N_gen_local,H,D] + attention_meta.multiview_action_attention_plan = build_multiview_action_attention_plan( + packed_seq.vision.token_shapes, packed_seq.num_action_tokens_per_supertoken, gen.device + ) + return args, kwargs ar_metadata = getattr(packed_seq, "multiview_transfer_ar_metadata", None) rolling_request = getattr(packed_seq, "rolling_replay_request", None) teacher_forcing_pass = getattr(packed_seq, "teacher_forcing_pass", None) @@ -299,6 +313,8 @@ def forward( previous_packed_seq = self._active_packed_seq self._active_packed_seq = packed_seq try: + if self.teacher_forcing_multiview_threeway and getattr(packed_seq, "multiview_action_conditioning", False): + kwargs["bounded_cp_output_gather"] = True if ( self.teacher_forcing_maskless and video_temporal_causal is not False diff --git a/cosmos_framework/model/generator/mot/flex_attention_memory_test.py b/cosmos_framework/model/generator/mot/flex_attention_memory_test.py new file mode 100644 index 000000000..cfb1e27d2 --- /dev/null +++ b/cosmos_framework/model/generator/mot/flex_attention_memory_test.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Periodic GC must reclaim unreachable inputs from fresh Flex/Flash compilations.""" + +import gc +import time +import weakref +from collections.abc import Callable +from typing import Any + +import pytest +import torch +from torch.nn.attention.flex_attention import BlockMask, create_block_mask +from torch.nn.attention.flex_attention import flex_attention as torch_flex_attention + +from cosmos_framework.callbacks.manual_gc import ManualGarbageCollection +from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig +from cosmos_framework.model.generator.mot import flex_attention as attention_module +from cosmos_framework.model.generator.mot.flex_attention import FlexBackend, resolve_flex_backend +from cosmos_framework.model.generator.mot.parallelize_unified_mot import _apply_selective_ac + + +def _mask(length: int, backend: FlexBackend) -> BlockMask: + q_len, kv_len = 256 * (1 + length % 3), 1152 + 128 * length + q_view = torch.arange(q_len, device="cuda") // 64 % 7 # [Q] + kv_view = torch.arange(kv_len, device="cuda") // 128 % 7 # [K] + + def predicate( + _b: torch.Tensor, # [] + _h: torch.Tensor, # [] + q_idx: torch.Tensor, # [] + kv_idx: torch.Tensor, # [] + ) -> torch.Tensor: # [] + return (q_view[q_idx] == kv_view[kv_idx]) | (kv_idx < 128) # [] + + return create_block_mask( + predicate, B=1, H=1, Q_LEN=q_len, KV_LEN=kv_len, device="cuda", BLOCK_SIZE=backend.block_size + ) + + +class _MemoryAttention(torch.nn.Module): + backend: FlexBackend + + def __init__(self, backend: FlexBackend) -> None: + super().__init__() + self.backend = backend + + def forward( + self, + q: torch.Tensor, # [1,Q,4,128] + k: torch.Tensor, # [1,K,2,128] + v: torch.Tensor, # [1,K,2,128] + block_mask: BlockMask, + ) -> torch.Tensor: # [1,Q,4,128] + # Reproduce non-leaf K/V concatenation, layout conversion and explicit + # checkpoint mask arguments without allocating a production-size model. + keys = torch.cat((k[:, :128], k[:, 128:]), dim=1) # [1,K,2,128] + values = torch.cat((v[:, :128], v[:, 128:]), dim=1) # [1,K,2,128] + return attention_module.flex_attention(q, keys, values, block_mask, self.backend) # [1,Q,4,128] + + +def _memory_step( + module: torch.nn.Module, + backend: FlexBackend, + length: int, +) -> list[weakref.ReferenceType[object]]: + """Return weak references only; no completed-step graph is intentionally kept.""" + mask = _mask(length, backend) + q_len, kv_len = mask.shape[-2:] + inputs = [ + torch.randn(1, count, heads, 128, device="cuda", dtype=torch.bfloat16).requires_grad_() + for count, heads in ((q_len, 4), (kv_len, 2), (kv_len, 2)) + ] # list[[1,S,H,128]] + output = module(*inputs, mask) # [1,Q,4,128] + loss = output.float().square().mean() # [] + loss.backward() + for value in inputs: + assert value.grad is not None and torch.isfinite(value.grad).all().item() + tensors = [value for value in vars(mask).values() if isinstance(value, torch.Tensor)] # list[[...]] + return [weakref.ref(value) for value in [*inputs, output, loss, mask, *tensors]] + + +# Cold FA4 compilation and full-generation GC need headroom beyond CI's 60s default. +# Run in the isolated GPU phase so other test workers do not compete for CPU time. +@pytest.mark.serial +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Fresh compiled Flash tensor lifetime requires CUDA") +@pytest.mark.parametrize("selective_ac", [False, True], ids=["eager", "selective"]) +@pytest.mark.parametrize( + "fresh_shape_each_step", + [ + pytest.param(False, marks=[pytest.mark.L0, pytest.mark.timeout(180)], id="smoke"), + pytest.param(True, marks=[pytest.mark.L1, pytest.mark.timeout(600)], id="stress"), + ], +) +def test_periodic_gc_releases_fresh_flash_inputs( + monkeypatch: pytest.MonkeyPatch, + selective_ac: bool, + fresh_shape_each_step: bool, + record_property: Callable[[str, Any], None], +) -> None: + torch.manual_seed(1729) + device = torch.device("cuda") + reason = attention_module.flash_backend_unavailable_reason(device) + if reason is not None: + pytest.skip(reason) + backend = resolve_flex_backend(device, preference="flex_flash") + # Keep the original regression compile budget so all 30 fresh shapes stay cached, + # independently of any training recipe. + monkeypatch.setattr(torch._dynamo.config, "recompile_limit", 512) + # Force static shape specializations to exercise GC with the shared wrapper. + compiled = torch.compile(torch_flex_attention, dynamic=False, fullgraph=True) + monkeypatch.setattr(attention_module, "_COMPILED_FLEX_ATTENTION", compiled) + attention = _MemoryAttention(backend) + module = ( + _apply_selective_ac(attention, ActivationCheckpointingConfig(mode="selective", save_ops_regex=["fmha"])) + if selective_ac + else attention + ) + # Preserve the original regression cadence without importing a Model 2 recipe. + policy = dict(every_n=10, warm_up=0, gc_level=2) + callback = ManualGarbageCollection(**policy) + was_enabled = gc.isenabled() + references: list[weakref.ReferenceType[object]] = [] + observations: list[dict[str, float | int]] = [] + measured_shapes: set[int] = set() + try: + # Isolate the shared code cache without releasing compiled graphs during measurement. + torch._dynamo.reset() + for _ in range(3): + _memory_step(module, backend, 0) + gc.collect(2) + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + for iteration in range(1, 31): + callback.on_training_step_start(None, {}, iteration) + assert not gc.isenabled() + # L0 compiles a new shape in each of the first two GC windows, then reuses one. + # L1 compiles a fresh shape every step and retains the growing compiler cache. + length = iteration if fresh_shape_each_step else 1 + ((iteration - 1) // policy["every_n"]) % 2 + measured_shapes.add(length) + references.extend(_memory_step(module, backend, length)) + if iteration % policy["every_n"] == 0: + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + start = time.monotonic() + callback.every_n_impl(None, None, {}, {}, None, iteration) + collection_seconds = time.monotonic() - start + torch.cuda.synchronize() + after = torch.cuda.memory_allocated() + live = sum(reference() is not None for reference in references) + observations.append( + { + "iteration": iteration, + "before_allocated_bytes": before, + "after_allocated_bytes": after, + "live_weak_references": live, + "collection_seconds": collection_seconds, + } + ) + assert live == 0, observations + assert after <= baseline + 1024**2, observations + record_property("selective_ac", selective_ac) + record_property("gc_policy", policy) + record_property("post_collection_memory", observations) + record_property("fresh_shapes", len(measured_shapes)) + record_property("measured_steps", 30) + finally: + torch._dynamo.reset() + gc.collect(2) + if was_enabled: + gc.enable() + else: + gc.disable() diff --git a/cosmos_framework/model/generator/mot/maskless_attention.py b/cosmos_framework/model/generator/mot/maskless_attention.py index b40343803..8adddf535 100644 --- a/cosmos_framework/model/generator/mot/maskless_attention.py +++ b/cosmos_framework/model/generator/mot/maskless_attention.py @@ -6,7 +6,8 @@ Replay roles determine each query group's keys. Grouping happens once before the decoder, over metadata runs rather than a quadratic token mask. Queries sharing the same keys share one varlen group, including across frames in a replay chunk. -The two sensor passes deliberately overlap, as in bidirectional maskless attention. +The sensor passes overlap when cross-view deduplication is disabled (legacy checkpoint +behavior). With deduplicate_cross_view enabled, same-view edges appear only once. The same-view partition re-reads each view's causal history once per chunk. Gathering it would copy those keys per group, so it always runs gapped, reading them in place; only the small diff --git a/cosmos_framework/model/generator/mot/multiview_action_attention.py b/cosmos_framework/model/generator/mot/multiview_action_attention.py new file mode 100644 index 000000000..572bdbdd6 --- /dev/null +++ b/cosmos_framework/model/generator/mot/multiview_action_attention.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Reusable same-instant RGB attention for multiview action sequences.""" + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch + +from cosmos_framework.model.attention import attention +from cosmos_framework.model.generator.mot.merge_bridge import MergeAttentionsBridge + + +@dataclass(frozen=True) +class MultiviewActionAttentionPlan: + """Frame-major RGB gathers from camera-major [actions, RGB] supertokens.""" + + rgb_indexes: torch.Tensor # [N_RGB]; empty for a single camera + offsets: torch.Tensor # [T+1] int32 + tokens_per_instant: int + gen_tokens: int + + +def build_multiview_action_attention_plan( + vision_token_shapes: Sequence[tuple[int, int, int]], + num_action_tokens: int, + device: torch.device, +) -> MultiviewActionAttentionPlan: + """Build one varlen segment per synchronized instant; actions never enter it. + + The plan only describes the packed action/RGB layout, so both teacher + forcing and self-forcing callers can reuse it. + """ + if not vision_token_shapes: + raise ValueError("Multiview action attention requires at least one camera.") + frames = vision_token_shapes[0][0] + if any(shape[0] != frames for shape in vision_token_shapes): + raise ValueError("Multiview action attention requires synchronized camera lengths.") + if num_action_tokens < 1: + raise ValueError("Multiview action attention requires RGB/action supertokens.") + per_view: list[torch.Tensor] = [] # each [T,S_RGB_view] + gen_tokens = 0 + for _, height, width in vision_token_shapes: + spatial = height * width + stride = num_action_tokens + spatial + per_view.append( + gen_tokens + torch.arange(frames)[:, None] * stride + num_action_tokens + torch.arange(spatial)[None, :] + ) # [T,S_RGB_view] + gen_tokens += frames * stride + tokens_per_instant = sum(shape[1] * shape[2] for shape in vision_token_shapes) + rgb_indexes = ( + torch.cat(per_view, dim=1).flatten().to(device=device) + if len(vision_token_shapes) > 1 + else torch.empty(0, dtype=torch.long, device=device) + ) # [N_RGB] + offsets = torch.arange(frames + 1, dtype=torch.int32, device=device) * tokens_per_instant # [T+1] + return MultiviewActionAttentionPlan(rgb_indexes, offsets, tokens_per_instant, gen_tokens) + + +def same_instant_multiview_rgb_attention( + query: torch.Tensor, # [N_gen_padded,H,D] + key: torch.Tensor, # [N_gen_padded,H_kv,D] + value: torch.Tensor, # [N_gen_padded,H_kv,D] + plan: MultiviewActionAttentionPlan, +) -> tuple[torch.Tensor, torch.Tensor]: # [1,N_gen,H,D], [1,N_gen,H] + """Attend live RGB at the same instant, retaining the merge backward contract. + + A view's own RGB participates in both this pass and its three-way pass, + matching the bidirectional decomposition's overlap. Callers skip this + component entirely for V=1. Clean history stays in the per-view TF path. + """ + indexes = plan.rgb_indexes + out, lse = attention( + query[indexes].unsqueeze(0), # [1,N_RGB,H,D] + key[indexes].unsqueeze(0), # [1,N_RGB,H_kv,D] + value[indexes].unsqueeze(0), # [1,N_RGB,H_kv,D] + cumulative_seqlen_Q=plan.offsets, + cumulative_seqlen_KV=plan.offsets, + max_seqlen_Q=plan.tokens_per_instant, + max_seqlen_KV=plan.tokens_per_instant, + backend="natten", + return_lse=True, + ) # [1,N_RGB,H,D], [1,N_RGB,H] + + def scatter( + output: torch.Tensor, normalizer: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: # [1,N_RGB,H,D], [1,N_RGB,H] -> full GEN + full_output = output.new_zeros(1, plan.gen_tokens, output.shape[2], output.shape[3]) # [1,N_gen,H,D] + full_lse = normalizer.new_full( + (1, plan.gen_tokens, normalizer.shape[2]), torch.finfo(normalizer.dtype).min + ) # [1,N_gen,H] + full_output[:, indexes] = output # [1,N_gen,H,D] + full_lse[:, indexes] = normalizer # [1,N_gen,H] + return full_output, full_lse + + def gather( + output: torch.Tensor, normalizer: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: # full GEN -> [1,N_RGB,H,D], [1,N_RGB,H] + return output[:, indexes], normalizer[:, indexes] # [1,N_RGB,H,D], [1,N_RGB,H] + + return MergeAttentionsBridge.apply(out, lse, scatter, gather) # [1,N_gen,H,D], [1,N_gen,H] diff --git a/cosmos_framework/model/generator/mot/rigrope.py b/cosmos_framework/model/generator/mot/rigrope.py new file mode 100644 index 000000000..2c2266560 --- /dev/null +++ b/cosmos_framework/model/generator/mot/rigrope.py @@ -0,0 +1,171 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""VRAD's active RopeRigPE, pinned to 979c4d32fe1da235c0a031ebf86d1767013db700. + +Ported from genradar/networks/position_embedding.py and datasets/utils.py. +By default direction and moment are independently normalized as in that source. +The spatial ablation instead scales the raw moment by a fixed distance in metres. +The fixed frequencies introduce buffers only, with no learned parameters. +""" + +import math + +import torch +from torch import nn +from torch.nn import functional as F + + +class RopeRigPE(nn.Module): + """Rotary Position Embedding from geometric (Plucker) features. + + Uses NeRF-style frequencies ``2^(linspace(min_freq_exp, max_freq_exp, n)) * pi`` so + that input features in [-1, 1] produce angles in + [2^min_freq_exp * pi, 2^max_freq_exp * pi]. The default [pi, 8 pi] matches the + angular range of grid-index rope3d (~26). + + Each coordinate dimension gets ``total_angles // feature_dim`` unique + frequencies. All dimensions are treated uniformly. + """ + + head_dim: int + feature_dim: int + min_freq_exponent: float + max_freq_exponent: float + freq_matrix: torch.Tensor # [8,D/2] + + def __init__( + self, + head_dim: int, + feature_dim: int = 8, + max_freq_exponent: float = 3.0, + min_freq_exponent: float = 0.0, + ) -> None: + super().__init__() + if head_dim <= 0 or head_dim % 2 or feature_dim != 8: + raise ValueError("RigRoPE requires a positive even head dimension and eight features") + if not (math.isfinite(min_freq_exponent) and math.isfinite(max_freq_exponent)): + raise ValueError("RigRoPE frequency exponents must be finite") + if min_freq_exponent > max_freq_exponent: + raise ValueError( + f"RigRoPE min_freq_exponent={min_freq_exponent} exceeds max_freq_exponent={max_freq_exponent}" + ) + self.head_dim = head_dim + self.feature_dim = feature_dim + self.min_freq_exponent = min_freq_exponent + self.max_freq_exponent = max_freq_exponent + freq_matrix = self._build_freq_matrix( + feature_dim, head_dim // 2, max_freq_exponent, min_freq_exponent + ) # [8,D/2] + # This matrix is fully determined by the constructor arguments. Keeping it out of the + # checkpoint lets a strict warm start from the pre-RigRoPE model initialize it locally. + self.register_buffer("freq_matrix", freq_matrix, persistent=False) + + @staticmethod + def _build_freq_matrix( + feature_dim: int, total_angles: int, max_freq_exponent: float, min_freq_exponent: float = 0.0 + ) -> torch.Tensor: + """Build NeRF-style ``[feature_dim, total_angles]`` frequency matrix. + + Angles are distributed evenly across all coordinate dimensions. + Each dimension gets its own set of unique ``2^e * pi`` frequencies. + """ + base_per_dim, extra = divmod(total_angles, feature_dim) + freq_matrix = torch.zeros(feature_dim, total_angles) # [8,D/2] + col = 0 + for dimension in range(feature_dim): + count = base_per_dim + int(dimension < extra) + exponents = torch.linspace(min_freq_exponent, max_freq_exponent, count) # [K] + freq_matrix[dimension, col : col + count] = (2.0**exponents) * math.pi # [K] + col += count + return freq_matrix.float() # [8,D/2] + + def reset_parameters(self) -> None: + """Re-initialise frequency buffer (needed after meta -> CUDA).""" + new = self._build_freq_matrix( + self.feature_dim, self.head_dim // 2, self.max_freq_exponent, self.min_freq_exponent + ) # [8,D/2] + self.freq_matrix.data.copy_(new.to(self.freq_matrix.device)) # [8,D/2] + + def forward(self, coords: torch.Tensor) -> torch.Tensor: # [...,8] -> [N,1,1,D] + """Compute TE-compatible RoPE; leading dimensions are flattened.""" + flat = coords.reshape(-1, self.feature_dim).float() # [N,8] + freq = self.freq_matrix.to(flat.device).float() # [8,D/2] + thetas = flat @ freq # [N,D/2] + emb = torch.cat([thetas, thetas], dim=-1) # [N,D] + return emb.view(flat.shape[0], 1, 1, self.head_dim).float() # [N,1,1,D] + + +def rig_features( + directions: torch.Tensor, + origins: torch.Tensor, + times: torch.Tensor, + *, + moment_scale_m: float | None = None, + include_time: bool = True, +) -> torch.Tensor: # directions, origins: [H,W,3]; times: [T]; returns [T,H,W,8] + """Plucker features with zero depth and an optional continuous, metric moment. + + Disabling time keeps its existing frequency slots at zero; it does not + redistribute those slots or alter the model's baseline mRoPE time encoding. + """ + if moment_scale_m is not None and (not math.isfinite(moment_scale_m) or moment_scale_m <= 0): + raise ValueError("RigRoPE moment scale must be finite and positive, in metres") + d = directions / (torch.norm(directions, dim=-1, keepdim=True) + 1e-8) # [H,W,3] + m = torch.cross(origins.expand_as(d), d, dim=-1) # [H,W,3] + if moment_scale_m is None: + m = m / (torch.norm(m, dim=-1, keepdim=True) + 1e-8) # [H,W,3] + else: + m = m / moment_scale_m # [H,W,3] + depth = torch.zeros_like(d[..., :1]) # [H,W,1] + spatial = torch.cat([d, m, depth], dim=-1) # [H,W,7] + spatial = spatial.unsqueeze(0).expand(len(times), -1, -1, -1) # [T,H,W,7] + encoded_times = times if include_time else torch.zeros_like(times) # [T] + timestamps = encoded_times[:, None, None, None].expand(-1, d.shape[0], d.shape[1], -1) # [T,H,W,1] + return torch.cat([spatial, timestamps], dim=-1) # [T,H,W,8] + + +def resize_rig_features(coords: torch.Tensor, shape: tuple[int, int, int]) -> torch.Tensor: + """VRAD trilinear feature interpolation; deliberately no normalization after it.""" + # coords: [V,T,H,W,8]; each physical view has an independent interpolation grid. + if tuple(coords.shape[1:4]) == shape: + return coords # [V,T,H,W,8] + channels = coords.permute(0, 4, 1, 2, 3).float() # [V,8,T,H,W] + resized = F.interpolate(channels, size=shape, mode="trilinear", align_corners=False) # [V,8,T',H',W'] + return resized.permute(0, 2, 3, 4, 1) # [V,T',H',W',8] + + +def spread_over_even_pairs(angles: torch.Tensor) -> torch.Tensor: # [N,D/2] -> [N,D] + """Place a half-width ``RopeRigPE`` output on the even split-half pairs of a full head. + + ``RopeRigPE(D/2)`` pairs channel ``j`` with ``j + D/4``; its pair ``j`` lands on the full + head's pair ``2j``, channels ``(2j, 2j + D/2)``. The odd pairs take angle zero, which leaves + them unrotated for the caller to fill. + """ + count, width = angles.shape + if width % 4: + raise ValueError(f"Even-pair RigRoPE needs a half-width angle table divisible by 4, got {width}") + half = angles.new_zeros(count, width) # [N,D/2], one angle per full-head pair + half[:, 0::2] = angles[:, : width // 2] + return torch.cat([half, half], dim=-1) # [N,D] + + +def apply_rig_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Apply split-half rotary to Q or K; never to V.""" + # x: [N,H,D]; cos, sin: [N,D] + first, second = x.chunk(2, dim=-1) # [N,H,D/2], [N,H,D/2] + rotated = torch.cat([-second, first], dim=-1) # [N,H,D] + return (x.float() * cos[:, None, :] + rotated.float() * sin[:, None, :]).to(x.dtype) # [N,H,D] + + +def apply_mrope_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Apply baseline split-half mRoPE in its original dtype after the CP gather.""" + # x: [N,H,D]; cos, sin: [N,D_rotary], allowing the baseline's partial rotary embedding. + rotary_dim = cos.shape[-1] + first, second = x[..., :rotary_dim].chunk(2, dim=-1) # each [N,H,D_rotary/2] + rotated = torch.cat([-second, first], dim=-1) # [N,H,D_rotary] + output = x[..., :rotary_dim] * cos[:, None, :] + rotated * sin[:, None, :] # [N,H,D_rotary] + if rotary_dim == x.shape[-1]: + return output # [N,H,D] + return torch.cat([output, x[..., rotary_dim:]], dim=-1) # [N,H,D] diff --git a/cosmos_framework/model/generator/mot/rigrope_test.py b/cosmos_framework/model/generator/mot/rigrope_test.py new file mode 100644 index 000000000..15b05939c --- /dev/null +++ b/cosmos_framework/model/generator/mot/rigrope_test.py @@ -0,0 +1,146 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""Analytic checks for the fixed VRAD frequency and Plucker conventions.""" + +import math + +import pytest +import torch + +from cosmos_framework.model.generator.mot.rigrope import ( + RopeRigPE, + apply_mrope_rotary, + apply_rig_rotary, + resize_rig_features, + rig_features, + spread_over_even_pairs, +) + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +@pytest.mark.parametrize("head_dim", [16, 32, 128, 130]) +def test_fixed_frequencies_and_no_parameters(head_dim: int) -> None: + encoding = RopeRigPE(head_dim) + assert list(encoding.parameters()) == [] + # The fixed matrix must not become a new required key in strict pre-RigRoPE warm starts. + assert list(encoding.state_dict()) == [] + counts = torch.count_nonzero(encoding.freq_matrix, dim=1).tolist() # [8] + assert counts == [head_dim // 16 + int(i < (head_dim // 2) % 8) for i in range(8)] + for row in encoding.freq_matrix: + frequencies = row[row != 0] # [K] + assert float(frequencies[0]) == pytest.approx(math.pi) + if len(frequencies) > 1: + assert float(frequencies[-1]) == pytest.approx(8 * math.pi) + restored = RopeRigPE(head_dim) + # Strict loading an old checkpoint with no RigRoPE key leaves the deterministic buffer intact. + restored.load_state_dict({}, strict=True) + torch.testing.assert_close(restored.freq_matrix, encoding.freq_matrix, rtol=0, atol=0) + + +def test_frequency_range_spans_the_configured_exponents() -> None: + encoding = RopeRigPE(128, max_freq_exponent=1.0, min_freq_exponent=-2.0) + for row in encoding.freq_matrix: + frequencies = row[row != 0] # [8] + assert float(frequencies[0]) == pytest.approx(math.pi / 4) + assert float(frequencies[-1]) == pytest.approx(2 * math.pi) + # Opposite-facing unit directions differ by 2 per component: the slowest slot must not wrap. + assert 2 * float(encoding.freq_matrix[encoding.freq_matrix != 0].min()) < 2 * math.pi + rebuilt = encoding.freq_matrix.clone() # [8,D/2] + encoding.freq_matrix.zero_() + encoding.reset_parameters() + torch.testing.assert_close(encoding.freq_matrix, rebuilt, rtol=0, atol=0) + + +@pytest.mark.parametrize("exponents", [(1.0, 0.0), (0.0, math.inf)]) +def test_frequency_range_rejects_inverted_or_infinite_exponents(exponents: tuple[float, float]) -> None: + low, high = exponents + with pytest.raises(ValueError, match="exponent"): + RopeRigPE(16, max_freq_exponent=high, min_freq_exponent=low) + + +def test_vrad_moment_normalization_zero_depth_and_time() -> None: + directions = torch.tensor([[[0.0, 2.0, 0.0], [3.0, 0.0, 0.0]]]) # [1,2,3] + origin = torch.tensor([5.0, 0.0, 0.0]) # [3] + times = torch.tensor([0.0, 0.1, 0.2]) # [3] + features = rig_features(directions, origin, times) # [3,1,2,8] + torch.testing.assert_close(features[0, 0, 0, :7], torch.tensor([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0])) + assert torch.count_nonzero(features[:, :, :, 6]) == 0 + assert torch.count_nonzero(features[:, :, 1, 3:6]) == 0 # Origin lies on the second ray. + torch.testing.assert_close(features[:, 0, 0, 7], times) + # VRAD intentionally removes metric baseline length by independently normalizing moment. + torch.testing.assert_close(features, rig_features(directions, origin * 10, times)) + + +def test_feature_interpolation_is_per_view_without_renormalization() -> None: + features = torch.zeros(2, 1, 2, 2, 8) # [V,T,H,W,8] + features[0, :, 0, :, 0] = 1 # [T,W] + features[0, :, 1, :, 1] = 1 # [T,W] + features[1, :, :, :, 2] = 1 # [T,H,W] + resized = resize_rig_features(features, (1, 1, 1)) # [2,1,1,1,8] + torch.testing.assert_close(resized[0, 0, 0, 0, :3], torch.tensor([0.5, 0.5, 0.0])) + torch.testing.assert_close(resized[1, 0, 0, 0, :3], torch.tensor([0.0, 0.0, 1.0])) + + +def test_rotary_relative_phase_and_gradients() -> None: + torch.manual_seed(19) + q = torch.randn(5, 2, 16, requires_grad=True) # [N,H,D] + k = torch.randn_like(q, requires_grad=True) # [N,H,D] + coords = torch.randn(5, 8) # [N,8] + angles = RopeRigPE(16)(coords).reshape(5, 16) # [N,D] + q_ = apply_rig_rotary(q, angles.cos(), angles.sin()) # [N,H,D] + k_ = apply_rig_rotary(k, angles.cos(), angles.sin()) # [N,H,D] + # Independent complex-number reference uses split-half pairs. + q_complex = torch.complex(q[..., :8], q[..., 8:]) # [N,H,D/2] + k_complex = torch.complex(k[..., :8], k[..., 8:]) # [N,H,D/2] + phase = torch.exp(1j * angles[..., :8])[:, None, :] # [N,1,D/2] + expected = ((q_complex[0] * phase[0]).conj() * (k_complex[3] * phase[3])).real.sum(-1) # [H] + actual = (q_[0] * k_[3]).sum(-1) # [H] + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) + expected_grad = torch.autograd.grad(expected.sum(), (q, k), retain_graph=True) # tuple[[N,H,D]] + actual_grad = torch.autograd.grad(actual.sum(), (q, k)) # tuple[[N,H,D]] + for actual_item, expected_item in zip(actual_grad, expected_grad, strict=True): + torch.testing.assert_close(actual_item, expected_item, atol=2e-6, rtol=2e-6) + torch.testing.assert_close(apply_rig_rotary(q, torch.ones_like(angles), torch.zeros_like(angles)), q) + + +def test_even_pair_spread_keeps_every_component_on_the_prope_pairs() -> None: + head_dim = 128 + encoding = RopeRigPE(head_dim // 2) + # Half width halves each descriptor component's frequencies rather than dropping components. + assert torch.count_nonzero(encoding.freq_matrix, dim=1).tolist() == [4] * 8 + coords = torch.randn(5, 8) # [N,8] + half = encoding(coords).reshape(5, head_dim // 2) # [N,D/2] + full = spread_over_even_pairs(half) # [N,D] + pairs = torch.arange(head_dim) % (head_dim // 2) % 2 == 0 # [D] + assert torch.count_nonzero(full[:, ~pairs]) == 0 + for pair in range(head_dim // 4): + torch.testing.assert_close(full[:, 2 * pair], half[:, pair], rtol=0, atol=0) + torch.testing.assert_close(full[:, 2 * pair + head_dim // 2], half[:, pair], rtol=0, atol=0) + with pytest.raises(ValueError, match="divisible by 4"): + spread_over_even_pairs(torch.zeros(5, 6)) + + +def test_even_pair_rigrope_beside_mrope_keeps_scores_relative() -> None: + torch.manual_seed(23) + head_dim = 32 + q, k = torch.randn(2, 1, 1, head_dim, dtype=torch.float64).unbind(0) # each [1,1,D] + pairs = torch.arange(head_dim) % (head_dim // 2) % 2 == 0 # [D] + rig = spread_over_even_pairs(torch.randn(2, head_dim // 2, dtype=torch.float64)) # [2,D] + mrope = torch.randn(2, head_dim // 2, dtype=torch.float64).repeat(1, 2) # [2,D] + + def score(rig_angles: torch.Tensor, mrope_angles: torch.Tensor) -> torch.Tensor: + def rotate(x: torch.Tensor, token: int) -> torch.Tensor: # [1,1,D] + rig_token, mrope_token = rig_angles[token : token + 1], mrope_angles[token : token + 1] # each [1,D] + geometry = apply_rig_rotary(x, rig_token.cos(), rig_token.sin()) + position = apply_mrope_rotary(x, mrope_token.cos(), mrope_token.sin()) + return torch.where(pairs, geometry, position) + + return (rotate(q, 0) * rotate(k, 1)).sum() + + rig_shift = spread_over_even_pairs(torch.randn(1, head_dim // 2, dtype=torch.float64)) # [1,D] + mrope_shift = torch.randn(1, head_dim // 2, dtype=torch.float64).repeat(1, 2) # [1,D] + # Shifting both tokens by one offset in either encoding leaves the score unchanged. + torch.testing.assert_close(score(rig + rig_shift, mrope + mrope_shift), score(rig, mrope), atol=1e-5, rtol=1e-5) diff --git a/cosmos_framework/model/generator/omni_mot_causal_model.py b/cosmos_framework/model/generator/omni_mot_causal_model.py index a9a2bd377..bf5cbddbc 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model.py @@ -15,7 +15,7 @@ import contextlib import copy import itertools -from collections.abc import Callable, Generator, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable from dataclasses import dataclass from typing import Any, Literal, cast from unittest.mock import patch @@ -30,15 +30,15 @@ import cosmos_framework.model.generator.omni_mot_model as omni_mot_model_module from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig from cosmos_framework.data.generator.augmentors.text_tokenizer import TEXT_SYSTEM_PROMPT_KEY -from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_broadcast_tensor_list from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel, _broadcast_seed, _per_view_caption_groups from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.model.generator.utils.memory import MemoryState -from cosmos_framework.data.generator.sequence_packing import PackedSequence, build_sequence_plans_from_data_batch +from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan, build_sequence_plans_from_data_batch from cosmos_framework.data.generator.sequence_packing.modality import compute_text_split_length from cosmos_framework.data.generator.sequence_packing.runtime import to_device_nonblocking from cosmos_framework.configs.base.defaults.causal_flex_attention import CausalFlexAttentionConfig from cosmos_framework.configs.base.defaults.replay_attention import ( + TEACHER_FORCING_KV_IMPLEMENTATIONS, TeacherForcingKVImplementation, TeacherForcingReplayPolicyConfig, ) @@ -70,15 +70,21 @@ ARMemoryState, DualKVCache, KVBufferPool, - MultiviewARMemoryState, TeacherForcingMemoryState, ) from cosmos_framework.model.generator.utils.kv_storage_backend import validate_kv_cache_dtype -from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend, sample_joint_transfer_ar +from cosmos_framework.model.generator.utils.multiview_ar import ( + MultiviewTransferARBackend, + generate_multiview_transfer_ar_chunk, + iter_robot_multiview_action_ar, + iter_samples_multiview_transfer_autoregressive, + sample_joint_transfer_ar, +) from cosmos_framework.model.generator.utils.nvfp4 import resolve_legacy_nvfp4_mode from cosmos_framework.data.generator.sequence_packing.autoregressive import ( pack_input_sequence_autoregressive, pack_input_sequence_autoregressive_batch, + pack_multiview_action_temporal_causal, resolve_text_system_prompt, ) from cosmos_framework.utils.generator.data_batch import condition_frame_indexes_vision_from_batch @@ -140,7 +146,7 @@ class reads them directly; they will migrate here once the corresponding # attention layout while it builds the network. teacher_forcing_kv_implementation: TeacherForcingKVImplementation = attrs.field( default="singleview_threeway_kv", - validator=attrs.validators.in_(("multiview_flex_kv", "multiview_maskless_kv", "singleview_threeway_kv")), + validator=attrs.validators.in_(TEACHER_FORCING_KV_IMPLEMENTATIONS), ) # Backend-neutral connectivity for clean replay and transfer control. All @@ -324,111 +330,6 @@ def _iter_ar_chunk_ranges(start_frame: int, num_frames: int, chunk_size: int): f = chunk_end -def _iter_multiview_logical_frames( - vision_latent: torch.Tensor, - *, - num_views: int, -) -> Generator[torch.Tensor, None, None]: # vision_latent: [B,C,V*T,H,W], yields [B,C,V,H,W] - """Yield one camera-major latent time step across all views at a time.""" - if num_views < 1: - raise ValueError(f"num_views must be positive, got {num_views}.") - if vision_latent.ndim != 5: - raise ValueError(f"Expected multiview latent rank 5, got shape {tuple(vision_latent.shape)}.") - flat_frame_count = int(vision_latent.shape[2]) - if flat_frame_count % num_views != 0: - raise ValueError(f"Camera-major latent length {flat_frame_count} must be divisible by num_views={num_views}.") - frames_per_view = flat_frame_count // num_views - vision_by_view = vision_latent.reshape( # [B,C,V,T,H,W] - vision_latent.shape[0], - vision_latent.shape[1], - num_views, - frames_per_view, - vision_latent.shape[3], - vision_latent.shape[4], - ) - for frame_idx in range(frames_per_view): - logical_frame = vision_by_view[:, :, :, frame_idx : frame_idx + 1].reshape( # [B,C,V,H,W] - vision_latent.shape[0], - vision_latent.shape[1], - num_views, - vision_latent.shape[3], - vision_latent.shape[4], - ) - yield logical_frame - - -def _multiview_conditioned_prefix_length( - target_condition_mask: torch.Tensor, - *, - num_views: int, - frames_per_view: int, -) -> int: - """Return a synchronized contiguous-prefix length from a camera-major mask.""" - expected_frames = num_views * frames_per_view - if num_views < 1 or frames_per_view < 1 or target_condition_mask.numel() != expected_frames: - raise ValueError( - "Expected one target condition value per camera-major frame; " - f"got shape {tuple(target_condition_mask.shape)} for V={num_views}, T={frames_per_view}." - ) - condition_grid = target_condition_mask.to(dtype=torch.bool).reshape(num_views, frames_per_view) # [V,T] - if not torch.equal(condition_grid, condition_grid[0:1].expand_as(condition_grid)): - raise ValueError("Multiview transfer AR requires synchronized target conditioning across every view.") - condition_indexes = torch.nonzero(condition_grid[0], as_tuple=False).squeeze(1) # [T_condition] - expected_indexes = torch.arange(condition_indexes.numel(), device=condition_indexes.device) # [T_condition] - if not torch.equal(condition_indexes, expected_indexes): - raise ValueError( - "Multiview transfer AR only supports contiguous prefix conditioning from frame 0; " - f"got condition frame indexes {condition_indexes.tolist()}." - ) - return condition_indexes.numel() - - -def _submit_multiview_conditioned_prefix( - target_latent: torch.Tensor, - *, - num_views: int, - frames_per_view: int, - condition_count: int, - output_frames: int, - on_clean_vision_chunk: Callable[[torch.Tensor], None] | None, -) -> torch.Tensor | None: # target_latent: [B,C,V*T,H,W], returns [B,C,V*T_prefix,H,W] or None - """Return the camera-major conditioned prefix and optionally submit it for decode.""" - if condition_count == 0: - return None - if num_views < 1 or frames_per_view < 1: - raise ValueError(f"Expected positive multiview dimensions, got V={num_views}, T={frames_per_view}.") - if target_latent.ndim != 5 or target_latent.shape[2] != num_views * frames_per_view: - raise ValueError( - "Expected camera-major target shape [B,C,V*T,H,W], " - f"got {tuple(target_latent.shape)} for V={num_views}, T={frames_per_view}." - ) - if not 0 <= condition_count <= frames_per_view: - raise ValueError(f"condition_count must be in [0, {frames_per_view}], got {condition_count}.") - if not 1 <= output_frames <= frames_per_view: - raise ValueError(f"output_frames must be in [1, {frames_per_view}], got {output_frames}.") - - prefix_frame_count = min(condition_count, output_frames) - batch_size, channels, _, height, width = target_latent.shape - target_by_view = target_latent.reshape( # [B,C,V,T,H,W] - batch_size, - channels, - num_views, - frames_per_view, - height, - width, - ) - conditioned_prefix = target_by_view[:, :, :, :prefix_frame_count].reshape( # [B,C,V*T_prefix,H,W] - batch_size, - channels, - num_views * prefix_frame_count, - height, - width, - ) - if on_clean_vision_chunk is not None: - on_clean_vision_chunk(conditioned_prefix) - return conditioned_prefix - - def _validate_kv_cache_dtype_supports_cuda_graphs(kv_cache_dtype: str | None, cuda_graph_path_active: bool) -> None: """Reject FP8 KV cache when the cuda-graph AR path is effectively active. @@ -546,7 +447,7 @@ def _resolve_teacher_forcing_replay_policy(value: Any) -> TeacherForcingReplayPo def _resolve_teacher_forcing_kv_implementation(value: Any) -> TeacherForcingKVImplementation: """Validate the public replay selector after LazyConfig resolution.""" - supported_implementations = ("multiview_flex_kv", "multiview_maskless_kv", "singleview_threeway_kv") + supported_implementations = TEACHER_FORCING_KV_IMPLEMENTATIONS if value not in supported_implementations: raise ValueError( f"teacher_forcing_kv_implementation must be one of {supported_implementations}, got {value!r}." @@ -675,6 +576,8 @@ def build_net( ) -> torch.nn.Module: """Resolve the selected teacher-forcing KV implementation and build it.""" uses_multiview_replay_kv = self._uses_multiview_replay_kv() + robot_threeway = self._uses_robot_multiview_threeway() + uses_interactive_network = uses_multiview_replay_kv or robot_threeway if self.config.causal_training_strategy not in _TEACHER_FORCING_REPLAY_STRATEGIES: return super().build_net(dtype, mp_policy=mp_policy, lora_enabled=lora_enabled) @@ -691,15 +594,29 @@ def build_net( maskless_replay = ( video_temporal_causal and self._get_teacher_forcing_kv_implementation() == "multiview_maskless_kv" ) + # LazyConfig leaves this nested value parented to ``self.config``. Keep + # an independent copy while temporarily hiding the opt-in below, or the + # first (regular) network build invalidates the value needed by the EMA + # network build. + multiview_action_conditioning = copy.deepcopy(getattr(self.config, "multiview_action_conditioning", None)) + if multiview_action_conditioning is not None and not robot_threeway: + raise ValueError("Bimanual multiview Stage 1/2 requires multiview_threeway_kv (three-way + varlen).") + if robot_threeway and ( + multiview_action_conditioning is None + or not video_temporal_causal + or replay_policy.multiview_attention_scope != "decomposed" + or replay_policy.decomposed_temporal_window_seconds is not None + ): + raise ValueError("multiview_threeway_kv requires causal robot packs and same-instant decomposed attention.") if maskless_replay: # Validate the execution device before network construction enters the meta device context. require_gapped_replay(torch.device(self.tensor_kwargs["device"])) - # One knob rather than two that had to agree: the pathway is what selects multiview - # attention, so there is no second flag to save and restore alongside it. + # Robot multiview keeps the production three-way layout and adds only + # one same-instant RGB varlen component in Interactive attention. self.config.joint_attn_implementation = "multiview" if uses_multiview_replay_kv else "three_way" if uses_multiview_replay_kv: - # Core validates temporal causality as a three-way-only layout. The - # replay mask supplies causality for this two-way path. + # Core validates temporal causality as a three-way-only layout. + # Interactive supplies AV replay below. self.config.video_temporal_causal = False self.config.multiview_attention.mask.attention_scope = replay_policy.multiview_attention_scope # Replay's duration is past-only; preserve it as signed bounds in the base config. @@ -709,12 +626,15 @@ def build_net( ) if maskless_replay: self.config.multiview_attention.backend = "maskless" - # Core's bidirectional geometry has no sliding-window fold. - # Replay replaces that geometry with its own same-instant - # visibility plan before decoder execution. self.config.multiview_attention.mask.decomposed_temporal_window_seconds = None try: - if uses_multiview_replay_kv: + if robot_threeway: + # Core's fixed-code validator is coupled to bidirectional + # maskless attention. Packing still reads the restored model + # config; hide the opt-in only while constructing a normal + # causal three-way network. + self.config.multiview_action_conditioning = None + if uses_interactive_network: with patch.object( omni_mot_model_module, "Cosmos3VFMNetwork", @@ -728,18 +648,20 @@ def build_net( # of the student/teacher selectors restored on the model below. net.config.multiview_attention_config = copy.deepcopy(self.config.multiview_attention) finally: + self.config.multiview_action_conditioning = multiview_action_conditioning self.config.video_temporal_causal = video_temporal_causal self.config.joint_attn_implementation = joint_attn_implementation self.config.multiview_attention.mask.attention_scope = attention_scope self.config.multiview_attention.mask.decomposed_temporal_window_seconds = decomposed_temporal_window_seconds self.config.multiview_attention.backend = multiview_backend - if uses_multiview_replay_kv: + if uses_interactive_network: net.config.video_temporal_causal = video_temporal_causal net.video_temporal_causal = video_temporal_causal setattr(net, "teacher_forcing_replay_policy", replay_policy) setattr(net, "teacher_forcing_frames_per_chunk", self.config.teacher_forcing_frames_per_chunk) setattr(net, "teacher_forcing_maskless", maskless_replay) + setattr(net, "teacher_forcing_multiview_threeway", robot_threeway) return net def maybe_convert_linears_to_nvfp4(self) -> None: @@ -778,7 +700,7 @@ def memory_init_training( Optionally removes the control item from transfer samples. The legacy path drops trailing latent frames to satisfy the chunkwise-TF constraint; - the multiview Flex path keeps partial camera-major tails and represents them + the multiview replay path keeps partial camera-major tails and represents them in mask metadata. Returns the default (non-rolling) memory info. Rolling KV-cache (segment-based) training is not supported. """ @@ -786,7 +708,7 @@ def memory_init_training( gen_data_clean = self._maybe_drop_teacher_forcing_transfer_control(gen_data_clean, data_batch) # Chunkwise TF requires every latent vision clip to satisfy (T - 1) % C == 0. # The legacy chunkwise path requires divisibility and therefore drops - # trailing latent frames in lockstep across modalities. The multiview Flex + # trailing latent frames in lockstep across modalities. The multiview replay # path represents camera-major partial tails explicitly in its mask metadata. if not self._uses_multiview_replay_kv(): gen_data_clean = self._truncate_for_chunkwise_tf(gen_data_clean) @@ -890,9 +812,53 @@ def _uses_multiview_replay_kv(self) -> bool: "multiview_maskless_kv", ) + def _uses_robot_multiview_threeway(self) -> bool: + """Whether this run adds robot RGB varlen attention to normal three-way TF.""" + implementation = self._get_teacher_forcing_kv_implementation() + _validate_teacher_forcing_kv_strategy(implementation, self.config.causal_training_strategy) + return ( + self.config.causal_training_strategy in _TEACHER_FORCING_REPLAY_STRATEGIES + and implementation == "multiview_threeway_kv" + ) + @override - def _pack_input_sequence(self, *args: Any, **kwargs: Any) -> PackedSequence: + def _pack_input_sequence( + self, + sequence_plans: list[SequencePlan], + input_text_indexes: list[list[int]], + gen_data_clean: GenerationDataClean, + input_timesteps: torch.Tensor, + include_end_of_generation_token: bool = False, + skip_text_tokens: bool = False, + initial_mrope_temporal_offset: int | float = 0, + ) -> PackedSequence: """Keep the standard multiview layout when replay supplies causality.""" + args = (sequence_plans, input_text_indexes, gen_data_clean, input_timesteps) + kwargs = dict( + include_end_of_generation_token=include_end_of_generation_token, + skip_text_tokens=skip_text_tokens, + initial_mrope_temporal_offset=initial_mrope_temporal_offset, + ) + multiview_action_conditioning = getattr(self.config, "multiview_action_conditioning", None) + if self._uses_robot_multiview_threeway(): + assert multiview_action_conditioning is not None + if gen_data_clean.num_vision_items_per_sample is not None: + if skip_text_tokens or include_end_of_generation_token or initial_mrope_temporal_offset: + raise ValueError("Robot multiview TF keeps sample-level text and whole-clip training positions.") + return pack_multiview_action_temporal_causal( + *args, + config=multiview_action_conditioning, + special_tokens=self.llm_special_tokens, + patch_spatial=self.config.diffusion_expert_config.patch_spatial, + temporal_factor=self.tokenizer_vision_gen.temporal_compression_factor, + temporal_margin=self.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin, + reset_spatial_ids=self.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids, + enable_fps_modulation=self.config.diffusion_expert_config.enable_fps_modulation, + base_fps=float(self.config.diffusion_expert_config.base_fps), + ) + # As in Stage 0, an ordinary robot sample is not represented as a + # one-view multiview sample. Keep its causal supertoken pack intact. + return super()._pack_input_sequence(*args, **kwargs) if not self._uses_multiview_replay_kv(): return super()._pack_input_sequence(*args, **kwargs) video_temporal_causal = self.config.video_temporal_causal @@ -993,11 +959,12 @@ def _truncate_for_chunkwise_tf(self, gen_data_clean: GenerationDataClean) -> Gen * vision latents ``x0_tokens_vision[i]`` ``[B,C,T,H,W]`` -> ``keep_t`` frames * vision pixels ``raw_state_vision[i]`` -> ``get_pixel_num_frames(keep_t)`` - * action latents ``x0_tokens_action[j]`` ``[(T-1)*tcf, D]`` -> ``(keep_t-1)*tcf`` - rows (``tcf`` action tokens interleaved per vision frame; latent frame 0 - is the null-action conditioning frame and is not stored, hence - ``(T-1)*tcf`` rows). ``raw_state_action`` aliases ``x0_tokens_action`` - (action is not VAE-encoded), so it follows automatically. + * action latents ``x0_tokens_action[j]`` contain one or more camera-major + ``[(T-1)*tcf, D]`` view blocks. Each block is independently truncated to + ``(keep_t-1)*tcf`` rows (``tcf`` action tokens interleaved per vision + frame; latent frame 0 is the null-action conditioning frame and is not + stored). ``raw_state_action`` aliases ``x0_tokens_action`` (action is not + VAE-encoded), so it follows automatically. Sound is packed as a separate, non-chunked split on its own latent rate; it is not part of the vision-supertoken chunk grid and needs no truncation @@ -1037,21 +1004,34 @@ def chunk_aligned_frames(latent_t: int) -> int: if gen_data_clean.x0_tokens_action is not None: tcf = self.tokenizer_vision_gen.temporal_compression_factor or 1 action_domain_ids = gen_data_clean.action_domain_id + action_view_counts = getattr(gen_data_clean, "num_views_per_action_item", None) if action_domain_ids is not None and len(action_domain_ids) != len(gen_data_clean.x0_tokens_action): raise ValueError( "Action-domain metadata must have one entry per action tensor before chunkwise-TF truncation; " f"got {len(action_domain_ids)} domain entries for {len(gen_data_clean.x0_tokens_action)} actions." ) + if action_view_counts is not None and len(action_view_counts) != len(gen_data_clean.x0_tokens_action): + raise ValueError( + "Action-view metadata must have one entry per action tensor before chunkwise-TF truncation; " + f"got {len(action_view_counts)} view counts for {len(gen_data_clean.x0_tokens_action)} actions." + ) for j, act in enumerate(gen_data_clean.x0_tokens_action): if act is None: continue - rows = act.shape[-2] # (T-1)*tcf - if rows % tcf != 0: + num_views = int(action_view_counts[j]) if action_view_counts is not None else 1 + rows = act.shape[-2] # num_views*(T-1)*tcf, camera-major + if num_views < 1 or rows % num_views != 0: + raise ValueError( + f"Action latents entry {j} has {rows} rows that cannot be divided into " + f"num_views_per_action_item={num_views} camera-major blocks." + ) + rows_per_view = rows // num_views + if rows_per_view % tcf != 0: raise ValueError( - f"Action latents entry {j} has {rows} rows along the temporal axis, not a multiple of " + f"Action latents entry {j} has {rows_per_view} rows per view, not a multiple of " f"temporal_compression_factor={tcf}; cannot chunk-align for chunkwise TF." ) - keep_rows = C * ((rows // tcf) // C) * tcf # (keep_t - 1) * tcf + keep_rows_per_view = C * ((rows_per_view // tcf) // C) * tcf # (keep_t - 1) * tcf if action_domain_ids is not None: domain_ids = action_domain_ids[j] flat_domain_ids = domain_ids.reshape(-1) @@ -1060,10 +1040,20 @@ def chunk_aligned_frames(latent_t: int) -> int: f"Action-domain entry {j} must be scalar or have one ID per real action row before " f"chunkwise-TF truncation; got {flat_domain_ids.numel()} IDs for {rows} rows." ) - if flat_domain_ids.numel() == rows and keep_rows < rows: - action_domain_ids[j] = flat_domain_ids[:keep_rows].contiguous() - if keep_rows < rows: - gen_data_clean.x0_tokens_action[j] = act[..., :keep_rows, :].contiguous() + if flat_domain_ids.numel() == rows and keep_rows_per_view < rows_per_view: + action_domain_ids[j] = ( + flat_domain_ids.reshape(num_views, rows_per_view)[:, :keep_rows_per_view] + .reshape(-1) + .contiguous() + ) + if keep_rows_per_view < rows_per_view: + gen_data_clean.x0_tokens_action[j] = ( + act.reshape(*act.shape[:-2], num_views, rows_per_view, act.shape[-1])[ + ..., :keep_rows_per_view, : + ] + .reshape(*act.shape[:-2], num_views * keep_rows_per_view, act.shape[-1]) + .contiguous() + ) return gen_data_clean @@ -1153,7 +1143,7 @@ def _build_clean_tf_cache( clean_target_indexes: torch.Tensor | None = None if self._uses_multiview_replay_kv(): if packed_sequence.vision is None or packed_sequence.num_views_per_vision_item is None: - raise ValueError("Two-way Flex teacher forcing requires multiview vision metadata.") + raise ValueError("Two-way replay teacher forcing requires multiview vision metadata.") # Match the packer's per-sample order: RGB items, then LiDAR items. # A clean LiDAR target must keep its original role in the replay mask. original_condition_masks = [ @@ -1226,12 +1216,32 @@ def _validate_teacher_forcing_pack(self, packed_seq: PackedSequence) -> None: f"or aligned [control, target] items; got {packed_seq.num_vision_items_per_sample} " f"with {len(packed_seq.sample_lens)} logical samples." ) + if getattr(packed_seq, "multiview_action_conditioning", False) and self._uses_robot_multiview_threeway(): + if ( + packed_seq.vision is None + or packed_seq.action is None + or len(packed_seq.sample_lens) != 1 + or len(packed_seq.vision.token_shapes) != len(packed_seq.action.token_shapes) + or packed_seq.sound is not None + or packed_seq.lidar is not None + or packed_seq.radar is not None + or packed_seq.num_action_tokens_per_supertoken != self.tokenizer_vision_gen.temporal_compression_factor + ): + raise ValueError("Robot multiview teacher forcing requires RGB/action supertokens.") + if any( + action_shape != (vision_shape[0] * packed_seq.num_action_tokens_per_supertoken,) + for vision_shape, action_shape in zip(packed_seq.vision.token_shapes, packed_seq.action.token_shapes) + ): + raise ValueError("Each robot view needs one full action group per latent frame.") + if any(not mask.bool().all() for mask in packed_seq.action.condition_mask): + raise ValueError("Robot multiview teacher forcing requires conditioned actions.") + return if self._uses_multiview_replay_kv(): if packed_seq.vision is None or packed_seq.action is not None or packed_seq.sound is not None: - raise ValueError("Two-way Flex teacher forcing supports RGB and optional LiDAR generation batches.") + raise ValueError("Two-way replay teacher forcing supports RGB and optional LiDAR generation batches.") if packed_seq.num_views_per_vision_item is None: raise ValueError( - "Two-way Flex teacher forcing requires per-camera VAE metadata; enable " + "Two-way replay teacher forcing requires per-camera VAE metadata; enable " "enable_per_camera_vae_encoding on the dataset." ) return @@ -1995,6 +2005,44 @@ def iter_samples_from_batch_autoregressive( self._reset_ar_forward_cuda_graph_runtime_for_generation() reset_ar_post_saturation_runtime_for_generation(self) + if mode == "forward_dynamics" and data_batch.get("_robot_multiview_action_ar", False): + if self.config.compile.enabled: + raise ValueError("Robot multiview AR requires eager attention; run with --no-use-torch-compile.") + sequence_plans = build_sequence_plans_from_data_batch( + data_batch=data_batch, + input_video_key=self.input_video_key, + input_image_key=self.input_image_key, + ) + caption_groups = _per_view_caption_groups(data_batch[self.input_caption_key]) + self._apply_inference_caption_plan(sequence_plans, caption_groups) + gen_data_clean = self.get_data_and_condition( + data_batch, + vision_condition_indexes=[[0]], + retain_raw_state_vision=False, + ) + self._release_inference_raw_vision(data_batch, gen_data_clean) + cond_text_tokens, uncond_text_tokens = self._get_inference_text_tokens( + data_batch, + has_negative_prompt, + caption_groups, + ) + yield from iter_robot_multiview_action_ar( + self, + plans=sequence_plans, + data=gen_data_clean, + conditional_text=cond_text_tokens, + unconditional_text=uncond_text_tokens, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + normalize_cfg=normalize_cfg, + sampler_mode=cast(Literal["rf", "distilled"], sampler_mode), + distilled_num_steps=distilled_num_steps, + max_num_frames=max_num_frames, + ) + return + if mode == "video_transfer" and self._uses_multiview_replay_kv(): if self.config.compile.enabled: raise ValueError("Multiview transfer AR requires eager attention; run with --no-use-torch-compile.") @@ -2797,62 +2845,24 @@ def _generate_multiview_transfer_ar_chunk( memory_seq_len: int, ) -> torch.Tensor: # [1,C,V*chunk_len,H,W] """Denoise one synchronized multiview chunk against the cached K/V suffix.""" - cond_memory = MultiviewARMemoryState( - num_layers=self.net.num_hidden_layers, - memory_seq_len=memory_seq_len, - cache=cond_cache, - ) - uncond_memory = ( - MultiviewARMemoryState( - num_layers=self.net.num_hidden_layers, - memory_seq_len=memory_seq_len, - cache=uncond_cache, - ) - if uncond_cache is not None - else None - ) - - def run_branch( - pack: PackedSequence, - noise_vision: torch.Tensor, # [1,C,V*chunk_len,H,W] - timestep: torch.Tensor, # [1,1] - branch: _ARBranch, - ) -> torch.Tensor: # [1,C,V*chunk_len,H,W] - OmniMoTCausalModel._set_ar_vision_noise(self, pack, noise_vision, timestep) - cfgp_enabled = self.parallel_dims is not None and self.parallel_dims.cfgp_enabled - # Under CFGP each rank has one branch-local cache in ``cond_memory``. - memory = cond_memory if cfgp_enabled or branch == "conditional" else uncond_memory - assert memory is not None - output = self.denoise(data_batch_packed=pack, memory=memory) - return torch.stack(output["preds_vision"]) # [1,C,V*chunk_len,H,W] - - def velocity_fn(noise_x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: - return OmniMoTCausalModel._predict_ar_velocity_with_cfg( - self, - noise_x=noise_x, - timestep=timestep, - vision_shape=curr_vision_latent.shape, - packed_seq=cond_pack, - packed_seq_uncond=uncond_pack, - guidance=guidance, - normalize_cfg=normalize_cfg, - run_branch=run_branch, - ) # [1,N_tokens_flat] - - initial_noise = curr_vision_latent.flatten(start_dim=1) # [1,N_tokens_flat] - denoised_flat = OmniMoTCausalModel._run_ar_sampler( + return generate_multiview_transfer_ar_chunk( self, - velocity_fn, - initial_noise, - sampler_mode=sampler_mode, + cond_pack=cond_pack, + uncond_pack=uncond_pack, + cond_cache=cond_cache, + uncond_cache=uncond_cache, + curr_vision_latent=curr_vision_latent, + guidance=guidance, num_steps=num_steps, shift=shift, seed=seed, - sample_idx=chunk_start, + chunk_start=chunk_start, num_frames=num_frames, + normalize_cfg=normalize_cfg, + sampler_mode=sampler_mode, distilled_num_steps=distilled_num_steps, - ) # [1,N_tokens_flat] - return denoised_flat.reshape(curr_vision_latent.shape) # [1,C,V*chunk_len,H,W] + memory_seq_len=memory_seq_len, + ) @torch.no_grad() def _iter_samples_multiview_transfer_autoregressive( @@ -2873,261 +2883,22 @@ def _iter_samples_multiview_transfer_autoregressive( has_negative_prompt: bool, ) -> Generator[dict[str, Any], torch.Tensor | None, None]: """Generate the target for a two-item camera-major multiview transfer sample.""" - if not self._uses_multiview_replay_kv(): - raise ValueError("Multiview transfer AR requires replayed two-way Flex teacher-forcing configuration.") - if self.config.action_gen or self.config.sound_gen: - raise ValueError("Multiview transfer AR supports vision-only models.") - sequence_plans = build_sequence_plans_from_data_batch( + yield from iter_samples_multiview_transfer_autoregressive( + self, data_batch=data_batch, - input_video_key=self.input_video_key, - input_image_key=self.input_image_key, - ) - if len(sequence_plans) != 1: - raise ValueError(f"Multiview transfer AR requires batch_size=1, got {len(sequence_plans)} samples.") - caption_groups = _per_view_caption_groups(data_batch[self.input_caption_key]) - self._apply_inference_caption_plan(sequence_plans, caption_groups) - vision_condition_indexes = [sequence_plans[0].condition_frame_indexes_vision] - gen_data_clean = self.get_data_and_condition( - data_batch, - vision_condition_indexes=vision_condition_indexes, - retain_raw_state_vision=False, - ) - self._release_inference_raw_vision(data_batch, gen_data_clean) - if gen_data_clean.x0_tokens_vision is None or len(gen_data_clean.x0_tokens_vision) != 2: - num_items = 0 if gen_data_clean.x0_tokens_vision is None else len(gen_data_clean.x0_tokens_vision) - raise ValueError(f"Multiview transfer AR requires [control, target], got {num_items} vision items.") - # Match joint inference: all CP shards must read the owner's encoded - # controls and conditioned targets, even when local VAE results differ. - if self._get_teacher_forcing_kv_implementation() == "multiview_maskless_kv": - context_parallel_broadcast_tensor_list( - gen_data_clean.x0_tokens_vision, self.parallel_dims - ) # each [1,C,V*T,H,W] - if gen_data_clean.num_vision_items_per_sample != [2]: - raise ValueError( - "Multiview transfer AR requires one sample with exactly two vision items; " - f"got {gen_data_clean.num_vision_items_per_sample}." - ) - if gen_data_clean.num_views_per_vision_item is None or len(gen_data_clean.num_views_per_vision_item) != 2: - raise ValueError("Multiview transfer AR requires per-camera VAE metadata for both vision items.") - num_views = gen_data_clean.num_views_per_vision_item[0] - if gen_data_clean.num_views_per_vision_item != [num_views, num_views]: - raise ValueError( - f"Control and target must use the same views, got {gen_data_clean.num_views_per_vision_item}." - ) - control_latent, target_latent = gen_data_clean.x0_tokens_vision - if control_latent.shape != target_latent.shape: - raise ValueError( - f"Control and target latent shapes must match, got {control_latent.shape} and {target_latent.shape}." - ) - if target_latent.shape[2] % num_views != 0: - raise ValueError(f"Target latent_t={target_latent.shape[2]} is not divisible by {num_views} views.") - frames_per_view = target_latent.shape[2] // num_views - output_frames = frames_per_view if max_num_frames is None else min(frames_per_view, max_num_frames) - if output_frames < 1: - raise ValueError(f"max_num_frames must leave at least one frame, got {max_num_frames}.") - if sync_num_frames_across_ranks and dist.is_available() and dist.is_initialized(): - frame_count = torch.tensor([output_frames], device=target_latent.device, dtype=torch.long) # [1] - dist.all_reduce(frame_count, op=dist.ReduceOp.MIN, group=sync_process_group) - output_frames = int(frame_count.item()) - - cond_text_tokens, uncond_text_tokens = self._get_inference_text_tokens( - data_batch, - has_negative_prompt, - caption_groups, - ) - cfgp_enabled = self.parallel_dims is not None and self.parallel_dims.cfgp_enabled - if cfgp_enabled: - seed = _broadcast_seed([seed], self.parallel_dims.cfgp_mesh.get_group(), self.parallel_dims.cfgp_rank)[0] - cfg_active = guidance != 1.0 or cfgp_enabled - if getattr(self.config, "rolling_kv_cache_chunks", None) is not None: - from cosmos_framework.model.generator.utils.rolling_kv.rolling_transfer_ar import ( - iter_rolling_transfer_ar, - ) - - if normalize_cfg and cfg_active: - raise ValueError("Rolling replay requires normalize_cfg=False when guidance is active.") - if sampler_mode not in ("rf", "distilled"): - raise ValueError("Rolling replay sampler_mode must be rf or distilled.") - for part in iter_rolling_transfer_ar( - self, - plans=sequence_plans, - data=gen_data_clean, - conditional_text=cond_text_tokens, - unconditional_text=uncond_text_tokens, - guidance=guidance, - seed=seed, - num_steps=num_steps, - shift=shift, - sampler_mode=sampler_mode, - distilled_num_steps=distilled_num_steps, - output_frames=output_frames, - on_clean_vision_chunk=on_clean_vision_chunk, - prompt_schedule=data_batch.get("rolling_prompt_schedule"), - ): - for frame in _iter_multiview_logical_frames(part["vision"], num_views=num_views): - yield {"vision": frame} - return - fps_vision = gen_data_clean.fps_vision.tolist() if gen_data_clean.fps_vision is not None else [24.0] - backend = self._make_multiview_transfer_ar_backend() - - def build_prefill_pack( - text_tokens: list[list[int]], - *, - materialized_target_frame_ranges: Sequence[tuple[int, int]] | None = None, - ) -> PackedSequence: - return backend.build_prefill_pack( - sequence_plans=sequence_plans, - gen_data_clean=gen_data_clean, - text_tokens=text_tokens, - materialized_target_frame_ranges=materialized_target_frame_ranges, - ) - - cond_prefill = build_prefill_pack( - cond_text_tokens, - materialized_target_frame_ranges=[], - ) - assert cond_prefill.vision is not None - target_condition_mask = cond_prefill.vision.condition_mask[1] # [V*T,1,1] - condition_count = _multiview_conditioned_prefix_length( - target_condition_mask, - num_views=num_views, - frames_per_view=frames_per_view, - ) - generated_target = target_latent.to(**self.tensor_kwargs).clone() # [1,C,V*T,H,W] - gen_data_clean.x0_tokens_vision[1] = generated_target - materialized_condition_count = min(condition_count, output_frames) - session = backend.create_session( - prefill_pack=cond_prefill, - num_views=num_views, - frames_per_view=frames_per_view, - condition_count=materialized_condition_count, - cfg_active=cfg_active, - cfgp_enabled=cfgp_enabled, - text_view_ids=sequence_plans[0].text_view_ids, - ) - - controls_read_rgb = self._get_teacher_forcing_replay_policy().controls_read_strict_past_clean_rgb - if not controls_read_rgb: - uncond_prefill = None - if cfg_active: - assert uncond_text_tokens is not None - uncond_prefill = build_prefill_pack( - uncond_text_tokens, - materialized_target_frame_ranges=[], - ) - backend.capture_control_cache( - session=session, - conditional_pack=cond_prefill, - unconditional_pack=uncond_prefill, - ) - - conditioned_prefix = _submit_multiview_conditioned_prefix( - generated_target, - num_views=num_views, - frames_per_view=frames_per_view, - condition_count=condition_count, - output_frames=output_frames, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + normalize_cfg=normalize_cfg, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + sync_num_frames_across_ranks=sync_num_frames_across_ranks, + sync_process_group=sync_process_group, + max_num_frames=max_num_frames, on_clean_vision_chunk=on_clean_vision_chunk, + has_negative_prompt=has_negative_prompt, ) - if conditioned_prefix is not None: - for prefix_frame in _iter_multiview_logical_frames(conditioned_prefix, num_views=num_views): - yield {"vision": prefix_frame} - for chunk_start, chunk_end in _iter_ar_chunk_ranges( - condition_count, output_frames, self.config.teacher_forcing_frames_per_chunk - ): - chunk_len = chunk_end - chunk_start - if controls_read_rgb: - backend.prime_control_cache( - session=session, - sequence_plans=sequence_plans, - gen_data_clean=gen_data_clean, - conditional_text_tokens=cond_text_tokens, - unconditional_text_tokens=uncond_text_tokens, - ) - memory_layout = backend.build_memory_layout(session) - noise_generator = torch.Generator(device=target_latent.device).manual_seed(seed + chunk_start) - chunk_noise = torch.empty( - ( - 1, - target_latent.shape[1], - num_views * chunk_len, - target_latent.shape[3], - target_latent.shape[4], - ), - device=target_latent.device, - dtype=self.tensor_kwargs["dtype"], - ).normal_(generator=noise_generator) # [1,C,V*chunk_len,H,W] - cond_pack = backend.build_current_pack( - vision_latent=chunk_noise, - text_tokens=cond_text_tokens, - text_view_ids=session.text_view_ids, - vision_view_ids=session.target_view_ids, - fps_vision=fps_vision, - num_views=num_views, - frames_per_view=frames_per_view, - chunk_start=chunk_start, - memory_layout=memory_layout, - current_role="current_target", - ) - uncond_pack = ( - backend.build_current_pack( - vision_latent=chunk_noise, - text_tokens=uncond_text_tokens, - text_view_ids=session.text_view_ids, - vision_view_ids=session.target_view_ids, - fps_vision=fps_vision, - num_views=num_views, - frames_per_view=frames_per_view, - chunk_start=chunk_start, - memory_layout=memory_layout, - current_role="current_target", - ) - if cfg_active - else None - ) - denoised_chunk = self._generate_multiview_transfer_ar_chunk( - cond_pack=cond_pack, - uncond_pack=uncond_pack, - cond_cache=session.conditional_cache, - uncond_cache=session.unconditional_cache, - curr_vision_latent=chunk_noise, - guidance=guidance, - num_steps=num_steps, - shift=shift, - seed=seed, - chunk_start=chunk_start, - num_frames=output_frames, - normalize_cfg=normalize_cfg, - sampler_mode=sampler_mode, - distilled_num_steps=distilled_num_steps, - memory_seq_len=session.memory_seq_len, - ) - backend.scatter_chunk( - generated_target, - denoised_chunk, - num_views=num_views, - frames_per_view=frames_per_view, - chunk_start=chunk_start, - chunk_end=chunk_end, - ) - if on_clean_vision_chunk is not None: - on_clean_vision_chunk(denoised_chunk) - - if chunk_end < output_frames: - backend.commit_clean_chunk( - session=session, - denoised_chunk=denoised_chunk.to(**self.tensor_kwargs), # [1,C,V*chunk_len,H,W] - chunk_start=chunk_start, - chunk_end=chunk_end, - conditional_text_tokens=cond_text_tokens, - unconditional_text_tokens=uncond_text_tokens, - fps_vision=fps_vision, - ) - - # Expose progress after the chunk is ready for future AR steps. Each - # event represents one latent time step across every synchronized view. - for output_frame in _iter_multiview_logical_frames(denoised_chunk, num_views=num_views): - yield {"vision": output_frame} def _cast_ar_action_tokens_to_projection_dtype(self, packed_seq: PackedSequence) -> None: """Cast AR action tokens to the dtype expected by the action projection.""" diff --git a/cosmos_framework/model/generator/omni_mot_causal_model_test.py b/cosmos_framework/model/generator/omni_mot_causal_model_test.py index 66101cc24..2148ca934 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model_test.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model_test.py @@ -229,6 +229,12 @@ def test_teacher_forcing_kv_implementation_default_and_validation() -> None: ) with pytest.raises(ValueError): OmniMoTCausalModelConfig(teacher_forcing_kv_implementation="unknown") + assert ( + OmniMoTCausalModelConfig( + teacher_forcing_kv_implementation="multiview_threeway_kv" + ).teacher_forcing_kv_implementation + == "multiview_threeway_kv" + ) @pytest.mark.L0 @@ -294,6 +300,7 @@ def test_target_only_teacher_forcing_requires_one_logical_sample( model = MagicMock() model.config.video_temporal_causal = True model.config.teacher_forcing_target_only_no_text_pass2 = True + model._uses_robot_multiview_threeway.return_value = False model._uses_multiview_replay_kv.return_value = False packed_sequence = PackedSequence( sample_lens=sample_lens, @@ -463,7 +470,12 @@ def test_teacher_forcing_kv_implementation_validates_lazy_config_value() -> None """LazyConfig must not bypass validation for the public replay selector.""" from cosmos_framework.model.generator.omni_mot_causal_model import _resolve_teacher_forcing_kv_implementation - for implementation in ("multiview_flex_kv", "multiview_maskless_kv", "singleview_threeway_kv"): + for implementation in ( + "multiview_flex_kv", + "multiview_maskless_kv", + "multiview_threeway_kv", + "singleview_threeway_kv", + ): assert _resolve_teacher_forcing_kv_implementation(implementation) == implementation with pytest.raises(ValueError, match="teacher_forcing_kv_implementation"): _resolve_teacher_forcing_kv_implementation("unknown") @@ -691,6 +703,7 @@ def test_target_only_replay_accepts_pack_after_control_dropout(self, dropout_rat model = self._make_model(dropout_rate) model.config.video_temporal_causal = True model.config.teacher_forcing_target_only_no_text_pass2 = True + model._uses_robot_multiview_threeway.return_value = False model._uses_multiview_replay_kv.return_value = False gen_data_clean = self._make_data() result = OmniMoTCausalModel._maybe_drop_teacher_forcing_transfer_control( @@ -1074,94 +1087,38 @@ def test_chunkwise_tf_truncates_framewise_action_domain_ids_with_real_actions() @pytest.mark.L0 @pytest.mark.CPU -def test_multiview_clean_callback_includes_partial_conditioned_prefix() -> None: - """A partial prefix is submitted once in camera-major view order.""" - from cosmos_framework.model.generator.omni_mot_causal_model import ( - _submit_multiview_conditioned_prefix, - ) - - target_latent = torch.arange(8, dtype=torch.float32).reshape(1, 1, 8, 1, 1) # [B,C,V*T,H,W] - callback_chunks: list[torch.Tensor] = [] - - _submit_multiview_conditioned_prefix( - target_latent, - num_views=2, - frames_per_view=4, - condition_count=2, - output_frames=4, - on_clean_vision_chunk=callback_chunks.append, - ) - - expected = torch.tensor([0.0, 1.0, 4.0, 5.0]).reshape(1, 1, 4, 1, 1) # [B,C,V*T_prefix,H,W] - assert len(callback_chunks) == 1 - torch.testing.assert_close(callback_chunks[0], expected) - - -@pytest.mark.L0 -@pytest.mark.CPU -def test_multiview_transfer_ar_rejects_sparse_condition_frames() -> None: - """Sparse target conditions must not be misread as a ground-truth prefix.""" - from cosmos_framework.model.generator.omni_mot_causal_model import ( - _multiview_conditioned_prefix_length, - ) - - sparse_mask = torch.tensor([1, 0, 1, 1, 0, 1], dtype=torch.bool) # [V*T] - - with pytest.raises(ValueError, match="contiguous prefix"): - _multiview_conditioned_prefix_length( - sparse_mask, - num_views=2, - frames_per_view=3, - ) - +def test_chunkwise_tf_truncates_each_camera_major_action_view() -> None: + """Each replicated action view is cropped instead of treating V*T as one timeline.""" + from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel -@pytest.mark.L0 -@pytest.mark.CPU -def test_multiview_clean_callback_includes_fully_conditioned_output() -> None: - """A fully conditioned target still submits one decodable callback chunk.""" - from cosmos_framework.model.generator.omni_mot_causal_model import ( - _submit_multiview_conditioned_prefix, + model = object.__new__(OmniMoTCausalModel) + model.config = SimpleNamespace( + teacher_forcing_frames_per_chunk=2, + causal_training_strategy="teacher_forcing", ) - - target_latent = torch.arange(8, dtype=torch.float32).reshape(1, 1, 8, 1, 1) # [B,C,V*T,H,W] - callback_chunks: list[torch.Tensor] = [] - - _submit_multiview_conditioned_prefix( - target_latent, - num_views=2, - frames_per_view=4, - condition_count=4, - output_frames=4, - on_clean_vision_chunk=callback_chunks.append, + model.tokenizer_vision_gen = SimpleNamespace( + temporal_compression_factor=4, + get_pixel_num_frames=lambda latent_frames: 1 + (latent_frames - 1) * 4, ) - - assert len(callback_chunks) == 1 - torch.testing.assert_close(callback_chunks[0], target_latent) - - -@pytest.mark.L0 -@pytest.mark.CPU -def test_multiview_clean_callback_truncates_conditioned_prefix_to_output_frames() -> None: - """A fully conditioned target honors max_num_frames before decode submission.""" - from cosmos_framework.model.generator.omni_mot_causal_model import ( - _submit_multiview_conditioned_prefix, + actions = torch.arange(3 * 900 * 2, dtype=torch.float32).reshape(3 * 900, 2) + domain_ids = torch.arange(3 * 900, dtype=torch.long) + gen_data = SimpleNamespace( + is_image_batch=False, + x0_tokens_vision=[torch.zeros(1, 2, 226, 1, 1) for _ in range(3)], + raw_state_vision=[torch.zeros(1, 3, 901, 1, 1) for _ in range(3)], + x0_tokens_action=[actions.clone()], + action_domain_id=[domain_ids.clone()], + num_views_per_action_item=[3], ) - target_latent = torch.arange(8, dtype=torch.float32).reshape(1, 1, 8, 1, 1) # [B,C,V*T,H,W] - callback_chunks: list[torch.Tensor] = [] - - _submit_multiview_conditioned_prefix( - target_latent, - num_views=2, - frames_per_view=4, - condition_count=4, - output_frames=2, - on_clean_vision_chunk=callback_chunks.append, - ) + result = OmniMoTCausalModel._truncate_for_chunkwise_tf(model, gen_data) - expected = torch.tensor([0.0, 1.0, 4.0, 5.0]).reshape(1, 1, 4, 1, 1) # [B,C,V*T_output,H,W] - assert len(callback_chunks) == 1 - torch.testing.assert_close(callback_chunks[0], expected) + expected_actions = actions.reshape(3, 900, 2)[:, :896].reshape(3 * 896, 2) + expected_domains = domain_ids.reshape(3, 900)[:, :896].reshape(3 * 896) + assert all(vision.shape[2] == 225 for vision in result.x0_tokens_vision) + assert all(video.shape[-3] == 897 for video in result.raw_state_vision) + torch.testing.assert_close(result.x0_tokens_action[0], expected_actions) + torch.testing.assert_close(result.action_domain_id[0], expected_domains) @pytest.mark.L0 @@ -1318,7 +1275,7 @@ def generate_chunk(**kwargs: object) -> torch.Tensor: # returns [B,C,V*T_chunk, with ( patch( - "cosmos_framework.model.generator.omni_mot_causal_model.build_sequence_plans_from_data_batch", + "cosmos_framework.model.generator.utils.multiview_ar.build_sequence_plans_from_data_batch", return_value=[sequence_plan], ), patch( @@ -1459,8 +1416,10 @@ def test_multiview_transfer_ar_pack_sets_metadata_and_aligned_view_positions() - def test_multiview_transfer_prefill_is_clean_without_extending_conditioned_prefix(condition_count: int) -> None: """Recomputed RGB history omits diffusion embeddings without changing the rollout prefix.""" from cosmos_framework.data.generator.sequence_packing import ModalityData - from cosmos_framework.model.generator.omni_mot_causal_model import _multiview_conditioned_prefix_length - from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend + from cosmos_framework.model.generator.utils.multiview_ar import ( + MultiviewTransferARBackend, + multiview_conditioned_prefix_length, + ) target_mask = (torch.arange(5) < condition_count).repeat(2).reshape(10, 1, 1) # [V*T,1,1] control_mask = torch.ones_like(target_mask) # [V*T,1,1] @@ -1495,7 +1454,7 @@ def test_multiview_transfer_prefill_is_clean_without_extending_conditioned_prefi assert result.vision.mse_loss_indexes.numel() == 0 assert result.vision.timesteps.numel() == 0 assert all(index.numel() == 0 for index in result.vision.noisy_frame_indexes) - actual_count = _multiview_conditioned_prefix_length(result.vision.condition_mask[1], num_views=2, frames_per_view=5) + actual_count = multiview_conditioned_prefix_length(result.vision.condition_mask[1], num_views=2, frames_per_view=5) assert actual_count == condition_count session = backend.create_session( prefill_pack=result, @@ -1867,7 +1826,7 @@ def get_text_tokens( model._get_inference_text_tokens.side_effect = get_text_tokens with patch( - "cosmos_framework.model.generator.omni_mot_causal_model.build_sequence_plans_from_data_batch", + "cosmos_framework.model.generator.utils.multiview_ar.build_sequence_plans_from_data_batch", return_value=[sequence_plan], ): iterator = OmniMoTCausalModel._iter_samples_multiview_transfer_autoregressive( diff --git a/cosmos_framework/model/generator/tokenizers/lidar/geometry.py b/cosmos_framework/model/generator/tokenizers/lidar/geometry.py new file mode 100644 index 000000000..7ab0f6b31 --- /dev/null +++ b/cosmos_framework/model/generator/tokenizers/lidar/geometry.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Calibrated geometry helpers for Pandar128 LiDAR range maps.""" + +from __future__ import annotations + +import numpy as np + +from cosmos_framework.model.generator.tokenizers.lidar.range_projection import LidarRangeProjectionConfig + +# Native beam order used by the MADS / LiDARGen 128x3600 range-map format. +PANDAR128_ELEVATIONS_DEG = ( + 14.436, + 13.535, + 13.082, + 12.624, + 12.165, + 11.702, + 11.239, + 10.771, + 10.305, + 9.83, + 9.356, + 8.88, + 8.401, + 7.921, + 7.438, + 6.953, + 6.467, + 5.978, + 5.487, + 4.996, + 4.501, + 4.007, + 3.509, + 3.013, + 2.512, + 2.013, + 1.885, + 1.761, + 1.637, + 1.511, + 1.386, + 1.258, + 1.13, + 1.008, + 0.88, + 0.756, + 0.63, + 0.505, + 0.379, + 0.251, + 0.124, + 0.0, + -0.129, + -0.254, + -0.38, + -0.506, + -0.632, + -0.76, + -0.887, + -1.012, + -1.141, + -1.266, + -1.393, + -1.519, + -1.646, + -1.773, + -1.901, + -2.027, + -2.155, + -2.282, + -2.409, + -2.535, + -2.663, + -2.789, + -2.916, + -3.044, + -3.172, + -3.299, + -3.425, + -3.552, + -3.68, + -3.806, + -3.933, + -4.062, + -4.19, + -4.318, + -4.444, + -4.571, + -4.699, + -4.824, + -4.951, + -5.081, + -5.209, + -5.336, + -5.463, + -5.589, + -5.718, + -5.843, + -5.968, + -6.1, + -6.607, + -7.117, + -7.624, + -8.134, + -8.64, + -9.149, + -9.652, + -10.16, + -10.665, + -11.17, + -11.672, + -12.174, + -12.673, + -13.173, + -13.67, + -14.166, + -14.66, + -15.154, + -15.645, + -16.135, + -16.622, + -17.106, + -17.592, + -18.072, + -18.548, + -19.03, + -19.501, + -19.978, + -20.445, + -20.918, + -21.379, + -21.848, + -22.304, + -22.768, + -23.219, + -23.678, + -24.123, + -25.016, +) + + +def _elevation_radians() -> np.ndarray: # [H] + """Beam elevations of the calibrated Pandar128 rows, in radians.""" + return np.deg2rad(np.asarray(PANDAR128_ELEVATIONS_DEG, dtype=np.float64)) + + +def _azimuth_radians(projection: LidarRangeProjectionConfig) -> np.ndarray: # [W] + """Column azimuths of a semantic-width range map, in radians.""" + return np.deg2rad( + np.linspace( + projection.azimuth_start_degrees, + projection.azimuth_end_degrees, + projection.semantic_width, + endpoint=projection.azimuth_endpoint, + dtype=np.float64, + ) + ) + + +def pandar128_ray_directions( + *, + range_projection: LidarRangeProjectionConfig, +) -> np.ndarray: # [H,W,3] + """Unit ray direction for every cell of a semantic-width range map. + + Shares the elevation calibration and ``π → -π`` azimuth convention with + :func:`range_map_to_xyz`, so scaling these directions by a metric range + reproduces the same sensor-frame XYZ that unprojection produces -- but + densely, and as a plain multiplication a loss can backpropagate through. + """ + if range_projection.semantic_height != len(PANDAR128_ELEVATIONS_DEG): + raise ValueError( + "Range-map height must match the calibrated Pandar128 elevations, got " + f"{range_projection.semantic_height} != {len(PANDAR128_ELEVATIONS_DEG)}" + ) + elevation = _elevation_radians()[:, None] # [H,1] + azimuth = _azimuth_radians(range_projection)[None, :] # [1,W] + cos_elevation = np.cos(elevation) # [H,1] + x = cos_elevation * np.cos(azimuth) # [H,W] + y = cos_elevation * np.sin(azimuth) # [H,W] + z = np.broadcast_to(np.sin(elevation), x.shape) # [H,W] + return np.stack((x, y, z), axis=-1).astype(np.float32) # [H,W,3] + + +def range_map_to_xyz( + metric_range: np.ndarray, + valid: np.ndarray, + *, + range_projection: LidarRangeProjectionConfig | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """Unproject a Pandar128 range map to sensor-frame XYZ. + + Returns ``xyz`` with X forward, Y left, Z up and the flattened indices of + valid rays. Azimuth follows the tested LiDARGen convention, ``π → -π``. + When ``range_projection`` is omitted the map width is treated as the + azimuth canvas, so BEV helpers can unproject already-resized clips. + """ + if metric_range.ndim != 2 or valid.shape != metric_range.shape: + raise ValueError( + f"Expected metric/valid [H,W] with matching shapes, got {metric_range.shape} and {valid.shape}" + ) + # A map that already sits at display width is its own azimuth canvas. + # Matching native_width to W lets BEV helpers unproject toy/downsampled + # clips (16, 64, ...) that are not divisors of the 3600-column native grid. + width = int(metric_range.shape[1]) + projection = range_projection or LidarRangeProjectionConfig( + native_width=width, + semantic_width=width, + model_width=width, + ) + if metric_range.shape != (projection.semantic_height, projection.semantic_width): + raise ValueError( + "Range-map shape must match the configured Pandar128 elevations and semantic projection, " + f"got {metric_range.shape} != {(projection.semantic_height, projection.semantic_width)}" + ) + + mask = valid.astype(bool) & np.isfinite(metric_range) & (metric_range > 0) + if not mask.any(): + return np.zeros((0, 3), dtype=np.float32), np.zeros((0,), dtype=np.int64) + + elevation = _elevation_radians()[:, None] + azimuth = _azimuth_radians(projection)[None, :] + phi = np.broadcast_to(elevation, metric_range.shape)[mask] + theta = np.broadcast_to(azimuth, metric_range.shape)[mask] + distance = metric_range[mask].astype(np.float64) + cos_phi = np.cos(phi) + xyz = np.stack( + ( + distance * cos_phi * np.cos(theta), + distance * cos_phi * np.sin(theta), + distance * np.sin(phi), + ), + axis=1, + ).astype(np.float32) + return xyz, np.flatnonzero(mask) diff --git a/cosmos_framework/model/generator/utils/camera_rays.py b/cosmos_framework/model/generator/utils/camera_rays.py new file mode 100644 index 000000000..e74a545d3 --- /dev/null +++ b/cosmos_framework/model/generator/utils/camera_rays.py @@ -0,0 +1,214 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""Unproject native calibration pixels to camera rays for every supported lens model. + +Pixel coordinates use the calibration's own integer-center convention. Every inverse is +checked by reprojection so a token center outside a lens model's invertible domain fails +closed instead of producing a plausible but wrong ray. +""" + +import math +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import Any + +import torch + +from cosmos_framework.model.generator.utils.ftheta import ftheta_unproject + +# Lens models RigRoPE can turn into rays; ``_validate_calibration`` owns their coefficient layouts. +SUPPORTED_RAY_CAMERA_MODELS = frozenset({"ftheta", "pinhole", "fisheye"}) +_NEWTON_ITERATIONS = 40 +_REPROJECTION_TOLERANCE = 1e-9 +_JACOBIAN_STEP = 1e-7 + + +def _matmul3(a: Sequence[Sequence[float]], b: Sequence[Sequence[float]]) -> list[list[float]]: + return [[sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)] for i in range(3)] + + +@lru_cache(maxsize=64) +def _opencv_tilt_matrix(tau_x: float, tau_y: float) -> tuple[tuple[float, ...], ...]: # returns [3,3] + """OpenCV's tilted-sensor projection (``computeTiltProjectionMatrix``) for angles in radians.""" + cx, sx, cy, sy = math.cos(tau_x), math.sin(tau_x), math.cos(tau_y), math.sin(tau_y) + rotation = _matmul3( + [[cy, 0.0, -sy], [0.0, 1.0, 0.0], [sy, 0.0, cy]], [[1.0, 0.0, 0.0], [0.0, cx, sx], [0.0, -sx, cx]] + ) + projection = [ + [rotation[2][2], 0.0, -rotation[0][2]], + [0.0, rotation[2][2], -rotation[1][2]], + [0.0, 0.0, 1.0], + ] + return tuple(tuple(row) for row in _matmul3(projection, rotation)) + + +def _opencv_distort(points: torch.Tensor, coefficients: Sequence[float]) -> torch.Tensor: # [...,2] -> [...,2] + """Apply OpenCV distortion (4, 5, 8, 12, or 14 coefficients) to normalized points. + + The layout is ``k1 k2 p1 p2 [k3 [k4 k5 k6 [s1 s2 s3 s4 [tau_x tau_y]]]]``: radial, + tangential, thin-prism and tilted-sensor terms, as in ``cv2.projectPoints``. + """ + k1, k2, p1, p2, k3, k4, k5, k6, s1, s2, s3, s4, tau_x, tau_y = (*coefficients, *([0.0] * 10))[:14] + x, y = points[..., 0], points[..., 1] # [...], [...] + r2 = x * x + y * y # [...] + r4 = r2 * r2 # [...] + radial = (1 + r2 * (k1 + r2 * (k2 + r2 * k3))) / (1 + r2 * (k4 + r2 * (k5 + r2 * k6))) # [...] + xd = x * radial + 2 * p1 * x * y + p2 * (r2 + 2 * x * x) + s1 * r2 + s2 * r4 # [...] + yd = y * radial + p1 * (r2 + 2 * y * y) + 2 * p2 * x * y + s3 * r2 + s4 * r4 # [...] + if tau_x or tau_y: + tilt = _opencv_tilt_matrix(tau_x, tau_y) + depth = tilt[2][0] * xd + tilt[2][1] * yd + tilt[2][2] # [...] + xd, yd = ( + (tilt[0][0] * xd + tilt[0][1] * yd + tilt[0][2]) / depth, + (tilt[1][0] * xd + tilt[1][1] * yd + tilt[1][2]) / depth, + ) # [...], [...] + return torch.stack((xd, yd), dim=-1) # [...,2] + + +def _opencv_undistort(distorted: torch.Tensor, coefficients: Sequence[float]) -> torch.Tensor: # [...,2] -> [...,2] + """Invert OpenCV distortion by Newton iteration from the distorted point.""" + if len(coefficients) not in (4, 5, 8, 12, 14): + raise ValueError( + f"RigRoPE supports OpenCV distortion with 4, 5, 8, 12, or 14 coefficients, got {len(coefficients)}" + ) + points = distorted.clone() # [...,2] + step_x = distorted.new_tensor([_JACOBIAN_STEP, 0.0]) # [2] + step_y = distorted.new_tensor([0.0, _JACOBIAN_STEP]) # [2] + for _ in range(_NEWTON_ITERATIONS): + residual = _opencv_distort(points, coefficients) - distorted # [...,2] + column_x = (_opencv_distort(points + step_x, coefficients) - _opencv_distort(points - step_x, coefficients)) / ( + 2 * _JACOBIAN_STEP + ) # [...,2] + column_y = (_opencv_distort(points + step_y, coefficients) - _opencv_distort(points - step_y, coefficients)) / ( + 2 * _JACOBIAN_STEP + ) # [...,2] + det = column_x[..., 0] * column_y[..., 1] - column_y[..., 0] * column_x[..., 1] # [...] + if (det <= 0).any(): + raise ValueError("Token centers fall outside the invertible OpenCV distortion domain") + dx = (column_y[..., 1] * residual[..., 0] - column_y[..., 0] * residual[..., 1]) / det # [...] + dy = (column_x[..., 0] * residual[..., 1] - column_x[..., 1] * residual[..., 0]) / det # [...] + points = points - torch.stack((dx, dy), dim=-1) # [...,2] + error = torch.linalg.vector_norm(_opencv_distort(points, coefficients) - distorted, dim=-1) # [...] + if not torch.isfinite(points).all() or (error > _REPROJECTION_TOLERANCE).any(): + raise ValueError("OpenCV undistortion did not converge at the token centers") + return points # [...,2] + + +def _equidistant_unproject(distorted: torch.Tensor, coefficients: Sequence[float]) -> torch.Tensor: + """Invert OpenCV fisheye ``theta_d = theta (1 + k1 theta^2 + ... + k4 theta^8)``; [...,2] -> [...,3].""" + if len(coefficients) != 4: + raise ValueError(f"Equidistant fisheye calibration requires 4 coefficients, got {len(coefficients)}") + k1, k2, k3, k4 = coefficients + theta_d = torch.linalg.vector_norm(distorted, dim=-1) # [...] + theta = theta_d.clone() # [...] + for _ in range(_NEWTON_ITERATIONS): + t2 = theta * theta # [...] + value = theta * (1 + t2 * (k1 + t2 * (k2 + t2 * (k3 + t2 * k4)))) - theta_d # [...] + slope = 1 + t2 * (3 * k1 + t2 * (5 * k2 + t2 * (7 * k3 + t2 * 9 * k4))) # [...] + if (slope <= 0).any(): + raise ValueError("Token centers fall outside the monotonic equidistant fisheye domain") + theta = theta - value / slope # [...] + t2 = theta * theta # [...] + error = (theta * (1 + t2 * (k1 + t2 * (k2 + t2 * (k3 + t2 * k4)))) - theta_d).abs() # [...] + if not torch.isfinite(theta).all() or (error > _REPROJECTION_TOLERANCE).any(): + raise ValueError("Equidistant fisheye inversion did not converge at the token centers") + if (theta < 0).any() or (theta >= math.pi).any(): + raise ValueError("Equidistant fisheye token ray angles must lie in [0, pi)") + scale = torch.sin(theta) / theta_d.clamp_min(1e-15) # [...] + rays = torch.cat((distorted * scale[..., None], torch.cos(theta)[..., None]), dim=-1) # [...,3] + optical_axis = rays.new_tensor([0.0, 0.0, 1.0]) # [3] + return torch.where((theta_d <= 1e-12)[..., None], optical_axis, rays) # [...,3] + + +def unproject_calibration(pixels: torch.Tensor, calibration: Mapping[str, Any]) -> torch.Tensor: + """Return unit camera rays (OpenCV axes) for native pixels; pixels [...,2] -> rays [...,3].""" + if pixels.shape[-1] != 2 or not pixels.is_floating_point() or not torch.isfinite(pixels).all(): + raise ValueError("Calibration pixels must be finite floating-point coordinates ending in two values") + model = calibration["camera_model"] + coefficients = [float(value) for value in calibration["distortion_coefficients"]] + if model == "ftheta": + return ftheta_unproject(pixels, coefficients) # [...,3] + original_dtype = pixels.dtype + pixels = pixels.double() # [...,2] + focal = pixels.new_tensor([calibration["fx_px"], calibration["fy_px"]]) # [2] + center = pixels.new_tensor([calibration["cx_px"], calibration["cy_px"]]) # [2] + distorted = (pixels - center) / focal # [...,2] + if model == "pinhole": + distortion = calibration["distortion_model"] + if distortion == "none": + points = distorted # [...,2] + elif distortion == "opencv": + points = _opencv_undistort(distorted, coefficients) # [...,2] + else: + raise ValueError(f"Unsupported pinhole distortion model for RigRoPE rays: {distortion}") + rays = torch.cat((points, torch.ones_like(points[..., :1])), dim=-1) # [...,3] + rays = rays / torch.linalg.vector_norm(rays, dim=-1, keepdim=True) # [...,3] + elif model == "fisheye": + if calibration["distortion_model"] != "equidistant": + raise ValueError( + f"Unsupported fisheye distortion model for RigRoPE rays: {calibration['distortion_model']}" + ) + rays = _equidistant_unproject(distorted, coefficients) # [...,3] + else: + raise ValueError(f"Unsupported camera model for RigRoPE rays: {model}") + return rays.to(dtype=original_dtype) # [...,3] + + +def token_pixel_grid(height: int, width: int, stride: int) -> torch.Tensor: # returns [H,W,3] + """Homogeneous centers of ``stride``-pixel tokens in integer-center pixel coordinates.""" + ys = (torch.arange(height, dtype=torch.float64) + 0.5) * stride - 0.5 # [H] + xs = (torch.arange(width, dtype=torch.float64) + 0.5) * stride - 0.5 # [W] + yy, xx = torch.meshgrid(ys, xs, indexing="ij") # [H,W], [H,W] + return torch.stack((xx, yy, torch.ones_like(xx)), dim=-1) # [H,W,3] + + +def _calibration_key(calibration: Mapping[str, Any]) -> tuple[Any, ...]: + return ( + calibration["camera_model"], + calibration["distortion_model"], + tuple(float(value) for value in calibration["distortion_coefficients"]), + *(float(calibration[key]) for key in ("fx_px", "fy_px", "cx_px", "cy_px")), + ) + + +@lru_cache(maxsize=256) +def _validate_cached( + key: tuple[Any, ...], affine_values: tuple[float, ...], image_hw: tuple[int, int], stride: int +) -> None: + model, distortion, coefficients, fx, fy, cx, cy = key + calibration = { + "camera_model": model, + "distortion_model": distortion, + "distortion_coefficients": coefficients, + "fx_px": fx, + "fy_px": fy, + "cx_px": cx, + "cy_px": cy, + } + height, width = image_hw + pixels = token_pixel_grid(math.ceil(height / stride), math.ceil(width / stride), stride) # [H,W,3] + affine = torch.tensor(affine_values, dtype=torch.float64).reshape(3, 3) # [3,3] + if not torch.isfinite(affine).all() or abs(float(torch.linalg.det(affine))) < 1e-12: + raise ValueError("image_from_calibration must be finite and invertible") + native = pixels @ torch.linalg.inv(affine).T # [H,W,3] + unproject_calibration(native[..., :2] / native[..., 2:], calibration) # [H,W,3] + + +def validate_token_geometry( + calibration: Mapping[str, Any], + image_from_calibration: torch.Tensor, # [3,3] + *, + image_hw: tuple[int, int], + stride: int, +) -> None: + """Validate the pre-interpolation latent-center rays RigRoPE will consume.""" + if calibration["camera_model"] not in SUPPORTED_RAY_CAMERA_MODELS: + raise ValueError(f"Unsupported camera model for RigRoPE rays: {calibration['camera_model']}") + _validate_cached( + _calibration_key(calibration), + tuple(float(value) for value in image_from_calibration.flatten()), + image_hw, + stride, + ) diff --git a/cosmos_framework/model/generator/utils/camera_rays_test.py b/cosmos_framework/model/generator/utils/camera_rays_test.py new file mode 100644 index 000000000..f07d57732 --- /dev/null +++ b/cosmos_framework/model/generator/utils/camera_rays_test.py @@ -0,0 +1,130 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +from typing import Any + +import cv2 +import numpy as np +import pytest +import torch + +from cosmos_framework.model.generator.utils.camera_rays import ( + token_pixel_grid, + unproject_calibration, + validate_token_geometry, +) +from cosmos_framework.model.generator.utils.ftheta import ftheta_unproject + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + +# DeformCam360 and GrandTour sidecar values. +_OPENCV_COEFFICIENTS = [-0.38748353719711304, 0.11847664415836334, -0.00010451977868797258, -0.0036122221499681473] +_FISHEYE_COEFFICIENTS = [-0.03843622792783702, -0.002634177248264113, 0.002422949539367233, -0.001382691294699419] +# OpenCV's 12- and 14-term layouts append thin-prism s1..s4, then the sensor tilt (tau_x, tau_y) in radians. +_OPENCV_RATIONAL_COEFFICIENTS = [*_OPENCV_COEFFICIENTS, 0.01, 0.002, -0.001, 0.0005] +_OPENCV_THIN_PRISM_COEFFICIENTS = [*_OPENCV_RATIONAL_COEFFICIENTS, 0.002, -0.001, 0.0015, -0.0005] +_OPENCV_TILTED_COEFFICIENTS = [*_OPENCV_THIN_PRISM_COEFFICIENTS, 0.01, -0.02] + + +def _calibration(camera_model: str, distortion_model: str, coefficients: list[float]) -> dict[str, Any]: + return { + "camera_model": camera_model, + "distortion_model": distortion_model, + "distortion_coefficients": coefficients, + "fx_px": 700.0, + "fy_px": 690.0, + "cx_px": 640.0, + "cy_px": 360.0, + } + + +def _camera_points() -> np.ndarray: # returns [N,3] + theta = np.linspace(0.0, 0.9, 7) # [A] + phi = np.linspace(0.0, 2 * np.pi, 9, endpoint=False) # [B] + tt, pp = np.meshgrid(theta, phi, indexing="ij") # [A,B], [A,B] + return np.stack((np.sin(tt) * np.cos(pp), np.sin(tt) * np.sin(pp), np.cos(tt)), axis=-1).reshape(-1, 3) # [N,3] + + +def _matrix(calibration: dict[str, Any]) -> np.ndarray: # returns [3,3] + return np.array( + [[calibration["fx_px"], 0, calibration["cx_px"]], [0, calibration["fy_px"], calibration["cy_px"]], [0, 0, 1]] + ) + + +@pytest.mark.parametrize( + ("camera_model", "distortion_model", "coefficients"), + [ + ("pinhole", "none", []), + ("pinhole", "opencv", _OPENCV_COEFFICIENTS), + ("pinhole", "opencv", [*_OPENCV_COEFFICIENTS, 0.01]), + ("pinhole", "opencv", _OPENCV_RATIONAL_COEFFICIENTS), + ("pinhole", "opencv", _OPENCV_THIN_PRISM_COEFFICIENTS), + ("pinhole", "opencv", _OPENCV_TILTED_COEFFICIENTS), + ("fisheye", "equidistant", _FISHEYE_COEFFICIENTS), + ], +) +def test_unprojection_inverts_opencv_projection( + camera_model: str, distortion_model: str, coefficients: list[float] +) -> None: + calibration = _calibration(camera_model, distortion_model, coefficients) + points = _camera_points() # [N,3] + distortion = np.array(coefficients or [0.0, 0.0, 0.0, 0.0], dtype=np.float64) # [K] + if camera_model == "fisheye": + pixels, _ = cv2.fisheye.projectPoints( + points[:, None], np.zeros(3), np.zeros(3), _matrix(calibration), distortion + ) + else: + pixels, _ = cv2.projectPoints(points, np.zeros(3), np.zeros(3), _matrix(calibration), distortion) + rays = unproject_calibration(torch.from_numpy(pixels[:, 0]), calibration) # [N,3] + torch.testing.assert_close(rays, torch.from_numpy(points), atol=1e-7, rtol=0) + + +def test_ftheta_dispatch_matches_native_unprojection() -> None: + coefficients = [16.0, 16.0, 32.0, 32.0, 0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0] + pixels = torch.tensor([[16.0, 16.0], [8.0, 24.0], [30.0, 2.0]], dtype=torch.float64) # [N,2] + calibration = {"camera_model": "ftheta", "distortion_model": "ftheta", "distortion_coefficients": coefficients} + torch.testing.assert_close(unproject_calibration(pixels, calibration), ftheta_unproject(pixels, coefficients)) + + +def test_optical_axis_and_principal_point_agree_across_models() -> None: + center = torch.tensor([[640.0, 360.0]], dtype=torch.float64) # [1,2] + for calibration in ( + _calibration("pinhole", "none", []), + _calibration("pinhole", "opencv", _OPENCV_COEFFICIENTS), + _calibration("fisheye", "equidistant", _FISHEYE_COEFFICIENTS), + ): + torch.testing.assert_close(unproject_calibration(center, calibration), torch.tensor([[0.0, 0.0, 1.0]]).double()) + + +def test_opencv_outside_invertible_domain_fails_closed() -> None: + calibration = _calibration("pinhole", "opencv", [-0.5, 0.0, 0.0, 0.0]) + # r - 0.5 r^3 peaks at 0.544; a distorted radius of 0.6 has no undistorted preimage. + pixels = torch.tensor([[640.0 + 0.6 * 700.0, 360.0]], dtype=torch.float64) # [1,2] + with pytest.raises(ValueError, match="OpenCV"): + unproject_calibration(pixels, calibration) + + +def test_equidistant_outside_monotonic_domain_fails_closed() -> None: + calibration = _calibration("fisheye", "equidistant", [-0.5, 0.0, 0.0, 0.0]) + pixels = torch.tensor([[640.0 + 1.5 * 700.0, 360.0]], dtype=torch.float64) # [1,2] + with pytest.raises(ValueError, match="fisheye"): + unproject_calibration(pixels, calibration) + + +def test_token_grid_uses_integer_pixel_centers() -> None: + grid = token_pixel_grid(2, 3, 16) # [2,3,3] + torch.testing.assert_close(grid[0, 0], torch.tensor([7.5, 7.5, 1.0]).double()) + torch.testing.assert_close(grid[1, 2], torch.tensor([39.5, 23.5, 1.0]).double()) + + +def test_token_geometry_validation_rejects_unsupported_and_out_of_domain_models() -> None: + affine = torch.eye(3, dtype=torch.float64) # [3,3] + validate_token_geometry(_calibration("pinhole", "none", []), affine, image_hw=(720, 1280), stride=16) + with pytest.raises(ValueError, match="Unsupported camera model"): + validate_token_geometry(_calibration("orthographic", "none", []), affine, image_hw=(720, 1280), stride=16) + # The image corners reach a distorted radius beyond the model's invertible range. + with pytest.raises(ValueError, match="OpenCV"): + validate_token_geometry( + _calibration("pinhole", "opencv", [-0.9, 0.0, 0.0, 0.0]), affine, image_hw=(720, 1280), stride=16 + ) diff --git a/cosmos_framework/model/generator/utils/camera_relative_pose.py b/cosmos_framework/model/generator/utils/camera_relative_pose.py new file mode 100644 index 000000000..543b9d732 --- /dev/null +++ b/cosmos_framework/model/generator/utils/camera_relative_pose.py @@ -0,0 +1,420 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""Prepare camera extrinsics for attention, outside compiled model regions. + +CameraGeometry poses are OpenCV camera-to-world matrices in meters. Attention +uses reference-to-camera matrices, centered on the same primary camera at each +instant. The dense camera branch only compares cameras at the same instant, so +this removes large world translations without changing its relative transforms. +Intrinsics and lens distortion are not used by this extrinsics-only ablation. +""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from cosmos_framework.utils.easy_io import easy_io +from cosmos_framework.data.generator.multiview.camera_geometry import ( + PER_FRAME_CAMERA_TO_RIG_KEY, + _validate_calibration, +) +from cosmos_framework.model.generator.utils.rigrope_geometry import ( + MomentNormalization, + PoseWorldFrame, + rigrope_geometry_available, +) +from cosmos_framework.utils.generator.image_resize import get_vision_data_resolution + + +def invert_rigid_transform(matrix: torch.Tensor) -> torch.Tensor: # [...,4,4] -> [...,4,4] + rotation = matrix[..., :3, :3].transpose(-1, -2) # [...,3,3] + translation = -(rotation @ matrix[..., :3, 3:4]) # [...,3,1] + upper = torch.cat((rotation, translation), dim=-1) # [...,3,4] + return torch.cat((upper, matrix[..., 3:4, :]), dim=-2) # [...,4,4] + + +def validate_camera_poses(poses: torch.Tensor) -> None: # [V,T,4,4] + """Reject invalid SE(3) matrices before the compiled attention path.""" + if poses.ndim != 4 or poses.shape[-2:] != (4, 4) or min(poses.shape[:2]) < 1: + raise ValueError(f"camera_to_world must have shape [V,T,4,4], got {tuple(poses.shape)}") + if not torch.isfinite(poses).all(): + raise ValueError("camera_to_world must contain finite values") + bottom = poses.new_tensor([0.0, 0.0, 0.0, 1.0]).expand_as(poses[..., 3, :]) # [V,T,4] + rotation = poses[..., :3, :3] # [V,T,3,3] + identity = torch.eye(3, dtype=poses.dtype, device=poses.device).expand_as(rotation) # [V,T,3,3] + if not torch.allclose(poses[..., 3, :], bottom, atol=1e-6, rtol=0): + raise ValueError("camera_to_world must have homogeneous bottom row [0,0,0,1]") + if not torch.allclose(rotation.transpose(-1, -2) @ rotation, identity, atol=1e-4, rtol=0): + raise ValueError("camera_to_world rotations must be orthonormal") + if not torch.allclose(torch.linalg.det(rotation), poses.new_ones(poses.shape[:2]), atol=1e-4, rtol=0): + raise ValueError("camera_to_world rotations must have determinant +1") + + +def _latent_frame_indices( + tokenizer: Any, frames: int, latent_count: int, pixel_hw: tuple[int, int] +) -> torch.Tensor: # [F] + """The source frame each latent frame's pose is read at: its tokenizer patch's right edge.""" + resolution = get_vision_data_resolution(pixel_hw) + positions = tokenizer.get_latent_temporal_positions( + num_pixel_frames=frames, resolution=resolution, num_latent_frames=latent_count + ) # [F] or None + if positions is None: + indices = torch.tensor( + [tokenizer.get_pixel_num_frames(index + 1, resolution=resolution) - 1 for index in range(latent_count)], + dtype=torch.long, + ) # [F] + else: + indices = (positions.detach().cpu() * tokenizer.temporal_compression_factor).round().long() # [F] + # Noncausal tokenizer padding repeats the first/last real frame. + return indices.clamp(0, frames - 1) # [F] + + +def prepare_camera_relative_poses( + geometry: Sequence[Mapping[str, Any] | None] | None, + *, + pixel_shapes: Sequence[tuple[int, int, int]], + latent_frames: Sequence[int], + num_views: Sequence[int], + items_per_sample: Sequence[int], + tokenizer: Any, +) -> list[torch.Tensor | None]: # each tensor: [V,F,4,4] + """Select poses at tokenizer patch right edges; duplicate them for control/target. + + ``pixel_shapes`` contains camera-major (V*T,H,W) shapes. ``geometry`` already + describes the final sampled frames, so its time axis indexes those frames, not + the original video's frame IDs. Missing poses are an error for multiview samples; + single-camera samples take the exact existing attention path and need none. + """ + if sum(items_per_sample) != len(num_views) or not (len(num_views) == len(pixel_shapes) == len(latent_frames)): + raise ValueError("Camera pose metadata must match the flattened vision item layout") + if geometry is not None and len(geometry) != len(items_per_sample): + raise ValueError("camera_geometry must contain one record per logical sample") + result: list[torch.Tensor | None] = [] + cursor = 0 + for sample, count in enumerate(items_per_sample): + item_views = num_views[cursor : cursor + count] + if not item_views or max(item_views) <= 1: + result.extend([None] * count) + cursor += count + continue + record = None if geometry is None else geometry[sample] + if record is None: + raise ValueError( + f"Sample {sample}: use_camera_relative_pose_emb requires camera_geometry; " + "enable load_camera_geometry in the loader or provide inference camera poses" + ) + keys = list(record["camera_keys"]) + if len(keys) != len(set(keys)) or any(views != len(keys) for views in item_views): + raise ValueError(f"Sample {sample}: camera_geometry camera order/count does not match the sampled rig") + poses = torch.as_tensor(record["camera_to_world"], dtype=torch.float64).detach() # [V,T,4,4] + valid = torch.as_tensor(record["pose_valid"], device=poses.device, dtype=torch.bool) # [V,T] + if poses.ndim != 4 or poses.shape[0] != len(keys) or valid.shape != poses.shape[:2]: + raise ValueError(f"Sample {sample}: camera pose shapes disagree with camera_keys") + reference = _reference_view(keys) + for item in range(cursor, cursor + count): + views = num_views[item] + total_frames, height, width = pixel_shapes[item] + if total_frames % views or latent_frames[item] % views: + raise ValueError("Camera pixel and latent frame counts must be divisible by the view count") + frames, latent_count = total_frames // views, latent_frames[item] // views + if poses.shape[1] != frames: + raise ValueError(f"Sample {sample}: camera_geometry has {poses.shape[1]} frames, video has {frames}") + indices = _latent_frame_indices(tokenizer, frames, latent_count, (height, width)).to(poses.device) # [F] + if indices.shape != (latent_count,) or not valid[:, indices].all(): + raise ValueError(f"Sample {sample}: camera poses are missing at selected latent frames") + selected = poses[:, indices] # [V,F,4,4] + validate_camera_poses(selected) + anchor = selected[reference : reference + 1] # [1,F,4,4] + relative = (invert_rigid_transform(selected) @ anchor).float() # [V,F,4,4] + result.append(relative) + cursor += count + return result + + +def _reference_view(keys: Sequence[str]) -> int: + # Prefer front-wide, then front fisheye; other rigs choose a stable physical + # identity so reordering the selected cameras does not change the reference. + reference_key = next( + (key for key in ("camera_front_wide_120fov", "camera_front_fisheye_200fov") if key in keys), + min(keys), + ) + return list(keys).index(reference_key) + + +def prepare_prope_cross_view_poses( + geometry: Sequence[Mapping[str, Any] | None] | None, + *, + pixel_shapes: Sequence[tuple[int, int, int]], + latent_frames: Sequence[int], + num_views: Sequence[int], + items_per_sample: Sequence[int], + tokenizer: Any, + allow_missing: bool, + translation_normalization: MomentNormalization = "fixed", + translation_scale_m: float = 25.0, + translation_floor_m: float = 0.05, + pose_world_frame: PoseWorldFrame = "scene", +) -> list[torch.Tensor | None]: # each tensor: [V,F,4,4] + """Reference-to-camera transforms for ``geometry_position_encoding="prope_cross_view"``. + + Extrinsics come from where RigRoPE reads them: a pose-world rig's per-frame + ``camera_to_world`` at the latent frames, or a static rig's ``camera_to_rig``. With + ``allow_missing``, a sample gets geometry exactly when RigRoPE would give it some, so the + two modes train the same posed population; the rest are ``None`` and take the identity. + ``pose_world_frame="rig_per_frame"`` anchors each frame on the reference camera at that + frame, which drops the motion the cameras share across the band's instants, as RigRoPE's + mode of that name does; ``"scene"`` anchors every frame on the reference camera at the + first latent frame and so keeps it. Translations are divided by ``translation_scale_m``, + or under ``"per_sample_rms"`` by the RMS camera distance from the per-frame (or + first-frame) centroid, floored at ``translation_floor_m``. + """ + if translation_normalization not in ("fixed", "per_sample_rms"): + raise ValueError(f"Unknown PRoPE translation normalization {translation_normalization!r}") + if pose_world_frame not in ("scene", "rig_per_frame"): + raise ValueError(f"Unknown PRoPE pose-world frame {pose_world_frame!r}") + for value, label in ((translation_scale_m, "scale"), (translation_floor_m, "floor")): + if not math.isfinite(value) or value <= 0: + raise ValueError(f"PRoPE translation {label} must be finite and positive, in metres") + if sum(items_per_sample) != len(num_views) or not (len(num_views) == len(pixel_shapes) == len(latent_frames)): + raise ValueError("Camera pose metadata must match the flattened vision item layout") + records: Sequence[Mapping[str, Any] | None] = ( + [None] * len(items_per_sample) if geometry is None and allow_missing else (geometry or []) + ) + if len(records) != len(items_per_sample): + raise ValueError("camera_geometry must contain one record per logical sample") + result: list[torch.Tensor | None] = [] + cursor = 0 + for sample, (record, count) in enumerate(zip(records, items_per_sample, strict=True)): + item_views = num_views[cursor : cursor + count] + multiview = bool(item_views) and max(item_views) > 1 + if not multiview or (allow_missing and not rigrope_geometry_available(record, has_lidar=False)): + result.extend([None] * count) + cursor += count + continue + if record is None: + raise ValueError(f"Sample {sample}: prope_cross_view requires camera_geometry for multiview samples") + keys = list(record["camera_keys"]) + if len(keys) != len(set(keys)) or any(views != len(keys) for views in item_views): + raise ValueError(f"Sample {sample}: camera_geometry camera order/count does not match the sampled rig") + rig = record.get("rig_geometry") or record["provenance"]["cameras"][0].get("rig_geometry") + if rig is None: + raise ValueError(f"Sample {sample}: missing calibrated physical sensor rig") + per_frame = bool(rig.get(PER_FRAME_CAMERA_TO_RIG_KEY, False)) + reference = _reference_view(keys) + for item in range(cursor, cursor + count): + views = num_views[item] + total_frames, height, width = pixel_shapes[item] + if total_frames % views or latent_frames[item] % views: + raise ValueError("Camera pixel and latent frame counts must be divisible by the view count") + frames, latent_count = total_frames // views, latent_frames[item] // views + if per_frame: + poses = torch.as_tensor(record["camera_to_world"], dtype=torch.float64).cpu() # [V,T,4,4] + valid = torch.as_tensor(record["pose_valid"], dtype=torch.bool).cpu() # [V,T] + indices = _latent_frame_indices(tokenizer, frames, latent_count, (height, width)) # [F] + if poses.shape != (views, frames, 4, 4) or not valid[:, indices].all(): + raise ValueError(f"Sample {sample}: a pose-world rig needs valid poses on every latent frame") + selected = poses[:, indices] # [V,F,4,4] + else: + static = torch.tensor([rig["camera_to_rig"][key] for key in keys], dtype=torch.float64) # [V,4,4] + selected = static[:, None].expand(-1, latent_count, -1, -1) # [V,F,4,4] + validate_camera_poses(selected) + centers = selected[..., :3, 3] # [V,F,3] + if per_frame and pose_world_frame == "scene": + anchor = selected[reference : reference + 1, :1] # [1,1,4,4] + spread = centers - centers[:, :1].mean(dim=0, keepdim=True) # [V,F,3] + else: + anchor = selected[reference : reference + 1] # [1,F,4,4] + spread = centers - centers.mean(dim=0, keepdim=True) # [V,F,3] + unit = translation_scale_m + if translation_normalization == "per_sample_rms": + unit = max(float(spread.square().sum(dim=-1).mean().sqrt()), translation_floor_m) + relative = invert_rigid_transform(selected) @ anchor # [V,F,4,4] + relative[..., :3, 3] /= unit + result.append(relative.float()) + cursor += count + return result + + +def load_camera_pose_geometry( + camera_keys: Sequence[str], + pose_paths: Sequence[str | None], + *, + num_frames: int, + start_frame: int = 0, + fps: float | None = None, + image_hw: tuple[int, int] | None = None, + require_calibration_mapping: bool = False, + require_sensor_rig: bool = False, +) -> dict[str, Any]: + """Read explicit per-camera pose JSONs for inference/validation. + + Each JSON contains ``camera_to_world`` as [T,4,4] (or one static [4,4]) + and optional ``pose_valid`` [T] (a single flag for a static pose). Frames must + match the associated video's FPS and origin, unless ``source_frame_indices`` + maps the window onto the sidecar's own frames at its ``fps``; the returned + ``provenance["source_fps"]`` is then that source clock. The selected camera + order is preserved. No implied identity fallback. + """ + if not camera_keys or len(camera_keys) != len(pose_paths) or len(set(camera_keys)) != len(camera_keys): + raise ValueError("Pose paths must match unique camera keys in selected order") + if num_frames < 1 or start_frame < 0: + raise ValueError("Camera pose windows require num_frames > 0 and start_frame >= 0") + all_poses: list[torch.Tensor] = [] + all_valid: list[torch.Tensor] = [] + calibrations: list[dict[str, Any]] = [] + affines: list[torch.Tensor] = [] # each [3,3] + source_starts: list[int] = [] + sensor_rig: dict[str, Any] | None = None + selected_source_indices: torch.Tensor | None = None + selected_frame_times_ns: torch.Tensor | None = None + # A source mapping's frame IDs count frames of the sidecar's own clock, not the output video's. + mapped_source_fps: list[float] = [] + uses_source_mapping: bool | None = None + for key, path in zip(camera_keys, pose_paths, strict=True): + if path is None: + raise ValueError(f"Camera {key!r} needs pose_path for use_camera_relative_pose_emb") + payload = easy_io.load(path) + poses = torch.as_tensor(payload["camera_to_world"], dtype=torch.float64) # [T,4,4] or [4,4] + if poses.ndim not in (2, 3) or poses.shape[-2:] != (4, 4) or poses.numel() == 0: + raise ValueError(f"Camera {key!r}: camera_to_world must have shape [T,4,4] or [4,4]") + if require_sensor_rig: + current_rig = payload.get("rig_geometry") + if current_rig is None or key not in current_rig.get("camera_to_rig", {}): + # PAIBench direct pose sidecars are canonical camera-to-world streams. For + # non-MADS samples these may be moving cameras, so the shared world is the + # only correct rig and each selected frame carries its own camera transform. + current_rig = { + "camera_to_rig": {key: (poses if poses.ndim == 2 else poses[0]).tolist()}, + "per_frame_camera_to_rig": True, + } + if sensor_rig is not None: + if current_rig.get("per_frame_camera_to_rig") and sensor_rig.get("per_frame_camera_to_rig"): + sensor_rig["camera_to_rig"].update(current_rig["camera_to_rig"]) + elif current_rig != sensor_rig: + raise ValueError("Selected pose sidecars contain different physical rigs") + else: + sensor_rig = current_rig + if payload.get("camera_key", key) != key: + raise ValueError(f"Camera {key!r}: pose_path belongs to a different camera") + source_frame_indices = payload.get("source_frame_indices") + # frame_indices and frame_times_ns are shared across views, so one convention must hold for all. + if uses_source_mapping is not None and uses_source_mapping != (source_frame_indices is not None): + raise ValueError("Selected cameras must all provide source_frame_indices, or none of them") + uses_source_mapping = source_frame_indices is not None + if source_frame_indices is not None: + if "fps" not in payload or not float(payload["fps"]) > 0: + raise ValueError(f"Camera {key!r}: source_frame_indices require a positive source fps") + source_indices = torch.as_tensor(source_frame_indices, dtype=torch.int64) + if source_indices.ndim != 1 or len(source_indices) < start_frame + num_frames: + raise ValueError(f"Camera {key!r}: source_frame_indices do not cover the requested window") + source_indices = source_indices[start_frame : start_frame + num_frames] + if selected_source_indices is not None and not torch.equal(selected_source_indices, source_indices): + raise ValueError("Selected cameras have different source-frame mappings") + selected_source_indices = source_indices + elif fps is not None and "fps" in payload and not math.isclose(float(payload["fps"]), fps): + raise ValueError(f"Camera {key!r}: pose FPS {payload['fps']} does not match video FPS {fps}") + if poses.ndim == 2: + # One static transform holds at every frame, so it needs no pose table covering the + # window, however large its absolute source-frame IDs. + static_valid = torch.as_tensor(payload.get("pose_valid", True), dtype=torch.bool) # [] + if static_valid.ndim != 0: + raise ValueError(f"Camera {key!r}: a static camera_to_world takes a single pose_valid flag") + all_poses.append(poses.expand(num_frames, -1, -1)) # [F,4,4] + all_valid.append(static_valid.expand(num_frames)) # [F] + else: + required_pose_frames = ( + int(source_indices.max()) + 1 if source_frame_indices is not None else start_frame + num_frames + ) + valid = torch.as_tensor(payload.get("pose_valid", [True] * len(poses)), dtype=torch.bool) # [T] + if poses.shape[0] < required_pose_frames or valid.shape != poses.shape[:1]: + raise ValueError(f"Camera {key!r}: pose sequence does not cover the requested video window") + window = ( + source_indices if source_frame_indices is not None else slice(start_frame, start_frame + num_frames) + ) + all_poses.append(poses[window]) # [F,4,4] + all_valid.append(valid[window]) # [F] + if source_frame_indices is not None: + source_starts.append(int(source_indices[0])) + source_fps = float(payload["fps"]) + if mapped_source_fps and not math.isclose(source_fps, mapped_source_fps[0]): + raise ValueError("Selected cameras have different source FPS") + mapped_source_fps.append(source_fps) + frame_times = (source_indices.double() * (1e9 / source_fps)).round().long() + if selected_frame_times_ns is not None and not torch.equal(selected_frame_times_ns, frame_times): + raise ValueError("Selected cameras have different frame-time mappings") + selected_frame_times_ns = frame_times + else: + source_starts.append(int(payload.get("source_frame_start", 0)) + start_frame) + if image_hw is not None: + # Caption and relative-ray geometry require native calibration and the + # affine into staged media. Pose-only JSONs remain valid for extrinsics. + if require_calibration_mapping and not payload.get("calibration_mapping_verified", False): + raise ValueError( + f"Camera {key!r}: relative-ray inference requires a verified calibration-to-media mapping" + ) + if not all(field in payload for field in ("calibration", "image_hw", "image_from_calibration")): + raise ValueError( + f"Camera {key!r}: calibrated geometry requires calibration and image metadata in pose_path" + ) + try: + _validate_calibration(payload["calibration"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError(f"Camera {key!r}: invalid calibration in pose_path: {error}") from error + calibrations.append(payload["calibration"]) + source_h, source_w = payload["image_hw"] + height, width = image_hw + if min(source_h, source_w, height, width) <= 0: + raise ValueError("Camera image dimensions must be positive") + ratio = max(width / source_w, height / source_h) + resize_h, resize_w = math.ceil(ratio * source_h), math.ceil(ratio * source_w) + top, left = round((resize_h - height) / 2), round((resize_w - width) / 2) + sx, sy = resize_w / source_w, resize_h / source_h + resize = torch.tensor( # [3,3] + [[sx, 0, (sx - 1) / 2 - left], [0, sy, (sy - 1) / 2 - top], [0, 0, 1]], dtype=torch.float64 + ) + affine = torch.as_tensor(payload["image_from_calibration"], dtype=torch.float64) # [3,3] + if affine.shape != (3, 3) or not torch.isfinite(affine).all(): + raise ValueError(f"Camera {key!r}: invalid image_from_calibration") + affines.append(resize @ affine) # [3,3] + if len(set(source_starts)) != 1: + raise ValueError("Selected camera pose sidecars have different source frame origins") + geometry: dict[str, Any] = { + "camera_keys": list(camera_keys), + "camera_to_world": torch.stack(all_poses), # [V,F,4,4] + "pose_valid": torch.stack(all_valid), # [V,F] + } + if not geometry["pose_valid"].all(): + raise ValueError("Selected camera pose window contains missing poses") + validate_camera_poses(geometry["camera_to_world"]) + if image_hw is not None: + geometry.update( + calibration=calibrations, + image_from_calibration=torch.stack(affines), # [V,3,3] + image_hw=torch.tensor([image_hw] * len(camera_keys)), # [V,2] + frame_indices=( + selected_source_indices[None].repeat(len(camera_keys), 1) + if selected_source_indices is not None + else torch.tensor(source_starts)[:, None] + torch.arange(num_frames)[None, :] + ), # [V,F] + ) + if require_sensor_rig: + if fps is None or fps <= 0 or image_hw is None: + raise ValueError("RigRoPE inference requires FPS and calibrated image dimensions") + geometry.update( + rig_geometry=sensor_rig, + frame_times_ns=( + selected_frame_times_ns[None].repeat(len(camera_keys), 1) + if selected_frame_times_ns is not None + else (torch.arange(num_frames, dtype=torch.float64) * (1e9 / fps)) + .round() + .long()[None] + .repeat(len(camera_keys), 1) + ), # [V,F] + time_valid=torch.ones((len(camera_keys), num_frames), dtype=torch.bool), # [V,F] + provenance={"source_fps": mapped_source_fps or [fps] * len(camera_keys)}, + ) + return geometry diff --git a/cosmos_framework/model/generator/utils/camera_relative_pose_test.py b/cosmos_framework/model/generator/utils/camera_relative_pose_test.py new file mode 100644 index 000000000..d01bb9004 --- /dev/null +++ b/cosmos_framework/model/generator/utils/camera_relative_pose_test.py @@ -0,0 +1,165 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +import json +import math +from pathlib import Path +from typing import Any + +import pytest +import torch + +from cosmos_framework.model.generator.utils.camera_relative_pose import load_camera_pose_geometry + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def _poses(num_frames: int) -> list[list[list[float]]]: + """Identity rotations translating along x by the frame index.""" + poses = torch.eye(4, dtype=torch.float64).repeat(num_frames, 1, 1) # [T,4,4] + poses[:, 0, 3] = torch.arange(num_frames, dtype=torch.float64) + return poses.tolist() + + +def _write_sidecar(tmp_path: Path, key: str, **fields: Any) -> str: + path = tmp_path / f"{key}.json" + path.write_text(json.dumps({"camera_key": key, "camera_to_world": _poses(10), **fields})) + return str(path) + + +def test_source_frame_indices_select_the_mapped_pose_rows(tmp_path: Path) -> None: + mapping = [0, 2, 4, 6, 8] + paths = [_write_sidecar(tmp_path, key, fps=10.0, source_frame_indices=mapping) for key in ("front", "rear")] + + geometry = load_camera_pose_geometry(["front", "rear"], paths, num_frames=3, start_frame=1) + + assert geometry["camera_to_world"][:, :, 0, 3].tolist() == [[2.0, 4.0, 6.0], [2.0, 4.0, 6.0]] + + +def test_source_frame_indices_must_be_all_or_none(tmp_path: Path) -> None: + paths = [ + _write_sidecar(tmp_path, "front", fps=10.0, source_frame_indices=[0, 2, 4, 6, 8]), + _write_sidecar(tmp_path, "rear", fps=10.0), + ] + + with pytest.raises(ValueError, match="all provide source_frame_indices, or none"): + load_camera_pose_geometry(["front", "rear"], paths, num_frames=3) + + +def _static_pose() -> torch.Tensor: # returns [4,4] + static = torch.eye(4, dtype=torch.float64) # [4,4] + static[:3, 3] = torch.tensor([1.0, 2.0, 3.0]) + return static + + +def _calibrated_fields(**calibration: float) -> dict[str, Any]: + """The image metadata RigRoPE inference requires beside the poses; ``calibration`` overrides fields.""" + pinhole = { + "camera_model": "pinhole", + "distortion_model": "none", + "distortion_coefficients": [], + "width_px": 1280, + "height_px": 720, + "fx_px": 900.0, + "fy_px": 900.0, + "cx_px": 639.5, + "cy_px": 359.5, + **calibration, + } + return {"calibration": pinhole, "image_hw": [720, 1280], "image_from_calibration": torch.eye(3).tolist()} + + +def test_static_pose_covers_absolute_source_frame_indices(tmp_path: Path) -> None: + static = _static_pose() # [4,4] + path = _write_sidecar( + tmp_path, "front", camera_to_world=static.tolist(), fps=10.0, source_frame_indices=[100, 102, 104] + ) + + geometry = load_camera_pose_geometry(["front"], [path], num_frames=3) + + torch.testing.assert_close(geometry["camera_to_world"], static.expand(1, 3, 4, 4)) + + +def test_static_pose_takes_a_single_pose_valid_flag(tmp_path: Path) -> None: + invalid = _write_sidecar(tmp_path, "front", camera_to_world=_static_pose().tolist(), pose_valid=False) + with pytest.raises(ValueError, match="contains missing poses"): + load_camera_pose_geometry(["front"], [invalid], num_frames=3) + + per_frame = _write_sidecar(tmp_path, "rear", camera_to_world=_static_pose().tolist(), pose_valid=[True] * 3) + with pytest.raises(ValueError, match="single pose_valid flag"): + load_camera_pose_geometry(["rear"], [per_frame], num_frames=3) + + +def test_static_pose_without_rig_geometry_makes_a_4x4_camera_to_rig(tmp_path: Path) -> None: + static = _static_pose() # [4,4] + path = _write_sidecar(tmp_path, "front", camera_to_world=static.tolist(), **_calibrated_fields()) + + geometry = load_camera_pose_geometry( + ["front"], [path], num_frames=3, fps=10.0, image_hw=(720, 1280), require_sensor_rig=True + ) + + assert geometry["rig_geometry"] == {"camera_to_rig": {"front": static.tolist()}, "per_frame_camera_to_rig": True} + + +@pytest.mark.parametrize( + ("mapping", "expected_source_fps"), + [({"fps": 30.0, "source_frame_indices": [2, 4, 6, 8]}, 30.0), ({}, 24.0)], + ids=["mapped", "unmapped"], +) +def test_rig_provenance_carries_the_clock_frame_indices_count( + tmp_path: Path, mapping: dict[str, Any], expected_source_fps: float +) -> None: + """RigRoPE times LiDAR source-frame IDs against ``frame_indices`` at ``provenance["source_fps"]``. + + A mapped sidecar counts its own 30 FPS frames, so 24 FPS output must not relabel that clock. + """ + paths = [_write_sidecar(tmp_path, key, **mapping, **_calibrated_fields()) for key in ("front", "rear")] + + geometry = load_camera_pose_geometry( + ["front", "rear"], paths, num_frames=3, fps=24.0, image_hw=(720, 1280), require_sensor_rig=True + ) + + assert geometry["provenance"]["source_fps"] == [expected_source_fps] * 2 + frame_indices = geometry["frame_indices"].double() # [V,F] + frame_times_s = geometry["frame_times_ns"].double() / 1e9 # [V,F] + torch.testing.assert_close( + frame_times_s - frame_times_s[:, :1], (frame_indices - frame_indices[:, :1]) / expected_source_fps + ) + + +def test_mapped_cameras_must_share_a_source_fps(tmp_path: Path) -> None: + paths = [ + _write_sidecar(tmp_path, "front", fps=30.0, source_frame_indices=[0, 1, 2]), + _write_sidecar(tmp_path, "rear", fps=15.0, source_frame_indices=[0, 1, 2]), + ] + + with pytest.raises(ValueError, match="different source FPS"): + load_camera_pose_geometry(["front", "rear"], paths, num_frames=3) + + +@pytest.mark.parametrize( + "invalid", [{"fx_px": 0.0}, {"fy_px": -900.0}, {"cx_px": math.inf}], ids=["zero_fx", "negative_fy", "inf_cx"] +) +def test_calibrated_inference_geometry_rejects_invalid_calibration(tmp_path: Path, invalid: dict[str, float]) -> None: + path = _write_sidecar(tmp_path, "front", camera_to_world=_static_pose().tolist(), **_calibrated_fields(**invalid)) + + with pytest.raises(ValueError, match="'front': invalid calibration in pose_path"): + load_camera_pose_geometry(["front"], [path], num_frames=3, fps=10.0, image_hw=(720, 1280)) + + +def test_calibrated_inference_geometry_rejects_malformed_calibration(tmp_path: Path) -> None: + fields = _calibrated_fields() + del fields["calibration"]["fx_px"] + path = _write_sidecar(tmp_path, "front", camera_to_world=_static_pose().tolist(), **fields) + + with pytest.raises(ValueError, match="'front': invalid calibration in pose_path"): + load_camera_pose_geometry(["front"], [path], num_frames=3, fps=10.0, image_hw=(720, 1280)) + + +@pytest.mark.parametrize("fields", [{}, {"fps": 0.0}], ids=["missing", "zero"]) +def test_source_frame_indices_require_a_positive_source_fps(tmp_path: Path, fields: dict[str, Any]) -> None: + paths = [_write_sidecar(tmp_path, "front", source_frame_indices=[0, 2, 4, 6, 8], **fields)] + + with pytest.raises(ValueError, match="require a positive source fps"): + load_camera_pose_geometry(["front"], paths, num_frames=3) diff --git a/cosmos_framework/model/generator/utils/ftheta.py b/cosmos_framework/model/generator/utils/ftheta.py new file mode 100644 index 000000000..8731821b6 --- /dev/null +++ b/cosmos_framework/model/generator/utils/ftheta.py @@ -0,0 +1,114 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""Native F-theta unprojection, independent of the attention encoding.""" + +import math +from collections.abc import Sequence +from functools import lru_cache + +import numpy as np +import torch + + +def _polynomial(values: torch.Tensor, coefficients: Sequence[float]) -> torch.Tensor: # [...], [K] -> [...] + result = torch.zeros_like(values) # [...] + for coefficient in reversed(coefficients): + result = result * values + coefficient # [...] + return result + + +@lru_cache(maxsize=256) +def _forward_domain(coefficients: tuple[float, ...]) -> float: + """Return the first positive turning point, bounded by the ray sphere.""" + derivative = [index * value for index, value in enumerate(coefficients)][1:] + if derivative[0] <= 0: + raise ValueError("F-theta forward polynomial must increase at the optical axis") + roots = np.polynomial.polynomial.polyroots(derivative) + boundaries = [float(root.real) for root in roots if abs(root.imag) < 1e-8 and root.real > 0] + return min([math.pi - 1e-8, *boundaries]) + + +def ftheta_unproject( + pixels: torch.Tensor, coefficients: Sequence[float] +) -> torch.Tensor: # pixels: [...,2], returns [...,3] + """Unproject native pixel centers using either direction of the provider polynomial. + + The optional final three coefficients describe A=[[c,d],[e,1]]. Forward + polynomials are inverted on their first monotonic interval rather than using + an equivalent pinhole K or an unconstrained fitted inverse polynomial. + """ + values = tuple(float(value) for value in coefficients) + if len(values) not in (11, 14) or not all(math.isfinite(value) for value in values): + raise ValueError("F-theta calibration requires 11 or 14 finite coefficients") + if pixels.shape[-1] != 2 or min(values[2:4]) <= 0: + raise ValueError("F-theta pixels must end in two coordinates and image dimensions must be positive") + if not pixels.is_floating_point() or not torch.isfinite(pixels).all(): + raise ValueError("F-theta pixels must have finite floating-point coordinates") + original_dtype = pixels.dtype + pixels = pixels.double() # [...,2] + center = pixels.new_tensor(values[:2]) # [2] + c, d, e = values[11:] if len(values) == 14 else (1.0, 0.0, 0.0) + if abs(c - d * e) < 1e-12: + raise ValueError("F-theta affine lens matrix is singular") + inverse_affine = pixels.new_tensor([[1.0, -d], [-e, c]]) / (c - d * e) # [2,2] + offset = (pixels - center) @ inverse_affine.T # [...,2] + radius = torch.linalg.vector_norm(offset, dim=-1) # [...] + polynomial = values[4:10] + if values[10] > 0: + angle = _polynomial(radius, polynomial) # [...] + derivative = _polynomial(radius, [index * value for index, value in enumerate(polynomial)][1:]) # [...] + if (derivative <= 0).any(): + raise ValueError("Token centers fall outside the monotonic F-theta backward domain") + else: + upper_angle = _forward_domain(polynomial) + upper = torch.full_like(radius, upper_angle) # [...] + lower = torch.zeros_like(radius) # [...] + if (radius > _polynomial(upper, polynomial) + 1e-7).any() or (radius < polynomial[0] - 1e-7).any(): + raise ValueError("Token centers fall outside the invertible F-theta forward domain") + for _ in range(48): + midpoint = (lower + upper) * 0.5 # [...] + below = _polynomial(midpoint, polynomial) < radius # [...] + lower = torch.where(below, midpoint, lower) # [...] + upper = torch.where(below, upper, midpoint) # [...] + angle = (lower + upper) * 0.5 # [...] + if not torch.isfinite(angle).all() or (angle < -1e-8).any() or (angle >= math.pi).any(): + raise ValueError("F-theta token ray angles must lie in [0, pi)") + scale = torch.sin(angle) / radius.clamp_min(1e-15) # [...] + rays = torch.cat((offset * scale[..., None], torch.cos(angle)[..., None]), dim=-1) # [...,3] + optical_axis = rays.new_tensor([0.0, 0.0, 1.0]) # [3] + rays = torch.where((radius <= 1e-12)[..., None], optical_axis, rays) # [...,3] + return rays.to(dtype=original_dtype) # [...,3] + + +@lru_cache(maxsize=128) +def _validate_cached( + coefficients: tuple[float, ...], affine_values: tuple[float, ...], image_hw: tuple[int, int], stride: int +) -> None: + height, width = image_hw + ys = (torch.arange(math.ceil(height / stride), dtype=torch.float64) + 0.5) * stride # [H] + xs = (torch.arange(math.ceil(width / stride), dtype=torch.float64) + 0.5) * stride # [W] + yy, xx = torch.meshgrid(ys, xs, indexing="ij") # [H,W], [H,W] + pixels = torch.stack((xx, yy, torch.ones_like(xx)), dim=-1) # [H,W,3] + affine = torch.tensor(affine_values, dtype=torch.float64).reshape(3, 3) # [3,3] + if not torch.isfinite(affine).all() or abs(float(torch.linalg.det(affine))) < 1e-12: + raise ValueError("image_from_calibration must be finite and invertible") + native = pixels @ torch.linalg.inv(affine).T # [H,W,3] + ftheta_unproject(native[..., :2] / native[..., 2:], coefficients) # [H,W,3] + + +def validate_ftheta_token_geometry( + coefficients: Sequence[float], + image_from_calibration: torch.Tensor, # [3,3] + *, + image_hw: tuple[int, int], + stride: int, +) -> None: + """Validate pre-interpolation latent-center rays consumed by RigRoPE.""" + _validate_cached( + tuple(float(value) for value in coefficients), + tuple(float(value) for value in image_from_calibration.flatten()), + image_hw, + stride, + ) diff --git a/cosmos_framework/model/generator/utils/multiview_ar.py b/cosmos_framework/model/generator/utils/multiview_ar.py index 21fb1761c..da09145d8 100644 --- a/cosmos_framework/model/generator/utils/multiview_ar.py +++ b/cosmos_framework/model/generator/utils/multiview_ar.py @@ -6,16 +6,19 @@ from __future__ import annotations import math -from collections.abc import Sequence +from collections.abc import Callable, Generator, Sequence from dataclasses import dataclass, fields, replace from typing import TYPE_CHECKING, Any, Literal, Protocol import torch +import torch.distributed as dist from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_broadcast_tensor_list +from cosmos_framework.model.generator.omni_mot_model import _broadcast_seed, _per_view_caption_groups from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.model.generator.utils.memory import MemoryState -from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan +from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan, build_sequence_plans_from_data_batch +from cosmos_framework.utils.generator.multiview import slice_camera_major_frames from cosmos_framework.model.generator.mot.causal_flex_attention import ( _ROLE_CLEAN_TARGET, _ROLE_CURRENT_TARGET, @@ -116,7 +119,7 @@ class MultiviewTransferARReplayContext: class MultiviewTransferARBackend: - """Own multiview Transfer packing, Flex KV layout, cache writes, and replay state.""" + """Own multiview Transfer packing, replay KV layout, cache writes, and replay state.""" def __init__(self, host: MultiviewTransferARHost) -> None: self.host = host @@ -222,7 +225,7 @@ def _current_pack_text_tokens( @staticmethod def build_memory_layout(session: MultiviewTransferARSession) -> MultiviewTransferARMemoryLayout: - """Resolve the current fixed-slot Flex memory layout.""" + """Resolve the current fixed-slot replay memory layout.""" return build_multiview_transfer_ar_memory_layout( token_shapes=list(session.token_shapes), target_condition_mask=session.target_condition_mask, @@ -914,16 +917,6 @@ def joint_ar_chunks( return chunks -def _camera_prefix( - latent: torch.Tensor, # [1,C,V*T,H,W] - num_views: int, - end: int, -) -> torch.Tensor: # [1,C,V*end,H,W] - batch, channels, total_frames, height, width = latent.shape - grid = latent.reshape(batch, channels, num_views, total_frames // num_views, height, width) # [1,C,V,T,H,W] - return grid[:, :, :, :end].reshape(batch, channels, num_views * end, height, width) # [1,C,V*end,H,W] - - def _conditioned_target( latent: torch.Tensor, # [1,C,T,H,W] mask: torch.Tensor, # [T,1,1] @@ -933,18 +926,6 @@ def _conditioned_target( return torch.where(condition, latent, torch.zeros_like(latent)) # [1,C,T,H,W] -def _prefix_condition_count(mask: torch.Tensor, num_views: int) -> int: # mask: [V*T,1,1] - """Reject unsynchronized or non-prefix conditioning in this joint sampler.""" - grid = mask.to(dtype=torch.bool).reshape(num_views, -1) # [V,T] - if not torch.equal(grid, grid[:1].expand_as(grid)): - raise ValueError("Joint AR currently requires the same conditioned prefix for every RGB view") - count = int(grid[0].sum().item()) - expected = torch.arange(grid.shape[1], device=grid.device) < count # [T] - if not torch.equal(grid[0], expected): - raise ValueError("Joint AR requires a contiguous conditioned target prefix") - return count - - def _validate_persistent_chunks(chunks: list[JointARChunk], vision_condition: int, lidar_condition: int) -> None: """Check whether persistent K/V can replace replay for the complete chunk schedule. @@ -989,8 +970,10 @@ def _prefix_data( raw_state_vision=None, raw_state_lidar=None, x0_tokens_vision=[ - _camera_prefix(data.x0_tokens_vision[0], num_views, chunk.vision_prefix_end), # [1,Cv,V*Tv_prefix,Hv,Wv] - _camera_prefix(vision_target, num_views, chunk.vision_prefix_end), # [1,Cv,V*Tv_prefix,Hv,Wv] + slice_camera_major_frames( + data.x0_tokens_vision[0], num_views, 0, chunk.vision_prefix_end + ), # [1,Cv,V*Tv_prefix,Hv,Wv] + slice_camera_major_frames(vision_target, num_views, 0, chunk.vision_prefix_end), # [1,Cv,V*Tv_prefix,Hv,Wv] ], x0_tokens_lidar=[ data.x0_tokens_lidar[0][:, :, : chunk.lidar_prefix_end], # [1,Cl,Tl_prefix,Hl,Wl] @@ -1021,6 +1004,217 @@ def _build_replay_branch( return pack, memory +def _robot_multiview_action_prefix( + data: GenerationDataClean, + generated_views: list[torch.Tensor], + *, + num_views: int, + frames_per_view: int, + end: int, + temporal_factor: int, +) -> GenerationDataClean: + """Materialize one complete robot prefix in the exact training layout.""" + assert data.x0_tokens_action is not None and len(data.x0_tokens_action) == 1 + action = data.x0_tokens_action[0] + actions_per_view = (frames_per_view - 1) * temporal_factor + prefix_actions_per_view = (end - 1) * temporal_factor + prefix_action = torch.cat( + [ + action[view * actions_per_view : view * actions_per_view + prefix_actions_per_view] + for view in range(num_views) + ], + dim=0, + ) + prefix_domains = data.action_domain_id + if prefix_domains is not None: + domain = prefix_domains[0].reshape(-1) + if domain.numel() == num_views * actions_per_view: + domain = torch.cat( + [ + domain[view * actions_per_view : view * actions_per_view + prefix_actions_per_view] + for view in range(num_views) + ] + ) + prefix_domains = [domain] + return replace( + data, + raw_state_vision=None, + raw_state_action=[prefix_action], + x0_tokens_vision=[view[:, :, :end] for view in generated_views], + x0_tokens_action=[prefix_action], + action_domain_id=prefix_domains, + ) + + +def iter_robot_multiview_action_ar( + host: OmniMoTCausalModel, + *, + plans: list[SequencePlan], + data: GenerationDataClean, + conditional_text: list[list[int]], + unconditional_text: list[list[int]] | None, + guidance: float, + seed: int, + num_steps: int, + shift: float, + normalize_cfg: bool, + sampler_mode: Literal["rf", "distilled"], + distilled_num_steps: int | None, + max_num_frames: int | None, +) -> Generator[dict[str, Any], None, None]: + """Generate heterogeneous robot camera views with full-prefix replay. + + Each denoising chunk uses the same per-view three-way replay plus + same-instant inter-view RGB attention as Stage-1 training. Rebuilding the + complete prefix is deliberate for this no-training sanity path: it avoids + introducing a second, approximate cache layout while preserving the latest + replay policy exactly. + """ + if len(plans) != 1 or not plans[0].has_vision or not plans[0].has_action: + raise ValueError("Robot multiview AR requires one RGB+action forward-dynamics sample.") + if plans[0].condition_frame_indexes_vision != [0]: + raise ValueError( + "Robot multiview AR requires synchronized frame 0 as its only RGB condition; " + f"got {plans[0].condition_frame_indexes_vision}." + ) + if not host._uses_robot_multiview_threeway() or host.config.compile.enabled: + raise ValueError("Robot multiview AR requires eager robot three-way attention.") + if host._get_teacher_forcing_kv_implementation() != "multiview_threeway_kv": + raise ValueError( + "Robot multiview AR requires the Stage-1 multiview_threeway_kv path; " + "legacy maskless/Flex replay is intentionally unsupported." + ) + if host.config.multiview_action_conditioning is None: + raise ValueError("Robot multiview AR requires multiview_action_conditioning.") + if data.x0_tokens_vision is None or data.x0_tokens_action is None: + raise ValueError("Robot multiview AR requires encoded RGB and action streams.") + + num_views = len(host.config.multiview_action_conditioning.view_codes) + if data.num_vision_items_per_sample != [num_views] or len(data.x0_tokens_vision) != num_views: + raise ValueError( + f"Robot multiview AR requires one RGB item per configured view ({num_views}); " + f"got {data.num_vision_items_per_sample}." + ) + frames_per_view = data.x0_tokens_vision[0].shape[2] + if any(view.ndim != 5 or view.shape[0] != 1 or view.shape[2] != frames_per_view for view in data.x0_tokens_vision): + raise ValueError("Robot multiview AR views must be aligned [1,C,T,H,W] tensors.") + temporal_factor = host.tokenizer_vision_gen.temporal_compression_factor or 4 + actions_per_view = (frames_per_view - 1) * temporal_factor + action = data.x0_tokens_action[0] + if action.ndim != 2 or action.shape[0] != num_views * actions_per_view: + raise ValueError( + "Robot multiview AR requires camera-major full-rate actions: " + f"got {tuple(action.shape)}, expected rows={num_views * actions_per_view}." + ) + output_frames = frames_per_view if max_num_frames is None else min(frames_per_view, max_num_frames) + if output_frames < 1: + raise ValueError(f"max_num_frames must leave at least one frame, got {max_num_frames}.") + chunk_size = int(host.config.teacher_forcing_frames_per_chunk) + if output_frames > 1 and (output_frames - 1) % chunk_size: + raise ValueError( + f"Robot multiview AR must match the Stage-1 [1,C,C,...] partition: got T={output_frames}, C={chunk_size}." + ) + if normalize_cfg and guidance != 1.0: + raise ValueError("Robot multiview replay sanity currently requires normalize_cfg=False.") + + generated_views = [torch.zeros_like(view, dtype=torch.float32) for view in data.x0_tokens_vision] + for generated, source in zip(generated_views, data.x0_tokens_vision, strict=True): + generated[:, :, :1] = source[:, :, :1].to(device=generated.device, dtype=generated.dtype) + yield {"vision_items": [view[:, :, :1] for view in generated_views]} + if output_frames == 1: + return + + for chunk_start in range(1, output_frames, chunk_size): + chunk_end = min(chunk_start + chunk_size, output_frames) + chunk_len = chunk_end - chunk_start + forward_data = _robot_multiview_action_prefix( + data, + generated_views, + num_views=num_views, + frames_per_view=frames_per_view, + end=chunk_end, + temporal_factor=temporal_factor, + ) + cond_pack, cond_memory = _build_replay_branch(host, plans, forward_data, conditional_text) + uncond_branch = ( + _build_replay_branch(host, plans, forward_data, unconditional_text) + if guidance != 1.0 and unconditional_text is not None + else None + ) + + view_shapes = [(1, view.shape[1], chunk_len, view.shape[3], view.shape[4]) for view in generated_views] + view_sizes = [math.prod(shape) for shape in view_shapes] + generator = torch.Generator(device=generated_views[0].device).manual_seed(seed + chunk_start) + initial_noise = torch.randn( + (1, sum(view_sizes)), + device=generated_views[0].device, + dtype=torch.float32, + generator=generator, + ) + + def branch_velocity( + pack: PackedSequence, + memory: TeacherForcingMemoryState, + noise: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + noisy_views: list[torch.Tensor] = [] + offset = 0 + for view, shape, size in zip(forward_data.x0_tokens_vision or [], view_shapes, view_sizes, strict=True): + noisy = view.clone() + noisy[:, :, chunk_start:chunk_end] = noise[:, offset : offset + size].reshape(shape) + noisy_views.append(noisy) + offset += size + host._update_inference_pack_template( + packed_sequence=pack, + noise_x_vision=noisy_views, + noise_x_lidar=None, + noise_x_radar=None, + noise_x_action=None, + noise_x_sound=None, + timestep=timestep, + ) + prediction = host.denoise(data_batch_packed=pack, memory=memory) + predicted_views = prediction["preds_vision"] + if len(predicted_views) != num_views: + raise ValueError(f"Expected {num_views} RGB predictions, got {len(predicted_views)}.") + return torch.cat( + [prediction[:, :, chunk_start:chunk_end].reshape(1, -1) for prediction in predicted_views], + dim=1, + ) + + def velocity(noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + conditional = branch_velocity(cond_pack, cond_memory, noise, timestep) + if uncond_branch is None: + return conditional + unconditional = branch_velocity(*uncond_branch, noise, timestep) + return unconditional + guidance * (conditional - unconditional) + + denoised = host._run_ar_sampler( + velocity, + initial_noise, + sampler_mode=sampler_mode, + num_steps=num_steps, + shift=shift, + seed=seed, + sample_idx=chunk_start, + num_frames=output_frames, + distilled_num_steps=distilled_num_steps, + schedule_frame_idx=chunk_start, + ) + if not torch.isfinite(denoised).all(): + raise FloatingPointError(f"Robot multiview AR produced nonfinite latents at chunk {chunk_start}.") + offset = 0 + clean_chunks: list[torch.Tensor] = [] + for generated, shape, size in zip(generated_views, view_shapes, view_sizes, strict=True): + clean = denoised[:, offset : offset + size].reshape(shape) + generated[:, :, chunk_start:chunk_end] = clean + clean_chunks.append(clean) + offset += size + for local_frame in range(chunk_len): + yield {"vision_items": [chunk[:, :, local_frame : local_frame + 1] for chunk in clean_chunks]} + + def _current_chunk_data( data: GenerationDataClean, vision_target: torch.Tensor, # [1,Cv,V*Tv,Hv,Wv] @@ -1228,8 +1422,8 @@ def sample_joint_transfer_ar( rgb_control, rgb_item, lidar_control, lidar_item = sensor_items[0] if not (rgb_control.is_control and lidar_control.is_control) or rgb_item.is_control or lidar_item.is_control: raise ValueError("Joint AR sensor control/target roles differ from the training layout") - vision_condition = _prefix_condition_count(initial_pack.vision.condition_mask[1], num_views) - lidar_condition = _prefix_condition_count(initial_pack.lidar.condition_mask[1], 1) + vision_condition = multiview_conditioned_prefix_length(initial_pack.vision.condition_mask[1], num_views) + lidar_condition = multiview_conditioned_prefix_length(initial_pack.lidar.condition_mask[1], 1) vision_target = _conditioned_target(data.x0_tokens_vision[1], initial_pack.vision.condition_mask[1]).to( device=host.tensor_kwargs["device"], dtype=torch.float32 ) # [1,Cv,V*Tv,Hv,Wv] @@ -1421,3 +1615,455 @@ def velocity(noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: # [1 current = _current_chunk_data(data, vision_target, lidar_target, num_views, chunk) # No later query reads the last chunk; do not run an unused clean forward. _commit_persistent_joint_chunk(host, plans, current, persistent_branches, chunk) + + +def _iter_multiview_logical_frames( + vision_latent: torch.Tensor, + *, + num_views: int, +) -> Generator[torch.Tensor, None, None]: # vision_latent: [B,C,V*T,H,W], yields [B,C,V,H,W] + """Yield one camera-major latent time step across all views at a time.""" + if num_views < 1: + raise ValueError(f"num_views must be positive, got {num_views}.") + if vision_latent.ndim != 5: + raise ValueError(f"Expected multiview latent rank 5, got shape {tuple(vision_latent.shape)}.") + flat_frame_count = int(vision_latent.shape[2]) + if flat_frame_count % num_views != 0: + raise ValueError(f"Camera-major latent length {flat_frame_count} must be divisible by num_views={num_views}.") + frames_per_view = flat_frame_count // num_views + vision_by_view = vision_latent.reshape( # [B,C,V,T,H,W] + vision_latent.shape[0], + vision_latent.shape[1], + num_views, + frames_per_view, + vision_latent.shape[3], + vision_latent.shape[4], + ) + for frame_idx in range(frames_per_view): + logical_frame = vision_by_view[:, :, :, frame_idx : frame_idx + 1].reshape( # [B,C,V,H,W] + vision_latent.shape[0], + vision_latent.shape[1], + num_views, + vision_latent.shape[3], + vision_latent.shape[4], + ) + yield logical_frame + + +def multiview_conditioned_prefix_length( + target_condition_mask: torch.Tensor, + num_views: int, + *, + frames_per_view: int | None = None, +) -> int: + """Return a synchronized contiguous-prefix length from a camera-major mask.""" + if frames_per_view is None: + frames_per_view = target_condition_mask.numel() // num_views if num_views > 0 else 0 + expected_frames = num_views * frames_per_view + if num_views < 1 or frames_per_view < 1 or target_condition_mask.numel() != expected_frames: + raise ValueError( + "Expected one target condition value per camera-major frame; " + f"got shape {tuple(target_condition_mask.shape)} for V={num_views}, T={frames_per_view}." + ) + condition_grid = target_condition_mask.to(dtype=torch.bool).reshape(num_views, frames_per_view) # [V,T] + if not torch.equal(condition_grid, condition_grid[0:1].expand_as(condition_grid)): + raise ValueError("Multiview transfer AR requires synchronized target conditioning across every view.") + condition_indexes = torch.nonzero(condition_grid[0], as_tuple=False).squeeze(1) # [T_condition] + expected_indexes = torch.arange(condition_indexes.numel(), device=condition_indexes.device) # [T_condition] + if not torch.equal(condition_indexes, expected_indexes): + raise ValueError( + "Multiview transfer AR only supports contiguous prefix conditioning from frame 0; " + f"got condition frame indexes {condition_indexes.tolist()}." + ) + return condition_indexes.numel() + + +def _submit_multiview_conditioned_prefix( + target_latent: torch.Tensor, + *, + num_views: int, + frames_per_view: int, + condition_count: int, + output_frames: int, + on_clean_vision_chunk: Callable[[torch.Tensor], None] | None, +) -> torch.Tensor | None: # target_latent: [B,C,V*T,H,W], returns [B,C,V*T_prefix,H,W] or None + """Return the camera-major conditioned prefix and optionally submit it for decode.""" + if condition_count == 0: + return None + if num_views < 1 or frames_per_view < 1: + raise ValueError(f"Expected positive multiview dimensions, got V={num_views}, T={frames_per_view}.") + if target_latent.ndim != 5 or target_latent.shape[2] != num_views * frames_per_view: + raise ValueError( + "Expected camera-major target shape [B,C,V*T,H,W], " + f"got {tuple(target_latent.shape)} for V={num_views}, T={frames_per_view}." + ) + if not 0 <= condition_count <= frames_per_view: + raise ValueError(f"condition_count must be in [0, {frames_per_view}], got {condition_count}.") + if not 1 <= output_frames <= frames_per_view: + raise ValueError(f"output_frames must be in [1, {frames_per_view}], got {output_frames}.") + + prefix_frame_count = min(condition_count, output_frames) + conditioned_prefix = slice_camera_major_frames( + target_latent, num_views, 0, prefix_frame_count + ) # [B,C,V*T_prefix,H,W] + if on_clean_vision_chunk is not None: + on_clean_vision_chunk(conditioned_prefix) + return conditioned_prefix + + +def generate_multiview_transfer_ar_chunk( + host: OmniMoTCausalModel, + *, + cond_pack: PackedSequence, + uncond_pack: PackedSequence | None, + cond_cache: list[tuple[torch.Tensor, torch.Tensor] | None], + uncond_cache: list[tuple[torch.Tensor, torch.Tensor] | None] | None, + curr_vision_latent: torch.Tensor, # [1,C,V*chunk_len,H,W] + guidance: float, + num_steps: int, + shift: float, + seed: int, + chunk_start: int, + num_frames: int, + normalize_cfg: bool, + sampler_mode: str, + distilled_num_steps: int | None, + memory_seq_len: int, +) -> torch.Tensor: # [1,C,V*chunk_len,H,W] + """Denoise one synchronized multiview chunk against the cached K/V suffix.""" + # Imported lazily because the shared model dispatches into this backend. + from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + + cond_memory = MultiviewARMemoryState( + num_layers=host.net.num_hidden_layers, + memory_seq_len=memory_seq_len, + cache=cond_cache, + ) + uncond_memory = ( + MultiviewARMemoryState( + num_layers=host.net.num_hidden_layers, + memory_seq_len=memory_seq_len, + cache=uncond_cache, + ) + if uncond_cache is not None + else None + ) + + def run_branch( + pack: PackedSequence, + noise_vision: torch.Tensor, # [1,C,V*chunk_len,H,W] + timestep: torch.Tensor, # [1,1] + branch: _ARBranch, + ) -> torch.Tensor: # [1,C,V*chunk_len,H,W] + OmniMoTCausalModel._set_ar_vision_noise(host, pack, noise_vision, timestep) + cfgp_enabled = host.parallel_dims is not None and host.parallel_dims.cfgp_enabled + # Under CFGP each rank has one branch-local cache in ``cond_memory``. + memory = cond_memory if cfgp_enabled or branch == "conditional" else uncond_memory + assert memory is not None + output = host.denoise(data_batch_packed=pack, memory=memory) + return torch.stack(output["preds_vision"]) # [1,C,V*chunk_len,H,W] + + def velocity_fn(noise_x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + return OmniMoTCausalModel._predict_ar_velocity_with_cfg( + host, + noise_x=noise_x, + timestep=timestep, + vision_shape=curr_vision_latent.shape, + packed_seq=cond_pack, + packed_seq_uncond=uncond_pack, + guidance=guidance, + normalize_cfg=normalize_cfg, + run_branch=run_branch, + ) # [1,N_tokens_flat] + + initial_noise = curr_vision_latent.flatten(start_dim=1) # [1,N_tokens_flat] + denoised_flat = OmniMoTCausalModel._run_ar_sampler( + host, + velocity_fn, + initial_noise, + sampler_mode=sampler_mode, + num_steps=num_steps, + shift=shift, + seed=seed, + sample_idx=chunk_start, + num_frames=num_frames, + distilled_num_steps=distilled_num_steps, + ) # [1,N_tokens_flat] + return denoised_flat.reshape(curr_vision_latent.shape) # [1,C,V*chunk_len,H,W] + + +@torch.no_grad() +def iter_samples_multiview_transfer_autoregressive( + host: OmniMoTCausalModel, + *, + data_batch: dict, + guidance: float, + seed: int, + num_steps: int, + shift: float, + normalize_cfg: bool, + sampler_mode: str, + distilled_num_steps: int | None, + sync_num_frames_across_ranks: bool, + sync_process_group: Any | None, + max_num_frames: int | None, + on_clean_vision_chunk: Callable[[torch.Tensor], None] | None, + has_negative_prompt: bool, +) -> Generator[dict[str, Any], torch.Tensor | None, None]: + """Generate the target for a two-item camera-major multiview transfer sample.""" + # Imported lazily because the shared model dispatches into this backend. + from cosmos_framework.model.generator.omni_mot_causal_model import _iter_ar_chunk_ranges + + if not host._uses_multiview_replay_kv(): + raise ValueError("Multiview transfer AR requires replayed two-way multiview teacher-forcing configuration.") + if host.config.action_gen or host.config.sound_gen: + raise ValueError("Multiview transfer AR supports vision-only models.") + sequence_plans = build_sequence_plans_from_data_batch( + data_batch=data_batch, + input_video_key=host.input_video_key, + input_image_key=host.input_image_key, + ) + if len(sequence_plans) != 1: + raise ValueError(f"Multiview transfer AR requires batch_size=1, got {len(sequence_plans)} samples.") + caption_groups = _per_view_caption_groups(data_batch[host.input_caption_key]) + host._apply_inference_caption_plan(sequence_plans, caption_groups) + vision_condition_indexes = [sequence_plans[0].condition_frame_indexes_vision] + gen_data_clean = host.get_data_and_condition( + data_batch, + vision_condition_indexes=vision_condition_indexes, + retain_raw_state_vision=False, + ) + host._release_inference_raw_vision(data_batch, gen_data_clean) + if gen_data_clean.x0_tokens_vision is None or len(gen_data_clean.x0_tokens_vision) != 2: + num_items = 0 if gen_data_clean.x0_tokens_vision is None else len(gen_data_clean.x0_tokens_vision) + raise ValueError(f"Multiview transfer AR requires [control, target], got {num_items} vision items.") + # Match joint inference: all CP shards must read the owner's encoded + # controls and conditioned targets, even when local VAE results differ. + if host._get_teacher_forcing_kv_implementation() == "multiview_maskless_kv": + context_parallel_broadcast_tensor_list( + gen_data_clean.x0_tokens_vision, host.parallel_dims + ) # each [1,C,V*T,H,W] + if gen_data_clean.num_vision_items_per_sample != [2]: + raise ValueError( + "Multiview transfer AR requires one sample with exactly two vision items; " + f"got {gen_data_clean.num_vision_items_per_sample}." + ) + if gen_data_clean.num_views_per_vision_item is None or len(gen_data_clean.num_views_per_vision_item) != 2: + raise ValueError("Multiview transfer AR requires per-camera VAE metadata for both vision items.") + num_views = gen_data_clean.num_views_per_vision_item[0] + if gen_data_clean.num_views_per_vision_item != [num_views, num_views]: + raise ValueError(f"Control and target must use the same views, got {gen_data_clean.num_views_per_vision_item}.") + control_latent, target_latent = gen_data_clean.x0_tokens_vision + if control_latent.shape != target_latent.shape: + raise ValueError( + f"Control and target latent shapes must match, got {control_latent.shape} and {target_latent.shape}." + ) + if target_latent.shape[2] % num_views != 0: + raise ValueError(f"Target latent_t={target_latent.shape[2]} is not divisible by {num_views} views.") + frames_per_view = target_latent.shape[2] // num_views + output_frames = frames_per_view if max_num_frames is None else min(frames_per_view, max_num_frames) + if output_frames < 1: + raise ValueError(f"max_num_frames must leave at least one frame, got {max_num_frames}.") + if sync_num_frames_across_ranks and dist.is_available() and dist.is_initialized(): + frame_count = torch.tensor([output_frames], device=target_latent.device, dtype=torch.long) # [1] + dist.all_reduce(frame_count, op=dist.ReduceOp.MIN, group=sync_process_group) + output_frames = int(frame_count.item()) + + cond_text_tokens, uncond_text_tokens = host._get_inference_text_tokens( + data_batch, + has_negative_prompt, + caption_groups, + ) + cfgp_enabled = host.parallel_dims is not None and host.parallel_dims.cfgp_enabled + if cfgp_enabled: + seed = _broadcast_seed([seed], host.parallel_dims.cfgp_mesh.get_group(), host.parallel_dims.cfgp_rank)[0] + cfg_active = guidance != 1.0 or cfgp_enabled + if getattr(host.config, "rolling_kv_cache_chunks", None) is not None: + from cosmos_framework.model.generator.utils.rolling_kv.rolling_transfer_ar import ( + iter_rolling_transfer_ar, + ) + + if normalize_cfg and cfg_active: + raise ValueError("Rolling replay requires normalize_cfg=False when guidance is active.") + if sampler_mode not in ("rf", "distilled"): + raise ValueError("Rolling replay sampler_mode must be rf or distilled.") + for part in iter_rolling_transfer_ar( + host, + plans=sequence_plans, + data=gen_data_clean, + conditional_text=cond_text_tokens, + unconditional_text=uncond_text_tokens, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + output_frames=output_frames, + on_clean_vision_chunk=on_clean_vision_chunk, + prompt_schedule=data_batch.get("rolling_prompt_schedule"), + ): + for frame in _iter_multiview_logical_frames(part["vision"], num_views=num_views): + yield {"vision": frame} + return + fps_vision = gen_data_clean.fps_vision.tolist() if gen_data_clean.fps_vision is not None else [24.0] + backend = host._make_multiview_transfer_ar_backend() + + def build_prefill_pack( + text_tokens: list[list[int]], + *, + materialized_target_frame_ranges: Sequence[tuple[int, int]] | None = None, + ) -> PackedSequence: + return backend.build_prefill_pack( + sequence_plans=sequence_plans, + gen_data_clean=gen_data_clean, + text_tokens=text_tokens, + materialized_target_frame_ranges=materialized_target_frame_ranges, + ) + + cond_prefill = build_prefill_pack( + cond_text_tokens, + materialized_target_frame_ranges=[], + ) + assert cond_prefill.vision is not None + target_condition_mask = cond_prefill.vision.condition_mask[1] # [V*T,1,1] + condition_count = multiview_conditioned_prefix_length( + target_condition_mask, + num_views=num_views, + frames_per_view=frames_per_view, + ) + generated_target = target_latent.to(**host.tensor_kwargs).clone() # [1,C,V*T,H,W] + gen_data_clean.x0_tokens_vision[1] = generated_target + materialized_condition_count = min(condition_count, output_frames) + session = backend.create_session( + prefill_pack=cond_prefill, + num_views=num_views, + frames_per_view=frames_per_view, + condition_count=materialized_condition_count, + cfg_active=cfg_active, + cfgp_enabled=cfgp_enabled, + text_view_ids=sequence_plans[0].text_view_ids, + ) + + controls_read_rgb = host._get_teacher_forcing_replay_policy().controls_read_strict_past_clean_rgb + if not controls_read_rgb: + uncond_prefill = None + if cfg_active: + assert uncond_text_tokens is not None + uncond_prefill = build_prefill_pack( + uncond_text_tokens, + materialized_target_frame_ranges=[], + ) + backend.capture_control_cache( + session=session, + conditional_pack=cond_prefill, + unconditional_pack=uncond_prefill, + ) + + conditioned_prefix = _submit_multiview_conditioned_prefix( + generated_target, + num_views=num_views, + frames_per_view=frames_per_view, + condition_count=condition_count, + output_frames=output_frames, + on_clean_vision_chunk=on_clean_vision_chunk, + ) + if conditioned_prefix is not None: + for prefix_frame in _iter_multiview_logical_frames(conditioned_prefix, num_views=num_views): + yield {"vision": prefix_frame} + for chunk_start, chunk_end in _iter_ar_chunk_ranges( + condition_count, output_frames, host.config.teacher_forcing_frames_per_chunk + ): + chunk_len = chunk_end - chunk_start + if controls_read_rgb: + backend.prime_control_cache( + session=session, + sequence_plans=sequence_plans, + gen_data_clean=gen_data_clean, + conditional_text_tokens=cond_text_tokens, + unconditional_text_tokens=uncond_text_tokens, + ) + memory_layout = backend.build_memory_layout(session) + noise_generator = torch.Generator(device=target_latent.device).manual_seed(seed + chunk_start) + chunk_noise = torch.empty( + ( + 1, + target_latent.shape[1], + num_views * chunk_len, + target_latent.shape[3], + target_latent.shape[4], + ), + device=target_latent.device, + dtype=host.tensor_kwargs["dtype"], + ).normal_(generator=noise_generator) # [1,C,V*chunk_len,H,W] + cond_pack = backend.build_current_pack( + vision_latent=chunk_noise, + text_tokens=cond_text_tokens, + text_view_ids=session.text_view_ids, + vision_view_ids=session.target_view_ids, + fps_vision=fps_vision, + num_views=num_views, + frames_per_view=frames_per_view, + chunk_start=chunk_start, + memory_layout=memory_layout, + current_role="current_target", + ) + uncond_pack = ( + backend.build_current_pack( + vision_latent=chunk_noise, + text_tokens=uncond_text_tokens, + text_view_ids=session.text_view_ids, + vision_view_ids=session.target_view_ids, + fps_vision=fps_vision, + num_views=num_views, + frames_per_view=frames_per_view, + chunk_start=chunk_start, + memory_layout=memory_layout, + current_role="current_target", + ) + if cfg_active + else None + ) + denoised_chunk = host._generate_multiview_transfer_ar_chunk( + cond_pack=cond_pack, + uncond_pack=uncond_pack, + cond_cache=session.conditional_cache, + uncond_cache=session.unconditional_cache, + curr_vision_latent=chunk_noise, + guidance=guidance, + num_steps=num_steps, + shift=shift, + seed=seed, + chunk_start=chunk_start, + num_frames=output_frames, + normalize_cfg=normalize_cfg, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + memory_seq_len=session.memory_seq_len, + ) + backend.scatter_chunk( + generated_target, + denoised_chunk, + num_views=num_views, + frames_per_view=frames_per_view, + chunk_start=chunk_start, + chunk_end=chunk_end, + ) + if on_clean_vision_chunk is not None: + on_clean_vision_chunk(denoised_chunk) + + if chunk_end < output_frames: + backend.commit_clean_chunk( + session=session, + denoised_chunk=denoised_chunk.to(**host.tensor_kwargs), # [1,C,V*chunk_len,H,W] + chunk_start=chunk_start, + chunk_end=chunk_end, + conditional_text_tokens=cond_text_tokens, + unconditional_text_tokens=uncond_text_tokens, + fps_vision=fps_vision, + ) + + # Expose progress after the chunk is ready for future AR steps. Each + # event represents one latent time step across every synchronized view. + for output_frame in _iter_multiview_logical_frames(denoised_chunk, num_views=num_views): + yield {"vision": output_frame} diff --git a/cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py b/cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py new file mode 100644 index 000000000..d5408d7a1 --- /dev/null +++ b/cosmos_framework/model/generator/utils/prope_cross_view_poses_test.py @@ -0,0 +1,146 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +from typing import Any + +import pytest +import torch + +from cosmos_framework.data.generator.multiview.camera_geometry import PER_FRAME_CAMERA_TO_RIG_KEY +from cosmos_framework.model.generator.utils.camera_relative_pose import ( + invert_rigid_transform, + prepare_prope_cross_view_poses, +) +from cosmos_framework.model.generator.utils.rigrope_geometry_test import ( + Tokenizer, + _pose_world_record, + _rotation, + _transform_scene, +) + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def _prepare( + records: list[dict[str, Any] | None], *, views: int = 2, allow_missing: bool = True, **kwargs: Any +) -> list[torch.Tensor | None]: # each [V,F,4,4] + return prepare_prope_cross_view_poses( + records, + pixel_shapes=[(3 * views, 32, 32)] * len(records), + latent_frames=[3 * views] * len(records), + num_views=[views] * len(records), + items_per_sample=[1] * len(records), + tokenizer=Tokenizer(), + allow_missing=allow_missing, + **kwargs, + ) + + +def _only(result: list[torch.Tensor | None]) -> torch.Tensor: # [V,F,4,4] + assert len(result) == 1 and result[0] is not None + return result[0].double() + + +def _relative(poses: torch.Tensor) -> torch.Tensor: # [V,F,4,4] -> [V,F,V,F,4,4] + """Every token pair's ``P_q P_k^-1``, which is all the attention reads of the poses.""" + return poses[:, :, None, None] @ invert_rigid_transform(poses)[None, None] + + +def test_scene_frame_relative_poses_are_the_camera_to_camera_transforms() -> None: + record = _pose_world_record() + cameras = record["camera_to_world"] # [V,F,4,4] + poses = _only(_prepare([record], translation_scale_m=1.0)) # [V,F,4,4] + expected = invert_rigid_transform(cameras)[:, :, None, None] @ cameras[None, None] # [V,F,V,F,4,4] + torch.testing.assert_close(_relative(poses), expected, atol=1e-6, rtol=0) + # "exterior" sorts first and anchors the first latent frame. + torch.testing.assert_close(poses[0, 0], torch.eye(4, dtype=torch.float64), atol=1e-6, rtol=0) + + +def test_rig_per_frame_anchors_each_frame_on_the_reference_camera() -> None: + record = _pose_world_record() + cameras = record["camera_to_world"] # [V,F,4,4] + poses = _only(_prepare([record], translation_scale_m=1.0, pose_world_frame="rig_per_frame")) # [V,F,4,4] + torch.testing.assert_close(poses[0], torch.eye(4, dtype=torch.float64).expand(3, 4, 4), atol=1e-6, rtol=0) + torch.testing.assert_close(poses, invert_rigid_transform(cameras) @ cameras[:1], atol=1e-6, rtol=0) + + +@pytest.mark.parametrize("pose_world_frame", ["scene", "rig_per_frame"]) +def test_rms_unit_makes_poses_invariant_to_the_scene_frame_and_scale(pose_world_frame: str) -> None: + record = _pose_world_record(world_offset=4.0) + moved = _transform_scene(record, _rotation(0.7, -0.4, 1.9), scale=3.0) + options = dict(translation_normalization="per_sample_rms", pose_world_frame=pose_world_frame) + torch.testing.assert_close( + _only(_prepare([moved], **options)), _only(_prepare([record], **options)), atol=1e-5, rtol=0 + ) + # A fixed unit keeps the scene's metres. + assert not torch.allclose(_only(_prepare([moved])), _only(_prepare([record])), atol=1e-3) + + +def test_rms_unit_is_the_rms_camera_distance_from_each_frames_centroid() -> None: + record = _pose_world_record() + cameras = record["camera_to_world"] # [V,F,4,4] + centers = cameras[..., :3, 3] # [V,F,3] + unit = float((centers - centers.mean(dim=0, keepdim=True)).square().sum(-1).mean().sqrt()) + expected = invert_rigid_transform(cameras) @ cameras[:1] # [V,F,4,4] + expected[..., :3, 3] /= unit + poses = _only( + _prepare([record], translation_normalization="per_sample_rms", pose_world_frame="rig_per_frame") + ) # [V,F,4,4] + torch.testing.assert_close(poses, expected, atol=1e-6, rtol=0) + # The floor bounds the unit from below for a rig whose cameras nearly coincide. + floored = _only( + _prepare( + [record], + translation_normalization="per_sample_rms", + translation_floor_m=10 * unit, + pose_world_frame="rig_per_frame", + ) + ) # [V,F,4,4] + torch.testing.assert_close(floored[..., :3, 3], expected[..., :3, 3] / 10, atol=1e-6, rtol=0) + + +def test_reordering_the_cameras_reorders_the_poses_alone() -> None: + record = _pose_world_record() + flipped = dict(record) + flipped["camera_keys"] = record["camera_keys"][::-1] + flipped["calibration"] = record["calibration"][::-1] + for key in ("camera_to_world", "pose_valid", "frame_times_ns", "time_valid", "image_from_calibration"): + flipped[key] = record[key].flip(0) + for frame in ("scene", "rig_per_frame"): + torch.testing.assert_close( + _only(_prepare([flipped], pose_world_frame=frame)), + _only(_prepare([record], pose_world_frame=frame)).flip(0), + atol=1e-6, + rtol=0, + ) + + +def test_static_rig_reads_camera_to_rig_on_every_frame() -> None: + record = _pose_world_record() + rig = record["rig_geometry"] + static = {**record, "rig_geometry": {"camera_to_rig": rig["camera_to_rig"]}} + extrinsics = torch.tensor( + [rig["camera_to_rig"][key] for key in record["camera_keys"]], dtype=torch.float64 + ) # [V,4,4] + poses = _only(_prepare([static], translation_scale_m=1.0)) # [V,F,4,4] + expected = (invert_rigid_transform(extrinsics) @ extrinsics[:1])[:, None].expand(-1, 3, -1, -1) # [V,F,4,4] + torch.testing.assert_close(poses, expected, atol=1e-6, rtol=0) + + +def test_missing_geometry_and_single_view_samples_take_no_poses() -> None: + record = _pose_world_record() + unposed = {**record, "pose_valid": torch.zeros((2, 3), dtype=torch.bool)} + result = _prepare([record, None, unposed]) + assert result[0] is not None and result[1] is None and result[2] is None + assert _prepare([None], views=1, allow_missing=False) == [None] + with pytest.raises(ValueError, match="requires camera_geometry"): + _prepare([None], allow_missing=False) + + +def test_pose_world_rig_with_invalid_poses_fails_closed_without_a_fallback() -> None: + record = _pose_world_record() + unposed = {**record, "pose_valid": torch.zeros((2, 3), dtype=torch.bool)} + assert unposed["rig_geometry"][PER_FRAME_CAMERA_TO_RIG_KEY] + with pytest.raises(ValueError, match="valid poses on every latent frame"): + _prepare([unposed], allow_missing=False) diff --git a/cosmos_framework/model/generator/utils/rigrope_geometry.py b/cosmos_framework/model/generator/utils/rigrope_geometry.py new file mode 100644 index 000000000..5983dece2 --- /dev/null +++ b/cosmos_framework/model/generator/utils/rigrope_geometry.py @@ -0,0 +1,406 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +"""Adapt native camera and Pandar128 calibration to VRAD's eight ray features. + +Only calibration and capture coordinates enter this path. RGB values, LiDAR +ranges, and target depth are deliberately absent from the encoding. Ego-motion is +absent for static-extrinsic rigs; a pose-world rig (a scene-fixed frame shared by +every camera) instead carries each camera's per-frame pose, so moving cameras such +as wrist mounts keep their true rays. A pose-world rig can instead be read about +itself at every encoded frame (``pose_world_frame="rig_per_frame"``), which drops +the motion the cameras share and keeps only their motion against each other. +""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, Literal + +import torch +from torch.nn import functional as F + +from cosmos_framework.data.generator.multiview.camera_geometry import ( + PER_FRAME_CAMERA_TO_RIG_KEY, + _validate_calibration, +) +from cosmos_framework.data.generator.multiview.camera_geometry_contract import validate_camera_to_world_matrix +from cosmos_framework.model.generator.mot.rigrope import rig_features +from cosmos_framework.model.generator.utils.camera_rays import ( + SUPPORTED_RAY_CAMERA_MODELS, + token_pixel_grid, + unproject_calibration, +) +from cosmos_framework.utils.generator.image_resize import get_vision_data_resolution +from cosmos_framework.model.generator.tokenizers.lidar.geometry import pandar128_ray_directions +from cosmos_framework.model.generator.tokenizers.lidar.range_projection import LidarRangeProjectionConfig + +MomentNormalization = Literal["fixed", "per_sample_rms"] +PoseWorldFrame = Literal["scene", "rig_per_frame"] + +# Relabels an OpenCV camera vector (x right, y down, z forward) as forward-left-up. +_OPENCV_TO_FLU = torch.tensor([[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]], dtype=torch.float64) # [3,3] + + +def camera_features( + record: Mapping[str, Any], + rig: Mapping[str, Any], + *, + pixel_shape: tuple[int, int, int], + latent_shape: tuple[int, int, int], + num_views: int, + tokenizer: Any, + moment_scale_m: float | None = None, + include_time: bool = True, + canonical_frame: bool = False, + moment_normalization: MomentNormalization = "fixed", + moment_scale_floor_m: float = 0.05, + pose_world_frame: PoseWorldFrame = "scene", +) -> torch.Tensor: # returns [V,T,H,W,8] + """Sample native rays at VAE centers in retained camera order, before DiT resizing. + + ``canonical_frame`` rotates a pose-world sample into its first view's first encoded camera, + relabelled forward-left-up, and centers a static rig on its cameras; see + ``MultiviewAttentionConfig.rigrope_canonical_frame``. ``moment_normalization="per_sample_rms"`` + replaces ``moment_scale_m`` with the sample's RMS camera distance from the moment origin, + floored at ``moment_scale_floor_m``. ``pose_world_frame="rig_per_frame"`` instead expresses + each encoded frame of a pose-world sample in that frame's first-view camera, relabelled + forward-left-up, about that frame's camera centroid, which makes ``canonical_frame`` moot for + it; see ``MultiviewAttentionConfig.rigrope_pose_world_frame``. + """ + if moment_normalization not in ("fixed", "per_sample_rms"): + raise ValueError(f"Unknown RigRoPE moment normalization {moment_normalization!r}") + if pose_world_frame not in ("scene", "rig_per_frame"): + raise ValueError(f"Unknown RigRoPE pose-world frame {pose_world_frame!r}") + if not math.isfinite(moment_scale_floor_m) or moment_scale_floor_m <= 0: + raise ValueError("RigRoPE moment scale floor must be finite and positive, in metres") + total_frames, height, width = pixel_shape + latent_total, latent_h, latent_w = latent_shape + if total_frames % num_views or latent_total % num_views: + raise ValueError("RigRoPE camera frames must be divisible by the view count") + frames, latent_t = total_frames // num_views, latent_total // num_views + keys = record["camera_keys"] + if len(keys) != num_views or len(set(keys)) != num_views: + raise ValueError("RigRoPE camera keys must match the physical views in packed order") + stride = int(tokenizer.spatial_compression_factor) + pixels = token_pixel_grid(latent_h, latent_w, stride) # [H,W,3] + resolution = get_vision_data_resolution((height, width)) + positions = tokenizer.get_latent_temporal_positions( + num_pixel_frames=frames, resolution=resolution, num_latent_frames=latent_t + ) # [T] or None + if positions is None: + indices = torch.tensor( + [tokenizer.get_pixel_num_frames(i + 1, resolution=resolution) - 1 for i in range(latent_t)], + dtype=torch.long, + ) # [T] + else: + indices = (positions.detach().cpu() * tokenizer.temporal_compression_factor).round().long() # [T] + indices = indices.clamp(0, frames - 1) # [T] + timestamps = torch.as_tensor(record["frame_times_ns"], dtype=torch.float64).cpu() # [V,F] + time_valid = torch.as_tensor(record["time_valid"], dtype=torch.bool).cpu() # [V,F] + if timestamps.shape != (num_views, frames) or not time_valid.all(): + raise ValueError("RigRoPE requires valid timestamps on the selected camera window") + origin_ns = float(timestamps.min()) + per_frame = bool(rig.get(PER_FRAME_CAMERA_TO_RIG_KEY, False)) + if per_frame: + poses = torch.as_tensor(record["camera_to_world"], dtype=torch.float64).cpu() # [V,F,4,4] + pose_valid = torch.as_tensor(record["pose_valid"], dtype=torch.bool).cpu() # [V,F] + if poses.shape != (num_views, frames, 4, 4) or not pose_valid[:, indices].all(): + raise ValueError("A pose-world RigRoPE rig requires valid poses on every encoded camera frame") + frame_poses = poses[:, indices] # [V,T,4,4] + rotations = frame_poses[..., :3, :3] # [V,T,3,3] + if pose_world_frame == "rig_per_frame": + to_rig = _OPENCV_TO_FLU @ rotations[0].transpose(-1, -2) # [T,3,3], scene -> frame-t reference FLU + centers = frame_poses[..., :3, 3] - frame_poses[..., :3, 3].mean(dim=0) # [V,T,3] + rotations = to_rig @ rotations # [V,T,3,3] + centers = torch.einsum("tij,vtj->vti", to_rig, centers) # [V,T,3] + else: + # Center on the sample's first encoded frame so moments do not depend on the world origin. + centers = frame_poses[..., :3, 3] - frame_poses[:, 0, :3, 3].mean(dim=0) # [V,T,3] + if canonical_frame and pose_world_frame == "scene": + to_canonical = _OPENCV_TO_FLU @ rotations[0, 0].T # [3,3], scene -> reference camera FLU + rotations = to_canonical @ rotations # [V,T,3,3] + centers = centers @ to_canonical.T # [V,T,3] + else: + transforms = torch.tensor([rig["camera_to_rig"][key] for key in keys], dtype=torch.float64) # [V,4,4] + rotations = transforms[:, None, :3, :3] # [V,1,3,3] + centers = transforms[:, None, :3, 3] # [V,1,3] + if canonical_frame: + centers = centers - centers.mean(dim=0, keepdim=True) # [V,1,3] + if moment_normalization == "per_sample_rms": + moment_scale_m = max(float(centers.square().sum(dim=-1).mean().sqrt()), moment_scale_floor_m) + rotations, centers = rotations.float(), centers.float() # [V,T|1,3,3], [V,T|1,3] + result: list[torch.Tensor] = [] + for view, key in enumerate(keys): + calibration = record["calibration"][view] + if calibration is None or calibration["camera_model"] not in SUPPORTED_RAY_CAMERA_MODELS: + raise ValueError(f"{key}: RigRoPE requires native calibration with a supported camera model") + # Fail-closed callers never pass a record through ``rigrope_geometry_available``, and an + # invalid focal length would unproject to NaN rays rather than fail. + try: + _validate_calibration(calibration) + except (KeyError, TypeError, ValueError) as error: + raise ValueError(f"{key}: invalid RigRoPE calibration: {error}") from error + affine = torch.as_tensor(record["image_from_calibration"][view], dtype=torch.float64).cpu() # [3,3] + native = pixels @ torch.linalg.inv(affine).T # [H,W,3] + rays = unproject_calibration(native[..., :2] / native[..., 2:], calibration).float() # [H,W,3] + times = ((timestamps[view, indices] - origin_ns) * 1e-9).float() # [T] + if per_frame: + features = torch.cat( + [ + rig_features( + rays @ rotation.T, + center, + times[frame : frame + 1], + moment_scale_m=moment_scale_m, + include_time=include_time, + ) # [1,H,W,8] + for frame, (rotation, center) in enumerate(zip(rotations[view], centers[view], strict=True)) + ] + ) # [T,H,W,8] + else: + features = rig_features( + rays @ rotations[view, 0].T, + centers[view, 0], + times, + moment_scale_m=moment_scale_m, + include_time=include_time, + ) # [T,H,W,8] + result.append(features) + return torch.stack(result) # [V,T,H,W,8] + + +def lidar_features( + rig: Mapping[str, Any], + *, + pixel_shape: tuple[int, int, int], + latent_shape: tuple[int, int, int], + tokenizer: Any, + projection: LidarRangeProjectionConfig, + times: torch.Tensor, # [F], clip-relative seconds + moment_scale_m: float | None = None, + include_time: bool = True, +) -> torch.Tensor: # returns [1,T,H,W,8] + """Use calibrated angular rays including the tokenizer's actual circular padding.""" + frames, height, width = pixel_shape + latent_t, latent_h, latent_w = latent_shape + if projection.model_width_transform != "circular_pad" or (height, width) != ( + projection.semantic_height, + projection.model_width, + ): + raise ValueError("RigRoPE LiDAR requires the configured native angular grid with circular padding") + rays = torch.from_numpy(pandar128_ray_directions(range_projection=projection)) # [H,W_semantic,3] + padding = projection.model_width - projection.semantic_width + left = padding // 2 + columns = (torch.arange(width) - left) % projection.semantic_width # [W_model] + padded = rays[:, columns].permute(2, 0, 1).unsqueeze(0) # [1,3,H,W_model] + spatial = tokenizer.spatial_compression if hasattr(tokenizer, "spatial_compression") else None + spatial = spatial() if callable(spatial) else spatial + stride_h, stride_w = spatial or (tokenizer.spatial_compression_factor,) * 2 + ys = (torch.arange(latent_h, dtype=torch.float32) + 0.5) * stride_h # [H_lat] + xs = (torch.arange(latent_w, dtype=torch.float32) + 0.5) * stride_w # [W_lat] + yy, xx = torch.meshgrid(ys, xs, indexing="ij") # [H_lat,W_lat], [H_lat,W_lat] + grid = torch.stack((2 * xx / (width - 1) - 1, 2 * yy / (height - 1) - 1), dim=-1) # [H_lat,W_lat,2] + sampled = F.grid_sample( + padded, grid[None], mode="bilinear", padding_mode="border", align_corners=True + ) # [1,3,H_lat,W_lat] + sampled = sampled[0].permute(1, 2, 0) # [H_lat,W_lat,3] + transform = torch.tensor(rig["lidar_to_rig"], dtype=torch.float32) # [4,4] + directions = sampled @ transform[:3, :3].T # [H_lat,W_lat,3] + indices = torch.tensor([tokenizer.get_pixel_num_frames(i + 1) - 1 for i in range(latent_t)]) # [T] + if len(times) != frames or int(indices.max()) >= frames: + raise ValueError( + f"RigRoPE LiDAR timestamps must cover the encoded sweeps: times={tuple(times.shape)}, " + f"pixel_frames={frames}, latent_frames={latent_t}, last_index={int(indices.max())}" + ) + return rig_features( + directions, transform[:3, 3], times[indices], moment_scale_m=moment_scale_m, include_time=include_time + ).unsqueeze(0) # [1,T,H_lat,W_lat,8] + + +def rigrope_geometry_available(record: Mapping[str, Any] | None, *, has_lidar: bool) -> bool: + """Gate a whole sample on absent metadata; reject malformed supplied calibration first.""" + if record is None: + return False + if not isinstance(record, Mapping): + raise ValueError("Supplied RigRoPE geometry must be a mapping") + keys = record.get("camera_keys") + calibration = record.get("calibration") + available = keys is not None and calibration is not None + if keys is not None and (not keys or len(keys) != len(set(keys))): + raise ValueError("Supplied camera keys must identify distinct physical cameras") + if calibration is not None: + if keys is not None and len(calibration) != len(keys): + raise ValueError("Supplied calibration must follow the selected camera order") + for cal in calibration: + if cal is None: + available = False + continue + _validate_calibration(cal) + if cal["camera_model"] not in SUPPORTED_RAY_CAMERA_MODELS: + raise ValueError(f"RigRoPE cannot unproject camera model {cal['camera_model']!r}") + provenance = record.get("provenance", {}) + rig = record.get("rig_geometry") + if rig is None: + cameras = provenance.get("cameras", []) + rig = cameras[0].get("rig_geometry") if cameras else None + if rig is None: + available = False + else: + if not isinstance(rig, Mapping) or not isinstance(rig.get("camera_to_rig"), Mapping): + raise ValueError("Malformed supplied sensor rig") + for key, transform in rig["camera_to_rig"].items(): + validate_camera_to_world_matrix(tuple(tuple(row) for row in transform), label=key) + if rig.get("lidar_to_rig") is not None: + validate_camera_to_world_matrix(tuple(tuple(row) for row in rig["lidar_to_rig"]), label="LiDAR") + if keys is not None and any(key not in rig["camera_to_rig"] for key in keys): + available = False + if has_lidar and rig.get("lidar_to_rig") is None: + available = False + if rig.get(PER_FRAME_CAMERA_TO_RIG_KEY, False): + pose_valid = record.get("pose_valid") + if pose_valid is None or record.get("camera_to_world") is None or not torch.as_tensor(pose_valid).all(): + available = False + arrays: dict[str, torch.Tensor] = {} + for key in ("image_from_calibration", "frame_times_ns", "time_valid", "frame_indices"): + value = record.get(key) + if value is None: + if key != "frame_indices" or has_lidar: + available = False + continue + tensor = torch.as_tensor(value) # [V,3,3] or [V,F] + expected_ndim = 3 if key == "image_from_calibration" else 2 + if tensor.ndim != expected_ndim or (keys is not None and tensor.shape[0] != len(keys)): + raise ValueError(f"Malformed supplied {key} shape: {tuple(tensor.shape)}") + if not torch.isfinite(tensor).all(): + raise ValueError(f"Nonfinite supplied {key}") + arrays[key] = tensor + affine = arrays.get("image_from_calibration") # [V,3,3] or None + if affine is not None: + if affine.shape[1:] != (3, 3) or (torch.linalg.det(affine.double()).abs() < 1e-12).any(): + raise ValueError("Supplied image-from-calibration affine must be invertible") + times = arrays.get("frame_times_ns") # [V,F] or None + if times is not None and (times.shape[1] == 0 or (times[:, 1:] < times[:, :-1]).any()): + raise ValueError("Supplied camera timestamps must be nonempty and ordered") + validity = arrays.get("time_valid") # [V,F] or None + if validity is not None: + if times is not None and validity.shape != times.shape: + raise ValueError("Supplied timestamp validity must match camera timestamps") + if not validity.all(): + available = False + fps = provenance.get("source_fps") + if fps is not None: + rates = torch.as_tensor(fps) # [V] + if rates.ndim != 1 or not torch.isfinite(rates).all() or (rates <= 0).any(): + raise ValueError("Supplied camera source FPS must be finite and positive") + if keys is not None and rates.shape[0] != len(keys): + raise ValueError("Supplied source FPS must follow camera order") + elif has_lidar: + available = False + return available + + +def prepare_rigrope_features( + data_batch: Mapping[str, Any], + *, + raw_vision: Sequence[torch.Tensor], # each [1,C,V*F,H,W] + latent_vision: Sequence[torch.Tensor], # each [1,C,V*T,H_lat,W_lat] + raw_lidar: Sequence[torch.Tensor] | None, # each [1,C,F,H,W] + latent_lidar: Sequence[torch.Tensor] | None, # each [1,C,T,H_lat,W_lat] + num_views: Sequence[int], + vision_counts: Sequence[int], + lidar_counts: Sequence[int], + camera_tokenizer: Any, + lidar_tokenizer: Any, + projection: LidarRangeProjectionConfig | None, + allow_missing: bool = False, + moment_scale_m: float | None = None, + include_time: bool = True, + cross_view_only: bool = False, + canonical_frame: bool = False, + moment_normalization: MomentNormalization = "fixed", + moment_scale_floor_m: float = 0.05, + pose_world_frame: PoseWorldFrame = "scene", +) -> tuple[list[torch.Tensor | None], list[torch.Tensor | None]]: + """Prepare independent geometry for every logical sample and repeated control item.""" + if (canonical_frame or moment_normalization != "fixed" or pose_world_frame != "scene") and any(lidar_counts): + # LiDAR descriptors would have to share each sample's camera frame and moment unit. + raise ValueError("RigRoPE canonical frames and per-sample moment units are camera-only") + records = data_batch.get("camera_geometry") + if records is None and (allow_missing or cross_view_only): + # Single-view camera samples need no geometry; cross-view samples still validate below. + records = [None] * len(vision_counts) + if records is None or len(records) != len(vision_counts) or len(lidar_counts) != len(vision_counts): + raise ValueError("RigRoPE needs one camera_geometry record per logical sample") + vision_result: list[torch.Tensor | None] = [] + lidar_result: list[torch.Tensor | None] = [] + vi = li = 0 + for sample, (record, vc, lc) in enumerate(zip(records, vision_counts, lidar_counts, strict=True)): + # A camera-only single-view sample (including its control) has no cross-view pass. + # It needs neither calibration nor descriptors, even inside a mixed packed batch. + if cross_view_only and lc == 0 and all(views == 1 for views in num_views[vi : vi + vc]): + vision_result.extend([None] * vc) + vi += vc + continue + if allow_missing: + available = rigrope_geometry_available(record, has_lidar=lc > 0) + if lc and data_batch.get("lidar_frame_indices") is None: + available = False + if not available: + vision_result.extend([None] * vc) + lidar_result.extend([None] * lc) + vi += vc + li += lc + continue + if record is None: + raise ValueError(f"Sample {sample}: missing RigRoPE camera geometry") + rig = record.get("rig_geometry") or record["provenance"]["cameras"][0].get("rig_geometry") + if rig is None: + raise ValueError(f"Sample {sample}: missing calibrated physical sensor rig") + for item in range(vi, vi + vc): + features = camera_features( + record, + rig, + pixel_shape=tuple(raw_vision[item].shape[-3:]), + latent_shape=tuple(latent_vision[item].shape[-3:]), + num_views=num_views[item], + tokenizer=camera_tokenizer, + moment_scale_m=moment_scale_m, + include_time=include_time, + canonical_frame=canonical_frame, + moment_normalization=moment_normalization, + moment_scale_floor_m=moment_scale_floor_m, + pose_world_frame=pose_world_frame, + ) # [V,T,H,W,8] + vision_result.append(features) + for item in range(li, li + lc): + if raw_lidar is None or latent_lidar is None or projection is None or lidar_tokenizer is None: + raise ValueError("RigRoPE LiDAR item counts disagree with the encoded stream") + frame_ids = data_batch.get("lidar_frame_indices") + if frame_ids is None: + raise ValueError("RigRoPE needs selected LiDAR source-frame IDs to align capture time") + ids = torch.as_tensor(frame_ids[sample], dtype=torch.float64).cpu() # [1,F] or [F] + # The joint packer preserves the singleton batch axis for tensor-origin metadata. + if ids.ndim == 2 and ids.shape[0] == 1: + ids = ids[0] # [F] + if ids.ndim != 1: + raise ValueError(f"Sample {sample}: expected LiDAR source-frame IDs [F] or [1,F], got {ids.shape}") + origin = int(torch.as_tensor(record["frame_indices"])[0, 0]) + source_fps = float(record["provenance"]["source_fps"][0]) + times = ((ids - origin) / source_fps).float() # [F] + features = lidar_features( + rig, + pixel_shape=tuple(raw_lidar[item].shape[-3:]), + latent_shape=tuple(latent_lidar[item].shape[-3:]), + tokenizer=lidar_tokenizer, + projection=projection, + times=times, + moment_scale_m=moment_scale_m, + include_time=include_time, + ) # [1,T,H,W,8] + lidar_result.append(features) + vi += vc + li += lc + if vi != len(raw_vision) or li != len(raw_lidar or []): + raise ValueError("RigRoPE item counts must cover all encoded sensor items") + return vision_result, lidar_result diff --git a/cosmos_framework/model/generator/utils/rigrope_geometry_test.py b/cosmos_framework/model/generator/utils/rigrope_geometry_test.py new file mode 100644 index 000000000..d029a4c26 --- /dev/null +++ b/cosmos_framework/model/generator/utils/rigrope_geometry_test.py @@ -0,0 +1,448 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# ----------------------------------------------------------------------------- + +import math +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from cosmos_framework.data.generator.multiview.camera_geometry import PER_FRAME_CAMERA_TO_RIG_KEY +from cosmos_framework.model.generator.utils.rigrope_geometry import ( + camera_features, + lidar_features, + prepare_rigrope_features, + rigrope_geometry_available, +) +from cosmos_framework.model.generator.tokenizers.lidar.geometry import pandar128_ray_directions +from cosmos_framework.model.generator.tokenizers.lidar.range_projection import V1P2_TRANSFER_RANGE_PROJECTION + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +class Tokenizer: + spatial_compression_factor: int = 16 + temporal_compression_factor: int = 1 + + def get_latent_temporal_positions(self, **kwargs: Any) -> torch.Tensor: + return torch.arange(kwargs["num_latent_frames"], dtype=torch.float32) # [T] + + def get_pixel_num_frames(self, count: int, **kwargs: Any) -> int: + return count + + +def _record(origin_y: float) -> dict[str, Any]: + transform = torch.eye(4) # [4,4] + transform[1, 3] = origin_y # [] + rig = {"camera_to_rig": {"front": transform.tolist()}, "lidar_to_rig": transform.tolist()} + return { + "camera_keys": ["front"], + "rig_geometry": rig, + "provenance": {"source_fps": [30.0]}, + "frame_indices": torch.tensor([[12, 13, 14]]), # [1,F] + "frame_times_ns": torch.tensor([[0, 33333333, 66666667]]), # [1,F] + "time_valid": torch.ones((1, 3), dtype=torch.bool), # [1,F] + "calibration": [ + { + "stream_id": "front", + "camera_model": "ftheta", + "distortion_model": "ftheta", + "distortion_coefficients": [16, 16, 32, 32, 0, 20, 0, 0, 0, 0, 0], + "width_px": 32, + "height_px": 32, + "fx_px": 20.0, + "fy_px": 20.0, + "cx_px": 16.0, + "cy_px": 16.0, + } + ], + "image_from_calibration": torch.eye(3)[None], # [1,3,3] + } + + +def test_single_camera_uses_intrinsics_extrinsics_and_capture_time() -> None: + record = _record(2.0) + result = camera_features( + record, + record["rig_geometry"], + pixel_shape=(3, 32, 32), + latent_shape=(3, 2, 2), + num_views=1, + tokenizer=Tokenizer(), + ) # [1,3,2,2,8] + assert result.shape == (1, 3, 2, 2, 8) + # VAE center (7.5,7.5) in integer-center pixels has radius sqrt(144.5) and forward f=20. + angle = (144.5**0.5) / 20 + assert float(result[0, 0, 0, 0, 2]) == pytest.approx(math.cos(angle), abs=1e-6) + assert float(result[0, 1, 0, 0, 7]) == pytest.approx(1 / 30) + assert torch.count_nonzero(result[..., 6]) == 0 + assert not torch.equal(result[0, 0, 0, 0, :3], result[0, 0, 1, 1, :3]) + + +def _pose_world_record(world_offset: float = 0.0) -> dict[str, Any]: + """Two pinhole cameras in a scene frame; ``wrist`` translates and turns every frame.""" + poses = torch.eye(4, dtype=torch.float64).repeat(2, 3, 1, 1) # [V,F,4,4] + poses[0, :, 0, 3] = 1.0 # [F] + for frame in range(3): + angle = 0.2 * frame + poses[1, frame, :3, :3] = torch.tensor( + [[math.cos(angle), 0.0, math.sin(angle)], [0.0, 1.0, 0.0], [-math.sin(angle), 0.0, math.cos(angle)]] + ) # [3,3] + poses[1, frame, 1, 3] = 0.5 * frame # [] + poses[..., :3, 3] += world_offset # [V,F,3] + pinhole = { + "camera_model": "pinhole", + "distortion_model": "none", + "distortion_coefficients": [], + "width_px": 32, + "height_px": 32, + "fx_px": 16.0, + "fy_px": 16.0, + "cx_px": 15.5, + "cy_px": 15.5, + } + rig = { + "camera_to_rig": {"exterior": poses[0, 0].tolist(), "wrist": poses[1, 0].tolist()}, + PER_FRAME_CAMERA_TO_RIG_KEY: True, + } + return { + "camera_keys": ["exterior", "wrist"], + "rig_geometry": rig, + "provenance": {"source_fps": [30.0, 30.0]}, + "camera_to_world": poses, + "pose_valid": torch.ones((2, 3), dtype=torch.bool), # [V,F] + "frame_times_ns": torch.tensor([[0, 33333333, 66666667]] * 2), # [V,F] + "time_valid": torch.ones((2, 3), dtype=torch.bool), # [V,F] + "calibration": [pinhole, pinhole], + "image_from_calibration": torch.eye(3).repeat(2, 1, 1), # [V,3,3] + } + + +def _pose_world_features(record: dict[str, Any], **kwargs: Any) -> torch.Tensor: # returns [V,T,H,W,8] + return camera_features( + record, + record["rig_geometry"], + pixel_shape=(6, 32, 32), + latent_shape=(6, 2, 2), + num_views=2, + tokenizer=Tokenizer(), + **kwargs, + ) + + +def _transform_scene(record: dict[str, Any], rotation: torch.Tensor, scale: float = 1.0) -> dict[str, Any]: + """The same capture expressed in another scene frame: rotated by ``rotation``, positions scaled.""" + poses = record["camera_to_world"].clone() # [V,F,4,4] + poses[..., :3, :3] = rotation @ poses[..., :3, :3] # [V,F,3,3] + poses[..., :3, 3] = scale * poses[..., :3, 3] @ rotation.T # [V,F,3] + return {**record, "camera_to_world": poses} + + +def _rotation(yaw: float, pitch: float, roll: float) -> torch.Tensor: # returns [3,3] float64 + cz, sz, cy, sy, cx, sx = (f(a) for a in (yaw, pitch, roll) for f in (math.cos, math.sin)) + rz = torch.tensor([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]], dtype=torch.float64) + ry = torch.tensor([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]], dtype=torch.float64) + rx = torch.tensor([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]], dtype=torch.float64) + return rz @ ry @ rx + + +def test_canonical_pose_world_rig_is_invariant_to_the_scene_rotation() -> None: + record = _pose_world_record() + turned = _transform_scene(record, _rotation(0.7, -0.4, 1.9)) + torch.testing.assert_close( + _pose_world_features(record, canonical_frame=True), + _pose_world_features(turned, canonical_frame=True), + atol=1e-5, + rtol=0, + ) + # Without it the scene axes leak into every direction and moment. + assert not torch.allclose(_pose_world_features(record), _pose_world_features(turned), atol=1e-3) + + +def test_canonical_pose_world_rig_puts_the_reference_camera_on_forward_left_up_axes() -> None: + record = _transform_scene(_pose_world_record(), _rotation(0.3, 0.2, -0.5)) + features = _pose_world_features(record, canonical_frame=True, moment_scale_m=1.0) # [2,3,2,2,8] + # The top-left token's ray is half a focal length left (+y) and up (+z) of forward (+x). + torch.testing.assert_close( + features[0, 0, 0, 0, :3], torch.tensor([1.0, 0.5, 0.5]) / math.sqrt(1.5), atol=1e-6, rtol=0 + ) + # The wrist camera sits 1 m to the exterior camera's left (-x in OpenCV, +y here) at frame 0, + # and the origin is their midpoint, so its moment is (0,0.5,0) x d. + ray = features[1, 0, 0, 0, :3] # [3] + expected = torch.linalg.cross(torch.tensor([0.0, 0.5, 0.0]), ray) # [3] + torch.testing.assert_close(features[1, 0, 0, 0, 3:6], expected, atol=1e-6, rtol=0) + + +def test_canonical_static_rig_keeps_its_axes_and_centers_on_its_cameras() -> None: + record = _pose_world_record() + camera_to_rig = {} + for view, key in enumerate(record["camera_keys"]): + transform = record["camera_to_world"][view, 0].clone() # [4,4] + transform[2, 3] += 3.0 # the rig origin sits 3 m behind both cameras + camera_to_rig[key] = transform.tolist() + static = {**record, "rig_geometry": {"camera_to_rig": camera_to_rig}} + plain = _pose_world_features(static, moment_scale_m=1.0) # [2,3,2,2,8] + canonical = _pose_world_features(static, moment_scale_m=1.0, canonical_frame=True) # [2,3,2,2,8] + torch.testing.assert_close(canonical[..., :3], plain[..., :3]) + # Taking the centroid (0.5,0,3) off each origin takes (0.5,0,3) x d off each moment. + shift = torch.linalg.cross(torch.tensor([0.5, 0.0, 3.0]).expand_as(plain[..., :3]), plain[..., :3]) + torch.testing.assert_close(canonical[..., 3:6], plain[..., 3:6] - shift, atol=1e-5, rtol=0) + + +def test_per_sample_moment_unit_is_the_rms_camera_distance_and_scale_invariant() -> None: + record = _pose_world_record() + options: dict[str, Any] = dict(canonical_frame=True, moment_normalization="per_sample_rms") + features = _pose_world_features(record, **options) # [2,3,2,2,8] + # Frame-0 origin (0.5,0,0); centers (0.5,0,0) and (-0.5,0.5t,0) for t=0,1,2, before the axis relabel. + rms = math.sqrt((3 * 0.25 + sum(0.25 + (0.5 * t) ** 2 for t in range(3))) / 6) + torch.testing.assert_close( + features, _pose_world_features(record, canonical_frame=True, moment_scale_m=rms), atol=1e-6, rtol=0 + ) + # A capture ten times larger rotates by the same angles. + enlarged = _transform_scene(record, torch.eye(3, dtype=torch.float64), scale=10.0) + torch.testing.assert_close(_pose_world_features(enlarged, **options), features, atol=1e-5, rtol=0) + + +def test_per_sample_moment_unit_is_floored_for_nearly_coincident_cameras() -> None: + record = _transform_scene(_pose_world_record(), torch.eye(3, dtype=torch.float64), scale=0.001) + torch.testing.assert_close( + _pose_world_features(record, moment_normalization="per_sample_rms", moment_scale_floor_m=0.05), + _pose_world_features(record, moment_scale_m=0.05), + ) + + +_OPENCV_TO_FLU = torch.tensor([[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]], dtype=torch.float64) # [3,3] + + +def _move_rig(record: dict[str, Any], stride_m: float) -> dict[str, Any]: + """Carry every camera on one body that drives ``stride_m`` and turns per frame.""" + poses = record["camera_to_world"].clone() # [V,F,4,4] + for frame in range(poses.shape[1]): + body = torch.eye(4, dtype=torch.float64) # [4,4] + body[:3, :3] = _rotation(0.3 * frame, 0.1 * frame, -0.05 * frame) # [3,3] + body[:3, 3] = torch.tensor([stride_m * frame, 2.0 + 0.1 * stride_m * frame, 0.3 * frame]) # [3] + poses[:, frame] = body @ poses[:, frame] # [V,4,4] + return {**record, "camera_to_world": poses} + + +def _rigid_record() -> dict[str, Any]: + """``_pose_world_record`` with the wrist camera fixed to the exterior one at its frame-0 pose.""" + record = _pose_world_record() + poses = record["camera_to_world"].clone() # [V,F,4,4] + poses[1] = poses[1, :1] # [F,4,4] + return {**record, "camera_to_world": poses} + + +def test_rig_per_frame_reads_a_moving_rigid_rig_as_its_static_extrinsics() -> None: + moving = _move_rig(_rigid_record(), stride_m=15.0) + # The static extrinsics that frame: the exterior camera relabelled forward-left-up, about the centroid. + frame0 = _rigid_record()["camera_to_world"][:, 0] # [V,4,4] + to_rig = _OPENCV_TO_FLU @ frame0[0, :3, :3].T # [3,3] + camera_to_rig = {} + for view, key in enumerate(moving["camera_keys"]): + transform = torch.eye(4, dtype=torch.float64) # [4,4] + transform[:3, :3] = to_rig @ frame0[view, :3, :3] # [3,3] + transform[:3, 3] = to_rig @ (frame0[view, :3, 3] - frame0[:, :3, 3].mean(dim=0)) # [3] + camera_to_rig[key] = transform.tolist() + static = {**moving, "rig_geometry": {"camera_to_rig": camera_to_rig}} + options: dict[str, Any] = dict(moment_scale_m=1.0) + torch.testing.assert_close( + _pose_world_features(moving, pose_world_frame="rig_per_frame", **options), + _pose_world_features(static, canonical_frame=True, **options), + atol=1e-5, + rtol=0, + ) + # Read in one frame, the drive and turn leak into every later frame's directions and moments. + scene = _pose_world_features(moving, canonical_frame=True, **options) # [2,3,2,2,8] + assert not torch.allclose(scene[:, 1:], _pose_world_features(static, canonical_frame=True, **options)[:, 1:]) + + +def test_rig_per_frame_drops_shared_motion_and_keeps_relative_motion() -> None: + record = _pose_world_record() + features = _pose_world_features(record, pose_world_frame="rig_per_frame") # [2,3,2,2,8] + torch.testing.assert_close( + _pose_world_features(_move_rig(record, stride_m=15.0), pose_world_frame="rig_per_frame"), + features, + atol=1e-5, + rtol=0, + ) + # The wrist turns and slides against the exterior camera, so its descriptors still change. + assert not torch.allclose(features[1, 0, ..., :3], features[1, 2, ..., :3], atol=1e-3) + assert not torch.allclose(features[1, 0, ..., 3:6], features[1, 2, ..., 3:6], atol=1e-3) + + +def test_rig_per_frame_rms_unit_is_the_rig_spread_not_its_path() -> None: + moving = _move_rig(_rigid_record(), stride_m=15.0) + options: dict[str, Any] = dict(moment_normalization="per_sample_rms") + # Both cameras sit 0.5 m from their midpoint at every frame. + torch.testing.assert_close( + _pose_world_features(moving, pose_world_frame="rig_per_frame", **options), + _pose_world_features(moving, pose_world_frame="rig_per_frame", moment_scale_m=0.5), + atol=1e-6, + rtol=0, + ) + assert not torch.allclose( + _pose_world_features(moving, canonical_frame=True, **options), + _pose_world_features(moving, canonical_frame=True, moment_scale_m=0.5), + atol=1e-3, + ) + + +def test_rig_per_frame_leaves_static_rigs_alone() -> None: + record = _record(2.0) + kwargs: dict[str, Any] = dict(pixel_shape=(3, 32, 32), latent_shape=(3, 2, 2), num_views=1, tokenizer=Tokenizer()) + torch.testing.assert_close( + camera_features(record, record["rig_geometry"], pose_world_frame="rig_per_frame", **kwargs), + camera_features(record, record["rig_geometry"], **kwargs), + atol=0, + rtol=0, + ) + + +@pytest.mark.parametrize( + "options", + [dict(canonical_frame=True), dict(moment_normalization="per_sample_rms"), dict(pose_world_frame="rig_per_frame")], +) +def test_canonical_frames_refuse_lidar_items(options: dict[str, Any]) -> None: + with pytest.raises(ValueError, match="camera-only"): + prepare_rigrope_features( + {"camera_geometry": [_record(2.0)]}, + raw_vision=[torch.empty(1, 3, 3, 32, 32)], + latent_vision=[torch.empty(1, 4, 3, 2, 2)], + raw_lidar=[torch.empty(1, 3, 3, 128, 1808)], + latent_lidar=[torch.empty(1, 4, 3, 8, 113)], + num_views=[1], + vision_counts=[1], + lidar_counts=[1], + camera_tokenizer=Tokenizer(), + lidar_tokenizer=Tokenizer(), + projection=V1P2_TRANSFER_RANGE_PROJECTION, + **options, + ) + + +def test_pose_world_rig_follows_each_encoded_frame_pose() -> None: + record = _pose_world_record() + features = _pose_world_features(record) # [2,3,2,2,8] + # The token center (7.5, 7.5) sits half a focal length up-left of the principal point. + ray = torch.tensor([-0.5, -0.5, 1.0]) / math.sqrt(1.5) # [3] + for frame in range(3): + rotation = record["camera_to_world"][1, frame, :3, :3].float() # [3,3] + torch.testing.assert_close(features[1, frame, 0, 0, :3], rotation @ ray, atol=1e-6, rtol=0) + # The fixed camera keeps its direction while the moving one turns. + torch.testing.assert_close(features[0, 0, ..., :3], features[0, 2, ..., :3]) + assert not torch.allclose(features[1, 0, ..., :3], features[1, 2, ..., :3]) + assert not torch.allclose(features[1, 0, ..., 3:6], features[1, 2, ..., 3:6]) + + +def test_pose_world_rig_is_invariant_to_the_world_origin() -> None: + torch.testing.assert_close( + _pose_world_features(_pose_world_record()), _pose_world_features(_pose_world_record(world_offset=5000.0)) + ) + + +def test_pose_world_rig_requires_encoded_frame_poses() -> None: + record = _pose_world_record() + record["pose_valid"][1, 1] = False + assert not rigrope_geometry_available(record, has_lidar=False) + with pytest.raises(ValueError, match="pose-world"): + _pose_world_features(record) + + +@pytest.mark.parametrize( + "invalid", [{"fx_px": 0.0}, {"fy_px": -1.0}, {"cx_px": math.nan}], ids=["zero_fx", "negative_fy", "nan_cx"] +) +def test_required_geometry_rejects_invalid_calibration_before_unprojecting(invalid: dict[str, float]) -> None: + """A fail-closed caller skips ``rigrope_geometry_available``, so the features check calibration.""" + record = _record(2.0) + record["calibration"] = [{**record["calibration"][0], **invalid}] + with pytest.raises(ValueError, match="front: invalid RigRoPE calibration"): + prepare_rigrope_features( + {"camera_geometry": [record]}, + raw_vision=[torch.empty(1, 3, 3, 32, 32)], + latent_vision=[torch.empty(1, 4, 3, 2, 2)], + raw_lidar=None, + latent_lidar=None, + num_views=[1], + vision_counts=[1], + lidar_counts=[0], + camera_tokenizer=Tokenizer(), + lidar_tokenizer=Tokenizer(), + projection=V1P2_TRANSFER_RANGE_PROJECTION, + ) + + +def test_mixed_samples_and_control_items_do_not_share_geometry() -> None: + records = [_record(2.0), _record(-2.0)] + pixels = [torch.empty(1, 3, 3, 32, 32)] * 4 # each [1,C,F,H,W] + latents = [torch.empty(1, 4, 3, 2, 2)] * 4 # each [1,C,T,H,W] + cameras, lidars = prepare_rigrope_features( + {"camera_geometry": records}, + raw_vision=pixels, + latent_vision=latents, + raw_lidar=None, + latent_lidar=None, + num_views=[1] * 4, + vision_counts=[2, 2], + lidar_counts=[0, 0], + camera_tokenizer=Tokenizer(), + lidar_tokenizer=Tokenizer(), + projection=V1P2_TRANSFER_RANGE_PROJECTION, + ) + assert lidars == [] + torch.testing.assert_close(cameras[0], cameras[1]) + torch.testing.assert_close(cameras[2], cameras[3]) + torch.testing.assert_close(cameras[0][..., :3], cameras[2][..., :3]) + torch.testing.assert_close(cameras[0][..., 3:6], -cameras[2][..., 3:6]) + + +def test_lidar_angular_grid_wrap_and_nonidentity_rig_frame() -> None: + record = _record(2.0) + tokenizer = SimpleNamespace(spatial_compression=(16, 16), get_pixel_num_frames=lambda count: count) + features = lidar_features( + record["rig_geometry"], + pixel_shape=(3, 128, 1808), + latent_shape=(3, 8, 113), + tokenizer=tokenizer, + projection=V1P2_TRANSFER_RANGE_PROJECTION, + times=torch.tensor([0.0, 0.1, 0.2]), + ) # [1,T,H,W,8] + rays = torch.from_numpy(pandar128_ray_directions(range_projection=V1P2_TRANSFER_RANGE_PROJECTION)) # [128,1800,3] + torch.testing.assert_close(features[0, 0, 0, 0, :3], rays[8, 4], atol=1e-6, rtol=1e-6) + torch.testing.assert_close(features[0, 0, 0, -1, :3], rays[8, 1796], atol=1e-6, rtol=1e-6) + assert torch.count_nonzero(features[..., 6]) == 0 + assert torch.count_nonzero(features[..., 3:6]) > 0 + torch.testing.assert_close(features[0, :, 0, 0, 7], torch.tensor([0.0, 0.1, 0.2])) + + +@pytest.mark.parametrize("batched_ids", [False, True]) +def test_packed_lidar_timestamps_preserve_ragged_samples_and_controls(batched_ids: bool) -> None: + records = [_record(2.0), _record(-2.0)] + ids = [torch.tensor([12, 15, 18]), torch.tensor([12, 15])] # [F0], [F1] + if batched_ids: + # JointDataLoader._get_next_sample retains this axis for tensor-origin metadata. + ids = [value[None] for value in ids] # [1,F0], [1,F1] + pixels = [torch.empty(1, 3, 3, 32, 32)] * 4 # each [1,C,F,H,W] + latents = [torch.empty(1, 4, 3, 2, 2)] * 4 # each [1,C,T,H,W] + raw_lidar = [torch.empty(1, 3, frames, 128, 1808) for frames in (3, 3, 2, 2)] # [1,3,F,H,W] + lidar_latents = [torch.empty(1, 4, frames, 8, 113) for frames in (3, 3, 2, 2)] # [1,4,T,H_lat,W_lat] + _, features = prepare_rigrope_features( + {"camera_geometry": records, "lidar_frame_indices": ids}, + raw_vision=pixels, + latent_vision=latents, + raw_lidar=raw_lidar, + latent_lidar=lidar_latents, + num_views=[1] * 4, + vision_counts=[2, 2], + lidar_counts=[2, 2], + camera_tokenizer=Tokenizer(), + lidar_tokenizer=SimpleNamespace(spatial_compression=(16, 16), get_pixel_num_frames=lambda count: count), + projection=V1P2_TRANSFER_RANGE_PROJECTION, + ) + torch.testing.assert_close(features[0][0, :, 0, 0, 7], torch.tensor([0.0, 0.1, 0.2])) + torch.testing.assert_close(features[2][0, :, 0, 0, 7], torch.tensor([0.0, 0.1])) + torch.testing.assert_close(features[0], features[1]) + torch.testing.assert_close(features[2], features[3]) + torch.testing.assert_close(features[0][:, :2, ..., 3:6], -features[2][..., 3:6]) diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py index 0eee69b0f..31cdbc0d5 100644 --- a/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py @@ -15,6 +15,7 @@ from cosmos_framework.utils import log from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan +from cosmos_framework.utils.generator.multiview import slice_camera_major_frames from cosmos_framework.model.generator.teacher_forcing import mark_modality_as_clean_condition from cosmos_framework.model.generator.utils.multiview_ar import JointARChunk, joint_ar_chunks from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( @@ -30,19 +31,6 @@ from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel -def _camera_slice( - latent: torch.Tensor, - num_views: int, - start: int, - stop: int, # latent: [1,C,V*T,H,W] -) -> torch.Tensor: # [1,C,V*(stop-start),H,W] - batch, channels, total, height, width = latent.shape - grid = latent.reshape(batch, channels, num_views, total // num_views, height, width) # [1,C,V,T,H,W] - return grid[:, :, :, start:stop].reshape( - batch, channels, num_views * (stop - start), height, width - ) # [1,C,V*Tc,H,W] - - def _initial_target( target: torch.Tensor, num_views: int, @@ -80,8 +68,8 @@ def _chunk_data( raw_state_vision=None, raw_state_lidar=None, x0_tokens_vision=[ - _camera_slice(data.x0_tokens_vision[0], num_views, start, stop), - _camera_slice(vision, num_views, start, stop), + slice_camera_major_frames(data.x0_tokens_vision[0], num_views, start, stop - start), + slice_camera_major_frames(vision, num_views, start, stop - start), ], # each [1,Cv,V*Tvc,Hv,Wv] x0_tokens_lidar=lidar_items, ) diff --git a/cosmos_framework/utils/generator/multiview.py b/cosmos_framework/utils/generator/multiview.py index e4991e098..6f853aa22 100644 --- a/cosmos_framework/utils/generator/multiview.py +++ b/cosmos_framework/utils/generator/multiview.py @@ -364,6 +364,20 @@ def normalize_multiview_control_weights(weights: list[float]) -> list[float]: return [weight / total_weight for weight in weights] +def slice_camera_major_frames( + clip: torch.Tensor, # [C,V*T,H,W] or [B,C,V*T,H,W] + n_views: int, + start: int, + length: int, +) -> torch.Tensor: # [C,V*length,H,W] or [B,C,V*length,H,W] + """Trim the same temporal window from every camera-major view slice.""" + *leading, total_frames, height, width = clip.shape + frames_per_view = total_frames // n_views + clip = clip.reshape(*leading, n_views, frames_per_view, height, width) # [...,V,T,H,W] + clip = clip[..., start : start + length, :, :] # [...,V,length,H,W] + return clip.reshape(*leading, n_views * length, height, width) # [...,V*length,H,W] + + def slice_multiview_view_frames( frames: torch.Tensor, *, diff --git a/cosmos_framework/utils/generator/multiview_test.py b/cosmos_framework/utils/generator/multiview_test.py index 0545e0e4d..c49eb3acd 100644 --- a/cosmos_framework/utils/generator/multiview_test.py +++ b/cosmos_framework/utils/generator/multiview_test.py @@ -11,9 +11,28 @@ iter_multiview_video_by_view, normalize_multiview_control_weights, pad_multiview_view_video, + slice_camera_major_frames, slice_multiview_view_frames, ) + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("batched", [False, True]) +def test_slice_camera_major_frames_preserves_each_views_temporal_window(batched: bool) -> None: + storage = torch.arange(2 * 3 * 15 * 2 * 4).reshape(2, 3, 15, 2, 4) # [B,C,V*T,H,2*W] + clip = storage[..., ::2] if batched else storage[0, ..., ::2] # [B,C,V*T,H,W] or [C,V*T,H,W] + for views in (1, 3): + frames = clip.shape[-3] // views + for start, length in ((0, frames), (0, 1), (1, 2), (frames, 0)): + expected = torch.cat( + [clip[..., view * frames + start : view * frames + start + length, :, :] for view in range(views)], + dim=-3, + ) # [...,V*length,H,W] + actual = slice_camera_major_frames(clip, views, start, length) # [...,V*length,H,W] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + # --------------------------------------------------------------------------- # decode_multiview_latent_per_view # --------------------------------------------------------------------------- diff --git a/cosmos_framework/utils/generator/video_preprocess.py b/cosmos_framework/utils/generator/video_preprocess.py index d5dfbc686..7bb8472d5 100644 --- a/cosmos_framework/utils/generator/video_preprocess.py +++ b/cosmos_framework/utils/generator/video_preprocess.py @@ -25,6 +25,79 @@ def resize_and_center_crop_tensor(frames: torch.Tensor, target_h: int, target_w: return torchvision_F.center_crop(resized_frames, [target_h, target_w]) # [...,target_h,target_w] +def resize_and_center_crop_affine(source_hw: tuple[int, int], target_hw: tuple[int, int]) -> torch.Tensor: + """Source-to-output pixel affine of ``resize_and_center_crop_tensor`` in integer-center coordinates.""" + (orig_h, orig_w), (target_h, target_w) = source_hw, target_hw + if min(orig_h, orig_w, target_h, target_w) <= 0: + raise ValueError(f"Image sizes must be positive, got source={source_hw}, target={target_hw}.") + scaling_ratio = max(target_w / orig_w, target_h / orig_h) + resize_h = int(math.ceil(scaling_ratio * orig_h)) + resize_w = int(math.ceil(scaling_ratio * orig_w)) + # Matches torchvision center_crop offsets, which round half to even. + top, left = int(round((resize_h - target_h) / 2.0)), int(round((resize_w - target_w) / 2.0)) + sx, sy = resize_w / orig_w, resize_h / orig_h + return torch.tensor( # [3,3] + [[sx, 0.0, (sx - 1) / 2 - left], [0.0, sy, (sy - 1) / 2 - top], [0.0, 0.0, 1.0]], dtype=torch.float64 + ) + + +def _resize_and_crop_window( + source_hw: tuple[int, int], target_hw: tuple[int, int], zoom: float, crop_center: tuple[float, float] +) -> tuple[int, int, int, int]: + """Resized size and crop origin ``(resize_h, resize_w, top, left)`` of ``resize_and_crop_tensor``.""" + (orig_h, orig_w), (target_h, target_w) = source_hw, target_hw + if min(orig_h, orig_w, target_h, target_w) <= 0: + raise ValueError(f"Image sizes must be positive, got source={source_hw}, target={target_hw}.") + if not (math.isfinite(zoom) and zoom >= 1.0): + raise ValueError(f"Crop zoom must be finite and at least 1, got {zoom}.") + if not all(0.0 <= fraction <= 1.0 for fraction in crop_center): + raise ValueError(f"Crop center fractions must lie in [0, 1], got {crop_center}.") + scaling_ratio = max(target_w / orig_w, target_h / orig_h) * zoom + resize_h = int(math.ceil(scaling_ratio * orig_h)) + resize_w = int(math.ceil(scaling_ratio * orig_w)) + # Round half to even, as torchvision center_crop does, so (0.5, 0.5) is its center crop. + top = int(round(crop_center[0] * (resize_h - target_h))) + left = int(round(crop_center[1] * (resize_w - target_w))) + return resize_h, resize_w, top, left + + +def resize_and_crop_tensor( + frames: torch.Tensor, + target_h: int, + target_w: int, + *, + zoom: float = 1.0, + crop_center: tuple[float, float] = (0.5, 0.5), +) -> torch.Tensor: + """Aspect-ratio-preserving resize to ``zoom`` times the covering size, then a target-size crop. + + ``crop_center`` places the crop along the slack in each axis, (top, left) fractions from 0 + (first rows or columns) to 1 (last); ``zoom=1`` at (0.5, 0.5) is ``resize_and_center_crop_tensor``. + """ + if frames.dim() < 2: + raise ValueError(f"Expected image-like tensor with trailing H,W dimensions, got shape {tuple(frames.shape)}.") + resize_h, resize_w, top, left = _resize_and_crop_window( + (int(frames.shape[-2]), int(frames.shape[-1])), (target_h, target_w), zoom, crop_center + ) + resized_frames = torchvision_F.resize(frames, [resize_h, resize_w], antialias=True) # [...,resize_h,resize_w] + return resized_frames[..., top : top + target_h, left : left + target_w] # [...,target_h,target_w] + + +def resize_and_crop_affine( + source_hw: tuple[int, int], + target_hw: tuple[int, int], + *, + zoom: float = 1.0, + crop_center: tuple[float, float] = (0.5, 0.5), +) -> torch.Tensor: + """Source-to-output pixel affine of ``resize_and_crop_tensor`` in integer-center coordinates.""" + resize_h, resize_w, top, left = _resize_and_crop_window(source_hw, target_hw, zoom, crop_center) + sx, sy = resize_w / source_hw[1], resize_h / source_hw[0] + return torch.tensor( # [3,3] + [[sx, 0.0, (sx - 1) / 2 - left], [0.0, sy, (sy - 1) / 2 - top], [0.0, 0.0, 1.0]], dtype=torch.float64 + ) + + def tensor_to_pil_images(video_tensor: torch.Tensor, *, channels_first: bool | None = None) -> list[Image.Image]: """Convert a video tensor of shape (C, T, H, W) or (T, C, H, W) into a list of PIL images.