digital-garbageandClaude Opus 5.5
ef84df4286
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
42 lines
1.8 KiB
Python
42 lines
1.8 KiB
Python
"""The SD1.5 module: recognising the file, and a pipeline that builds into a real graph."""
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
UNET = "model.diffusion_model.input_blocks.0.0.weight"
|
|
CLIP_L = "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def registered(comfyui):
|
|
import nodes as comfy_nodes
|
|
|
|
async def load():
|
|
await comfy_nodes.init_extra_nodes(init_custom_nodes=False)
|
|
await comfy_nodes.load_custom_node(str(Path(__file__).resolve().parents[4]), module_parent="custom_nodes")
|
|
|
|
asyncio.run(load())
|
|
return comfy_nodes
|
|
|
|
|
|
def test_sd1_is_recognised_and_sd2_sdxl_h3_are_not(comfyui):
|
|
from modules.models import sd15
|
|
assert sd15.detect({UNET, CLIP_L})
|
|
assert sd15.probe_traits({UNET, CLIP_L}) == ["spatial_latent", "predict_eps"]
|
|
assert sd15.detect({UNET, "cond_stage_model.transformer.embeddings.token_embedding.weight"}) # an older SD1.x file's CLIP naming
|
|
assert not sd15.detect({UNET, "cond_stage_model.model.token_embedding.weight"}) # SD2 (OpenCLIP)
|
|
assert not sd15.detect({UNET, "conditioner.embedders.0.transformer.text_model.embeddings.token_embedding.weight"}) # SDXL
|
|
assert not sd15.detect({"video_patch_proj.weight", "audio_patch_proj.weight"}) # H3
|
|
assert sd15.probe_traits({UNET}) == []
|
|
|
|
|
|
def test_the_pipeline_builds_a_real_graph_and_is_offered(registered):
|
|
from core import graph, routes
|
|
from modules.models.sd15 import text_to_image
|
|
|
|
_prompt, problems = graph.build(text_to_image())
|
|
assert [p for p in problems if "ckpt_name" not in p and "is '', which is not one of" not in p] == []
|
|
assert "sd15_text_to_image" in {p["id"] for _s, make in routes.modules().providers("pipeline_presets") for p in make()}
|