| from __future__ import annotations |
|
|
| from math import isqrt |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| FACTOR_NAMES = ("line", "color", "texture", "layout") |
|
|
|
|
| class ValueOnlyResampler(nn.Module): |
| """Position-free, bias-free resampling whose output is zero for zero values.""" |
|
|
| def __init__(self, width: int, tokens: int = 144, heads: int = 8) -> None: |
| super().__init__() |
| if width % heads: |
| raise ValueError("width must be divisible by heads") |
| self.heads = heads |
| self.head_dim = width // heads |
| self.queries = nn.Parameter(torch.empty(tokens, width)) |
| self.q_proj = nn.Linear(width, width, bias=False) |
| self.k_proj = nn.Linear(width, width, bias=False) |
| self.v_proj = nn.Linear(width, width, bias=False) |
| self.out_proj = nn.Linear(width, width, bias=False) |
| nn.init.normal_(self.queries, std=width**-0.5) |
|
|
| def forward(self, values: torch.Tensor) -> torch.Tensor: |
| batch, source_tokens, width = values.shape |
| target_tokens = self.queries.shape[0] |
| queries = self.queries.unsqueeze(0).expand(batch, -1, -1) |
| q = self.q_proj(queries).view(batch, target_tokens, self.heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(values).view(batch, source_tokens, self.heads, self.head_dim).transpose(1, 2) |
| result = F.scaled_dot_product_attention(q, k, v) |
| return self.out_proj(result.transpose(1, 2).reshape(batch, target_tokens, width)) |
|
|
|
|
| class SharedAxisStyleEncoder(nn.Module): |
| """One VAE field and one shared axis space for embedding, routing, and transfer.""" |
|
|
| def __init__( |
| self, |
| dino_dim: int, |
| style_dim: int = 1024, |
| axis_count: int = 256, |
| embedding_dim: int = 1024, |
| anima_dim: int = 0, |
| resolution: str = "s12", |
| set_layers: int = 2, |
| set_heads: int = 8, |
| direct_context_axes: bool = False, |
| axis_aligned: bool = False, |
| disable_token_film: bool = False, |
| translation_invariant_stats: bool = False, |
| ) -> None: |
| super().__init__() |
| if resolution not in {"s12", "s24r"}: |
| raise ValueError(f"unknown spatial resolution: {resolution}") |
| if axis_count % set_heads: |
| raise ValueError("axis_count must be divisible by set_heads") |
| self.resolution = resolution |
| self.style_dim = style_dim |
| self.axis_count = axis_count |
| self.embedding_dim = embedding_dim |
| self.set_layers = set_layers |
| self.set_heads = set_heads |
| self.direct_context_axes = direct_context_axes |
| self.axis_aligned = axis_aligned |
| self.disable_token_film = disable_token_film |
| self.translation_invariant_stats = translation_invariant_stats |
| if disable_token_film and not axis_aligned: |
| raise ValueError("disable_token_film requires axis_aligned=True") |
| if translation_invariant_stats and not axis_aligned: |
| raise ValueError("translation-invariant statistics require axis_aligned=True") |
| last_stride = 2 if resolution == "s12" else 1 |
| self.stem = nn.Sequential( |
| nn.Conv2d(16, 256, 3, stride=2, padding=1), |
| nn.GELU(), |
| nn.Conv2d(256, 512, 3, stride=2, padding=1), |
| nn.GELU(), |
| nn.Conv2d(512, style_dim, 3, stride=last_stride, padding=1), |
| ) |
| self.token_norm = nn.LayerNorm(style_dim, elementwise_affine=False) |
| self.dino_norm = nn.LayerNorm(dino_dim) |
| self.dino_film = nn.Linear(dino_dim, 2 * style_dim, bias=False) |
| self.dino_relation = nn.Linear(dino_dim, axis_count, bias=False) |
| nn.init.zeros_(self.dino_film.weight) |
|
|
| self.axis_dictionary = nn.Parameter(torch.empty(style_dim, axis_count)) |
| nn.init.orthogonal_(self.axis_dictionary) |
| self.membership_logits = nn.Parameter(torch.zeros(len(FACTOR_NAMES), axis_count)) |
| self.moment_projection = nn.Linear(2 * axis_count, axis_count, bias=False) |
| relation_layer = nn.TransformerEncoderLayer( |
| axis_count, |
| nhead=set_heads, |
| dim_feedforward=2 * axis_count, |
| dropout=0.0, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.relation_encoder = nn.TransformerEncoder( |
| relation_layer, |
| set_layers, |
| enable_nested_tensor=False, |
| ) |
| self.axis_reliability = nn.Linear(axis_count, axis_count) |
| self.face_reliability = nn.Linear(axis_count, 1) |
| self.embedding_head = nn.Sequential( |
| nn.LayerNorm(axis_count), |
| nn.Linear(axis_count, embedding_dim), |
| ) |
| self.resampler = ( |
| ValueOnlyResampler(style_dim, tokens=144, heads=set_heads) |
| if resolution == "s24r" |
| else None |
| ) |
| |
| self.anima_norm = nn.LayerNorm(anima_dim) if anima_dim else None |
| self.anima_film = nn.Linear(anima_dim, 2 * style_dim, bias=False) if anima_dim else None |
| self.anima_relation = nn.Linear(anima_dim, axis_count, bias=False) if anima_dim else None |
| if self.anima_film is not None: |
| nn.init.zeros_(self.anima_film.weight) |
| if axis_aligned: |
| self.axis_moment_weights = nn.Parameter(torch.empty(axis_count, 2)) |
| nn.init.normal_(self.axis_moment_weights, std=axis_count**-0.5) |
| if translation_invariant_stats: |
| self.axis_power_weights = nn.Parameter(torch.zeros(axis_count, 4)) |
|
|
| def model_config(self) -> dict[str, int | str | bool]: |
| return { |
| "style_dim": self.style_dim, |
| "axis_count": self.axis_count, |
| "embedding_dim": self.embedding_dim, |
| "anima_dim": self.anima_norm.normalized_shape[0] if self.anima_norm is not None else 0, |
| "resolution": self.resolution, |
| "set_layers": self.set_layers, |
| "set_heads": self.set_heads, |
| "direct_context_axes": self.direct_context_axes, |
| "axis_aligned": self.axis_aligned, |
| "disable_token_film": self.disable_token_film, |
| "translation_invariant_stats": self.translation_invariant_stats, |
| } |
|
|
| def _field( |
| self, |
| latent: torch.Tensor, |
| scale: torch.Tensor, |
| shift: torch.Tensor, |
| axis_shift: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| batch, references = latent.shape[:2] |
| field = self.stem(latent.flatten(0, 1).float()) |
| field = field.flatten(2).transpose(1, 2).reshape(batch, references, -1, field.shape[1]) |
| field = self.token_norm(field) |
| field = field * (1 + scale.unsqueeze(2)) + shift.unsqueeze(2) |
| axes = torch.einsum("brnd,dk->brnk", field, self.axis_dictionary.to(field.dtype)) |
| if axis_shift is not None: |
| axes = axes + axis_shift.unsqueeze(2) |
| return field, axes |
|
|
| def _moments(self, axes: torch.Tensor) -> torch.Tensor: |
| mean = axes.mean(dim=2) |
| log_std = torch.log(axes.std(dim=2, unbiased=False).clamp_min(1e-6)) |
| if self.axis_aligned: |
| coordinates = ( |
| torch.stack((mean, log_std), dim=-1) |
| * self.axis_moment_weights.to(mean.dtype) |
| ).sum(-1) |
| if self.translation_invariant_stats: |
| power = self._power_statistics(axes).to(mean.dtype) |
| coordinates = coordinates + ( |
| power * self.axis_power_weights.to(mean.dtype) |
| ).sum(-1) |
| return coordinates |
| moments = torch.cat((mean, log_std), dim=-1) |
| return self.moment_projection(moments) |
|
|
| def _power_statistics(self, axes: torch.Tensor) -> torch.Tensor: |
| """Phase-free radial and directional power for each aligned spatial axis.""" |
| side = isqrt(axes.shape[2]) |
| if side * side != axes.shape[2]: |
| raise ValueError(f"unexpected spatial token count: {axes.shape[2]}") |
| fy = torch.fft.fftfreq(side, device=axes.device)[:, None] |
| fx = torch.fft.rfftfreq(side, device=axes.device)[None, :] |
| radius2 = fy.square() + fx.square() |
| radius = radius2.sqrt() |
| basis = torch.stack(( |
| ((radius > 0) & (radius <= 0.25)).float(), |
| ((radius > 0.25) & (radius <= 0.5)).float(), |
| (fx.square() - fy.square()) / radius2.clamp_min(1e-12), |
| 4 * fx.square() * fy.square() / radius2.square().clamp_min(1e-12), |
| )) |
| rfft_weight = torch.ones(side // 2 + 1, device=axes.device) |
| rfft_weight[1:-1] = 2 |
| maps = axes.float().permute(0, 1, 3, 2).reshape( |
| *axes.shape[:2], |
| axes.shape[-1], |
| side, |
| side, |
| ) |
| maps = maps - maps.mean(dim=(-2, -1), keepdim=True) |
| power = torch.fft.rfft2(maps, norm="ortho").abs().square() |
| power = power * rfft_weight[None, None, None, None, :] |
| power = power / power.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-12) |
| return torch.einsum("brkhw,phw->brkp", power, basis) |
|
|
| def factor_coordinates(self, axis_coordinates: torch.Tensor) -> torch.Tensor: |
| """Apply overlapping factor memberships without creating factor branches.""" |
| memberships = torch.sigmoid(self.membership_logits).to(axis_coordinates.dtype) |
| return axis_coordinates.unsqueeze(-2) * memberships |
|
|
| def encode_vae_axes(self, latents: torch.Tensor) -> torch.Tensor: |
| """Encode VAE latents without DINO/Anima context for the frozen transfer teacher.""" |
| if latents.ndim != 5 or latents.shape[2] != 16: |
| raise ValueError(f"expected latents [B,R,16,H,W], got {tuple(latents.shape)}") |
| batch, references = latents.shape[:2] |
| zeros = torch.zeros( |
| batch, |
| references, |
| self.axis_dictionary.shape[0], |
| device=latents.device, |
| dtype=latents.dtype, |
| ) |
| _, axes = self._field(latents, zeros, zeros) |
| return self._moments(axes) |
|
|
| def _user_axis_gate(self, user_weights: torch.Tensor) -> torch.Tensor: |
| memberships = torch.sigmoid(self.membership_logits) |
| factors = user_weights[..., : len(FACTOR_NAMES)].unsqueeze(-1) |
| return 1 - torch.prod(1 - factors * memberships[None, None], dim=-2) |
|
|
| def _controlled_tokens( |
| self, |
| axes: torch.Tensor, |
| beta: torch.Tensor, |
| view_gate: torch.Tensor, |
| ) -> torch.Tensor: |
| gated = axes * beta.unsqueeze(2) * view_gate[:, :, None, None] |
| tokens = torch.einsum( |
| "brnk,dk->brnd", |
| gated, |
| self.axis_dictionary.to(gated.dtype), |
| ) |
| if self.resampler is not None: |
| batch, references, token_count, width = tokens.shape |
| tokens = self.resampler(tokens.reshape(batch * references, token_count, width)) |
| tokens = tokens.reshape(batch, references, -1, width) |
| return tokens |
|
|
| def soft_orthogonality_loss(self) -> torch.Tensor: |
| dictionary = F.normalize(self.axis_dictionary, dim=0) |
| gram = dictionary.T @ dictionary |
| return (gram - torch.eye(gram.shape[0], device=gram.device, dtype=gram.dtype)).square().mean() |
|
|
| def forward( |
| self, |
| full_latents: torch.Tensor, |
| dino: torch.Tensor, |
| reference_valid: torch.Tensor, |
| face_latents: torch.Tensor | None = None, |
| face_valid: torch.Tensor | None = None, |
| user_weights: torch.Tensor | None = None, |
| mode: str = "auto", |
| overall_style_gain: float | torch.Tensor = 1.0, |
| build_memory: bool = True, |
| anima: torch.Tensor | None = None, |
| ) -> dict[str, torch.Tensor]: |
| if full_latents.ndim != 5 or full_latents.shape[2] != 16: |
| raise ValueError(f"expected full latents [B,R,16,H,W], got {tuple(full_latents.shape)}") |
| if not reference_valid.any(dim=1).all(): |
| raise ValueError("every sample needs at least one valid reference") |
| if mode not in {"auto", "assisted", "manual"}: |
| raise ValueError(f"unknown routing mode: {mode}") |
|
|
| batch, references = full_latents.shape[:2] |
| if user_weights is None: |
| user_weights = torch.ones( |
| batch, |
| references, |
| len(FACTOR_NAMES) + 1, |
| device=full_latents.device, |
| dtype=full_latents.dtype, |
| ) |
| context = self.dino_norm(dino.float()) |
| scale, shift = self.dino_film(context).chunk(2, dim=-1) |
| relation_context = self.dino_relation(context) |
| if self.anima_norm is not None: |
| if anima is None: |
| raise ValueError("Anima features are required when anima_dim is configured") |
| anima_context = self.anima_norm(anima.float().flatten(start_dim=2)) |
| anima_scale, anima_shift = self.anima_film(anima_context).chunk(2, dim=-1) |
| scale = scale + anima_scale |
| shift = shift + anima_shift |
| relation_context = relation_context + self.anima_relation(anima_context) |
| elif anima is not None: |
| raise ValueError("Anima features were provided but anima_dim=0") |
| if self.disable_token_film: |
| scale = torch.zeros_like(scale) |
| shift = torch.zeros_like(shift) |
| axis_shift = ( |
| relation_context |
| if self.axis_aligned and self.direct_context_axes |
| else None |
| ) |
| full_field, full_axes = self._field( |
| full_latents, |
| scale, |
| shift, |
| axis_shift, |
| ) |
| full_coordinates = self._moments(full_axes) |
| if self.direct_context_axes and not self.axis_aligned: |
| full_coordinates = full_coordinates + relation_context |
|
|
| if face_latents is None: |
| face_valid = torch.zeros_like(reference_valid) |
| face_axes = None |
| face_coordinates = torch.zeros_like(full_coordinates) |
| else: |
| if face_valid is None: |
| raise ValueError("face_valid is required with face_latents") |
| _, face_axes = self._field( |
| face_latents, |
| scale, |
| shift, |
| axis_shift, |
| ) |
| face_coordinates = self._moments(face_axes) |
|
|
| preliminary = ( |
| full_coordinates |
| + face_coordinates * face_valid.unsqueeze(-1) |
| + (0 if self.direct_context_axes else relation_context) |
| ) |
| contextual = self.relation_encoder(preliminary, src_key_padding_mask=~reference_valid) |
| reliability = torch.sigmoid(self.axis_reliability(contextual)) |
| predicted_face = torch.sigmoid(self.face_reliability(contextual).squeeze(-1)) |
|
|
| if mode == "auto": |
| face_weight = predicted_face |
| elif mode == "assisted": |
| face_weight = predicted_face * user_weights[..., -1] |
| else: |
| face_weight = user_weights[..., -1] |
| face_weight = face_weight * face_valid * reference_valid |
| coordinates = ( |
| full_coordinates + face_weight.unsqueeze(-1) * face_coordinates |
| ) / (1 + face_weight.unsqueeze(-1)) |
|
|
| axis_gate = self._user_axis_gate(user_weights.float()) |
| valid = reference_valid.unsqueeze(-1).to(axis_gate.dtype) |
| if mode == "auto": |
| weights = reliability * valid |
| gamma = torch.ones(batch, self.axis_dictionary.shape[1], device=weights.device, dtype=weights.dtype) |
| elif mode == "assisted": |
| weights = reliability * axis_gate * valid |
| gamma = axis_gate.amax(dim=1) |
| else: |
| weights = axis_gate * valid |
| gamma = axis_gate.amax(dim=1) |
|
|
| denominator = weights.sum(dim=1, keepdim=True) |
| alpha = torch.where( |
| denominator > 0, |
| weights / denominator.clamp_min(torch.finfo(weights.dtype).tiny), |
| torch.zeros_like(weights), |
| ) |
| beta = alpha * gamma.unsqueeze(1) |
| mixed = gamma * torch.einsum("brk,brk->bk", alpha, coordinates) |
|
|
| gain = torch.as_tensor(overall_style_gain, device=mixed.device, dtype=mixed.dtype) |
| while gain.ndim < mixed.ndim: |
| gain = gain.unsqueeze(-1) |
| mixed = mixed * gain |
| beta = beta * gain.unsqueeze(1) if gain.ndim == 2 else beta * gain |
|
|
| pre_embedding = self.embedding_head(mixed) |
| output = { |
| "embedding": F.normalize(pre_embedding, dim=-1), |
| "pre_embedding": pre_embedding, |
| "style_present": mixed.ne(0).any(dim=-1), |
| "reference_axis_coordinates": coordinates, |
| "axis_coordinates": mixed, |
| "axis_reliability": reliability, |
| "axis_weights": weights, |
| "axis_alpha": alpha, |
| "axis_gamma": gamma, |
| "factor_memberships": torch.sigmoid(self.membership_logits), |
| "reference_factor_coordinates": self.factor_coordinates(coordinates), |
| "factor_coordinates": self.factor_coordinates(mixed), |
| "face_weight": face_weight, |
| } |
| if build_memory: |
| full_tokens = self._controlled_tokens( |
| full_axes, |
| beta, |
| reference_valid.to(beta.dtype), |
| ) |
| memories = [full_tokens.flatten(1, 2)] |
| memory_masks = [ |
| reference_valid[:, :, None].expand(-1, -1, full_tokens.shape[2]).flatten(1) |
| ] |
| if face_axes is not None: |
| face_tokens = self._controlled_tokens(face_axes, beta, face_weight) |
| memories.append(face_tokens.flatten(1, 2)) |
| memory_masks.append( |
| (reference_valid & face_valid)[:, :, None] |
| .expand(-1, -1, face_tokens.shape[2]) |
| .flatten(1) |
| ) |
| output.update({ |
| "style_memory": torch.cat(memories, dim=1), |
| "style_memory_valid": torch.cat(memory_masks, dim=1), |
| "style_condition": torch.einsum( |
| "bk,dk->bd", |
| mixed, |
| self.axis_dictionary.to(mixed.dtype), |
| ), |
| }) |
| return output |
|
|
|
|
| def load_shared_axis_state( |
| model: SharedAxisStyleEncoder, |
| state: dict[str, torch.Tensor], |
| ) -> None: |
| """Load a checkpoint, deriving aligned per-axis moments from a legacy dense head.""" |
| state = dict(state) |
| if model.axis_aligned and "axis_moment_weights" not in state: |
| dense = state["moment_projection.weight"] |
| axes = model.axis_count |
| state["axis_moment_weights"] = torch.stack(( |
| dense.diagonal(), |
| dense[:, axes:].diagonal(), |
| ), dim=-1) |
| if model.translation_invariant_stats and "axis_power_weights" not in state: |
| state["axis_power_weights"] = torch.zeros_like(model.axis_power_weights) |
| model.load_state_dict(state) |
|
|
