mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-04-02 02:38:52 -03:00
fix(recipe): Import LoRAs from Civitai image URLs using modelVersionIds (#868)
When importing recipes from Civitai image URLs, the API returns modelVersionIds at the root level instead of inside the meta object. This caused LoRA information to not be recognized and imported. Changes: - analysis_service.py: Merge modelVersionIds from image_info into metadata - civitai_image.py: Add modelVersionIds field recognition and processing logic - test_civitai_image_parser.py: Add test for modelVersionIds handling
This commit is contained in:
@@ -42,6 +42,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
"height",
|
||||
"Model",
|
||||
"Model hash",
|
||||
"modelVersionIds",
|
||||
)
|
||||
return any(key in payload for key in civitai_image_fields)
|
||||
|
||||
@@ -429,6 +430,65 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
|
||||
result["loras"].append(lora_entry)
|
||||
|
||||
# Process modelVersionIds from Civitai image API
|
||||
# These are model version IDs returned at root level when meta doesn't contain resources
|
||||
if "modelVersionIds" in metadata and isinstance(
|
||||
metadata["modelVersionIds"], list
|
||||
):
|
||||
for version_id in metadata["modelVersionIds"]:
|
||||
version_id_str = str(version_id)
|
||||
|
||||
# Skip if we've already added this LoRA by version ID
|
||||
if version_id_str in added_loras:
|
||||
continue
|
||||
|
||||
# Initialize lora entry with version ID
|
||||
lora_entry = {
|
||||
"id": version_id,
|
||||
"modelId": 0,
|
||||
"name": "Unknown LoRA",
|
||||
"version": "",
|
||||
"type": "lora",
|
||||
"weight": 1.0,
|
||||
"existsLocally": False,
|
||||
"thumbnailUrl": "/loras_static/images/no-preview.png",
|
||||
"baseModel": "",
|
||||
"size": 0,
|
||||
"downloadUrl": "",
|
||||
"isDeleted": False,
|
||||
}
|
||||
|
||||
# Fetch model info from Civitai
|
||||
if metadata_provider and version_id_str:
|
||||
try:
|
||||
civitai_info = (
|
||||
await metadata_provider.get_model_version_info(
|
||||
version_id_str
|
||||
)
|
||||
)
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts,
|
||||
)
|
||||
|
||||
if populated_entry is None:
|
||||
continue # Skip invalid LoRA types
|
||||
|
||||
lora_entry = populated_entry
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error fetching Civitai info for model version {version_id}: {e}"
|
||||
)
|
||||
|
||||
# Track this LoRA for deduplication
|
||||
if version_id_str:
|
||||
added_loras[version_id_str] = len(result["loras"])
|
||||
|
||||
result["loras"].append(lora_entry)
|
||||
|
||||
# If we found LoRA hashes in the metadata but haven't already
|
||||
# populated entries for them, fall back to creating LoRAs from
|
||||
# the hashes section. Some Civitai image responses only include
|
||||
|
||||
@@ -143,6 +143,12 @@ class RecipeAnalysisService:
|
||||
):
|
||||
metadata = metadata["meta"]
|
||||
|
||||
# Include modelVersionIds from root level if available
|
||||
# Civitai API returns modelVersionIds at root level, not in meta
|
||||
model_version_ids = image_info.get("modelVersionIds")
|
||||
if model_version_ids and isinstance(metadata, dict):
|
||||
metadata["modelVersionIds"] = model_version_ids
|
||||
|
||||
# Validate that metadata contains meaningful recipe fields
|
||||
# If not, treat as None to trigger EXIF extraction from downloaded image
|
||||
if isinstance(metadata, dict) and not self._has_recipe_fields(metadata):
|
||||
|
||||
@@ -184,7 +184,10 @@ async def test_parse_metadata_populates_checkpoint_and_rewrites_thumbnails(monke
|
||||
assert result["model"] is not None
|
||||
assert result["model"]["name"] == "Checkpoint Example"
|
||||
assert result["model"]["type"] == "checkpoint"
|
||||
assert result["model"]["thumbnailUrl"] == "https://image.civitai.com/checkpoints/width=450,optimized=true"
|
||||
assert (
|
||||
result["model"]["thumbnailUrl"]
|
||||
== "https://image.civitai.com/checkpoints/width=450,optimized=true"
|
||||
)
|
||||
assert result["model"]["modelId"] == 111
|
||||
assert result["model"]["size"] == 1024 * 1024
|
||||
assert result["model"]["hash"] == "ffaa0011"
|
||||
@@ -192,5 +195,106 @@ async def test_parse_metadata_populates_checkpoint_and_rewrites_thumbnails(monke
|
||||
|
||||
assert result["loras"]
|
||||
assert result["loras"][0]["name"] == "Example Lora Model"
|
||||
assert result["loras"][0]["thumbnailUrl"] == "https://image.civitai.com/loras/width=450,optimized=true"
|
||||
assert (
|
||||
result["loras"][0]["thumbnailUrl"]
|
||||
== "https://image.civitai.com/loras/width=450,optimized=true"
|
||||
)
|
||||
assert result["loras"][0]["hash"] == "abc123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_metadata_handles_modelVersionIds(monkeypatch):
|
||||
"""Test that modelVersionIds from Civitai image API are properly processed."""
|
||||
lora_info_1 = {
|
||||
"id": 2398829,
|
||||
"modelId": 123456,
|
||||
"model": {"name": "Dance LoRA 1", "type": "lora"},
|
||||
"name": "Version 1.0",
|
||||
"images": [{"url": "https://image.civitai.com/lora1/original=true"}],
|
||||
"baseModel": "SDXL",
|
||||
"downloadUrl": "https://civitai.com/lora1/download",
|
||||
"files": [
|
||||
{
|
||||
"type": "Model",
|
||||
"primary": True,
|
||||
"sizeKB": 10240,
|
||||
"name": "dance_lora_1.safetensors",
|
||||
"hashes": {"SHA256": "aabbccdd0011"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
lora_info_2 = {
|
||||
"id": 2398838,
|
||||
"modelId": 123457,
|
||||
"model": {"name": "Style LoRA 2", "type": "lora"},
|
||||
"name": "Version 2.0",
|
||||
"images": [{"url": "https://image.civitai.com/lora2/original=true"}],
|
||||
"baseModel": "SDXL",
|
||||
"downloadUrl": "https://civitai.com/lora2/download",
|
||||
"files": [
|
||||
{
|
||||
"type": "Model",
|
||||
"primary": True,
|
||||
"sizeKB": 20480,
|
||||
"name": "style_lora_2.safetensors",
|
||||
"hashes": {"SHA256": "aabbccdd0022"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
async def fake_metadata_provider():
|
||||
class Provider:
|
||||
async def get_model_version_info(self, version_id):
|
||||
if version_id == "2398829":
|
||||
return lora_info_1, None
|
||||
if version_id == "2398838":
|
||||
return lora_info_2, None
|
||||
return None, "Model not found"
|
||||
|
||||
return Provider()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.recipes.parsers.civitai_image.get_default_metadata_provider",
|
||||
fake_metadata_provider,
|
||||
)
|
||||
|
||||
parser = CivitaiApiMetadataParser()
|
||||
|
||||
# This simulates the metadata from Civitai image API where modelVersionIds
|
||||
# is at the root level and meta only contains basic prompt info
|
||||
metadata = {
|
||||
"id": 109882763,
|
||||
"meta": {
|
||||
"id": 109882763,
|
||||
"meta": {"prompt": "A woman does the hip bump dance."},
|
||||
},
|
||||
"modelVersionIds": [2398829, 2398838],
|
||||
}
|
||||
|
||||
assert parser.is_metadata_matching(metadata)
|
||||
|
||||
result = await parser.parse_metadata(metadata)
|
||||
|
||||
# Verify both LoRAs were created from modelVersionIds
|
||||
assert len(result["loras"]) == 2
|
||||
|
||||
# Check first LoRA
|
||||
lora1 = result["loras"][0]
|
||||
assert lora1["id"] == 2398829
|
||||
assert lora1["name"] == "Dance LoRA 1"
|
||||
assert lora1["type"] == "lora"
|
||||
assert lora1["hash"] == "aabbccdd0011"
|
||||
assert lora1["baseModel"] == "SDXL"
|
||||
assert (
|
||||
lora1["thumbnailUrl"]
|
||||
== "https://image.civitai.com/lora1/width=450,optimized=true"
|
||||
)
|
||||
|
||||
# Check second LoRA
|
||||
lora2 = result["loras"][1]
|
||||
assert lora2["id"] == 2398838
|
||||
assert lora2["name"] == "Style LoRA 2"
|
||||
assert lora2["type"] == "lora"
|
||||
assert lora2["hash"] == "aabbccdd0022"
|
||||
assert lora2["baseModel"] == "SDXL"
|
||||
|
||||
Reference in New Issue
Block a user