File size: 6,224 Bytes
47ea602 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """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]
|