Spaces:
Running on Zero
Running on Zero
| """FCP geometry and checkpoint identity checks. | |
| Excluded checkpoint families are recognized for rejection, without retaining their | |
| model construction or execution branches. Public inputs use future_prediction. | |
| """ | |
| from collections.abc import Mapping, Sequence | |
| from dataclasses import dataclass | |
| from omegaconf import OmegaConf | |
| PREDICTIVE_SOURCE_LAYERS_METADATA_BUFFER = "predictive_multi_chunk_source_layers_metadata" | |
| PREDICTIVE_OFFSET_SOURCE_MODE_MARKER = "predictive_offset_source_mode_marker" | |
| PREDICTIVE_CHAINED_SOURCE_MODE_MARKER_VALUE = 1 | |
| PREDICTIVE_SOURCE_UNSHARED_SCOPE_MARKER = "predictive_source_unshared_scope_marker" | |
| PREDICTIVE_SOURCE_CONDITIONING_SCOPE_MARKER_VALUE = 1 | |
| PREDICTIVE_AUX_SHARED_ACROSS_OFFSETS_MARKER = "predictive_aux_shared_across_offsets_marker" | |
| PREDICTIVE_DIRECT_MODULE_NAMES = ( | |
| "patch_embedding", | |
| "time_embedding", | |
| "time_projection", | |
| "text_embedding", | |
| "head", | |
| ) | |
| PREDICTIVE_SHARED_DIRECT_MODULES_NAMESPACE = "predictive_multi_chunk_shared_direct_modules" | |
| def is_predictive_source_layers_metadata_key(key: str) -> bool: | |
| """Return whether ``key`` stores the ordered MLP source-layer signature.""" | |
| return key == PREDICTIVE_SOURCE_LAYERS_METADATA_BUFFER or key.endswith( | |
| f".{PREDICTIVE_SOURCE_LAYERS_METADATA_BUFFER}" | |
| ) | |
| def is_predictive_offset_source_mode_marker_key(key: str) -> bool: | |
| """Return whether ``key`` stores chained offset-source identity.""" | |
| return key == PREDICTIVE_OFFSET_SOURCE_MODE_MARKER or key.endswith( | |
| f".{PREDICTIVE_OFFSET_SOURCE_MODE_MARKER}" | |
| ) | |
| def is_predictive_source_unshared_scope_marker_key(key: str) -> bool: | |
| """Return whether ``key`` stores source-conditioning scope identity.""" | |
| return key == PREDICTIVE_SOURCE_UNSHARED_SCOPE_MARKER or key.endswith( | |
| f".{PREDICTIVE_SOURCE_UNSHARED_SCOPE_MARKER}" | |
| ) | |
| def is_predictive_aux_shared_across_offsets_marker_key(key: str) -> bool: | |
| """Return whether ``key`` stores cross-offset direct-module identity.""" | |
| return key == PREDICTIVE_AUX_SHARED_ACROSS_OFFSETS_MARKER or key.endswith( | |
| f".{PREDICTIVE_AUX_SHARED_ACROSS_OFFSETS_MARKER}" | |
| ) | |
| def checkpoint_has_predictive_aux_state(state_dict: Mapping) -> bool: | |
| """Return whether a checkpoint contains any predictive aux namespace.""" | |
| return any( | |
| ".predictive_multi_chunk_branches." in key | |
| or key.startswith("predictive_multi_chunk_branches.") | |
| or key.startswith("model.predictive_multi_chunk_branches.") | |
| or f".{PREDICTIVE_SHARED_DIRECT_MODULES_NAMESPACE}." in key | |
| or key.startswith(f"{PREDICTIVE_SHARED_DIRECT_MODULES_NAMESPACE}.") | |
| or key.startswith(f"model.{PREDICTIVE_SHARED_DIRECT_MODULES_NAMESPACE}.") | |
| or is_predictive_source_layers_metadata_key(key) | |
| or is_predictive_offset_source_mode_marker_key(key) | |
| or is_predictive_source_unshared_scope_marker_key(key) | |
| or is_predictive_aux_shared_across_offsets_marker_key(key) | |
| for key in state_dict | |
| ) | |
| def validate_aux_shared_across_offsets_checkpoint( | |
| state_dict: Mapping, | |
| *, | |
| label: str, | |
| ) -> bool: | |
| """Require exact cross-offset direct-module checkpoint-family identity.""" | |
| source_is_predictive = checkpoint_has_predictive_aux_state(state_dict) | |
| if not source_is_predictive: | |
| return False | |
| marker_keys = [ | |
| key for key in state_dict if is_predictive_aux_shared_across_offsets_marker_key(key) | |
| ] | |
| if len(marker_keys) > 1: | |
| raise RuntimeError( | |
| f"{label} contains duplicate predictive aux_shared_across_offsets markers: " | |
| f"{marker_keys}." | |
| ) | |
| checkpoint_value = 0 | |
| if marker_keys: | |
| marker = state_dict[marker_keys[0]] | |
| if ( | |
| not hasattr(marker, "detach") | |
| or marker.numel() != 1 | |
| or marker.ndim != 0 | |
| or marker.is_floating_point() | |
| or marker.is_complex() | |
| or str(marker.dtype) == "torch.bool" | |
| ): | |
| raise RuntimeError(f"{label} predictive aux_shared_across_offsets marker is malformed.") | |
| checkpoint_value = marker.detach().cpu().item() | |
| max_value = (1 << len(PREDICTIVE_DIRECT_MODULE_NAMES)) - 1 | |
| if ( | |
| isinstance(checkpoint_value, bool) | |
| or checkpoint_value <= 0 | |
| or checkpoint_value > max_value | |
| ): | |
| raise RuntimeError( | |
| f"{label} predictive aux_shared_across_offsets marker has unknown " | |
| f"value {checkpoint_value!r}." | |
| ) | |
| if checkpoint_value != 0: | |
| raise RuntimeError( | |
| f"{label} predictive aux_shared_across_offsets mismatch: " | |
| f"checkpoint_marker={checkpoint_value}, runtime_marker=0." | |
| ) | |
| return True | |
| def validate_offset_source_mode_checkpoint( | |
| state_dict: Mapping, | |
| *, | |
| label: str, | |
| ) -> bool: | |
| """Require exact direct/chained identity for predictive warm starts.""" | |
| source_is_predictive = checkpoint_has_predictive_aux_state(state_dict) | |
| if not source_is_predictive: | |
| return False | |
| marker_keys = [key for key in state_dict if is_predictive_offset_source_mode_marker_key(key)] | |
| if len(marker_keys) > 1: | |
| raise RuntimeError( | |
| f"{label} contains duplicate predictive offset-source mode markers: {marker_keys}." | |
| ) | |
| checkpoint_mode = "direct" | |
| if marker_keys: | |
| marker = state_dict[marker_keys[0]] | |
| if ( | |
| not hasattr(marker, "detach") | |
| or marker.numel() != 1 | |
| or marker.ndim != 0 | |
| or marker.is_floating_point() | |
| or marker.is_complex() | |
| ): | |
| raise RuntimeError(f"{label} predictive offset-source mode marker is malformed.") | |
| value = marker.detach().cpu().item() | |
| if isinstance(value, bool) or value != PREDICTIVE_CHAINED_SOURCE_MODE_MARKER_VALUE: | |
| raise RuntimeError( | |
| f"{label} predictive offset-source mode marker has unknown value {value!r}." | |
| ) | |
| checkpoint_mode = "chained" | |
| if checkpoint_mode != "direct": | |
| raise RuntimeError( | |
| f"{label} predictive offset-source mode mismatch: " | |
| f"checkpoint={checkpoint_mode}, runtime=direct." | |
| ) | |
| return True | |
| def validate_source_unshared_scope_checkpoint( | |
| state_dict: Mapping, | |
| *, | |
| label: str, | |
| ) -> bool: | |
| """Require exact legacy/source-conditioning checkpoint-family identity.""" | |
| source_is_predictive = checkpoint_has_predictive_aux_state(state_dict) | |
| if not source_is_predictive: | |
| return False | |
| marker_keys = [key for key in state_dict if is_predictive_source_unshared_scope_marker_key(key)] | |
| if len(marker_keys) > 1: | |
| raise RuntimeError( | |
| f"{label} contains duplicate predictive source-unshared scope markers: {marker_keys}." | |
| ) | |
| checkpoint_scope = "source_conditioning" if marker_keys else None | |
| if marker_keys: | |
| marker = state_dict[marker_keys[0]] | |
| if ( | |
| not hasattr(marker, "detach") | |
| or marker.numel() != 1 | |
| or marker.ndim != 0 | |
| or marker.is_floating_point() | |
| or marker.is_complex() | |
| or str(marker.dtype) == "torch.bool" | |
| ): | |
| raise RuntimeError(f"{label} predictive source-unshared scope marker is malformed.") | |
| value = marker.detach().cpu().item() | |
| if isinstance(value, bool) or value != PREDICTIVE_SOURCE_CONDITIONING_SCOPE_MARKER_VALUE: | |
| raise RuntimeError( | |
| f"{label} predictive source-unshared scope marker has unknown value {value!r}." | |
| ) | |
| if checkpoint_scope is not None: | |
| raise RuntimeError( | |
| f"{label} predictive source-unshared scope mismatch: " | |
| f"checkpoint={checkpoint_scope or 'legacy'}, " | |
| f"runtime=legacy." | |
| ) | |
| return True | |
| def validate_mlp_source_layer_checkpoint( | |
| state_dict: Mapping, | |
| *, | |
| source_layers: Sequence[int], | |
| label: str, | |
| ) -> bool: | |
| """Validate ordered source-layer metadata for an MLP target. | |
| Returns true when the source checkpoint is predictive. A checkpoint with | |
| no predictive aux state is a valid non-predictive warm start and returns | |
| false without requiring metadata. | |
| """ | |
| source_is_predictive = checkpoint_has_predictive_aux_state(state_dict) | |
| if not source_is_predictive: | |
| return source_is_predictive | |
| metadata_keys = [key for key in state_dict if is_predictive_source_layers_metadata_key(key)] | |
| if len(metadata_keys) != 1: | |
| raise RuntimeError( | |
| f"{label} must contain exactly one ordered source-layer metadata tensor " | |
| f"for an MLP target; found {metadata_keys}." | |
| ) | |
| metadata = state_dict[metadata_keys[0]] | |
| if not hasattr(metadata, "detach"): | |
| raise RuntimeError(f"{label} source-layer metadata must be a tensor.") | |
| if metadata.ndim != 1: | |
| raise RuntimeError(f"{label} source-layer metadata must be a one-dimensional tensor.") | |
| metadata_values = metadata.detach().cpu().tolist() | |
| if any(isinstance(value, bool) or not isinstance(value, int) for value in metadata_values): | |
| raise RuntimeError(f"{label} source-layer metadata must contain integers.") | |
| checkpoint_layers = list(metadata_values) | |
| expected_layers = [int(value) for value in source_layers] | |
| if checkpoint_layers != expected_layers: | |
| raise RuntimeError( | |
| f"{label} source_layers mismatch: checkpoint={checkpoint_layers}, " | |
| f"runtime={expected_layers}." | |
| ) | |
| return True | |
| class PredictiveChunkLayout: | |
| num_frames: int | |
| block_size: int | |
| independent_first_frame: bool | |
| regular_frame_start: int | |
| num_regular_chunks: int | |
| def first_regular_chunk_index(self): | |
| return int(self.independent_first_frame) | |
| def valid_pair_count(self, future_offset): | |
| return self.num_regular_chunks - future_offset | |
| def regular_pair_indices(self, future_offset): | |
| return tuple( | |
| (self.first_regular_chunk_index + i, self.first_regular_chunk_index + i + future_offset) | |
| for i in range(self.valid_pair_count(future_offset)) | |
| ) | |
| def build_predictive_chunk_layout( | |
| *, num_frames, block_size, independent_first_frame, consumer_name="FCP" | |
| ): | |
| start = int(independent_first_frame) | |
| if block_size <= 0 or num_frames <= start or (num_frames - start) % block_size: | |
| raise ValueError( | |
| f"{consumer_name}: incompatible chunk layout T={num_frames}, " | |
| f"block={block_size}, independent_first_frame={independent_first_frame}" | |
| ) | |
| return PredictiveChunkLayout( | |
| num_frames, block_size, independent_first_frame, start, (num_frames - start) // block_size | |
| ) | |
| def read_predictive_multi_chunk_config(args, model): | |
| # The public resolver validates names and value domains before model loading. | |
| cfg = OmegaConf.to_container(args.future_prediction, resolve=True) | |
| if not cfg["enabled"]: | |
| raise ValueError("FCP requires future_prediction.enabled=true") | |
| layout = build_predictive_chunk_layout( | |
| num_frames=args.image_or_video_shape[1], | |
| block_size=args.num_frame_per_block, | |
| independent_first_frame=args.independent_first_frame, | |
| ) | |
| layers = cfg["source_layers"] | |
| if layers[-1] >= len(model.blocks): | |
| raise ValueError("future_prediction.source_layers exceeds the loaded backbone") | |
| depth = cfg["mini_dit_depth"] | |
| init_source = cfg.get("init_source") | |
| if cfg["init_mode"] == "copy": | |
| if init_source is None: | |
| if len(layers) != 1: | |
| raise ValueError("future_prediction.init_source is required for multi-layer copy") | |
| init_source = list(range(layers[0] + 1, layers[0] + 1 + depth)) | |
| if ( | |
| len(init_source) != depth | |
| or len(set(init_source)) != depth | |
| or any(i < 0 or i >= len(model.blocks) for i in init_source) | |
| ): | |
| raise ValueError("future_prediction.init_source is incompatible with loaded blocks") | |
| elif "init_source" in cfg: | |
| raise ValueError("future_prediction.init_source must be absent for random initialization") | |
| hidden_dim = round(cfg["source_layer_fusion"]["hidden_ratio"] * model.dim) | |
| if hidden_dim <= 0: | |
| raise ValueError("future_prediction.source_layer_fusion.hidden_ratio gives zero width") | |
| if any(layout.valid_pair_count(offset) <= 0 for offset in cfg["future_offsets"]): | |
| raise ValueError("future_prediction.future_offsets leaves no regular pair") | |
| return { | |
| **cfg, | |
| "init_source": init_source, | |
| "source_layer_fusion": {**cfg["source_layer_fusion"], "hidden_dim": hidden_dim}, | |
| "source_context_streams": cfg["source_context"]["streams"], | |
| "num_frame_per_block": args.num_frame_per_block, | |
| "independent_first_frame": args.independent_first_frame, | |
| "num_chunks": layout.num_regular_chunks, | |
| } | |
| def configure_predictive_multi_chunk_from_config(model, config): | |
| model.configure_predictive_multi_chunk(config) | |