From 26d105c439a8b062c37d6cbda66fc05944c3da6a Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 11 Jun 2025 22:06:16 +0800 Subject: [PATCH] Enhance Civitai model handling: add get_model_version method for detailed metadata retrieval, update routes to utilize new method, and improve URL handling in context menu for model re-linking. --- py/services/civitai_client.py | 63 +++++ py/utils/routes_common.py | 61 ++-- refs/output.json | 265 +++++++++++++++++- .../ContextMenu/ModelContextMenuMixin.js | 26 +- templates/components/modals.html | 8 +- 5 files changed, 386 insertions(+), 37 deletions(-) diff --git a/py/services/civitai_client.py b/py/services/civitai_client.py index 59abd7d7..bec4f5d4 100644 --- a/py/services/civitai_client.py +++ b/py/services/civitai_client.py @@ -224,6 +224,69 @@ class CivitaiClient: except Exception as e: logger.error(f"Error fetching model versions: {e}") return None + + async def get_model_version(self, model_id: str, version_id: str = "") -> Optional[Dict]: + """Get specific model version with additional metadata + + Args: + model_id: The Civitai model ID + version_id: Optional specific version ID to retrieve + + Returns: + Optional[Dict]: The model version data with additional fields or None if not found + """ + try: + session = await self._ensure_fresh_session() + async with session.get(f"{self.base_url}/models/{model_id}") as response: + if response.status != 200: + return None + + data = await response.json() + model_versions = data.get('modelVersions', []) + + # Find matching version + matched_version = None + + if version_id: + # If version_id provided, find exact match + for version in model_versions: + if str(version.get('id')) == str(version_id): + matched_version = version + break + else: + # If no version_id then use the first version + matched_version = model_versions[0] if model_versions else None + + # If no match found, return None + if not matched_version: + return None + + # Build result with modified fields + result = matched_version.copy() # Copy to avoid modifying original + + # Replace index with modelId + if 'index' in result: + del result['index'] + result['modelId'] = model_id + + # Add model field with metadata from top level + result['model'] = { + "name": data.get("name"), + "type": data.get("type"), + "nsfw": data.get("nsfw", False), + "poi": data.get("poi", False), + "description": data.get("description"), + "tags": data.get("tags", []) + } + + # Add creator field from top level + result['creator'] = data.get("creator") + + return result + + except Exception as e: + logger.error(f"Error fetching model version: {e}") + return None async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]: """Fetch model version metadata from Civitai diff --git a/py/utils/routes_common.py b/py/utils/routes_common.py index bcfdbec8..6cdf32e4 100644 --- a/py/utils/routes_common.py +++ b/py/utils/routes_common.py @@ -47,13 +47,30 @@ class ModelRouteUtils: if civitai_metadata.get('model', {}).get('name'): local_metadata['model_name'] = civitai_metadata['model']['name'] - # Fetch additional model metadata (description and tags) if we have model ID - model_id = civitai_metadata['modelId'] - if model_id: - model_metadata, _ = await client.get_model_metadata(str(model_id)) - if (model_metadata): - local_metadata['modelDescription'] = model_metadata.get('description', '') - local_metadata['tags'] = model_metadata.get('tags', []) + # Extract model metadata directly from civitai_metadata if available + model_metadata = None + + if 'model' in civitai_metadata and civitai_metadata.get('model'): + # Data is already available in the response from get_model_version + model_metadata = { + 'description': civitai_metadata.get('model', {}).get('description', ''), + 'tags': civitai_metadata.get('model', {}).get('tags', []), + 'creator': civitai_metadata.get('creator', {}) + } + + # If we have modelId and don't have enough metadata, fetch additional data + if not model_metadata or not model_metadata.get('description'): + model_id = civitai_metadata.get('modelId') + if model_id: + fetched_metadata, _ = await client.get_model_metadata(str(model_id)) + if fetched_metadata: + model_metadata = fetched_metadata + + # Update local metadata with the model information + if model_metadata: + local_metadata['modelDescription'] = model_metadata.get('description', '') + local_metadata['tags'] = model_metadata.get('tags', []) + if 'creator' in model_metadata and model_metadata['creator']: local_metadata['civitai']['creator'] = model_metadata['creator'] # Update base model @@ -607,7 +624,7 @@ class ModelRouteUtils: @staticmethod async def handle_relink_civitai(request: web.Request, scanner) -> web.Response: - """Handle CivitAI metadata re-linking request by model version ID + """Handle CivitAI metadata re-linking request by model ID and/or version ID Args: request: The aiohttp request @@ -619,10 +636,11 @@ class ModelRouteUtils: try: data = await request.json() file_path = data.get('file_path') + model_id = data.get('model_id') model_version_id = data.get('model_version_id') - if not file_path or not model_version_id: - return web.json_response({"success": False, "error": "Both file_path and model_version_id are required"}, status=400) + if not file_path or not model_id: + return web.json_response({"success": False, "error": "Both file_path and model_id are required"}, status=400) metadata_path = os.path.splitext(file_path)[0] + '.metadata.json' @@ -632,24 +650,24 @@ class ModelRouteUtils: # Create a client for fetching from Civitai client = await CivitaiClient.get_instance() try: - # Fetch metadata by model version ID - civitai_metadata, error = await client.get_model_version_info(model_version_id) + # Fetch metadata using get_model_version which includes more comprehensive data + civitai_metadata = await client.get_model_version(model_id, model_version_id) if not civitai_metadata: - error_msg = error or "Model version not found on CivitAI" + error_msg = f"Model version not found on CivitAI for ID: {model_id}" + if model_version_id: + error_msg += f" with version: {model_version_id}" return web.json_response({"success": False, "error": error_msg}, status=404) - # Find the primary model file to get the correct SHA256 hash + # Try to find the primary model file to get the SHA256 hash primary_model_file = None for file in civitai_metadata.get('files', []): if file.get('primary', False) and file.get('type') == 'Model': primary_model_file = file break - if not primary_model_file or not primary_model_file.get('hashes', {}).get('SHA256'): - return web.json_response({"success": False, "error": "No SHA256 hash found in model metadata"}, status=404) - - # Update the SHA256 hash in local metadata (convert to lowercase) - local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower() + # Update the SHA256 hash in local metadata if available + if primary_model_file and primary_model_file.get('hashes', {}).get('SHA256'): + local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower() # Update metadata with CivitAI information await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, client) @@ -659,8 +677,9 @@ class ModelRouteUtils: return web.json_response({ "success": True, - "message": f"Model successfully re-linked to Civitai version {model_version_id}", - "hash": local_metadata['sha256'] + "message": f"Model successfully re-linked to Civitai model {model_id}" + + (f" version {model_version_id}" if model_version_id else ""), + "hash": local_metadata.get('sha256', '') }) finally: diff --git a/refs/output.json b/refs/output.json index 775870c5..aa5a9702 100644 --- a/refs/output.json +++ b/refs/output.json @@ -1,11 +1,258 @@ { - "loras": " ", - "prompt": "in the style of ck-rw, aorun, scales, makeup, bare shoulders, pointy ears, dress, claws, in the style of cksc, artist:moriimee, in the style of cknc, masterpiece, best quality, good quality, very aesthetic, absurdres, newest, 8K, depth of field, focused subject, close up, stylized, in gold and neon shades, wabi sabi, 1girl, rainbow angel wings, looking at viewer, dynamic angle, from below, from side, relaxing", - "negative_prompt": "bad quality, worst quality, worst detail, sketch ,signature, watermark, patreon logo, nsfw", - "steps": "20", - "sampler": "euler_ancestral", - "cfg_scale": "8", - "seed": "241", - "size": "832x1216", - "clip_skip": "2" + "id": 649516, + "name": "Cynthia -シロナ - Pokemon Diamond and Pearl - PDXL LORA", + "description": "

Warning: Without Adetailer eyes are fucked (rainbow color and artefact)

Trained on Pony Diffusion V6 XL with 63 pictures.
Best result with weight between : 0.8-1.

Basic prompts : 1girl, cynthia \\(pokemon\\), blonde hair, hair over one eye, very long hair, grey eyes, eyelashes, hair ornament 
Outfit prompts : fur collar, black coat, fur-trimmed coat, long sleeves, black pants, black shirt, high heels

Reviews are really appreciated, i love to see the community use my work, that's why I share it.
If you like my work, you can tip me here.

Got a specific request ? I'm open for commission on my kofi or fiverr gig *! If you provide enough data, OCs are accepted

", + "allowNoCredit": true, + "allowCommercialUse": [ + "Image", + "RentCivit" + ], + "allowDerivatives": true, + "allowDifferentLicense": true, + "type": "LORA", + "minor": false, + "sfwOnly": false, + "poi": false, + "nsfw": false, + "nsfwLevel": 29, + "availability": "Public", + "cosmetic": null, + "supportsGeneration": true, + "stats": { + "downloadCount": 811, + "favoriteCount": 0, + "thumbsUpCount": 175, + "thumbsDownCount": 0, + "commentCount": 4, + "ratingCount": 0, + "rating": 0, + "tippedAmountCount": 10 + }, + "creator": { + "username": "Konan", + "image": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/7cd552a1-60fe-4baf-a0e4-f7d5d5381711/width=96/Konan.jpeg" + }, + "tags": [ + "anime", + "character", + "cynthia", + "woman", + "pokemon", + "pokegirl" + ], + "modelVersions": [ + { + "id": 726676, + "index": 0, + "name": "v1.0", + "baseModel": "Pony", + "createdAt": "2024-08-16T01:13:16.099Z", + "publishedAt": "2024-08-16T01:14:44.984Z", + "status": "Published", + "availability": "Public", + "nsfwLevel": 29, + "trainedWords": [ + "1girl, cynthia \\(pokemon\\), blonde hair, hair over one eye, very long hair, grey eyes, eyelashes, hair ornament", + "fur collar, black coat, fur-trimmed coat, long sleeves, black pants, black shirt, high heels" + ], + "covered": true, + "stats": { + "downloadCount": 811, + "ratingCount": 0, + "rating": 0, + "thumbsUpCount": 175, + "thumbsDownCount": 0 + }, + "files": [ + { + "id": 641092, + "sizeKB": 56079.65234375, + "name": "CynthiaXL.safetensors", + "type": "Model", + "pickleScanResult": "Success", + "pickleScanMessage": "No Pickle imports", + "virusScanResult": "Success", + "virusScanMessage": null, + "scannedAt": "2024-08-16T01:17:19.087Z", + "metadata": { + "format": "SafeTensor" + }, + "hashes": {}, + "downloadUrl": "https://civitai.com/api/download/models/726676", + "primary": true + } + ], + "images": [ + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/b346d757-2b59-4aeb-9f09-3bee2724519d/width=1248/24511993.jpeg", + "nsfwLevel": 1, + "width": 1248, + "height": 1824, + "hash": "UqNc==RP.9s+~pxvIst7kWWBWBjY%MWBt7WB", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/fc132ac0-cc1c-4b68-a1d7-5b97b0996ac2/width=1248/24511997.jpeg", + "nsfwLevel": 1, + "width": 1248, + "height": 1824, + "hash": "UMGSS+?tTw.60MIX9cbb~WxHRRR-NEtLRiR%", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/7b3237d1-e672-466a-85d0-cc5dd42ab130/width=1160/24512001.jpeg", + "nsfwLevel": 4, + "width": 1160, + "height": 1696, + "hash": "U9NA6f~o00%h00wvIYt74:ER-=D%5600DiE1", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/ccd7d11d-4fa9-4434-85a1-fb999312e60d/width=1248/24511991.jpeg", + "nsfwLevel": 1, + "width": 1248, + "height": 1824, + "hash": "UyNTg.j?~qxu?aoLRkj]%MfkM{jZaya}a#ax", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/1743be6d-7fe5-4b55-9f19-c931618fa259/width=1248/24511996.jpeg", + "nsfwLevel": 4, + "width": 1248, + "height": 1824, + "hash": "UGOC~n^+?w~6Tx_4oM^$yYEkMds74:9F#*xY", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/91693c98-d037-4489-882c-100eb26019a0/width=1160/24512010.jpeg", + "nsfwLevel": 4, + "width": 1160, + "height": 1696, + "hash": "UJI}kp^-Kl%hXAIX4;Nf^+M|9GRP0Mt8%L%2", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/49c7a294-ac5b-4832-98e5-2acd0f1a8782/width=1248/24512017.jpeg", + "nsfwLevel": 4, + "width": 1248, + "height": 1824, + "hash": "UML;8Qn|9G%3mnWA4nWFMf%N?Hae~qog-oNF", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/d7b442f2-6ead-4a7a-9578-54d9ec2ff148/width=1248/24512015.jpeg", + "nsfwLevel": 1, + "width": 1248, + "height": 1824, + "hash": "UPGR#kt8xw%M0LWC9bWC?wxtR*NLM^jrxWM|", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/d840f1e9-3dd3-4531-b83a-1ba2c6b7feaa/width=1160/24512004.jpeg", + "nsfwLevel": 8, + "width": 1160, + "height": 1696, + "hash": "ULNm1i_39wi^*I%hDiM_tlo#xuV?^kNIxCs,", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/520387ae-c176-43e3-92bd-5cd2a672475e/width=1248/24512012.jpeg", + "nsfwLevel": 4, + "width": 1248, + "height": 1824, + "hash": "URM%l.%M.9Ip~poIkExu_3V@M|xuD%oJM{D*", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/9ea28b94-f326-4776-83ff-851cc203c627/width=1248/24511988.jpeg", + "nsfwLevel": 1, + "width": 1248, + "height": 1824, + "hash": "U-PZloog_Nxut6j]WXWB-;j?IVa#ofaxj]j]", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + }, + { + "url": "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/2e749dbb-7d5a-48f1-8e29-fea5022a5fe9/width=1248/24522268.jpeg", + "nsfwLevel": 16, + "width": 1248, + "height": 1824, + "hash": "UPLgtm9Z0z=|0yRRE2-A9rWAoNE1~DwOr=t7", + "type": "image", + "minor": false, + "poi": false, + "hasMeta": true, + "hasPositivePrompt": true, + "onSite": false, + "remixOfId": null + } + ], + "downloadUrl": "https://civitai.com/api/download/models/726676" + } + ] } \ No newline at end of file diff --git a/static/js/components/ContextMenu/ModelContextMenuMixin.js b/static/js/components/ContextMenu/ModelContextMenuMixin.js index 1f404370..a5a0c69b 100644 --- a/static/js/components/ContextMenu/ModelContextMenuMixin.js +++ b/static/js/components/ContextMenu/ModelContextMenuMixin.js @@ -273,10 +273,10 @@ export const ModelContextMenuMixin = { // Create new bound handler this._boundRelinkHandler = async () => { const url = urlInput.value.trim(); - const modelVersionId = this.extractModelVersionId(url); + const { modelId, modelVersionId } = this.extractModelVersionId(url); - if (!modelVersionId) { - errorDiv.textContent = 'Invalid URL format. Must include modelVersionId parameter.'; + if (!modelId) { + errorDiv.textContent = 'Invalid URL format. Must include model ID.'; return; } @@ -297,6 +297,7 @@ export const ModelContextMenuMixin = { }, body: JSON.stringify({ file_path: filePath, + model_id: modelId, model_version_id: modelVersionId }) }); @@ -331,15 +332,30 @@ export const ModelContextMenuMixin = { // Show modal modalManager.showModal('relinkCivitaiModal'); + + // Auto-focus the URL input field after modal is shown + setTimeout(() => urlInput.focus(), 50); }, extractModelVersionId(url) { try { + // Handle all three URL formats: + // 1. https://civitai.com/models/649516 + // 2. https://civitai.com/models/649516?modelVersionId=726676 + // 3. https://civitai.com/models/649516/cynthia-pokemon-diamond-and-pearl-pdxl-lora?modelVersionId=726676 + const parsedUrl = new URL(url); + + // Extract model ID from path + const pathMatch = parsedUrl.pathname.match(/\/models\/(\d+)/); + const modelId = pathMatch ? pathMatch[1] : null; + + // Extract model version ID from query parameters const modelVersionId = parsedUrl.searchParams.get('modelVersionId'); - return modelVersionId; + + return { modelId, modelVersionId }; } catch (e) { - return null; + return { modelId: null, modelVersionId: null }; } }, diff --git a/templates/components/modals.html b/templates/components/modals.html index 3b35104f..989b9319 100644 --- a/templates/components/modals.html +++ b/templates/components/modals.html @@ -598,10 +598,14 @@
- +
- The URL must include the modelVersionId parameter. + Paste any Civitai model URL. Supported formats:
+ • https://civitai.com/models/649516
+ • https://civitai.com/models/649516?modelVersionId=726676
+ • https://civitai.com/models/649516/model-name?modelVersionId=726676
+ Note: If no modelVersionId is provided, the latest version will be used.