| """ |
| Adaptive Normalization Modules for ChronFlux |
| Includes AdaIN (Adaptive Instance Normalization) and AdaLN (Adaptive Layer Normalization) |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| NIGHT_SKY_INTENSITY = -0.95 |
| WARM_BIAS_RGB = [0.2, 0.0, -0.2] |
|
|
|
|
| class AdaIN(nn.Module): |
| """ |
| Adaptive Instance Normalization (AdaIN) for style transfer. |
| Used for generating pseudo-style targets during training with extensive augmentation. |
| |
| Reference: Huang & Belongie, "Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization" |
| """ |
| def __init__(self, eps=1e-5): |
| super().__init__() |
| self.eps = eps |
|
|
| def calc_mean_std(self, feat): |
| size = feat.size() |
| N, C = size[:2] |
| feat_view = feat.view(N, C, -1) |
|
|
| feat_std, feat_mean = torch.std_mean(feat_view, dim=2, unbiased=False) |
|
|
| feat_std = (feat_std + self.eps).view(N, C, 1, 1) |
| feat_mean = feat_mean.view(N, C, 1, 1) |
|
|
| return feat_mean, feat_std |
|
|
| def forward(self, content, style): |
| """ |
| Apply AdaIN: transfers style statistics to content |
| |
| Args: |
| content: Content features [B, C, H, W] |
| style: Style features [B, C, H, W] |
| |
| Returns: |
| Stylized content with style's mean and std |
| """ |
| size = content.size() |
| content_mean, content_std = self.calc_mean_std(content) |
| style_mean, style_std = self.calc_mean_std(style) |
|
|
| normalized = (content - content_mean) / content_std |
| stylized = normalized * style_std + style_mean |
|
|
| return stylized |
|
|
|
|
| class AdaINTransfer(nn.Module): |
| """ |
| Extended AdaIN for pseudo-style target generation with augmentation support. |
| Used specifically for training pipeline to create robust style targets. |
| """ |
| def __init__(self, eps=1e-5): |
| super().__init__() |
| self.adain = AdaIN(eps=eps) |
|
|
| self.register_buffer( |
| 'warm_bias_template', |
| torch.tensor(WARM_BIAS_RGB, dtype=torch.float32).view(1, 3, 1, 1) |
| ) |
|
|
| def adaptive_night_transfer(self, content, style, sky_mask=None): |
| """ |
| Transfer night style to content with optional sky masking |
| |
| Args: |
| content: Day images [B, C, H, W] |
| style: Night images [B, C, H, W] |
| sky_mask: Binary mask for sky regions [B, 1, H, W] |
| |
| Returns: |
| Night-styled content |
| """ |
| target = self.adain(content, style) |
|
|
| if sky_mask is not None: |
| target = target * (1 - sky_mask) + (NIGHT_SKY_INTENSITY * sky_mask) |
|
|
| return target.clamp(-0.95, 0.95) |
|
|
| def adaptive_sunset_transfer(self, content, style_night, sky_mask=None): |
| """ |
| Create sunset style by channel manipulation of night style |
| |
| Args: |
| content: Day images [B, C, H, W] |
| style_night: Night images for base style [B, C, H, W] |
| sky_mask: Binary mask for sky regions [B, 1, H, W] |
| |
| Returns: |
| Sunset-styled content with warm tones |
| """ |
| content_mean, content_std = self.adain.calc_mean_std(content) |
| style_mean, style_std = self.adain.calc_mean_std(style_night) |
|
|
| mu_swap = style_mean.clone() |
| mu_swap[:, 0] = style_mean[:, 2] |
| mu_swap[:, 2] = style_mean[:, 0] |
|
|
| std_swap = style_std.clone() |
| std_swap[:, 0] = style_std[:, 2] |
| std_swap[:, 2] = style_std[:, 0] |
|
|
| B = content.shape[0] |
| warm_bias = self.warm_bias_template.to(content.dtype).expand(B, 3, 1, 1) |
|
|
| mu_final = mu_swap * 0.7 + warm_bias * 0.3 |
| std_final = std_swap * 1.4 |
|
|
| target = (content - content_mean) / content_std * std_final + mu_final |
|
|
| if sky_mask is not None: |
| sky_color = mu_final.view(B, 3, 1, 1) |
| target = target * (1 - sky_mask) + (target * 0.4 + sky_color * 0.6) * sky_mask |
|
|
| return target.clamp(-0.95, 0.95) |
|
|
|
|
| class AdaLNModulation(nn.Module): |
| """ |
| Adaptive Layer Normalization (AdaLN) Modulation. |
| Generates scale, shift, and gate parameters for layer-wise adaptive normalization. |
| Provides better control over style consistency and layer-wise adjustments. |
| """ |
| def __init__(self, hidden_size, components=6): |
| super().__init__() |
| self.silu = nn.SiLU() |
| self.linear = nn.Linear(hidden_size, components * hidden_size, bias=True) |
| self.components = components |
|
|
| nn.init.zeros_(self.linear.weight) |
| nn.init.zeros_(self.linear.bias) |
|
|
| def forward(self, vec): |
| """ |
| Generate modulation parameters |
| |
| Args: |
| vec: Conditioning vector [B, hidden_size] |
| |
| Returns: |
| Tuple of (shift, scale, gate) parameters for each component |
| """ |
| chunk = self.linear(self.silu(vec)) |
| return chunk.chunk(self.components, dim=1) |
|
|
|
|
| class AdaLNBlock(nn.Module): |
| """ |
| Base AdaLN normalization block with scale, shift, and gate. |
| Can be used as a building block for various architectures. |
| """ |
| def __init__(self, hidden_size, use_gate=True): |
| super().__init__() |
| self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| self.use_gate = use_gate |
|
|
| def forward(self, x, scale, shift, gate=None): |
| """ |
| Apply AdaLN: normalize, then scale and shift, optionally gate |
| |
| Args: |
| x: Input features [B, L, D] |
| scale: Scale parameter [B, D] |
| shift: Shift parameter [B, D] |
| gate: Optional gate parameter [B, D] |
| |
| Returns: |
| Modulated features |
| """ |
| x_norm = self.norm(x) |
| x_mod = x_norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
| if self.use_gate and gate is not None: |
| x_mod = gate.unsqueeze(1) * x_mod |
|
|
| return x_mod |
|
|
|
|
| class AdaLNZero(nn.Module): |
| """ |
| AdaLN-Zero: AdaLN with zero-initialized output scaling. |
| Particularly useful for residual connections and stable training. |
| """ |
| def __init__(self, hidden_size): |
| super().__init__() |
| self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| self.scale = nn.Parameter(torch.zeros(hidden_size)) |
| self.shift = nn.Parameter(torch.zeros(hidden_size)) |
|
|
| def forward(self, x, modulation_scale, modulation_shift): |
| """ |
| Apply AdaLN with zero-initialization |
| |
| Args: |
| x: Input features [B, L, D] |
| modulation_scale: Additional scale from conditioning [B, D] |
| modulation_shift: Additional shift from conditioning [B, D] |
| |
| Returns: |
| Modulated features |
| """ |
| x_norm = self.norm(x) |
| x_mod = x_norm * (self.scale + modulation_scale.unsqueeze(1)) + (self.shift + modulation_shift.unsqueeze(1)) |
| return x_mod |
|
|
|
|
| class StyleAdaLN(nn.Module): |
| """ |
| Style-aware AdaLN that combines style and conditioning vectors. |
| Particularly useful for style transfer with additional controls. |
| """ |
| def __init__(self, hidden_size, style_dim): |
| super().__init__() |
| self.style_proj = nn.Linear(style_dim, hidden_size) |
| self.adaln = AdaLNModulation(hidden_size, components=2) |
| self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
|
|
| def forward(self, x, style_vec, cond_vec=None): |
| """ |
| Apply style-aware normalization |
| |
| Args: |
| x: Input features [B, L, D] |
| style_vec: Style vector [B, style_dim] |
| cond_vec: Optional additional conditioning [B, D] |
| |
| Returns: |
| Style-modulated features |
| """ |
| style_emb = self.style_proj(style_vec) |
|
|
| if cond_vec is not None: |
| combined = style_emb + cond_vec |
| else: |
| combined = style_emb |
|
|
| scale, shift = self.adaln(combined) |
|
|
| x_norm = self.norm(x) |
| x_mod = x_norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
| return x_mod |
|
|
|
|
| __all__ = [ |
| 'AdaIN', |
| 'AdaINTransfer', |
| 'AdaLNModulation', |
| 'AdaLNBlock', |
| 'AdaLNZero', |
| 'StyleAdaLN', |
| ] |
