ThereforeGames
7bdbde715c
Merge serializable live UNet configuration into saved metadata so architecture settings such as MiniMax H3 PDD output-head counts survive reload. Add regression coverage and document the fix in the changelog and advanced usage guide.
116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
import importlib
|
|
import json
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
|
|
COMFY_ROOT = Path(__file__).resolve().parents[3]
|
|
CUSTOM_NODES_ROOT = COMFY_ROOT / "custom_nodes"
|
|
PACKAGE_NAME = Path(__file__).resolve().parents[1].name
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
sys.path.insert(0, str(CUSTOM_NODES_ROOT))
|
|
|
|
model_save = importlib.import_module(f"{PACKAGE_NAME}.int8_model_save")
|
|
|
|
|
|
class ModelSaveTests(unittest.TestCase):
|
|
def test_default_filename_uses_quantized_model_scope(self):
|
|
filename_input = model_save.INT8ModelSave.INPUT_TYPES()["required"]["filename_prefix"]
|
|
self.assertEqual(filename_input[1]["default"], "quantized_models/Quantized_Model")
|
|
|
|
def test_unquantized_model_is_rejected(self):
|
|
model_patcher = SimpleNamespace(
|
|
model=torch.nn.Sequential(torch.nn.Linear(4, 4)),
|
|
object_patches={},
|
|
object_patches_backup={},
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Place Enable Quantization on MODEL"):
|
|
model_save._validate_quantized_model(model_patcher)
|
|
|
|
def test_pending_toolkit_quantized_object_patch_is_accepted(self):
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module.weight = torch.nn.Parameter(
|
|
torch.ones((4, 4), dtype=torch.int8),
|
|
requires_grad=False,
|
|
)
|
|
quantized_module._is_quantized = True
|
|
quantized_module._quant_format = "int8_tensorwise"
|
|
model_patcher = SimpleNamespace(
|
|
model=torch.nn.Sequential(torch.nn.Linear(4, 4)),
|
|
object_patches={"diffusion_model.block.linear": quantized_module},
|
|
object_patches_backup={},
|
|
)
|
|
|
|
model_save._validate_quantized_model(model_patcher)
|
|
|
|
def test_pending_grouped_int8_object_patches_are_accepted(self):
|
|
for quant_format in ("asym_w4a8_int8", "w6a8_int8"):
|
|
with self.subTest(quant_format=quant_format):
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module._is_quantized = True
|
|
quantized_module._quant_format = quant_format
|
|
model_patcher = SimpleNamespace(
|
|
model=torch.nn.Sequential(torch.nn.Linear(4, 4)),
|
|
object_patches={"diffusion_model.block.linear": quantized_module},
|
|
object_patches_backup={},
|
|
)
|
|
|
|
model_save._validate_quantized_model(model_patcher)
|
|
|
|
def test_saved_metadata_preserves_live_model_configuration(self):
|
|
model_patcher = SimpleNamespace(
|
|
model=SimpleNamespace(
|
|
model_config=SimpleNamespace(
|
|
unet_config={
|
|
"image_model": "minimax_h3",
|
|
"latents_dim": 24,
|
|
"pdd_steps": 8,
|
|
"non_serializable": object(),
|
|
},
|
|
),
|
|
),
|
|
)
|
|
metadata = {
|
|
"config": json.dumps({
|
|
"vae": {"channels": 3},
|
|
"transformer": {"custom": "preserved", "pdd_steps": 1},
|
|
}),
|
|
}
|
|
|
|
model_save._merge_model_config_metadata(model_patcher, metadata)
|
|
config = json.loads(metadata["config"])
|
|
|
|
self.assertEqual(config["vae"], {"channels": 3})
|
|
self.assertEqual(config["transformer"]["custom"], "preserved")
|
|
self.assertEqual(config["transformer"]["image_model"], "minimax_h3")
|
|
self.assertEqual(config["transformer"]["latents_dim"], 24)
|
|
self.assertEqual(config["transformer"]["pdd_steps"], 8)
|
|
self.assertNotIn("non_serializable", config["transformer"])
|
|
|
|
def test_dynamic_lora_is_rejected(self):
|
|
quantized_module = torch.nn.Linear(4, 4, bias=False)
|
|
quantized_module.weight = torch.nn.Parameter(
|
|
torch.ones((4, 4), dtype=torch.int8),
|
|
requires_grad=False,
|
|
)
|
|
quantized_module._is_quantized = True
|
|
quantized_module._quant_format = "int8_tensorwise"
|
|
model_patcher = SimpleNamespace(
|
|
model=torch.nn.Sequential(torch.nn.Linear(4, 4)),
|
|
object_patches={"diffusion_model.block.linear": quantized_module},
|
|
object_patches_backup={},
|
|
get_attachment=lambda key: (("Dynamic", "example.safetensors", 1.0),),
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "cannot serialize Dynamic LoRAs"):
|
|
model_save._validate_quantized_model(model_patcher)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|