From 12c88835f2285c14e3901db073a0ba9a65c0fc8d Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Wed, 24 Sep 2025 09:16:02 +0800
Subject: [PATCH] refactor: enhance model version retrieval logic in
CivitaiClient, fixes #460
---
py/services/civitai_client.py | 84 ++++++++++----
refs/civitai_api_model_by_modelId.json | 110 +++++++++++++++++++
tests/routes/test_base_model_routes_smoke.py | 2 +-
3 files changed, 176 insertions(+), 20 deletions(-)
create mode 100644 refs/civitai_api_model_by_modelId.json
diff --git a/py/services/civitai_client.py b/py/services/civitai_client.py
index ccadbcb5..bb6004ed 100644
--- a/py/services/civitai_client.py
+++ b/py/services/civitai_client.py
@@ -1,4 +1,5 @@
import os
+import copy
import logging
import asyncio
from typing import Optional, Dict, Tuple, List
@@ -189,31 +190,76 @@ class CivitaiClient:
)
if not success:
return None
-
+
model_versions = data.get('modelVersions', [])
-
- # Step 2: Determine the version_id to use
- target_version_id = version_id
- if target_version_id is None:
- target_version_id = model_versions[0].get('id')
-
- # Step 3: Get detailed version info using the version_id
- success, version = await downloader.make_request(
- 'GET',
- f"{self.base_url}/model-versions/{target_version_id}",
- use_auth=True
- )
- if not success:
+ if not model_versions:
+ logger.warning(f"No model versions found for model {model_id}")
return None
-
+
+ # Step 2: Determine the target version entry to use
+ target_version = None
+ if version_id is not None:
+ target_version = next(
+ (item for item in model_versions if item.get('id') == version_id),
+ None
+ )
+ if target_version is None:
+ logger.warning(
+ f"Version {version_id} not found for model {model_id}, defaulting to first version"
+ )
+ if target_version is None:
+ target_version = model_versions[0]
+
+ target_version_id = target_version.get('id')
+
+ # Step 3: Get detailed version info using the SHA256 hash
+ model_hash = None
+ for file_info in target_version.get('files', []):
+ if file_info.get('type') == 'Model' and file_info.get('primary'):
+ model_hash = file_info.get('hashes', {}).get('SHA256')
+ if model_hash:
+ break
+
+ version = None
+ if model_hash:
+ success, version = await downloader.make_request(
+ 'GET',
+ f"{self.base_url}/model-versions/by-hash/{model_hash}",
+ use_auth=True
+ )
+ if not success:
+ logger.warning(
+ f"Failed to fetch version by hash for model {model_id} version {target_version_id}: {version}"
+ )
+ version = None
+ else:
+ logger.warning(
+ f"No primary model hash found for model {model_id} version {target_version_id}"
+ )
+
+ if version is None:
+ version = copy.deepcopy(target_version)
+ version.pop('index', None)
+ version['modelId'] = model_id
+ version['model'] = {
+ 'name': data.get('name'),
+ 'type': data.get('type'),
+ 'nsfw': data.get('nsfw'),
+ 'poi': data.get('poi')
+ }
+
# Step 4: Enrich version_info with model data
# Add description and tags from model data
- version['model']['description'] = data.get("description")
- version['model']['tags'] = data.get("tags", [])
-
+ model_info = version.get('model')
+ if not isinstance(model_info, dict):
+ model_info = {}
+ version['model'] = model_info
+ model_info['description'] = data.get("description")
+ model_info['tags'] = data.get("tags", [])
+
# Add creator from model data
version['creator'] = data.get("creator")
-
+
return version
# Case 3: Neither model_id nor version_id provided
diff --git a/refs/civitai_api_model_by_modelId.json b/refs/civitai_api_model_by_modelId.json
new file mode 100644
index 00000000..2cd20f20
--- /dev/null
+++ b/refs/civitai_api_model_by_modelId.json
@@ -0,0 +1,110 @@
+{
+ "id": 1231067,
+ "name": "Vivid Impressions Storybook Style",
+ "description": "If you'd like to support me, feel free to visit my Ko-Fi page. ❤️
Please share your images using the \"+add post\" button below. It supports the creators. Thanks! 💕
Trigger word: ppstorybook
Strength: 0.8, experiment as you like
",
+ "allowNoCredit": true,
+ "allowCommercialUse": [
+ "Image",
+ "RentCivit",
+ "Rent",
+ "Sell"
+ ],
+ "allowDerivatives": true,
+ "allowDifferentLicense": true,
+ "type": "LORA",
+ "minor": false,
+ "sfwOnly": false,
+ "poi": false,
+ "nsfw": false,
+ "nsfwLevel": 1,
+ "availability": "Public",
+ "cosmetic": null,
+ "supportsGeneration": true,
+ "stats": {
+ "downloadCount": 2183,
+ "favoriteCount": 0,
+ "thumbsUpCount": 416,
+ "thumbsDownCount": 0,
+ "commentCount": 12,
+ "ratingCount": 0,
+ "rating": 0,
+ "tippedAmountCount": 360
+ },
+ "creator": {
+ "username": "PixelPawsAI",
+ "image": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/f3a1aa7c-0159-4dd8-884a-1e7ceb350f96/width=96/PixelPawsAI.jpeg"
+ },
+ "tags": [
+ "style",
+ "illustration",
+ "storybook"
+ ],
+ "modelVersions": [
+ {
+ "id": 1387174,
+ "index": 0,
+ "name": "v1.0",
+ "baseModel": "Flux.1 D",
+ "baseModelType": "Standard",
+ "createdAt": "2025-02-08T11:15:47.197Z",
+ "publishedAt": "2025-02-08T11:29:04.487Z",
+ "status": "Published",
+ "availability": "Public",
+ "nsfwLevel": 1,
+ "trainedWords": [
+ "ppstorybook"
+ ],
+ "covered": true,
+ "stats": {
+ "downloadCount": 2183,
+ "ratingCount": 0,
+ "rating": 0,
+ "thumbsUpCount": 416,
+ "thumbsDownCount": 0
+ },
+ "files": [
+ {
+ "id": 1289799,
+ "sizeKB": 18829.1484375,
+ "name": "pp-storybook_rank2_bf16.safetensors",
+ "type": "Model",
+ "pickleScanResult": "Success",
+ "pickleScanMessage": "No Pickle imports",
+ "virusScanResult": "Success",
+ "virusScanMessage": null,
+ "scannedAt": "2025-02-08T11:21:04.247Z",
+ "metadata": {
+ "format": "SafeTensor"
+ },
+ "hashes": {
+ "AutoV1": "F414C813",
+ "AutoV2": "9753338AB6",
+ "SHA256": "9753338AB693CA82BF89ED77A5D1912879E40051463EC6E330FB9866CE798668",
+ "CRC32": "A65AE7B3",
+ "BLAKE3": "A5F8AB95AC2486345E4ACCAE541FF19D97ED53EFB0A7CC9226636975A0437591",
+ "AutoV3": "34A22376739D"
+ },
+ "downloadUrl": "https://civitai.com/api/download/models/1387174",
+ "primary": true
+ }
+ ],
+ "images": [
+ {
+ "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/42b875cf-c62b-41fa-a349-383b7f074351/original=true/56547310.jpeg",
+ "nsfwLevel": 1,
+ "width": 832,
+ "height": 1216,
+ "hash": "U5IiO6s-4Vn+0~EO^5xa00VsL#IU_O?E7yWC",
+ "type": "image",
+ "minor": false,
+ "poi": false,
+ "hasMeta": true,
+ "hasPositivePrompt": true,
+ "onSite": false,
+ "remixOfId": null
+ }
+ ],
+ "downloadUrl": "https://civitai.com/api/download/models/1387174"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/tests/routes/test_base_model_routes_smoke.py b/tests/routes/test_base_model_routes_smoke.py
index 136bd0a8..7f7e0c9a 100644
--- a/tests/routes/test_base_model_routes_smoke.py
+++ b/tests/routes/test_base_model_routes_smoke.py
@@ -31,7 +31,7 @@ from py_local.routes.base_model_routes import BaseModelRoutes
from py_local.services.model_file_service import AutoOrganizeResult
from py_local.services.service_registry import ServiceRegistry
from py_local.services.websocket_manager import ws_manager
-from py_local.utils.routes_common import ExifUtils
+from py_local.utils.exif_utils import ExifUtils
from py_local.config import config