Support URL CB asset uploads

app.py

@@ -20,7 +20,7 @@ from datetime import datetime
 from functools import lru_cache
 from io import BytesIO
 from pathlib import Path
-from urllib.parse import quote
+from urllib.parse import quote, unquote, urlparse
 from zoneinfo import ZoneInfo
 
 ASYNCIO_FD_ERROR = "Invalid file descriptor: -1"
@@ -197,6 +197,15 @@ MODEL_LOCATION_CHOICES = [
     (kind.replace("_", " ").title(), kind)
     for kind in MODEL_KINDS
 ]
+DOWNLOAD_HEADERS = {
+    "User-Agent": (
+        "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:140.0) "
+        "Gecko/20100101 Firefox/140.0"
+    ),
+    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
+    "Accept-Language": "en-US,en;q=0.5",
+    "Referer": "https://civitai.com/",
+}
 DEFAULT_NEGATIVE = (
     "(censored, mosaic censoring, bar censor:1.1), bad quality, worst quality, "
     "worst detail, bad anatomy, extra fingers, extra toes, extra legs, 4 toes, "
@@ -502,54 +511,90 @@ def bucket_numbers(kind, anima):
     return numbers
 
 
-def upload_bucket_assets(files, kind, anima, password):
+def bucket_url_filename(response):
+    name = response.headers.get_filename()
+    if not name:
+        name = Path(urlparse(response.geturl()).path).name
+    name = unquote(name).replace("\\", "/").rsplit("/", 1)[-1].strip()
+    match = MODEL_NUMBER.match(name)
+    return match.group(2) if match else name
+
+
+def upload_bucket_assets(files, url, kind, anima, password):
     if not valid_pass(password):
         raise gr.Error("Invalid password")
     if kind not in MODEL_KINDS:
         raise gr.Error("Invalid CB location")
-    if not files:
-        raise gr.Error("Select at least one file")
+    files = files or []
+    url = (url or "").strip()
+    if not files and not url:
+        raise gr.Error("Select a file or enter a URL")
 
-    assets = []
-    for file in files:
-        source = Path(file)
-        match = MODEL_NUMBER.match(source.name)
-        number, name = (int(match.group(1)), match.group(2)) \
-            if match else (None, source.name)
-        if source.suffix.casefold() not in MODEL_SUFFIXES:
-            raise gr.Error(f"Unsupported model file: {name}")
-        if anima and not is_anima_asset(name):
-            name = f"{ANIMA_PREFIX}_{name}"
-        assets.append((source, number, name))
-
-    with model_upload_lock:
-        used = bucket_numbers(kind, anima)
-        reserved = {
-            number for _, number, _ in assets
-            if number is not None and number not in used
-        }
-        next_number = max(used, default=0) + 1
-        additions = []
-        paths = []
-        for source, number, name in assets:
-            if number is not None and number in reserved:
-                reserved.remove(number)
-            else:
-                while next_number in used or next_number in reserved:
+    temp = None
+    try:
+        assets = []
+        for file in files:
+            source = Path(file)
+            match = MODEL_NUMBER.match(source.name)
+            number, name = (int(match.group(1)), match.group(2)) \
+                if match else (None, source.name)
+            if source.suffix.casefold() not in MODEL_SUFFIXES:
+                raise gr.Error(f"Unsupported model file: {name}")
+            if anima and not is_anima_asset(name):
+                name = f"{ANIMA_PREFIX}_{name}"
+            assets.append((source, number, name))
+
+        if url:
+            parsed = urlparse(url)
+            if parsed.scheme not in ("http", "https") or not parsed.netloc:
+                raise gr.Error("Enter a valid URL")
+            request = urllib.request.Request(url, headers=DOWNLOAD_HEADERS)
+            with urllib.request.urlopen(request) as response:
+                name = bucket_url_filename(response)
+                suffix = Path(name).suffix.casefold()
+                if suffix not in MODEL_SUFFIXES:
+                    raise gr.Error("URL did not return a model file")
+                if anima and not is_anima_asset(name):
+                    name = f"{ANIMA_PREFIX}_{name}"
+                with tempfile.NamedTemporaryFile(
+                    suffix=suffix,
+                    delete=False,
+                ) as file:
+                    temp = Path(file.name)
+                    shutil.copyfileobj(response, file)
+            assets.append((temp, None, name))
+
+        with model_upload_lock:
+            used = bucket_numbers(kind, anima)
+            reserved = {
+                number for _, number, _ in assets
+                if number is not None and number not in used
+            }
+            next_number = max(used, default=0) + 1
+            additions = []
+            paths = []
+            for source, number, name in assets:
+                if number is not None and number in reserved:
+                    reserved.remove(number)
+                else:
+                    while next_number in used or next_number in reserved:
+                        next_number += 1
+                    number = next_number
                     next_number += 1
-                number = next_number
-                next_number += 1
-            used.add(number)
-            path = f"{kind}/{number}_{name}"
-            additions.append((source, path))
-            paths.append(path)
-        batch_bucket_files(
-            BUCKET_ID,
-            add=additions,
-            token=IMAGE_TOKEN,
-        )
-        index_bucket_models()
-    return "Added " + ", ".join(paths)
+                used.add(number)
+                path = f"{kind}/{number}_{name}"
+                additions.append((source, path))
+                paths.append(path)
+            batch_bucket_files(
+                BUCKET_ID,
+                add=additions,
+                token=IMAGE_TOKEN,
+            )
+            index_bucket_models()
+        return "Added " + ", ".join(paths)
+    finally:
+        if temp:
+            temp.unlink(missing_ok=True)
 
 
 def model_kind(name):
@@ -3612,6 +3657,7 @@ with gr.Blocks(title="Image generation") as demo:
                         file_types=list(MODEL_SUFFIXES),
                         type="filepath",
                     )
+                    model_url_input = gr.Textbox(label="URL")
                     with gr.Row():
                         model_location_input = gr.Dropdown(
                             MODEL_LOCATION_CHOICES,
@@ -3622,7 +3668,7 @@ with gr.Blocks(title="Image generation") as demo:
                             False,
                             label="Anima model",
                         )
-                        model_upload_button = gr.Button("Upload")
+                        model_upload_button = gr.Button("Add")
                     model_upload_status = gr.Textbox(
                         label="Status",
                         interactive=False,
@@ -3990,6 +4036,7 @@ with gr.Blocks(title="Image generation") as demo:
         upload_bucket_assets,
         [
             model_files_input,
+            model_url_input,
             model_location_input,
             anima_model_input,
             pass_input,