77 lines
4.0 KiB
Python
77 lines
4.0 KiB
Python
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
spec = importlib.util.spec_from_file_location("fl_lora", Path(__file__).parents[1] / "nodes/utility/FL_ModelDifferenceLoraSave.py")
|
|
m = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(m)
|
|
|
|
|
|
def model(weight, bias=None):
|
|
network = torch.nn.Module()
|
|
network.diffusion_model = torch.nn.Module()
|
|
network.diffusion_model.linear = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=bias is not None)
|
|
network.diffusion_model.linear.weight = torch.nn.Parameter(weight.clone())
|
|
if bias is not None:
|
|
network.diffusion_model.linear.bias = torch.nn.Parameter(bias.clone())
|
|
return SimpleNamespace(model=network, backup={}, hook_backup={}, patches={})
|
|
|
|
|
|
class ModelDifferenceLoraTests(unittest.TestCase):
|
|
def test_export_roundtrip_sign_bias_and_no_input_mutation(self):
|
|
base = model(torch.randn(4, 5), torch.randn(4))
|
|
tuned = model(base.model.diffusion_model.linear.weight.detach() + torch.randn(4, 5),
|
|
base.model.diffusion_model.linear.bias.detach() + 0.25)
|
|
before = tuned.model.diffusion_model.linear.weight.detach().clone()
|
|
with TemporaryDirectory() as root, patch.object(m.folder_paths, "get_folder_paths", return_value=[root]):
|
|
path, = m.FL_ModelDifferenceLoraSave().save(tuned, base, "Krea2/test", 4, "cpu")
|
|
sd = load_file(path)
|
|
key = "diffusion_model.linear"
|
|
delta = sd[key + ".lora_up.weight"].float() @ sd[key + ".lora_down.weight"].float()
|
|
torch.testing.assert_close(base.model.diffusion_model.linear.weight + delta, before, atol=.003, rtol=.003)
|
|
torch.testing.assert_close(sd[key + ".bias.diff"].float(), torch.full((4,), .25))
|
|
key_map = {key: key + ".weight", key + ".bias": key + ".bias"}
|
|
patches = m.comfy.lora.load_lora(sd, key_map)
|
|
self.assertEqual(set(patches), set(key_map.values()))
|
|
with safe_open(path, framework="pt") as f:
|
|
self.assertAlmostEqual(json.loads(f.metadata()["fl_extraction"])["retained_weight_energy"], 1, places=5)
|
|
second, = m.FL_ModelDifferenceLoraSave().save(tuned, base, "Krea2/test", 4, "cpu")
|
|
self.assertNotEqual(path, second)
|
|
torch.testing.assert_close(tuned.model.diffusion_model.linear.weight, before)
|
|
|
|
def test_lowrank_and_convolution_reconstruction(self):
|
|
for shape in ((96, 80), (96, 5, 4, 4)):
|
|
diff = (torch.randn(96, 3) @ torch.randn(3, 80)).reshape(shape)
|
|
up, down = m.factorize_difference(diff, 3)
|
|
torch.testing.assert_close(up.flatten(1) @ down.flatten(1), diff.flatten(1), atol=1e-4, rtol=1e-4)
|
|
|
|
def test_patches_and_backup_are_applied_before_subtraction(self):
|
|
target = model(torch.full((3, 4), 99.))
|
|
key = "diffusion_model.linear.weight"
|
|
target.backup[key] = SimpleNamespace(weight=torch.ones(3, 4))
|
|
target.patches[key] = [(0.5, (torch.full((3, 4), 4.),), 1., None, None)]
|
|
torch.testing.assert_close(m.materialize_weight(target, key, torch.device("cpu")), torch.full((3, 4), 3.))
|
|
|
|
def test_invalid_models_and_paths_do_not_save(self):
|
|
base = model(torch.zeros(4, 5))
|
|
with TemporaryDirectory() as root, patch.object(m.folder_paths, "get_folder_paths", return_value=[root]):
|
|
for tuned, prefix in ((base, "same"), (model(torch.zeros(3, 5)), "shape"),
|
|
(model(torch.full((4, 5), float("nan"))), "nan"),
|
|
(base, "../escape"), (base, "C:/escape")):
|
|
with self.subTest(prefix=prefix), self.assertRaises(ValueError):
|
|
m.FL_ModelDifferenceLoraSave().save(tuned, base, prefix, 4, "cpu")
|
|
self.assertEqual(list(Path(root).rglob("*.safetensors")), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|