File size: 7,463 Bytes
275b5a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 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)
|