fix(CivitAI): correct parse URL with query params

Author: bigcat88

comfy_cli/command/models/models.py

@@ -1,8 +1,9 @@
+import contextlib
 import os
 import pathlib
 import sys
 from typing import Annotated, Optional
-from urllib.parse import unquote, urlparse
+from urllib.parse import parse_qs, unquote, urlparse
 
 import requests
 import typer
@@ -74,39 +75,65 @@ def check_huggingface_url(url: str) -> tuple[bool, Optional[str], Optional[str],
     return True, repo_id, filename, folder_name, branch_name
 
 
-def check_civitai_url(url: str) -> tuple[bool, bool, int, int]:
+def check_civitai_url(url: str) -> tuple[bool, bool, Optional[int], Optional[int]]:
     """
     Returns:
-        is_civitai_model_url: True if the url is a civitai model url
-        is_civitai_api_url: True if the url is a civitai api url
-        model_id: The model id or None if it's api url
-        version_id: The version id or None if it doesn't have version id info
+        is_civitai_model_url: True if the url is a civitai *web* model url (e.g. /models/12345)
+        is_civitai_api_url: True if the url is a civitai *api* url useful for resolving downloads
+        model_id: The model id (for /models/*), else None
+        version_id: The version id (for /api/download/models/* or ?modelVersionId=), else None
     """
-    prefix = "civitai.com"
     try:
-        if prefix in url:
-            # URL is civitai api download url: https://civitai.com/api/download/models/12345
-            if "civitai.com/api/download" in url:
-                # This is a direct download link
-                version_id = url.strip("/").split("/")[-1]
-                return False, True, None, int(version_id)
-
-            # URL is civitai web url (e.g.
-            #   - https://civitai.com/models/43331
-            #   - https://civitai.com/models/43331/majicmix-realistic
-            subpath = url[url.find(prefix) + len(prefix) :].strip("/")
-            url_parts = subpath.split("?")
-            if len(url_parts) > 1:
-                model_id = url_parts[0].split("/")[1]
-                version_id = url_parts[1].split("=")[1]
-                return True, False, int(model_id), int(version_id)
-            else:
-                model_id = subpath.split("/")[1]
-                return True, False, int(model_id), None
-    except (ValueError, IndexError):
+        parsed = urlparse(url)
+        host = (parsed.hostname or "").lower()
+        if host != "civitai.com" and not host.endswith(".civitai.com"):
+            return False, False, None, None
+        p_parts = [p for p in parsed.path.split("/") if p]
+        query = parse_qs(parsed.query)
+
+        if len(p_parts) >= 4 and p_parts[0] == "api":
+            # Case 1: /api/download/models/<version_id>
+            # e.g. https://civitai.com/api/download/models/1617665?type=Model&format=SafeTensor
+            if p_parts[1] == "download" and p_parts[2] == "models":
+                try:
+                    version_id = int(p_parts[3])
+                    return False, True, None, version_id
+                except ValueError:
+                    return False, True, None, None
+
+            # Case 2: /api/v1/model-versions/<version_id>
+            if p_parts[1] == "v1" and p_parts[2] in ("model-versions", "modelVersions"):
+                try:
+                    version_id = int(p_parts[3])
+                    return False, True, None, version_id
+                except ValueError:
+                    return False, True, None, None
+
+        # Case 3: /models/<model_id>[/*] with optional ?modelVersionId=<id>
+        # e.g. https://civitai.com/models/43331
+        #      https://civitai.com/models/43331/majicmix-realistic?modelVersionId=485088
+        if len(p_parts) >= 2 and p_parts[0] == "models":
+            try:
+                model_id = int(p_parts[1])
+            except ValueError:
+                return False, False, None, None
+            version_id = None
+            mv = query.get("modelVersionId")
+            if mv and len(mv) > 0:
+                with contextlib.suppress(ValueError):
+                    version_id = int(mv[0])
+            if version_id is None:
+                mv = query.get("version")
+                if mv and len(mv) > 0:
+                    with contextlib.suppress(ValueError):
+                        version_id = int(mv[0])
+            return True, False, model_id, version_id
+
+        return False, False, None, None
+
+    except Exception:
         print("Error parsing CivitAI model URL")
-
-    return False, False, None, None
+        return False, False, None, None
 
 
 def request_civitai_model_version_api(version_id: int, headers: Optional[dict] = None):

tests/comfy_cli/command/models/test_models.py

@@ -11,6 +11,11 @@ def test_valid_model_url_with_version():
     assert check_civitai_url(url) == (True, False, 43331, None)
 
 
+def test_valid_model_url_with_version_and_additional_segments():
+    url = "https://civitai.com/models/43331/majicmix-realistic/extra"
+    assert check_civitai_url(url) == (True, False, 43331, None)
+
+
 def test_valid_model_url_with_query():
     url = "https://civitai.com/models/43331?version=12345"
     assert check_civitai_url(url) == (True, False, 43331, 12345)
@@ -31,8 +36,73 @@ def test_malformed_url():
     assert check_civitai_url(url) == (False, False, None, None)
 
 
+def test_invalid_model_id_url():
+    url = "https://civitai.com/models/invalid_id"
+    assert check_civitai_url(url) == (False, False, None, None)
+
+
 def test_malformed_query_url():
     url = "https://civitai.com/models/43331?version="
+    assert check_civitai_url(url) == (True, False, 43331, None)
+
+
+def test_model_url_with_model_version_id_query():
+    url = "https://civitai.com/models/43331?modelVersionId=485088"
+    assert check_civitai_url(url) == (True, False, 43331, 485088)
+
+
+def test_model_url_with_model_version_id_invalid():
+    url = "https://civitai.com/models/43331?modelVersionId=abc"
+    assert check_civitai_url(url) == (True, False, 43331, None)
+
+
+def test_valid_api_v1_model_versions_url():
+    url = "https://civitai.com/api/v1/model-versions/1617665"
+    assert check_civitai_url(url) == (False, True, None, 1617665)
+
+
+def test_valid_api_v1_model_versions_camelcase_segment():
+    url = "https://civitai.com/api/v1/modelVersions/1617665"
+    assert check_civitai_url(url) == (False, True, None, 1617665)
+
+
+def test_valid_api_download_with_query_params():
+    url = "https://civitai.com/api/download/models/1617665?type=Model&format=SafeTensor"
+    assert check_civitai_url(url) == (False, True, None, 1617665)
+
+
+def test_api_download_trailing_slash_is_ok():
+    url = "https://civitai.com/api/download/models/1617665/"
+    assert check_civitai_url(url) == (False, True, None, 1617665)
+
+
+def test_api_download_non_numeric_id_models_version():
+    url = "https://civitai.com/api/v1/modelVersions/notanumber"
+    assert check_civitai_url(url) == (False, True, None, None)
+
+
+def test_api_download_non_numeric_id():
+    url = "https://civitai.com/api/download/models/notanumber"
+    assert check_civitai_url(url) == (False, True, None, None)
+
+
+def test_model_url_with_slug_and_query():
+    url = "https://civitai.com/models/43331/majicmix-realistic?modelVersionId=485088"
+    assert check_civitai_url(url) == (True, False, 43331, 485088)
+
+
+def test_www_subdomain_is_accepted():
+    url = "https://www.civitai.com/models/43331?version=12345"
+    assert check_civitai_url(url) == (True, False, 43331, 12345)
+
+
+def test_completly_mailformed_civitai_url():
+    url = "https://civitai.com/"
+    assert check_civitai_url(url) == (False, False, None, None)
+
+
+def test_non_evil_civitai_url():
+    url = "https://evilcivitai.com/models/43331?version=12345"
     assert check_civitai_url(url) == (False, False, None, None)