import math from einops import rearrange import torch from torch import nn import numpy as np import torch.distributed from .base import ContinuousTransformer class ClusterModel(nn.Module): """Codebook xluster model based on euclidean distance.""" def __init__(self, cluster_centers: torch.Tensor): super().__init__() assert len(cluster_centers.shape) == 3 self._n_codebooks = cluster_centers.shape[0] self._n_clusters = cluster_centers.shape[1] self._dim = cluster_centers.shape[2] self.register_buffer("embed", cluster_centers) @property def n_codebooks(self): return self._n_codebooks @property def n_clusters(self): return self._n_clusters @property def dim(self): return self._dim @torch.no_grad() def quantize(self, x, codebook_idx): embed = self.embed[codebook_idx].t() x = x.float() dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True)) embed_ind = dist.max(dim=-1).indices return embed_ind @torch.no_grad() def encode(self, x, n_codebooks=None): assert len(x.shape) == 3 # batch, time, dim if n_codebooks is None: n_codebooks = self.n_codebooks shape = x.shape # pre-process x = rearrange(x, "... d -> (...) d") # quantize resid = x.clone() quantized_vector_list = [] for n_codebook in range(n_codebooks): embed_ind = self.quantize(resid, n_codebook) resid -= torch.nn.functional.embedding(embed_ind, self.embed[n_codebook]) # post-process embed_ind = embed_ind.view(*shape[:-1]) quantized_vector_list.append(embed_ind) quantized_vectors = torch.stack(quantized_vector_list) quantized_vectors = torch.swapaxes(quantized_vectors, 0, 1) # batch, codebooks, clusters return quantized_vectors @torch.no_grad() def decode(self, embed_ind): assert len(embed_ind.shape) == 3 # batch, (codebooks), clusters embeddings = None for n_codebook in range(embed_ind.shape[1]): codebook_embs = torch.nn.functional.embedding( embed_ind[:, n_codebook, :], self.embed[n_codebook] ) if embeddings is None: embeddings = codebook_embs else: embeddings += codebook_embs return embeddings class Discriminator(torch.nn.Module): def __init__(self, hidden_dim): super().__init__() self.discriminator = torch.nn.Sequential( torch.nn.Conv1d(hidden_dim, 128, kernel_size=4, stride=2), torch.nn.SiLU(), torch.nn.GroupNorm(32, 128), torch.nn.Conv1d(128, 128, kernel_size=4, stride=2), torch.nn.SiLU(), torch.nn.GroupNorm(32, 128), torch.nn.Conv1d(128, 128, kernel_size=4, stride=2), torch.nn.SiLU(), torch.nn.GroupNorm(32, 128), torch.nn.Conv1d(128, 128, kernel_size=4, stride=2), torch.nn.SiLU(), ) self.linear = torch.nn.Linear(128, 1) def forward(self, x): x = self.discriminator(x.permute(0, 2, 1)) x = x.permute(0, 2, 1) x = self.linear(x) return x class FourierFeatures(nn.Module): def __init__(self, in_features, out_features): super().__init__() assert out_features % 2 == 0 self.weight = nn.Parameter(torch.zeros(out_features // 2, in_features)) def forward(self, input): f = 2 * math.pi * input @ self.weight.T return torch.cat([f.cos(), f.sin()], dim=-1) class SemanticConditioner(nn.Module): def __init__( self, vocab_size: int = 4001, input_dim: int = 768, output_dim: int = 1536, use_rvq: bool = False, n_codebooks: int = 1, ): super().__init__() self.vocab_size = vocab_size self.use_rvq = use_rvq self.n_codebooks = n_codebooks if self.use_rvq: # RVQ will be initialized later in _initialize_embedding cluster_centers = torch.randn(n_codebooks, vocab_size - 1, input_dim) * 0.02 self.rvq = ClusterModel(cluster_centers) self.rvq_pad_embedding = nn.Parameter(torch.randn(1, cluster_centers.shape[2]) * 0.02) else: self.embedding = torch.nn.Embedding(self.vocab_size, input_dim) self.proj_out = nn.Linear(input_dim, output_dim) def _initialize_embedding(self, cluster_centers): """Initialize embedding layer with centroids (unified logic for both modes)""" if self.use_rvq: # RVQ mode: keep original cluster centers intact for proper dequantization # Use original cluster centers without modification # This ensures dequantization works correctly with the original clusters # update the cluster centers in rvq self.rvq.embed.data.copy_(torch.from_numpy(cluster_centers).float()) self.rvq.embed.requires_grad = False else: # Non-RVQ mode: use only first codebook (original behavior) first_codebook = cluster_centers[0] # Shape: [n_clusters, dim] # Common logic: create embedding with first codebook + pad token centroids_with_pad = np.concatenate( [first_codebook, np.random.randn(1, first_codebook.shape[1])], axis=0 ) assert self.embedding.weight.shape == centroids_with_pad.shape self.embedding.weight.data.copy_(torch.tensor(centroids_with_pad, dtype=torch.float32)) def forward(self, x, noise=None): if self.use_rvq: # x should be multi-level quantized codes # x shape: [batch, n_codebooks, sequence] or [batch, sequence, n_codebooks] # Assume input is [batch, sequence, n_codebooks] for consistency with standard mode if len(x.shape) == 3 and x.shape[-1] > 1: # x is [batch, sequence, n_codebooks] - transpose to [batch, n_codebooks, sequence] x = x.permute(0, 2, 1) # Detect pad tokens before dequantization # Pad tokens are marked as vocab_size - 1 (4000 for vocab_size=4001) pad_mask = (x == self.vocab_size - 1).any(dim=1) # [batch, sequence] # Create a copy of x for dequantization, replacing pad tokens with valid cluster indices x_for_decode = x.clone() # Replace pad tokens (vocab_size - 1) with cluster index 0 for dequantization # This ensures we can still decode the non-pad tokens properly x_for_decode = torch.where(x == self.vocab_size - 1, torch.zeros_like(x), x) # Decode quantized codes back to continuous vectors continuous_vectors = self.rvq.decode(x_for_decode) # Shape: [batch, sequence, features] # Replace padded positions with learnable pad embedding # Use type_as() for torch.compile compatibility - ensures dtype/device match # Always expand (even if not used) to avoid dynamic control flow issues with torch.compile rvq_pad_embedding = self.rvq_pad_embedding.expand( continuous_vectors.shape[0], continuous_vectors.shape[1], -1 ).type_as(continuous_vectors) # Replace padded positions with learnable RVQ pad embedding # torch.where handles the case where pad_mask is all False efficiently continuous_vectors_with_pad = torch.where( pad_mask.unsqueeze(-1), rvq_pad_embedding, continuous_vectors ) x = continuous_vectors_with_pad else: # Standard mode: x is already token indices x = self.embedding(x) if noise is not None: x = x + noise x = self.proj_out(x) return x # [batch, sequence, channels] class TextConditioner(nn.Module): def __init__( self, vocab_size: int = 60001, output_dim: int = 1536, ): super().__init__() self.embedding = torch.nn.Embedding(vocab_size, output_dim) def forward(self, x): x = self.embedding(x) return x # [batch, sequence, channels] class LatentContextConditioner(nn.Module): def __init__( self, io_channels: int = 128, output_dim: int = 1536, ): super().__init__() self.proj_out = nn.Linear(io_channels, output_dim) def forward(self, x): x = self.proj_out(x) return x # [batch, sequence, channels] class ActNorm1d(nn.Module): """ Per-channel affine: y = scale * x + bias Data-dependent init on first forward. Optional DDP sync of mean/std. Expects x shape [B, C, T]. """ def __init__(self, num_channels: int, eps: float = 1e-6, sync_init: bool = True): super().__init__() self.bias = nn.Parameter(torch.zeros(1, num_channels, 1)) self.scale = nn.Parameter(torch.ones(1, num_channels, 1)) self.eps = eps self.initialized = False self.sync_init = sync_init @torch.no_grad() def _ddp_sync_stats(self, mean: torch.Tensor, var: torch.Tensor): # mean/var are [1, C, 1]. All-reduce to global stats if DDP is initialized. if torch.distributed.is_available() and torch.distributed.is_initialized(): world = torch.distributed.get_world_size() torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM) torch.distributed.all_reduce(var, op=torch.distributed.ReduceOp.SUM) mean /= world var /= world return mean, var @torch.no_grad() def _initialize(self, x: torch.Tensor): # x: [B, C, T] (any dtype) -> compute stats in fp32 xf = x.detach().to(torch.float32) mean = xf.mean(dim=(0, 2), keepdim=True) # [1, C, 1] var = xf.var(dim=(0, 2), keepdim=True, unbiased=False) # [1, C, 1] if self.sync_init: mean, var = self._ddp_sync_stats(mean, var) std = torch.sqrt(var + self.eps) self.bias.data.copy_(-mean) # so y = (x - mean) / std at start self.scale.data.copy_(1.0 / std) self.initialized = True def forward(self, x: torch.Tensor): if not self.initialized: self._initialize(x) # keep math in input dtype; params are learned thereafter return x * self.scale.to(x.dtype) + self.bias.to(x.dtype) class DiffusionTransformer(nn.Module): def __init__( self, io_hz=25, io_channels=128, embed_dim=1536, depth=24, n_heads=24, qk_norm=True, block_size=750, cond_semantic_n_vocab=4001, cond_semantic_len=750, cond_text_n_vocab=60001, cond_text_len=2560, ctx_len=0, infill_ctx_len=0, stem_ctx_len=0, use_fine_guidance=False, shared_ctx=False, use_actnorm=False, actnorm_sync_init=True, use_mmdit=False, cond_text_embed_len=None, init_variant="baseline", use_rvq=False, n_codebooks=1, ): super().__init__() self.io_hz = io_hz self.io_channels = io_channels self.embed_dim = embed_dim self.depth = depth self.n_heads = n_heads self.qk_norm = qk_norm self.block_size = block_size self.cond_semantic_n_vocab = cond_semantic_n_vocab self.cond_semantic_len = cond_semantic_len self.cond_text_n_vocab = cond_text_n_vocab self.cond_text_len = cond_text_len self.ctx_len = ctx_len self.infill_ctx_len = infill_ctx_len if infill_ctx_len is not None else 0 self.use_fine_guidance = use_fine_guidance self.stem_ctx_len = stem_ctx_len if stem_ctx_len is not None else 0 self.shared_ctx = shared_ctx self.use_actnorm = use_actnorm self.actnorm_sync_init = actnorm_sync_init self.use_mmdit = use_mmdit self.cond_text_embed_len = cond_text_embed_len or cond_text_len self.init_variant = init_variant self.use_rvq = use_rvq self.n_codebooks = n_codebooks self.semantic_conditioner = SemanticConditioner( vocab_size=cond_semantic_n_vocab, output_dim=embed_dim, use_rvq=use_rvq, n_codebooks=n_codebooks, ) self.text_conditioner = TextConditioner(vocab_size=cond_text_n_vocab, output_dim=embed_dim) if self.ctx_len > 0 or self.shared_ctx: self.ctx_conditioner = LatentContextConditioner( io_channels=io_channels, output_dim=embed_dim ) self.vae_pad_embed = nn.Parameter(torch.randn(io_channels) * 0.02) if self.infill_ctx_len > 0 or self.shared_ctx: self.infill_ctx_conditioner = LatentContextConditioner( io_channels=io_channels, output_dim=embed_dim ) self.vae_infill_pad_embed = nn.Parameter(torch.randn(io_channels) * 0.02) if self.shared_ctx: self.vox_conditioner = LatentContextConditioner( io_channels=io_channels, output_dim=embed_dim ) self.vae_vox_pad_embed = nn.Parameter(torch.randn(io_channels) * 0.02) if self.stem_ctx_len > 0: self.stem_ctx_conditioner = LatentContextConditioner( io_channels=io_channels, output_dim=embed_dim ) self.vae_stem_pad_embed = nn.Parameter(torch.randn(io_channels) * 0.02) if self.shared_ctx: self.vae_default_embed = nn.Parameter(torch.randn(embed_dim) * 0.02) # Timestep embeddings timestep_features_dim = 256 self.timestep_features = FourierFeatures(1, timestep_features_dim) self.to_timestep_embed = nn.Sequential( nn.Linear(timestep_features_dim, embed_dim, bias=True), nn.SiLU(), nn.Linear(embed_dim, embed_dim, bias=True), ) # Transformer assert embed_dim % n_heads == 0 dim_heads = embed_dim // n_heads # Set up attention block sizes # Standard structure: conditioning + timestep + audio attn_block_sizes = [ self.cond_text_len + self.cond_semantic_len + self.ctx_len + self.infill_ctx_len + self.stem_ctx_len, 1 + self.block_size + (self.cond_text_embed_len if self.use_mmdit else 0), ] self.transformer = ContinuousTransformer( embed_dim, depth, dim_heads=dim_heads, dim_in=io_channels, dim_out=io_channels, qk_norm=qk_norm, attn_block_sizes=attn_block_sizes, ) self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) # MMDiT multimodal processing (optional) if self.use_mmdit: # Projection layer to convert audio tokens to embed_dim for multimodal processing self.audio_to_embed = nn.Linear(io_channels, embed_dim) # Projection layer to convert back from embed_dim to io_channels self.embed_to_audio = nn.Linear(embed_dim, io_channels) # ---- ActNorm modules (only if enabled) ---- if self.use_actnorm: self.pre_act = ActNorm1d(io_channels, sync_init=actnorm_sync_init) self.post_act = ActNorm1d(io_channels, sync_init=actnorm_sync_init) @property def device(self): return next(self.parameters()).device @property def dtype(self): return next(self.parameters()).dtype def disable_compile(self): """ Explicitly disable torch.compile for this model and all submodules. Call this on the critic model to ensure it doesn't use compiled code. """ # Mark all modules to not be compiled for module in self.modules(): if hasattr(module, "_is_compiled"): module._is_compiled = False # Also clear any dynamo state if hasattr(module, "_dynamo_cache"): module._dynamo_cache = None return self def _initialize(self, semantic_clusters_filepath=None): """Initialize model weights and load cluster centers if provided""" # Load cluster centers from data config if available cluster_centers = None if semantic_clusters_filepath is not None: cluster_centers = np.load(semantic_clusters_filepath) print( f"Loaded cluster centers: shape={cluster_centers.shape}, dtype={cluster_centers.dtype}" ) print(f"Cluster centers range: min={cluster_centers.min()}, max={cluster_centers.max()}") print(f"Number of clusters per codebook: {cluster_centers.shape[1]}") print(f"Expected vocab size: {self.cond_semantic_n_vocab}") # Initialize semantic conditioner with cluster centers self.semantic_conditioner._initialize_embedding(cluster_centers) # Note: RVQ is initialized inside _initialize_embedding for RVQ mode # Weight initialization if self.init_variant == "baseline": # ---- your existing init (unchanged) ---- nn.init.zeros_(self.preprocess_conv.weight) nn.init.zeros_(self.postprocess_conv.weight) nn.init.normal_(self.timestep_features.weight) for layer in self.transformer.layers: nn.init.zeros_(layer.ff.ff[1].weight) nn.init.zeros_(layer.ff.ff[1].bias) nn.init.zeros_(layer.self_attn.to_out.weight) with torch.no_grad(): for p in self.parameters(): p *= 0.5 return elif self.init_variant == "stability_v1": # ---- stability_v1 (small delta, high impact) ---- # 1) Identity 1x1 convs def _init_identity_1x1_(conv): assert conv.kernel_size == (1,) assert conv.in_channels == conv.out_channels with torch.no_grad(): conv.weight.zero_() eye = torch.eye(conv.in_channels, device=conv.weight.device, dtype=conv.weight.dtype) conv.weight.copy_(eye.view(conv.in_channels, conv.in_channels, 1)) if conv.bias is not None: conv.bias.zero_() _init_identity_1x1_(self.preprocess_conv) _init_identity_1x1_(self.postprocess_conv) # 2) Keep your existing per-block zero on last MLP + attn out (already good) for layer in self.transformer.layers: nn.init.zeros_(layer.ff.ff[1].weight) nn.init.zeros_(layer.ff.ff[1].bias) nn.init.zeros_(layer.self_attn.to_out.weight) if getattr(layer.self_attn.to_out, "bias", None) is not None: nn.init.zeros_(layer.self_attn.to_out.bias) # 3) Depth-aware residual scaling (single line per projection) res_scale = (1.0 / (2.0 * self.depth)) ** 0.5 with torch.no_grad(): for layer in self.transformer.layers: layer.ff.ff[1].weight.mul_(res_scale) if layer.ff.ff[1].bias is not None: layer.ff.ff[1].bias.mul_(res_scale) layer.self_attn.to_out.weight.mul_(res_scale) if getattr(layer.self_attn.to_out, "bias", None) is not None: layer.self_attn.to_out.bias.mul_(res_scale) # 4) Conditioner heads no-op + pad zeros nn.init.zeros_(self.semantic_conditioner.proj_out.weight) nn.init.zeros_(self.semantic_conditioner.proj_out.bias) # nn.init.zeros_(self.text_conditioner.proj_out.weight) # nn.init.zeros_(self.text_conditioner.proj_out.bias) if hasattr(self, "ctx_conditioner"): nn.init.zeros_(self.ctx_conditioner.proj_out.weight) nn.init.zeros_(self.ctx_conditioner.proj_out.bias) if hasattr(self, "infill_ctx_conditioner"): nn.init.zeros_(self.infill_ctx_conditioner.proj_out.weight) nn.init.zeros_(self.infill_ctx_conditioner.proj_out.bias) if hasattr(self, "stem_ctx_conditioner"): nn.init.zeros_(self.stem_ctx_conditioner.proj_out.weight) nn.init.zeros_(self.stem_ctx_conditioner.proj_out.bias) # pads → true "no signal" if hasattr(self, "vae_pad_embed"): nn.init.zeros_(self.vae_pad_embed) if hasattr(self, "vae_infill_pad_embed"): nn.init.zeros_(self.vae_infill_pad_embed) if hasattr(self, "vae_vox_pad_embed"): nn.init.zeros_(self.vae_vox_pad_embed) if hasattr(self, "vae_stem_pad_embed"): nn.init.zeros_(self.vae_stem_pad_embed) # optional default embed stays small-N(0,0.02) so it's distinguishable if used if hasattr(self, "vae_default_embed"): nn.init.normal_(self.vae_default_embed, 0.0, 0.02) # 5) Neutral MMDiT adapters (if enabled) if self.use_mmdit: nn.init.xavier_uniform_(self.audio_to_embed.weight) if self.audio_to_embed.bias is not None: nn.init.zeros_(self.audio_to_embed.bias) nn.init.zeros_(self.embed_to_audio.weight) if self.embed_to_audio.bias is not None: nn.init.zeros_(self.embed_to_audio.bias) # Keep your original timestep_features, or reduce its std later if needed nn.init.normal_(self.timestep_features.weight, mean=0.0, std=1.0) else: raise ValueError(f"Invalid init variant: {self.init_variant}") def _forward( self, x, t, cross_attn_cond=None, text_embeds=None, beta=None, ): assert cross_attn_cond is not None prepend_inputs = self._prepare_prepend_inputs(t, cross_attn_cond, beta) return self._apply_transformer(x, prepend_inputs, text_embeds=text_embeds) def _prepare_prepend_inputs(self, t, cross_attn_cond, beta): timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None])) prepend_inputs = timestep_embed.unsqueeze(1) if cross_attn_cond is not None: prepend_inputs = torch.cat((cross_attn_cond, prepend_inputs), dim=1) if beta is not None: assert self.use_fine_guidance timestep_embed += self.to_beta_embed(self.beta_features(beta[:, None])) return prepend_inputs def _apply_transformer(self, x, prepend_inputs, text_embeds=None): """Apply transformer for normal forward pass (torch.compile friendly).""" if self.use_actnorm: x = self.pre_act(x) x = self.preprocess_conv(x) + x x = rearrange(x, "b c t -> b t c") if self.use_mmdit and text_embeds is not None: audio_embeds = self.audio_to_embed(x) multimodal_embeds = torch.cat([text_embeds, audio_embeds], dim=1) output = self.transformer(multimodal_embeds, prepend_embeds=prepend_inputs) audio_output = output[:, text_embeds.shape[1] :, :] x = self.embed_to_audio(audio_output) else: x = self.transformer(x, prepend_embeds=prepend_inputs) x = rearrange(x, "b t c -> b c t") x = self.postprocess_conv(x) + x if self.use_actnorm: x = self.post_act(x) return x def _apply_transformer_intermediate(self, x, prepend_inputs, text_embeds, intermediate_layer_idx): """Apply transformer and extract intermediate features (for distillation only).""" if self.use_actnorm: x = self.pre_act(x) x = self.preprocess_conv(x) + x x = rearrange(x, "b c t -> b t c") if self.use_mmdit and text_embeds is not None: audio_embeds = self.audio_to_embed(x) multimodal_embeds = torch.cat([text_embeds, audio_embeds], dim=1) return self.transformer.forward_intermediate( multimodal_embeds, prepend_embeds=prepend_inputs, intermediate_layer_idx=intermediate_layer_idx, ) else: return self.transformer.forward_intermediate( x, prepend_embeds=prepend_inputs, intermediate_layer_idx=intermediate_layer_idx, ) def forward( self, x, t, text_codes=None, semantic_codes=None, ctx_vae=None, ctx_mask=None, infill_ctx_vae=None, infill_ctx_mask=None, semantic_noise=None, vox_vae=None, vox_mask=None, stem_ctx_vae=None, stem_ctx_mask=None, beta=None, intermediate_layer_idx=None, ): assert text_codes is not None # assert semantic_codes is not None assert torch.all(text_codes < self.cond_text_n_vocab), text_codes.max() # special mask check for shared_ctx if self.shared_ctx: assert ctx_mask is not None assert infill_ctx_mask is not None # assert vox_mask is not None assert ctx_mask.shape == infill_ctx_mask.shape # == vox_mask.shape # there are three valid cases: # 1. ctx_mask is all ones and infill_ctx_mask is all zeros # 2. ctx_mask is all zeros and infill_ctx_mask is all ones # 3. vox_mask is all ones and both ctx and infill ctx are all zeros # 4. all masks are all zeros (first chunk condition) # Check each batch item without a loop ctx_all_ones = ctx_mask.all(dim=1) # [batch_size] ctx_all_zeros = (~ctx_mask).all(dim=1) # [batch_size] infill_all_ones = infill_ctx_mask.all(dim=1) # [batch_size] infill_all_zeros = (~infill_ctx_mask).all(dim=1) # [batch_size] if vox_mask is not None: assert vox_mask.shape == ctx_mask.shape == infill_ctx_mask.shape vox_all_ones = vox_mask.all(dim=1) # [batch_size] vox_all_zeros = (~vox_mask).all(dim=1) # [batch_size] else: vox_all_ones = False vox_all_zeros = True case1 = ctx_all_ones & infill_all_zeros & vox_all_zeros case2 = ctx_all_zeros & infill_all_ones & vox_all_zeros case3 = ctx_all_zeros & infill_all_zeros & vox_all_ones case4 = ctx_all_zeros & infill_all_zeros & vox_all_zeros valid_masks = case1 | case2 | case3 | case4 # combine all valid cases if not torch.all(valid_masks): print("Invalid mask configuration for shared context") print("case1:", case1) print("case2:", case2) print("case3:", case3) print("case4:", case4) assert torch.all(valid_masks), "Invalid mask configuration for shared context" # assert torch.all(semantic_codes < self.cond_semantic_n_vocab) cross_attn_cond, text_embeds = self._build_conditioning( x.dtype, text_codes, semantic_codes, ctx_vae, ctx_mask, infill_ctx_vae, infill_ctx_mask, vox_vae, vox_mask, stem_ctx_vae, stem_ctx_mask, ) if intermediate_layer_idx is not None: prepend_inputs = self._prepare_prepend_inputs(t, cross_attn_cond, beta) return self._apply_transformer_intermediate( x, prepend_inputs, text_embeds=text_embeds, intermediate_layer_idx=intermediate_layer_idx, ) return self._forward( x, t, cross_attn_cond=cross_attn_cond, text_embeds=text_embeds, beta=beta, ) def _build_conditioning( self, dtype, text_codes, semantic_codes, ctx_vae, ctx_mask, infill_ctx_vae, infill_ctx_mask, vox_vae, vox_mask, stem_ctx_vae, stem_ctx_mask, ): text_embeds = self.text_conditioner(text_codes) semantic_embeds = self.semantic_conditioner(semantic_codes) cross_attn_cond = torch.concat([text_embeds, semantic_embeds], dim=1) def process_context(vae, mask, conditioner, pad_embed): expanded_vae_pad = pad_embed.view(1, -1, 1).expand(vae.shape[0], -1, vae.shape[2]) inverted_mask = (~mask).float().unsqueeze(1) vae_input = vae * mask.float().unsqueeze(1) + expanded_vae_pad * inverted_mask vae_input = vae_input.permute(0, 2, 1).to(dtype) return conditioner(vae_input) def process_shared_context(vae, mask, conditioner, pad_embed): expanded_vae_pad = pad_embed.view(1, -1, 1).expand(vae.shape[0], -1, vae.shape[2]) vae_input = vae + expanded_vae_pad vae_input = vae_input.permute(0, 2, 1).to(dtype) embeds = conditioner(vae_input) return embeds * mask.type_as(embeds).unsqueeze(2) if ctx_vae is not None and infill_ctx_vae is not None and self.shared_ctx: ctx_embeds = process_shared_context( ctx_vae, ctx_mask, self.ctx_conditioner, self.vae_pad_embed, ) infill_ctx_embeds = process_shared_context( infill_ctx_vae, infill_ctx_mask, self.infill_ctx_conditioner, self.vae_infill_pad_embed, ) if vox_mask is not None: vox_embeds = process_shared_context( vox_vae, vox_mask, self.vox_conditioner, self.vae_vox_pad_embed, ) else: vox_embeds = torch.zeros_like(ctx_embeds) vox_mask = torch.zeros_like(ctx_mask) expanded_default_embed = self.vae_default_embed.view(1, 1, -1).repeat( ctx_embeds.shape[0], ctx_embeds.shape[1], 1 ) default_embed_mask = ( ((ctx_mask == 0) & (infill_ctx_mask == 0) & (vox_mask == 0)) .type_as(expanded_default_embed) .unsqueeze(2) ) default_embeds = expanded_default_embed * default_embed_mask shared_ctx_embeds = ctx_embeds + infill_ctx_embeds + vox_embeds + default_embeds cross_attn_cond = torch.concat([cross_attn_cond, shared_ctx_embeds], dim=1) else: if ctx_vae is not None: ctx_embeds = process_context( ctx_vae, ctx_mask, self.ctx_conditioner, self.vae_pad_embed, ) cross_attn_cond = torch.concat([cross_attn_cond, ctx_embeds], dim=1) if infill_ctx_vae is not None: infill_ctx_embeds = process_context( infill_ctx_vae, infill_ctx_mask, self.infill_ctx_conditioner, self.vae_infill_pad_embed, ) cross_attn_cond = torch.concat([cross_attn_cond, infill_ctx_embeds], dim=1) if stem_ctx_vae is not None: stem_ctx_embeds = process_context( stem_ctx_vae, stem_ctx_mask, self.stem_ctx_conditioner, self.vae_stem_pad_embed ) cross_attn_cond = torch.concat([cross_attn_cond, stem_ctx_embeds], dim=1) return cross_attn_cond, text_embeds def extract_intermediate_features( self, x, t, text_codes=None, semantic_codes=None, ctx_vae=None, ctx_mask=None, infill_ctx_vae=None, infill_ctx_mask=None, semantic_noise=None, vox_vae=None, vox_mask=None, stem_ctx_vae=None, stem_ctx_mask=None, beta=None, intermediate_layer_idx=None, ): """ Extract intermediate features for distillation. This method is NOT compiled and should only be used during distillation training. NOTE: This method should only be called on models that were NOT compiled. The critic_model should be loaded with compile=False. """ if intermediate_layer_idx is None: raise ValueError("intermediate_layer_idx must be provided") return self.forward( x, t, text_codes=text_codes, semantic_codes=semantic_codes, ctx_vae=ctx_vae, ctx_mask=ctx_mask, infill_ctx_vae=infill_ctx_vae, infill_ctx_mask=infill_ctx_mask, vox_vae=vox_vae, vox_mask=vox_mask, stem_ctx_vae=stem_ctx_vae, stem_ctx_mask=stem_ctx_mask, beta=beta, intermediate_layer_idx=intermediate_layer_idx, ) def forward_inference( self, x, t, cross_attn_cond=None, no_text_cross_attn_cond=None, no_tag_cross_attn_cond=None, no_sem_cross_attn_cond=None, no_ctx_cross_attn_cond=None, no_infill_ctx_cross_attn_cond=None, no_stem_ctx_cross_attn_cond=None, text_cfg_scale=1.0, tag_cfg_scale=1.0, sem_cfg_scale=1.0, ctx_cfg_scale=1.0, infill_ctx_cfg_scale=1.0, stem_ctx_cfg_scale=1.0, ): assert cross_attn_cond is not None assert no_text_cross_attn_cond is not None has_text_cfg = text_cfg_scale != 1.0 and no_text_cross_attn_cond is not None has_tag_cfg = tag_cfg_scale != 1.0 and no_tag_cross_attn_cond is not None has_sem_cfg = sem_cfg_scale != 1.0 and no_sem_cross_attn_cond is not None has_ctx_cfg = ctx_cfg_scale != 1.0 and no_ctx_cross_attn_cond is not None has_infill_ctx_cfg = infill_ctx_cfg_scale != 1.0 and no_infill_ctx_cross_attn_cond is not None has_stem_ctx_cfg = stem_ctx_cfg_scale != 1.0 and no_stem_ctx_cross_attn_cond is not None n_rep = ( 1 + int(has_text_cfg) + int(has_tag_cfg) + int(has_sem_cfg) + int(has_ctx_cfg) + int(has_infill_ctx_cfg) + int(has_stem_ctx_cfg) ) # Classifier-free guidance if self.use_fine_guidance: beta = 1 / (1 + text_cfg_scale) # Concatenate conditioned and unconditioned inputs on the batch dimension batch_inputs = torch.cat([x] * n_rep, dim=0) batch_timestep = torch.cat([t] * n_rep, dim=0) # Handle CFG for cross-attention conditioning if has_text_cfg: cross_attn_cond = torch.cat( [ cross_attn_cond, no_text_cross_attn_cond, ], dim=0, ) if has_tag_cfg: cross_attn_cond = torch.cat([cross_attn_cond, no_tag_cross_attn_cond], dim=0) if has_sem_cfg: cross_attn_cond = torch.cat([cross_attn_cond, no_sem_cross_attn_cond], dim=0) if has_ctx_cfg: cross_attn_cond = torch.cat([cross_attn_cond, no_ctx_cross_attn_cond], dim=0) if has_infill_ctx_cfg: cross_attn_cond = torch.cat([cross_attn_cond, no_infill_ctx_cross_attn_cond], dim=0) if has_stem_ctx_cfg: cross_attn_cond = torch.cat([cross_attn_cond, no_stem_ctx_cross_attn_cond], dim=0) batch_output = self._forward( batch_inputs, batch_timestep, cross_attn_cond=cross_attn_cond, text_embeds=None, # TODO: Extract text embeddings from cross_attn_cond for CFG ) # Classifier-free guidance cfg_list = [] if has_text_cfg: cfg_list.append(text_cfg_scale) if has_tag_cfg: cfg_list.append(tag_cfg_scale) if has_sem_cfg: cfg_list.append(sem_cfg_scale) if has_ctx_cfg: cfg_list.append(ctx_cfg_scale) if has_infill_ctx_cfg: cfg_list.append(infill_ctx_cfg_scale) if has_stem_ctx_cfg: cfg_list.append(stem_ctx_cfg_scale) if len(cfg_list) == 4: cond_output, no_out_1, no_out_2, no_out_3, no_out_4 = torch.chunk(batch_output, 5, dim=0) batch_output = ( cond_output + (cond_output - no_out_1) * (cfg_list[0] - 1.0) + (cond_output - no_out_2) * (cfg_list[1] - 1.0) + (cond_output - no_out_3) * (cfg_list[2] - 1.0) + (cond_output - no_out_4) * (cfg_list[3] - 1.0) ) elif len(cfg_list) == 3: cond_output, no_out_1, no_out_2, no_out_3 = torch.chunk(batch_output, 4, dim=0) batch_output = ( cond_output + (cond_output - no_out_1) * (cfg_list[0] - 1.0) + (cond_output - no_out_2) * (cfg_list[1] - 1.0) + (cond_output - no_out_3) * (cfg_list[2] - 1.0) ) elif len(cfg_list) == 2: cond_output, no_out_1, no_out_2 = torch.chunk(batch_output, 3, dim=0) batch_output = ( cond_output + (cond_output - no_out_1) * (cfg_list[0] - 1.0) + (cond_output - no_out_2) * (cfg_list[1] - 1.0) ) elif len(cfg_list) == 1: cond_output, no_out_1 = torch.chunk(batch_output, 2, dim=0) batch_output = cond_output + (cond_output - no_out_1) * (cfg_list[0] - 1.0) else: batch_output = batch_output return batch_output