Krea 2 LoRA trainer (HF Jobs backend)

README.md

@@ -1,13 +1,56 @@
 ---
-title: Krea2 Lora Trainer
-emoji: 👁
+title: Krea 2 LoRA Trainer
+emoji: 🎨
 colorFrom: indigo
 colorTo: yellow
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.13'
+python_version: '3.12'
 app_file: app.py
-pinned: false
+hardware: cpu-basic
+pinned: true
+hf_oauth: true
+hf_oauth_scopes:
+- read-repos
+- write-repos
+- manage-repos
+- jobs
+short_description: Train Krea 2 LoRAs on your images via HF Jobs
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Krea 2 LoRA Trainer
+
+Train a **DreamBooth-LoRA for Krea 2** from your own images, entirely on Hugging Face
+infrastructure:
+
+- **Sign in with Hugging Face** — the dataset, the job, and the pushed LoRA all run under
+  **your** account and billing (no pasted tokens);
+- the **Space** (this app, `cpu-basic`) collects your images + hyperparameters and submits a job;
+- training runs on **HF Jobs** using the diffusers Krea 2 trainer
+  (`examples/dreambooth/train_dreambooth_lora_krea2.py`);
+- the LoRA is **trained on Krea 2 RAW** and **validated / inferred on Krea 2 Turbo**, then pushed
+  to the Hub model repo you choose.
+
+You only pay for the Job's actual GPU runtime.
+
+## How tokens are used
+
+The Krea 2 weights are currently **gated**. The job downloads them with the Space's `KREA_TOKEN`
+secret and passes them to the trainer as **local dirs** — so your own token never needs Krea
+access, and the Krea token never touches your repos. Your token (from sign-in) is used only for
+your dataset and the pushed LoRA.
+
+> Set the `KREA_TOKEN` secret to a token with access to `krea/Krea-2-Raw` + `krea/Krea-2-Turbo`.
+
+## diffusers version
+
+The trainer lives in diffusers PR #14046 (branch `krea2-lora`). Once it is merged, set the
+`DIFFUSERS_REF` Space **variable** to `main` (or a release tag).
+
+## Usage
+
+1. Sign in with Hugging Face.
+2. Name your LoRA, set a trigger word / concept, and upload 4–30 images.
+3. (Optional) caption each image; blanks fall back to the trigger.
+4. Tweak hyperparameters if you like, pick a GPU flavor, and **Submit training job**.
+5. Copy the job id into the **Monitor** tab and **Refresh** to stream logs.

app.py

@@ -0,0 +1,270 @@
+"""Krea 2 LoRA Trainer — HF Space (HF Jobs backend).
+
+Sign in with Hugging Face, upload a handful of images (4–30 is ideal), optionally caption them,
+set a trigger word, and submit a DreamBooth-LoRA training job to HF Jobs. The job trains on
+**Krea 2 RAW**, validates / infers on **Krea 2 Turbo**, and pushes the LoRA to your Hub — all
+under your account. The Space runs on `cpu-basic`; the GPU work happens on HF Jobs.
+
+The gated Krea 2 weights are downloaded inside the job with the Space's `KREA_TOKEN` secret;
+the user's token is only ever used for their own dataset + the pushed LoRA.
+"""
+
+from __future__ import annotations
+
+import re
+
+import gradio as gr
+
+import jobs
+
+MAX_IMAGES = 40
+MAX_LOG = 60_000
+
+LR_SCHEDULERS = ["constant", "cosine", "linear", "constant_with_warmup", "polynomial"]
+OPTIMIZERS = ["adamW", "prodigy"]
+QUANT_CHOICES = [
+    ("None — bf16 (best quality, most VRAM)", "none"),
+    ("FP8 (faster compute, needs GPU ≥ 8.9)", "fp8"),
+    ("4-bit NF4 / QLoRA (lowest VRAM)", "4bit"),
+]
+
+FLAVOR_GUIDE = """**Which GPU?** You're billed per-minute of actual runtime. Krea 2 is a 12B DiT.
+
+| Flavor | VRAM | Best for |
+|---|---|---|
+| `l40sx1` | 48 GB | cheapest — pair with **4-bit NF4** quantization |
+| `a100-large` | 80 GB | **recommended (default)** — bf16 with offload + cached latents |
+| `h200` | 141 GB | fastest / highest resolution |
+
+First run is slow to start: it downloads the gated Krea 2 RAW + Turbo weights before training.
+"""
+
+TRIGGER_HELP = (
+    "A **trigger** that anchors your concept and is used as the default caption for every image. "
+    "For an object/character a rare token like `TOK` works; for a **style**, a descriptive phrase "
+    "(e.g. *“hand-drawn children's book illustration”*) works better than a random token."
+)
+CAPTION_HELP = (
+    "<details><summary>ℹ️ <b>Captioning tips</b></summary>\n\n"
+    "- Captions are <b>optional</b> — blank ones fall back to your trigger.\n"
+    "- For a <b>style</b> LoRA, describe what you do <i>not</i> want baked in (subject, scene) and "
+    "<i>omit</i> the stylistic parts you want learned, then keep the style trigger phrase.\n"
+    "- For an <b>object/character</b>, a trigger word plus the right class noun is enough.\n"
+    "</details>"
+)
+
+
+def _signin_state(profile: gr.OAuthProfile | None):
+    if profile is None:
+        return (
+            "🔒 **You're not signed in.** Use **Sign in with Hugging Face** (top-right) — the dataset, "
+            "the training job, and the resulting LoRA all run under **your** account.",
+            gr.update(),
+        )
+    return (
+        f"✅ Signed in as **{profile.username}** — the job and the pushed LoRA live under your account.",
+        gr.update(placeholder=f"e.g. {profile.username}/my-krea2-lora"),
+    )
+
+
+def load_captioning(images, instance_prompt):
+    """Reveal one (image, caption) row per uploaded image; prefill captions with the trigger."""
+    n = len(images) if images else 0
+    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)]  # captioning_area
+    for i in range(MAX_IMAGES):
+        visible = i < n
+        updates.append(gr.update(visible=visible))  # row
+        updates.append(gr.update(value=images[i] if visible else None, visible=visible))  # image
+        updates.append(gr.update(value=(instance_prompt or "") if visible else None, visible=visible))
+    return updates
+
+
+def gather_dataset(images, *captions):
+    """Pair uploaded image paths with their caption textbox values into a list of [path, caption]."""
+    images = images or []
+    return [[img, (captions[i] if i < len(captions) else "")] for i, img in enumerate(images)]
+
+
+def start_training(
+    dataset_rows, lora_name, instance_prompt, validation_prompt, rank, lora_alpha, max_train_steps,
+    learning_rate, lr_scheduler, resolution, repeats, train_batch_size, gradient_accumulation_steps,
+    seed, optimizer, use_8bit_adam, cache_latents, gradient_checkpointing, offload, quantization,
+    lora_layers, validation_epochs, hub_model_id, flavor, timeout,
+    profile: gr.OAuthProfile | None = None, oauth_token: gr.OAuthToken | None = None,
+):
+    if oauth_token is None or profile is None:
+        return "❌ Please **sign in with Hugging Face** first (top-right).", "", ""
+    if not dataset_rows:
+        return "❌ Upload at least one image.", "", ""
+    if not (lora_name or "").strip() and not (hub_model_id or "").strip():
+        return "❌ Give your LoRA a name (or a Hub model id).", "", ""
+
+    image_paths = [r[0] for r in dataset_rows]
+    captions = [r[1] for r in dataset_rows]
+    params = {
+        "lora_name": lora_name, "hub_model_id": hub_model_id,
+        "instance_prompt": instance_prompt, "validation_prompt": validation_prompt,
+        "rank": rank, "lora_alpha": lora_alpha, "max_train_steps": max_train_steps,
+        "learning_rate": learning_rate, "lr_scheduler": lr_scheduler, "resolution": resolution,
+        "repeats": repeats, "train_batch_size": train_batch_size,
+        "gradient_accumulation_steps": gradient_accumulation_steps, "seed": seed,
+        "optimizer": optimizer, "use_8bit_adam": bool(use_8bit_adam),
+        "cache_latents": bool(cache_latents), "gradient_checkpointing": bool(gradient_checkpointing),
+        "offload": bool(offload), "quantization": quantization, "lora_layers": lora_layers,
+        "validation_epochs": validation_epochs, "hf_token": oauth_token.token,
+    }
+    try:
+        res = jobs.submit(params, image_paths, captions, flavor=flavor, timeout=timeout)
+    except Exception as e:  # noqa: BLE001
+        return f"❌ Submission failed: {e}", "", ""
+
+    status = f"✅ Job submitted on **{flavor}**, running as **{profile.username}**."
+    link = (
+        f"**Job:** [{res['job_id']}]({res['url']})  \n"
+        f"**Dataset:** `{res['dataset_repo']}`  \n"
+        f"**LoRA will be pushed to:** `{res['hub_model_id']}`"
+    )
+    return status, link, res["job_id"]
+
+
+def refresh(job_id, oauth_token: gr.OAuthToken | None = None):
+    if not (job_id or "").strip():
+        return "Enter a job id.", ""
+    token = oauth_token.token if oauth_token else ""
+    st = jobs.job_status(job_id.strip(), token)
+    logs = jobs.job_logs(job_id.strip(), token)
+    return f"**Status:** `{st}`", logs[-MAX_LOG:] if len(logs) > MAX_LOG else logs
+
+
+with gr.Blocks(title="Krea 2 LoRA Trainer") as demo:
+    with gr.Row(equal_height=True):
+        gr.Markdown(
+            "# 🎨 Krea 2 LoRA Trainer\n"
+            "Train a LoRA on your own images — trains on **Krea 2 RAW**, validates on **Turbo**, "
+            "runs on **HF Jobs**, pushed to your Hub.",
+        )
+        gr.LoginButton(scale=0, min_width=220)
+
+    banner = gr.Markdown()
+
+    with gr.Tabs():
+        with gr.Tab("Train"):
+            lora_name = gr.Textbox(
+                label="LoRA name", placeholder="e.g. my-watercolor-style",
+                info="Used for your output model repo (you/<name>) and the dataset repo.",
+            )
+            instance_prompt = gr.Textbox(
+                label="Trigger word / concept", value="TOK", info=TRIGGER_HELP,
+            )
+            images = gr.File(
+                label="Upload your images (4–30 ideal)", file_count="multiple",
+                file_types=["image"], height=220,
+            )
+
+            with gr.Column(visible=False) as captioning_area:
+                gr.Markdown("**Captions** (optional) — edit per image, or leave as the trigger.")
+                gr.Markdown(CAPTION_HELP)
+                caption_rows, caption_imgs, caption_txts = [], [], []
+                for i in range(MAX_IMAGES):
+                    with gr.Row(visible=False) as row:
+                        img = gr.Image(
+                            show_label=False, interactive=False, height=90, width=90, scale=0,
+                        )
+                        cap = gr.Textbox(show_label=False, scale=4, container=False,
+                                         placeholder="caption for this image")
+                    caption_rows.append(row)
+                    caption_imgs.append(img)
+                    caption_txts.append(cap)
+
+            with gr.Accordion("Advanced options", open=False):
+                with gr.Row():
+                    rank = gr.Number(label="LoRA rank", value=32, precision=0,
+                                     info="Authors recommend 32; raise it for long runs / high-frequency styles.")
+                    lora_alpha = gr.Number(label="LoRA alpha", value=32, precision=0,
+                                           info="Keep equal to rank (scale 1.0).")
+                with gr.Row():
+                    max_train_steps = gr.Number(label="Training steps", value=1000, precision=0)
+                    learning_rate = gr.Number(label="Learning rate", value=3e-4,
+                                              info="3e-4 ~ 7e-4 with constant works well.")
+                with gr.Row():
+                    lr_scheduler = gr.Dropdown(LR_SCHEDULERS, value="constant", label="LR scheduler")
+                    resolution = gr.Number(label="Resolution", value=1024, precision=0)
+                with gr.Row():
+                    repeats = gr.Number(label="Dataset repeats", value=1, precision=0)
+                    seed = gr.Number(label="Seed", value=0, precision=0)
+                lora_layers = gr.Textbox(
+                    label="Target layers (optional)", placeholder="to_q,to_k,to_v,to_out.0,to_gate",
+                    info="Comma-separated. Blank = the recommended full set. For long runs, narrow to "
+                         "attention and raise the rank.",
+                )
+                with gr.Accordion("Memory / performance", open=False):
+                    with gr.Row():
+                        quantization = gr.Dropdown(QUANT_CHOICES, value="none", label="Quantization")
+                        optimizer = gr.Dropdown(OPTIMIZERS, value="adamW", label="Optimizer")
+                    with gr.Row():
+                        use_8bit_adam = gr.Checkbox(label="8-bit Adam", value=True)
+                        cache_latents = gr.Checkbox(label="Cache latents", value=True)
+                    with gr.Row():
+                        gradient_checkpointing = gr.Checkbox(label="Gradient checkpointing", value=True)
+                        offload = gr.Checkbox(label="CPU offload (VAE + text encoder)", value=False)
+                    with gr.Row():
+                        train_batch_size = gr.Number(label="Batch size", value=1, precision=0)
+                        gradient_accumulation_steps = gr.Number(label="Grad accumulation", value=1, precision=0)
+                with gr.Accordion("Validation", open=False):
+                    validation_prompt = gr.Textbox(
+                        label="Validation prompt", placeholder="(defaults to your trigger)",
+                        info="Generated on Turbo every N epochs to preview progress.",
+                    )
+                    validation_epochs = gr.Number(label="Validate every N epochs", value=25, precision=0)
+
+            with gr.Group():
+                hub_model_id = gr.Textbox(
+                    label="Output Hub model id (optional)", placeholder="you/my-krea2-lora",
+                    info="Where the trained LoRA is pushed. Blank = you/<lora-name>.",
+                )
+                with gr.Row():
+                    flavor = gr.Dropdown(jobs.FLAVORS, value=jobs.DEFAULT_FLAVOR, label="GPU flavor")
+                    timeout = gr.Textbox(label="Timeout", value="3h",
+                                         info="Max job runtime. First run downloads the model first.")
+                with gr.Accordion("GPU guide", open=False):
+                    gr.Markdown(FLAVOR_GUIDE)
+
+            submit_btn = gr.Button("🚀 Submit training job", variant="primary", size="lg")
+            status = gr.Markdown()
+            joblink = gr.Markdown()
+
+        with gr.Tab("Monitor"):
+            with gr.Row():
+                job_id = gr.Textbox(label="Job id", scale=3)
+                refresh_btn = gr.Button("🔄 Refresh", scale=1)
+            mon_status = gr.Markdown()
+            mon_logs = gr.Textbox(label="Job logs", lines=22, autoscroll=True, max_lines=22)
+
+    dataset_state = gr.State([])
+
+    # outputs must match load_captioning's interleaved returns: area, then (row, img, cap) per image
+    caption_outputs = [captioning_area]
+    for r, im, c in zip(caption_rows, caption_imgs, caption_txts):
+        caption_outputs += [r, im, c]
+
+    demo.load(_signin_state, inputs=None, outputs=[banner, hub_model_id])
+    images.change(load_captioning, inputs=[images, instance_prompt], outputs=caption_outputs)
+
+    submit_btn.click(
+        gather_dataset, inputs=[images, *caption_txts], outputs=dataset_state,
+    ).then(
+        start_training,
+        inputs=[dataset_state, lora_name, instance_prompt, validation_prompt, rank, lora_alpha,
+                max_train_steps, learning_rate, lr_scheduler, resolution, repeats, train_batch_size,
+                gradient_accumulation_steps, seed, optimizer, use_8bit_adam, cache_latents,
+                gradient_checkpointing, offload, quantization, lora_layers, validation_epochs,
+                hub_model_id, flavor, timeout],
+        outputs=[status, joblink, job_id],
+    )
+    refresh_btn.click(refresh, inputs=[job_id], outputs=[mon_status, mon_logs])
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=4).launch()

jobs.py

@@ -0,0 +1,265 @@
+"""HF Jobs backend for the Krea 2 LoRA trainer Space.
+
+Per training request the Space (cpu-basic, no GPU/torch):
+  1. stages the uploaded images + a `metadata.jsonl` (per-image captions),
+  2. pushes them to a private HF **dataset** repo under the signed-in user,
+  3. generates a self-contained UV job script,
+  4. submits it with `HfApi.run_uv_job(... token=<user>)` so the job runs + is billed
+     to the signed-in user and the trained LoRA is pushed to their Hub.
+
+Token split (important):
+  * The **gated Krea 2 weights** (`krea/Krea-2-Raw`, `krea/Krea-2-Turbo`) are not public.
+    They are pre-downloaded *inside the job* with the Space's `KREA_TOKEN` secret and passed
+    to the trainer as **local dirs**, so `from_pretrained` needs no Krea auth.
+  * Everything else (dataset download, `create_repo`/`upload_folder` of the LoRA) uses the
+    job's ambient `HF_TOKEN` env = the **signed-in user's** token. The Krea token never
+    touches the user's repos and the user's token never needs Krea access.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+import tempfile
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+# The Krea 2 LoRA trainer lives in diffusers PR #14046 (branch `krea2-lora`). Once it is merged,
+# set the `DIFFUSERS_REF` Space variable to `main` (or a release tag) — no code change needed.
+DIFFUSERS_REF = os.environ.get("DIFFUSERS_REF", "krea2-lora")
+BASE_MODEL_RAW = "krea/Krea-2-Raw"      # non-distilled base — train LoRA on this
+BASE_MODEL_TURBO = "krea/Krea-2-Turbo"  # 8-step distilled — validate / infer on this
+
+DEFAULT_FLAVOR = "a100-large"
+FLAVORS = ["l40sx1", "a100-large", "h200"]
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+
+# Deterministic on-job paths the trainer reads from (baked into the CLI args below).
+JOB_RAW = "/tmp/krea/raw"
+JOB_TURBO = "/tmp/krea/turbo"
+JOB_DATA = "/tmp/data"
+JOB_OUT = "/tmp/out"
+JOB_BNB = "/tmp/bnb.json"
+
+
+def slug(name: str) -> str:
+    s = re.sub(r"[^a-zA-Z0-9-]+", "-", (name or "").strip()).strip("-").lower()
+    return s or "krea2-lora"
+
+
+def _namespace(token: str) -> str:
+    from huggingface_hub import whoami  # noqa: PLC0415
+    return whoami(token=token)["name"]
+
+
+def build_metadata(image_paths: list[str], captions: list[str], instance_prompt: str) -> list[dict]:
+    """One `metadata.jsonl` row per image: {file_name, prompt}. Empty captions fall back to the
+    instance prompt (the trigger sentence). Files are renamed to a stable `0000.ext` order."""
+    rows = []
+    fallback = (instance_prompt or "a photo").strip()
+    for i, p in enumerate(image_paths):
+        cap = ""
+        if i < len(captions) and captions[i]:
+            cap = str(captions[i]).strip()
+        rows.append({"file_name": f"{i:04d}{Path(p).suffix.lower()}", "prompt": cap or fallback})
+    return rows
+
+
+def build_train_args(params: dict, hub_model_id: str) -> list[str]:
+    """Turn UI params into the `train_dreambooth_lora_krea2.py` CLI. Krea weights are passed as
+    local dirs (pre-downloaded in the job); the dataset is a local imagefolder (image/prompt cols)."""
+    instance_prompt = (params.get("instance_prompt") or "TOK").strip()
+    val_prompt = (params.get("validation_prompt") or instance_prompt).strip()
+    args = [
+        "--pretrained_model_name_or_path", JOB_RAW,
+        "--validation_model_path", JOB_TURBO,
+        "--dataset_name", JOB_DATA,
+        "--image_column", "image",
+        "--caption_column", "prompt",
+        "--instance_prompt", instance_prompt,
+        "--output_dir", JOB_OUT,
+        "--mixed_precision", "bf16",
+        "--resolution", str(int(params["resolution"])),
+        "--train_batch_size", str(int(params["train_batch_size"])),
+        "--gradient_accumulation_steps", str(int(params["gradient_accumulation_steps"])),
+        "--repeats", str(int(params["repeats"])),
+        "--rank", str(int(params["rank"])),
+        "--lora_alpha", str(int(params["lora_alpha"])),
+        "--learning_rate", str(float(params["learning_rate"])),
+        "--lr_scheduler", str(params["lr_scheduler"]),
+        "--lr_warmup_steps", "0",
+        "--max_train_steps", str(int(params["max_train_steps"])),
+        "--optimizer", str(params["optimizer"]),
+        "--seed", str(int(params["seed"])),
+        "--validation_prompt", val_prompt,
+        "--validation_epochs", str(int(params["validation_epochs"])),
+        "--num_validation_images", "2",
+        "--push_to_hub",
+        "--hub_model_id", hub_model_id,
+    ]
+    if params.get("lora_layers"):
+        args += ["--lora_layers", str(params["lora_layers"]).strip()]
+    if params.get("gradient_checkpointing", True):
+        args += ["--gradient_checkpointing"]
+    if params.get("cache_latents", True):
+        args += ["--cache_latents"]
+    if params.get("offload"):
+        args += ["--offload"]
+    if params.get("use_8bit_adam") and str(params["optimizer"]).lower() == "adamw":
+        args += ["--use_8bit_adam"]
+    quant = params.get("quantization", "none")
+    if quant == "fp8":
+        args += ["--do_fp8_training"]
+    elif quant == "4bit":
+        args += ["--bnb_quantization_config_path", JOB_BNB]
+    return args
+
+
+# --------------------------------------------------------------------------------------
+# UV job script (runs on HF Jobs GPU hardware)
+# --------------------------------------------------------------------------------------
+
+JOB_SCRIPT_TEMPLATE = '''# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "git+https://github.com/huggingface/diffusers.git@{ref}",
+#   "torch",
+#   "torchvision",
+#   "transformers>=4.41.2",
+#   "accelerate>=0.31.0",
+#   "peft>=0.11.1",
+#   "datasets",
+#   "bitsandbytes",
+#   "prodigyopt",
+#   "ftfy",
+#   "sentencepiece",
+#   "hf_transfer",
+#   "huggingface_hub[hf-xet]",
+# ]
+# ///
+"""Auto-generated Krea 2 DreamBooth-LoRA job. Trains on Krea 2 RAW, validates on Turbo."""
+import json, os, subprocess, sys, urllib.request
+from pathlib import Path
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+REF = "{ref}"
+RAW, TURBO, DATA, BNB = "{raw}", "{turbo}", "{data}", "{bnb}"
+DATASET_REPO = {dataset_repo!r}
+QUANT = {quant!r}
+TRAIN_ARGS = {train_args}
+KREA_TOKEN = os.environ["KREA_TOKEN"]  # gated Krea weights ONLY
+
+SCRIPT_URL = (
+    "https://raw.githubusercontent.com/huggingface/diffusers/"
+    + REF + "/examples/dreambooth/train_dreambooth_lora_krea2.py"
+)
+
+
+def main():
+    from huggingface_hub import snapshot_download
+
+    print("=== 1/4 download gated Krea 2 weights (krea token) ===", flush=True)
+    snapshot_download("krea/Krea-2-Raw", local_dir=RAW, token=KREA_TOKEN)
+    snapshot_download("krea/Krea-2-Turbo", local_dir=TURBO, token=KREA_TOKEN)
+
+    print("=== 2/4 download dataset (user token / HF_TOKEN env) ===", flush=True)
+    snapshot_download(DATASET_REPO, repo_type="dataset", local_dir=DATA)
+
+    if QUANT == "4bit":
+        Path(BNB).write_text(json.dumps({{
+            "load_in_4bit": True, "bnb_4bit_quant_type": "nf4",
+            "bnb_4bit_compute_dtype": "bfloat16",
+        }}))
+
+    print("=== 3/4 fetch trainer script @ " + REF + " ===", flush=True)
+    urllib.request.urlretrieve(SCRIPT_URL, "/tmp/train_dreambooth_lora_krea2.py")
+
+    print("=== 4/4 accelerate launch (pushes LoRA to the Hub) ===", flush=True)
+    cmd = [sys.executable, "-m", "accelerate.commands.launch",
+           "/tmp/train_dreambooth_lora_krea2.py", *TRAIN_ARGS]
+    print(">>> " + " ".join(cmd), flush=True)
+    subprocess.run(cmd, check=True)
+    print("=== DONE ===", flush=True)
+
+
+if __name__ == "__main__":
+    main()
+'''
+
+
+def submit(params: dict, image_paths: list[str], captions: list[str],
+           flavor: str, timeout: str) -> dict:
+    """Stage dataset → push private dataset repo → generate UV script → submit job.
+    Returns {job_id, url, dataset_repo, hub_model_id}."""
+    token = (params.get("hf_token") or "").strip()
+    if not token:
+        raise ValueError("Missing user token (sign in with Hugging Face).")
+    if not os.environ.get("KREA_TOKEN"):
+        raise RuntimeError("Space is missing the KREA_TOKEN secret (gated Krea 2 access).")
+
+    ns = _namespace(token)
+    name = slug(params.get("lora_name", ""))
+    dataset_repo = f"{ns}/{name}-dataset"
+    hub_model_id = (params.get("hub_model_id") or "").strip() or f"{ns}/{name}"
+
+    api = HfApi(token=token)
+    tmp = Path(tempfile.mkdtemp(prefix="krea2-"))
+    try:
+        # stage images under stable names + metadata.jsonl
+        rows = build_metadata(image_paths, captions, params.get("instance_prompt", ""))
+        for row, src in zip(rows, image_paths):
+            shutil.copy(src, tmp / row["file_name"])
+        (tmp / "metadata.jsonl").write_text(
+            "\n".join(json.dumps(r) for r in rows) + "\n"
+        )
+
+        # push the dataset (private, user namespace)
+        api.create_repo(dataset_repo, repo_type="dataset", private=True, exist_ok=True, token=token)
+        api.upload_folder(repo_id=dataset_repo, repo_type="dataset", folder_path=str(tmp), token=token)
+
+        # render the job script
+        train_args = build_train_args(params, hub_model_id)
+        script = JOB_SCRIPT_TEMPLATE.format(
+            ref=DIFFUSERS_REF, raw=JOB_RAW, turbo=JOB_TURBO, data=JOB_DATA, bnb=JOB_BNB,
+            dataset_repo=dataset_repo, quant=params.get("quantization", "none"),
+            train_args=json.dumps(train_args),
+        )
+        script_path = tmp / "job_train.py"
+        script_path.write_text(script)
+
+        job = api.run_uv_job(
+            str(script_path),
+            flavor=flavor,
+            timeout=timeout,
+            # HF_TOKEN = user (push + dataset); KREA_TOKEN = gated Krea weights only.
+            secrets={"HF_TOKEN": token, "KREA_TOKEN": os.environ["KREA_TOKEN"]},
+            token=token,
+        )
+        job_id = getattr(job, "id", "") or ""
+        url = getattr(job, "url", "") or (f"https://huggingface.co/jobs/{ns}/{job_id}" if job_id else "")
+        return {"job_id": job_id, "url": url, "dataset_repo": dataset_repo, "hub_model_id": hub_model_id}
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
+def job_logs(job_id: str, token: str = "") -> str:
+    try:
+        return "\n".join(HfApi(token=token).fetch_job_logs(job_id=job_id, token=token))
+    except Exception as e:  # noqa: BLE001
+        return f"(could not fetch logs: {e})"
+
+
+def job_status(job_id: str, token: str = "") -> str:
+    try:
+        job = HfApi(token=token).inspect_job(job_id=job_id, token=token)
+        status = getattr(job, "status", None)
+        stage = getattr(status, "stage", None)
+        if stage is None and isinstance(status, dict):
+            stage = status.get("stage")
+        return str(stage or status or "UNKNOWN")
+    except Exception as e:  # noqa: BLE001
+        return f"UNKNOWN ({e})"

requirements.txt

@@ -0,0 +1,3 @@
+gradio[oauth]>=6.18
+huggingface_hub[hf-xet]>=1.5
+hf_transfer