PRD.md

# PRD — Klein LoRA Studio

> A "signature look" restyle studio built on FLUX.2 [klein] 4B. Upload a photo,
> apply a curated style, and see base vs **your own trained LoRA** side by side —
> so the aesthetic is consistent, recognizable, and genuinely *yours*, not a
> prompt anyone could copy.

| Field | Value |
| --- | --- |
| **Hackathon** | [Build Small](https://huggingface.co/build-small-hackathon) (Hugging Face × Gradio) |
| **Track** | 🏡 Backyard AI (practical) — also reads as 🍄 Thousand Token Wood depending on framing |
| **Partner kit** | Black Forest Labs — FLUX.2 [klein] |
| **Model** | `black-forest-labs/FLUX.2-klein-4B` (4B, Apache 2.0) + a custom style LoRA trained on `FLUX.2-klein-base-4B` |
| **Badges targeted** | 🏅 Tiny Titan (≤4B), 🎨 Off Brand (custom UI), 🎬 Best Demo, 🧩 Bonus Quest Champion |
| **Deadline** | **June 15, 2026 · 23:59 UTC** |
| **Starter** | Forked from `stephenbtl/klein-build-small-starter` (this repo) |

---

## 1. Summary

Klein LoRA Studio turns the starter's generic "Image → Image" tab into a focused
**personal restyle studio**. The user picks a *named look* (e.g. "Cottagecore
Kitchen", "Mid-century Living Room", "Riso Print"), drops in a photo, and gets a
side-by-side **before → after**. The differentiator from a plain prompt box is
that the app's signature look is delivered by a **LoRA we train** on ~20–40
reference images, so every output is consistent and the style is locked to *this*
app — exactly the 🎯 fine-tuning angle the partner kit is built around.

This is the LoRA-forward sibling of the other two concept branches. It leans on
the starter's existing `lora_compare` machinery (base ↔ LoRA at the same seed)
and elevates it from a developer comparison tool into a polished consumer app.

## 2. The problem & user

**User:** someone non-technical with a pile of photos and a specific aesthetic in
mind — "what would my living room look like in Scandinavian?", "restyle my whole
garden set as cottagecore", "make these portraits look like my illustration
style." Plain text-to-image can't keep a *consistent* look across a set, and
prompt-engineering a style is fiddly and copyable.

**Insight (from `STARTER_IDEAS.md` #3 + `PROMPTING.md`):** narrow beats broad.
Hide the prompt engineering behind named presets, make the output *consistent*
(that is what a LoRA buys you), and always show a before/after — it demos far
better than a single image.

## 3. Why it fits "Build Small" (rule → how we satisfy it)

| Rule / badge | How this app delivers it |
| --- | --- |
| **REQ-01 ≤ 32B** | klein 4B is 4B params. The LoRA adds negligible params. ✅ |
| **REQ-02 Gradio Space in org** | Forked Gradio Space; deploy into the Build Small HF org. ✅ |
| **REQ-03 Demo video** | Before/after restyle of a real room + the "apply look across a set" flow. |
| **REQ-04 Social post** | One post, linked from README. |
| **REQ-05 ZeroGPU ≤10 apps/user** | Single Space, ZeroGPU (`zero-a10g`). ✅ |
| **REQ-06 README tags + write-up** | See §9 for the exact YAML block. |
| 🏅 **Tiny Titan (≤4B)** | Whole experience runs on klein **4B** — qualifies for the ≤4B badge. |
| 🎨 **Off Brand** | Replace the 6-tab dev UI with a single-purpose two-pane "studio" layout. |
| 🎬 **Best Demo** | Side-by-side restyle is inherently demo-friendly. |

## 4. Scope

### MVP (must ship by deadline)
1. **One-screen studio UI** (not the starter's 6 tabs): input photo on the left,
   restyled result on the right, a row of **named look presets** underneath.
2. **Preset = curated edit prompt + the signature LoRA** applied at a fixed
   scale. User clicks a look; no prompt writing required.
3. **Before → after** display (original and result side by side), with a "swap"
   or slider if time allows.
4. **The signature style LoRA** trained and loaded by default (see §6). Ship at
   least **one** strong trained look; presets beyond it can be curated edit
   prompts layered on the same LoRA.

### Stretch
- **"Save this look → apply to a set":** batch-restyle several uploads so a whole
  photo set is consistent (the thing a plain prompt can't give them).
- **Strength slider** (LoRA scale 0.0–1.5, reuse `l_scale`) exposed as "subtle →
  strong."
- **2×2 variation grid** (4 seeds) so the user picks their favourite.

### Out of scope
- Training inside the Space (training is offline via AI Toolkit — see §6).
- Text-to-image from scratch (this app is restyle-first; keep a hidden T2I path
  only for preset previews if useful).

## 5. Technical design (grounded in the starter `app.py`)

The starter already contains everything we need; this is mostly **subtraction and
re-skinning**, not new ML code.

- **Pipeline:** `Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=torch.bfloat16)`, built on CPU at module scope and moved to CUDA inside the `@GPU` call — keep the starter's ZeroGPU pattern verbatim (`import spaces` *before* torch; `get_pipe()` moves to cuda).
- **Restyle call:** reuse `img2img()` / `lora_compare()` logic. Editing footgun:
  always call `pipe(prompt=…, image=…)` **by keyword** — `image` is the first
  positional arg.
- **Sizing:** keep `klein_size(w, h)` — snaps to multiples of 16 under klein's
  4096-patch ceiling; resize input with `Image.LANCZOS`.
- **Steps / guidance:** distilled 4B → `num_inference_steps=4`. The LoRA is
  trained on **base** (50 steps) but loads on distilled for fast demos with mild
  drift (per `TRAIN_A_LORA.md`). For the cleanest result, optionally switch a
  "High quality" toggle to `FLUX.2-klein-base-4B` at 50 steps, guidance 4.0.
- **LoRA loading:** reuse `_apply_lora()` / `_unload_lora()`:
  `pipe.load_lora_weights(path, adapter_name="user")` then
  `pipe.set_adapters(["user"], adapter_weights=[scale])`. Always unload in a
  `finally` so the shared pipe stays clean.
- **Presets data structure:** a `LOOKS = {name: {"prompt": <edit prompt>, "lora_scale": <float>}}` dict, mirroring the starter's `HUB_LORAS` pattern.
- **Deps:** unchanged `requirements.txt` (diffusers main for `Flux2KleinPipeline`, `peft` for `load_lora_weights`, `spaces` for ZeroGPU).

### Files to change
| File | Change |
| --- | --- |
| `app.py` | Replace tabbed UI with the single studio layout; add `LOOKS`; load signature LoRA at startup; wire preset clicks → `img2img`+LoRA. Keep the ZeroGPU + `klein_size` helpers. |
| `README.md` | New frontmatter (title/short_description/tags), submission write-up, demo + social links (§9). |
| `configs/my_lora_klein_4b.yaml` | Set `trigger_word`, `name`, dataset path for the signature look (§6). |
| `examples/` | Add 2–3 before/after pairs for the chosen look (used in the demo + gr.Examples). |

## 6. The signature LoRA (the differentiator)

Train **one** style LoRA that becomes the app's identity (e.g. a clean editorial
illustration look, or a specific interior aesthetic). Following `TRAIN_A_LORA.md`:

- **Base model:** `FLUX.2-klein-base-4B` (Apache 2.0). Distilled won't fine-tune.
- **Data:** 15–40 images sharing one look, ≥1024px, one `.txt` caption each.
  **Caption content, not style.** Start every caption with a coined trigger
  (`ZK_LOOK`).
- **Config:** `configs/my_lora_klein_4b.yaml` — change the 3 `<<< CHANGE >>>`
  lines + sample prompts. **Keep `arch: "flux2_klein_4b"`** (omitting it crashes
  ai-toolkit, issue #691).
- **Train:** AI Toolkit on RunPod RTX 4090 (~$0.50, ~30–40 min, ~1800 steps).
  **Watch the sample images, not the loss** — pick the best checkpoint (~step
  750–1500), not the last.
- **Use:** download the `.safetensors`, commit it (or load from the Hub), and
  apply it as the default adapter in the studio.

## 7. Demo & social plan (REQ-03 / REQ-04)

- **Demo video (2–4 min):** open the studio, restyle a real living-room photo
  through 2–3 looks, then the "apply look to a whole set" flow showing
  consistency. Sell it — "no humility" (per the field guide trail map).
- **Social post:** before/after grid + one-line pitch + Space link; link the post
  back from the README.

## 8. Risks & mitigations

| Risk | Mitigation |
| --- | --- |
| ZeroGPU cold-start / 60s budget | Build pipeline on CPU at module scope (starter already does this); keep distilled 4-step path as default. |
| LoRA drift on distilled | Offer a "High quality" base-4B/50-step toggle; train at sensible strength. |
| Style leaks into captions → weak LoRA | Caption content only; skim `.txt` for style words (they leak ~25%). |
| Subject not preserved in restyle | Use edit prompts that describe *the change*, not the whole scene (`PROMPTING.md`). |
| Deadline today | MVP = studio UI + one trained look + before/after. Everything else is stretch. |

## 9. Submission checklist (REQ-01 → REQ-06)

- [ ] **REQ-01** Every model ≤32B — klein 4B + LoRA. ✅ (also ≤4B → Tiny Titan)
- [ ] **REQ-02** Gradio Space uploaded into the Build Small HF org.
- [ ] **REQ-03** Demo video recorded and linked.
- [ ] **REQ-04** One social post, linked from README.
- [ ] **REQ-05** ≤10 ZeroGPU apps for this user.
- [ ] **REQ-06** README YAML tagged + idea/tech write-up.

**README frontmatter to apply at submission (REQ-06):**
```yaml
title: Klein LoRA Studio
short_description: Restyle photos with your own trained klein 4B signature look
sdk: gradio
app_file: app.py
license: apache-2.0
suggested_hardware: zero-a10g
models:
  - black-forest-labs/FLUX.2-klein-4B
  - black-forest-labs/FLUX.2-klein-base-4B
tags:
  - build-small-hackathon
  - backyard-ai          # track
  - tiny-titan           # ≤4B badge
  - off-brand            # custom UI badge
  - best-demo            # badge
  - flux
  - image-to-image
  - lora
```

## 10. Definition of done

- The Space loads on ZeroGPU with no token/gating and restyles an uploaded photo
  through at least one named look powered by the signature LoRA.
- Before → after is visible in one screen.
- README has the track + badge tags, a short idea/tech write-up, and links to the
  demo video and social post.
- Submitted into the Build Small org before 23:59 UTC.