Create civitai.py

backend/civitai.py

@@ -0,0 +1,29 @@
+import requests
+import os
+import hashlib
+import shutil
+
+def download_civitai_model(model_version_id: str, api_key: str) -> str:
+    api_url = f"https://civitai.com/api/v1/model-versions/{model_version_id}"
+    headers = {"Authorization": f"Bearer {api_key}"}
+    resp = requests.get(api_url, headers=headers)
+    resp.raise_for_status()
+    data = resp.json()
+    download_url = data["downloadUrl"]
+    cache_dir = "/data/civitai_cache"
+    os.makedirs(cache_dir, exist_ok=True)
+    filename = os.path.basename(download_url)
+    local_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(local_path):
+        with requests.get(download_url, headers=headers, stream=True) as r:
+            r.raise_for_status()
+            with open(local_path, "wb") as f:
+                shutil.copyfileobj(r.raw, f)
+    # For safetensors or zip, extract if zip?
+    if local_path.endswith(".zip"):
+        import zipfile
+        extract_dir = local_path + "_extracted"
+        with zipfile.ZipFile(local_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+        return extract_dir
+    return local_path