diff --git a/py/services/checkpoint_scanner.py b/py/services/checkpoint_scanner.py index af4318e5..02402798 100644 --- a/py/services/checkpoint_scanner.py +++ b/py/services/checkpoint_scanner.py @@ -410,6 +410,10 @@ class CheckpointScanner(ModelScanner): return None + def resolve_sub_type_for_path(self, file_path: Optional[str]) -> Optional[str]: + """Resolve sub_type from the configured root that contains the file.""" + return self._resolve_sub_type(self._find_root_for_file(file_path)) + def adjust_metadata(self, metadata, file_path, root_path): """Adjust metadata during scanning to set sub_type.""" sub_type = self._resolve_sub_type(root_path) @@ -419,9 +423,7 @@ class CheckpointScanner(ModelScanner): def adjust_cached_entry(self, entry: Dict[str, Any]) -> Dict[str, Any]: """Adjust entries loaded from the persisted cache to ensure sub_type is set.""" - sub_type = self._resolve_sub_type( - self._find_root_for_file(entry.get("file_path")) - ) + sub_type = self.resolve_sub_type_for_path(entry.get("file_path")) if sub_type: entry["sub_type"] = sub_type return entry diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 4df3a296..22555254 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -1339,6 +1339,14 @@ class ModelScanner: """Hook for subclasses: adjust entries loaded from the persisted cache.""" return entry + def resolve_sub_type_for_path(self, file_path: Optional[str]) -> Optional[str]: + """Hook for subclasses: resolve the location-derived sub_type for a file. + + Returns ``None`` when the model type has no location-derived sub-types + (the default), in which case any stored value is left untouched. + """ + return None + @staticmethod def _normalize_path_value(path: Optional[str]) -> str: if not path: @@ -1869,6 +1877,20 @@ class ModelScanner: except Exception as e: logger.error(f"Error moving metadata file: {e}") + if metadata is not None: + # sub_type is derived from the model's location (e.g. a file + # moved from a checkpoints root into a unet root becomes a + # diffusion_model). Persist the recalculated value into the + # moved metadata file so later metadata-driven cache syncs + # do not revert the cache entry to the stale sub_type. + new_sub_type = self.resolve_sub_type_for_path(target_file) + if new_sub_type and metadata.get('sub_type') != new_sub_type: + metadata['sub_type'] = new_sub_type + try: + await MetadataManager.save_metadata(moved_metadata_path, metadata) + except Exception as e: + logger.error(f"Error persisting sub_type for moved model: {e}") + update_result = await self.update_single_model_cache(source_path, target_file, metadata, recalculate_type=True) return { @@ -2064,6 +2086,11 @@ class ModelScanner: file_path_override=file_path, ) + # Location-derived fields (e.g. the checkpoint sub_type) must be + # re-resolved from the file path rather than trusting the on-disk + # metadata snapshot, which may predate a cross-root move. + desired_entry = self.adjust_cached_entry(desired_entry) + # Ensure sha256 is populated (defensive — metadata should have it) if ( not desired_entry.get("sha256") diff --git a/static/js/managers/MoveManager.js b/static/js/managers/MoveManager.js index f6d55f34..da1656ca 100644 --- a/static/js/managers/MoveManager.js +++ b/static/js/managers/MoveManager.js @@ -329,7 +329,11 @@ class MoveManager { const results = await apiClient.moveBulkModels(this.bulkFilePaths, targetPath, this.useDefaultPath); movedFiles = (results || []) .filter(r => r.success) - .map(r => ({ original_file_path: r.original_file_path, new_file_path: r.new_file_path })); + .map(r => ({ + original_file_path: r.original_file_path, + new_file_path: r.new_file_path, + sub_type: r.cache_entry?.sub_type + })); // Deselect moving items and exit bulk mode this.bulkFilePaths.forEach(path => bulkManager.deselectItem(path)); @@ -340,7 +344,11 @@ class MoveManager { if (result) { movedFiles.push({ original_file_path: result.original_file_path || this.currentFilePath, - new_file_path: result.new_file_path + new_file_path: result.new_file_path, + // The backend recalculates location-derived fields + // (e.g. checkpoint -> diffusion_model) during the move; + // carry them so the card re-renders with the new type. + sub_type: result.cache_entry?.sub_type }); } @@ -379,24 +387,28 @@ class MoveManager { } if (stillVisible) { + const newData = { + file_path: moved.new_file_path, + folder: newRelativeFolder + }; + if (moved.sub_type) newData.sub_type = moved.sub_type; pathsToUpdate.push({ originalPath: moved.original_file_path, - newData: { - file_path: moved.new_file_path, - folder: newRelativeFolder - } + newData }); } else { pathsToRemove.push(moved.original_file_path); } } else { // No folder filter active — items remain visible, just update path + const newData = { + file_path: moved.new_file_path, + folder: this._getRelativeFolder(moved.new_file_path) + }; + if (moved.sub_type) newData.sub_type = moved.sub_type; pathsToUpdate.push({ originalPath: moved.original_file_path, - newData: { - file_path: moved.new_file_path, - folder: this._getRelativeFolder(moved.new_file_path) - } + newData }); } } diff --git a/tests/frontend/managers/MoveManager.test.js b/tests/frontend/managers/MoveManager.test.js index 1fd78b66..912f6e5c 100644 --- a/tests/frontend/managers/MoveManager.test.js +++ b/tests/frontend/managers/MoveManager.test.js @@ -17,7 +17,8 @@ vi.mock('../../../static/js/state/index.js', () => ({ } } } - } + }, + getCurrentPageState: vi.fn(() => ({ activeFolder: null, searchOptions: {} })) })); vi.mock('../../../static/js/managers/ModalManager.js', () => ({ @@ -162,4 +163,75 @@ describe('MoveManager', () => { true ); }); + + it('should propagate the recalculated sub_type from the move response to the card', async () => { + // Setup state: moving a checkpoint into the unet root + moveManager.useDefaultPath = false; + moveManager.bulkFilePaths = null; + moveManager.currentFilePath = '/models/checkpoints/model.safetensors'; + moveManager.modelRoots = ['/models/checkpoints', '/models/unet']; + document.getElementById('moveModelRoot').innerHTML = ''; + document.getElementById('moveModelRoot').value = '/models/unet'; + moveManager.folderTreeManager.selectedPath = ''; + + const updateSingleItem = vi.fn(); + state.virtualScroller = { + updateSingleItem, + removeMultipleItemsByFilePath: vi.fn() + }; + + mockApiClient.moveSingleModel = vi.fn().mockResolvedValue({ + success: true, + original_file_path: '/models/checkpoints/model.safetensors', + new_file_path: '/models/unet/model.safetensors', + cache_entry: { sub_type: 'diffusion_model' } + }); + + try { + await moveManager.moveModel(); + + expect(updateSingleItem).toHaveBeenCalledWith( + '/models/checkpoints/model.safetensors', + expect.objectContaining({ + file_path: '/models/unet/model.safetensors', + sub_type: 'diffusion_model' + }) + ); + } finally { + delete state.virtualScroller; + } + }); + + it('should omit sub_type from the card update when the response has no cache entry', async () => { + moveManager.useDefaultPath = false; + moveManager.bulkFilePaths = null; + moveManager.currentFilePath = '/models/loras/a.safetensors'; + moveManager.modelRoots = ['/models/loras']; + document.getElementById('moveModelRoot').innerHTML = ''; + document.getElementById('moveModelRoot').value = '/models/loras'; + moveManager.folderTreeManager.selectedPath = ''; + + const updateSingleItem = vi.fn(); + state.virtualScroller = { + updateSingleItem, + removeMultipleItemsByFilePath: vi.fn() + }; + + mockApiClient.moveSingleModel = vi.fn().mockResolvedValue({ + success: true, + original_file_path: '/models/loras/a.safetensors', + new_file_path: '/models/loras/b/a.safetensors' + }); + + try { + await moveManager.moveModel(); + + expect(updateSingleItem).toHaveBeenCalledWith( + '/models/loras/a.safetensors', + expect.not.objectContaining({ sub_type: expect.anything() }) + ); + } finally { + delete state.virtualScroller; + } + }); }); diff --git a/tests/services/test_checkpoint_scanner_sub_type.py b/tests/services/test_checkpoint_scanner_sub_type.py index bb723d7a..6dff90d5 100644 --- a/tests/services/test_checkpoint_scanner_sub_type.py +++ b/tests/services/test_checkpoint_scanner_sub_type.py @@ -1,10 +1,15 @@ """Tests for CheckpointScanner sub_type resolution.""" +import json +import os import pytest import asyncio +from pathlib import Path from unittest.mock import MagicMock, patch from py.services.checkpoint_scanner import CheckpointScanner +from py.services.model_cache import ModelCache +from py.services.model_hash_index import ModelHashIndex from py.utils.models import CheckpointMetadata @@ -142,3 +147,150 @@ class TestCheckpointScannerSubType: config_module.config.checkpoints_roots = original_checkpoints_roots if original_unet_roots is not None: config_module.config.unet_roots = original_unet_roots + + +def _make_move_scanner(ckpt_root: Path, unet_root: Path) -> CheckpointScanner: + """Create a CheckpointScanner wired for move/sync tests without async init.""" + scanner = object.__new__(CheckpointScanner) + scanner.model_type = "checkpoint" + scanner.model_class = CheckpointMetadata + scanner.file_extensions = {".safetensors"} + scanner._cache = None + scanner._cache_version = 0 + scanner._hash_index = ModelHashIndex() + scanner._tags_count = {} + scanner._excluded_models = [] + scanner._is_initializing = False + scanner._persistent_cache = MagicMock() + scanner._name_display_mode = "model_name" + scanner._cancel_requested = False + scanner._all_folders_ttl_cache = None + roots = [str(ckpt_root), str(unet_root)] + scanner.get_model_roots = lambda: roots + return scanner + + +def _set_config_roots(monkeypatch, ckpt_root: Path, unet_root: Path) -> None: + from py import config as config_module + + monkeypatch.setattr( + config_module.config, "checkpoints_roots", [str(ckpt_root)] + ) + monkeypatch.setattr(config_module.config, "unet_roots", [str(unet_root)]) + monkeypatch.setattr(config_module.config, "extra_checkpoints_roots", []) + monkeypatch.setattr(config_module.config, "extra_unet_roots", []) + + +def _write_model(root: Path, name: str, sub_type: str) -> str: + model_path = root / f"{name}.safetensors" + model_path.write_bytes(b"fake") + (root / f"{name}.metadata.json").write_text( + json.dumps( + { + "file_path": str(model_path).replace(os.sep, "/"), + "file_name": name, + "model_name": name, + "sha256": "abc123", + "sub_type": sub_type, + "hash_status": "completed", + "tags": [], + } + ) + ) + return str(model_path).replace(os.sep, "/") + + +@pytest.mark.asyncio +async def test_move_to_unet_root_updates_sub_type_in_cache_and_metadata( + tmp_path, monkeypatch +): + """Moving a checkpoint into a unet root must recalculate sub_type and + persist it into the moved .metadata.json, so later metadata-driven cache + syncs cannot revert the cache entry to the stale sub_type.""" + ckpt_root = tmp_path / "checkpoints" + unet_root = tmp_path / "unet" + ckpt_root.mkdir() + unet_root.mkdir() + _set_config_roots(monkeypatch, ckpt_root, unet_root) + + scanner = _make_move_scanner(ckpt_root, unet_root) + source = _write_model(ckpt_root, "mymodel", "checkpoint") + + scanner._cache = ModelCache( + raw_data=[ + { + "file_path": source, + "file_name": "mymodel", + "model_name": "mymodel", + "folder": "", + "sha256": "abc123", + "sub_type": "checkpoint", + "tags": [], + } + ], + folders=[""], + ) + + result = await scanner.move_model(source, str(unet_root).replace(os.sep, "/")) + assert result is not None + + cache = await scanner.get_cached_data() + entry = next( + (e for e in cache.raw_data if e.get("file_name") == "mymodel"), None + ) + assert entry is not None + assert entry["sub_type"] == "diffusion_model" + + moved_metadata = json.loads( + (unet_root / "mymodel.metadata.json").read_text() + ) + assert moved_metadata["sub_type"] == "diffusion_model" + + +@pytest.mark.asyncio +async def test_sync_cache_from_metadata_does_not_revert_sub_type( + tmp_path, monkeypatch +): + """An opportunistic sync from a stale .metadata.json (sub_type predating a + cross-root move) must not overwrite the location-derived cache sub_type.""" + ckpt_root = tmp_path / "checkpoints" + unet_root = tmp_path / "unet" + ckpt_root.mkdir() + unet_root.mkdir() + _set_config_roots(monkeypatch, ckpt_root, unet_root) + + scanner = _make_move_scanner(ckpt_root, unet_root) + file_path = _write_model(unet_root, "mymodel", "diffusion_model") + + scanner._cache = ModelCache( + raw_data=[ + { + "file_path": file_path, + "file_name": "mymodel", + "model_name": "mymodel", + "folder": "", + "sha256": "abc123", + "sub_type": "diffusion_model", + "tags": [], + } + ], + folders=[""], + ) + + # Stale metadata snapshot: still says 'checkpoint' (as before a move). + stale_metadata = { + "file_path": file_path, + "file_name": "mymodel", + "model_name": "mymodel Renamed", + "sha256": "abc123", + "sub_type": "checkpoint", + "hash_status": "completed", + "tags": [], + } + + changed = await scanner.sync_cache_from_metadata(file_path, stale_metadata) + assert changed is True # other fields (model_name) did change + + entry = scanner._cache.raw_data[0] + assert entry["sub_type"] == "diffusion_model" + assert entry["model_name"] == "mymodel Renamed"