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

A newer version of the Gradio SDK is available: 6.29.1

Upgrade

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.

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:

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:

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:

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")