Transfer_and_Test/src/models_controlnet_two_condition.py

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)