File size: 5,869 Bytes
1ec327b | 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 | """COMFYUI_ROOT=/path/to/ComfyUI python -m pytest code/tests/test_comfy_adapter.py"""
import importlib.util
import os
from pathlib import Path
import sys
import pytest
import torch
from safetensors.torch import save_file
sys.path.insert(0, os.environ['COMFYUI_ROOT'])
from comfy.cli_args import args
args.cpu = True
import comfy.ops
from comfy.ldm.moss_image.model import MOSSImage
from comfy.model_patcher import ModelPatcher
spec = importlib.util.spec_from_file_location('rapid_nodes', Path(__file__).parents[2] / 'comfyui/rapid-moss/nodes.py')
rapid = importlib.util.module_from_spec(spec)
spec.loader.exec_module(rapid)
@pytest.fixture
def adapter(tmp_path, monkeypatch):
torch.manual_seed(123)
dim, ffn, rank = 128, 192, 3
lora = {}
for name in rapid.PROJ:
inp, out = (ffn, dim) if name == 'ffn_down' else (dim, 2 * ffn if name == 'ffn_up_gate' else dim)
prefix = 'diffusion_model.layers.0.proj.' + name
lora[prefix + '.lora_down.weight'] = torch.randn(rank, inp) * 0.03
lora[prefix + '.lora_up.weight'] = torch.randn(out, rank) * 0.03
lora[prefix + '.alpha'] = torch.tensor(float(rank))
time = {
'time_embedding.0.weight': torch.randn(dim, 256) * .01,
'time_embedding.0.bias': torch.randn(dim) * .01,
'time_embedding.2.weight': torch.randn(dim, dim) * .01,
'time_embedding.2.bias': torch.randn(dim) * .01,
'shared_block_modulation.bias': torch.randn(6 * dim) * .01,
'video_head_modulation.bias': torch.randn(2 * dim) * .01,
}
save_file(lora, tmp_path / 'lora.safetensors')
save_file(time, tmp_path / 'time_embedding.safetensors')
save_file({'latent_mix.Q': torch.eye(128), 'step_codes': torch.tensor([1000, 968, 926, 871])},
tmp_path / 'dataspace.safetensors')
monkeypatch.setattr(rapid, '_adapter_dirs', lambda: {'test': str(tmp_path)})
return lora
def make_model(fused):
holder = torch.nn.Module()
holder.diffusion_model = MOSSImage(dim=128, ffn_dim=192, num_heads=1, num_layers=1,
fused_qkv=fused, dtype=torch.float32, operations=comfy.ops.manual_cast)
for p in holder.parameters():
torch.nn.init.normal_(p, std=.02)
return ModelPatcher(holder, torch.device('cpu'), torch.device('cpu'))
@pytest.mark.parametrize('fused', [False, True])
def test_lora_apply_and_restore(adapter, fused):
model = make_model(fused)
layer = model.model.diffusion_model.layers[0]
base_weights = {k: m.weight.detach().clone() for k, m in layer.proj.items()}
patched, _, sigmas = rapid.RapidMOSSAdapter().apply(model, 'test', 4, .7)
assert sigmas.tolist() == pytest.approx([1, .968, .926, .871, 0])
assert not model.object_patches
patched.patch_model(load_weights=False)
for name, module in layer.proj.items():
x = torch.randn(1, 5, module.weight.shape[1])
expected = torch.nn.functional.linear(x, base_weights[name])
names = ('q', 'k', 'v') if name == 'qkv' else (name,)
deltas = []
for proj in names:
prefix = 'diffusion_model.layers.0.proj.' + proj
deltas.append(.7 * (x @ adapter[prefix + '.lora_down.weight'].T) @ adapter[prefix + '.lora_up.weight'].T)
expected += torch.cat(deltas, dim=-1)
torch.testing.assert_close(module(x), expected)
torch.testing.assert_close(module.weight, base_weights[name], rtol=0, atol=0)
patched.unpatch_model()
for name, module in layer.proj.items():
x = torch.randn(1, 2, module.weight.shape[1])
torch.testing.assert_close(module(x), torch.nn.functional.linear(x, base_weights[name]))
def test_sampler_matches_dataspace_recurrence():
torch.manual_seed(4)
z = torch.randn(1, 3, 2, 2)
q = torch.randn(3, 3) * .1
codes = torch.tensor([1000, 968, 926, 871])
sigmas = torch.cat((codes.float() / 1000, torch.zeros(1)))
calls = []
def denoise(x, sigma):
calls.append(sigma.clone())
return x - sigma[:, None, None, None] * (x * .3 + sigma[:, None, None, None])
expected = torch.zeros_like(z)
for code in codes:
x = z + torch.einsum('oc,bchw->bohw', q, expected)
expected = z - (x * .3 + code / 1000)
actual = rapid.sample_dataspace(denoise, z, sigmas, Q=q, codes=codes)
torch.testing.assert_close(actual, expected)
assert len(calls) == 4
def test_reject_too_many_steps(adapter):
with pytest.raises(ValueError, match='4 step codes'):
rapid.RapidMOSSAdapter().apply(make_model(True), 'test', 5, 1.)
@pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
@pytest.mark.parametrize('device', ['cpu', pytest.param('cuda', marks=pytest.mark.skipif(
os.environ.get('RAPID_MOSS_TEST_CUDA') != '1', reason='set RAPID_MOSS_TEST_CUDA=1 for GPU checks'))])
def test_native_forward_fused_matches_split(adapter, dtype, device):
model = make_model(False)
fused = make_model(True)
sd = model.model.state_dict()
prefix = 'diffusion_model.layers.0.proj.'
sd[prefix + 'qkv.weight'] = torch.cat([sd.pop(prefix + n + '.weight') for n in ('q', 'k', 'v')])
fused.model.load_state_dict(sd)
x = torch.randn(1, 128, 2, 3, device=device, dtype=dtype)
ctx = torch.randn(1, 4, 4096, device=device, dtype=dtype)
t = torch.tensor([.968], device=device)
outputs = []
for m in (model, fused):
m.model.to(device=device, dtype=dtype)
patched, _, _ = rapid.RapidMOSSAdapter().apply(m, 'test', 4, 1.)
patched.patch_model(load_weights=False)
with torch.no_grad():
outputs.append(m.model.diffusion_model(x.clone(), t, ctx))
patched.unpatch_model()
assert all(torch.isfinite(out).all() for out in outputs)
assert all(out.abs().max() < 10 for out in outputs)
torch.testing.assert_close(*outputs, rtol=.01, atol=.001)
|