Upload app.py with huggingface_hub

app.py

@@ -13,6 +13,8 @@ Three tokens, three jobs:
 
 from __future__ import annotations
 
+import os
+
 import gradio as gr
 
 import caption
@@ -20,6 +22,7 @@ import jobs
 
 MAX_IMAGES = 40
 MAX_LOG = 60_000
+IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tiff", ".tif")
 
 LR_SCHEDULERS = ["constant", "cosine", "linear", "constant_with_warmup", "polynomial"]
 OPTIMIZERS = ["adamW", "prodigy"]
@@ -91,12 +94,34 @@ CAPTION_TIP = (
     "trigger token.\n"
     "- <b>✨ Auto-caption</b> fills these for you; edit anything you like. Blank captions fall "
     "back to the trigger.\n"
+    "- Already have captions? Upload a <code>.txt</code> next to each image with the same name "
+    "(e.g. <code>cat.jpg</code> + <code>cat.txt</code>) and they're filled in automatically.\n"
     "</details>"
 )
 
 
-def _paths(images) -> list[str]:
-    return [p if isinstance(p, str) else getattr(p, "name", p) for p in (images or [])]
+def _paths(files) -> list[str]:
+    return [p if isinstance(p, str) else getattr(p, "name", p) for p in (files or [])]
+
+
+def _split_uploads(files):
+    """Split a mixed upload into image paths and a {basename: caption} map from .txt sidecars.
+
+    A .txt file whose name matches an image (e.g. cat.jpg + cat.txt) prefills that caption.
+    """
+    images, captions = [], {}
+    for p in _paths(files):
+        ext = os.path.splitext(p)[1].lower()
+        base = os.path.splitext(os.path.basename(p))[0]
+        if ext == ".txt":
+            try:
+                with open(p, encoding="utf-8") as f:
+                    captions[base] = f.read().strip()
+            except Exception:  # noqa: BLE001
+                pass
+        elif ext in IMAGE_EXTS:
+            images.append(p)
+    return images, captions
 
 
 def _signin_state(profile: gr.OAuthProfile | None):
@@ -120,38 +145,46 @@ def on_concept_change(concept: str):
     )
 
 
-def load_captioning(images, trigger):
-    """Reveal the captioning + training UI once images are uploaded; one row per image."""
-    paths = _paths(images)
+def load_captioning(files, trigger):
+    """Reveal the captioning + training UI once images are uploaded; one row per image.
+
+    If a matching `<name>.txt` sidecar was uploaded, its content prefills that image's caption.
+    """
+    paths, txt_caps = _split_uploads(files)
     n = len(paths)
     if n > MAX_IMAGES:
         raise gr.Error(f"For now, up to {MAX_IMAGES} images are supported (got {n}).")
     updates = [gr.update(visible=n > 0), gr.update(visible=n > 0)]  # captioning_area, post_upload
     for i in range(MAX_IMAGES):
         visible = i < n
+        if visible:
+            base = os.path.splitext(os.path.basename(paths[i]))[0]
+            cap_val = txt_caps.get(base) or (trigger or "")
+        else:
+            cap_val = None
         updates.append(gr.update(visible=visible))  # row
         updates.append(gr.update(value=paths[i] if visible else None, visible=visible))  # image
-        updates.append(gr.update(value=(trigger or "") if visible else None, visible=visible))  # cap
+        updates.append(gr.update(value=cap_val, visible=visible))  # caption
     return updates
 
 
-def ai_suggest_trigger(images, concept_type):
-    paths = _paths(images)
+def ai_suggest_trigger(files, concept_type):
+    paths, _ = _split_uploads(files)
     if not paths:
         raise gr.Error("Upload images first.")
     return caption.suggest_trigger(paths, concept_type)
 
 
-def ai_caption_all(images, concept_type, trigger):
-    paths = _paths(images)
+def ai_caption_all(files, concept_type, trigger):
+    paths, _ = _split_uploads(files)
     if not paths:
         raise gr.Error("Upload images first.")
     caps = [caption.caption_one(p, concept_type, trigger) for p in paths]
     return [gr.update(value=caps[i]) if i < len(caps) else gr.update() for i in range(MAX_IMAGES)]
 
 
-def gather_dataset(images, *captions):
-    paths = _paths(images)
+def gather_dataset(files, *captions):
+    paths, _ = _split_uploads(files)
     return [[img, (captions[i] if i < len(captions) else "")] for i, img in enumerate(paths)]
 
 
@@ -240,7 +273,7 @@ with gr.Blocks(title="Krea 2 LoRA Trainer") as demo:
                 with gr.Row():
                     images = gr.File(
                         label="Upload your images (4–30 ideal)", file_count="multiple",
-                        file_types=["image"], interactive=True, scale=1,
+                        file_types=["image", ".txt"], interactive=True, scale=1,
                     )
                     with gr.Column(scale=3, visible=False) as captioning_area:
                         gr.Markdown("## Trigger & captioning")