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 (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):

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.