fix(civitai): support civitai.red URLs (#897)

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

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",