Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
"""Every weight dtype, compute dtype and attention the Diffusion Model Loader offers, on a real (tiny) H3 file.
|
|
|
|
The file is written from ComfyUI's own `model_base.MiniMaxH3` at a few thousand parameters, so ComfyUI's real
|
|
detection and loading path reads it. Each combination must come back as a model, not an exception.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def tiny_file(comfyui, tmp_path_factory):
|
|
import comfy.supported_models as sm
|
|
import folder_paths
|
|
import safetensors.torch as st
|
|
torch.manual_seed(0)
|
|
config = sm.MiniMaxH3(dict(
|
|
image_model="minimax_h3", hidden_size=64, num_layers=4, token_refiner_num_layers=1, num_attention_heads=2,
|
|
attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64,
|
|
time_embed_dim=32, rope_inv_freq_len=4))
|
|
config.set_inference_dtype(torch.float32, None)
|
|
base = config.get_model({}, device=torch.device("cpu"))
|
|
for p in base.diffusion_model.parameters():
|
|
torch.nn.init.normal_(p, std=0.05)
|
|
folder = tmp_path_factory.mktemp("unet")
|
|
st.save_file({k: v.contiguous() for k, v in base.diffusion_model.state_dict().items()}, str(folder / "tiny_h3.safetensors"))
|
|
folder_paths.add_model_folder_path("diffusion_models", str(folder))
|
|
return "tiny_h3.safetensors"
|
|
|
|
|
|
def _choices():
|
|
from modules.loaders.common import COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices
|
|
return list(WEIGHT_DTYPES), list(COMPUTE_DTYPES), list(attention_choices())
|
|
|
|
|
|
def pytest_generate_tests(metafunc):
|
|
if "combo" in metafunc.fixturenames:
|
|
w, c, a = _choices() if _importable() else ([], [], [])
|
|
metafunc.parametrize("combo", [(x, y, z) for x in w for y in c for z in a], ids=lambda t: "-".join(t))
|
|
|
|
|
|
def _importable():
|
|
try:
|
|
from modules.loaders.common import COMPUTE_DTYPES # noqa: F401
|
|
return True
|
|
except Exception: # noqa: BLE001 -- collection must not fail without ComfyUI
|
|
return False
|
|
|
|
|
|
def test_every_dtype_and_attention_combination_loads(tiny_file, combo):
|
|
from modules.loaders.diffusion_model.nodes import FunUIDiffusionModelLoader
|
|
weight, compute, attention = combo
|
|
model, status = FunUIDiffusionModelLoader.execute(tiny_file, weight, compute, attention).result
|
|
assert type(model.model).__name__ == "MiniMaxH3"
|
|
assert weight in status
|