From bdc8dec8603cecd1f7c08308af476063be790389 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Thu, 16 Apr 2026 08:54:12 +0800 Subject: [PATCH] fix(civitai): support civitai.red URLs (#897) --- locales/en.json | 8 +-- locales/zh-CN.json | 8 +-- py/routes/handlers/recipe_handlers.py | 10 +-- py/utils/civitai_utils.py | 68 ++++++++++++++++++- .../ContextMenu/ModelContextMenuMixin.js | 21 +----- static/js/managers/DownloadManager.js | 14 ++-- static/js/utils/civitaiUtils.js | 52 ++++++++++++++ .../components/modelContextMenuMixin.test.js | 21 ++++++ .../managers/downloadManager.history.test.js | 10 +++ tests/frontend/utils/civitaiUtils.test.js | 42 ++++++++++++ tests/routes/test_recipe_routes.py | 33 +++++++++ tests/services/test_batch_import_service.py | 1 + tests/utils/test_civitai_utils.py | 46 ++++++++++++- 13 files changed, 294 insertions(+), 40 deletions(-) create mode 100644 tests/frontend/components/modelContextMenuMixin.test.js diff --git a/locales/en.json b/locales/en.json index 696ccc08..4acaff89 100644 --- a/locales/en.json +++ b/locales/en.json @@ -685,9 +685,9 @@ "title": "Import a recipe from image or URL", "urlLocalPath": "URL / Local Path", "uploadImage": "Upload Image", - "urlSectionDescription": "Input a Civitai image URL or local file path to import as a recipe.", + "urlSectionDescription": "Input a Civitai image URL from civitai.com or civitai.red, or a local file path, to import as a recipe.", "imageUrlOrPath": "Image URL or File Path:", - "urlPlaceholder": "https://civitai.com/images/... or C:/path/to/image.png", + "urlPlaceholder": "https://civitai.com/images/... or https://civitai.red/images/... or C:/path/to/image.png", "fetchImage": "Fetch Image", "uploadSectionDescription": "Upload an image with LoRA metadata to import as a recipe.", "selectImage": "Select Image", @@ -1090,9 +1090,9 @@ }, "proceedText": "Only proceed if you're sure this is what you want.", "urlLabel": "Civitai Model URL:", - "urlPlaceholder": "https://civitai.com/models/649516/model-name?modelVersionId=726676", + "urlPlaceholder": "https://civitai.com/models/649516/model-name?modelVersionId=726676 or https://civitai.red/models/649516/model-name?modelVersionId=726676", "helpText": { - "title": "Paste any Civitai model URL. Supported formats:", + "title": "Paste any Civitai model URL from civitai.com or civitai.red. Supported formats:", "format1": "https://civitai.com/models/649516", "format2": "https://civitai.com/models/649516?modelVersionId=726676", "format3": "https://civitai.com/models/649516/model-name?modelVersionId=726676", diff --git a/locales/zh-CN.json b/locales/zh-CN.json index 62ed3c3b..93c4f630 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -685,9 +685,9 @@ "title": "从图片或 URL 导入配方", "urlLocalPath": "URL / 本地路径", "uploadImage": "上传图片", - "urlSectionDescription": "输入 Civitai 图片 URL 或本地文件路径以导入为配方。", + "urlSectionDescription": "输入来自 civitai.com 或 civitai.red 的 Civitai 图片 URL，或本地文件路径以导入为配方。", "imageUrlOrPath": "图片 URL 或文件路径：", - "urlPlaceholder": "https://civitai.com/images/... 或 C:/path/to/image.png", + "urlPlaceholder": "https://civitai.com/images/... 或 https://civitai.red/images/... 或 C:/path/to/image.png", "fetchImage": "获取图片", "uploadSectionDescription": "上传带有 LoRA 元数据的图片以导入为配方。", "selectImage": "选择图片", @@ -1090,9 +1090,9 @@ }, "proceedText": "仅在你确定需要此操作时继续。", "urlLabel": "Civitai 模型 URL：", - "urlPlaceholder": "https://civitai.com/models/649516/model-name?modelVersionId=726676", + "urlPlaceholder": "https://civitai.com/models/649516/model-name?modelVersionId=726676 或 https://civitai.red/models/649516/model-name?modelVersionId=726676", "helpText": { - "title": "粘贴任意 Civitai 模型 URL。支持格式：", + "title": "粘贴任意来自 civitai.com 或 civitai.red 的 Civitai 模型 URL。支持格式：", "format1": "https://civitai.com/models/649516", "format2": "https://civitai.com/models/649516?modelVersionId=726676", "format3": "https://civitai.com/models/649516/model-name?modelVersionId=726676", diff --git a/py/routes/handlers/recipe_handlers.py b/py/routes/handlers/recipe_handlers.py index 6120cb17..28025fa6 100644 --- a/py/routes/handlers/recipe_handlers.py +++ b/py/routes/handlers/recipe_handlers.py @@ -26,7 +26,7 @@ from ...services.recipes import ( RecipeValidationError, ) from ...services.metadata_service import get_default_metadata_provider -from ...utils.civitai_utils import rewrite_preview_url +from ...utils.civitai_utils import extract_civitai_image_id, rewrite_preview_url from ...utils.exif_utils import ExifUtils from ...recipes.merger import GenParamsMerger from ...recipes.enrichment import RecipeEnricher @@ -1196,13 +1196,13 @@ class RecipeManagementHandler: temp_path = temp_file.name download_url = image_url image_info = None - civitai_match = re.match(r"https://civitai\.com/images/(\d+)", image_url) - if civitai_match: + civitai_image_id = extract_civitai_image_id(image_url) + if civitai_image_id: if civitai_client is None: raise RecipeDownloadError( "Civitai client unavailable for image download" ) - image_info = await civitai_client.get_image_info(civitai_match.group(1)) + image_info = await civitai_client.get_image_info(civitai_image_id) if not image_info: raise RecipeDownloadError( "Failed to fetch image information from Civitai" @@ -1236,7 +1236,7 @@ class RecipeManagementHandler: return ( file_obj.read(), extension, - image_info.get("meta") if civitai_match and image_info else None, + image_info.get("meta") if civitai_image_id and image_info else None, ) except RecipeDownloadError: raise diff --git a/py/utils/civitai_utils.py b/py/utils/civitai_utils.py index 76426e55..0a24f84d 100644 --- a/py/utils/civitai_utils.py +++ b/py/utils/civitai_utils.py @@ -2,10 +2,12 @@ from __future__ import annotations +import re from typing import Any, Dict, Iterable, Mapping, Sequence -from urllib.parse import urlparse, urlunparse +from urllib.parse import parse_qs, urlparse, urlunparse +_SUPPORTED_CIVITAI_PAGE_HOSTS = frozenset({"civitai.com", "civitai.red"}) _DEFAULT_ALLOW_COMMERCIAL_USE: Sequence[str] = ("Sell",) _LICENSE_DEFAULTS: Dict[str, Any] = { "allowNoCredit": True, @@ -17,6 +19,67 @@ _COMMERCIAL_ALLOWED_VALUES = {"sell", "rent", "rentcivit", "image"} _COMMERCIAL_SHIFT = 1 +def is_supported_civitai_page_host(hostname: str | None) -> bool: + """Return whether the hostname is a supported Civitai page domain.""" + + if not hostname: + return False + return hostname.lower() in _SUPPORTED_CIVITAI_PAGE_HOSTS + + +def _parse_supported_civitai_page_url(url: str | None): + if not url: + return None + + try: + parsed = urlparse(url) + except ValueError: + return None + + if parsed.scheme not in {"http", "https"}: + return None + + if not is_supported_civitai_page_host(parsed.hostname): + return None + + return parsed + + +def extract_civitai_model_url_parts( + url: str | None, +) -> tuple[str | None, str | None]: + """Extract model and version identifiers from a supported Civitai model URL.""" + + parsed = _parse_supported_civitai_page_url(url) + if parsed is None: + return None, None + + path_match = re.search(r"/models/(\d+)", parsed.path) + if not path_match: + return None, None + + model_id = path_match.group(1) + + query_params = parse_qs(parsed.query) + version_values = query_params.get("modelVersionId") or [] + version_id = version_values[0] if version_values else None + return model_id, version_id + + +def extract_civitai_image_id(url: str | None) -> str | None: + """Extract the image identifier from a supported Civitai image page URL.""" + + parsed = _parse_supported_civitai_page_url(url) + if parsed is None: + return None + + path_match = re.search(r"/images/(\d+)", parsed.path) + if not path_match: + return None + + return path_match.group(1) + + def _normalize_commercial_values(value: Any) -> Sequence[str]: """Return a normalized list of commercial permissions preserving source values.""" @@ -199,6 +262,9 @@ def rewrite_preview_url( __all__ = [ "build_license_flags", + "extract_civitai_image_id", + "extract_civitai_model_url_parts", + "is_supported_civitai_page_host", "resolve_license_payload", "resolve_license_info", "rewrite_preview_url", diff --git a/static/js/components/ContextMenu/ModelContextMenuMixin.js b/static/js/components/ContextMenu/ModelContextMenuMixin.js index a8e6057f..66e810b9 100644 --- a/static/js/components/ContextMenu/ModelContextMenuMixin.js +++ b/static/js/components/ContextMenu/ModelContextMenuMixin.js @@ -6,6 +6,7 @@ import { bulkManager } from '../../managers/BulkManager.js'; import { MODEL_CONFIG } from '../../api/apiConfig.js'; import { translate } from '../../utils/i18nHelpers.js'; import { getNsfwLevelSelector } from '../shared/NsfwLevelSelector.js'; +import { extractCivitaiModelUrlParts } from '../../utils/civitaiUtils.js'; // Mixin with shared functionality for LoraContextMenu and CheckpointContextMenu export const ModelContextMenuMixin = { @@ -154,25 +155,7 @@ export const ModelContextMenuMixin = { }, 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 { modelId, modelVersionId }; - } catch (e) { - return { modelId: null, modelVersionId: null }; - } + return extractCivitaiModelUrlParts(url); }, parseModelId(value) { diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js index 2d52c95d..52ed7085 100644 --- a/static/js/managers/DownloadManager.js +++ b/static/js/managers/DownloadManager.js @@ -6,6 +6,7 @@ import { getModelApiClient, resetAndReload } from '../api/modelApiFactory.js'; import { getStorageItem, setStorageItem } from '../utils/storageHelpers.js'; import { FolderTreeManager } from '../components/FolderTreeManager.js'; import { translate } from '../utils/i18nHelpers.js'; +import { extractCivitaiModelUrlParts } from '../utils/civitaiUtils.js'; export class DownloadManager { constructor() { @@ -197,21 +198,22 @@ export class DownloadManager { } extractModelId(url) { - const versionMatch = url.match(/modelVersionId=(\d+)/i); - this.modelVersionId = versionMatch ? versionMatch[1] : null; - const civarchiveMatch = url.match(/https?:\/\/(?:www\.)?(?:civitaiarchive|civarchive)\.com\/models\/(\d+)/i); if (civarchiveMatch) { + const versionMatch = url.match(/modelVersionId=(\d+)/i); + this.modelVersionId = versionMatch ? versionMatch[1] : null; this.source = 'civarchive'; return civarchiveMatch[1]; } - const civitaiMatch = url.match(/https?:\/\/(?:www\.)?civitai\.com\/models\/(\d+)/i); - if (civitaiMatch) { + const { modelId, modelVersionId } = extractCivitaiModelUrlParts(url); + if (modelId) { + this.modelVersionId = modelVersionId; this.source = null; - return civitaiMatch[1]; + return modelId; } + this.modelVersionId = null; this.source = null; return null; } diff --git a/static/js/utils/civitaiUtils.js b/static/js/utils/civitaiUtils.js index fb3422af..4a4f2ee7 100644 --- a/static/js/utils/civitaiUtils.js +++ b/static/js/utils/civitaiUtils.js @@ -13,6 +13,11 @@ export const OptimizationMode = { THUMBNAIL: 'thumbnail', }; +const SUPPORTED_CIVITAI_PAGE_HOSTS = new Set([ + 'civitai.com', + 'civitai.red', +]); + /** * Rewrite Civitai preview URLs to use optimized renditions. * Mirrors the backend's rewrite_preview_url() function from py/utils/civitai_utils.py @@ -119,3 +124,50 @@ export function isCivitaiUrl(url) { return false; } } + +export function isSupportedCivitaiPageHost(hostname) { + if (!hostname) { + return false; + } + + return SUPPORTED_CIVITAI_PAGE_HOSTS.has(hostname.toLowerCase()); +} + +export function extractCivitaiModelUrlParts(url) { + if (!url) { + return { modelId: null, modelVersionId: null }; + } + + try { + const parsedUrl = new URL(url); + if (!isSupportedCivitaiPageHost(parsedUrl.hostname)) { + return { modelId: null, modelVersionId: null }; + } + + const pathMatch = parsedUrl.pathname.match(/\/models\/(\d+)/); + const modelId = pathMatch ? pathMatch[1] : null; + const modelVersionId = parsedUrl.searchParams.get('modelVersionId'); + + return { modelId, modelVersionId }; + } catch (e) { + return { modelId: null, modelVersionId: null }; + } +} + +export function extractCivitaiImageId(url) { + if (!url) { + return null; + } + + try { + const parsedUrl = new URL(url); + if (!isSupportedCivitaiPageHost(parsedUrl.hostname)) { + return null; + } + + const pathMatch = parsedUrl.pathname.match(/\/images\/(\d+)/); + return pathMatch ? pathMatch[1] : null; + } catch (e) { + return null; + } +} diff --git a/tests/frontend/components/modelContextMenuMixin.test.js b/tests/frontend/components/modelContextMenuMixin.test.js new file mode 100644 index 00000000..00736ed1 --- /dev/null +++ b/tests/frontend/components/modelContextMenuMixin.test.js @@ -0,0 +1,21 @@ +import { describe, expect, it } from 'vitest'; + +import { ModelContextMenuMixin } from '../../../static/js/components/ContextMenu/ModelContextMenuMixin.js'; + +describe('ModelContextMenuMixin.extractModelVersionId', () => { + it('accepts civitai.red model URLs', () => { + expect( + ModelContextMenuMixin.extractModelVersionId( + 'https://civitai.red/models/65423/nijimecha-artstyle?modelVersionId=777' + ) + ).toEqual({ modelId: '65423', modelVersionId: '777' }); + }); + + it('rejects model-like URLs from unsupported hosts', () => { + expect( + ModelContextMenuMixin.extractModelVersionId( + 'https://example.com/models/65423?modelVersionId=777' + ) + ).toEqual({ modelId: null, modelVersionId: null }); + }); +}); diff --git a/tests/frontend/managers/downloadManager.history.test.js b/tests/frontend/managers/downloadManager.history.test.js index f34e03ae..7abb7071 100644 --- a/tests/frontend/managers/downloadManager.history.test.js +++ b/tests/frontend/managers/downloadManager.history.test.js @@ -136,4 +136,14 @@ describe('DownloadManager version history badges', () => { expect(items[1].querySelector('.local-path')?.textContent).toContain('/models/still-local.safetensors'); expect(items[1].querySelector('.downloaded-badge')).toBeNull(); }); + + it('extracts model and version ids from civitai.red URLs', () => { + const manager = new DownloadManager(); + + expect( + manager.extractModelId('https://civitai.red/models/65423/nijimecha-artstyle?modelVersionId=777') + ).toBe('65423'); + expect(manager.modelVersionId).toBe('777'); + expect(manager.source).toBeNull(); + }); }); diff --git a/tests/frontend/utils/civitaiUtils.test.js b/tests/frontend/utils/civitaiUtils.test.js index d14637b4..b7cef6d9 100644 --- a/tests/frontend/utils/civitaiUtils.test.js +++ b/tests/frontend/utils/civitaiUtils.test.js @@ -4,7 +4,10 @@ import { getOptimizedUrl, getShowcaseUrl, getThumbnailUrl, + extractCivitaiImageId, + extractCivitaiModelUrlParts, isCivitaiUrl, + isSupportedCivitaiPageHost, OptimizationMode } from '../../../static/js/utils/civitaiUtils.js'; @@ -217,4 +220,43 @@ describe('civitaiUtils', () => { expect(isCivitaiUrl('not-a-url')).toBe(false); }); }); + + describe('isSupportedCivitaiPageHost', () => { + it('accepts civitai.com and civitai.red page hosts', () => { + expect(isSupportedCivitaiPageHost('civitai.com')).toBe(true); + expect(isSupportedCivitaiPageHost('civitai.red')).toBe(true); + }); + + it('rejects unrelated hosts', () => { + expect(isSupportedCivitaiPageHost('www.civitai.com')).toBe(false); + expect(isSupportedCivitaiPageHost('www.civitai.red')).toBe(false); + expect(isSupportedCivitaiPageHost('example.com')).toBe(false); + expect(isSupportedCivitaiPageHost('')).toBe(false); + expect(isSupportedCivitaiPageHost(null)).toBe(false); + }); + }); + + describe('extractCivitaiModelUrlParts', () => { + it('extracts model and version ids from civitai.red model URLs', () => { + expect( + extractCivitaiModelUrlParts('https://civitai.red/models/65423/name?modelVersionId=98765') + ).toEqual({ modelId: '65423', modelVersionId: '98765' }); + }); + + it('rejects model-like URLs from unsupported hosts', () => { + expect( + extractCivitaiModelUrlParts('https://example.com/models/65423?modelVersionId=98765') + ).toEqual({ modelId: null, modelVersionId: null }); + }); + }); + + describe('extractCivitaiImageId', () => { + it('extracts image ids from civitai.red image URLs', () => { + expect(extractCivitaiImageId('https://civitai.red/images/126920345')).toBe('126920345'); + }); + + it('rejects image-like URLs from unsupported hosts', () => { + expect(extractCivitaiImageId('https://example.com/images/126920345')).toBe(null); + }); + }); }); diff --git a/tests/routes/test_recipe_routes.py b/tests/routes/test_recipe_routes.py index 1fd09ae1..94cd0384 100644 --- a/tests/routes/test_recipe_routes.py +++ b/tests/routes/test_recipe_routes.py @@ -635,6 +635,39 @@ async def test_import_remote_video_recipe(monkeypatch, tmp_path: Path) -> None: assert call["extension"] == ".mp4" +async def test_import_remote_recipe_supports_civitai_red(monkeypatch, tmp_path: Path) -> None: + async def fake_get_default_metadata_provider(): + return SimpleNamespace(get_model_version_info=lambda id: ({}, None)) + + monkeypatch.setattr( + "py.recipes.enrichment.get_default_metadata_provider", + fake_get_default_metadata_provider, + ) + + async with recipe_harness(monkeypatch, tmp_path) as harness: + harness.civitai.image_info["126920345"] = { + "id": 126920345, + "url": "https://image.civitai.com/x/y/original=true/sample.jpeg", + "type": "image", + } + + response = await harness.client.get( + "/api/lm/recipes/import-remote", + params={ + "image_url": "https://civitai.red/images/126920345", + "name": "Red Recipe", + "resources": json.dumps([]), + "base_model": "Flux", + }, + ) + + payload = await response.json() + assert response.status == 200 + assert payload["success"] is True + assert harness.downloader.urls + assert "width=450,optimized=true" in harness.downloader.urls[0] + + async def test_analyze_uploaded_image_error_path(monkeypatch, tmp_path: Path) -> None: async with recipe_harness(monkeypatch, tmp_path) as harness: harness.analysis.raise_for_uploaded = RecipeValidationError( diff --git a/tests/services/test_batch_import_service.py b/tests/services/test_batch_import_service.py index 003a0c12..49b9193e 100644 --- a/tests/services/test_batch_import_service.py +++ b/tests/services/test_batch_import_service.py @@ -581,6 +581,7 @@ class TestInputValidation: assert service._validate_url("https://example.com/image.png") is True assert service._validate_url("http://example.com/image.png") is True assert service._validate_url("https://civitai.com/images/123") is True + assert service._validate_url("https://civitai.red/images/123") is True def test_validate_invalid_url(self, service): assert service._validate_url("not-a-url") is False diff --git a/tests/utils/test_civitai_utils.py b/tests/utils/test_civitai_utils.py index 840b17d3..7d2836fd 100644 --- a/tests/utils/test_civitai_utils.py +++ b/tests/utils/test_civitai_utils.py @@ -1,4 +1,11 @@ -from py.utils.civitai_utils import build_license_flags, resolve_license_info, resolve_license_payload +from py.utils.civitai_utils import ( + build_license_flags, + extract_civitai_image_id, + extract_civitai_model_url_parts, + is_supported_civitai_page_host, + resolve_license_info, + resolve_license_payload, +) def test_resolve_license_payload_defaults(): @@ -78,3 +85,40 @@ def test_build_license_flags_parses_aggregate_inside_list(): flags = build_license_flags(source) expected_flags = (1 << 0) | (7 << 1) | (1 << 5) assert flags == expected_flags + + +def test_supported_civitai_page_hosts_include_red(): + assert is_supported_civitai_page_host("civitai.com") is True + assert is_supported_civitai_page_host("civitai.red") is True + assert is_supported_civitai_page_host("www.civitai.com") is False + assert is_supported_civitai_page_host("www.civitai.red") is False + assert is_supported_civitai_page_host("example.com") is False + + +def test_extract_civitai_model_url_parts_supports_red(): + model_id, version_id = extract_civitai_model_url_parts( + "https://civitai.red/models/65423/nijimecha-artstyle?modelVersionId=777" + ) + + assert model_id == "65423" + assert version_id == "777" + + +def test_extract_civitai_model_url_parts_rejects_non_civitai_host(): + model_id, version_id = extract_civitai_model_url_parts( + "https://example.com/models/65423?modelVersionId=777" + ) + + assert model_id is None + assert version_id is None + + +def test_extract_civitai_image_id_supports_red(): + assert ( + extract_civitai_image_id("https://civitai.red/images/126920345") + == "126920345" + ) + + +def test_extract_civitai_image_id_rejects_non_civitai_host(): + assert extract_civitai_image_id("https://example.com/images/126920345") is None