krea2train / agent.md
multimodalart's picture
multimodalart HF Staff
Add agent recipe (agent.md) for curl-based LoRA training
f99c6cc verified
|
Raw History Blame Contribute Delete
3.87 kB
# 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 the training job
Download the ready-made job script and submit it on HF Jobs under the user's token:
```bash
export HF_TOKEN=hf_... # the user's write token
curl -sO https://huggingface.co/spaces/multimodalart/krea2-lora-trainer/raw/main/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://huggingface.co/spaces/multimodalart/krea2-lora-trainer/raw/main/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.
## 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")
```