src/models/discriminator.py

"""
Projected discriminator for fine-tuning PixelGen with an adversarial loss.

Design (mirrors StyleGAN-T, https://github.com/autonomousvision/stylegan-t,
but adapted for this project):
    - Backbone: a *frozen* DINOv2 (ViT-B/14) feature extractor.  Intermediate
      patch-token features from several blocks are pulled out as multi-scale
      perceptual representations.
    - Heads: a small CNN head (1-D conv + residual block + 1x1 cls conv) is
      attached on top of each selected DINO layer.  These heads are the only
      part of the discriminator that is updated by the optimizer.
    - Text conditioning: the per-token text embedding produced by the Qwen3
      encoder is mean-pooled and projected to a low-dim cmap vector that is
      inner-producted with the head's per-location features (projected
      discriminator conditioning, same trick as StyleGAN-T / ProjectedGAN).

Trainable parameters: ``ProjectedDiscriminator.heads`` and
``ProjectedDiscriminator.text_proj``.
Frozen parameters: the DINOv2 backbone passed in via ``dino_encoder``.
"""

from typing import List, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from torchvision.transforms import RandomCrop

from src.utils.no_grad import freeze_model


# -----------------------------------------------------------------------------
# Differentiable augmentation (DiffAugment), trimmed copy of the version used by
# StyleGAN-T (training/diffaug.py) so we don't have to pull in their package.
# -----------------------------------------------------------------------------
def _rand_brightness(x: torch.Tensor) -> torch.Tensor:
    return x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)


def _rand_saturation(x: torch.Tensor) -> torch.Tensor:
    x_mean = x.mean(dim=1, keepdim=True)
    return (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean


def _rand_contrast(x: torch.Tensor) -> torch.Tensor:
    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
    return (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean


def _rand_translation(x: torch.Tensor, ratio: float = 0.125) -> torch.Tensor:
    shift_x = int(x.size(2) * ratio + 0.5)
    shift_y = int(x.size(3) * ratio + 0.5)
    tx = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
    ty = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
    grid_b, grid_x, grid_y = torch.meshgrid(
        torch.arange(x.size(0), dtype=torch.long, device=x.device),
        torch.arange(x.size(2), dtype=torch.long, device=x.device),
        torch.arange(x.size(3), dtype=torch.long, device=x.device),
        indexing="ij",
    )
    grid_x = torch.clamp(grid_x + tx + 1, 0, x.size(2) + 1)
    grid_y = torch.clamp(grid_y + ty + 1, 0, x.size(3) + 1)
    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
    return x_pad.permute(0, 2, 3, 1).contiguous()[grid_b, grid_x, grid_y].permute(0, 3, 1, 2)


def _rand_cutout(x: torch.Tensor, ratio: float = 0.2) -> torch.Tensor:
    cut_h, cut_w = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
    off_x = torch.randint(0, x.size(2) + (1 - cut_h % 2), size=[x.size(0), 1, 1], device=x.device)
    off_y = torch.randint(0, x.size(3) + (1 - cut_w % 2), size=[x.size(0), 1, 1], device=x.device)
    grid_b, grid_x, grid_y = torch.meshgrid(
        torch.arange(x.size(0), dtype=torch.long, device=x.device),
        torch.arange(cut_h, dtype=torch.long, device=x.device),
        torch.arange(cut_w, dtype=torch.long, device=x.device),
        indexing="ij",
    )
    grid_x = torch.clamp(grid_x + off_x - cut_h // 2, min=0, max=x.size(2) - 1)
    grid_y = torch.clamp(grid_y + off_y - cut_w // 2, min=0, max=x.size(3) - 1)
    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
    mask[grid_b, grid_x, grid_y] = 0
    return x * mask.unsqueeze(1)


_AUGMENT_FNS = {
    "color": [_rand_brightness, _rand_saturation, _rand_contrast],
    "translation": [_rand_translation],
    "cutout": [_rand_cutout],
}


def diff_augment(x: torch.Tensor, policy: str = "color,translation,cutout") -> torch.Tensor:
    if not policy:
        return x
    for p in policy.split(","):
        for f in _AUGMENT_FNS[p]:
            x = f(x)
    return x.contiguous()


# -----------------------------------------------------------------------------
# Head blocks (mirroring StyleGAN-T's DiscHead but with stock PyTorch ops).
# -----------------------------------------------------------------------------
class _BatchNormLocal(nn.Module):
    """Identical to StyleGAN-T's BatchNormLocal: virtual-batch BN over groups."""

    def __init__(self, num_features: int, virtual_bs: int = 8, eps: float = 1e-5):
        super().__init__()
        self.virtual_bs = virtual_bs
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = x.size()
        groups = int(np.ceil(x.size(0) / self.virtual_bs))
        x = x.view(groups, -1, x.size(-2), x.size(-1))
        mean = x.mean([1, 3], keepdim=True)
        var = x.var([1, 3], keepdim=True, unbiased=False)
        x = (x - mean) / torch.sqrt(var + self.eps)
        x = x * self.weight[None, :, None] + self.bias[None, :, None]
        return x.view(shape)


def _spec_conv1d(in_c: int, out_c: int, kernel_size: int) -> nn.Module:
    layer = nn.Conv1d(in_c, out_c, kernel_size=kernel_size,
                      padding=kernel_size // 2, padding_mode="circular")
    return spectral_norm(layer, n_power_iterations=1, eps=1e-12)


def _block(channels: int, kernel_size: int) -> nn.Sequential:
    return nn.Sequential(
        _spec_conv1d(channels, channels, kernel_size),
        _BatchNormLocal(channels),
        nn.LeakyReLU(0.2, inplace=True),
    )


class _Residual(nn.Module):
    def __init__(self, fn: nn.Module):
        super().__init__()
        self.fn = fn

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return (self.fn(x) + x) / np.sqrt(2.0)


class DiscHead(nn.Module):
    """Per-layer discriminator head -- projected-conditional logit producer."""

    def __init__(self, channels: int, c_dim: int, cmap_dim: int = 64):
        super().__init__()
        self.c_dim = c_dim
        self.cmap_dim = cmap_dim
        self.main = nn.Sequential(
            _block(channels, kernel_size=1),
            _Residual(_block(channels, kernel_size=9)),
        )
        if c_dim > 0:
            self.cmapper = nn.Linear(c_dim, cmap_dim)
            self.cls = spectral_norm(
                nn.Conv1d(channels, cmap_dim, kernel_size=1, padding=0),
                n_power_iterations=1, eps=1e-12,
            )
        else:
            self.cmapper = None
            self.cls = spectral_norm(
                nn.Conv1d(channels, 1, kernel_size=1, padding=0),
                n_power_iterations=1, eps=1e-12,
            )

    def forward(self, x: torch.Tensor, c: Optional[torch.Tensor]) -> torch.Tensor:
        h = self.main(x)
        out = self.cls(h)
        if self.cmapper is not None and c is not None:
            cmap = self.cmapper(c).unsqueeze(-1)
            out = (out * cmap).sum(1, keepdim=True) * (1.0 / np.sqrt(self.cmap_dim))
        return out


# -----------------------------------------------------------------------------
# Projected discriminator: frozen DINOv2 + trainable heads.
# -----------------------------------------------------------------------------
class ProjectedDiscriminator(nn.Module):
    """Projected discriminator built on top of a *frozen* DINOv2 backbone.

    Args:
        dino_encoder: an already-constructed ``src.models.encoder.DINOv2`` (or a
            module exposing ``get_intermediate_feats(x, n=...)`` returning a
            tuple of ``(B, N, D)`` patch-token tensors).  Its parameters are
            frozen by this class.
        dino_layers: which DINOv2 blocks to tap (0-indexed).
        embed_dim: DINOv2 patch-token dimension (768 for ViT-B/14).
        text_dim: token-wise text embedding dim coming from the conditioner.
        cmap_dim: dimension of the projected text vector inside each head.
        diffaug: whether to apply DiffAugment to the discriminator input.
        diffaug_policy: DiffAugment policy string.
        p_crop: probability of random-cropping when the input image is larger
            than ``crop_resolution``.
        crop_resolution: side length used by the optional random crop.
        input_resolution: the resolution the discriminator resizes input to
            before feeding to DINOv2 (set to ``None`` to skip the resize and
            let DINOv2 handle arbitrary sizes itself).
    """

    def __init__(
        self,
        dino_encoder: nn.Module,
        dino_layers: List[int] = (2, 5, 8, 11),
        embed_dim: int = 768,
        text_dim: int = 2048,
        cmap_dim: int = 64,
        diffaug: bool = True,
        diffaug_policy: str = "color,translation,cutout",
        p_crop: float = 0.5,
        crop_resolution: int = 224,
        input_resolution: Optional[int] = 224,
    ):
        super().__init__()
        self.dino = dino_encoder
        freeze_model(self.dino)
        self.dino_layers = list(dino_layers)
        self.embed_dim = embed_dim
        self.cmap_dim = cmap_dim
        self.diffaug = diffaug
        self.diffaug_policy = diffaug_policy
        self.p_crop = p_crop
        self.crop_resolution = crop_resolution
        self.input_resolution = input_resolution

        self.heads = nn.ModuleList([
            DiscHead(embed_dim, c_dim=cmap_dim, cmap_dim=cmap_dim)
            for _ in self.dino_layers
        ])

        self.text_proj = nn.Sequential(
            nn.LayerNorm(text_dim),
            nn.Linear(text_dim, cmap_dim),
        )

    def train(self, mode: bool = True):
        super().train(mode)
        # backbone always in eval mode (BN/dropout off, no gradients)
        self.dino.eval()
        for p in self.dino.parameters():
            p.requires_grad_(False)
        return self

    def trainable_parameters(self):
        for p in self.heads.parameters():
            yield p
        for p in self.text_proj.parameters():
            yield p

    # --- text conditioning helpers --------------------------------------------------
    @staticmethod
    def _pool_text(c_text: torch.Tensor, valid_length: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Pool per-token text embeddings to (B, D)."""
        if c_text.ndim == 2:
            return c_text
        if valid_length is None:
            return c_text.mean(dim=1)
        # masked mean
        B, T, _ = c_text.shape
        mask = torch.arange(T, device=c_text.device)[None, :] < valid_length[:, None]
        mask = mask.to(c_text.dtype).unsqueeze(-1)
        denom = mask.sum(dim=1).clamp_min(1.0)
        return (c_text * mask).sum(dim=1) / denom

    # --- main forward ---------------------------------------------------------------
    def forward(
        self,
        x: torch.Tensor,
        c_text: torch.Tensor,
        valid_length: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Return concatenated per-head logits (B, K) where K = sum of per-head outputs.

        Args:
            x: image tensor in ``[-1, 1]`` of shape ``(B, 3, H, W)``.
            c_text: text embeddings ``(B, T, D)`` or pooled ``(B, D)``.
            valid_length: optional ``(B,)`` long tensor of valid-token counts
                for masked mean pooling of ``c_text``.
        """
        if self.diffaug:
            x = diff_augment(x, policy=self.diffaug_policy)

        # to [0, 1]
        x = x.add(1).div(2).clamp(0.0, 1.0)

        # optional random crop / resize so DINO sees a fixed size
        if self.input_resolution is not None:
            if x.size(-1) > self.crop_resolution and np.random.random() < self.p_crop:
                x = RandomCrop(self.crop_resolution)(x)
            if x.size(-1) != self.input_resolution:
                x = F.interpolate(x, size=self.input_resolution, mode="area")

        # DINOv2 features (frozen): tuple of (B, N, D)
        feats = self.dino.get_intermediate_feats(
            x, resize=False, n=self.dino_layers,
            reshape=False, return_class_token=False,
        )

        # pooled text condition -> (B, cmap_dim)
        c_pooled = self._pool_text(c_text.float(), valid_length)
        c_pooled = self.text_proj(c_pooled)

        logits = []
        for feat, head in zip(feats, self.heads):
            feat_t = feat.float().transpose(1, 2).contiguous()  # (B, D, N)
            logits.append(head(feat_t, c_pooled).reshape(x.size(0), -1))
        return torch.cat(logits, dim=1)


class DenoiserBackboneDiscriminator(nn.Module):
    """DMD2-style discriminator that reuses a trainable denoiser backbone.

    The backbone should support ``classify_mode=True`` and return its deepest
    hidden feature. For PixelGen's DiT this is a token tensor ``(B, N, C)``,
    which is mean-pooled and fed to a small randomly initialized MLP head.
    """

    def __init__(
        self,
        backbone: nn.Module,
        hidden_size: int = 1536,
        head_hidden_size: Optional[int] = None,
        head_depth: int = 2,
        train_backbone: bool = False,
        backbone_lr_scale: float = 1.0,
    ):
        super().__init__()
        self.backbone = backbone
        self.train_backbone = train_backbone
        self.backbone_lr_scale = backbone_lr_scale
        head_hidden_size = head_hidden_size or hidden_size

        layers = [
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, head_hidden_size),
            nn.SiLU(),
        ]
        for _ in range(max(head_depth - 1, 0)):
            layers.extend([
                nn.Linear(head_hidden_size, head_hidden_size),
                nn.LayerNorm(head_hidden_size),
                nn.SiLU(),
            ])
        layers.append(nn.Linear(head_hidden_size, 1))
        self.cls_pred_branch = nn.Sequential(*layers)
        self._initialized_from_denoiser = False
        self._set_backbone_requires_grad(self.train_backbone)

    def initialize_from_denoiser(self, denoiser: nn.Module, force: bool = False):
        if self._initialized_from_denoiser and not force:
            return
        missing, unexpected = self.backbone.load_state_dict(denoiser.state_dict(), strict=False)
        if missing or unexpected:
            print(
                "Initialized D discriminator backbone from denoiser with "
                f"missing={len(missing)} unexpected={len(unexpected)}"
            )
        else:
            print("Initialized D discriminator backbone from denoiser.")
        self._initialized_from_denoiser = True
        self._set_backbone_requires_grad(self.train_backbone)

    def set_backbone_trainable(self, trainable: bool):
        self.train_backbone = trainable
        self._set_backbone_requires_grad(trainable)

    def _set_backbone_requires_grad(self, trainable: bool):
        for param in self.backbone.parameters():
            param.requires_grad_(trainable)
        if trainable:
            self.backbone.train(self.training)
        else:
            self.backbone.eval()

    def set_trainable(self, requires_grad: bool):
        for param in self.cls_pred_branch.parameters():
            param.requires_grad_(requires_grad)
        self._set_backbone_requires_grad(requires_grad and self.train_backbone)

    def optimizer_param_groups(self):
        groups = [
            {
                "params": [p for p in self.cls_pred_branch.parameters() if p.requires_grad],
                "name": "discriminator_head",
            },
        ]
        backbone_params = [p for p in self.backbone.parameters() if p.requires_grad]
        if backbone_params:
            groups.append(
                {
                    "params": backbone_params,
                    "lr_scale": self.backbone_lr_scale,
                    "name": "discriminator_backbone",
                }
            )
        return groups

    def train(self, mode: bool = True):
        super().train(mode)
        if not self.train_backbone:
            self.backbone.eval()
        return self

    @staticmethod
    def _pool_feature(rep: torch.Tensor) -> torch.Tensor:
        if rep.ndim == 3:
            return rep.float().mean(dim=1)
        if rep.ndim == 4:
            return rep.float().mean(dim=(2, 3))
        if rep.ndim == 2:
            return rep.float()
        raise ValueError(f"Unsupported discriminator feature shape: {tuple(rep.shape)}")

    def forward(
        self,
        x: torch.Tensor,
        c_text: torch.Tensor,
        t: Optional[torch.Tensor] = None,
        valid_length: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        del valid_length
        if t is None:
            t = torch.ones([x.shape[0]], device=x.device, dtype=torch.float32)
        rep = self.backbone(x, t, c_text, classify_mode=True)
        pooled = self._pool_feature(rep)
        return self.cls_pred_branch(pooled)


# -----------------------------------------------------------------------------
# Hinge loss helpers (StyleGAN-T uses hinge loss in training/loss.py).
# -----------------------------------------------------------------------------
def discriminator_hinge_loss(real_logits: torch.Tensor, fake_logits: torch.Tensor):
    loss_real = F.relu(1.0 - real_logits).mean()
    loss_fake = F.relu(1.0 + fake_logits).mean()
    return loss_real, loss_fake


def generator_hinge_loss(fake_logits: torch.Tensor) -> torch.Tensor:
    return (-fake_logits).mean()