DemonGatanjieuandClaude Opus 5.5
f65f12152b
Checked against the headers of 315 models in ComfyUI-Manager's list: a Hunyuan-DiT checkpoint and LBM were read as VAEs, SUPIR as a ControlNet, XLabs' Flux LoRAs and IP-Adapter as diffusion models, a whole CLIP as a text encoder only; also PhotoMaker, LoRAs of weight differences or OFT, and checkpoints without a VAE. Files two loaders read from different folders (SUPIR, a whole CLIP) are not sure. A lone drawing model in checkpoints loads there (the checkpoint loader reads it as a diffusion model). Newly told: diffusers UNets, Stable Cascade, SDXL T2I-Adapters, RCAN/SPAN/compact upscalers. 301 of 315 now agree. "Looks like no drawing model" said what it could not tell: now "Kind not told". Kept header answers carry the rules version and are read again when it changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
378 lines
25 KiB
Python
378 lines
25 KiB
Python
"""Model import: telling what a model file is from its header (model_kind.py), finding the
|
|
dropped file on this computer, and putting it into the models folder without replacing
|
|
anything (api/model_import.py). No network; files go to temporary folders."""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
import model_identity # noqa: E402
|
|
import model_kind # noqa: E402
|
|
from api import model_import # noqa: E402
|
|
|
|
|
|
def tensors(*names, shapes=None, meta=None):
|
|
header = {name: {"dtype": "F16", "shape": (shapes or {}).get(name, [4, 4]), "data_offsets": [0, 0]} for name in names}
|
|
if meta:
|
|
header["__metadata__"] = meta
|
|
return header
|
|
|
|
|
|
def safetensors_bytes(header):
|
|
raw = json.dumps(header).encode("utf-8")
|
|
return struct.pack("<Q", len(raw)) + raw + b"\0" * 64
|
|
|
|
|
|
def gguf_bytes(arch, names=()):
|
|
def text(value):
|
|
raw = value.encode("utf-8")
|
|
return struct.pack("<Q", len(raw)) + raw
|
|
data = b"GGUF" + struct.pack("<I", 3) + struct.pack("<Q", len(names)) + struct.pack("<Q", 2)
|
|
data += text("general.architecture") + struct.pack("<I", 8) + text(arch)
|
|
data += text("tokenizer.tokens") + struct.pack("<I", 9) + struct.pack("<I", 8) + struct.pack("<Q", 2) + text("a") + text("b")
|
|
for name in names:
|
|
data += text(name) + struct.pack("<I", 2) + struct.pack("<QQ", 4, 4) + struct.pack("<I", 1) + struct.pack("<Q", 0)
|
|
return data
|
|
|
|
|
|
def kind_of(header):
|
|
return model_kind.classify("safetensors", header)
|
|
|
|
|
|
class KindTests(unittest.TestCase):
|
|
def test_loras_and_their_base_models(self):
|
|
sdxl = kind_of(tensors("lora_unet_input_blocks_4_1_proj_in.lora_down.weight", "lora_te1_text_model_encoder_layers_0_mlp_fc1.alpha"))
|
|
self.assertEqual((sdxl["kind"], sdxl["base"], sdxl["sure"]), ("loras", "SDXL", True))
|
|
# Diffusers names say nothing: the cross-attention width does (768 = SD 1.5).
|
|
sd15 = kind_of(tensors("unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.lora_A.weight",
|
|
shapes={"unet.down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.lora_A.weight": [16, 768]}))
|
|
self.assertEqual((sd15["kind"], sd15["base"]), ("loras", "SD 1.5"))
|
|
flux = kind_of(tensors("lora_unet_double_blocks_0_img_attn_proj.lora_up.weight"))
|
|
self.assertEqual((flux["kind"], flux["base"]), ("loras", "Flux.1 D"))
|
|
# Anima (Cosmos-Predict2 with a Qwen3 adapter): Civitai calls it a Checkpoint, it goes in diffusion models.
|
|
anima = kind_of(tensors("net.blocks.0.self_attn.q_proj.weight", "net.blocks.0.cross_attn.k_proj.weight",
|
|
"net.x_embedder.proj.1.weight", "net.llm_adapter.blocks.0.mlp.0.weight"))
|
|
self.assertEqual((anima["kind"], anima["base"], anima["sure"]), ("diffusion_models", "Anima", True))
|
|
bare = kind_of(tensors("blocks.0.self_attn.q_proj.weight", "blocks.0.cross_attn.k_proj.weight",
|
|
"blocks.0.adaln_modulation_self_attn.1.weight", "llm_adapter.blocks.0.mlp.0.weight"))
|
|
self.assertEqual((bare["kind"], bare["base"]), ("diffusion_models", "Anima"))
|
|
wan = kind_of(tensors("diffusion_model.blocks.0.self_attn.q.lora_A.weight", "diffusion_model.blocks.0.cross_attn.k.lora_B.weight"))
|
|
self.assertEqual(wan["base"], "Wan Video")
|
|
# SDXL's families only show in what it was trained on.
|
|
illustrious = kind_of(tensors("lora_unet_output_blocks_0_1_proj_in.lora_down.weight",
|
|
meta={"ss_base_model_version": "sdxl_base_v1-0", "ss_sd_model_name": "illustriousXL_v01.safetensors"}))
|
|
self.assertEqual(illustrious["base"], "Illustrious")
|
|
lycoris = kind_of(tensors("lora_unet_down_blocks_0_attentions_0_proj_in.hada_w1_a"))
|
|
self.assertEqual(lycoris["kind"], "loras")
|
|
|
|
def test_whole_models_and_their_parts(self):
|
|
checkpoint = kind_of(tensors("model.diffusion_model.input_blocks.0.0.weight", "model.diffusion_model.label_emb.0.0.weight",
|
|
"first_stage_model.encoder.down.0.block.0.conv1.weight", "conditioner.embedders.1.model.ln_final.weight"))
|
|
self.assertEqual((checkpoint["kind"], checkpoint["base"]), ("checkpoints", "SDXL"))
|
|
merged = kind_of(tensors("model.diffusion_model.input_blocks.0.0.weight", "first_stage_model.decoder.up.0.block.0.conv1.weight",
|
|
"lora_te_text_model_encoder_layers_0_mlp_fc1.lora_down.weight"))
|
|
self.assertEqual(merged["kind"], "checkpoints") # a merge that kept a LoRA's keys
|
|
flux = kind_of(tensors("double_blocks.0.img_attn.qkv.weight", "img_in.weight", "single_blocks.0.linear1.weight"))
|
|
self.assertEqual((flux["kind"], flux["base"]), ("diffusion_models", "Flux.1 D"))
|
|
# Qwen-Image also has img_in: not taken for Flux.
|
|
qwen = kind_of(tensors("transformer_blocks.0.img_mod.1.weight", "img_in.weight", "txt_norm.weight"))
|
|
self.assertEqual((qwen["kind"], qwen["base"]), ("diffusion_models", "Qwen"))
|
|
vae = kind_of(tensors("encoder.down.0.block.0.conv1.weight", "decoder.up.0.block.0.conv1.weight"))
|
|
self.assertEqual((vae["kind"], vae["base"], vae["sure"]), ("vae", "", True))
|
|
t5 = kind_of(tensors("encoder.block.0.layer.0.SelfAttention.q.weight", "shared.weight"))
|
|
self.assertEqual(t5["kind"], "text_encoders")
|
|
clip = kind_of(tensors("text_model.encoder.layers.0.mlp.fc1.weight"))
|
|
self.assertEqual(clip["kind"], "text_encoders")
|
|
vision = kind_of(tensors("vision_model.encoder.layers.0.mlp.fc1.weight"))
|
|
self.assertEqual(vision["kind"], "clip_vision")
|
|
control = kind_of(tensors("control_model.input_hint_block.0.weight", "control_model.zero_convs.0.0.weight"))
|
|
self.assertEqual(control["kind"], "controlnet")
|
|
# Model patches: ControlNets by name, but only "Load Model Patch" reads them (models/model_patches).
|
|
lllite = kind_of(tensors("lllite_conditioning1.conv1.weight", "lllite_modules.0.down.weight"))
|
|
self.assertEqual(lllite, {"kind": "model_patches", "base": "", "sure": True, "form": "lllite"})
|
|
blockwise = kind_of(tensors("controlnet_blocks.0.y_rms.weight", "controlnet_blocks.0.img_mod.weight", "img_in.weight"))
|
|
self.assertEqual((blockwise["kind"], blockwise["form"]), ("model_patches", ""))
|
|
# An IP-Adapter carrying LoRA layers of its own (the Anima IP-Adapter) is no LoRA.
|
|
anima_ip = kind_of(tensors("blocks.0.adaln_ip.1.weight", "blocks.0.ip_k_proj.weight",
|
|
"lora.base_model.model.blocks.0.cross_attn.k_proj.lora_A.weight"))
|
|
self.assertEqual(anima_ip["kind"], "ipadapter")
|
|
self.assertEqual(kind_of(tensors("image_proj.proj.weight", "ip_adapter.1.to_k_ip.weight"))["kind"], "ipadapter")
|
|
|
|
def test_layouts_once_taken_for_another_kind(self):
|
|
# Found against ComfyUI-Manager's model list (real headers) and the loaders' own formats.
|
|
hunyuan_dit = kind_of(tensors("model.blocks.0.attn1.Wqkv.weight", "text_encoders.mt5xl.transformer.shared.weight",
|
|
"vae.encoder.down.0.block.0.conv1.weight"))
|
|
self.assertEqual(hunyuan_dit["kind"], "checkpoints") # not its VAE
|
|
no_vae = kind_of(tensors("model.diffusion_model.input_blocks.0.0.weight",
|
|
"cond_stage_model.transformer.text_model.encoder.layers.0.mlp.fc1.weight"))
|
|
self.assertEqual(no_vae["kind"], "checkpoints") # saved without its VAE: still a checkpoint
|
|
comfy_lora = kind_of(tensors("diffusion_model.input_blocks.0.0.lora_up.weight",
|
|
"text_encoders.clip_l.transformer.text_model.encoder.layers.0.mlp.fc1.lora_up.weight"))
|
|
self.assertEqual(comfy_lora["kind"], "loras")
|
|
lbm = kind_of(tensors("denoiser.down_blocks.0.resnets.0.conv1.weight", "vae.vae_model.encoder.down_blocks.0.resnets.0.conv1.weight"))
|
|
self.assertEqual((lbm["kind"], lbm["sure"]), ("", False)) # a VAE inside another model: not a VAE file
|
|
supir = kind_of(tensors("model.control_model.input_hint_block.0.weight", "model.diffusion_model.input_blocks.0.0.weight",
|
|
"first_stage_model.denoise_encoder.conv_in.weight"))
|
|
self.assertEqual((supir["kind"], supir["sure"]), ("model_patches", False)) # read from model_patches or checkpoints
|
|
whole_clip = kind_of(tensors("vision_model.encoder.layers.0.mlp.fc1.weight", "text_model.encoder.layers.0.mlp.fc1.weight"))
|
|
self.assertEqual((whole_clip["kind"], whole_clip["sure"]), ("clip_vision", False))
|
|
xlabs_lora = kind_of(tensors("double_blocks.0.processor.proj_lora1.down.weight", "double_blocks.0.processor.qkv_lora1.up.weight"))
|
|
self.assertEqual(xlabs_lora["kind"], "loras")
|
|
xlabs_ip = kind_of(tensors("double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight", "ip_adapter_proj_model.proj.weight"))
|
|
self.assertEqual(xlabs_ip["kind"], "ipadapter")
|
|
for names in (("diffusion_model.blocks.0.self_attn.q.diff", "diffusion_model.blocks.0.self_attn.q.diff_b"),
|
|
("lora_unet_down_blocks_0_attentions_0_proj_in.oft_blocks",),
|
|
("lora_unet_x.a1.weight", "lora_unet_x.b1.weight")): # weight differences, OFT, GLoRA
|
|
self.assertEqual(kind_of(tensors(*names))["kind"], "loras", names)
|
|
photomaker = kind_of(tensors("id_encoder.vision_model.encoder.layers.0.mlp.fc1.weight", "lora_weights.unet.x.lora.down.weight"))
|
|
self.assertEqual(photomaker["kind"], "")
|
|
hunyuan_video = kind_of(tensors("model.model.double_blocks.0.img_attn.qkv.weight", "model.model.txt_in.input_embedder.weight"))
|
|
self.assertEqual(hunyuan_video["kind"], "diffusion_models")
|
|
|
|
def test_kinds_once_not_told(self):
|
|
# Told now, rather than "kind not told" (which the scan page offers to skip).
|
|
unet = ("conv_in.weight", "down_blocks.0.resnets.0.conv1.weight", "mid_block.attentions.0.proj_in.weight",
|
|
"up_blocks.0.resnets.0.conv1.weight", "time_embedding.linear_1.weight") # diffusers: IC-Light, Kolors
|
|
self.assertEqual(kind_of(tensors(*unet))["kind"], "diffusion_models")
|
|
motion = kind_of(tensors("down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight",
|
|
"mid_block.motion_modules.0.temporal_transformer.proj_in.weight", "time_embedding.linear_1.weight"))
|
|
self.assertEqual(motion["kind"], "") # AnimateDiff, not a diffusion model
|
|
vae = kind_of(tensors("encoder.down_blocks.0.resnets.0.conv1.weight", "decoder.mid_block.resnets.0.conv1.weight"))
|
|
self.assertEqual(vae["kind"], "vae")
|
|
cascade = kind_of(tensors("down_blocks.0.0.channelwise.0.weight", "down_downscalers.1.1.weight", "clf.1.weight"))
|
|
self.assertEqual(cascade["kind"], "diffusion_models")
|
|
self.assertEqual(kind_of(tensors("adapter.body.0.resnets.0.block1.weight", "adapter.conv_in.weight"))["kind"], "controlnet")
|
|
for names in (("head.0.weight", "body.0.body.0.body.0.weight", "tail.0.0.weight"), # RCAN
|
|
("conv_1.sk.weight", "block_1.c1_r.sk.weight"), # SPAN
|
|
("body.0.weight", "body.1.weight", "body.2.bias")): # compact Real-ESRGAN
|
|
found = kind_of(tensors(*names))
|
|
self.assertEqual((found["kind"], found["sure"]), ("upscale_models", False), names)
|
|
|
|
def test_embeddings(self):
|
|
self.assertEqual(kind_of(tensors("clip_l", "clip_g")), {"kind": "embeddings", "base": "SDXL", "sure": True, "form": ""})
|
|
sd15 = kind_of(tensors("emb_params", shapes={"emb_params": [8, 768]}))
|
|
self.assertEqual((sd15["kind"], sd15["base"]), ("embeddings", "SD 1.5"))
|
|
|
|
def test_unknown_files_are_not_guessed(self):
|
|
self.assertEqual(kind_of(tensors("something.weight")), {"kind": "", "base": "", "sure": False, "form": ""})
|
|
self.assertEqual(model_kind.classify(None, None), {"kind": "", "base": "", "sure": False})
|
|
|
|
def test_gguf(self):
|
|
flux = model_kind.classify("gguf", model_kind.parse_gguf(gguf_bytes("flux", ["double_blocks.0.img_attn.qkv.weight"])))
|
|
self.assertEqual((flux["kind"], flux["base"]), ("diffusion_models", "Flux.1 D"))
|
|
self.assertEqual(model_kind.parse_gguf(gguf_bytes("flux", ["a.weight"]))["keys"], ["a.weight"])
|
|
self.assertEqual(model_kind.classify("gguf", model_kind.parse_gguf(gguf_bytes("t5encoder")))["kind"], "text_encoders")
|
|
# Cut short: what fits is still read.
|
|
self.assertEqual(model_kind.parse_gguf(gguf_bytes("wan")[:70])["meta"].get("general.architecture"), "wan")
|
|
self.assertIsNone(model_kind.parse_gguf(b"not a gguf"))
|
|
|
|
def test_the_drawing_models_parameters(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
checkpoint = os.path.join(temp, "c.safetensors") # its text encoder and VAE are not counted
|
|
Path(checkpoint).write_bytes(safetensors_bytes(tensors(
|
|
"model.diffusion_model.input_blocks.0.0.weight", "model.diffusion_model.out.2.weight",
|
|
"cond_stage_model.transformer.x.weight", "first_stage_model.decoder.x.weight",
|
|
shapes={"model.diffusion_model.input_blocks.0.0.weight": [320, 4, 3, 3], "model.diffusion_model.out.2.weight": [4, 320]},
|
|
meta={"a": "b"})))
|
|
self.assertEqual(model_kind.diffusion_params(checkpoint), 320 * 4 * 9 + 4 * 320)
|
|
unet = os.path.join(temp, "u.safetensors")
|
|
Path(unet).write_bytes(safetensors_bytes(tensors("double_blocks.0.x.weight", "vae.y.weight", shapes={"double_blocks.0.x.weight": [3072, 3072]})))
|
|
self.assertEqual(model_kind.diffusion_params(unet), 3072 * 3072)
|
|
gguf = os.path.join(temp, "g.gguf")
|
|
Path(gguf).write_bytes(gguf_bytes("flux", ["a.weight", "b.weight"]))
|
|
self.assertEqual(model_kind.diffusion_params(gguf), 2 * 4 * 4)
|
|
Path(gguf).write_bytes(gguf_bytes("flux", ["a.weight", "b.weight"])[:-10]) # cut short: not a size to go by
|
|
self.assertEqual(model_kind.diffusion_params(gguf), 0)
|
|
self.assertEqual(model_kind.diffusion_params(os.path.join(temp, "missing.safetensors")), 0)
|
|
|
|
def test_files_on_disk_and_the_scanner(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
lora = os.path.join(temp, "x.safetensors")
|
|
Path(lora).write_bytes(safetensors_bytes(tensors("lora_unet_double_blocks_0_img_attn_proj.lora_up.weight")))
|
|
self.assertEqual(model_kind.inspect_file(lora)["kind"], "loras")
|
|
self.assertEqual(model_identity.infer_base_model_from_header(lora), "Flux.1 D")
|
|
pickle = os.path.join(temp, "x.ckpt")
|
|
Path(pickle).write_bytes(b"\x80\x02whatever")
|
|
self.assertEqual(model_kind.inspect_file(pickle), {"kind": "", "base": "", "sure": False, "format": None})
|
|
self.assertEqual(model_identity.infer_base_model_from_header(pickle), "Unknown")
|
|
broken = os.path.join(temp, "y.safetensors")
|
|
Path(broken).write_bytes(struct.pack("<Q", 10 ** 12))
|
|
self.assertEqual(model_kind.read_head(broken), (None, None))
|
|
|
|
|
|
class ImportTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
base = os.path.realpath(self.temp.name)
|
|
self.loras = os.path.join(base, "models", "loras")
|
|
self.downloads = os.path.join(base, "Downloads")
|
|
os.makedirs(self.loras)
|
|
os.makedirs(os.path.join(self.downloads, "a", "b"))
|
|
roots = {"loras": [self.loras], "vae": [os.path.join(base, "models", "vae")]}
|
|
self.patches = [
|
|
mock.patch.object(model_import.folder_paths, "folder_names_and_paths", {name: (paths, set()) for name, paths in roots.items()}),
|
|
mock.patch.object(model_import.folder_paths, "get_folder_paths", side_effect=lambda t: roots.get(t, [])),
|
|
mock.patch.object(model_import, "add_entry"),
|
|
mock.patch.object(model_import, "forget_file_lists"),
|
|
mock.patch.object(model_import, "search_folders", return_value=[("downloads", self.downloads)]),
|
|
]
|
|
for patch in self.patches:
|
|
patch.start()
|
|
self.data = safetensors_bytes(tensors("lora_unet_double_blocks_0_img_attn_proj.lora_up.weight"))
|
|
|
|
def tearDown(self):
|
|
for patch in reversed(self.patches):
|
|
patch.stop()
|
|
self.temp.cleanup()
|
|
|
|
def dropped(self, *parts):
|
|
path = os.path.join(self.downloads, *parts)
|
|
Path(path).write_bytes(self.data)
|
|
return path, os.stat(path).st_mtime * 1000
|
|
|
|
def test_the_dropped_file_is_found_by_name_size_and_time(self):
|
|
path, mtime = self.dropped("a", "Cool.safetensors")
|
|
size = len(self.data)
|
|
self.assertEqual(model_import.find_original("cool.safetensors", size, mtime), ("downloads", path))
|
|
self.assertIsNone(model_import.find_original("cool.safetensors", size + 1, mtime))
|
|
self.assertIsNone(model_import.find_original("cool.safetensors", size, mtime - 60000))
|
|
# In subfolders two deep, not deeper.
|
|
self.dropped("a", "b", "two.safetensors")
|
|
self.assertIsNotNone(model_import.find_original("two.safetensors", size, 0))
|
|
os.makedirs(os.path.join(self.downloads, "a", "b", "c"))
|
|
self.dropped("a", "b", "c", "three.safetensors")
|
|
self.assertIsNone(model_import.find_original("three.safetensors", size, 0))
|
|
|
|
def test_inspect_found_and_not_found(self):
|
|
path, mtime = self.dropped("cool.safetensors")
|
|
found = model_import.inspect("cool.safetensors", len(self.data), mtime, b"", local=True)
|
|
self.assertEqual((found["kind"], found["base"], found["where"]), ("loras", "Flux.1 D", "downloads"))
|
|
self.assertTrue(found["token"])
|
|
# Asked from another computer: never looked for, judged from the bytes sent.
|
|
remote = model_import.inspect("cool.safetensors", len(self.data), mtime, self.data, local=False)
|
|
self.assertNotIn("token", remote)
|
|
self.assertEqual(remote["kind"], "loras")
|
|
# A type this ComfyUI has no folder for is not offered.
|
|
control = safetensors_bytes(tensors("control_model.input_hint_block.0.weight"))
|
|
self.assertEqual(model_import.inspect("c.safetensors", len(control), 0, control, local=False)["kind"], "")
|
|
|
|
def test_place_moves_and_never_replaces(self):
|
|
path, mtime = self.dropped("cool.safetensors")
|
|
token = model_import.inspect("cool.safetensors", len(self.data), mtime, b"", local=True)["token"]
|
|
Path(self.loras, "Flux").mkdir()
|
|
Path(self.loras, "Flux", "cool.safetensors").write_bytes(b"older")
|
|
placed = model_import.place({"token": token, "type": "loras", "root": 0, "rel": "Flux/cool.safetensors"})
|
|
self.assertEqual(placed["rel"], "Flux/cool (2).safetensors")
|
|
self.assertFalse(os.path.exists(path))
|
|
self.assertEqual(Path(self.loras, "Flux", "cool.safetensors").read_bytes(), b"older")
|
|
self.assertEqual(Path(self.loras, "Flux", "cool (2).safetensors").read_bytes(), self.data)
|
|
with self.assertRaises(model_import.Refused) as refused:
|
|
model_import.place({"token": token, "type": "loras", "root": 0, "rel": "x.safetensors"})
|
|
self.assertEqual(refused.exception.code, "gone")
|
|
|
|
def test_place_keeps_the_original_when_asked(self):
|
|
path, mtime = self.dropped("cool.safetensors")
|
|
token = model_import.inspect("cool.safetensors", len(self.data), mtime, b"", local=True)["token"]
|
|
model_import.place({"token": token, "type": "loras", "root": 0, "rel": "cool.safetensors", "keep": True})
|
|
self.assertTrue(os.path.exists(path))
|
|
self.assertEqual(Path(self.loras, "cool.safetensors").read_bytes(), self.data)
|
|
self.assertFalse(Path(self.loras, "cool.safetensors.part").exists())
|
|
|
|
def test_a_failed_copy_leaves_nothing_behind(self):
|
|
path, _ = self.dropped("cool.safetensors")
|
|
dest = os.path.join(self.loras, "cool.safetensors")
|
|
with mock.patch.object(model_import.shutil, "copyfile", side_effect=OSError("disk full")):
|
|
with self.assertRaises(OSError):
|
|
model_import.move_into(path, dest, keep=True)
|
|
self.assertTrue(os.path.exists(path))
|
|
self.assertFalse(os.path.exists(dest))
|
|
|
|
def test_a_garbled_copy_is_caught_and_the_original_stays(self):
|
|
path, _ = self.dropped("cool.safetensors")
|
|
dest = os.path.join(self.loras, "cool.safetensors")
|
|
garble = lambda source, target: Path(target).write_bytes(b"x" * len(Path(source).read_bytes())) # right size, wrong bytes
|
|
with mock.patch.object(model_import, "_same_drive", return_value=False), \
|
|
mock.patch.object(model_import.shutil, "copyfile", side_effect=garble):
|
|
with self.assertRaises(model_import.Refused) as refused:
|
|
model_import.move_into(path, dest, keep=False)
|
|
self.assertEqual(refused.exception.code, "incomplete")
|
|
self.assertEqual(Path(path).read_bytes(), self.data)
|
|
self.assertFalse(os.path.exists(dest) or os.path.exists(dest + ".part"))
|
|
|
|
def test_a_file_that_appeared_meanwhile_is_never_replaced(self):
|
|
# free_path found the name free; a file took it before the move: the move fails, both stay.
|
|
path, _ = self.dropped("cool.safetensors")
|
|
dest = os.path.join(self.loras, "cool.safetensors")
|
|
Path(dest).write_bytes(b"theirs")
|
|
for keep in (False, True):
|
|
with self.assertRaises(FileExistsError):
|
|
model_import.move_into(path, dest, keep=keep)
|
|
self.assertEqual((Path(path).read_bytes(), Path(dest).read_bytes()), (self.data, b"theirs"))
|
|
self.assertFalse(os.path.exists(dest + ".part"))
|
|
|
|
def test_destinations_stay_inside_the_models_folder(self):
|
|
for data, code in (({"type": "checkpoints", "rel": "a.safetensors"}, "bad_type"),
|
|
({"type": "loras", "root": 3, "rel": "a.safetensors"}, "bad_root"),
|
|
({"type": "loras", "rel": "../a.safetensors"}, "bad_path"),
|
|
({"type": "loras", "rel": "a.txt"}, "bad_name")):
|
|
with self.assertRaises(model_import.Refused) as refused:
|
|
model_import.destination(data)
|
|
self.assertEqual(refused.exception.code, code)
|
|
|
|
def test_upload(self):
|
|
async def run():
|
|
app = web.Application()
|
|
model_import.register_routes(app)
|
|
async with TestClient(TestServer(app)) as client:
|
|
ok = await client.put("/anomalous/import/upload", params={"type": "loras", "root": "0", "rel": "up/new.safetensors",
|
|
"size": str(len(self.data))}, data=self.data)
|
|
short = await client.put("/anomalous/import/upload", params={"type": "loras", "root": "0", "rel": "up/short.safetensors",
|
|
"size": str(len(self.data) + 5)}, data=self.data)
|
|
return ok.status, await ok.json(), short.status, await short.json()
|
|
status, body, short_status, short_body = asyncio.run(run())
|
|
self.assertEqual((status, body["rel"]), (200, "up/new.safetensors"))
|
|
self.assertEqual(Path(self.loras, "up", "new.safetensors").read_bytes(), self.data)
|
|
self.assertEqual((short_status, short_body["error"]), (400, "incomplete"))
|
|
self.assertEqual(sorted(os.listdir(os.path.join(self.loras, "up"))), ["new.safetensors"])
|
|
|
|
def test_upload_only_from_this_computer(self):
|
|
async def run():
|
|
app = web.Application()
|
|
model_import.register_routes(app)
|
|
async with TestClient(TestServer(app)) as client: # through a proxy: not this computer
|
|
resp = await client.put("/anomalous/import/upload", headers={"X-Forwarded-For": "192.168.1.20"},
|
|
params={"type": "loras", "root": "0", "rel": "far.safetensors", "size": str(len(self.data))},
|
|
data=self.data)
|
|
return resp.status, await resp.json()
|
|
self.assertEqual(asyncio.run(run()), (403, {"error": "not_local"}))
|
|
self.assertFalse(Path(self.loras, "far.safetensors").exists())
|
|
|
|
def test_identify_reports_civitai_and_a_copy_already_there(self):
|
|
path, mtime = self.dropped("cool.safetensors")
|
|
token = model_import.inspect("cool.safetensors", len(self.data), mtime, b"", local=True)["token"]
|
|
record = {"found": True, "base_model": "Illustrious", "model_name": "Cool", "version_name": "v2", "page": "p"}
|
|
with mock.patch.object(model_import, "civitai_by_hash", return_value=record), \
|
|
mock.patch.object(model_import, "_same_hash", return_value="Flux/old.safetensors"):
|
|
found = model_import.identify(token)
|
|
self.assertEqual(found["civitai"]["version_name"], "v2")
|
|
self.assertEqual(found["duplicate"], "Flux/old.safetensors")
|
|
with mock.patch.object(model_import, "civitai_by_hash", side_effect=model_import.SourceUnreachable("down")), \
|
|
mock.patch.object(model_import, "_same_hash", return_value=None):
|
|
self.assertEqual(model_import.identify(token)["reason"], "network")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|