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! 💕

If you like my LoRA, please like, comment, or donate some Buzz. Much appreciated! ❤️

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