agent: host /agent recipe + /train_job.py + /api/train; mount via FastAPI

app.py

@@ -37,6 +37,122 @@ CONCEPT_CHOICES = [
     ("⚙️ Custom", "custom"),
 ]
 
+SPACE_HOST = os.environ.get("SPACE_HOST", "multimodalart-krea2-lora-trainer.hf.space")
+AGENT_URL = f"https://{SPACE_HOST}/agent"
+
+_AGENT_DOC = r'''# Train a Krea 2 LoRA — agent recipe
+
+You can train a custom **Krea 2** LoRA on Hugging Face Jobs. Krea 2 weights are public, so you
+only need the USER's own write-scoped HF token (https://huggingface.co/settings/tokens). Every
+step below runs under that token and is billed to that account.
+
+## Inputs to collect from the user
+- **lora_name** — a unique name. The result is pushed to `<user>/<lora_name>`.
+- **trigger** — a style phrase (e.g. `"in the style of TOK"`) or, for a subject/face/object, a
+  rare token (e.g. `"TOK person"`). It is included in every training caption.
+- **concept** — one of: `style` | `character` | `object` | `face`.
+- **images** — 4–30 images. The user will give you ONE of:
+  (a) an existing HF dataset repo with `image` + `prompt` columns → use it directly, skip Step 1;
+  (b) a local folder of images → build the dataset (Step 1);
+  (c) a request like *"find impasto-style paintings"* → gather the images yourself first, then
+  build the dataset (Step 1).
+
+## Step 1 — build the dataset (skip if the user already has one)
+Create an imagefolder dataset under the user's account; each row is an image + a caption that
+contains the trigger.
+
+```python
+import json, os
+from pathlib import Path
+from huggingface_hub import HfApi
+
+TOKEN   = os.environ["HF_TOKEN"]                 # the user's write token
+REPO    = "<user>/<lora_name>-dataset"
+TRIGGER = "<trigger>"
+images   = ["img1.png", "img2.png"]              # local paths you gathered
+captions = ["a cat on a sofa, " + TRIGGER, "a cat by a window, " + TRIGGER]  # 1 per image
+
+api = HfApi(token=TOKEN)
+api.create_repo(REPO, repo_type="dataset", private=True, exist_ok=True)
+tmp = Path("ds"); tmp.mkdir(exist_ok=True)
+rows = []
+for i, (src, cap) in enumerate(zip(images, captions)):
+    fn = f"{i:03d}{Path(src).suffix.lower()}"
+    (tmp / fn).write_bytes(Path(src).read_bytes())
+    rows.append({"file_name": fn, "prompt": cap})
+(tmp / "metadata.jsonl").write_text("\n".join(json.dumps(r) for r in rows))
+api.upload_folder(repo_id=REPO, repo_type="dataset", folder_path=str(tmp))
+print("dataset:", REPO)
+```
+
+No captions yet? Use the trigger alone as every caption — fine for a single-concept style.
+
+## Step 2 — launch training (pick ONE)
+
+### Option A — run the job yourself (recommended)
+Submit the job on HF Jobs under the user's token:
+
+```bash
+export HF_TOKEN=hf_...                     # the user's write token
+curl -sO https://__HOST__/train_job.py
+hf jobs uv run --flavor rtx-pro-6000 --timeout 1h -s HF_TOKEN=$HF_TOKEN train_job.py -- \
+    --dataset <user>/<lora_name>-dataset \
+    --lora-name <lora_name> \
+    --trigger "<trigger>" \
+    --concept style
+```
+
+Python equivalent:
+
+```python
+from huggingface_hub import HfApi
+api = HfApi(token=TOKEN)
+job = api.run_uv_job(
+    "https://__HOST__/train_job.py",
+    flavor="rtx-pro-6000", timeout="1h",
+    secrets={"HF_TOKEN": TOKEN},
+    script_args=["--dataset", "<user>/<lora_name>-dataset",
+                 "--lora-name", "<lora_name>",
+                 "--trigger", "<trigger>", "--concept", "style"],
+)
+print(job.url)
+```
+
+Useful flags: `--steps` (1000), `--rank` (32), `--resolution` (1024), `--learning-rate` (3e-4),
+`--quantization {none,fp8,4bit}`, `--no-gallery`, `--num-gallery` (3). Run with `--help` for all.
+
+### Option B — let this Space submit the job for you
+POST the user's token + the dataset repo; the Space submits the job (billed to that token):
+
+```bash
+curl -sX POST https://__HOST__/api/train -H 'Content-Type: application/json' -d '{
+  "hf_token": "hf_...",
+  "dataset_repo": "<user>/<lora_name>-dataset",
+  "lora_name": "<lora_name>",
+  "trigger": "<trigger>",
+  "concept": "style"
+}'
+# -> {"job_id": "...", "url": "https://huggingface.co/jobs/...", "hub_model_id": "<user>/<lora_name>"}
+```
+
+## Result
+~40 min on `rtx-pro-6000` (1000 steps, regional torch.compile). The LoRA is pushed to
+`<user>/<lora_name>` with a preview gallery + README. Use it:
+
+```python
+import torch
+from diffusers import Krea2Pipeline
+pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", torch_dtype=torch.bfloat16).to("cuda")
+pipe.load_lora_weights("<user>/<lora_name>")
+image = pipe("<trigger>, a fox in a snowy forest", num_inference_steps=8, guidance_scale=0.0).images[0]
+image.save("out.png")
+```
+'''
+
+
+def build_agent_doc() -> str:
+    return _AGENT_DOC.replace("__HOST__", SPACE_HOST)
+
 CSS = """
 #main_title{text-align:center}
 #main_title h1{font-size:2.25rem;margin-bottom:0}
@@ -46,6 +162,8 @@ CSS = """
 .login_logout{width:100% !important}
 #login{width:100% !important;margin:0 auto}
 #signin_banner{text-align:center}
+.agent_hint{text-align:center;opacity:.7;font-size:.9em;margin:.4em 0 -.2em}
+.agent_curl{max-width:380px;margin:0 auto}
 #main_ui.locked{pointer-events:none;opacity:.55;filter:grayscale(.3)}
 
 /* captioning panel */
@@ -292,6 +410,9 @@ with gr.Blocks(title="Krea 2 LoRA Trainer") as demo:
         gr.Column(scale=2, min_width=0)
         with gr.Column(scale=0, min_width=380):
             gr.LoginButton(elem_id="login", elem_classes=["login_logout"])
+            gr.Markdown("…or hand it to your coding agent:", elem_classes=["agent_hint"])
+            gr.Code(f"curl {AGENT_URL}", language="shell", interactive=False,
+                    show_label=False, elem_classes=["agent_curl"])
         gr.Column(scale=2, min_width=0)
     banner = gr.Markdown(elem_id="signin_banner")
 
@@ -455,4 +576,31 @@ with gr.Blocks(title="Krea 2 LoRA Trainer") as demo:
 
 
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=4).launch(theme=THEME, css=CSS)
+    import uvicorn
+    from fastapi import FastAPI, Request
+    from fastapi.responses import JSONResponse, PlainTextResponse
+
+    fastapi_app = FastAPI()
+
+    @fastapi_app.get("/agent", response_class=PlainTextResponse)
+    def _agent_doc():
+        return build_agent_doc()
+
+    @fastapi_app.get("/train_job.py", response_class=PlainTextResponse)
+    def _train_job():
+        return jobs.build_agent_job_script()
+
+    @fastapi_app.post("/api/train")
+    async def _api_train(request: Request):
+        try:
+            payload = await request.json()
+        except Exception:  # noqa: BLE001
+            return JSONResponse({"error": "invalid JSON body"}, status_code=400)
+        try:
+            return JSONResponse(jobs.submit_from_dataset(payload))
+        except Exception as e:  # noqa: BLE001
+            return JSONResponse({"error": str(e)}, status_code=400)
+
+    demo.queue(default_concurrency_limit=4)
+    app = gr.mount_gradio_app(fastapi_app, demo, path="/", theme=THEME, css=CSS)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)))