| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import Tensor, nn |
| import torch.nn.functional as F |
|
|
| from .model import SolPixOutput |
|
|
|
|
| @dataclass(frozen=True) |
| class FlowMatchingConfig: |
| """Logit-normal timestep sampling used for the rectified-flow objective.""" |
|
|
| logit_mean: float = 0.0 |
| logit_std: float = 1.0 |
|
|
|
|
| class REPAHead(nn.Module): |
| """Optional train-only projection from SolPix features to teacher features.""" |
|
|
| def __init__(self, model_width: int = 512, teacher_width: int = 768): |
| super().__init__() |
| self.projection = nn.Sequential( |
| nn.LayerNorm(model_width), |
| nn.Linear(model_width, teacher_width), |
| ) |
|
|
| def forward(self, features: Tensor) -> Tensor: |
| return F.normalize(self.projection(features), dim=-1) |
|
|
|
|
| def sample_logit_normal_time( |
| batch_size: int, |
| device: torch.device, |
| config: FlowMatchingConfig | None = None, |
| ) -> Tensor: |
| config = config or FlowMatchingConfig() |
| logits = torch.randn(batch_size, device=device) * config.logit_std + config.logit_mean |
| return logits.sigmoid() |
|
|
|
|
| def flow_matching_loss( |
| model: nn.Module, |
| clean_latents: Tensor, |
| text_embeddings: Tensor, |
| text_mask: Tensor | None = None, |
| *, |
| noise: Tensor | None = None, |
| time: Tensor | None = None, |
| config: FlowMatchingConfig | None = None, |
| ) -> tuple[Tensor, Tensor, Tensor]: |
| """Compute rectified-flow MSE for x_t=(1-t)x0+t*noise, v=noise-x0. |
| |
| Returns ``(loss, predicted_velocity, time)`` so callers can log timesteps and |
| combine this term with auxiliary objectives such as REPA. |
| """ |
| if clean_latents.ndim != 4: |
| raise ValueError("clean_latents must have shape [batch, channels, height, width]") |
| batch = clean_latents.shape[0] |
| noise = torch.randn_like(clean_latents) if noise is None else noise |
| if noise.shape != clean_latents.shape: |
| raise ValueError("noise must have the same shape as clean_latents") |
| time = sample_logit_normal_time(batch, clean_latents.device, config) if time is None else time |
| if time.ndim == 0: |
| time = time.expand(batch) |
| if time.shape != (batch,): |
| raise ValueError("time must be a scalar or a [batch] tensor") |
| broadcast_time = time.to(clean_latents.dtype).reshape(batch, 1, 1, 1) |
| noisy_latents = (1 - broadcast_time) * clean_latents + broadcast_time * noise |
| target_velocity = noise - clean_latents |
| predicted_velocity = model(noisy_latents, time, text_embeddings, text_mask) |
| loss = F.mse_loss(predicted_velocity.float(), target_velocity.float()) |
| return loss, predicted_velocity, time |
|
|
|
|
| def representation_alignment_loss( |
| output: SolPixOutput, |
| teacher_features: Tensor, |
| head: REPAHead, |
| token_mask: Tensor | None = None, |
| ) -> Tensor: |
| """Cosine alignment loss for clean-image teacher features (REPA-style).""" |
| if output.repa_features is None: |
| raise ValueError("call SolPix(..., return_features=True) to obtain REPA features") |
| prediction = head(output.repa_features) |
| target = F.normalize(teacher_features.detach(), dim=-1) |
| if prediction.shape != target.shape: |
| raise ValueError( |
| f"projected model features {tuple(prediction.shape)} do not match teacher features {tuple(target.shape)}" |
| ) |
| per_token = 1 - (prediction * target).sum(dim=-1) |
| if token_mask is None: |
| return per_token.mean() |
| if token_mask.shape != per_token.shape: |
| raise ValueError("token_mask must have shape [batch, image_tokens]") |
| weights = token_mask.to(per_token.dtype) |
| return (per_token * weights).sum() / weights.sum().clamp_min(1) |
|
|
