mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-11 00:47:14 +08:00
55a15f87ce
* feat(assets): add namespaced model type tags * fix(assets): mark path-derived upload tags automatic * fix(assets): merge duplicate scan specs * test(assets): make duplicate path normalization portable * feat(assets): add loader_path as the authoritative loader locator (#14796) * fix(assets): filter model_type tags by bucket extension sets Buckets sharing a base directory (e.g. diffusion_models and a custom unet_gguf) tagged every file in the directory regardless of whether the bucket could load it, so .safetensors files were tagged model_type:unet_gguf and vice versa. Carry each bucket's registered extension set through get_comfy_models_folders and only emit a model_type tag when the file extension matches, keeping the empty-set match-all convention from folder_paths.filter_files_extensions. Files under a model base matching no bucket now keep only the models tag instead of every directory-matching model_type tag. * feat(assets): replace response file_path with persisted loader_path The old file_path response field was a namespaced storage locator (models/checkpoints/foo.safetensors): not an absolute path, not unique identity, and not the value a loader consumes. Nothing needs that shape on the wire (hash/ID-based locating is the long-term direction), so it is dropped rather than renamed; the storage-root matching stays internal, powering display_name. What loaders DO need is the in-root loader path (category dropped: models/checkpoints/foo/bar.safetensors -> foo/bar.safetensors). Serve it as a first-class loader_path field, persisted on asset_references (migration 0006) and written by every ingest pipeline at insert, so responses read the column verbatim. Like the model_type tags, loader_path is a seed-time derivative of the model folder registry, maintained by the same scan lifecycle (new files seed fresh values, pruning retires rows whose bucket disappeared). Rows predating the column serve a null loader_path; databases from before this stack already need recreating for the base branch's tag changes. loader_path resolves every registered base including extra_model_paths entries; display_name only the canonical storage roots. A file can therefore be loadable with no display name (extra-path models) or the reverse (unregistered files under the models root), and loader_path is null exactly when no loader can resolve the file. * test(assets): lock loader_path matrix (asymmetry, null, persist/read) Cover the behaviour that has no production change but is easy to regress: the extra-path asymmetry (loadable but no storage namespace), null loader_path persistence for orphan files, and the response reading the stored column with a compute fallback for un-backfilled rows. * fix(assets): persist subfolder-qualified loader_path for ingested outputs ingest_existing_file built its seed spec with the file's basename, so outputs saved into a subfolder persisted loader_path (and the user_metadata filename that preview URLs split for their subfolder param) as just the basename: the served locator pointed at a file that does not exist at that path. Scanner and seeder specs already derive fname via compute_loader_path; use the same derivation here. * fix(assets): only extension-matching buckets contribute a loader_path The model-base match in get_asset_category_and_relative_path ignored each bucket's extension set, so a file inside a registered base whose extension the bucket cannot load (e.g. a .txt uploaded into model_type:checkpoints) advertised a loader_path that no loader list would ever resolve, while the tag side of the same stack already excluded it. Apply the extension check used for backend tags (empty set accepts any extension), keeping loader_path null exactly when no loader can resolve the file. * fix(assets): refresh loader_path when re-ingesting an existing reference upsert_reference only wrote loader_path on the INSERT branch, so re-ingesting an existing reference (an output overwritten in place, or a file re-registered after its loader_path derivation changed) kept the stale or NULL value forever. Write it on the UPDATE branch too, with a null-safe change guard so a loader_path difference alone is enough to trigger the update, and identical values stay a no-op. * fix(assets): repair semantic merge breakage from #14796 and master Two textually-clean but semantically-broken merges: - routes.py lost its folder_paths import when #14796's import block superseded the base's, while the content-type hardening added via the base's master merge still calls folder_paths.is_dangerous_content_type. - master's SVG download-hardening test uploads with the pre-namespacing bare checkpoints tag, which this branch's destination validation rejects; use model_type:checkpoints. --------- Co-authored-by: guill <jacob.e.segal@gmail.com>
320 lines
12 KiB
Python
320 lines
12 KiB
Python
"""Tests for bulk ingest services."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.assets.database.models import Asset, AssetReference
|
|
from app.assets.database.queries import get_reference_tags
|
|
from app.assets.scanner import build_asset_specs
|
|
from app.assets.services.bulk_ingest import SeedAssetSpec, batch_insert_seed_assets
|
|
|
|
|
|
class TestBatchInsertSeedAssets:
|
|
def test_populates_mime_type_for_model_files(self, session: Session, temp_dir: Path):
|
|
"""Verify mime_type is stored in the Asset table for model files."""
|
|
file_path = temp_dir / "model.safetensors"
|
|
file_path.write_bytes(b"fake safetensors content")
|
|
|
|
specs: list[SeedAssetSpec] = [
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 24,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Test Model",
|
|
"tags": ["models"],
|
|
"fname": "model.safetensors",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": "application/safetensors",
|
|
}
|
|
]
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
|
|
# Verify Asset has mime_type populated
|
|
assets = session.query(Asset).all()
|
|
assert len(assets) == 1
|
|
assert assets[0].mime_type == "application/safetensors"
|
|
|
|
def test_mime_type_none_when_not_provided(self, session: Session, temp_dir: Path):
|
|
"""Verify mime_type is None when not provided in spec."""
|
|
file_path = temp_dir / "unknown.bin"
|
|
file_path.write_bytes(b"binary data")
|
|
|
|
specs: list[SeedAssetSpec] = [
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 11,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Unknown File",
|
|
"tags": [],
|
|
"fname": "unknown.bin",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": None,
|
|
}
|
|
]
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
|
|
assets = session.query(Asset).all()
|
|
assert len(assets) == 1
|
|
assert assets[0].mime_type is None
|
|
|
|
def test_various_model_mime_types(self, session: Session, temp_dir: Path):
|
|
"""Verify various model file types get correct mime_type."""
|
|
test_cases = [
|
|
("model.safetensors", "application/safetensors"),
|
|
("model.pt", "application/pytorch"),
|
|
("model.ckpt", "application/pickle"),
|
|
("model.gguf", "application/gguf"),
|
|
]
|
|
|
|
specs: list[SeedAssetSpec] = []
|
|
for filename, mime_type in test_cases:
|
|
file_path = temp_dir / filename
|
|
file_path.write_bytes(b"content")
|
|
specs.append(
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 7,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": filename,
|
|
"tags": [],
|
|
"fname": filename,
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": mime_type,
|
|
}
|
|
)
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == len(test_cases)
|
|
|
|
for filename, expected_mime in test_cases:
|
|
ref = session.query(AssetReference).filter_by(name=filename).first()
|
|
assert ref is not None
|
|
asset = session.query(Asset).filter_by(id=ref.asset_id).first()
|
|
assert asset.mime_type == expected_mime, f"Expected {expected_mime} for {filename}, got {asset.mime_type}"
|
|
|
|
def test_duplicate_paths_merge_tags_before_insert(
|
|
self, session: Session, temp_dir: Path
|
|
):
|
|
"""Overlapping model-folder registrations can emit the same path twice."""
|
|
file_path = temp_dir / "shared.safetensors"
|
|
file_path.write_bytes(b"shared model")
|
|
|
|
specs: list[SeedAssetSpec] = [
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 12,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Shared Model",
|
|
"tags": ["models", "model_type:checkpoints"],
|
|
"fname": "shared.safetensors",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": "application/safetensors",
|
|
},
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 12,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Shared Model",
|
|
"tags": ["models", "model_type:diffusion_models"],
|
|
"fname": "shared.safetensors",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": "application/safetensors",
|
|
},
|
|
]
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
assert result.won_paths == 1
|
|
refs = session.query(AssetReference).all()
|
|
assert len(refs) == 1
|
|
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
|
|
"models",
|
|
"model_type:checkpoints",
|
|
"model_type:diffusion_models",
|
|
}
|
|
|
|
def test_duplicate_paths_are_merged_after_abspath_normalization(
|
|
self, session: Session, temp_dir: Path, monkeypatch
|
|
):
|
|
"""The scanner may emit equivalent paths with different spelling."""
|
|
file_path = temp_dir / "same-file.safetensors"
|
|
file_path.write_bytes(b"shared model")
|
|
monkeypatch.chdir(temp_dir)
|
|
relative_path = file_path.name
|
|
absolute_path = os.path.abspath(relative_path)
|
|
|
|
specs: list[SeedAssetSpec] = [
|
|
{
|
|
"abs_path": relative_path,
|
|
"size_bytes": 12,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Shared Model",
|
|
"tags": ["models", "model_type:checkpoints"],
|
|
"fname": "same-file.safetensors",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": "application/safetensors",
|
|
},
|
|
{
|
|
"abs_path": absolute_path,
|
|
"size_bytes": 12,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "Shared Model",
|
|
"tags": ["models", "model_type:diffusion_models"],
|
|
"fname": "same-file.safetensors",
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": "application/safetensors",
|
|
},
|
|
]
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
assert result.won_paths == 1
|
|
refs = session.query(AssetReference).all()
|
|
assert len(refs) == 1
|
|
assert refs[0].file_path == absolute_path
|
|
# loader_path is persisted from the spec's fname (compute_loader_path).
|
|
assert refs[0].loader_path == "same-file.safetensors"
|
|
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
|
|
"models",
|
|
"model_type:checkpoints",
|
|
"model_type:diffusion_models",
|
|
}
|
|
|
|
def test_scanner_duplicate_shared_model_paths_keep_all_model_type_tags(
|
|
self, session: Session, temp_dir: Path
|
|
):
|
|
"""Shared extra model roots make scanner collection emit duplicate paths."""
|
|
shared_root = temp_dir / "shared"
|
|
input_dir = temp_dir / "input"
|
|
output_dir = temp_dir / "output"
|
|
temp_root = temp_dir / "temp"
|
|
for directory in (shared_root, input_dir, output_dir, temp_root):
|
|
directory.mkdir()
|
|
file_path = shared_root / "dual_use_model.safetensors"
|
|
file_path.write_bytes(b"shared model")
|
|
|
|
with (
|
|
patch("app.assets.services.path_utils.folder_paths") as mock_fp,
|
|
patch(
|
|
"app.assets.services.path_utils.get_comfy_models_folders",
|
|
return_value=[
|
|
("checkpoints", [str(shared_root)], {".safetensors"}),
|
|
("diffusion_models", [str(shared_root)], {".safetensors"}),
|
|
],
|
|
),
|
|
):
|
|
mock_fp.get_input_directory.return_value = str(input_dir)
|
|
mock_fp.get_output_directory.return_value = str(output_dir)
|
|
mock_fp.get_temp_directory.return_value = str(temp_root)
|
|
|
|
specs, tag_pool, skipped = build_asset_specs(
|
|
paths=[str(file_path), str(file_path)],
|
|
existing_paths=set(),
|
|
enable_metadata_extraction=False,
|
|
compute_hashes=False,
|
|
)
|
|
|
|
assert skipped == 0
|
|
assert len(specs) == 2
|
|
assert tag_pool == {
|
|
"models",
|
|
"model_type:checkpoints",
|
|
"model_type:diffusion_models",
|
|
}
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
assert result.won_paths == 1
|
|
refs = session.query(AssetReference).all()
|
|
assert len(refs) == 1
|
|
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
|
|
"models",
|
|
"model_type:checkpoints",
|
|
"model_type:diffusion_models",
|
|
}
|
|
|
|
def test_loader_path_persisted_as_null_when_fname_is_none(
|
|
self, session: Session, temp_dir: Path
|
|
):
|
|
"""A file with no in-root loader path (fname=None, e.g. an orphan under
|
|
models_root) persists loader_path as NULL rather than a synthesized value."""
|
|
file_path = temp_dir / "orphan.bin"
|
|
file_path.write_bytes(b"x")
|
|
|
|
specs: list[SeedAssetSpec] = [
|
|
{
|
|
"abs_path": str(file_path),
|
|
"size_bytes": 1,
|
|
"mtime_ns": 1234567890000000000,
|
|
"info_name": "orphan.bin",
|
|
"tags": [],
|
|
"fname": None,
|
|
"metadata": None,
|
|
"hash": None,
|
|
"mime_type": None,
|
|
}
|
|
]
|
|
|
|
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
|
|
|
|
assert result.inserted_refs == 1
|
|
refs = session.query(AssetReference).all()
|
|
assert len(refs) == 1
|
|
assert refs[0].file_path == str(file_path)
|
|
assert refs[0].loader_path is None
|
|
|
|
|
|
class TestMetadataExtraction:
|
|
def test_extracts_mime_type_for_model_files(self, temp_dir: Path):
|
|
"""Verify metadata extraction returns correct mime_type for model files."""
|
|
from app.assets.services.metadata_extract import extract_file_metadata
|
|
|
|
file_path = temp_dir / "model.safetensors"
|
|
file_path.write_bytes(b"fake safetensors content")
|
|
|
|
meta = extract_file_metadata(str(file_path))
|
|
|
|
assert meta.content_type == "application/safetensors"
|
|
|
|
def test_mime_type_for_various_model_formats(self, temp_dir: Path):
|
|
"""Verify various model file types get correct mime_type from metadata."""
|
|
from app.assets.services.metadata_extract import extract_file_metadata
|
|
|
|
test_cases = [
|
|
("model.safetensors", "application/safetensors"),
|
|
("model.sft", "application/safetensors"),
|
|
("model.pt", "application/pytorch"),
|
|
("model.pth", "application/pytorch"),
|
|
("model.ckpt", "application/pickle"),
|
|
("model.pkl", "application/pickle"),
|
|
("model.gguf", "application/gguf"),
|
|
]
|
|
|
|
for filename, expected_mime in test_cases:
|
|
file_path = temp_dir / filename
|
|
file_path.write_bytes(b"content")
|
|
|
|
meta = extract_file_metadata(str(file_path))
|
|
|
|
assert meta.content_type == expected_mime, f"Expected {expected_mime} for {filename}, got {meta.content_type}"
|