mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-23 00:09:32 +08:00
fix(assets): move model asset on model_type: edit to stay loader-coherent
Under supports_model_type_tags, a model_type: edit on a filesystem-backed model asset now relocates and re-registers the file to the folder matching the new model_type:<folder_name>, so the file location stays coherent with the label instead of a label-only relabel that left the file loader-orphaned. Triggered off the POST /tags add (apply_tags): on a registered folder change the file moves (preserving subfolder + filename), path-derived system tags are re-derived, and filename metadata refreshed, all in one transaction with file rollback on failure. Collisions are rejected (never clobber); an unregistered model_type: stays a permissive label; non-model and hash-only references are untouched. Shared on-disk folders keep their plural twins.
This commit is contained in:
parent
990242a07d
commit
bf898cc552
@ -26,6 +26,7 @@ from app.assets.seeder import ScanInProgressError, asset_seeder
|
||||
from app.assets.services import (
|
||||
DependencyMissingError,
|
||||
HashMismatchError,
|
||||
ModelMoveError,
|
||||
apply_tags,
|
||||
asset_exists,
|
||||
create_from_hash,
|
||||
@ -623,6 +624,11 @@ async def add_asset_tags(request: web.Request) -> web.Response:
|
||||
already_present=result.already_present,
|
||||
total_tags=result.total_tags,
|
||||
)
|
||||
except ModelMoveError as me:
|
||||
# A model_type: edit that couldn't be applied coherently (unknown
|
||||
# folder, or destination collision). The FE re-adds the prior
|
||||
# model_type: tag on a non-2xx, so the asset stays coherent.
|
||||
return _build_error_response(me.status, me.code, me.message, {"id": reference_id})
|
||||
except PermissionError as pe:
|
||||
return _build_error_response(403, "FORBIDDEN", str(pe), {"id": reference_id})
|
||||
except ValueError as ve:
|
||||
|
||||
@ -22,9 +22,11 @@ from app.assets.services.file_utils import (
|
||||
from app.assets.services.ingest import (
|
||||
DependencyMissingError,
|
||||
HashMismatchError,
|
||||
ModelMoveError,
|
||||
create_from_hash,
|
||||
ingest_existing_file,
|
||||
register_output_files,
|
||||
relocate_model_asset_for_model_type_tags,
|
||||
upload_from_temp_path,
|
||||
)
|
||||
from app.assets.database.queries import (
|
||||
@ -62,8 +64,10 @@ __all__ = [
|
||||
"HashMismatchError",
|
||||
"IngestResult",
|
||||
"ListAssetsResult",
|
||||
"ModelMoveError",
|
||||
"RegisterAssetResult",
|
||||
"RemoveTagsResult",
|
||||
"relocate_model_asset_for_model_type_tags",
|
||||
"TagUsage",
|
||||
"UploadResult",
|
||||
"UserMetadata",
|
||||
|
||||
@ -2,11 +2,13 @@ import contextlib
|
||||
import logging
|
||||
import mimetypes
|
||||
import os
|
||||
import shutil
|
||||
from typing import Any, Sequence
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
import app.assets.services.hashing as hashing
|
||||
from app.assets.database.models import AssetReference
|
||||
from app.assets.database.queries import (
|
||||
add_tags_to_reference,
|
||||
count_active_siblings,
|
||||
@ -20,6 +22,7 @@ from app.assets.database.queries import (
|
||||
list_references_by_asset_id,
|
||||
reference_exists,
|
||||
remove_missing_tag_for_asset_id,
|
||||
remove_tags_from_reference,
|
||||
set_reference_metadata,
|
||||
set_reference_system_metadata,
|
||||
set_reference_tags,
|
||||
@ -35,7 +38,9 @@ from app.assets.services.image_dimensions import extract_image_dimensions
|
||||
from app.assets.services.path_utils import (
|
||||
compute_relative_filename,
|
||||
get_backend_system_tags_from_path,
|
||||
get_model_base_for_folder,
|
||||
get_name_and_tags_from_asset_path,
|
||||
model_folders_for_path,
|
||||
resolve_destination_from_tags,
|
||||
validate_path_within_base,
|
||||
)
|
||||
@ -453,6 +458,186 @@ class DependencyMissingError(Exception):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ModelMoveError(Exception):
|
||||
"""A model_type: edit could not be applied coherently (BE-1641).
|
||||
|
||||
Carries an HTTP-ish ``status``/``code`` so the route can surface a precise
|
||||
4xx (rather than the generic 404 the bare ValueError path produces). The FE
|
||||
edit-type flow compensates on any non-2xx by re-adding the prior
|
||||
``model_type:`` tag, so a reject here leaves the asset coherent.
|
||||
"""
|
||||
|
||||
def __init__(self, code: str, message: str, status: int = 409):
|
||||
self.code = code
|
||||
self.message = message
|
||||
self.status = status
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
def _move_file(src: str, dst: str) -> None:
|
||||
"""Relocate a file, falling back to a cross-device copy+unlink.
|
||||
|
||||
``os.replace`` is atomic but fails with ``EXDEV`` when src and dst live on
|
||||
different filesystems (``extra_model_paths`` may point at another mount), so
|
||||
fall back to ``shutil.move`` there.
|
||||
"""
|
||||
try:
|
||||
os.replace(src, dst)
|
||||
except OSError:
|
||||
shutil.move(src, dst)
|
||||
|
||||
|
||||
def relocate_model_asset_for_model_type_tags(
|
||||
session: Session,
|
||||
ref: AssetReference,
|
||||
requested_tags: Sequence[str],
|
||||
origin: str = "automatic",
|
||||
) -> bool:
|
||||
"""Move a filesystem-backed model asset to match an added ``model_type:`` tag.
|
||||
|
||||
BE-1641 / spec-drift §2: under the ``supports_model_type_tags`` contract a
|
||||
``model_type:`` edit must stay coherent on *edit*, not just upload. When a
|
||||
``model_type:<folder>`` tag is applied to a filesystem-backed model asset
|
||||
whose file is not already under that folder, move the file to the folder's
|
||||
base and re-derive the path-based system tags so location and label agree.
|
||||
|
||||
Mutates ``ref`` and its tags in-place within ``session`` (the caller owns
|
||||
the commit). Returns True if a physical move happened, False otherwise
|
||||
(non-filesystem / hash-only / non-model asset, no ``model_type:`` added, or
|
||||
the target folder already covers the current path — the shared-dir case in
|
||||
spec-drift §1).
|
||||
|
||||
Raises:
|
||||
ModelMoveError: the target folder is unknown, or the destination is
|
||||
already occupied (collision) — never clobbers (Q2).
|
||||
"""
|
||||
if not ref.file_path:
|
||||
# API-created / hash-only reference: nothing on disk to move. Labels
|
||||
# stay labels (matches AC scope: "filesystem-backed model asset").
|
||||
return False
|
||||
|
||||
requested_folders = [
|
||||
t.split(":", 1)[1]
|
||||
for t in normalize_tags(list(requested_tags))
|
||||
if t.startswith("model_type:") and t.split(":", 1)[1]
|
||||
]
|
||||
if not requested_folders:
|
||||
return False
|
||||
|
||||
old_path = os.path.abspath(ref.file_path)
|
||||
current_folders = set(model_folders_for_path(old_path))
|
||||
if not current_folders:
|
||||
# Not under any model base (e.g. an input/output asset). A model_type:
|
||||
# label here is meaningless for placement; leave it as a plain label.
|
||||
return False
|
||||
|
||||
# The FE emits exactly one model_type: per edit; if several are requested,
|
||||
# the last one wins deterministically.
|
||||
target_folder = requested_folders[-1]
|
||||
|
||||
# Shared on-disk dir (spec-drift §1): the path already covers the target
|
||||
# folder, so re-deriving would keep both twins — no physical move needed.
|
||||
if target_folder in current_folders:
|
||||
return False
|
||||
|
||||
# An unregistered folder_name cannot correspond to any real on-disk
|
||||
# location, so there is nothing to relocate. Keep Core's established
|
||||
# permissive model_type: labeling (spec-drift §3: Core is local/trusted and
|
||||
# does not reject model_type: labels) — store it as a plain label, don't
|
||||
# reject. A genuine edit-type action always targets a registered folder_name
|
||||
# from the discovery vocabulary, so this only affects manual label adds.
|
||||
try:
|
||||
new_base = get_model_base_for_folder(target_folder)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
rel = compute_relative_filename(old_path)
|
||||
if not rel:
|
||||
raise ModelMoveError(
|
||||
"INVALID_MODEL_PATH",
|
||||
f"cannot determine relative path for model asset: {old_path}",
|
||||
status=400,
|
||||
)
|
||||
new_path = os.path.abspath(os.path.join(new_base, rel))
|
||||
try:
|
||||
validate_path_within_base(new_path, new_base)
|
||||
except ValueError as e:
|
||||
raise ModelMoveError("INVALID_MODEL_PATH", str(e), status=400)
|
||||
|
||||
if new_path == old_path:
|
||||
return False
|
||||
|
||||
# Q2: collision -> reject, never clobber. Cover both an on-disk file and a
|
||||
# reference that already owns the destination path.
|
||||
if os.path.exists(new_path):
|
||||
raise ModelMoveError(
|
||||
"DESTINATION_EXISTS", f"destination already exists: {new_path}"
|
||||
)
|
||||
if get_reference_by_file_path(session, new_path) is not None:
|
||||
raise ModelMoveError(
|
||||
"DESTINATION_EXISTS", f"destination already registered: {new_path}"
|
||||
)
|
||||
|
||||
os.makedirs(os.path.dirname(new_path), exist_ok=True)
|
||||
_move_file(old_path, new_path)
|
||||
try:
|
||||
_reregister_moved_reference(session, ref, new_path, origin=origin)
|
||||
except Exception:
|
||||
# Never half-move: roll the file back before surfacing the failure.
|
||||
with contextlib.suppress(Exception):
|
||||
_move_file(new_path, old_path)
|
||||
raise
|
||||
return True
|
||||
|
||||
|
||||
def _reregister_moved_reference(
|
||||
session: Session,
|
||||
ref: AssetReference,
|
||||
new_path: str,
|
||||
origin: str = "automatic",
|
||||
) -> None:
|
||||
"""Point ``ref`` at ``new_path`` and reconcile path-derived system tags.
|
||||
|
||||
Re-derives ``models`` + ``model_type:*`` from the new location, drops any
|
||||
stale ``model_type:`` no longer justified by the path, and refreshes the
|
||||
relative ``filename`` metadata. User labels are left untouched.
|
||||
"""
|
||||
# Bytes are unchanged by a move; only the location and mtime can shift
|
||||
# (shutil.move's cross-device fallback re-stats).
|
||||
_size_bytes, mtime_ns = get_size_and_mtime_ns(new_path)
|
||||
ref.file_path = new_path
|
||||
ref.mtime_ns = mtime_ns
|
||||
ref.updated_at = get_utc_now()
|
||||
session.flush()
|
||||
|
||||
derived = get_backend_system_tags_from_path(new_path)
|
||||
derived_model_types = {t for t in derived if t.startswith("model_type:")}
|
||||
|
||||
current = set(get_reference_tags(session, reference_id=ref.id))
|
||||
stale = {
|
||||
t for t in current if t.startswith("model_type:")
|
||||
} - derived_model_types
|
||||
if stale:
|
||||
remove_tags_from_reference(
|
||||
session, reference_id=ref.id, tags=sorted(stale)
|
||||
)
|
||||
add_tags_to_reference(
|
||||
session,
|
||||
reference_id=ref.id,
|
||||
tags=derived,
|
||||
origin=origin,
|
||||
create_if_missing=True,
|
||||
)
|
||||
|
||||
_update_metadata_with_filename(
|
||||
session,
|
||||
reference_id=ref.id,
|
||||
file_path=new_path,
|
||||
current_metadata=ref.user_metadata,
|
||||
user_metadata={},
|
||||
)
|
||||
|
||||
|
||||
def upload_from_temp_path(
|
||||
temp_path: str,
|
||||
name: str | None = None,
|
||||
|
||||
@ -25,6 +25,44 @@ def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
|
||||
return targets
|
||||
|
||||
|
||||
def get_model_base_for_folder(folder_name: str) -> str:
|
||||
"""Resolve a registered model ``folder_name`` to its canonical base path.
|
||||
|
||||
This is the single-valued reverse lookup (model_type -> on-disk base) the
|
||||
upload destination and the edit-type move (BE-1641) share: resolve the one
|
||||
named folder to its first configured base. Never fans out.
|
||||
|
||||
Raises:
|
||||
ValueError: the folder_name is not registered, or has no base path.
|
||||
"""
|
||||
model_folder_paths = dict(get_comfy_models_folders())
|
||||
try:
|
||||
bases = model_folder_paths[folder_name]
|
||||
except KeyError:
|
||||
raise ValueError(f"unknown model category '{folder_name}'")
|
||||
if not bases:
|
||||
raise ValueError(f"no base path configured for category '{folder_name}'")
|
||||
return os.path.abspath(bases[0])
|
||||
|
||||
|
||||
def model_folders_for_path(path: str) -> list[str]:
|
||||
"""Return every model folder_name whose registered base covers ``path``.
|
||||
|
||||
Set-valued (spec-drift §1): a path shared by >=2 registered folders (e.g.
|
||||
``diffusion_models`` and ``unet_gguf`` both registering ``models/unet``)
|
||||
belongs to all of them. Purely lexical — does not require the file to exist.
|
||||
Empty when ``path`` is not under any model base (e.g. an input/output file).
|
||||
"""
|
||||
fp_path = Path(os.path.abspath(path))
|
||||
out: list[str] = []
|
||||
for folder_name, bases in get_comfy_models_folders():
|
||||
for base in bases:
|
||||
if fp_path.is_relative_to(os.path.abspath(base)):
|
||||
out.append(folder_name)
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def _validate_subfolder(subfolder: str | None) -> list[str]:
|
||||
if not subfolder:
|
||||
return []
|
||||
@ -65,14 +103,7 @@ def resolve_destination_from_tags(
|
||||
folder_name = model_type_tags[0].split(":", 1)[1]
|
||||
if not folder_name:
|
||||
raise ValueError("models uploads require exactly one model_type:<folder_name> tag")
|
||||
model_folder_paths = dict(get_comfy_models_folders())
|
||||
try:
|
||||
bases = model_folder_paths[folder_name]
|
||||
except KeyError:
|
||||
raise ValueError(f"unknown model category '{folder_name}'")
|
||||
if not bases:
|
||||
raise ValueError(f"no base path configured for category '{folder_name}'")
|
||||
base_dir = os.path.abspath(bases[0])
|
||||
base_dir = get_model_base_for_folder(folder_name)
|
||||
elif root == "input":
|
||||
base_dir = os.path.abspath(folder_paths.get_input_directory())
|
||||
else:
|
||||
|
||||
@ -9,6 +9,7 @@ from app.assets.database.queries import (
|
||||
remove_tags_from_reference,
|
||||
)
|
||||
from app.assets.database.queries.tags import list_tag_counts_for_filtered_assets
|
||||
from app.assets.services.ingest import relocate_model_asset_for_model_type_tags
|
||||
from app.assets.services.schemas import TagUsage
|
||||
from app.database.db import create_session
|
||||
|
||||
@ -22,6 +23,12 @@ def apply_tags(
|
||||
with create_session() as session:
|
||||
ref_row = get_reference_with_owner_check(session, reference_id, owner_id)
|
||||
|
||||
# BE-1641: a flag-on model_type: edit on a filesystem-backed model asset
|
||||
# must MOVE/re-register the file so its location stays coherent with the
|
||||
# label (not a label-only relabel). Runs before the label add so the
|
||||
# path-derived system tags are reconciled first; may raise ModelMoveError.
|
||||
relocate_model_asset_for_model_type_tags(session, ref_row, tags)
|
||||
|
||||
result = add_tags_to_reference(
|
||||
session,
|
||||
reference_id=reference_id,
|
||||
|
||||
338
tests-unit/assets_test/services/test_model_type_move.py
Normal file
338
tests-unit/assets_test/services/test_model_type_move.py
Normal file
@ -0,0 +1,338 @@
|
||||
"""Tests for the BE-1641 edit-type MOVE/re-register behavior.
|
||||
|
||||
A flag-on ``model_type:`` edit on a filesystem-backed model asset must move the
|
||||
file to the folder that matches the new ``model_type:<folder_name>`` so the file
|
||||
location stays coherent with the label (not a label-only relabel). The move runs
|
||||
off the ``POST /tags`` add path (``apply_tags``).
|
||||
"""
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset, AssetReference
|
||||
from app.assets.database.queries import (
|
||||
add_tags_to_reference,
|
||||
ensure_tags_exist,
|
||||
get_reference_tags,
|
||||
)
|
||||
from app.assets.helpers import get_utc_now
|
||||
from app.assets.services import ModelMoveError, apply_tags
|
||||
from app.assets.services.ingest import relocate_model_asset_for_model_type_tags
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_dirs():
|
||||
"""Temp model/input/output/temp dirs with a shared on-disk model folder.
|
||||
|
||||
``diffusion_models`` and ``unet_gguf`` both register the same ``models/unet``
|
||||
dir (spec-drift §1 plural membership), so a file there belongs to both.
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as root:
|
||||
root_path = Path(root)
|
||||
checkpoints = root_path / "models" / "checkpoints"
|
||||
loras = root_path / "models" / "loras"
|
||||
unet = root_path / "models" / "unet"
|
||||
input_dir = root_path / "input"
|
||||
output_dir = root_path / "output"
|
||||
temp_dir = root_path / "temp"
|
||||
for d in (checkpoints, loras, unet, input_dir, output_dir, temp_dir):
|
||||
d.mkdir(parents=True)
|
||||
|
||||
folders = [
|
||||
("checkpoints", [str(checkpoints)]),
|
||||
("loras", [str(loras)]),
|
||||
("diffusion_models", [str(unet)]),
|
||||
("unet_gguf", [str(unet)]),
|
||||
]
|
||||
|
||||
with patch("app.assets.services.path_utils.folder_paths") as mock_fp, patch(
|
||||
"app.assets.services.path_utils.get_comfy_models_folders",
|
||||
return_value=folders,
|
||||
):
|
||||
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_dir)
|
||||
yield {
|
||||
"checkpoints": checkpoints,
|
||||
"loras": loras,
|
||||
"unet": unet,
|
||||
"input": input_dir,
|
||||
"output": output_dir,
|
||||
"temp": temp_dir,
|
||||
}
|
||||
|
||||
|
||||
def _make_fs_ref(
|
||||
session: Session,
|
||||
file_path: Path,
|
||||
tags: list[str],
|
||||
*,
|
||||
contents: bytes = b"weights",
|
||||
name: str | None = None,
|
||||
owner_id: str = "",
|
||||
) -> AssetReference:
|
||||
"""Create an on-disk file plus a filesystem-backed reference carrying ``tags``."""
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_path.write_bytes(contents)
|
||||
|
||||
asset = Asset(hash=f"blake3:{file_path.name}", size_bytes=len(contents))
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
|
||||
now = get_utc_now()
|
||||
ref = AssetReference(
|
||||
owner_id=owner_id,
|
||||
name=name or file_path.name,
|
||||
asset_id=asset.id,
|
||||
file_path=str(file_path),
|
||||
mtime_ns=os.stat(file_path).st_mtime_ns,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
session.add(ref)
|
||||
session.flush()
|
||||
|
||||
if tags:
|
||||
ensure_tags_exist(session, tags)
|
||||
add_tags_to_reference(session, reference_id=ref.id, tags=tags)
|
||||
session.commit()
|
||||
return ref
|
||||
|
||||
|
||||
def _tags_after(session: Session, reference_id: str) -> set[str]:
|
||||
session.expire_all()
|
||||
return set(get_reference_tags(session, reference_id))
|
||||
|
||||
|
||||
class TestMoveHappyPath:
|
||||
def test_checkpoint_to_lora_moves_file_and_reconciles_tags(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "model.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
|
||||
dst = model_dirs["loras"] / "model.safetensors"
|
||||
assert not src.exists()
|
||||
assert dst.exists()
|
||||
assert dst.read_bytes() == b"weights"
|
||||
|
||||
session.expire_all()
|
||||
moved = session.get(AssetReference, ref.id)
|
||||
assert moved.file_path == str(dst)
|
||||
|
||||
tags = _tags_after(session, ref.id)
|
||||
assert "model_type:loras" in tags
|
||||
assert "model_type:checkpoints" not in tags
|
||||
assert "models" in tags
|
||||
|
||||
def test_preserves_subfolder_structure(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "sub" / "nested" / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
|
||||
dst = model_dirs["loras"] / "sub" / "nested" / "m.safetensors"
|
||||
assert not src.exists()
|
||||
assert dst.exists()
|
||||
session.expire_all()
|
||||
assert session.get(AssetReference, ref.id).file_path == str(dst)
|
||||
|
||||
def test_preserves_user_labels(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "labelled.safetensors"
|
||||
ref = _make_fs_ref(
|
||||
session, src, ["models", "model_type:checkpoints", "favorite", "sdxl"]
|
||||
)
|
||||
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
|
||||
tags = _tags_after(session, ref.id)
|
||||
assert {"favorite", "sdxl"} <= tags
|
||||
|
||||
def test_refreshes_filename_metadata(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "deep" / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
|
||||
session.expire_all()
|
||||
moved = session.get(AssetReference, ref.id)
|
||||
# filename is relative to the category root, so the subfolder survives.
|
||||
assert moved.user_metadata["filename"] == "deep/m.safetensors"
|
||||
|
||||
|
||||
class TestNoMoveCases:
|
||||
def test_same_type_is_noop(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
moved = relocate_model_asset_for_model_type_tags(
|
||||
session, ref, ["model_type:checkpoints"]
|
||||
)
|
||||
assert moved is False
|
||||
assert src.exists()
|
||||
|
||||
def test_shared_dir_does_not_move(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
# File in the shared unet dir belongs to BOTH diffusion_models and
|
||||
# unet_gguf; editing to the sibling that shares the dir must not move.
|
||||
src = model_dirs["unet"] / "g.gguf"
|
||||
ref = _make_fs_ref(
|
||||
session,
|
||||
src,
|
||||
["models", "model_type:diffusion_models", "model_type:unet_gguf"],
|
||||
)
|
||||
|
||||
moved = relocate_model_asset_for_model_type_tags(
|
||||
session, ref, ["model_type:unet_gguf"]
|
||||
)
|
||||
assert moved is False
|
||||
assert src.exists()
|
||||
|
||||
def test_hash_only_reference_is_label_only(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
asset = Asset(hash="blake3:hashonly", size_bytes=10)
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
now = get_utc_now()
|
||||
ref = AssetReference(
|
||||
name="hashonly",
|
||||
asset_id=asset.id,
|
||||
file_path=None,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
session.add(ref)
|
||||
session.commit()
|
||||
|
||||
moved = relocate_model_asset_for_model_type_tags(
|
||||
session, ref, ["model_type:loras"]
|
||||
)
|
||||
assert moved is False
|
||||
|
||||
def test_non_model_filesystem_asset_is_label_only(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["input"] / "photo.png"
|
||||
ref = _make_fs_ref(session, src, ["input"])
|
||||
|
||||
moved = relocate_model_asset_for_model_type_tags(
|
||||
session, ref, ["model_type:loras"]
|
||||
)
|
||||
assert moved is False
|
||||
assert src.exists()
|
||||
|
||||
def test_no_model_type_tag_is_noop(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
moved = relocate_model_asset_for_model_type_tags(
|
||||
session, ref, ["favorite"]
|
||||
)
|
||||
assert moved is False
|
||||
assert src.exists()
|
||||
|
||||
|
||||
class TestUnknownFolderIsLabelOnly:
|
||||
def test_unknown_folder_is_stored_as_label_not_rejected(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
# Core stays permissive about model_type: LABELS (spec-drift §3): an
|
||||
# unregistered folder_name can't map to a real location, so it's a plain
|
||||
# label, not a move and not a reject.
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
result = apply_tags(reference_id=ref.id, tags=["model_type:bogus"])
|
||||
|
||||
assert src.exists() # not moved
|
||||
assert "model_type:bogus" in result.total_tags
|
||||
# The real path-derived type is untouched.
|
||||
assert "model_type:checkpoints" in result.total_tags
|
||||
|
||||
|
||||
class TestRejects:
|
||||
def test_collision_rejected_409_and_not_clobbered(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
# Pre-existing file at the destination must not be overwritten.
|
||||
dst = model_dirs["loras"] / "m.safetensors"
|
||||
dst.write_bytes(b"existing-lora")
|
||||
|
||||
with pytest.raises(ModelMoveError) as ei:
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
assert ei.value.status == 409
|
||||
assert ei.value.code == "DESTINATION_EXISTS"
|
||||
assert src.exists()
|
||||
assert dst.read_bytes() == b"existing-lora"
|
||||
|
||||
def test_collision_with_registered_reference_rejected(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
|
||||
# Another reference already owns the destination path (no on-disk file).
|
||||
dst = model_dirs["loras"] / "m.safetensors"
|
||||
other = Asset(hash="blake3:other", size_bytes=5)
|
||||
session.add(other)
|
||||
session.flush()
|
||||
now = get_utc_now()
|
||||
session.add(
|
||||
AssetReference(
|
||||
name="m.safetensors",
|
||||
asset_id=other.id,
|
||||
file_path=str(dst),
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
|
||||
with pytest.raises(ModelMoveError) as ei:
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
assert ei.value.code == "DESTINATION_EXISTS"
|
||||
assert src.exists()
|
||||
|
||||
|
||||
class TestRollback:
|
||||
def test_file_rolled_back_when_db_update_fails(
|
||||
self, mock_create_session, session: Session, model_dirs
|
||||
):
|
||||
src = model_dirs["checkpoints"] / "m.safetensors"
|
||||
ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"])
|
||||
dst = model_dirs["loras"] / "m.safetensors"
|
||||
|
||||
with patch(
|
||||
"app.assets.services.ingest.get_size_and_mtime_ns",
|
||||
side_effect=OSError("boom"),
|
||||
):
|
||||
with pytest.raises(OSError, match="boom"):
|
||||
apply_tags(reference_id=ref.id, tags=["model_type:loras"])
|
||||
|
||||
# The half-move must be undone: source restored, destination clean.
|
||||
assert src.exists()
|
||||
assert not dst.exists()
|
||||
@ -259,3 +259,92 @@ def test_tags_prefix_treats_underscore_literal(
|
||||
assert tag_ok in names, f"Expected {tag_ok} to be returned for prefix '{base}_'"
|
||||
assert tag_bad not in names, f"'{tag_bad}' must not match — '_' is not a wildcard"
|
||||
assert body["total"] == 1
|
||||
|
||||
|
||||
def _upload_model(http, api_base, name, tags):
|
||||
"""Upload a fresh filesystem-backed model asset, returning the response dict."""
|
||||
content = uuid.uuid4().bytes + b"W" * (4096 - 16)
|
||||
files = {"file": (name, content, "application/octet-stream")}
|
||||
form = {"tags": json.dumps(tags), "name": name, "user_metadata": json.dumps({})}
|
||||
r = http.post(api_base + "/api/assets", files=files, data=form, timeout=120)
|
||||
body = r.json()
|
||||
assert r.status_code == 201, body
|
||||
return body
|
||||
|
||||
|
||||
def test_edit_type_moves_file_and_reregisters(
|
||||
http: requests.Session,
|
||||
api_base: str,
|
||||
comfy_tmp_base_dir,
|
||||
):
|
||||
"""BE-1641: a flag-on model_type: edit MOVEs the file to the new folder.
|
||||
|
||||
Mirrors the FE edit-type flow (remove old model_type:, add new) and asserts
|
||||
both the tag set and the on-disk location end up coherent.
|
||||
"""
|
||||
|
||||
name = f"edit_type_{uuid.uuid4().hex[:8]}.safetensors"
|
||||
asset = _upload_model(http, api_base, name, ["models", "model_type:checkpoints", "unit-tests"])
|
||||
aid = asset["id"]
|
||||
|
||||
models_dir = comfy_tmp_base_dir / "models"
|
||||
rel = asset["user_metadata"]["filename"]
|
||||
src_path = models_dir / "checkpoints" / rel
|
||||
dst_path = models_dir / "loras" / rel
|
||||
assert src_path.is_file(), f"uploaded model should be under checkpoints: {src_path}"
|
||||
|
||||
# FE edit-type: DELETE old model_type:, then POST the new one.
|
||||
rd = http.delete(
|
||||
f"{api_base}/api/assets/{aid}/tags",
|
||||
json={"tags": ["model_type:checkpoints"]},
|
||||
timeout=120,
|
||||
)
|
||||
assert rd.status_code == 200, rd.json()
|
||||
|
||||
ra = http.post(
|
||||
f"{api_base}/api/assets/{aid}/tags",
|
||||
json={"tags": ["model_type:loras"]},
|
||||
timeout=120,
|
||||
)
|
||||
assert ra.status_code == 200, ra.json()
|
||||
|
||||
# File relocated on disk, no stale copy left behind.
|
||||
assert dst_path.is_file(), f"file should have moved to loras: {dst_path}"
|
||||
assert not src_path.exists(), "stale checkpoints copy must be gone"
|
||||
|
||||
# Tag set is coherent with the new location.
|
||||
rg = http.get(f"{api_base}/api/assets/{aid}", timeout=120)
|
||||
detail = rg.json()
|
||||
assert rg.status_code == 200, detail
|
||||
tags = set(detail["tags"])
|
||||
assert "model_type:loras" in tags
|
||||
assert "model_type:checkpoints" not in tags
|
||||
assert "models" in tags
|
||||
|
||||
http.delete(f"{api_base}/api/assets/{aid}", timeout=30)
|
||||
|
||||
|
||||
def test_edit_type_unknown_folder_is_label_only(
|
||||
http: requests.Session,
|
||||
api_base: str,
|
||||
):
|
||||
"""An unregistered model_type: target stays a plain label (no move/reject).
|
||||
|
||||
Core stays permissive about model_type: labels; only a registered folder
|
||||
triggers the edit-type move. The file must not move.
|
||||
"""
|
||||
name = f"edit_bad_{uuid.uuid4().hex[:8]}.safetensors"
|
||||
asset = _upload_model(http, api_base, name, ["models", "model_type:checkpoints", "unit-tests"])
|
||||
aid = asset["id"]
|
||||
|
||||
r = http.post(
|
||||
f"{api_base}/api/assets/{aid}/tags",
|
||||
json={"tags": ["model_type:does_not_exist"]},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
assert "model_type:does_not_exist" in body["total_tags"]
|
||||
assert "model_type:checkpoints" in body["total_tags"]
|
||||
|
||||
http.delete(f"{api_base}/api/assets/{aid}", timeout=30)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user