File size: 3,998 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 | """CPU check of rapid_moss.model against the release MOSSImageTransformer.
python tests/test_equiv.py MODEL_DIR LODESTONES_DIR
Runs everything in fp32 (rapid_moss.model.BF16 patched) so the logic must match to ~1e-5:
plain forward, negative-branch skip 22, the lodestones generator path (fp32 time path +
unmerged LoRA hooks, as its model card describes), discriminator taps and a gradient check.
"""
import sys
import time
import torch
from safetensors.torch import load_file
import rapid_moss.model as M
from rapid_moss.nets import Discriminator, FakeScore, Generator
torch.manual_seed(0)
model_dir, lode_dir = sys.argv[1], sys.argv[2]
M.BF16 = torch.float32
release = M.load_release_module(model_dir)
t0 = time.time()
tr = release.MOSSImageTransformer.from_pretrained(f'{model_dir}/transformer', torch_dtype=torch.float32,
local_files_only=True)
model = M.Moss(tr, release)
print(f'load {time.time() - t0:.1f}s')
h, w, L = 6, 10, 12
latent = torch.randn(1, 128, h, w)
emb = torch.randn(L, 4096) * 2
ctx = model.context(emb)
def rel(a, b):
return ((a - b).norm() / b.norm()).item()
with torch.no_grad():
for t in (999.0, 500.0, 37.0):
ref = tr(latent, t, emb[None]).sample.float()
got = model(latent, t, ctx)
print(f'forward t={t}: rel {rel(got, ref):.2e}')
assert rel(got, ref) < 1e-4
ref = tr(latent, 600.0, emb[None], skip_layers=[22]).sample.float()
got = model(latent, 600.0, ctx, skip=(22,))
print(f'skip22: rel {rel(got, ref):.2e} (vs no-skip {rel(got, model(latent, 600.0, ctx)):.2e})')
assert rel(got, ref) < 1e-4
# lodestones generator: release model + replaced time path + LoRA hooks (model card recipe)
gen = Generator(model, 128, 128, 4)
gen.load_lodestones(lode_dir)
te = load_file(f'{lode_dir}/time_embedding.safetensors')
lora = load_file(f'{lode_dir}/lora.safetensors')
import copy
tr2 = copy.deepcopy(tr)
tr2.load_state_dict(te, strict=False)
for i, layer in enumerate(tr2.layers):
for name, lin in layer.proj.items():
p = f'diffusion_model.layers.{i}.proj.{name}'
down, up, alpha = lora[p + '.lora_down.weight'], lora[p + '.lora_up.weight'], lora[p + '.alpha']
scale = float(alpha) / down.shape[0]
lin.register_forward_hook(lambda m, inp, out, d=down, u=up, s=scale: out + s * (inp[0] @ d.t()) @ u.t())
Q = load_file(f'{lode_dir}/dataspace.safetensors')['latent_mix.Q']
codes = load_file(f'{lode_dir}/dataspace.safetensors')['step_codes']
z = torch.randn(1, 128, h, w)
x0_ref = torch.zeros_like(z)
x0_got = torch.zeros_like(z)
for r in range(3):
x = z + torch.einsum('oc,bchw->bohw', Q, x0_ref)
x0_ref = z - tr2(x, codes[r], emb[None]).sample.float()
x0_got = gen.step(model, z, x0_got, r, ctx)
print(f'generator r={r}: rel {rel(x0_got, x0_ref):.2e}')
assert rel(x0_got, x0_ref) < 1e-4
del tr2
disc = Discriminator(model.dim, (5, 11), 256)
feats = model(latent, 300.0, ctx, taps=disc.taps)
print('taps', {k: tuple(v.shape) for k, v in feats.items()}, 'logits', tuple(disc(feats, h, w).shape))
# gradients: generator step through checkpointed blocks, fake score, D input-gradient
fake = FakeScore(model, 16, 16, 4)
x0 = gen.step(model, z, torch.randn_like(z), 1, ctx, ckpt=True)
x0.backward(torch.randn_like(x0))
nz = {n: p.grad.abs().sum().item() > 0 for n, p in gen.named_parameters()}
print('gen grads nonzero:', all(nz.values()), [n for n, v in nz.items() if not v][:5])
assert all(nz.values())
v = fake(model, latent, 400.0, ctx, 2, ckpt=True)
v.square().mean().backward()
print('fake r_emb grad rows nonzero:', (fake.r_emb.grad.abs().sum(1) > 0).tolist())
leaf = latent.clone().requires_grad_(True)
g, = torch.autograd.grad(-disc(model(leaf, 300.0, ctx, taps=disc.taps, ckpt=True), h, w).mean(), leaf)
print('D input grad norm', g.norm().item())
assert g.norm() > 0
print('OK')
|