code/rapid_moss/nets.py

"""Trainable parts: data-space generator, fake score, DMDX discriminator heads; lodestones-format I/O."""
from __future__ import annotations

from pathlib import Path

import torch
from safetensors.torch import load_file, save_file
from torch import nn
from torch.nn import functional as F

from .model import PROJ, LoRA, TimePath


class Generator(nn.Module):
    """Data-space iteration (arXiv 2610.12102) in the lodestones MOSS form:
    x = z + Q x0 (channel mix), v = net(x, code_r), x0' = z - v, x0 starts at zero."""

    def __init__(self, model, rank, alpha, steps):
        super().__init__()
        self.lora = LoRA(model.dims, model.depth, rank, alpha)
        self.time = TimePath(model.base_time.state(), trainable=True)
        channels = model.tr.config.latent_channels
        self.Q = nn.Parameter(torch.zeros(channels, channels))
        self.register_buffer('codes', torch.full((steps,), 1000.0))
        self.steps = steps

    def step(self, model, z, x0, r, ctx, ckpt=False):
        x = z + torch.einsum('oc,bchw->bohw', self.Q, x0)
        return z - model(x, self.codes[r], ctx, lora=self.lora, tp=self.time, ckpt=ckpt)

    @torch.no_grad()
    def rollout(self, model, z, ctx, n):
        """n generator calls from x0 = 0 (no grad); returns x0 after n steps."""
        x0 = torch.zeros_like(z)
        for r in range(n):
            x0 = self.step(model, z, x0, r, ctx)
        return x0

    @torch.no_grad()
    def load_lodestones(self, path):
        path = Path(path)
        ds = load_file(path / 'dataspace.safetensors')
        self.Q.copy_(ds['latent_mix.Q'])
        self.codes.copy_(ds['step_codes'][:self.steps])
        te = load_file(path / 'time_embedding.safetensors')
        for k in TimePath.NAMES:
            self.time[k].copy_(te[k])
        lora = load_file(path / 'lora.safetensors')
        for i in range(self.lora.A[PROJ[0]].shape[0]):
            for name in PROJ:
                p = f'diffusion_model.layers.{i}.proj.{name}'
                down, up = lora[p + '.lora_down.weight'], lora[p + '.lora_up.weight']
                alpha = float(lora[p + '.alpha'])
                if down.shape[0] != self.lora.rank:
                    raise ValueError(f'rank mismatch: file {down.shape[0]} vs {self.lora.rank}')
                # fold the file's alpha/rank into B so our own scale applies unchanged
                self.lora.A[name][i].copy_(down)
                self.lora.B[name][i].copy_(up * (alpha / down.shape[0]) / self.lora.scale)

    @torch.no_grad()
    def export(self, out_dir, meta=None):
        """Write lodestones-compatible files (kohya LoRA, fp32 time path, Q + step codes)."""
        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)
        rank = self.lora.rank
        lora = {}
        for i in range(self.lora.A[PROJ[0]].shape[0]):
            for name in PROJ:
                p = f'diffusion_model.layers.{i}.proj.{name}'
                lora[p + '.lora_down.weight'] = self.lora.A[name][i].float().contiguous()
                lora[p + '.lora_up.weight'] = (self.lora.B[name][i].float() * self.lora.scale).contiguous()
                lora[p + '.alpha'] = torch.tensor(float(rank))
        info = {'rank': str(rank), 'alpha': f'{float(rank)}', 'steps': str(self.steps), **(meta or {})}
        save_file(lora, out / 'lora.safetensors', metadata=info)
        save_file({k: v.float().contiguous() for k, v in self.time.state().items()}, out / 'time_embedding.safetensors')
        save_file({'latent_mix.Q': self.Q.float().contiguous(), 'step_codes': self.codes.float().contiguous()},
                  out / 'dataspace.safetensors', metadata=info)


class FakeScore(nn.Module):
    """Generator-distribution velocity: base + LoRA, conditioned on the refinement index r."""

    def __init__(self, model, rank, alpha, steps):
        super().__init__()
        self.lora = LoRA(model.dims, model.depth, rank, alpha)
        self.r_emb = nn.Parameter(torch.zeros(steps, model.dim))

    def __call__(self, model, x, t, ctx, r, ckpt=False):
        return model(x, t, ctx, lora=self.lora, time_extra=self.r_emb[r][None], ckpt=ckpt)


class Head(nn.Module):
    """DMDX 2D head: 4x4 stride-2 conv, GroupNorm(32), SiLU, then 1x1 to patch logits."""

    def __init__(self, dim, hidden):
        super().__init__()
        self.conv = nn.Conv2d(dim, hidden, 4, 2, 1)
        self.norm = nn.GroupNorm(32, hidden)
        self.out = nn.Conv2d(hidden, 1, 1)

    def forward(self, f):
        return self.out(F.silu(self.norm(self.conv(f))))


class Discriminator(nn.Module):
    """Trainable heads on frozen-teacher block features (backbone = base, no LoRA)."""

    def __init__(self, dim, taps, hidden):
        super().__init__()
        self.taps = tuple(sorted(taps))
        self.heads = nn.ModuleList([Head(dim, hidden) for _ in self.taps])

    def forward(self, feats, h, w):
        logits = []
        for i, head in zip(self.taps, self.heads):
            f = feats[i].float()
            f = f * torch.rsqrt(f.square().mean(-1, keepdim=True) + 1e-6)
            logits.append(head(f.transpose(1, 2).reshape(1, -1, h, w)).flatten(1))
        return torch.cat(logits, 1)


class Critic(nn.Module):
    """DMAD critic (arXiv 2610.02188), teacher head only: base + LoRA on blocks [start:], a conv head on the
    last block's image tokens giving one logit per sample (teacher sample vs generator sample)."""

    def __init__(self, model, rank, alpha, start, hidden):
        super().__init__()
        self.lora = LoRA(model.dims, model.depth, rank, alpha, start=start)
        self.last = model.depth - 1
        self.head = nn.Sequential(nn.Conv2d(model.dim, hidden, 4, 2, 1), nn.GroupNorm(32, hidden), nn.SiLU(),
                                  nn.Conv2d(hidden, hidden, 4, 2, 1), nn.GroupNorm(32, hidden), nn.SiLU())
        self.out = nn.Linear(hidden, 1)

    def forward(self, model, x, t, ctx, ckpt=False):
        _, _, h, w = x.shape
        f = model(x, t, ctx, lora=self.lora, taps=(self.last,), ckpt=ckpt)[self.last].float()
        f = f * torch.rsqrt(f.square().mean(-1, keepdim=True) + 1e-6)
        f = self.head(f.transpose(1, 2).reshape(1, -1, h, w))
        return self.out(f.mean((2, 3)))[:, 0]