# 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. from __future__ import annotations import copy from abc import ABC, abstractmethod from collections import defaultdict from contextlib import nullcontext from dataclasses import dataclass, fields from enum import Enum from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union from chronoedit._ext.imaginaire.utils.easy_io import easy_io import omegaconf import torch import torch.nn as nn from torch.distributed import ProcessGroup from chronoedit._ext.imaginaire.lazy_config import instantiate from chronoedit._ext.imaginaire.utils import log from chronoedit._src.utils.context_parallel import broadcast from chronoedit._ext.common.functional.batch_ops import batch_mul from chronoedit._ext.common.utils.count_params import count_params, disabled_train from chronoedit._src.datasets.data_sources.item_datasets_for_validation import get_itemdataset_option_local T = TypeVar("T", bound="BaseCondition") class DataType(str, Enum): IMAGE = "image" VIDEO = "video" MIX = "mix" def __str__(self) -> str: return self.value def broadcast_condition(condition: BaseCondition, process_group: Optional[ProcessGroup] = None) -> BaseCondition: """ Broadcast the condition from the minimum rank in the specified group(s). """ if condition.is_broadcasted: return condition kwargs = condition.to_dict(skip_underscore=False) for key, value in kwargs.items(): if value is not None: kwargs[key] = broadcast(value, process_group) kwargs["_is_broadcasted"] = True return type(condition)(**kwargs) @dataclass(frozen=True) class BaseCondition(ABC): """ Attributes: _is_broadcasted: Flag indicating if parallel broadcast splitting has been performed. This is an internal implementation detail. """ _is_broadcasted: bool = False def to_dict(self, skip_underscore: bool = True) -> Dict[str, Any]: """Converts the condition to a dictionary. Returns: Dictionary containing the condition's fields and values. """ # return {f.name: getattr(self, f.name) for f in fields(self) if not f.name.startswith("_")} return {f.name: getattr(self, f.name) for f in fields(self) if not (f.name.startswith("_") and skip_underscore)} @property def is_broadcasted(self) -> bool: return self._is_broadcasted def broadcast(self, process_group: torch.distributed.ProcessGroup) -> BaseCondition: """Broadcasts and splits the condition across the checkpoint parallelism group. For most condition, such asT2VCondition, we do not need split. Args: process_group: The process group for broadcast and split Returns: A new BaseCondition instance with the broadcasted and split condition. """ if self.is_broadcasted: return self return broadcast_condition(self, process_group) @dataclass(frozen=True) class T2VCondition(BaseCondition): crossattn_emb: Optional[torch.Tensor] = None data_type: DataType = DataType.VIDEO padding_mask: Optional[torch.Tensor] = None fps: Optional[torch.Tensor] = None def edit_data_type(self, data_type: DataType) -> T2VCondition: """Edit the data type of the condition. Args: data_type: The new data type. Returns: A new T2VCondition instance with the new data type. """ kwargs = self.to_dict(skip_underscore=False) kwargs["data_type"] = data_type return type(self)(**kwargs) @property def is_video(self) -> bool: return self.data_type == DataType.VIDEO @dataclass(frozen=True) class GR00TV1Img2VidCondition(T2VCondition): gt_first_frame: Optional[torch.Tensor] = None use_image_condition: bool = False condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None def edit_video_condition( self, x0_B_C_T_H_W, process_group: Optional[ProcessGroup] = None ) -> GR00TV1Img2VidCondition: """Edit the video condition to include the video mask information. Args: x0_B_C_T_H_W: The first frame of the video. Returns: A new GR00TV1Img2VidCondition instance with the video mask information. """ pg_size = 1 if process_group is None else process_group.size() kwargs = self.to_dict(skip_underscore=False) B, _, T, H, W = x0_B_C_T_H_W.shape condition_video_input_mask = torch.zeros((B, 1, T, H, W), dtype=x0_B_C_T_H_W.dtype, device=x0_B_C_T_H_W.device) if pg_size == 1 or process_group.rank() == 0: kwargs["gt_first_frame"] = x0_B_C_T_H_W[:, :, 0].detach() condition_video_input_mask[:, :, 0] += 1 kwargs["condition_video_input_mask_B_C_T_H_W"] = condition_video_input_mask return type(self)(**kwargs) class AbstractEmbModel(nn.Module): def __init__(self): super().__init__() self._is_trainable = None self._dropout_rate = None self._input_key = None self._return_dict = False @property def is_trainable(self) -> bool: return self._is_trainable @property def dropout_rate(self) -> Union[float, torch.Tensor]: return self._dropout_rate @property def input_key(self) -> str: return self._input_key @property def is_return_dict(self) -> bool: return self._return_dict @is_trainable.setter def is_trainable(self, value: bool): self._is_trainable = value @dropout_rate.setter def dropout_rate(self, value: Union[float, torch.Tensor]): self._dropout_rate = value @input_key.setter def input_key(self, value: str): self._input_key = value @is_return_dict.setter def is_return_dict(self, value: bool): self._return_dict = value @is_trainable.deleter def is_trainable(self): del self._is_trainable @dropout_rate.deleter def dropout_rate(self): del self._dropout_rate @input_key.deleter def input_key(self): del self._input_key @is_return_dict.deleter def is_return_dict(self): del self._return_dict def random_dropout_input( self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None, key: Optional[str] = None ) -> torch.Tensor: del key dropout_rate = dropout_rate if dropout_rate is not None else self.dropout_rate return batch_mul( torch.bernoulli((1.0 - dropout_rate) * torch.ones(in_tensor.shape[0])).type_as(in_tensor), in_tensor, ) def details(self) -> str: return "" def summary(self) -> str: input_key = self.input_key if self.input_key is not None else getattr(self, "input_keys", None) return ( f"{self.__class__.__name__} \n\tinput key: {input_key}" f"\n\tParam count: {count_params(self, False)} \n\tTrainable: {self.is_trainable}" f"\n\tDropout rate: {self.dropout_rate}" f"\n\t{self.details()}" ) class TextAttr(AbstractEmbModel): def __init__(self, input_key: List[str], dropout_rate: Optional[float] = 0.0): super().__init__() self._input_key = input_key self._dropout_rate = dropout_rate def forward(self, token: torch.Tensor): return {"crossattn_emb": token} def random_dropout_input( self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None, key: Optional[str] = None ) -> torch.Tensor: if key is not None and "mask" in key: return in_tensor return super().random_dropout_input(in_tensor, dropout_rate, key) def details(self) -> str: return "Output key: [crossattn_emb]" class TextAttrEmptyStringDrop(AbstractEmbModel): def __init__(self, input_key: List[str], dropout_rate: Optional[float] = 0.0): super().__init__() self._input_key = input_key self._dropout_rate = dropout_rate self.empty_prompt_data = None def forward(self, token: torch.Tensor): return {"crossattn_emb": token} def random_dropout_input( self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None, key: Optional[str] = None ) -> torch.Tensor: if key is not None and "mask" in key: return in_tensor del key if self.empty_prompt_data is None: self.empty_prompt_data = easy_io.load(get_itemdataset_option_local("empty_string_umt5").path) dropout_rate = dropout_rate if dropout_rate is not None else self.dropout_rate B = in_tensor.shape[0] # batch size # Create dropout mask: 1 -> keep in_tensor, 0 -> use empty_prompt_data keep_mask = torch.bernoulli((1.0 - dropout_rate) * torch.ones(B, device=in_tensor.device)).type_as(in_tensor) keep_mask = keep_mask.view(B, *[1] * (in_tensor.dim() - 1)) # broadcastable shape # Prepare empty_prompt_data with correct shape, dtype, and device empty_prompt = self.empty_prompt_data.to(dtype=in_tensor.dtype, device=in_tensor.device) # Repeat empty_prompt along batch dimension if needed if empty_prompt.shape[0] != B: if empty_prompt.shape[0] == 1: empty_prompt = empty_prompt.expand(B, *empty_prompt.shape[1:]) else: raise ValueError( f"empty_prompt_data batch size {empty_prompt.shape[0]} does not match in_tensor batch size {B}") # Mix using the dropout mask return keep_mask * in_tensor + (1.0 - keep_mask) * empty_prompt def details(self) -> str: return "Output key: [crossattn_emb]" class ReMapkey(AbstractEmbModel): def __init__( self, input_key: str, output_key: Optional[str] = None, dropout_rate: Optional[float] = 0.0, dtype: Optional[str] = None, ): super().__init__() self.output_key = output_key self.dtype = { None: None, "float": torch.float32, "bfloat16": torch.bfloat16, "half": torch.float16, "float16": torch.float16, "int": torch.int32, "long": torch.int64, }[ dtype ] self._input_key = input_key self._output_key = output_key self._dropout_rate = dropout_rate def forward(self, element: torch.Tensor) -> Dict[str, torch.Tensor]: key = self.output_key if self.output_key else self.input_key if isinstance(element, torch.Tensor): element = element.to(dtype=self.dtype) return {key: element} def details(self) -> str: key = self.output_key if self.output_key else self.input_key return f"Output key: {key} \n\tDtype: {self.dtype}" class BooleanFlag(AbstractEmbModel): def __init__(self, input_key: str, output_key: Optional[str] = None, dropout_rate: Optional[float] = 0.0): super().__init__() self._input_key = input_key self._dropout_rate = dropout_rate self.output_key = output_key def forward(self, *args, **kwargs) -> Dict[str, torch.Tensor]: del args, kwargs key = self.output_key if self.output_key else self.input_key return {key: self.flag} def random_dropout_input( self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None, key: Optional[str] = None ) -> torch.Tensor: del key dropout_rate = dropout_rate if dropout_rate is not None else self.dropout_rate self.flag = torch.bernoulli((1.0 - dropout_rate) * torch.ones(1)).bool().to(device=in_tensor.device) return in_tensor def details(self) -> str: key = self.output_key if self.output_key else self.input_key return f"Output key: {key} \n\t This is a boolean flag" class GeneralConditioner(nn.Module, ABC): """ An abstract module designed to handle various embedding models with conditional and unconditional configurations. This abstract base class initializes and manages a collection of embedders that can dynamically adjust their dropout rates based on conditioning. Attributes: KEY2DIM (dict): A mapping from output keys to dimensions used for concatenation. embedders (nn.ModuleDict): A dictionary containing all embedded models initialized and configured based on the provided configurations. Parameters: emb_models (Union[List, Any]): A dictionary where keys are embedder names and values are configurations for initializing the embedders. Example: See Edify4ConditionerConfig """ KEY2DIM = {"crossattn_emb": 1} def __init__(self, **emb_models: Union[List, Any]): super().__init__() self.embedders = nn.ModuleDict() for n, (emb_name, emb_config) in enumerate(emb_models.items()): embedder = instantiate(emb_config) # assert isinstance( # embedder, AbstractEmbModel # ), f"embedder model {embedder.__class__.__name__} has to inherit from AbstractEmbModel" embedder.is_trainable = getattr(emb_config, "is_trainable", True) embedder.dropout_rate = getattr(emb_config, "dropout_rate", 0.0) if not embedder.is_trainable: embedder.train = disabled_train for param in embedder.parameters(): param.requires_grad = False embedder.eval() log.info(f"Initialized embedder #{n}-{emb_name}: \n {embedder.summary()}") self.embedders[emb_name] = embedder @abstractmethod def forward( self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None, ) -> Any: """Should be implemented in subclasses to handle conditon datatype""" raise NotImplementedError def _forward( self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None, ) -> Dict: """ Processes the input batch through all configured embedders, applying conditional dropout rates if specified. Output tensors for each key are concatenated along the dimensions specified in KEY2DIM. Parameters: batch (Dict): The input data batch to process. override_dropout_rate (Optional[Dict[str, float]]): Optional dictionary to override default dropout rates per embedder key. Returns: Dict: A dictionary of output tensors concatenated by specified dimensions. Note: In case the network code is sensitive to the order of concatenation, you can either control the order via \ config file or make sure the embedders return a unique key for each output. """ output = defaultdict(list) if override_dropout_rate is None: override_dropout_rate = {} # make sure emb_name in override_dropout_rate is valid for emb_name in override_dropout_rate.keys(): assert emb_name in self.embedders, f"invalid name found {emb_name}" for emb_name, embedder in self.embedders.items(): embedding_context = nullcontext if embedder.is_trainable else torch.no_grad with embedding_context(): if isinstance(embedder.input_key, str): emb_out = embedder( embedder.random_dropout_input( batch[embedder.input_key], override_dropout_rate.get(emb_name, None) ) ) elif isinstance(embedder.input_key, (list, omegaconf.listconfig.ListConfig)): emb_out = embedder( *[ embedder.random_dropout_input(batch.get(k), override_dropout_rate.get(emb_name, None), k) for k in embedder.input_key ] ) else: raise KeyError( f"Embedder '{embedder.__class__.__name__}' requires an 'input_key' attribute to be defined as either a string or list of strings" ) for k, v in emb_out.items(): output[k].append(v) # Concatenate the outputs return {k: torch.cat(v, dim=self.KEY2DIM.get(k, -1)) for k, v in output.items()} def get_condition_uncondition( self, data_batch: Dict, ) -> Tuple[Any, Any]: """ Processes the provided data batch to generate two sets of outputs: conditioned and unconditioned. This method manipulates the dropout rates of embedders to simulate two scenarios — one where all conditions are applied (conditioned), and one where they are removed or reduced to the minimum (unconditioned). This method first sets the dropout rates to zero for the conditioned scenario to fully apply the embedders' effects. For the unconditioned scenario, it sets the dropout rates to 1 (or to 0 if the initial unconditional dropout rate is insignificant) to minimize the embedders' influences, simulating an unconditioned generation. Parameters: data_batch (Dict): The input data batch that contains all necessary information for embedding processing. The data is expected to match the required format and keys expected by the embedders. Returns: Tuple[Any, Any]: A tuple containing two condition: - The first one contains the outputs with all embedders fully applied (conditioned outputs). - The second one contains the outputs with embedders minimized or not applied (unconditioned outputs). """ cond_dropout_rates, dropout_rates = {}, {} for emb_name, embedder in self.embedders.items(): cond_dropout_rates[emb_name] = 0.0 dropout_rates[emb_name] = 1.0 if embedder.dropout_rate > 1e-4 else 0.0 condition: Any = self(data_batch, override_dropout_rate=cond_dropout_rates) un_condition: Any = self(data_batch, override_dropout_rate=dropout_rates) return condition, un_condition def get_condition_with_negative_prompt( self, data_batch: Dict, ) -> Tuple[Any, Any]: """ Similar functionality as get_condition_uncondition But use negative prompts for unconditon """ cond_dropout_rates, uncond_dropout_rates = {}, {} for emb_name, embedder in self.embedders.items(): cond_dropout_rates[emb_name] = 0.0 if isinstance(embedder, TextAttr) or isinstance(embedder, TextAttrEmptyStringDrop): uncond_dropout_rates[emb_name] = 0.0 else: uncond_dropout_rates[emb_name] = 1.0 if embedder.dropout_rate > 1e-4 else 0.0 data_batch_neg_prompt = copy.deepcopy(data_batch) if "neg_t5_text_embeddings" in data_batch_neg_prompt: if isinstance(data_batch_neg_prompt["neg_t5_text_embeddings"], torch.Tensor): data_batch_neg_prompt["t5_text_embeddings"] = data_batch_neg_prompt["neg_t5_text_embeddings"] condition: Any = self(data_batch, override_dropout_rate=cond_dropout_rates) un_condition: Any = self(data_batch_neg_prompt, override_dropout_rate=uncond_dropout_rates) return condition, un_condition class VideoConditioner(GeneralConditioner): def forward( self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None, ) -> T2VCondition: output = super()._forward(batch, override_dropout_rate) return T2VCondition(**output) class GR00TV1Img2VidConditioner(GeneralConditioner): def forward( self, batch: Dict, override_dropout_rate: Optional[Dict[str, float]] = None, ) -> GR00TV1Img2VidCondition: output = super()._forward(batch, override_dropout_rate) return GR00TV1Img2VidCondition(**output)