mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -03:00
feat(recipes): add manual checkpoint reconnect for broken recipe entries
Checkpoint entries that cannot be restored by download (deleted, unresolvable hash, or name-only remnants with no CivitAI identifiers) now get the same remediation chain LoRAs already had: - scanner: parameterized reconnect-suggestion ranking, update/restore/ set-hash-invalid for the checkpoint entry, and clear hashInvalid on rematch write-back (was only done for LoRAs) - persistence/handlers/routes: reconnect/restore/reconnect-suggestions/ mark-hash-invalid endpoints under /api/lm/recipe/checkpoint/* - modal: checkpoint reconnect UI (deleted/hash-invalid badges, inline form with suggestions, undo for reconnected entries); download failures mark the hash invalid only on explicit unresolvable signals (not found/deleted/404/410), matching the LoRA rule - css: checkpoint undo button shares the LoRA undo styles - i18n: the 14 new keys translated in all 9 locales
This commit is contained in:
Will Miao
@@ -95,6 +95,7 @@ vi.mock('../../../static/js/api/apiConfig.js', () => ({
|
||||
vi.mock('../../../static/js/managers/DownloadManager.js', () => ({
|
||||
downloadManager: {
|
||||
downloadVersionWithDefaults: downloadVersionWithDefaultsMock,
|
||||
_lastDownloadError: '',
|
||||
},
|
||||
}));
|
||||
|
||||
@@ -776,4 +777,321 @@ describe('RecipeModal resource item interactions', () => {
|
||||
lora_index: '0',
|
||||
});
|
||||
});
|
||||
|
||||
describe('checkpoint reconnect', () => {
|
||||
const brokenCheckpoint = {
|
||||
name: 'gone-checkpoint',
|
||||
file_name: 'gone',
|
||||
inLibrary: false,
|
||||
isDeleted: true,
|
||||
hash: 'a2a12bfa01',
|
||||
};
|
||||
const hashInvalidCheckpoint = {
|
||||
name: 'invalid-checkpoint',
|
||||
file_name: 'invalid',
|
||||
inLibrary: false,
|
||||
hashInvalid: true,
|
||||
hash: 'deadbeefcafe',
|
||||
};
|
||||
|
||||
function recipeWithCheckpoint(checkpoint) {
|
||||
return {
|
||||
...JSON.parse(JSON.stringify(recipeWithResources)),
|
||||
checkpoint: { ...checkpoint },
|
||||
};
|
||||
}
|
||||
|
||||
// Hydration re-fetches the recipe right after render and re-renders the
|
||||
// modal, so the mock must resolve the SAME broken-checkpoint recipe —
|
||||
// otherwise the fetch wipes isDeleted/hashInvalid back to the fixture.
|
||||
async function renderBrokenCheckpoint(recipeModal, checkpoint) {
|
||||
const isolated = recipeWithCheckpoint(checkpoint);
|
||||
fetchRecipeDetailsMock.mockResolvedValue(isolated);
|
||||
recipeModal.showRecipeDetails(isolated);
|
||||
await flushWiring();
|
||||
}
|
||||
|
||||
function mockCheckpointSuggestionsFetch(payload) {
|
||||
const requests = [];
|
||||
global.fetch = vi.fn(async (url, options) => {
|
||||
requests.push({ url: String(url), options });
|
||||
if (String(url).includes('/checkpoint/reconnect-suggestions')) {
|
||||
return { ok: true, json: async () => payload };
|
||||
}
|
||||
if (String(url).includes('/recipe/checkpoint/reconnect')) {
|
||||
return {
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
success: true,
|
||||
updated_checkpoint: {
|
||||
name: 'main-checkpoint',
|
||||
file_name: 'main',
|
||||
inLibrary: true,
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
return { ok: true, json: async () => ({}) };
|
||||
});
|
||||
return requests;
|
||||
}
|
||||
|
||||
it('renders a deleted checkpoint with a badge and reconnect affordance', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
await renderBrokenCheckpoint(recipeModal, brokenCheckpoint);
|
||||
|
||||
const item = document.querySelector('.checkpoint-item');
|
||||
expect(item.classList.contains('is-deleted')).toBe(true);
|
||||
expect(item.querySelector('.deleted-badge')).not.toBeNull();
|
||||
const reconnectButton = item.querySelector('.checkpoint-reconnect');
|
||||
expect(reconnectButton).not.toBeNull();
|
||||
// Deleted checkpoints lose the civitai link (their source page is gone)
|
||||
expect(item.querySelector('.recipe-lora-title a.recipe-civitai-link')).toBeNull();
|
||||
// The inline form is present but hidden until the button is pressed
|
||||
const container = item.querySelector('.lora-reconnect-container[data-lora-index="checkpoint"]');
|
||||
expect(container).not.toBeNull();
|
||||
expect(container.classList.contains('active')).toBe(false);
|
||||
});
|
||||
|
||||
it('renders a hash-invalid checkpoint with the unresolvable hash badge', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
await renderBrokenCheckpoint(recipeModal, hashInvalidCheckpoint);
|
||||
|
||||
const item = document.querySelector('.checkpoint-item');
|
||||
expect(item.querySelector('.invalid-hash-badge')).not.toBeNull();
|
||||
expect(item.querySelector('.checkpoint-reconnect')).not.toBeNull();
|
||||
});
|
||||
|
||||
it('renders reconnect for a name-only checkpoint with no download identifiers', async () => {
|
||||
// Importers can leave a checkpoint entry with nothing but a model name
|
||||
// (no hash / version id, so nothing was ever queryable on CivitAI).
|
||||
// It cannot be downloaded and is not marked deleted — reconnect is the
|
||||
// only remediation, so it must still surface.
|
||||
const recipeModal = await createRecipeModal();
|
||||
await renderBrokenCheckpoint(recipeModal, {
|
||||
type: 'checkpoint',
|
||||
modelName: 'meichidarkMix_meichidarkanimxlV1',
|
||||
inLibrary: false,
|
||||
});
|
||||
|
||||
const item = document.querySelector('.checkpoint-item');
|
||||
expect(item.querySelector('.checkpoint-download')).toBeNull();
|
||||
const reconnectButton = item.querySelector('.checkpoint-reconnect');
|
||||
expect(reconnectButton).not.toBeNull();
|
||||
const container = item.querySelector('.lora-reconnect-container[data-lora-index="checkpoint"]');
|
||||
expect(container).not.toBeNull();
|
||||
|
||||
reconnectButton.click();
|
||||
expect(container.classList.contains('active')).toBe(true);
|
||||
});
|
||||
|
||||
it('fetches checkpoint suggestions against the checkpoint endpoint', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const suggestionsPayload = {
|
||||
success: true,
|
||||
suggestions: [
|
||||
{
|
||||
file_name: 'main-checkpoint.safetensors',
|
||||
base_model: 'SD 1.5',
|
||||
preview_url: '/preview/main.png',
|
||||
score: 0.95,
|
||||
match_reason: 'same_version',
|
||||
target_name: 'main-checkpoint',
|
||||
},
|
||||
],
|
||||
};
|
||||
mockCheckpointSuggestionsFetch(suggestionsPayload);
|
||||
|
||||
await renderBrokenCheckpoint(recipeModal, brokenCheckpoint);
|
||||
document.querySelector('.checkpoint-reconnect').click();
|
||||
const container = document.querySelector('.lora-reconnect-container[data-lora-index="checkpoint"]');
|
||||
|
||||
expect(container.classList.contains('active')).toBe(true);
|
||||
expect(global.fetch).toHaveBeenCalledWith(
|
||||
'/api/lm/recipe/recipe-resources/checkpoint/reconnect-suggestions'
|
||||
);
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(container.querySelectorAll('.reconnect-suggestion').length).toBe(1);
|
||||
});
|
||||
expect(container.querySelector('.reconnect-suggestion-name').textContent)
|
||||
.toBe('main-checkpoint');
|
||||
});
|
||||
|
||||
it('reconnects the checkpoint via its own endpoint when a suggestion is clicked', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const requests = mockCheckpointSuggestionsFetch({
|
||||
success: true,
|
||||
suggestions: [
|
||||
{
|
||||
file_name: 'main-checkpoint.safetensors',
|
||||
target_name: 'main-checkpoint',
|
||||
match_reason: 'same_version',
|
||||
score: 0.95,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
await renderBrokenCheckpoint(recipeModal, brokenCheckpoint);
|
||||
document.querySelector('.checkpoint-reconnect').click();
|
||||
const container = document.querySelector('.lora-reconnect-container[data-lora-index="checkpoint"]');
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(container.querySelectorAll('.reconnect-suggestion').length).toBe(1);
|
||||
});
|
||||
container.querySelector('.reconnect-suggestion').click();
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(requests.some(r => r.url === '/api/lm/recipe/checkpoint/reconnect')).toBe(true);
|
||||
});
|
||||
const reconnectRequest = requests.find(r => r.url === '/api/lm/recipe/checkpoint/reconnect');
|
||||
expect(reconnectRequest.options.method).toBe('POST');
|
||||
expect(JSON.parse(reconnectRequest.options.body)).toEqual({
|
||||
recipe_id: 'recipe-resources',
|
||||
target_name: 'main-checkpoint',
|
||||
});
|
||||
await vi.waitFor(() => {
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.recipes.checkpointReconnectedSuccessfully',
|
||||
{},
|
||||
'success'
|
||||
);
|
||||
});
|
||||
expect(recipeModal.currentRecipe.checkpoint.inLibrary).toBe(true);
|
||||
});
|
||||
|
||||
it('warns when the checkpoint reconnect crossed base-model families', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
global.fetch = vi.fn(async (url) => {
|
||||
if (String(url).includes('/checkpoint/reconnect-suggestions')) {
|
||||
return { ok: true, json: async () => ({ success: true, suggestions: [] }) };
|
||||
}
|
||||
if (String(url).includes('/recipe/checkpoint/reconnect')) {
|
||||
return {
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
success: true,
|
||||
updated_checkpoint: { name: 'main', inLibrary: true },
|
||||
base_model_mismatch: { recipe_base_model: 'Illustrious', checkpoint_base_model: 'Pony' },
|
||||
}),
|
||||
};
|
||||
}
|
||||
return { ok: true, json: async () => ({}) };
|
||||
});
|
||||
|
||||
await renderBrokenCheckpoint(recipeModal, brokenCheckpoint);
|
||||
document.querySelector('.checkpoint-reconnect').click();
|
||||
const container = document.querySelector('.lora-reconnect-container[data-lora-index="checkpoint"]');
|
||||
const input = container.querySelector('.reconnect-input');
|
||||
input.value = 'main';
|
||||
container.querySelector('.reconnect-confirm-btn').click();
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.recipes.reconnectCheckpointBaseModelMismatch',
|
||||
{ recipe: 'Illustrious', checkpoint: 'Pony' },
|
||||
'warning'
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('offers undo for a reconnected checkpoint and restores via the API', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const isolatedRecipe = recipeWithCheckpoint(brokenCheckpoint);
|
||||
isolatedRecipe.checkpoint = {
|
||||
name: 'main-checkpoint',
|
||||
file_name: 'main',
|
||||
inLibrary: true,
|
||||
reconnectSnapshot: { name: 'gone-checkpoint', file_name: 'gone', isDeleted: true },
|
||||
};
|
||||
fetchRecipeDetailsMock.mockResolvedValue(isolatedRecipe);
|
||||
const requests = [];
|
||||
global.fetch = vi.fn(async (url, options) => {
|
||||
requests.push({ url: String(url), options });
|
||||
if (String(url).includes('/recipe/checkpoint/restore')) {
|
||||
return {
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
success: true,
|
||||
updated_checkpoint: { name: 'gone', inLibrary: false, isDeleted: true },
|
||||
}),
|
||||
};
|
||||
}
|
||||
return { ok: true, json: async () => ({}) };
|
||||
});
|
||||
recipeModal.showRecipeDetails(isolatedRecipe);
|
||||
await flushWiring();
|
||||
|
||||
const item = document.querySelector('.checkpoint-item');
|
||||
const undoButton = item.querySelector('.checkpoint-undo-reconnect');
|
||||
expect(undoButton).not.toBeNull();
|
||||
|
||||
undoButton.click();
|
||||
await vi.waitFor(() => {
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.recipes.checkpointRestored', {}, 'success');
|
||||
});
|
||||
const restoreRequest = requests.find(r => r.url === '/api/lm/recipe/checkpoint/restore');
|
||||
expect(restoreRequest.options.method).toBe('POST');
|
||||
expect(JSON.parse(restoreRequest.options.body)).toEqual({
|
||||
recipe_id: 'recipe-resources',
|
||||
});
|
||||
});
|
||||
|
||||
it('marks the checkpoint hash invalid only when the failure is unresolvable', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const { downloadManager } = await import('../../../static/js/managers/DownloadManager.js');
|
||||
const requests = [];
|
||||
global.fetch = vi.fn(async (url, options) => {
|
||||
requests.push({ url: String(url), options });
|
||||
return { ok: true, json: async () => ({ success: true }) };
|
||||
});
|
||||
|
||||
// Explicit "model removed" signal: the entry becomes a rematch/
|
||||
// reconnect candidate (same rule as the LoRA resolve "not found").
|
||||
// Use an isolated copy so the hashInvalid mutation does not leak into
|
||||
// the shared recipeWithResources fixture used by later tests.
|
||||
const isolatedRecipe = JSON.parse(JSON.stringify(recipeWithResources));
|
||||
fetchRecipeDetailsMock.mockResolvedValue(isolatedRecipe);
|
||||
downloadVersionWithDefaultsMock.mockResolvedValue(false);
|
||||
downloadManager._lastDownloadError = 'Model not found';
|
||||
recipeModal.showRecipeDetails(isolatedRecipe);
|
||||
await flushWiring();
|
||||
document.querySelector('.checkpoint-download').click();
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(requests.some(r => r.url === '/api/lm/recipe/checkpoint/mark-hash-invalid')).toBe(true);
|
||||
});
|
||||
const markRequest = requests.find(r => r.url === '/api/lm/recipe/checkpoint/mark-hash-invalid');
|
||||
expect(markRequest.options.method).toBe('POST');
|
||||
expect(JSON.parse(markRequest.options.body)).toEqual({ recipe_id: 'recipe-resources' });
|
||||
expect(recipeModal.currentRecipe.checkpoint.hashInvalid).toBe(true);
|
||||
});
|
||||
|
||||
it('does not mark the checkpoint hash invalid on transient download failures', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const { downloadManager } = await import('../../../static/js/managers/DownloadManager.js');
|
||||
const requests = [];
|
||||
global.fetch = vi.fn(async (url, options) => {
|
||||
requests.push({ url: String(url), options });
|
||||
return { ok: true, json: async () => ({ success: true }) };
|
||||
});
|
||||
|
||||
// Transport/API exceptions must NOT enroll the entry in the
|
||||
// remediation flow — transient failures are not evidence the model is
|
||||
// unrecoverable (mirrors the LoRA path).
|
||||
downloadVersionWithDefaultsMock.mockRejectedValue(new Error('Network timeout'));
|
||||
recipeModal.showRecipeDetails(recipeWithResources);
|
||||
await flushWiring();
|
||||
document.querySelector('.checkpoint-download').click();
|
||||
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
expect(requests.some(r => r.url === '/api/lm/recipe/checkpoint/mark-hash-invalid')).toBe(false);
|
||||
|
||||
// Business failure without an unresolvable signal also stays untouched.
|
||||
downloadVersionWithDefaultsMock.mockResolvedValue(false);
|
||||
downloadManager._lastDownloadError = 'Connection refused';
|
||||
await recipeModal.downloadCheckpoint(recipeModal.currentRecipe.checkpoint);
|
||||
expect(requests.some(r => r.url === '/api/lm/recipe/checkpoint/mark-hash-invalid')).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -311,6 +311,28 @@ class StubPersistenceService:
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
return SimpleNamespace(payload={"success": True}, status=200)
|
||||
|
||||
async def reconnect_checkpoint(
|
||||
self, *, recipe_scanner, recipe_id: str, target_name: str
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
return SimpleNamespace(payload={"success": True}, status=200)
|
||||
|
||||
async def restore_checkpoint(
|
||||
self, *, recipe_scanner, recipe_id: str
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
return SimpleNamespace(payload={"success": True}, status=200)
|
||||
|
||||
async def get_checkpoint_reconnect_suggestions(
|
||||
self, *, recipe_scanner, recipe_id: str, query: str | None = None
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
return SimpleNamespace(
|
||||
payload={"success": True, "suggestions": []}, status=200
|
||||
)
|
||||
|
||||
async def mark_checkpoint_hash_invalid(
|
||||
self, *, recipe_scanner, recipe_id: str, hash_invalid: bool = True
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
return SimpleNamespace(payload={"success": True}, status=200)
|
||||
|
||||
async def bulk_delete(
|
||||
self, *, recipe_scanner, recipe_ids: List[str]
|
||||
) -> SimpleNamespace: # pragma: no cover
|
||||
@@ -2050,3 +2072,99 @@ async def test_find_duplicates_forwards_include_prompt_and_assigns_unique_keys(
|
||||
assert len(groups) == 2
|
||||
assert {g["type"] for g in groups} == {"fingerprint", "source_path"}
|
||||
assert len({g["key"] for g in groups}) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpoint reconnect routes (manual remediation for recipe.checkpoint)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_checkpoint_reconnect_route(monkeypatch, tmp_path: Path) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/reconnect",
|
||||
json={"recipe_id": "r1", "target_name": "main"},
|
||||
)
|
||||
payload = await response.json()
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
|
||||
|
||||
async def test_checkpoint_reconnect_route_requires_target_name(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/reconnect",
|
||||
json={"recipe_id": "r1"},
|
||||
)
|
||||
assert response.status == 400
|
||||
|
||||
|
||||
async def test_checkpoint_restore_route(monkeypatch, tmp_path: Path) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/restore",
|
||||
json={"recipe_id": "r1"},
|
||||
)
|
||||
payload = await response.json()
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
|
||||
|
||||
async def test_checkpoint_restore_route_requires_recipe_id(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/restore",
|
||||
json={},
|
||||
)
|
||||
assert response.status == 400
|
||||
|
||||
|
||||
async def test_checkpoint_reconnect_suggestions_route(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.get(
|
||||
"/api/lm/recipe/r1/checkpoint/reconnect-suggestions?query=main"
|
||||
)
|
||||
payload = await response.json()
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert payload["suggestions"] == []
|
||||
|
||||
|
||||
async def test_checkpoint_reconnect_suggestions_route_without_query(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.get(
|
||||
"/api/lm/recipe/r1/checkpoint/reconnect-suggestions"
|
||||
)
|
||||
payload = await response.json()
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
|
||||
|
||||
async def test_checkpoint_mark_hash_invalid_route(monkeypatch, tmp_path: Path) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/mark-hash-invalid",
|
||||
json={"recipe_id": "r1"},
|
||||
)
|
||||
payload = await response.json()
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
|
||||
|
||||
async def test_checkpoint_mark_hash_invalid_route_requires_recipe_id(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
response = await harness.client.post(
|
||||
"/api/lm/recipe/checkpoint/mark-hash-invalid",
|
||||
json={},
|
||||
)
|
||||
assert response.status == 400
|
||||
|
||||
@@ -768,6 +768,251 @@ async def test_set_lora_entry_hash_invalid_persists_flag(tmp_path: Path, recipe_
|
||||
assert cleared_lora["hashInvalid"] is False
|
||||
|
||||
|
||||
async def test_update_checkpoint_entry_updates_cache_and_file(
|
||||
tmp_path: Path, recipe_scanner
|
||||
):
|
||||
scanner, _ = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
recipe_id = "recipe-ckpt-1"
|
||||
recipe_path = recipes_dir / f"{recipe_id}.recipe.json"
|
||||
original_checkpoint = {
|
||||
"name": "Old Model",
|
||||
"version": "v1",
|
||||
"id": 1,
|
||||
"type": "Checkpoint",
|
||||
"baseModel": "SDXL 1.0",
|
||||
"file_name": "old",
|
||||
"hash": "aaa",
|
||||
"isDeleted": True,
|
||||
}
|
||||
recipe_data = {
|
||||
"id": recipe_id,
|
||||
"file_path": str(tmp_path / "image.png"),
|
||||
"title": "Original",
|
||||
"modified": 0.0,
|
||||
"created_date": 0.0,
|
||||
"base_model": "SDXL 1.0",
|
||||
"checkpoint": dict(original_checkpoint),
|
||||
}
|
||||
recipe_path.write_text(json.dumps(recipe_data))
|
||||
await scanner.add_recipe(dict(recipe_data))
|
||||
|
||||
target_info = {
|
||||
"sha256": "abc123",
|
||||
"file_path": str(tmp_path / "checkpoints" / "main.safetensors"),
|
||||
"preview_url": "preview.png",
|
||||
"model_name": "Main Model",
|
||||
"base_model": "SDXL 1.0",
|
||||
"civitai": {"id": 42, "name": "v2"},
|
||||
}
|
||||
|
||||
updated_recipe, updated_checkpoint = await scanner.update_checkpoint_entry(
|
||||
recipe_id,
|
||||
target_name="main",
|
||||
target_checkpoint=target_info,
|
||||
)
|
||||
|
||||
# Write-back follows the pinned checkpoint key set, keeping the
|
||||
# user-entered file_name.
|
||||
assert updated_checkpoint["file_name"] == "main"
|
||||
assert updated_checkpoint["hash"] == "abc123"
|
||||
assert updated_checkpoint["isDeleted"] is False
|
||||
assert updated_checkpoint["hashInvalid"] is False
|
||||
assert updated_checkpoint["name"] == "Main Model"
|
||||
assert updated_checkpoint["version"] == "v2"
|
||||
assert updated_checkpoint["baseModel"] == "SDXL 1.0"
|
||||
assert updated_checkpoint["id"] == 42
|
||||
# The pre-reconnect state is snapshotted for undo
|
||||
assert updated_checkpoint["reconnectSnapshot"] == original_checkpoint
|
||||
assert "reconnectSnapshot" not in updated_checkpoint["reconnectSnapshot"]
|
||||
|
||||
with recipe_path.open("r", encoding="utf-8") as file_obj:
|
||||
persisted = json.load(file_obj)
|
||||
assert persisted["checkpoint"]["hash"] == "abc123"
|
||||
assert persisted["checkpoint"]["reconnectSnapshot"] == original_checkpoint
|
||||
|
||||
cache = await scanner.get_cached_data()
|
||||
cached_recipe = next(item for item in cache.raw_data if item["id"] == recipe_id)
|
||||
assert cached_recipe["checkpoint"]["hash"] == "abc123"
|
||||
|
||||
|
||||
async def test_update_checkpoint_entry_backfills_missing_display_keys(
|
||||
tmp_path: Path, recipe_scanner
|
||||
):
|
||||
scanner, _ = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
recipe_id = "recipe-ckpt-sparse"
|
||||
recipe_path = recipes_dir / f"{recipe_id}.recipe.json"
|
||||
# Parser-style sparse entry without name/version/baseModel keys
|
||||
recipe_data = {
|
||||
"id": recipe_id,
|
||||
"file_path": str(tmp_path / "image.png"),
|
||||
"title": "Sparse",
|
||||
"modified": 0.0,
|
||||
"created_date": 0.0,
|
||||
"checkpoint": {"file_name": "old", "hash": "aaa", "isDeleted": True},
|
||||
}
|
||||
recipe_path.write_text(json.dumps(recipe_data))
|
||||
await scanner.add_recipe(dict(recipe_data))
|
||||
|
||||
target_info = {
|
||||
"sha256": "abc123",
|
||||
"file_path": "/models/checkpoints/main.safetensors",
|
||||
"model_name": "Main Model",
|
||||
"base_model": "SDXL 1.0",
|
||||
"civitai": {"id": 42, "name": "v2"},
|
||||
}
|
||||
|
||||
_, updated_checkpoint = await scanner.update_checkpoint_entry(
|
||||
recipe_id,
|
||||
target_name="main",
|
||||
target_checkpoint=target_info,
|
||||
)
|
||||
|
||||
assert updated_checkpoint["name"] == "Main Model"
|
||||
assert updated_checkpoint["version"] == "v2"
|
||||
assert updated_checkpoint["baseModel"] == "SDXL 1.0"
|
||||
assert updated_checkpoint["modelVersionId"] == 42
|
||||
|
||||
|
||||
async def test_restore_checkpoint_entry_round_trip(tmp_path: Path, recipe_scanner):
|
||||
scanner, _ = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
recipe_id = "recipe-ckpt-restore"
|
||||
recipe_path = recipes_dir / f"{recipe_id}.recipe.json"
|
||||
original_checkpoint = {
|
||||
"name": "Old Model",
|
||||
"file_name": "old",
|
||||
"hash": "aaa",
|
||||
"isDeleted": True,
|
||||
}
|
||||
recipe_data = {
|
||||
"id": recipe_id,
|
||||
"file_path": str(tmp_path / "image.png"),
|
||||
"title": "Original",
|
||||
"modified": 0.0,
|
||||
"created_date": 0.0,
|
||||
"checkpoint": dict(original_checkpoint),
|
||||
}
|
||||
recipe_path.write_text(json.dumps(recipe_data))
|
||||
await scanner.add_recipe(dict(recipe_data))
|
||||
|
||||
target_info = {
|
||||
"sha256": "abc123",
|
||||
"file_path": "/models/checkpoints/main.safetensors",
|
||||
"model_name": "Main Model",
|
||||
"civitai": {"id": 42, "name": "v2"},
|
||||
}
|
||||
await scanner.update_checkpoint_entry(
|
||||
recipe_id, target_name="main", target_checkpoint=target_info
|
||||
)
|
||||
|
||||
restored_recipe, restored_checkpoint = await scanner.restore_checkpoint_entry(
|
||||
recipe_id
|
||||
)
|
||||
|
||||
assert restored_recipe["checkpoint"] == original_checkpoint
|
||||
assert "reconnectSnapshot" not in restored_recipe["checkpoint"]
|
||||
assert restored_checkpoint["file_name"] == "old"
|
||||
|
||||
with recipe_path.open("r", encoding="utf-8") as file_obj:
|
||||
persisted = json.load(file_obj)
|
||||
assert persisted["checkpoint"] == original_checkpoint
|
||||
|
||||
|
||||
async def test_restore_checkpoint_entry_without_snapshot_rejected(
|
||||
tmp_path: Path, recipe_scanner
|
||||
):
|
||||
scanner, _ = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
recipe_id = "recipe-ckpt-no-snapshot"
|
||||
recipe_path = recipes_dir / f"{recipe_id}.recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps({"id": recipe_id, "checkpoint": {"file_name": "plain"}})
|
||||
)
|
||||
|
||||
with pytest.raises(RecipeValidationError):
|
||||
await scanner.restore_checkpoint_entry(recipe_id)
|
||||
|
||||
|
||||
async def test_set_checkpoint_entry_hash_invalid_persists_flag(
|
||||
tmp_path: Path, recipe_scanner
|
||||
):
|
||||
scanner, _ = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
recipe_id = "hash-invalid-ckpt"
|
||||
recipe_path = recipes_dir / f"{recipe_id}.recipe.json"
|
||||
recipe_data = {
|
||||
"id": recipe_id,
|
||||
"file_path": str(tmp_path / "image.png"),
|
||||
"title": "Hash invalid",
|
||||
"modified": 0.0,
|
||||
"created_date": 0.0,
|
||||
"checkpoint": {"name": "Old", "file_name": "old", "hash": "a2a12bfa01"},
|
||||
}
|
||||
recipe_path.write_text(json.dumps(recipe_data))
|
||||
await scanner.add_recipe(dict(recipe_data))
|
||||
|
||||
updated_recipe, updated_checkpoint = await scanner.set_checkpoint_entry_hash_invalid(
|
||||
recipe_id, True
|
||||
)
|
||||
|
||||
assert updated_checkpoint["hashInvalid"] is True
|
||||
assert updated_recipe["checkpoint"]["hashInvalid"] is True
|
||||
with recipe_path.open("r", encoding="utf-8") as file_obj:
|
||||
persisted = json.load(file_obj)
|
||||
assert persisted["checkpoint"]["hashInvalid"] is True
|
||||
assert persisted["checkpoint"]["hash"] == "a2a12bfa01"
|
||||
|
||||
cache = await scanner.get_cached_data()
|
||||
cached_recipe = next(item for item in cache.raw_data if item["id"] == recipe_id)
|
||||
assert cached_recipe["checkpoint"]["hashInvalid"] is True
|
||||
|
||||
_, cleared_checkpoint = await scanner.set_checkpoint_entry_hash_invalid(
|
||||
recipe_id, False
|
||||
)
|
||||
assert cleared_checkpoint["hashInvalid"] is False
|
||||
|
||||
|
||||
async def test_find_local_checkpoints_by_name_uses_checkpoint_scanner(
|
||||
tmp_path: Path, monkeypatch
|
||||
):
|
||||
from py.services.recipe_scanner import RecipeScanner as RecipeScannerCls
|
||||
|
||||
class StubCheckpointScanner:
|
||||
async def find_models_by_name(self, name, *, base_model=None):
|
||||
return [
|
||||
{"file_name": f"{name}.safetensors", "base_model": base_model or ""}
|
||||
]
|
||||
|
||||
class StubLoraScannerForCkpt:
|
||||
async def get_cached_data(self):
|
||||
return SimpleNamespace(raw_data=[], version_index={})
|
||||
|
||||
RecipeScannerCls._instance = None
|
||||
scanner = RecipeScannerCls(
|
||||
lora_scanner=StubLoraScannerForCkpt(),
|
||||
checkpoint_scanner=StubCheckpointScanner(), # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
matches = await scanner.find_local_checkpoints_by_name("main")
|
||||
assert matches == [{"file_name": "main.safetensors", "base_model": ""}]
|
||||
|
||||
assert await scanner.find_local_checkpoints_by_name("") == []
|
||||
scanner._checkpoint_scanner = None
|
||||
assert await scanner.find_local_checkpoints_by_name("main") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_recipe_syntax_tokens_skips_unobtainable_loras(tmp_path: Path, recipe_scanner):
|
||||
scanner, _ = recipe_scanner
|
||||
|
||||
@@ -1676,3 +1676,306 @@ async def test_analyze_remote_image_meta_null_keeps_exif_loras(tmp_path, monkeyp
|
||||
assert loras[0]["hash"] == LORA_SHA256
|
||||
assert loras[0].get("isDeleted") in (None, False)
|
||||
assert "Daphne" in str(payload.get("gen_params", {}).get("prompt"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Checkpoint reconnect chain (manual remediation for recipe.checkpoint)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_persistence_service():
|
||||
return RecipePersistenceService(
|
||||
exif_utils=DummyExifUtils(),
|
||||
card_preview_width=512,
|
||||
logger=logging.getLogger("test"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect_checkpoint_distinguishes_ambiguous_mismatched_and_missing(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
models = [
|
||||
{
|
||||
"file_name": "realistic.safetensors",
|
||||
"folder": "sdxl",
|
||||
"file_path": "/models/checkpoints/sdxl/realistic.safetensors",
|
||||
"base_model": "SDXL 1.0",
|
||||
},
|
||||
{
|
||||
"file_name": "realistic.safetensors",
|
||||
"folder": "sd15",
|
||||
"file_path": "/models/checkpoints/sd15/realistic.safetensors",
|
||||
"base_model": "SD 1.5",
|
||||
},
|
||||
]
|
||||
|
||||
class DummyScanner:
|
||||
def __init__(self, recipe_path):
|
||||
self._recipe_path = recipe_path
|
||||
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
return str(self._recipe_path)
|
||||
|
||||
async def find_local_checkpoints_by_name(self, name, base_model=None):
|
||||
return ModelScanner.find_matching_models(models, name, base_model=base_model)
|
||||
|
||||
def write_recipe(base_model):
|
||||
recipe_path = tmp_path / "recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps({"id": "r1", "base_model": base_model, "checkpoint": {}})
|
||||
)
|
||||
return DummyScanner(recipe_path)
|
||||
|
||||
# Ambiguous bare name: two candidates survive (recipe base model unknown)
|
||||
scanner = write_recipe("")
|
||||
with pytest.raises(RecipeValidationError, match="include the folder path"):
|
||||
await service.reconnect_checkpoint(
|
||||
recipe_scanner=scanner, recipe_id="r1", target_name="realistic"
|
||||
)
|
||||
|
||||
# Confident base-model mismatch: the only candidate belongs to another family
|
||||
scanner = write_recipe("SD 1.5")
|
||||
with pytest.raises(RecipeValidationError, match="different base model"):
|
||||
await service.reconnect_checkpoint(
|
||||
recipe_scanner=scanner, recipe_id="r1", target_name="sdxl/realistic"
|
||||
)
|
||||
|
||||
# No candidate at all
|
||||
scanner = write_recipe("SDXL 1.0")
|
||||
with pytest.raises(RecipeNotFoundError, match="not found"):
|
||||
await service.reconnect_checkpoint(
|
||||
recipe_scanner=scanner, recipe_id="r1", target_name="missing"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect_checkpoint_family_compatible_succeeds_with_warning(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
pony_item = {
|
||||
"file_name": "main.safetensors",
|
||||
"folder": "",
|
||||
"file_path": "/models/checkpoints/main.safetensors",
|
||||
"base_model": "Pony",
|
||||
"sha256": "ab" * 32,
|
||||
}
|
||||
|
||||
recipe_path = tmp_path / "recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps({"id": "r1", "base_model": "Illustrious", "checkpoint": {}})
|
||||
)
|
||||
|
||||
class DummyScanner:
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
return str(recipe_path)
|
||||
|
||||
async def find_local_checkpoints_by_name(self, name, base_model=None):
|
||||
return [pony_item]
|
||||
|
||||
async def update_checkpoint_entry(self, recipe_id, *, target_name, target_checkpoint):
|
||||
assert target_checkpoint is pony_item
|
||||
return ({"id": "r1"}, {"file_name": target_checkpoint["file_name"]})
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
result = await service.reconnect_checkpoint(
|
||||
recipe_scanner=DummyScanner(), recipe_id="r1", target_name="main"
|
||||
)
|
||||
|
||||
assert result.payload["success"] is True
|
||||
assert result.payload["base_model_mismatch"] == {
|
||||
"recipe_base_model": "Illustrious",
|
||||
"checkpoint_base_model": "Pony",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect_checkpoint_exact_base_model_has_no_warning(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
item = {
|
||||
"file_name": "main.safetensors",
|
||||
"folder": "",
|
||||
"file_path": "/models/checkpoints/main.safetensors",
|
||||
"base_model": "SDXL 1.0",
|
||||
"sha256": "ab" * 32,
|
||||
}
|
||||
|
||||
recipe_path = tmp_path / "recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps({"id": "r1", "base_model": "SDXL 1.0", "checkpoint": {}})
|
||||
)
|
||||
|
||||
class DummyScanner:
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
return str(recipe_path)
|
||||
|
||||
async def find_local_checkpoints_by_name(self, name, base_model=None):
|
||||
return [item]
|
||||
|
||||
async def update_checkpoint_entry(self, recipe_id, *, target_name, target_checkpoint):
|
||||
return ({"id": "r1"}, {"file_name": target_checkpoint["file_name"]})
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
result = await service.reconnect_checkpoint(
|
||||
recipe_scanner=DummyScanner(), recipe_id="r1", target_name="main"
|
||||
)
|
||||
|
||||
assert result.payload["success"] is True
|
||||
assert "base_model_mismatch" not in result.payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_checkpoint_delegates_and_reports(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
class DummyScanner:
|
||||
async def restore_checkpoint_entry(self, recipe_id):
|
||||
assert recipe_id == "r1"
|
||||
return (
|
||||
{"id": "r1", "checkpoint": {"file_name": "old.safetensors"}},
|
||||
{"file_name": "old.safetensors"},
|
||||
)
|
||||
|
||||
async def find_recipes_by_fingerprint(self, fingerprint):
|
||||
return []
|
||||
|
||||
result = await service.restore_checkpoint(
|
||||
recipe_scanner=DummyScanner(), recipe_id="r1"
|
||||
)
|
||||
|
||||
assert result.payload["success"] is True
|
||||
assert result.payload["updated_checkpoint"]["file_name"] == "old.safetensors"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_checkpoint_reconnect_suggestions_loads_entry_and_delegates(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
recipe_path = tmp_path / "recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"id": "r1",
|
||||
"base_model": "SD 1.5",
|
||||
"checkpoint": {"file_name": "old.safetensors", "hash": "aaa"},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
class DummyScanner:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
assert recipe_id == "r1"
|
||||
return str(recipe_path)
|
||||
|
||||
async def suggest_checkpoint_reconnect_candidates(
|
||||
self, *, entry, recipe_base_model, query=None, limit=5
|
||||
):
|
||||
self.calls.append(
|
||||
{
|
||||
"entry": entry,
|
||||
"recipe_base_model": recipe_base_model,
|
||||
"query": query,
|
||||
}
|
||||
)
|
||||
return [
|
||||
{
|
||||
"file_name": "new.safetensors",
|
||||
"score": 1.0,
|
||||
"match_reason": "same_hash",
|
||||
"target_name": "new",
|
||||
}
|
||||
]
|
||||
|
||||
scanner = DummyScanner()
|
||||
result = await service.get_checkpoint_reconnect_suggestions(
|
||||
recipe_scanner=scanner, recipe_id="r1", query="new"
|
||||
)
|
||||
|
||||
assert result.payload["success"] is True
|
||||
assert result.payload["suggestions"][0]["target_name"] == "new"
|
||||
assert scanner.calls == [
|
||||
{
|
||||
"entry": {"file_name": "old.safetensors", "hash": "aaa"},
|
||||
"recipe_base_model": "SD 1.5",
|
||||
"query": "new",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_checkpoint_reconnect_suggestions_validates_recipe(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
class MissingScanner:
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
return str(tmp_path / "missing.json")
|
||||
|
||||
with pytest.raises(RecipeNotFoundError):
|
||||
await service.get_checkpoint_reconnect_suggestions(
|
||||
recipe_scanner=MissingScanner(), recipe_id="nope"
|
||||
)
|
||||
|
||||
recipe_path = tmp_path / "recipe.json"
|
||||
recipe_path.write_text(json.dumps({"id": "r1"}))
|
||||
|
||||
class EmptyScanner:
|
||||
async def get_recipe_json_path(self, recipe_id):
|
||||
return str(recipe_path)
|
||||
|
||||
with pytest.raises(RecipeValidationError, match="checkpoint"):
|
||||
await service.get_checkpoint_reconnect_suggestions(
|
||||
recipe_scanner=EmptyScanner(), recipe_id="r1"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_checkpoint_hash_invalid_delegates_and_reports(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
class DummyScanner:
|
||||
async def set_checkpoint_entry_hash_invalid(self, recipe_id, hash_invalid):
|
||||
assert recipe_id == "r1"
|
||||
assert hash_invalid is True
|
||||
return (
|
||||
{"id": "r1", "checkpoint": {"file_name": "m", "hashInvalid": True}},
|
||||
{"file_name": "m", "hashInvalid": True},
|
||||
)
|
||||
|
||||
result = await service.mark_checkpoint_hash_invalid(
|
||||
recipe_scanner=DummyScanner(), recipe_id="r1"
|
||||
)
|
||||
|
||||
assert result.payload["success"] is True
|
||||
assert result.payload["recipe_id"] == "r1"
|
||||
assert result.payload["hash_invalid"] is True
|
||||
assert result.payload["updated_checkpoint"]["hashInvalid"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_checkpoint_hash_invalid_can_clear_flag(tmp_path):
|
||||
service = _make_persistence_service()
|
||||
|
||||
class DummyScanner:
|
||||
async def set_checkpoint_entry_hash_invalid(self, recipe_id, hash_invalid):
|
||||
assert hash_invalid is False
|
||||
return (
|
||||
{"id": "r1", "checkpoint": {"file_name": "m", "hashInvalid": False}},
|
||||
{"file_name": "m", "hashInvalid": False},
|
||||
)
|
||||
|
||||
result = await service.mark_checkpoint_hash_invalid(
|
||||
recipe_scanner=DummyScanner(),
|
||||
recipe_id="r1",
|
||||
hash_invalid=False,
|
||||
)
|
||||
|
||||
assert result.payload["hash_invalid"] is False
|
||||
assert result.payload["updated_checkpoint"]["hashInvalid"] is False
|
||||
|
||||
Reference in New Issue
Block a user