mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
Compare commits
4 Commits
303833bbae
...
v1.2.2
| Author | SHA1 | Date | |
|---|---|---|---|
| 6e2185c182 | |||
| 41302e75ba | |||
| a17399d667 | |||
| e2d85a0a21 |
@@ -170,6 +170,10 @@ The system runs in two modes:
|
||||
- Route registrars organize endpoints by domain: `ModelRouteRegistrar`, `RecipeRouteRegistrar`, etc.
|
||||
- Request handlers in `py/routes/handlers/` implement route logic
|
||||
- All routes use aiohttp, return `web.json_response` or `web.Response`
|
||||
- Endpoints consumed by the companion browser extension (lm-civitai-extension)
|
||||
MUST also accept `GET` with query-string params: the extension is GET-only by
|
||||
convention (see its AGENTS.md), even for state-changing operations such as
|
||||
`GET /api/lm/recipe/{recipe_id}/reimport`
|
||||
|
||||
### Recipe System
|
||||
|
||||
|
||||
+427
-395
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,58 @@
|
||||
# CivitAI image imports can end up with 0 LoRAs
|
||||
|
||||
## Symptom
|
||||
|
||||
Importing a CivitAI image URL can produce a recipe with **zero LoRA
|
||||
entries**, even though the image page lists LoRAs in its resource panel.
|
||||
|
||||
Reported example: `https://civitai.red/images/140818889` was imported as a
|
||||
local recipe with 0 LoRAs, while the page shows 3 LoRAs. Some images (e.g.
|
||||
NSFW / higher browsing level) additionally require a login to view, so their
|
||||
data is not publicly reachable at all.
|
||||
|
||||
## Root cause
|
||||
|
||||
URL imports use only two data sources:
|
||||
|
||||
1. **CivitAI REST image API** — `GET /api/v1/images?imageId=<id>&nsfw=X&withMeta=true` → `meta`
|
||||
2. **Embedded image metadata** — EXIF/XMP read from the downloaded bytes
|
||||
|
||||
For the same image both sources can be empty, and the one source that does
|
||||
contain the data is never queried. Verified for image 140818889:
|
||||
|
||||
| Source | What it returned |
|
||||
|---|---|
|
||||
| REST image API | `meta` holds only a prompt; `modelVersionIds: []`; no `resources`/`hashes`; `baseModel: null` |
|
||||
| Downloaded image | PNG with **no EXIF/XMP** (the CDN URL ends in `.jpeg`, the body is PNG) |
|
||||
| Image page HTML | `__NEXT_DATA__` embeds the trpc `image.getGenerationData` result → full `resources` list: 3 LoRAs, each with `modelId`, `modelVersionId`, `modelName`, `modelType`, `versionName`, `baseModel` |
|
||||
|
||||
Key points:
|
||||
|
||||
- The page's resource panel is fed by an **internal, non-public trpc
|
||||
endpoint**, not by the public REST image API.
|
||||
- That internal endpoint is **login-gated** for some content — the
|
||||
"requires login" symptom.
|
||||
- Even with the version IDs in hand, `/model-versions/{id}` for these
|
||||
(Krea) versions returns **no `sha256`**, so an exact local-file hash match
|
||||
is impossible; only model/version identity is recoverable.
|
||||
|
||||
## Conclusion / status
|
||||
|
||||
0-LoRA imports are a data-source gap: public REST meta and image EXIF are
|
||||
both empty, while the only complete source (page generation data) is
|
||||
internal, sometimes login-gated, and not used by the importer.
|
||||
|
||||
Such imports **cannot be reliably auto-repaired/completed** by the backend
|
||||
alone. The old "Repair Metadata" feature only re-fetched the same incomplete
|
||||
REST meta and could not fix them; it was deprecated and has been removed.
|
||||
|
||||
**Fixed via the companion browser extension.** When the extension is
|
||||
installed with a valid license, it scrapes the image page's internal trpc
|
||||
generation data with the user's session and calls the payload-capable
|
||||
re-import endpoint (`POST /api/lm/recipe/{recipe_id}/reimport` with
|
||||
`image_url`/`name`/`resources`/`gen_params`/`base_model`/`tags` query
|
||||
params), which rebuilds the recipe from the caller-supplied metadata. The
|
||||
web UI delegates re-import of CivitAI-image-sourced recipes to the extension
|
||||
automatically (probe + `lm:reimport*` DOM events); without the extension,
|
||||
re-import silently falls back to the native path, which remains limited by
|
||||
the data-source gap documented above.
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "Keine Zuordnung für die {total} ausgewählten Rezepte erforderlich",
|
||||
"rematchFailed": "Zuordnung der ausgewählten Rezepte fehlgeschlagen: {message}",
|
||||
"reimporting": "Rezept wird aus Quelle neu importiert...",
|
||||
"reimportingViaExtension": "Rezept {current}/{total} wird über die Browser-Erweiterung neu importiert...",
|
||||
"reimportSuccess": "Rezept erfolgreich neu importiert",
|
||||
"reimportBulkComplete": "Neuimport abgeschlossen: {completed} importiert, {failed} fehlgeschlagen (von {total})",
|
||||
"reimportBulkFailed": "Neuimport einiger Rezepte fehlgeschlagen",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "No rematch needed for any of the {total} selected recipes",
|
||||
"rematchFailed": "Failed to rematch selected recipes: {message}",
|
||||
"reimporting": "Re-importing recipe from source...",
|
||||
"reimportingViaExtension": "Re-importing recipe {current}/{total} via browser extension...",
|
||||
"reimportSuccess": "Recipe re-imported successfully",
|
||||
"reimportBulkComplete": "Re-import complete: {completed} re-imported, {failed} failed (of {total})",
|
||||
"reimportBulkFailed": "Failed to re-import some recipes",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "Ninguna de las {total} recetas seleccionadas necesita reasociación",
|
||||
"rematchFailed": "Falló la reasociación de las recetas seleccionadas: {message}",
|
||||
"reimporting": "Reimportando receta desde origen...",
|
||||
"reimportingViaExtension": "Reimportando receta {current}/{total} mediante la extensión del navegador...",
|
||||
"reimportSuccess": "Receta reimportada exitosamente",
|
||||
"reimportBulkComplete": "Reimportación completa: {completed} reimportadas, {failed} fallidas (de {total})",
|
||||
"reimportBulkFailed": "Error al reimportar algunas recetas",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "Aucune des {total} Recipes sélectionnées ne nécessite de réassociation",
|
||||
"rematchFailed": "Échec de la réassociation des Recipes sélectionnées : {message}",
|
||||
"reimporting": "Ré-import de la Recipe depuis la source...",
|
||||
"reimportingViaExtension": "Ré-import de la Recipe {current}/{total} via l’extension du navigateur...",
|
||||
"reimportSuccess": "Recette ré-importée avec succès",
|
||||
"reimportBulkComplete": "Ré-import terminé : {completed} ré-importé(s), {failed} échec(s) (sur {total})",
|
||||
"reimportBulkFailed": "Échec du ré-import de certaines Recipes",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "אין צורך בהתאמה עבור {total} המתכונים שנבחרו",
|
||||
"rematchFailed": "ההתאמה מחדש של המתכונים שנבחרו נכשלה: {message}",
|
||||
"reimporting": "מייבא מתכון מחדש מהמקור...",
|
||||
"reimportingViaExtension": "מייבא מתכון מחדש {current}/{total} דרך תוסף הדפדפן...",
|
||||
"reimportSuccess": "המתכון יובא מחדש בהצלחה",
|
||||
"reimportBulkComplete": "ייבוא מחדש הושלם: {completed} יובאו, {failed} נכשלו (מתוך {total})",
|
||||
"reimportBulkFailed": "ייבוא מחדש של חלק מהמתכונים נכשל",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "選択した {total} 件のレシピは再マッチングの必要がありませんでした",
|
||||
"rematchFailed": "選択したレシピの再マッチングに失敗しました:{message}",
|
||||
"reimporting": "ソースからレシピを再インポート中...",
|
||||
"reimportingViaExtension": "ブラウザ拡張機能経由でレシピを再インポート中 ({current}/{total})...",
|
||||
"reimportSuccess": "レシピの再インポートが完了しました",
|
||||
"reimportBulkComplete": "再インポート完了:{completed} 件成功、{failed} 件失敗(合計 {total} 件)",
|
||||
"reimportBulkFailed": "一部のレシピの再インポートに失敗しました",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "선택한 {total}개 레시피는 재매칭이 필요하지 않습니다",
|
||||
"rematchFailed": "선택한 레시피 재매칭 실패: {message}",
|
||||
"reimporting": "소스에서 레시피를 다시 가져오는 중...",
|
||||
"reimportingViaExtension": "브라우저 확장 프로그램을 통해 레시피를 다시 가져오는 중 ({current}/{total})...",
|
||||
"reimportSuccess": "레시피를 다시 가져왔습니다",
|
||||
"reimportBulkComplete": "다시 가져오기 완료: {completed}개 성공, {failed}개 실패 (총 {total}개)",
|
||||
"reimportBulkFailed": "일부 레시피를 다시 가져오지 못했습니다",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "Ни один из {total} выбранных рецептов не требует сопоставления",
|
||||
"rematchFailed": "Не удалось сопоставить выбранные рецепты: {message}",
|
||||
"reimporting": "Переимпорт рецепта из источника...",
|
||||
"reimportingViaExtension": "Переимпорт рецепта {current}/{total} через расширение браузера...",
|
||||
"reimportSuccess": "Рецепт успешно переимпортирован",
|
||||
"reimportBulkComplete": "Переимпорт завершён: {completed} переимпортировано, {failed} ошибок (из {total})",
|
||||
"reimportBulkFailed": "Не удалось переимпортировать некоторые рецепты",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "{total} 个所选配方均无需重新匹配",
|
||||
"rematchFailed": "重新匹配所选配方失败:{message}",
|
||||
"reimporting": "正在从源重新导入配方...",
|
||||
"reimportingViaExtension": "正在通过浏览器扩展重新导入配方 {current}/{total}...",
|
||||
"reimportSuccess": "配方已从源重新导入成功",
|
||||
"reimportBulkComplete": "重新导入完成:{completed} 个已导入,{failed} 个失败(共 {total} 个)",
|
||||
"reimportBulkFailed": "重新导入某些配方失败",
|
||||
|
||||
@@ -2229,6 +2229,7 @@
|
||||
"rematchSkipped": "{total} 個所選配方均無需重新匹配",
|
||||
"rematchFailed": "重新匹配所選配方失敗:{message}",
|
||||
"reimporting": "正在從來源重新匯入配方...",
|
||||
"reimportingViaExtension": "正在透過瀏覽器擴充功能重新匯入配方 {current}/{total}...",
|
||||
"reimportSuccess": "配方已從來源重新匯入成功",
|
||||
"reimportBulkComplete": "重新匯入完成:{completed} 個已匯入,{failed} 個失敗(共 {total} 個)",
|
||||
"reimportBulkFailed": "重新匯入某些配方失敗",
|
||||
|
||||
@@ -1020,12 +1020,55 @@ class RecipeManagementHandler:
|
||||
persisted_source_path=persisted_source_path,
|
||||
)
|
||||
|
||||
async with self._import_semaphore:
|
||||
import_response = await self._do_import_from_url(
|
||||
source_path,
|
||||
recipe_scanner,
|
||||
target_dir=old_folder,
|
||||
)
|
||||
# Optional caller-supplied metadata payload (companion browser
|
||||
# extension re-import). Only honored for CivitAI image page
|
||||
# sources; everything else uses the native URL import below.
|
||||
params = request.rel_url.query
|
||||
payload_image_url = params.get("image_url")
|
||||
payload_name = params.get("name")
|
||||
payload_resources = params.get("resources")
|
||||
has_import_payload = bool(
|
||||
payload_image_url and payload_name and payload_resources
|
||||
)
|
||||
|
||||
import_response: web.Response | None = None
|
||||
if has_import_payload and image_id:
|
||||
try:
|
||||
async with self._import_semaphore:
|
||||
import_response = await self._import_remote_recipe_impl(
|
||||
image_url=payload_image_url,
|
||||
name=payload_name,
|
||||
resources_raw=payload_resources,
|
||||
gen_params_raw=params.get("gen_params"),
|
||||
tags_raw=params.get("tags"),
|
||||
base_model=params.get("base_model", "") or "",
|
||||
source_path=source_path,
|
||||
target_dir=old_folder,
|
||||
)
|
||||
except RecipeValidationError as exc:
|
||||
# Malformed resources/gen_params JSON: treat as "no
|
||||
# payload" and use the legacy URL re-import.
|
||||
self._logger.warning(
|
||||
"Ignoring malformed re-import payload for recipe %s "
|
||||
"(%s); falling back to source URL re-import",
|
||||
recipe_id,
|
||||
exc,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._logger.warning(
|
||||
"Payload-based re-import failed for recipe %s: %s; "
|
||||
"falling back to source URL re-import",
|
||||
recipe_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
if import_response is None:
|
||||
async with self._import_semaphore:
|
||||
import_response = await self._do_import_from_url(
|
||||
source_path,
|
||||
recipe_scanner,
|
||||
target_dir=old_folder,
|
||||
)
|
||||
|
||||
await self._persistence_service.delete_recipe(
|
||||
recipe_scanner=recipe_scanner, recipe_id=recipe_id
|
||||
@@ -1052,14 +1095,19 @@ class RecipeManagementHandler:
|
||||
exc,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"old_recipe_id": recipe_id,
|
||||
"recipe_id": new_recipe_id,
|
||||
"source_path": source_path,
|
||||
}
|
||||
response_body: Dict[str, Any] = {
|
||||
"success": True,
|
||||
"old_recipe_id": recipe_id,
|
||||
"recipe_id": new_recipe_id,
|
||||
"source_path": source_path,
|
||||
}
|
||||
loras_count = await self._count_recipe_loras(
|
||||
recipe_scanner, new_recipe_id
|
||||
)
|
||||
if loras_count is not None:
|
||||
response_body["loras_count"] = loras_count
|
||||
|
||||
return web.json_response(response_body)
|
||||
except RecipeNotFoundError as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=404)
|
||||
except RecipeValidationError as exc:
|
||||
@@ -1092,31 +1140,14 @@ class RecipeManagementHandler:
|
||||
if not resources_raw:
|
||||
raise RecipeValidationError("Missing required field: resources")
|
||||
|
||||
checkpoint_entry, lora_entries = self._parse_resources_payload(
|
||||
resources_raw
|
||||
)
|
||||
gen_params_request = self._parse_gen_params(params.get("gen_params"))
|
||||
|
||||
self._logger.info(
|
||||
"Remote recipe import received: url=%s, lora_count=%d",
|
||||
image_url,
|
||||
len(lora_entries),
|
||||
)
|
||||
self._logger.debug(
|
||||
" gen_params_keys=%s, checkpoint_keys=%s",
|
||||
sorted(gen_params_request.keys()) if gen_params_request else [],
|
||||
sorted(checkpoint_entry.keys()) if isinstance(checkpoint_entry, dict) else [],
|
||||
)
|
||||
|
||||
# Throttle concurrent imports to avoid starving ComfyUI's event loop
|
||||
async with self._import_semaphore:
|
||||
return await self._do_import_remote_recipe(
|
||||
return await self._import_remote_recipe_impl(
|
||||
image_url=image_url,
|
||||
name=name,
|
||||
lora_entries=lora_entries,
|
||||
checkpoint_entry=checkpoint_entry,
|
||||
gen_params_request=gen_params_request,
|
||||
tags=self._parse_tags(params.get("tags")),
|
||||
resources_raw=resources_raw,
|
||||
gen_params_raw=params.get("gen_params"),
|
||||
tags_raw=params.get("tags"),
|
||||
base_model=params.get("base_model", "") or "",
|
||||
source_path=params.get("source_path") or image_url,
|
||||
)
|
||||
@@ -1130,6 +1161,52 @@ class RecipeManagementHandler:
|
||||
)
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
|
||||
async def _import_remote_recipe_impl(
|
||||
self,
|
||||
*,
|
||||
image_url: str,
|
||||
name: str,
|
||||
resources_raw: str,
|
||||
gen_params_raw: Optional[str],
|
||||
tags_raw: Optional[str],
|
||||
base_model: str,
|
||||
source_path: str,
|
||||
target_dir: str | None = None,
|
||||
) -> web.Response:
|
||||
"""Payload-based remote import engine shared by import-remote and the
|
||||
extension-driven re-import path.
|
||||
|
||||
Parses the caller-supplied payloads and delegates to
|
||||
:meth:`_do_import_remote_recipe`. Raises ``RecipeValidationError`` on
|
||||
malformed payloads so callers can decide how to handle them (the
|
||||
re-import path falls back to the legacy URL import).
|
||||
"""
|
||||
checkpoint_entry, lora_entries = self._parse_resources_payload(resources_raw)
|
||||
gen_params_request = self._parse_gen_params(gen_params_raw)
|
||||
|
||||
self._logger.info(
|
||||
"Remote recipe import received: url=%s, lora_count=%d",
|
||||
image_url,
|
||||
len(lora_entries),
|
||||
)
|
||||
self._logger.debug(
|
||||
" gen_params_keys=%s, checkpoint_keys=%s",
|
||||
sorted(gen_params_request.keys()) if gen_params_request else [],
|
||||
sorted(checkpoint_entry.keys()) if isinstance(checkpoint_entry, dict) else [],
|
||||
)
|
||||
|
||||
return await self._do_import_remote_recipe(
|
||||
image_url=image_url,
|
||||
name=name,
|
||||
lora_entries=lora_entries,
|
||||
checkpoint_entry=checkpoint_entry,
|
||||
gen_params_request=gen_params_request,
|
||||
tags=self._parse_tags(tags_raw),
|
||||
base_model=base_model,
|
||||
source_path=source_path,
|
||||
target_dir=target_dir,
|
||||
)
|
||||
|
||||
async def _do_import_remote_recipe(
|
||||
self,
|
||||
*,
|
||||
@@ -1141,6 +1218,7 @@ class RecipeManagementHandler:
|
||||
tags: list[Any],
|
||||
base_model: str,
|
||||
source_path: str,
|
||||
target_dir: str | None = None,
|
||||
) -> web.Response:
|
||||
recipe_scanner = self._recipe_scanner_getter()
|
||||
if recipe_scanner is None:
|
||||
@@ -1304,6 +1382,7 @@ class RecipeManagementHandler:
|
||||
tags=tags,
|
||||
metadata=metadata,
|
||||
extension=extension,
|
||||
target_dir=target_dir,
|
||||
)
|
||||
return web.json_response(result.payload, status=result.status)
|
||||
|
||||
@@ -1768,6 +1847,25 @@ class RecipeManagementHandler:
|
||||
return []
|
||||
return [tag.strip() for tag in tag_text.split(",") if tag.strip()]
|
||||
|
||||
async def _count_recipe_loras(
|
||||
self, recipe_scanner: Any, recipe_id: Optional[str]
|
||||
) -> Optional[int]:
|
||||
"""Best-effort LoRA count for a freshly saved recipe (for the
|
||||
re-import response). Returns None when the recipe cannot be read."""
|
||||
if not recipe_id:
|
||||
return None
|
||||
try:
|
||||
recipe = await recipe_scanner.get_recipe_by_id(recipe_id)
|
||||
except Exception as exc:
|
||||
self._logger.debug(
|
||||
"Could not read new recipe %s for loras_count: %s",
|
||||
recipe_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
loras = (recipe or {}).get("loras")
|
||||
return len(loras) if isinstance(loras, list) else None
|
||||
|
||||
def _parse_gen_params(self, payload: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
if payload is None:
|
||||
return None
|
||||
|
||||
@@ -110,6 +110,11 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/recipe/{recipe_id}/reimport", "reimport_recipe"
|
||||
),
|
||||
# The companion browser extension only ever issues GET requests, so the
|
||||
# payload-based re-import variant must also be reachable via GET.
|
||||
RouteDefinition(
|
||||
"GET", "/api/lm/recipe/{recipe_id}/reimport", "reimport_recipe"
|
||||
),
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/recipe/{recipe_id}/send-workflow", "send_recipe_workflow"
|
||||
),
|
||||
|
||||
@@ -505,6 +505,50 @@ class CivitaiClient:
|
||||
logger.warning(f"Failed to fetch version by id {version_id}")
|
||||
return None
|
||||
|
||||
async def get_version_file_mini(
|
||||
self, version_id: int, file_id: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Fetch raw stored file info via the model-versions/mini endpoint.
|
||||
|
||||
The public REST API rewrites ``files[].name`` to
|
||||
``"{model}_{version}"`` for non-LoRA model types, so every
|
||||
precision variant of a multi-file version shares one name (#1100).
|
||||
The mini endpoint returns the raw ``ModelFile.name`` in
|
||||
``fileName``. ``file_id`` is mandatory: without it mini picks a
|
||||
file via its own primary-file logic, which can disagree with the
|
||||
REST ``primary`` flag.
|
||||
|
||||
Returns the mini payload dict on success, None on any failure.
|
||||
"""
|
||||
try:
|
||||
success, data = await self._make_request(
|
||||
"GET",
|
||||
f"{self.base_url}/model-versions/mini/{version_id}",
|
||||
params={"modelFileId": file_id},
|
||||
use_auth=True,
|
||||
)
|
||||
if success and isinstance(data, dict):
|
||||
return data
|
||||
if is_expected_offline_error(data):
|
||||
return None
|
||||
logger.debug(
|
||||
"Mini endpoint lookup failed for version %s file %s: %s",
|
||||
version_id,
|
||||
file_id,
|
||||
data,
|
||||
)
|
||||
return None
|
||||
except RateLimitError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug(
|
||||
"Error fetching mini info for version %s file %s: %s",
|
||||
version_id,
|
||||
file_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict[str, Any]]:
|
||||
if not model_hash:
|
||||
return None
|
||||
|
||||
@@ -35,6 +35,7 @@ from .service_registry import ServiceRegistry
|
||||
from .settings_manager import get_settings_manager
|
||||
from .metadata_service import get_default_metadata_provider, get_metadata_provider
|
||||
from .downloader import get_downloader, DownloadProgress, DownloadStreamControl
|
||||
from .errors import RateLimitError
|
||||
from .aria2_downloader import Aria2Error, get_aria2_downloader
|
||||
from .aria2_transfer_state import Aria2TransferStateStore
|
||||
from .download_queue_service import DownloadQueueService
|
||||
@@ -929,6 +930,42 @@ class DownloadManager:
|
||||
|
||||
return download_urls
|
||||
|
||||
async def _fetch_raw_file_name(
|
||||
self,
|
||||
metadata_provider,
|
||||
version_id: Optional[int],
|
||||
file_id: Any,
|
||||
) -> Optional[str]:
|
||||
"""Best-effort lookup of the raw stored filename via the CivitAI
|
||||
model-versions/mini endpoint (#1100). Returns None on any failure so
|
||||
the caller can fall back to the (possibly rewritten) REST name."""
|
||||
if version_id is None or file_id is None:
|
||||
return None
|
||||
fetch = getattr(metadata_provider, "get_version_file_mini", None)
|
||||
if fetch is None:
|
||||
return None
|
||||
try:
|
||||
mini_info = await fetch(int(version_id), int(file_id))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
except RateLimitError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug(
|
||||
"Mini endpoint lookup failed for version %s file %s: %s",
|
||||
version_id,
|
||||
file_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
if not isinstance(mini_info, dict):
|
||||
return None
|
||||
raw_name = mini_info.get("fileName")
|
||||
if not isinstance(raw_name, str) or not raw_name.strip():
|
||||
return None
|
||||
# Defensive: never let a path component slip into the filename.
|
||||
return os.path.basename(raw_name.strip()) or None
|
||||
|
||||
def _build_metadata_for_resume(
|
||||
self,
|
||||
*,
|
||||
@@ -1858,6 +1895,24 @@ class DownloadManager:
|
||||
if not download_urls:
|
||||
return {"success": False, "error": "No mirror URL found"}
|
||||
|
||||
# The public REST API rewrites files[].name to
|
||||
# "{model}_{version}" for non-LoRA model types, so every
|
||||
# precision variant of a multi-file version shares one name and
|
||||
# lands on disk with a random short-hash suffix. The mini
|
||||
# endpoint returns the raw stored filename (#1100). CivArchive
|
||||
# already serves raw names.
|
||||
if source != "civarchive":
|
||||
raw_file_name = await self._fetch_raw_file_name(
|
||||
metadata_provider, resolved_version_id, file_info.get("id")
|
||||
)
|
||||
if raw_file_name and raw_file_name != file_info.get("name"):
|
||||
logger.info(
|
||||
"[download] Using raw stored filename '%s' instead of REST name '%s'",
|
||||
raw_file_name,
|
||||
file_info.get("name"),
|
||||
)
|
||||
file_info = {**file_info, "name": raw_file_name}
|
||||
|
||||
# 3. Prepare download
|
||||
file_name = file_info.get("name", "")
|
||||
if not file_name:
|
||||
|
||||
@@ -169,6 +169,17 @@ class ModelMetadataProvider(ABC):
|
||||
"""Published model count for the user; None when unsupported."""
|
||||
return None
|
||||
|
||||
async def get_version_file_mini(
|
||||
self, version_id: int, file_id: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Fetch raw stored file info via CivitAI's model-versions/mini endpoint.
|
||||
|
||||
Only the CivitAI provider implements this (#1100); other providers
|
||||
already serve raw file names (CivArchive) or cannot resolve this
|
||||
lookup (SQLite), so the default is None.
|
||||
"""
|
||||
return None
|
||||
|
||||
class CivitaiModelMetadataProvider(ModelMetadataProvider):
|
||||
"""Provider that uses Civitai API for metadata"""
|
||||
|
||||
@@ -203,6 +214,11 @@ class CivitaiModelMetadataProvider(ModelMetadataProvider):
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
return await self.client.get_creator_model_count(username)
|
||||
|
||||
async def get_version_file_mini(
|
||||
self, version_id: int, file_id: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
return await self.client.get_version_file_mini(version_id, file_id)
|
||||
|
||||
class CivArchiveModelMetadataProvider(ModelMetadataProvider):
|
||||
"""Provider that uses CivArchive API for metadata"""
|
||||
|
||||
@@ -700,6 +716,37 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
continue
|
||||
return None
|
||||
|
||||
async def get_version_file_mini(
|
||||
self, version_id: int, file_id: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result = await self._call_with_rate_limit(
|
||||
label,
|
||||
provider.get_version_file_mini,
|
||||
version_id,
|
||||
file_id,
|
||||
)
|
||||
if result:
|
||||
return result
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"Provider %s failed for get_version_file_mini: %s", label, e
|
||||
)
|
||||
continue
|
||||
return None
|
||||
|
||||
def _iter_providers(self):
|
||||
return zip(self.providers, self._provider_labels)
|
||||
|
||||
@@ -791,6 +838,16 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
return await self._provider.get_creator_model_count(username)
|
||||
|
||||
async def get_version_file_mini(
|
||||
self, version_id: int, file_id: int
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
return await self._rate_limit_helper.run(
|
||||
self._label,
|
||||
self._provider.get_version_file_mini,
|
||||
version_id,
|
||||
file_id,
|
||||
)
|
||||
|
||||
class ModelMetadataProviderManager:
|
||||
"""Manager for selecting and using model metadata providers"""
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-lora-manager"
|
||||
description = "Revolutionize your workflow with the ultimate LoRA companion for ComfyUI!"
|
||||
version = "1.2.1"
|
||||
version = "1.2.2"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = [
|
||||
"aiohttp",
|
||||
|
||||
@@ -6,6 +6,7 @@ import { setSessionItem, removeSessionItem } from '../../utils/storageHelpers.js
|
||||
import { updateRecipeMetadata } from '../../api/recipeApi.js';
|
||||
import { state } from '../../state/index.js';
|
||||
import { moveManager } from '../../managers/MoveManager.js';
|
||||
import { probeExtension, delegateReimport, getCivitaiImageInfo } from '../../utils/extensionReimportBridge.js';
|
||||
|
||||
export class RecipeContextMenu extends BaseContextMenu {
|
||||
constructor() {
|
||||
@@ -355,6 +356,24 @@ export class RecipeContextMenu extends BaseContextMenu {
|
||||
return;
|
||||
}
|
||||
|
||||
// Recipes imported from a CivitAI image page can carry incomplete
|
||||
// metadata (0 LoRAs); the companion browser extension can re-import
|
||||
// them with the full page data. Fall back to the native path whenever
|
||||
// the extension is absent, unlicensed, or the delegation fails.
|
||||
const recipeItem = state.virtualScroller?.items?.find(item => item?.id === recipeId);
|
||||
const civitaiImage = getCivitaiImageInfo(recipeItem?.source_path);
|
||||
if (civitaiImage) {
|
||||
try {
|
||||
const probe = await probeExtension();
|
||||
if (probe?.supported && probe?.licenseValid) {
|
||||
await this.reimportViaExtension(recipeId, civitaiImage, recipeItem?.title || '');
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn('Extension re-import unavailable, using native path:', error);
|
||||
}
|
||||
}
|
||||
|
||||
state.loadingManager.showSimpleLoading('Re-importing recipe from source...');
|
||||
|
||||
try {
|
||||
@@ -377,6 +396,34 @@ export class RecipeContextMenu extends BaseContextMenu {
|
||||
showToast('recipes.contextMenu.reimport.failed', { message: error.message }, 'error');
|
||||
}
|
||||
}
|
||||
|
||||
// Re-import a single CivitAI-image recipe through the companion browser
|
||||
// extension. Throws on delegation failure so the caller can fall back to
|
||||
// the native path.
|
||||
async reimportViaExtension(recipeId, civitaiImage, title) {
|
||||
state.loadingManager.showSimpleLoading('Re-importing recipe via browser extension...');
|
||||
|
||||
try {
|
||||
const { failed } = await delegateReimport([{
|
||||
recipeId,
|
||||
imageId: civitaiImage.imageId,
|
||||
imageUrl: civitaiImage.imageUrl,
|
||||
title,
|
||||
}]);
|
||||
|
||||
state.loadingManager.hide();
|
||||
if (failed > 0) {
|
||||
showToast('recipes.contextMenu.reimport.failed', { message: 'Extension re-import failed' }, 'error');
|
||||
} else {
|
||||
showToast('toast.recipes.reimportSuccess', {}, 'success');
|
||||
}
|
||||
const { resetAndReload } = await import('../../api/recipeApi.js');
|
||||
resetAndReload(false, { preserveScroll: false });
|
||||
} catch (error) {
|
||||
state.loadingManager.hide();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mix in shared methods from ModelContextMenuMixin
|
||||
|
||||
@@ -2877,12 +2877,14 @@ class RecipeModal {
|
||||
|
||||
canDownloadLora(lora) {
|
||||
if (!lora) return false;
|
||||
const modelId = lora.modelId || lora.modelID || lora.model_id;
|
||||
const versionId = lora.id || lora.modelVersionId;
|
||||
// Direct download needs both identifiers; a hash alone is enough
|
||||
// because downloadRecipeLora resolves it to a version on demand —
|
||||
// the same fallback the bulk "download missing" flow uses.
|
||||
return !!((modelId && versionId) || lora.hash);
|
||||
// A bare CivitAI version id is enough: it uniquely pins the exact
|
||||
// file, and downloadRecipeLora resolves the owning model id from the
|
||||
// version on demand (the same fallback the bulk "download missing"
|
||||
// flow uses). A hash alone is likewise sufficient. A model id without
|
||||
// an exact version id is NOT enough — downloading the model's latest
|
||||
// version could silently mismatch the recipe's pinned version.
|
||||
return !!(versionId || lora.hash);
|
||||
}
|
||||
|
||||
renderCivitaiLink(url) {
|
||||
@@ -2991,6 +2993,9 @@ class RecipeModal {
|
||||
* Resolve the Civitai model/version identifiers needed for download.
|
||||
* Recipe LoRAs parsed from PNG metadata often carry only a hash; resolve
|
||||
* it through the same endpoint the bulk "download missing" flow uses.
|
||||
* Version-only entries (page-imported recipes whose CivitAI version has
|
||||
* no sha256) are resolved through the version endpoint, which returns
|
||||
* the owning model id.
|
||||
*/
|
||||
async resolveLoraDownloadIdentifiers(lora) {
|
||||
let modelId = lora.modelId || lora.modelID || lora.model_id;
|
||||
@@ -3001,21 +3006,41 @@ class RecipeModal {
|
||||
return { modelId, versionId, versionName };
|
||||
}
|
||||
|
||||
if (!lora.hash) {
|
||||
return null;
|
||||
// Hash-only entries (PNG/recipe-JSON imports): resolve the owning
|
||||
// model/version through the same endpoint the bulk "download
|
||||
// missing" flow uses.
|
||||
if (lora.hash) {
|
||||
const response = await fetch(`/api/lm/loras/civitai/model/hash/${lora.hash}`);
|
||||
const versionInfo = await response.json();
|
||||
if (versionInfo?.error) {
|
||||
return null;
|
||||
}
|
||||
|
||||
modelId = versionInfo.modelId || versionInfo.model?.id;
|
||||
versionId = versionInfo.id;
|
||||
versionName = versionInfo.name || versionName;
|
||||
|
||||
return modelId && versionId ? { modelId, versionId, versionName } : null;
|
||||
}
|
||||
|
||||
const response = await fetch(`/api/lm/loras/civitai/model/hash/${lora.hash}`);
|
||||
const versionInfo = await response.json();
|
||||
if (versionInfo?.error) {
|
||||
return null;
|
||||
// Version-only entries (page-imported recipes whose CivitAI versions
|
||||
// expose no sha256): the version id still pins the exact file, so
|
||||
// resolve the owning model id from the version endpoint on demand.
|
||||
if (versionId) {
|
||||
const response = await fetch(`/api/lm/loras/civitai/model/version/${versionId}`);
|
||||
const versionInfo = await response.json();
|
||||
if (!versionInfo || versionInfo?.error === 'Model not found') {
|
||||
return null;
|
||||
}
|
||||
|
||||
modelId = versionInfo.modelId || versionInfo.model?.id;
|
||||
versionId = versionInfo.id || versionId;
|
||||
versionName = versionInfo.name || versionName;
|
||||
|
||||
return modelId && versionId ? { modelId, versionId, versionName } : null;
|
||||
}
|
||||
|
||||
modelId = versionInfo.modelId || versionInfo.model?.id;
|
||||
versionId = versionInfo.id;
|
||||
versionName = versionInfo.name || versionName;
|
||||
|
||||
return modelId && versionId ? { modelId, versionId, versionName } : null;
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -10,6 +10,7 @@ import { createBaseModelPicker, inferBaseModelsFromFilepaths } from '../componen
|
||||
import { getPriorityTagSuggestions } from '../utils/priorityTagHelpers.js';
|
||||
import { eventManager } from '../utils/EventManager.js';
|
||||
import { translate } from '../utils/i18nHelpers.js';
|
||||
import { probeExtension, delegateReimport, getCivitaiImageInfo } from '../utils/extensionReimportBridge.js';
|
||||
import { getNsfwLevelSelector } from '../components/shared/NsfwLevelSelector.js';
|
||||
|
||||
export class BulkManager {
|
||||
@@ -857,17 +858,74 @@ export class BulkManager {
|
||||
`Re-importing recipe 1/${total}...`
|
||||
);
|
||||
|
||||
// Partition the selection: recipes sourced from a CivitAI image page
|
||||
// can be delegated to the companion browser extension (which scrapes
|
||||
// the full page metadata); everything else uses the native endpoint.
|
||||
const delegatable = [];
|
||||
const nativeFilePaths = [];
|
||||
for (const filePath of filePaths) {
|
||||
const recipeItem = recipeMap.get(filePath);
|
||||
const civitaiImage = getCivitaiImageInfo(recipeItem?.source_path);
|
||||
if (civitaiImage && recipeItem?.id) {
|
||||
delegatable.push({
|
||||
filePath,
|
||||
recipeId: recipeItem.id,
|
||||
imageId: civitaiImage.imageId,
|
||||
imageUrl: civitaiImage.imageUrl,
|
||||
title: recipeItem.title || '',
|
||||
});
|
||||
} else {
|
||||
nativeFilePaths.push(filePath);
|
||||
}
|
||||
}
|
||||
|
||||
// Probe once; on any probe/delegate failure the delegatable recipes
|
||||
// fall back to the native sequential loop below.
|
||||
if (delegatable.length > 0) {
|
||||
try {
|
||||
const probe = await probeExtension();
|
||||
if (probe?.supported && probe?.licenseValid) {
|
||||
const batchResult = await delegateReimport(
|
||||
delegatable.map(({ recipeId, imageId, imageUrl, title }) => ({
|
||||
recipeId, imageId, imageUrl, title,
|
||||
})),
|
||||
{
|
||||
onProgress: (progress) => {
|
||||
progressUI.updateProgress(
|
||||
Math.floor(((progress.current || 0) / total) * 100),
|
||||
progress.title || '',
|
||||
translate('toast.recipes.reimportingViaExtension', {
|
||||
current: progress.current || 0,
|
||||
total,
|
||||
})
|
||||
);
|
||||
},
|
||||
}
|
||||
);
|
||||
completed += batchResult.completed;
|
||||
failed += batchResult.failed;
|
||||
} else {
|
||||
nativeFilePaths.push(...delegatable.map(entry => entry.filePath));
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn('[reimportSelectedRecipes] extension delegation failed, using native path:', error);
|
||||
nativeFilePaths.push(...delegatable.map(entry => entry.filePath));
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
for (let i = 0; i < filePaths.length; i++) {
|
||||
const filePath = filePaths[i];
|
||||
const processedBeforeNative = completed + failed;
|
||||
for (let i = 0; i < nativeFilePaths.length; i++) {
|
||||
const filePath = nativeFilePaths[i];
|
||||
const recipeItem = recipeMap.get(filePath);
|
||||
const recipeId = recipeItem?.id;
|
||||
const recipeName = recipeItem?.title || recipeId || 'Unknown';
|
||||
const processed = processedBeforeNative + i;
|
||||
|
||||
progressUI.updateProgress(
|
||||
Math.floor((i / total) * 100),
|
||||
Math.floor((processed / total) * 100),
|
||||
recipeName,
|
||||
`Re-importing recipe ${Math.min(i + 1, total)}/${total}...`
|
||||
`Re-importing recipe ${Math.min(processed + 1, total)}/${total}...`
|
||||
);
|
||||
|
||||
if (!recipeId) {
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
/**
|
||||
* Bridge to the companion LoRA Manager browser extension.
|
||||
*
|
||||
* The extension can re-import recipes sourced from CivitAI image pages with
|
||||
* the complete page metadata (internal trpc data scraped with the user's
|
||||
* session), fixing recipes that the native import (REST API + EXIF only)
|
||||
* saved with 0 LoRAs.
|
||||
*
|
||||
* Protocol: DOM CustomEvents on `document`; `detail` is ALWAYS a JSON
|
||||
* string on both sides.
|
||||
*
|
||||
* LM page -> extension: `lm:reimportProbe`, detail `{}`.
|
||||
* extension -> LM page: `lm:reimportProbeResult`,
|
||||
* detail `{supported, licenseValid, extensionVersion?, reason?}`.
|
||||
* LM page -> extension: `lm:reimportViaExtension`,
|
||||
* detail `{requestId, recipes: [{recipeId, imageId, imageUrl, title}]}`.
|
||||
* extension -> LM page: `lm:reimportProgress`,
|
||||
* detail `{requestId, current, total, recipeId, title, status, message?}`.
|
||||
* extension -> LM page: `lm:reimportBatchDone`,
|
||||
* detail `{requestId, completed, failed}`.
|
||||
*/
|
||||
|
||||
const PROBE_EVENT = 'lm:reimportProbe';
|
||||
const PROBE_RESULT_EVENT = 'lm:reimportProbeResult';
|
||||
const REIMPORT_EVENT = 'lm:reimportViaExtension';
|
||||
const PROGRESS_EVENT = 'lm:reimportProgress';
|
||||
const BATCH_DONE_EVENT = 'lm:reimportBatchDone';
|
||||
|
||||
const DEFAULT_PROBE_TIMEOUT_MS = 500;
|
||||
// Generous batch timeout; any progress event resets it (heartbeat).
|
||||
const DEFAULT_REIMPORT_TIMEOUT_MS = 3 * 60 * 1000;
|
||||
|
||||
// Mirrors py/utils/civitai_utils.py (_SUPPORTED_CIVITAI_PAGE_HOSTS).
|
||||
const SUPPORTED_CIVITAI_PAGE_HOSTS = new Set([
|
||||
'civitai.com',
|
||||
'civitai.red',
|
||||
'civitai.green',
|
||||
]);
|
||||
|
||||
/**
|
||||
* Parse the JSON-string `detail` of a protocol event.
|
||||
* @param {CustomEvent} event
|
||||
* @returns {object|null} Parsed detail object, or null when absent/invalid.
|
||||
*/
|
||||
function parseDetail(event) {
|
||||
try {
|
||||
const detail = JSON.parse(event?.detail ?? 'null');
|
||||
return detail && typeof detail === 'object' ? detail : null;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Dispatch a protocol event with a JSON-stringified detail.
|
||||
* @param {string} type - Event name.
|
||||
* @param {object} payload - Detail payload (JSON-stringified).
|
||||
*/
|
||||
function dispatchProtocolEvent(type, payload) {
|
||||
document.dispatchEvent(
|
||||
new CustomEvent(type, { detail: JSON.stringify(payload ?? {}) })
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a correlation id for a re-import batch.
|
||||
* @returns {string}
|
||||
*/
|
||||
function generateRequestId() {
|
||||
if (globalThis.crypto?.randomUUID) {
|
||||
return globalThis.crypto.randomUUID();
|
||||
}
|
||||
return `lm-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Probe whether the companion extension is installed and usable.
|
||||
*
|
||||
* @param {{timeoutMs?: number}} [options]
|
||||
* @returns {Promise<{supported: boolean, licenseValid: boolean, extensionVersion?: string, reason?: string}|null>}
|
||||
* Resolves with the probe result, or null when the extension is absent or
|
||||
* too old to answer (timeout).
|
||||
*/
|
||||
export function probeExtension({ timeoutMs = DEFAULT_PROBE_TIMEOUT_MS } = {}) {
|
||||
return new Promise((resolve) => {
|
||||
let settled = false;
|
||||
const timer = setTimeout(() => finish(null), timeoutMs);
|
||||
|
||||
const finish = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
clearTimeout(timer);
|
||||
document.removeEventListener(PROBE_RESULT_EVENT, onResult);
|
||||
resolve(value);
|
||||
};
|
||||
|
||||
const onResult = (event) => {
|
||||
const detail = parseDetail(event);
|
||||
if (!detail) return;
|
||||
finish({
|
||||
supported: Boolean(detail.supported),
|
||||
licenseValid: Boolean(detail.licenseValid),
|
||||
extensionVersion: detail.extensionVersion,
|
||||
reason: detail.reason,
|
||||
});
|
||||
};
|
||||
|
||||
document.addEventListener(PROBE_RESULT_EVENT, onResult);
|
||||
dispatchProtocolEvent(PROBE_EVENT, {});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Delegate a batch of recipe re-imports to the companion extension.
|
||||
*
|
||||
* @param {Array<{recipeId: string, imageId: number, imageUrl: string, title: string}>} recipes
|
||||
* @param {{onProgress?: (progress: object) => void, timeoutMs?: number}} [options]
|
||||
* @returns {Promise<{completed: number, failed: number}>} Resolves on
|
||||
* `lm:reimportBatchDone`; rejects on timeout. Listeners are cleaned up in
|
||||
* all outcomes.
|
||||
*/
|
||||
export function delegateReimport(recipes, { onProgress, timeoutMs = DEFAULT_REIMPORT_TIMEOUT_MS } = {}) {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (!Array.isArray(recipes) || recipes.length === 0) {
|
||||
reject(new Error('delegateReimport requires a non-empty recipe list'));
|
||||
return;
|
||||
}
|
||||
|
||||
const requestId = generateRequestId();
|
||||
let settled = false;
|
||||
let timer = null;
|
||||
|
||||
const cleanup = () => {
|
||||
clearTimeout(timer);
|
||||
document.removeEventListener(PROGRESS_EVENT, onProgressEvent);
|
||||
document.removeEventListener(BATCH_DONE_EVENT, onBatchDone);
|
||||
};
|
||||
const succeed = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
resolve(value);
|
||||
};
|
||||
const fail = (error) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(error);
|
||||
};
|
||||
const armTimer = () => {
|
||||
clearTimeout(timer);
|
||||
timer = setTimeout(
|
||||
() => fail(new Error('Extension re-import timed out')),
|
||||
timeoutMs
|
||||
);
|
||||
};
|
||||
|
||||
const onProgressEvent = (event) => {
|
||||
const detail = parseDetail(event);
|
||||
if (!detail || detail.requestId !== requestId) return;
|
||||
// Heartbeat: any progress for this batch resets the timeout.
|
||||
armTimer();
|
||||
if (typeof onProgress === 'function') {
|
||||
try {
|
||||
onProgress(detail);
|
||||
} catch (error) {
|
||||
console.error('[extensionReimportBridge] onProgress callback failed:', error);
|
||||
}
|
||||
}
|
||||
};
|
||||
const onBatchDone = (event) => {
|
||||
const detail = parseDetail(event);
|
||||
if (!detail || detail.requestId !== requestId) return;
|
||||
succeed({
|
||||
completed: Number.isInteger(detail.completed) ? detail.completed : 0,
|
||||
failed: Number.isInteger(detail.failed) ? detail.failed : 0,
|
||||
});
|
||||
};
|
||||
|
||||
document.addEventListener(PROGRESS_EVENT, onProgressEvent);
|
||||
document.addEventListener(BATCH_DONE_EVENT, onBatchDone);
|
||||
armTimer();
|
||||
dispatchProtocolEvent(REIMPORT_EVENT, { requestId, recipes });
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract CivitAI image page info from a recipe source_path.
|
||||
* Mirrors py/utils/civitai_utils.py `extract_civitai_image_id`.
|
||||
*
|
||||
* @param {string|null} sourcePath - Recipe source_path.
|
||||
* @returns {{imageId: number, imageUrl: string}|null} Null when the path is
|
||||
* not a `/images/<id>` URL on civitai.com/.red/.green.
|
||||
*/
|
||||
export function getCivitaiImageInfo(sourcePath) {
|
||||
if (!sourcePath || typeof sourcePath !== 'string') {
|
||||
return null;
|
||||
}
|
||||
|
||||
let parsed;
|
||||
try {
|
||||
parsed = new URL(sourcePath);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
|
||||
return null;
|
||||
}
|
||||
if (!SUPPORTED_CIVITAI_PAGE_HOSTS.has(parsed.hostname.toLowerCase())) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const pathMatch = parsed.pathname.match(/\/images\/(\d+)/);
|
||||
if (!pathMatch) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return { imageId: Number(pathMatch[1]), imageUrl: sourcePath };
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
|
||||
const showToastMock = vi.fn();
|
||||
const showSimpleLoadingMock = vi.fn();
|
||||
const hideLoadingMock = vi.fn();
|
||||
const resetAndReloadMock = vi.fn();
|
||||
const probeExtensionMock = vi.fn();
|
||||
const delegateReimportMock = vi.fn();
|
||||
|
||||
const stateStub = {
|
||||
virtualScroller: { items: [] },
|
||||
loadingManager: {
|
||||
showSimpleLoading: showSimpleLoadingMock,
|
||||
hide: hideLoadingMock,
|
||||
},
|
||||
};
|
||||
|
||||
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
|
||||
showToast: showToastMock,
|
||||
copyToClipboard: vi.fn(),
|
||||
sendLoraToWorkflow: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/storageHelpers.js', () => ({
|
||||
setSessionItem: vi.fn(),
|
||||
removeSessionItem: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/recipeApi.js', () => ({
|
||||
updateRecipeMetadata: vi.fn(),
|
||||
resetAndReload: resetAndReloadMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/state/index.js', () => ({
|
||||
state: stateStub,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/managers/MoveManager.js', () => ({
|
||||
moveManager: { showMoveModal: vi.fn() },
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/ContextMenu/ModelContextMenuMixin.js', () => ({
|
||||
ModelContextMenuMixin: {
|
||||
handleCommonMenuActions: vi.fn(() => false),
|
||||
initNSFWSelector: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Keep the real getCivitaiImageInfo (gating logic under test); mock only the
|
||||
// extension communication.
|
||||
vi.mock('../../../static/js/utils/extensionReimportBridge.js', async (importOriginal) => {
|
||||
const actual = await importOriginal();
|
||||
return {
|
||||
...actual,
|
||||
probeExtension: probeExtensionMock,
|
||||
delegateReimport: delegateReimportMock,
|
||||
};
|
||||
});
|
||||
|
||||
describe('RecipeContextMenu.reimportRecipe extension delegation', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
document.body.innerHTML = `
|
||||
<div id="recipeContextMenu" class="context-menu" style="display: none;">
|
||||
<div class="context-menu-item" data-action="reimport"></div>
|
||||
</div>
|
||||
`;
|
||||
stateStub.virtualScroller.items = [
|
||||
{
|
||||
id: 'recipe-1',
|
||||
file_path: '/recipes/recipe-1.webp',
|
||||
title: 'Civitai Recipe',
|
||||
source_path: 'https://civitai.com/images/12345',
|
||||
},
|
||||
{
|
||||
id: 'recipe-2',
|
||||
file_path: '/recipes/recipe-2.webp',
|
||||
title: 'Local Recipe',
|
||||
source_path: '/data/imports/local.png',
|
||||
},
|
||||
];
|
||||
global.fetch = vi.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => ({ success: true, recipe_id: 'new-id', loras_count: 2 }),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete global.fetch;
|
||||
});
|
||||
|
||||
async function createMenu() {
|
||||
const { RecipeContextMenu } = await import(
|
||||
'../../../static/js/components/ContextMenu/RecipeContextMenu.js'
|
||||
);
|
||||
return new RecipeContextMenu();
|
||||
}
|
||||
|
||||
it('delegates to the extension for a CivitAI image source when licensed', async () => {
|
||||
probeExtensionMock.mockResolvedValue({ supported: true, licenseValid: true });
|
||||
delegateReimportMock.mockResolvedValue({ completed: 1, failed: 0 });
|
||||
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-1');
|
||||
|
||||
expect(delegateReimportMock).toHaveBeenCalledWith([{
|
||||
recipeId: 'recipe-1',
|
||||
imageId: 12345,
|
||||
imageUrl: 'https://civitai.com/images/12345',
|
||||
title: 'Civitai Recipe',
|
||||
}]);
|
||||
expect(global.fetch).not.toHaveBeenCalled();
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.recipes.reimportSuccess', {}, 'success');
|
||||
expect(resetAndReloadMock).toHaveBeenCalledWith(false, { preserveScroll: false });
|
||||
});
|
||||
|
||||
it('shows the failure toast when the extension reports failures', async () => {
|
||||
probeExtensionMock.mockResolvedValue({ supported: true, licenseValid: true });
|
||||
delegateReimportMock.mockResolvedValue({ completed: 0, failed: 1 });
|
||||
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-1');
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'recipes.contextMenu.reimport.failed',
|
||||
{ message: 'Extension re-import failed' },
|
||||
'error'
|
||||
);
|
||||
expect(global.fetch).not.toHaveBeenCalled();
|
||||
expect(resetAndReloadMock).toHaveBeenCalledWith(false, { preserveScroll: false });
|
||||
});
|
||||
|
||||
it('uses the native path for non-CivitAI sources without probing', async () => {
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-2');
|
||||
|
||||
expect(probeExtensionMock).not.toHaveBeenCalled();
|
||||
expect(delegateReimportMock).not.toHaveBeenCalled();
|
||||
expect(global.fetch).toHaveBeenCalledWith('/api/lm/recipe/recipe-2/reimport', {
|
||||
method: 'POST',
|
||||
});
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.recipes.reimportSuccess', {}, 'success');
|
||||
});
|
||||
|
||||
it('uses the native path when the extension is absent (probe timeout)', async () => {
|
||||
probeExtensionMock.mockResolvedValue(null);
|
||||
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-1');
|
||||
|
||||
expect(delegateReimportMock).not.toHaveBeenCalled();
|
||||
expect(global.fetch).toHaveBeenCalledWith('/api/lm/recipe/recipe-1/reimport', {
|
||||
method: 'POST',
|
||||
});
|
||||
});
|
||||
|
||||
it('uses the native path when the license is invalid', async () => {
|
||||
probeExtensionMock.mockResolvedValue({ supported: true, licenseValid: false });
|
||||
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-1');
|
||||
|
||||
expect(delegateReimportMock).not.toHaveBeenCalled();
|
||||
expect(global.fetch).toHaveBeenCalledWith('/api/lm/recipe/recipe-1/reimport', {
|
||||
method: 'POST',
|
||||
});
|
||||
});
|
||||
|
||||
it('falls back to the native path when delegation fails', async () => {
|
||||
probeExtensionMock.mockResolvedValue({ supported: true, licenseValid: true });
|
||||
delegateReimportMock.mockRejectedValue(new Error('Extension re-import timed out'));
|
||||
|
||||
const menu = await createMenu();
|
||||
await menu.reimportRecipe('recipe-1');
|
||||
|
||||
expect(global.fetch).toHaveBeenCalledWith('/api/lm/recipe/recipe-1/reimport', {
|
||||
method: 'POST',
|
||||
});
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.recipes.reimportSuccess', {}, 'success');
|
||||
});
|
||||
});
|
||||
@@ -153,6 +153,16 @@ const hashInvalidLora = {
|
||||
hashInvalid: true,
|
||||
};
|
||||
|
||||
// Mirrors the shape served for page-imported recipes whose CivitAI version
|
||||
// exposes no sha256: an exact modelVersionId but no modelId and no hash.
|
||||
const versionOnlyLora = {
|
||||
name: 'version-lora',
|
||||
modelName: 'Version Only LoRA',
|
||||
inLibrary: false,
|
||||
modelVersionId: 3221586,
|
||||
modelVersionName: 'V1 KREA-2',
|
||||
};
|
||||
|
||||
const recipeWithResources = {
|
||||
id: 'recipe-resources',
|
||||
file_path: '/recipes/resources.json',
|
||||
@@ -171,6 +181,7 @@ const recipeWithResources = {
|
||||
hashInvalidLora,
|
||||
{ name: 'mystery-lora', modelName: 'Mystery LoRA', inLibrary: false },
|
||||
hashOnlyLora,
|
||||
versionOnlyLora,
|
||||
],
|
||||
};
|
||||
|
||||
@@ -281,6 +292,57 @@ describe('RecipeModal resource item interactions', () => {
|
||||
);
|
||||
});
|
||||
|
||||
it('renders a download action (not reconnect) for a version-only LoRA', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
recipeModal.showRecipeDetails(recipeWithResources);
|
||||
await flushWiring();
|
||||
|
||||
const item = document.querySelector('[data-lora-index="6"]');
|
||||
expect(item).not.toBeNull();
|
||||
expect(item.classList.contains('missing-locally')).toBe(true);
|
||||
// Missing from the local library (badge) but still downloadable by its
|
||||
// exact CivitAI version id, so the row offers Download, not Reconnect.
|
||||
expect(item.querySelector('.missing-badge')).not.toBeNull();
|
||||
expect(item.querySelector('.lora-download')).not.toBeNull();
|
||||
expect(item.querySelector('.lora-reconnect')).toBeNull();
|
||||
});
|
||||
|
||||
it('downloads a version-only LoRA by resolving the model id from the version endpoint', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const requests = [];
|
||||
// Isolated copy keeps mutations out of the shared fixture.
|
||||
const isolatedRecipe = JSON.parse(JSON.stringify(recipeWithResources));
|
||||
fetchRecipeDetailsMock.mockResolvedValue(isolatedRecipe);
|
||||
global.fetch = vi.fn(async (url) => {
|
||||
requests.push(String(url));
|
||||
if (String(url).includes('/civitai/model/version/3221586')) {
|
||||
return {
|
||||
ok: true,
|
||||
json: async () => ({ id: 3221586, modelId: 56789, name: 'V1 KREA-2' }),
|
||||
};
|
||||
}
|
||||
return { ok: true, json: async () => ({}) };
|
||||
});
|
||||
recipeModal.showRecipeDetails(isolatedRecipe);
|
||||
await flushWiring();
|
||||
|
||||
const item = document.querySelector('[data-lora-index="6"]');
|
||||
item.querySelector('.lora-download').click();
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(downloadVersionWithDefaultsMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
expect(
|
||||
requests.some(u => u.includes('/civitai/model/version/3221586'))
|
||||
).toBe(true);
|
||||
expect(downloadVersionWithDefaultsMock).toHaveBeenCalledWith(
|
||||
'loras',
|
||||
56789,
|
||||
3221586,
|
||||
expect.objectContaining({ source: 'recipe-modal' })
|
||||
);
|
||||
});
|
||||
|
||||
it('does not navigate when a missing LoRA row is clicked', async () => {
|
||||
const recipeModal = await createRecipeModal();
|
||||
const navigateSpy = vi
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
import { describe, it, expect, vi, afterEach } from 'vitest';
|
||||
|
||||
import {
|
||||
probeExtension,
|
||||
delegateReimport,
|
||||
getCivitaiImageInfo,
|
||||
} from '../../../static/js/utils/extensionReimportBridge.js';
|
||||
|
||||
const dispatchedEvents = [];
|
||||
|
||||
function dispatchProtocolEvent(type, payload) {
|
||||
document.dispatchEvent(
|
||||
new CustomEvent(type, { detail: JSON.stringify(payload) })
|
||||
);
|
||||
}
|
||||
|
||||
// Installs a fake extension that answers probes with the given result.
|
||||
function installProbeResponder(result) {
|
||||
const listener = () => dispatchProtocolEvent('lm:reimportProbeResult', result);
|
||||
document.addEventListener('lm:reimportProbe', listener);
|
||||
return () => document.removeEventListener('lm:reimportProbe', listener);
|
||||
}
|
||||
|
||||
afterEach(() => {
|
||||
dispatchedEvents.length = 0;
|
||||
});
|
||||
|
||||
describe('probeExtension', () => {
|
||||
it('resolves null when no extension answers within the timeout', async () => {
|
||||
const result = await probeExtension({ timeoutMs: 20 });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('resolves the probe result when the extension answers', async () => {
|
||||
const uninstall = installProbeResponder({
|
||||
supported: true,
|
||||
licenseValid: true,
|
||||
extensionVersion: '1.2.3',
|
||||
});
|
||||
try {
|
||||
const result = await probeExtension({ timeoutMs: 1000 });
|
||||
expect(result).toEqual({
|
||||
supported: true,
|
||||
licenseValid: true,
|
||||
extensionVersion: '1.2.3',
|
||||
reason: undefined,
|
||||
});
|
||||
} finally {
|
||||
uninstall();
|
||||
}
|
||||
});
|
||||
|
||||
it('reports unsupported/unlicensed answers verbatim', async () => {
|
||||
const uninstall = installProbeResponder({
|
||||
supported: false,
|
||||
licenseValid: false,
|
||||
reason: 'license expired',
|
||||
});
|
||||
try {
|
||||
const result = await probeExtension({ timeoutMs: 1000 });
|
||||
expect(result.supported).toBe(false);
|
||||
expect(result.licenseValid).toBe(false);
|
||||
expect(result.reason).toBe('license expired');
|
||||
} finally {
|
||||
uninstall();
|
||||
}
|
||||
});
|
||||
|
||||
it('ignores malformed probe results and times out', async () => {
|
||||
const listener = () => {
|
||||
document.dispatchEvent(
|
||||
new CustomEvent('lm:reimportProbeResult', { detail: '{broken json' })
|
||||
);
|
||||
};
|
||||
document.addEventListener('lm:reimportProbe', listener);
|
||||
try {
|
||||
const result = await probeExtension({ timeoutMs: 20 });
|
||||
expect(result).toBeNull();
|
||||
} finally {
|
||||
document.removeEventListener('lm:reimportProbe', listener);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('delegateReimport', () => {
|
||||
const recipes = [
|
||||
{ recipeId: 'r1', imageId: 123, imageUrl: 'https://civitai.com/images/123', title: 'One' },
|
||||
{ recipeId: 'r2', imageId: 456, imageUrl: 'https://civitai.com/images/456', title: 'Two' },
|
||||
];
|
||||
|
||||
it('rejects immediately for an empty recipe list', async () => {
|
||||
await expect(delegateReimport([])).rejects.toThrow('non-empty');
|
||||
});
|
||||
|
||||
it('rejects on timeout when the extension never answers', async () => {
|
||||
await expect(
|
||||
delegateReimport(recipes, { timeoutMs: 20 })
|
||||
).rejects.toThrow('timed out');
|
||||
});
|
||||
|
||||
it('dispatches the batch with a requestId and resolves on batchDone', async () => {
|
||||
const progressEvents = [];
|
||||
let seenRequest = null;
|
||||
|
||||
const listener = (event) => {
|
||||
seenRequest = JSON.parse(event.detail);
|
||||
const { requestId } = seenRequest;
|
||||
// Progress for a DIFFERENT batch must be ignored.
|
||||
dispatchProtocolEvent('lm:reimportProgress', {
|
||||
requestId: 'other-batch',
|
||||
current: 99,
|
||||
total: 99,
|
||||
recipeId: 'nope',
|
||||
title: 'nope',
|
||||
status: 'success',
|
||||
});
|
||||
dispatchProtocolEvent('lm:reimportProgress', {
|
||||
requestId,
|
||||
current: 1,
|
||||
total: 2,
|
||||
recipeId: 'r1',
|
||||
title: 'One',
|
||||
status: 'success',
|
||||
});
|
||||
dispatchProtocolEvent('lm:reimportProgress', {
|
||||
requestId,
|
||||
current: 2,
|
||||
total: 2,
|
||||
recipeId: 'r2',
|
||||
title: 'Two',
|
||||
status: 'failed',
|
||||
message: 'boom',
|
||||
});
|
||||
dispatchProtocolEvent('lm:reimportBatchDone', {
|
||||
requestId,
|
||||
completed: 1,
|
||||
failed: 1,
|
||||
});
|
||||
};
|
||||
document.addEventListener('lm:reimportViaExtension', listener);
|
||||
try {
|
||||
const result = await delegateReimport(recipes, {
|
||||
onProgress: (progress) => progressEvents.push(progress),
|
||||
timeoutMs: 1000,
|
||||
});
|
||||
|
||||
expect(seenRequest.recipes).toEqual(recipes);
|
||||
expect(typeof seenRequest.requestId).toBe('string');
|
||||
expect(seenRequest.requestId.length).toBeGreaterThan(0);
|
||||
expect(result).toEqual({ completed: 1, failed: 1 });
|
||||
// Only this batch's progress events reach the callback.
|
||||
expect(progressEvents.map((p) => p.recipeId)).toEqual(['r1', 'r2']);
|
||||
expect(progressEvents[1].status).toBe('failed');
|
||||
} finally {
|
||||
document.removeEventListener('lm:reimportViaExtension', listener);
|
||||
}
|
||||
});
|
||||
|
||||
it('resets the timeout on every progress heartbeat', async () => {
|
||||
vi.useFakeTimers();
|
||||
let requestId = null;
|
||||
const listener = (event) => {
|
||||
requestId = JSON.parse(event.detail).requestId;
|
||||
};
|
||||
document.addEventListener('lm:reimportViaExtension', listener);
|
||||
try {
|
||||
const promise = delegateReimport(recipes, { timeoutMs: 1000 });
|
||||
|
||||
// At t=900ms a progress event arrives, pushing the deadline to t=1900ms.
|
||||
await vi.advanceTimersByTimeAsync(900);
|
||||
dispatchProtocolEvent('lm:reimportProgress', {
|
||||
requestId,
|
||||
current: 1,
|
||||
total: 2,
|
||||
recipeId: 'r1',
|
||||
title: 'One',
|
||||
status: 'started',
|
||||
});
|
||||
// t=1800ms: past the original deadline, still alive thanks to heartbeat.
|
||||
await vi.advanceTimersByTimeAsync(900);
|
||||
dispatchProtocolEvent('lm:reimportBatchDone', {
|
||||
requestId,
|
||||
completed: 2,
|
||||
failed: 0,
|
||||
});
|
||||
|
||||
await expect(promise).resolves.toEqual({ completed: 2, failed: 0 });
|
||||
} finally {
|
||||
document.removeEventListener('lm:reimportViaExtension', listener);
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('getCivitaiImageInfo', () => {
|
||||
it.each([
|
||||
'https://civitai.com/images/12345',
|
||||
'https://civitai.red/images/12345',
|
||||
'https://civitai.green/images/12345',
|
||||
'https://civitai.com/images/12345?foo=bar',
|
||||
])('extracts the image id from %s', (url) => {
|
||||
expect(getCivitaiImageInfo(url)).toEqual({ imageId: 12345, imageUrl: url });
|
||||
});
|
||||
|
||||
it.each([
|
||||
null,
|
||||
'',
|
||||
'not a url',
|
||||
'ftp://civitai.com/images/12345',
|
||||
'https://civitai.com/models/12345',
|
||||
'https://example.com/images/12345',
|
||||
'https://image.civitai.com/x/y/original=true/pic.png',
|
||||
])('returns null for %s', (url) => {
|
||||
expect(getCivitaiImageInfo(url)).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -2341,3 +2341,147 @@ async def test_get_recipe_detail_includes_recipe_json_path(
|
||||
assert response.status == 200
|
||||
payload = await response.json()
|
||||
assert "recipe_json_path" not in payload
|
||||
|
||||
|
||||
async def test_reimport_with_extension_payload_uses_payload_path(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""A re-import carrying the companion extension's metadata payload must
|
||||
use the payload-based import engine (caller-supplied LoRAs) instead of
|
||||
the legacy CivitAI image URL import, and report loras_count."""
|
||||
provider_calls: list[str | int] = []
|
||||
|
||||
class Provider:
|
||||
async def get_model_version_info(self, model_version_id):
|
||||
provider_calls.append(model_version_id)
|
||||
return {}, None
|
||||
|
||||
async def fake_get_default_metadata_provider():
|
||||
return Provider()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.recipes.enrichment.get_default_metadata_provider",
|
||||
fake_get_default_metadata_provider,
|
||||
)
|
||||
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
old_file = harness.tmp_dir / "recipes" / "sub" / "rec-ext.webp"
|
||||
harness.scanner.recipes["rec-ext"] = {
|
||||
"id": "rec-ext",
|
||||
"title": "Old title",
|
||||
"file_path": str(old_file),
|
||||
"tags": ["tag1"],
|
||||
"source_path": "https://civitai.com/images/12345",
|
||||
}
|
||||
harness.civitai.image_info["12345"] = {
|
||||
"id": 12345,
|
||||
"url": "https://image.civitai.com/x/y/original=true/pic.png",
|
||||
"type": "image",
|
||||
}
|
||||
harness.persistence.save_result = SimpleNamespace(
|
||||
payload={"success": True, "recipe_id": "new-rec-ext"}, status=200
|
||||
)
|
||||
# The freshly saved recipe as the scanner would see it (for loras_count).
|
||||
harness.scanner.recipes["new-rec-ext"] = {
|
||||
"id": "new-rec-ext",
|
||||
"loras": [{"file_name": "Painterly"}],
|
||||
}
|
||||
|
||||
resources = [
|
||||
{
|
||||
"type": "lora",
|
||||
"modelId": 20,
|
||||
"modelVersionId": 44,
|
||||
"modelName": "Painterly",
|
||||
"modelVersionName": "v2",
|
||||
"weight": 0.5,
|
||||
},
|
||||
]
|
||||
# The extension only issues GET requests (per its API convention).
|
||||
response = await harness.client.get(
|
||||
"/api/lm/recipe/rec-ext/reimport",
|
||||
params={
|
||||
"image_url": "https://civitai.com/images/12345",
|
||||
"name": "Extension Recipe",
|
||||
"resources": json.dumps(resources),
|
||||
"gen_params": json.dumps({"prompt": "from extension"}),
|
||||
"base_model": "Flux",
|
||||
},
|
||||
)
|
||||
payload = await response.json()
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert payload["old_recipe_id"] == "rec-ext"
|
||||
assert payload["recipe_id"] == "new-rec-ext"
|
||||
assert payload["loras_count"] == 1
|
||||
|
||||
save_call = harness.persistence.save_calls[-1]
|
||||
# Caller-supplied payload data wins: name, LoRAs, gen params.
|
||||
assert save_call["name"] == "Extension Recipe"
|
||||
assert save_call["metadata"]["loras"][0]["file_name"] == "Painterly"
|
||||
assert save_call["metadata"]["loras"][0]["weight"] == 0.5
|
||||
assert save_call["metadata"]["gen_params"]["prompt"] == "from extension"
|
||||
# Reimport semantics: original source_path and folder are preserved.
|
||||
assert save_call["metadata"]["source_path"] == "https://civitai.com/images/12345"
|
||||
assert save_call["target_dir"] == str(harness.tmp_dir / "recipes" / "sub")
|
||||
# The old recipe is deleted and user edits carried over.
|
||||
assert harness.persistence.delete_calls == ["rec-ext"]
|
||||
assert harness.persistence.update_calls[-1]["recipe_id"] == "new-rec-ext"
|
||||
assert harness.persistence.update_calls[-1]["updates"]["title"] == "Old title"
|
||||
assert harness.persistence.update_calls[-1]["updates"]["tags"] == ["tag1"]
|
||||
|
||||
|
||||
async def test_reimport_with_malformed_payload_falls_back_to_legacy(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""Malformed resources JSON must be treated as "no payload": the legacy
|
||||
source-URL import runs and the request still succeeds."""
|
||||
async def fake_get_default_metadata_provider():
|
||||
return SimpleNamespace(get_model_version_info=lambda id: ({}, None))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.recipes.enrichment.get_default_metadata_provider",
|
||||
fake_get_default_metadata_provider,
|
||||
)
|
||||
|
||||
async with recipe_harness(monkeypatch, tmp_path) as harness:
|
||||
harness.scanner.recipes["rec-bad"] = {
|
||||
"id": "rec-bad",
|
||||
"title": "Broken payload",
|
||||
"file_path": str(harness.tmp_dir / "recipes" / "rec-bad.webp"),
|
||||
"tags": [],
|
||||
"source_path": "https://civitai.com/images/12345",
|
||||
}
|
||||
harness.civitai.image_info["12345"] = {
|
||||
"id": 12345,
|
||||
"url": "https://image.civitai.com/x/y/original=true/pic.png",
|
||||
"type": "image",
|
||||
}
|
||||
harness.persistence.save_result = SimpleNamespace(
|
||||
payload={"success": True, "recipe_id": "legacy-new"}, status=200
|
||||
)
|
||||
harness.scanner.recipes["legacy-new"] = {"id": "legacy-new", "loras": []}
|
||||
|
||||
response = await harness.client.get(
|
||||
"/api/lm/recipe/rec-bad/reimport",
|
||||
params={
|
||||
"image_url": "https://civitai.com/images/12345",
|
||||
"name": "Ignored Name",
|
||||
"resources": "{not valid json",
|
||||
},
|
||||
)
|
||||
payload = await response.json()
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert payload["recipe_id"] == "legacy-new"
|
||||
assert payload["loras_count"] == 0
|
||||
|
||||
save_call = harness.persistence.save_calls[-1]
|
||||
# Legacy URL path: the payload name is ignored and the title is
|
||||
# derived from the (empty) metadata, and no caller LoRAs are used.
|
||||
assert save_call["name"] == "Civitai Image 12345"
|
||||
assert save_call["metadata"]["loras"] == []
|
||||
assert save_call["metadata"]["source_path"] == "https://civitai.com/images/12345"
|
||||
assert harness.persistence.delete_calls == ["rec-bad"]
|
||||
|
||||
@@ -818,3 +818,44 @@ async def test_get_model_by_hash_rejects_empty_placeholder_without_request(downl
|
||||
assert result is None
|
||||
assert error == "Model not found"
|
||||
assert requested == []
|
||||
|
||||
|
||||
async def test_get_version_file_mini_returns_payload(downloader):
|
||||
"""The mini endpoint returns the raw stored filename (#1100)."""
|
||||
client = await CivitaiClient.get_instance()
|
||||
|
||||
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
||||
assert method == "GET"
|
||||
assert url.endswith("/model-versions/mini/3284136")
|
||||
assert kwargs.get("params") == {"modelFileId": 3168412}
|
||||
assert use_auth is True
|
||||
return True, {"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
||||
|
||||
downloader.make_request = fake_make_request
|
||||
|
||||
result = await client.get_version_file_mini(3284136, 3168412)
|
||||
|
||||
assert result == {"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
||||
|
||||
|
||||
async def test_get_version_file_mini_returns_none_on_failure(downloader):
|
||||
client = await CivitaiClient.get_instance()
|
||||
|
||||
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
||||
return False, "Model file 2 not found in version 1"
|
||||
|
||||
downloader.make_request = fake_make_request
|
||||
|
||||
assert await client.get_version_file_mini(1, 2) is None
|
||||
|
||||
|
||||
async def test_get_version_file_mini_propagates_rate_limit(downloader):
|
||||
client = await CivitaiClient.get_instance()
|
||||
|
||||
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
||||
return False, RateLimitError("limited", retry_after=1.0)
|
||||
|
||||
downloader.make_request = fake_make_request
|
||||
|
||||
with pytest.raises(RateLimitError):
|
||||
await client.get_version_file_mini(1, 2)
|
||||
|
||||
@@ -2098,3 +2098,174 @@ async def test_discard_cleared_downloads_stops_tracking_and_preserves_files(
|
||||
# Partial files are preserved for a future resume from disk.
|
||||
assert save_path.exists()
|
||||
assert control_path.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_uses_raw_file_name_from_mini_endpoint(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""#1100: when the REST name is rewritten ("{model}_{version}"), the raw
|
||||
stored filename from the mini endpoint wins for the on-disk name."""
|
||||
manager = DownloadManager()
|
||||
get_settings_manager().settings["default_unet_root"] = str(tmp_path / "unet")
|
||||
metadata_provider.payload = {
|
||||
"id": 3284136,
|
||||
"model": {"type": "Checkpoint", "tags": ["realistic"]},
|
||||
"baseModel": "ZImageTurbo",
|
||||
"creator": {"username": "Author"},
|
||||
"files": [
|
||||
{
|
||||
"id": 3168412,
|
||||
"type": "Model",
|
||||
"primary": True,
|
||||
"name": "cyberrealisticZImage_v80.safetensors",
|
||||
"downloadUrl": "https://civitai.com/api/download/models/3284136?fileId=3168412",
|
||||
}
|
||||
],
|
||||
}
|
||||
metadata_provider.get_version_file_mini = AsyncMock(
|
||||
return_value={"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
||||
)
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured["download_urls"] = kwargs["download_urls"]
|
||||
captured["file_path"] = kwargs["metadata"].file_path
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=3284136,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True, result
|
||||
metadata_provider.get_version_file_mini.assert_awaited_once_with(3284136, 3168412)
|
||||
assert captured["file_path"].endswith("CyberRealistic_zit_v8.0_bf16.safetensors")
|
||||
# The file's own pinned downloadUrl is untouched.
|
||||
assert captured["download_urls"] == [
|
||||
"https://civitai.com/api/download/models/3284136?fileId=3168412"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_falls_back_to_rest_name_when_mini_fails(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""A failed/absent mini lookup must keep the previous behavior."""
|
||||
manager = DownloadManager()
|
||||
metadata_provider.payload = {
|
||||
"id": 42,
|
||||
"model": {"type": "Checkpoint", "tags": ["fantasy"]},
|
||||
"baseModel": "BaseModel",
|
||||
"creator": {"username": "Author"},
|
||||
"files": [
|
||||
{
|
||||
"id": 1001,
|
||||
"type": "Model",
|
||||
"primary": True,
|
||||
"name": "rewritten_v10.safetensors",
|
||||
"downloadUrl": "https://example.invalid/file.safetensors",
|
||||
}
|
||||
],
|
||||
}
|
||||
metadata_provider.get_version_file_mini = AsyncMock(return_value=None)
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured["file_path"] = kwargs["metadata"].file_path
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=42,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
assert captured["file_path"].endswith("rewritten_v10.safetensors")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_skips_mini_lookup_for_civarchive_source(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""CivArchive already serves raw stored names — no mini call."""
|
||||
manager = DownloadManager()
|
||||
mini_mock = AsyncMock(return_value={"fileName": "should_not_be_used.safetensors"})
|
||||
metadata_provider.get_version_file_mini = mini_mock
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_manager,
|
||||
"get_metadata_provider",
|
||||
AsyncMock(return_value=metadata_provider),
|
||||
)
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(self, **kwargs):
|
||||
captured["file_path"] = kwargs["metadata"].file_path
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=99,
|
||||
save_dir=str(tmp_path),
|
||||
use_default_paths=True,
|
||||
progress_callback=None,
|
||||
source="civarchive",
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
mini_mock.assert_not_called()
|
||||
assert captured["file_path"].endswith("file.safetensors")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_raw_file_name_edge_cases():
|
||||
"""_fetch_raw_file_name never raises and strips path components."""
|
||||
manager = DownloadManager()
|
||||
|
||||
provider = SimpleNamespace()
|
||||
|
||||
# Missing version id / file id short-circuit before any provider call.
|
||||
provider.get_version_file_mini = AsyncMock()
|
||||
assert await manager._fetch_raw_file_name(provider, None, 1) is None
|
||||
assert await manager._fetch_raw_file_name(provider, 1, None) is None
|
||||
provider.get_version_file_mini.assert_not_called()
|
||||
|
||||
# Provider without the method (older mocks / non-CivitAI providers).
|
||||
assert await manager._fetch_raw_file_name(object(), 1, 2) is None
|
||||
|
||||
# Non-dict payload, empty fileName.
|
||||
provider.get_version_file_mini = AsyncMock(return_value="oops")
|
||||
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|
||||
provider.get_version_file_mini = AsyncMock(return_value={"fileName": " "})
|
||||
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|
||||
|
||||
# Path components are stripped defensively.
|
||||
provider.get_version_file_mini = AsyncMock(
|
||||
return_value={"fileName": "../evil/model.safetensors"}
|
||||
)
|
||||
assert await manager._fetch_raw_file_name(provider, 1, 2) == "model.safetensors"
|
||||
|
||||
# Provider exceptions degrade to None.
|
||||
provider.get_version_file_mini = AsyncMock(side_effect=RuntimeError("boom"))
|
||||
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|
||||
|
||||
@@ -207,3 +207,62 @@ async def test_retry_helper_retries_normally_for_small_retry_after(monkeypatch):
|
||||
result, _ = await helper.run("test", succeeding)
|
||||
assert result == {"ok": True}
|
||||
assert calls == 2 # Retried once (small retry_after)
|
||||
|
||||
|
||||
class MiniCapableProvider(ModelMetadataProvider):
|
||||
"""Provider that serves raw file names via the mini endpoint (#1100)."""
|
||||
|
||||
def __init__(self, payload=None) -> None:
|
||||
self.payload = payload
|
||||
self.calls = []
|
||||
|
||||
async def get_model_by_hash(self, model_hash: str):
|
||||
return None, None
|
||||
|
||||
async def get_model_versions(self, model_id: str):
|
||||
return None
|
||||
|
||||
async def get_model_version(self, model_id=None, version_id=None):
|
||||
return None
|
||||
|
||||
async def get_model_version_info(self, version_id: str):
|
||||
return None, None
|
||||
|
||||
async def get_user_models(self, username: str, cursor=None):
|
||||
return None
|
||||
|
||||
async def get_version_file_mini(self, version_id: int, file_id: int):
|
||||
self.calls.append((version_id, file_id))
|
||||
return self.payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_provider_get_version_file_mini_defaults_to_none():
|
||||
provider = TrackingProvider()
|
||||
assert await provider.get_version_file_mini(1, 2) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_get_version_file_mini_returns_first_hit():
|
||||
primary = TrackingProvider() # base default: None
|
||||
secondary = MiniCapableProvider({"fileName": "raw.safetensors"})
|
||||
|
||||
fallback = FallbackMetadataProvider(
|
||||
[("primary", primary), ("secondary", secondary)],
|
||||
)
|
||||
|
||||
result = await fallback.get_version_file_mini(10, 20)
|
||||
|
||||
assert result == {"fileName": "raw.safetensors"}
|
||||
assert secondary.calls == [(10, 20)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_retrying_provider_delegates_get_version_file_mini():
|
||||
inner = MiniCapableProvider({"fileName": "raw.safetensors"})
|
||||
wrapper = RateLimitRetryingProvider(inner, label="inner")
|
||||
|
||||
result = await wrapper.get_version_file_mini(10, 20)
|
||||
|
||||
assert result == {"fileName": "raw.safetensors"}
|
||||
assert inner.calls == [(10, 20)]
|
||||
|
||||
Reference in New Issue
Block a user