mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-26 07:35:44 -03:00
fix: verify returned image ID matches requested ID in CivitAI API
Fix issue #870 where importing recipes from CivitAI image URLs would return the wrong image when the API response did not contain the requested image ID. The get_image_info() method now: - Iterates through all returned items to find matching ID - Returns None when no match is found and logs warning with returned IDs - Handles invalid (non-numeric) ID formats New test cases: - test_get_image_info_returns_matching_item - test_get_image_info_returns_none_when_id_mismatch - test_get_image_info_handles_invalid_id Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -490,14 +490,33 @@ class CivitaiClient:
|
|||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
url = f"{self.base_url}/images?imageId={image_id}&nsfw=X"
|
url = f"{self.base_url}/images?imageId={image_id}&nsfw=X"
|
||||||
|
requested_id = int(image_id)
|
||||||
|
|
||||||
logger.debug(f"Fetching image info for ID: {image_id}")
|
logger.debug(f"Fetching image info for ID: {image_id}")
|
||||||
success, result = await self._make_request("GET", url, use_auth=True)
|
success, result = await self._make_request("GET", url, use_auth=True)
|
||||||
|
|
||||||
if success:
|
if success:
|
||||||
if result and "items" in result and len(result["items"]) > 0:
|
if result and "items" in result and isinstance(result["items"], list):
|
||||||
logger.debug(f"Successfully fetched image info for ID: {image_id}")
|
items = result["items"]
|
||||||
return result["items"][0]
|
|
||||||
|
# First, try to find the item with matching ID
|
||||||
|
for item in items:
|
||||||
|
if isinstance(item, dict) and item.get("id") == requested_id:
|
||||||
|
logger.debug(f"Successfully fetched image info for ID: {image_id}")
|
||||||
|
return item
|
||||||
|
|
||||||
|
# No matching ID found - log warning with details about returned items
|
||||||
|
returned_ids = [
|
||||||
|
item.get("id") for item in items
|
||||||
|
if isinstance(item, dict) and "id" in item
|
||||||
|
]
|
||||||
|
logger.warning(
|
||||||
|
f"CivitAI API returned no matching image for requested ID {image_id}. "
|
||||||
|
f"Returned {len(items)} item(s) with IDs: {returned_ids}. "
|
||||||
|
f"This may indicate the image was deleted, hidden, or there is a database lag."
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
logger.warning(f"No image found with ID: {image_id}")
|
logger.warning(f"No image found with ID: {image_id}")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -505,6 +524,10 @@ class CivitaiClient:
|
|||||||
return None
|
return None
|
||||||
except RateLimitError:
|
except RateLimitError:
|
||||||
raise
|
raise
|
||||||
|
except ValueError as e:
|
||||||
|
error_msg = f"Invalid image ID format: {image_id}"
|
||||||
|
logger.error(error_msg)
|
||||||
|
return None
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_msg = f"Error fetching image info: {e}"
|
error_msg = f"Error fetching image info: {e}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
|
|||||||
@@ -484,9 +484,11 @@ async def test_get_model_version_info_success(monkeypatch, downloader):
|
|||||||
assert result["images"][0]["meta"]["other"] == "keep"
|
assert result["images"][0]["meta"]["other"] == "keep"
|
||||||
|
|
||||||
|
|
||||||
async def test_get_image_info_returns_first_item(monkeypatch, downloader):
|
async def test_get_image_info_returns_matching_item(monkeypatch, downloader):
|
||||||
|
"""When API returns multiple items, return the one matching the requested ID."""
|
||||||
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
||||||
return True, {"items": [{"id": 1}, {"id": 2}]}
|
# Requested ID is 42, but it's the second item in the response
|
||||||
|
return True, {"items": [{"id": 41}, {"id": 42, "name": "target"}, {"id": 43}]}
|
||||||
|
|
||||||
downloader.make_request = fake_make_request
|
downloader.make_request = fake_make_request
|
||||||
|
|
||||||
@@ -494,7 +496,25 @@ async def test_get_image_info_returns_first_item(monkeypatch, downloader):
|
|||||||
|
|
||||||
result = await client.get_image_info("42")
|
result = await client.get_image_info("42")
|
||||||
|
|
||||||
assert result == {"id": 1}
|
assert result == {"id": 42, "name": "target"}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_image_info_returns_none_when_id_mismatch(monkeypatch, downloader, caplog):
|
||||||
|
"""When API returns items but none match the requested ID, return None and log warning."""
|
||||||
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
||||||
|
# Requested ID is 999, but API returns different IDs (simulating deleted/hidden image)
|
||||||
|
return True, {"items": [{"id": 1}, {"id": 2}, {"id": 3}]}
|
||||||
|
|
||||||
|
downloader.make_request = fake_make_request
|
||||||
|
|
||||||
|
client = await CivitaiClient.get_instance()
|
||||||
|
|
||||||
|
result = await client.get_image_info("999")
|
||||||
|
|
||||||
|
assert result is None
|
||||||
|
# Verify warning was logged
|
||||||
|
assert "CivitAI API returned no matching image for requested ID 999" in caplog.text
|
||||||
|
assert "Returned 3 item(s) with IDs: [1, 2, 3]" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
async def test_get_image_info_handles_missing(monkeypatch, downloader):
|
async def test_get_image_info_handles_missing(monkeypatch, downloader):
|
||||||
@@ -508,3 +528,13 @@ async def test_get_image_info_handles_missing(monkeypatch, downloader):
|
|||||||
result = await client.get_image_info("42")
|
result = await client.get_image_info("42")
|
||||||
|
|
||||||
assert result is None
|
assert result is None
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_image_info_handles_invalid_id(monkeypatch, downloader, caplog):
|
||||||
|
"""When given a non-numeric image ID, return None and log error."""
|
||||||
|
client = await CivitaiClient.get_instance()
|
||||||
|
|
||||||
|
result = await client.get_image_info("not-a-number")
|
||||||
|
|
||||||
|
assert result is None
|
||||||
|
assert "Invalid image ID format" in caplog.text
|
||||||
|
|||||||
Reference in New Issue
Block a user