refactor: enhance model version retrieval logic in CivitaiClient, fixes #460

This commit is contained in:
Will Miao
2025-09-24 09:16:02 +08:00
parent 6f4453aaf3
commit 12c88835f2
3 changed files with 176 additions and 20 deletions

View File

@@ -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

View File

@@ -0,0 +1,110 @@
{
"id": 1231067,
"name": "Vivid Impressions Storybook Style",
"description": "<h3 id=\"if-you'd-like-to-support-me-feel-free-to-visit-my-ko-fi-page.-please-share-your-images-using-the-&quot;+add-post&quot;-button-below.-it-supports-the-creators.-thanks!-nnfwkvfly\">If you'd like to support me, feel free to visit my <a target=\"_blank\" rel=\"ugc\" href=\"https://ko-fi.com/pixelpawsai\">Ko-Fi</a> page. ❤️<br /><br />Please share your images using the \"<span style=\"color:rgb(250, 82, 82)\">+add post</span>\" button below. It supports the creators. Thanks! 💕</h3><h3 id=\"if-you-like-my-lora-please-like-comment-or-donate-some-buzz.-much-appreciated!-vyeqok3go\">If you like my LoRA, please<span style=\"color:rgb(230, 73, 128)\"> </span><span style=\"color:rgb(250, 82, 82)\">like</span>, <span style=\"color:rgb(250, 82, 82)\">comment</span>, or <span style=\"color:#fa5252\">donate some Buzz</span>. Much appreciated! ❤️</h3><h3 id=\"-lo912t8rj\"></h3><h3 id=\"trigger-word:-ppstorybook-wlggllim2\"><strong><span style=\"color:rgb(253, 126, 20)\">Trigger word: </span></strong>ppstorybook</h3><h3 id=\"strength:-0.8-experiment-as-you-like-luvhks6za\"><strong><span style=\"color:rgb(253, 126, 20)\">Strength: </span></strong>0.8, experiment as you like</h3>",
"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"
}
]
}

View File

@@ -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