Add CB asset URL downloader

app.py

@@ -20,6 +20,7 @@ from datetime import datetime
 from functools import lru_cache
 from io import BytesIO
 from pathlib import Path
+from urllib.parse import unquote, urlparse
 from zoneinfo import ZoneInfo
 
 ASYNCIO_FD_ERROR = "Invalid file descriptor: -1"
@@ -193,6 +194,10 @@ MODEL_SUFFIXES = (
 NUMBERED_MODEL_KINDS = ("checkpoints", "loras")
 MODEL_NUMBER = re.compile(r"^(\d+)_(.+)$")
 ANIMA_PREFIX = "anima"
+MODEL_LOCATION_CHOICES = [
+    (kind.replace("_", " ").title(), kind)
+    for kind in MODEL_KINDS
+]
 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, "
@@ -284,6 +289,7 @@ pool = ThreadPoolExecutor(max_workers=1)
 backup_pool = ThreadPoolExecutor(max_workers=1)
 matrix_pool = ThreadPoolExecutor(max_workers=1)
 model_pool = ThreadPoolExecutor(max_workers=1)
+model_download_lock = threading.Lock()
 matrix_grids = set()
 local_client = None
 local_client_lock = threading.Lock()
@@ -398,6 +404,11 @@ def download_assets(models):
         )
 
 
+def is_anima_asset(name):
+    name = name.casefold()
+    return name.startswith(ANIMA_PREFIX) and not name.startswith("animag")
+
+
 def index_bucket_models():
     models = {kind: {} for kind in MODEL_KINDS}
     items = [
@@ -414,7 +425,7 @@ def index_bucket_models():
         kind, separator, name = item.path.partition("/")
         match = MODEL_NUMBER.match(name)
         if separator and kind in counters and match:
-            anima = match.group(2).casefold().startswith(ANIMA_PREFIX)
+            anima = is_anima_asset(match.group(2))
             counters[kind][anima] = max(counters[kind][anima], int(match.group(1)))
 
     copies = []
@@ -424,7 +435,7 @@ def index_bucket_models():
         if not separator or kind not in models:
             continue
         if kind in counters and not MODEL_NUMBER.match(name):
-            anima = name.casefold().startswith(ANIMA_PREFIX)
+            anima = is_anima_asset(name)
             counters[kind][anima] += 1
             name = f"{counters[kind][anima]}_{name}"
             path = f"{kind}/{name}"
@@ -451,6 +462,78 @@ def index_bucket_models():
     return added
 
 
+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 next_bucket_number(kind, anima):
+    numbers = []
+    for item in list_bucket_tree(
+        BUCKET_ID,
+        prefix=f"{kind}/",
+        recursive=True,
+        token=False,
+    ):
+        if item.type != "file" or "/" in item.path.removeprefix(f"{kind}/"):
+            continue
+        match = MODEL_NUMBER.match(Path(item.path).name)
+        if match and is_anima_asset(match.group(2)) == anima:
+            numbers.append(int(match.group(1)))
+    return max(numbers, default=0) + 1
+
+
+def download_bucket_asset(url, kind, anima, password):
+    if not valid_pass(password):
+        raise gr.Error("Invalid password")
+    url = url.strip()
+    parsed = urlparse(url)
+    if parsed.scheme not in ("http", "https") or not parsed.netloc:
+        raise gr.Error("Enter a valid URL")
+    if kind not in MODEL_KINDS:
+        raise gr.Error("Invalid CB location")
+    if anima and kind == "checkpoints":
+        kind = "diffusion_models"
+
+    temp = None
+    try:
+        request = urllib.request.Request(
+            url,
+            headers={"User-Agent": "Mozilla/5.0"},
+        )
+        with urllib.request.urlopen(request) as response:
+            name = bucket_url_filename(response)
+            if Path(name).suffix.casefold() 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=Path(name).suffix,
+                delete=False,
+            ) as file:
+                temp = Path(file.name)
+                shutil.copyfileobj(response, file)
+
+        with model_download_lock:
+            number = next_bucket_number(kind, anima)
+            name = f"{number}_{name}"
+            path = f"{kind}/{name}"
+            batch_bucket_files(
+                BUCKET_ID,
+                add=[(temp, path)],
+                token=IMAGE_TOKEN,
+            )
+            index_bucket_models()
+        return f"Added {path}"
+    finally:
+        if temp:
+            temp.unlink(missing_ok=True)
+
+
 def model_kind(name):
     return "diffusion_models" if is_anima_model(name) else "checkpoints"
 
@@ -3480,6 +3563,23 @@ with gr.Blocks(title="Image generation") as demo:
                     )
                     model_state = gr.State(DEFAULT_MODEL)
                     refresh_button = gr.Button("Refresh", scale=1)
+                with gr.Accordion("Add CB asset", open=False):
+                    model_url_input = gr.Textbox(label="URL")
+                    with gr.Row():
+                        model_location_input = gr.Dropdown(
+                            MODEL_LOCATION_CHOICES,
+                            value="checkpoints",
+                            label="CB location",
+                        )
+                        anima_model_input = gr.Checkbox(
+                            False,
+                            label="Anima model",
+                        )
+                        model_download_button = gr.Button("Download")
+                    model_download_status = gr.Textbox(
+                        label="Status",
+                        interactive=False,
+                    )
                 prompt_input = gr.Textbox(label="Prompt", lines=6)
                 negative_input = gr.Textbox(
                     DEFAULT_NEGATIVE,
@@ -3839,6 +3939,18 @@ with gr.Blocks(title="Image generation") as demo:
         ],
         queue=False,
     )
+    model_download_button.click(
+        download_bucket_asset,
+        [
+            model_url_input,
+            model_location_input,
+            anima_model_input,
+            pass_input,
+        ],
+        model_download_status,
+        queue=False,
+        api_visibility="private",
+    )
 
     api_button.click(
         api_generate,