diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 052bbc970ca..47771221bde 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -6,6 +6,7 @@ import torch from pydantic import BaseModel, Field +from invokeai.app.util.path_safety import is_plain_filename from invokeai.backend.ltx2 import checkpoint_layout as ltx2_layout from invokeai.backend.model_manager.checkpoint_prefix import COMFYUI_KEY_PREFIXES from invokeai.backend.model_manager.configs.base import ( @@ -1183,7 +1184,7 @@ def _get_submodels_or_raise(cls, mod: ModelOnDisk) -> dict[SubModelType, Submode match class_name: case "CLIPTextModelWithProjection": model_type = ModelType.CLIPEmbed - path_or_prefix = (mod.path / key).resolve().as_posix() + path_or_prefix = key # We need to read the config to determine the variant of the CLIP model. clip_embed_config = get_config_dict_or_raise( @@ -1200,7 +1201,7 @@ def _get_submodels_or_raise(cls, mod: ModelOnDisk) -> dict[SubModelType, Submode ) case "SD3Transformer2DModel": model_type = ModelType.Main - path_or_prefix = (mod.path / key).resolve().as_posix() + path_or_prefix = key variant = None submodels[SubModelType(key)] = SubmodelDefinition( path_or_prefix=path_or_prefix, @@ -3452,6 +3453,10 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti # its loader needs, otherwise a partial pipeline is treated as self-contained # (is_self_contained_sdnq_pipeline) and the loaders later request fixed vae/ text_encoder/ # tokenizer/ subfolders that have nothing to load. + # The record stores the key and loading joins it onto wherever the model ends up, so only a + # bare folder name keeps the folder validated here the same one that is loaded later. + if not is_plain_filename(key): + continue component_path = mod.path / key if not _sdnq_component_dir_is_populated(component_path, submodel_type, _SDNQ_PIPELINE_ALWAYS_SDNQ_LOADED): continue @@ -3466,7 +3471,7 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti continue submodels[submodel_type] = SubmodelDefinition( - path_or_prefix=component_path.resolve().as_posix(), + path_or_prefix=key, model_type=model_type, variant=None, ) @@ -3599,6 +3604,10 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti # See the FLUX.2 _get_submodels note: only record a component whose folder actually holds # the files its loader needs, so a partial download with a complete model_index.json isn't # mis-classified as a self-contained SDNQ pipeline. + # The record stores the key and loading joins it onto wherever the model ends up, so only a + # bare folder name keeps the folder validated here the same one that is loaded later. + if not is_plain_filename(key): + continue component_path = mod.path / key if not _sdnq_component_dir_is_populated(component_path, submodel_type, _SDNQ_PIPELINE_ALWAYS_SDNQ_LOADED): continue @@ -3612,7 +3621,7 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti continue submodels[submodel_type] = SubmodelDefinition( - path_or_prefix=component_path.resolve().as_posix(), + path_or_prefix=key, model_type=model_type, variant=None, ) @@ -3753,6 +3762,10 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti # *should* have, so a partial download keeps a complete index over missing or empty # folders, and the advertised class is a claim about the folder rather than a fact. # Without these, a half-downloaded pipeline looks complete until generation time. + # The record stores the key and loading joins it onto wherever the model ends up, so only a + # bare folder name keeps the folder validated here the same one that is loaded later. + if not is_plain_filename(key): + continue component_path = mod.path / key if not _sdnq_component_dir_is_populated(component_path, submodel_type, _SDNQ_FLUX1_ALWAYS_SDNQ_LOADED): continue @@ -3764,7 +3777,7 @@ def _get_submodels(cls, mod: ModelOnDisk) -> dict[SubModelType, SubmodelDefiniti continue submodels[submodel_type] = SubmodelDefinition( - path_or_prefix=component_path.resolve().as_posix(), + path_or_prefix=key, model_type=model_type, variant=None, ) diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 3fae1d5a77b..5f68261da37 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -101,9 +101,21 @@ def resolve_submodel_path(config: AnyModelConfig, submodel_type: SubModelType, f `fallback` is used when the config carries no such entry — configs persisted before submodel discovery existed, and the layouts where the component folder *is* the model path. + + The recorded path is relative to the model folder (`config.path`, absolute by load time). + Records written before that convention hold the absolute path seen at identification — for an + installed model, the staging or source folder it was moved out of. Its last component is the + component key, so the model's own folder of that name wins; the recorded path is kept only when + the model has no such folder (a component that was a symlink was recorded by its target). """ discovered = (getattr(config, "submodels", None) or {}).get(submodel_type) - return Path(discovered.path_or_prefix) if discovered else fallback + if not discovered: + return fallback + recorded = Path(discovered.path_or_prefix) + rebased = Path(config.path) / (recorded.name if recorded.is_absolute() else recorded) + if recorded.is_absolute() and recorded.exists() and not rebased.exists(): + return recorded + return rebased def _device_supports_fp8_storage(device: torch.device, logger: Optional[Logger] = None) -> bool: diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 2639a1d3f51..4350e717b28 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -3748,3 +3748,27 @@ def test_restore_skips_a_marker_with_an_invalid_status_without_claiming_it_or_ab for tmpdir in (bad_tmpdir, good_tmpdir): active_install_sentinel_path(tmpdir).unlink(missing_ok=True) shutil.rmtree(tmpdir, ignore_errors=True) + + +def test_installed_pipeline_components_resolve_inside_the_installed_model( + mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig +) -> None: + """Pipelines are identified where they are staged (a `tmpinstall_*` folder for downloads) and + then moved into the models dir; the stored record must lead loading to the moved components.""" + from invokeai.backend.model_manager.load.load_default import resolve_submodel_path + from invokeai.backend.model_manager.taxonomy import SubModelType + from tests.backend.model_manager.configs.test_sdnq_flux1_submodel_discovery import _make_flux1_pipeline + + staging = mm2_app_config.models_path / "tmpinstall_test" / "pipeline" + _make_flux1_pipeline(staging, component_keys={"text_encoder": "clip_encoder"}) + + record = mm2_installer.record_store.get_model(mm2_installer.install_path(staging)) + # Stored relative to the model, so nothing in the record points back at the staging folder. + assert record.submodels is not None + assert record.submodels[SubModelType.TextEncoder].path_or_prefix == "clip_encoder" + + # The loader makes `path` absolute against the models dir before resolving components. + record.path = (mm2_app_config.models_path / record.path).resolve().as_posix() + text_encoder = resolve_submodel_path(record, SubModelType.TextEncoder, Path(record.path) / "text_encoder") + assert text_encoder.parent == Path(record.path) + assert (text_encoder / "config.json").is_file() diff --git a/tests/backend/model_manager/configs/test_sdnq_flux1_submodel_discovery.py b/tests/backend/model_manager/configs/test_sdnq_flux1_submodel_discovery.py index 7a89942f578..53a3e128b2d 100644 --- a/tests/backend/model_manager/configs/test_sdnq_flux1_submodel_discovery.py +++ b/tests/backend/model_manager/configs/test_sdnq_flux1_submodel_discovery.py @@ -9,6 +9,7 @@ """ import json +import shutil from pathlib import Path from unittest.mock import MagicMock @@ -17,6 +18,7 @@ from safetensors.torch import save_file from invokeai.backend.model_manager.configs.main import Main_SDNQ_Diffusers_FLUX_Config +from invokeai.backend.model_manager.load.load_default import resolve_submodel_path from invokeai.backend.model_manager.taxonomy import SubModelType _REQUIRED_FIELDS = { @@ -183,13 +185,45 @@ def test_nonstandard_index_keys_are_recorded_at_the_path_they_actually_live_at(t assert config.submodels is not None assert set(config.submodels) == _ALL_SLOTS - assert config.submodels[SubModelType.TextEncoder].path_or_prefix.endswith("clip_encoder") - assert config.submodels[SubModelType.VAE].path_or_prefix.endswith("autoencoder") + assert config.submodels[SubModelType.TextEncoder].path_or_prefix == "clip_encoder" + assert config.submodels[SubModelType.VAE].path_or_prefix == "autoencoder" # And the conventional folders genuinely do not exist, so a reconstructing loader would miss. assert not (root / "text_encoder").exists() assert not (root / "vae").exists() +def test_an_index_key_that_escapes_the_pipeline_is_not_recorded(tmp_path: Path) -> None: + """The record keeps the key and loading joins it onto wherever the model now lives, so a key such + as `../vae` would be validated beside the staging folder and then loaded beside the installed one.""" + root = _make_flux1_pipeline(tmp_path / "pipeline", component_keys={"vae": "../vae"}) + assert (tmp_path / "vae" / "config.json").is_file() + + config = _discover(root) + + assert config.submodels is not None + assert SubModelType.VAE not in config.submodels + + +def test_recorded_components_are_found_after_the_install_moves_the_pipeline(tmp_path: Path) -> None: + """Downloads are identified in a `tmpinstall_*` staging folder and then moved into the models + dir; the record must lead the loader to the components in their final location, not to wherever + identification happened to look (which here is left behind, as a relocated models dir can be).""" + keys = {"text_encoder": "clip_encoder"} + staging = _make_flux1_pipeline(tmp_path / "tmpinstall_abc" / "pipeline", component_keys=keys) + config = _discover(staging) + + installed = tmp_path / "models" / "test-key" + # Copied, not moved, so the staging copy survives and only the record can say which one is ours. + shutil.copytree(staging, installed) + config.path = installed.as_posix() + + assert config.submodels is not None + assert config.submodels[SubModelType.TextEncoder].path_or_prefix == "clip_encoder" + te_dir = resolve_submodel_path(config, SubModelType.TextEncoder, installed / "text_encoder") + assert te_dir == installed / "clip_encoder" + assert (te_dir / "config.json").is_file() + + def test_a_serialized_submodel_map_does_not_survive_a_deleted_component(tmp_path: Path) -> None: """Rehydration must revalidate against current files. diff --git a/tests/backend/model_manager/configs/test_sdnq_pipeline_submodel_discovery.py b/tests/backend/model_manager/configs/test_sdnq_pipeline_submodel_discovery.py index 81d0b21afb1..f80e6ac917c 100644 --- a/tests/backend/model_manager/configs/test_sdnq_pipeline_submodel_discovery.py +++ b/tests/backend/model_manager/configs/test_sdnq_pipeline_submodel_discovery.py @@ -191,10 +191,13 @@ def _mod(root: Path) -> MagicMock: def _assert_complete_pipeline(config) -> None: assert config.submodels is not None - assert SubModelType.Transformer in config.submodels - assert SubModelType.TextEncoder in config.submodels - assert SubModelType.Tokenizer in config.submodels - assert SubModelType.VAE in config.submodels + # Recorded relative to the pipeline folder, so the record survives the install moving it. + assert {slot: d.path_or_prefix for slot, d in config.submodels.items()} == { + SubModelType.Transformer: "transformer", + SubModelType.TextEncoder: "text_encoder", + SubModelType.Tokenizer: "tokenizer", + SubModelType.VAE: "vae", + } assert is_self_contained_sdnq_pipeline(config) @@ -426,6 +429,30 @@ def test_sdnq_pipeline_tokenizer_without_a_declared_class_is_still_recorded(tmp_ _assert_complete_pipeline(config) +@pytest.mark.parametrize( + ("factory", "root_name"), + [ + (Main_SDNQ_Diffusers_Flux2_Config, "flux2-escaping-key"), + (Main_SDNQ_Diffusers_ZImage_Config, "zimage-escaping-key"), + ], +) +def test_sdnq_pipeline_index_key_that_escapes_the_pipeline_is_not_recorded(tmp_path: Path, factory, root_name: str): + """The record keeps the key and loading joins it onto wherever the model now lives, so `../vae` + would be validated beside the staging folder and then loaded beside the installed one.""" + maker = _make_flux2_pipeline if factory is Main_SDNQ_Diffusers_Flux2_Config else _make_zimage_pipeline + root = maker(tmp_path / root_name, "Qwen3ForCausalLM") + index_path = root / "model_index.json" + index = json.loads(index_path.read_text(encoding="utf-8")) + index["../vae"] = index.pop("vae") + index_path.write_text(json.dumps(index), encoding="utf-8") + (root / "vae").rename(root.parent / "vae") + + config = factory.from_model_on_disk(_mod(root), {**_REQUIRED_FIELDS, "path": root.as_posix()}) + + assert config.submodels is not None + assert SubModelType.VAE not in config.submodels + + @pytest.mark.parametrize("suffix", [".gguf", ".bin", ".pt"]) @pytest.mark.parametrize( ("factory", "root_name"), diff --git a/tests/backend/model_manager/load/test_load_default_helpers.py b/tests/backend/model_manager/load/test_load_default_helpers.py index 0346d448cd2..de78cb1c695 100644 --- a/tests/backend/model_manager/load/test_load_default_helpers.py +++ b/tests/backend/model_manager/load/test_load_default_helpers.py @@ -50,16 +50,89 @@ class TestResolveSubmodelPath: def _config(self, submodels: dict | None) -> SimpleNamespace: return SimpleNamespace(path="/models/pipeline", submodels=submodels) - def test_it_prefers_the_path_discovery_recorded(self) -> None: - """`model_index.json` may name a component anything; discovery records the folder it found.""" - config = self._config( - {SubModelType.TextEncoder: SimpleNamespace(path_or_prefix="/models/pipeline/clip_encoder")} - ) + def test_a_relative_record_follows_the_model_folder(self) -> None: + """Components are recorded relative to the model, so moving the model moves them too.""" + config = self._config({SubModelType.TextEncoder: SimpleNamespace(path_or_prefix="clip_encoder")}) resolved = resolve_submodel_path(config, SubModelType.TextEncoder, Path("/models/pipeline/text_encoder")) assert resolved == Path("/models/pipeline/clip_encoder") + def test_it_rebases_a_legacy_absolute_record_after_the_pipeline_is_moved(self, tmp_path: Path) -> None: + """Older records hold the identification-time path: the download's `tmpinstall_*` staging dir.""" + staging_root = tmp_path / "tmpinstall_pipeline" + discovered_component = staging_root / "clip_encoder" + discovered_component.mkdir(parents=True) + + installed_root = tmp_path / "models" / "pipeline" + installed_root.parent.mkdir() + staging_root.rename(installed_root) + config = SimpleNamespace( + # The loader makes `config.path` absolute before any component is resolved. + path=installed_root.as_posix(), + submodels={SubModelType.TextEncoder: SimpleNamespace(path_or_prefix=discovered_component.as_posix())}, + ) + + resolved = resolve_submodel_path(config, SubModelType.TextEncoder, installed_root / "text_encoder") + + assert resolved == installed_root / "clip_encoder" + + def test_a_legacy_absolute_record_loses_to_the_models_own_component(self, tmp_path: Path) -> None: + """A copy left at the recorded location (an install source, a copied models dir) may have + been replaced since; the model's own folder is the one that belongs to this record.""" + elsewhere = tmp_path / "Downloads" / "pipeline" / "text_encoder" + own = tmp_path / "models" / "key" / "text_encoder" + elsewhere.mkdir(parents=True) + own.mkdir(parents=True) + config = SimpleNamespace( + path=str(tmp_path / "models" / "key"), + submodels={SubModelType.TextEncoder: SimpleNamespace(path_or_prefix=elsewhere.as_posix())}, + ) + + assert resolve_submodel_path(config, SubModelType.TextEncoder, Path("/unused")) == own + + def test_a_legacy_record_of_a_symlinked_component_keeps_its_target(self, tmp_path: Path) -> None: + """Identification resolved symlinks, so the record names the target, which the model folder lacks.""" + target = tmp_path / "shared" / "te_weights" + target.mkdir(parents=True) + model = tmp_path / "pipeline" + model.mkdir() + (model / "text_encoder").symlink_to(target) + config = SimpleNamespace( + path=str(model), + submodels={SubModelType.TextEncoder: SimpleNamespace(path_or_prefix=target.as_posix())}, + ) + + assert resolve_submodel_path(config, SubModelType.TextEncoder, Path("/unused")) == target + + def test_the_model_folder_is_the_root_whatever_the_fallback(self, tmp_path: Path) -> None: + """The transformer loaders fall back to the model folder itself when it has no `transformer/`, + so the fallback's parent is the models dir; the root must come from the config instead.""" + (tmp_path / "models" / "key" / "dit").mkdir(parents=True) + (tmp_path / "models" / "dit").mkdir() # a namesake beside the model must not be picked up + model = tmp_path / "models" / "key" + stale = (tmp_path / "models" / "tmpinstall_x" / "pipeline" / "dit").as_posix() + config = SimpleNamespace( + path=str(model), submodels={SubModelType.Transformer: SimpleNamespace(path_or_prefix=stale)} + ) + + assert resolve_submodel_path(config, SubModelType.Transformer, model) == model / "dit" + + def test_a_missing_component_is_reported_inside_the_model_folder(self, tmp_path: Path) -> None: + """When neither location exists, the failure names the installed model's folder rather than a + long-gone staging path, which Transformers would misreport as an invalid Hub repo id.""" + missing_component = tmp_path / "tmpinstall_x" / "clip_encoder" + installed_root = tmp_path / "pipeline" + installed_root.mkdir() + config = SimpleNamespace( + path=installed_root.as_posix(), + submodels={SubModelType.TextEncoder: SimpleNamespace(path_or_prefix=missing_component.as_posix())}, + ) + + resolved = resolve_submodel_path(config, SubModelType.TextEncoder, installed_root / "text_encoder") + + assert resolved == installed_root / "clip_encoder" + def test_it_falls_back_when_the_slot_was_not_discovered(self) -> None: """Configs persisted before submodel discovery existed carry no map.""" fallback = Path("/models/pipeline/text_encoder") @@ -105,9 +178,7 @@ def test_it_sizes_a_nonstandard_component_key(self, tmp_path: Path) -> None: config = SimpleNamespace( path=str(tmp_path), repo_variant=None, - submodels={ - SubModelType.TextEncoder: SimpleNamespace(path_or_prefix=(tmp_path / "clip_encoder").as_posix()) - }, + submodels={SubModelType.TextEncoder: SimpleNamespace(path_or_prefix="clip_encoder")}, ) size = self._loader().get_size_fs(config, tmp_path, SubModelType.TextEncoder)