from __future__ import annotations import copy import math import torch import torch.nn as nn from models import ContinuousConditionEmbedder, SWaG class TwoConditionControlNetSWaG(SWaG): """SWaG with a ControlNet-style two-condition branch for P/S arrivals.""" def __init__( self, condition_dropout_prob: float = 0.1, condition_normalization_length: float = 1.0, control_depth: int = 12, **kwargs, ): super().__init__(**kwargs) if not 0 < int(control_depth) <= int(kwargs.get('depth', 12)): raise ValueError('control_depth must be between 1 and model depth') hidden = self.timestep_embedder.mlp[-1].out_features frequency = self.timestep_embedder.frequency_size self.condition_dropout_prob = float(condition_dropout_prob) self.condition_normalization_length = float(condition_normalization_length) self.control_depth = int(control_depth) self.p_embedder = ContinuousConditionEmbedder(hidden, frequency) self.s_embedder = ContinuousConditionEmbedder(hidden, frequency) self.condition_mlp = nn.Sequential( nn.Linear(2 * hidden, hidden), nn.SiLU(), nn.Linear(hidden, hidden), ) self.control_blocks = nn.ModuleList([copy.deepcopy(block) for block in self.blocks[: self.control_depth]]) self.control_input_zero = nn.Linear(hidden, hidden) # The SWaG block uses small tokens as the query and mid/large tokens as # contexts. Keep explicit zero-initialized residual paths for all # three scales so the condition can affect the full multi-scale path. self.control_mid_input_zero = nn.Linear(hidden, hidden) self.control_large_input_zero = nn.Linear(hidden, hidden) self.control_zero = nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(self.control_depth)]) self.control_mid_zero = nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(self.control_depth)]) self.control_large_zero = nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(self.control_depth)]) for embedder in (self.p_embedder, self.s_embedder): nn.init.normal_(embedder.mlp[0].weight, std=0.02) nn.init.normal_(embedder.mlp[2].weight, std=0.02) nn.init.xavier_uniform_(self.condition_mlp[0].weight) nn.init.zeros_(self.condition_mlp[0].bias) nn.init.xavier_uniform_(self.condition_mlp[-1].weight) nn.init.zeros_(self.condition_mlp[-1].bias) nn.init.zeros_(self.control_input_zero.weight) nn.init.zeros_(self.control_input_zero.bias) for projection in ( self.control_mid_input_zero, self.control_large_input_zero, *self.control_zero, *self.control_mid_zero, *self.control_large_zero, ): nn.init.zeros_(projection.weight) nn.init.zeros_(projection.bias) def copy_backbone_blocks_to_control(self) -> None: for source, target in zip(self.blocks[: self.control_depth], self.control_blocks): target.load_state_dict(source.state_dict()) def _drop_mask(self, labels: torch.Tensor, force_drop_mask=None) -> torch.Tensor: if force_drop_mask is not None: mask = force_drop_mask.to(device=labels.device, dtype=torch.bool) if mask.shape != (labels.shape[0],): raise ValueError(f'Expected force_drop_mask shape [{labels.shape[0]}], got {tuple(mask.shape)}') return mask if self.training and self.condition_dropout_prob > 0: return torch.rand(labels.shape[0], device=labels.device) < self.condition_dropout_prob return torch.zeros(labels.shape[0], dtype=torch.bool, device=labels.device) def _condition_embedding(self, labels: torch.Tensor, drop_mask: torch.Tensor) -> torch.Tensor: if labels.shape != (labels.shape[0], 2): raise ValueError(f'Expected P/S labels shape [{labels.shape[0]}, 2], got {tuple(labels.shape)}') if not torch.isfinite(labels).all(): raise ValueError('P/S labels contain non-finite values') normalized = labels / self.condition_normalization_length return self.condition_mlp( torch.cat( ( self.p_embedder(normalized[:, 0], drop_mask), self.s_embedder(normalized[:, 1], drop_mask), ), dim=-1, ) ) def forward(self, x, timesteps, labels, force_drop_mask=None): if x.ndim != 3 or x.shape[1:] != (self.in_channels, self.length): raise ValueError(f'Expected x shape [N, {self.in_channels}, {self.length}], got {tuple(x.shape)}') if timesteps.ndim != 1 or timesteps.shape[0] != x.shape[0]: raise ValueError(f'Expected timesteps shape [{x.shape[0]}], got {tuple(timesteps.shape)}') if labels.shape != (x.shape[0], 2): raise ValueError(f'Expected labels shape [{x.shape[0]}, 2], got {tuple(labels.shape)}') if not torch.isfinite(labels).all(): raise ValueError('two-condition labels contain non-finite values') timestep = self.timestep_embedder(timesteps) drop_mask = self._drop_mask(labels, force_drop_mask) condition = self._condition_embedding(labels, drop_mask) small, mid, large = self.encoder(x) small = small + self.small_position.to(dtype=small.dtype) mid = mid + self.mid_position.to(dtype=mid.dtype) large = large + self.large_position.to(dtype=large.dtype) small, mid, large = self.token_preprocessor(small, mid, large, timestep) control_small = small + self.control_input_zero(condition).unsqueeze(1) control_mid = mid + self.control_mid_input_zero(condition).unsqueeze(1) control_large = large + self.control_large_input_zero(condition).unsqueeze(1) for index, block in enumerate(self.blocks): if index < self.control_depth: control_small = self.control_blocks[index](control_small, control_mid, control_large, timestep) small = small + self.control_zero[index](control_small) mid = mid + self.control_mid_zero[index](control_mid) large = large + self.control_large_zero[index](control_large) small = block(small, mid, large, timestep) decoded = self.decoder_projection(small.transpose(1, 2)) decoded = self.decoder_refine( torch.nn.functional.interpolate(decoded, size=self.length, mode='linear', align_corners=False) ) if self.skip_branch is not None: decoded = torch.cat([decoded, self.skip_branch(x)], dim=1) return self.output_head(self.fusion(decoded)) def forward_with_cfg(self, x, timesteps, labels, cfg_scale: float): keep = torch.zeros(x.shape[0], dtype=torch.bool, device=x.device) drop = torch.ones(x.shape[0], dtype=torch.bool, device=x.device) conditional = self(x, timesteps, labels, force_drop_mask=keep) unconditional = self(x, timesteps, labels, force_drop_mask=drop) eps_c, rest_c = conditional[:, : self.in_channels], conditional[:, self.in_channels :] eps_u = unconditional[:, : self.in_channels] guided_eps = eps_u + float(cfg_scale) * (eps_c - eps_u) return torch.cat([guided_eps, rest_c], dim=1)