# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Implements the forward op for training, validation, and inference.""" from typing import Any import torch from cosmos_predict1.tokenizer.training.datasets.utils import IMAGE_KEY, INPUT_KEY, MASK_KEY, RECON_KEY, VIDEO_KEY from cosmos_predict1.tokenizer.training.losses.continuous import RECON_CONSISTENCY_KEY, VIDEO_CONSISTENCY_LOSS from cosmos_predict1.utils import ema from cosmos_predict1.utils.lazy_config import LazyDict, instantiate from cosmos_predict1.utils.model import Model PREDICTION = "prediction" EMA_PREDICTION = "ema_prediction" class TokenizerModel(Model): def __init__(self, config) -> None: super().__init__() self.config = config self.network = instantiate(config.network) self.loss = instantiate(config.loss) self.metric = instantiate(config.metric) self.precision = getattr(torch, config.precision) if self.config.ema.enabled: self.ema = ema.EMAModelTracker( self, beta=self.config.ema.beta, torch_compile_buffer_renaming=self.config.ema.torch_compile_buffer_renaming, ) self.init_input_keys() def init_input_keys(self): self.image_key = IMAGE_KEY self.video_key = VIDEO_KEY def get_input_key(self, data_batch: dict[str, torch.Tensor]) -> str: if self.image_key in data_batch: return self.image_key elif self.video_key in data_batch: return self.video_key else: raise ValueError("Input key not found in data_batch.") def init_optimizer_scheduler( self, optimizer_config: LazyDict, scheduler_config: LazyDict ) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]: """Creates the optimizer and scheduler for the network. Args: optimizer_config: The optimizer config for the net. scheduler_config: The scheduler config for the net. Returns: optimizer (torch.optim.Optimizer): The net optimizer. scheduler (torch.optim.lr_scheduler.LRScheduler): The net optimization scheduler. """ optimizer_config.params = self.network.parameters() optimizer = instantiate(optimizer_config) scheduler_config.optimizer = optimizer scheduler = instantiate(scheduler_config) return optimizer, scheduler def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None: if self.config.ema.enabled: self.ema.to(dtype=torch.float32) self.network = self.network.to(dtype=self.precision, memory_format=memory_format) self.loss = self.loss.to(dtype=self.precision, memory_format=memory_format) def state_dict( self, destination: dict[str, Any] = None, prefix: str = "", keep_vars: bool = False ) -> dict[str, Any]: original_state_dict = super(TokenizerModel, self).state_dict(destination, prefix, keep_vars) # Filter out '.loss' and 'ema.loss-' keys from the state dict. filtered_state_dict = {k: v for k, v in original_state_dict.items() if not k.startswith("loss.")} filtered_state_dict = {k: v for k, v in filtered_state_dict.items() if not k.startswith("ema.loss-")} filtered_state_dict = { k: v for k, v in filtered_state_dict.items() if not k.startswith("network.encoder.patcher") } filtered_state_dict = { k: v for k, v in filtered_state_dict.items() if not k.startswith("network.decoder.unpatcher") } return filtered_state_dict def load_state_dict(self, state_dict: Any, strict: bool = True) -> None: own_state = self.state_dict() filtered_state_dict = {k: v for k, v in state_dict.items() if k in own_state} # Load only filtered state dict. super(TokenizerModel, self).load_state_dict(filtered_state_dict, strict=False) # If strict is True, ensure all parameters are loaded (except the excluded ones) missing_keys = set(own_state.keys()) - set(filtered_state_dict.keys()) if missing_keys and strict: raise KeyError(f"Missing keys in state_dict: {missing_keys}") def _on_before_network_forward(self, data_batch: dict[str, torch.Tensor]) -> None: consistency_loss = self.loss.loss_modules[VIDEO_CONSISTENCY_LOSS] if hasattr(consistency_loss, "enabled") and consistency_loss.enabled: _input_key = self.get_input_key(data_batch) if _input_key is self.video_key: data_batch[_input_key] = consistency_loss.shuffle(data_batch[_input_key]) return def _on_after_network_forward( self, data_batch: dict[str, torch.Tensor], output_batch: dict[str, torch.Tensor] ) -> None: consistency_loss = self.loss.loss_modules[VIDEO_CONSISTENCY_LOSS] if hasattr(consistency_loss, "enabled") and consistency_loss.enabled: _input_key = self.get_input_key(data_batch) if _input_key is self.video_key: data_batch[_input_key] = consistency_loss.unshuffle(data_batch[_input_key]) output_batch[RECON_CONSISTENCY_KEY] = torch.ones_like(output_batch[RECON_KEY]) * output_batch[RECON_KEY] output_batch[RECON_KEY] = consistency_loss.unshuffle(output_batch[RECON_KEY]) return def _network_forward(self, data_batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: # A callback proxy to modify the input before the forward pass. self._on_before_network_forward(data_batch) # Do the forward pass. tensor_batch = data_batch[self.get_input_key(data_batch)] output_batch = self.network(tensor_batch) output_batch = output_batch if self.network.training else output_batch._asdict() # A callback proxy to modify the output after the forward pass. self._on_after_network_forward(data_batch, output_batch) return output_batch def training_step( self, data_batch: dict[str, torch.Tensor], iteration: int, ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: _input_key = self.get_input_key(data_batch) output_dict = self._network_forward(data_batch) input_images, recon_images = data_batch[_input_key], output_dict[RECON_KEY] # pass loss_mask to loss computation inputs = {INPUT_KEY: input_images, MASK_KEY: data_batch.get("loss_mask", torch.ones_like(input_images))} loss_dict, loss_value = self.loss(inputs, output_dict, iteration) return dict({PREDICTION: recon_images, **loss_dict}), loss_value @torch.no_grad() def validation_step( self, data_batch: dict[str, torch.Tensor], iteration: int, ema_model: bool = False, ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: _input_key = self.get_input_key(data_batch) output_dict = self._network_forward(data_batch) input_images, recon_images = data_batch[_input_key], output_dict[RECON_KEY] # pass loss_mask to loss computation inputs = {INPUT_KEY: input_images, MASK_KEY: data_batch.get("loss_mask", torch.ones_like(input_images))} loss_dict, loss_value = self.loss(inputs, output_dict, iteration) metric_dict = self.metric(input_images, output_dict, iteration) loss_dict.update(metric_dict) prediction_key = EMA_PREDICTION if ema_model else PREDICTION return dict({prediction_key: recon_images, **loss_dict}), loss_value @torch.inference_mode() def forward(self, data_batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: _input_key = self.get_input_key(data_batch) output_dict = self._network_forward(data_batch) return dict({PREDICTION: output_dict[RECON_KEY]})