Krea 2 LoRA Trainer

Trains a LoRA for Krea 2 via fal.ai's fal-ai/krea-2-trainer API, downloads the result, and converts it to the layer names ComfyUI expects (via a vendored copy of ComfyUI-Pixaroma's _krea_lora_convert_helpers.py).

Pipeline for one job:

input/<set>/*.png (+*.txt captions)
        │  zip + upload
        ▼
   fal.ai queue (submit → poll by request_id → result)
        │  download lora_file
        ▼
outputs/<trigger_phrase>/<name>.safetensors            (as returned by fal.ai)
outputs/<trigger_phrase>/<name>_converted.safetensors   (renamed keys for ComfyUI)

The web UI shows a connection indicator (green/red dot + username + credit balance) at the top, a "Reconnect" button, and a way to paste in an API key without touching .env - see Connecting to fal.ai.

Architecture

The project is deliberately split into layers so it can run standalone or be mounted inside another (monolith) web backend:

• app/core/ — all business logic, no FastAPI and no hardcoded disk paths:
  • training_service.py — TrainingService, the job state machine (create_job, advance(job_id), list_jobs, cancel_job). This is what a monolith imports directly, without going through HTTP.
  • job_store.py — JobStore (Protocol) + JsonFileJobStore for standalone mode. A monolith plugs in its own DB-backed implementation with the same three methods.
  • storage.py — Storage (Protocol) + LocalStorage (the input// outputs/ folders on disk). A monolith can swap in its own storage (S3, etc.).
  • fal_gateway.py — a thin wrapper around the official fal_client package (upload/submit/poll/result/cancel), plus check_connection() (calls fal.ai's account/billing API to verify the key and read back username + credit balance) and set_api_key() to swap credentials at runtime.
  • settings_store.py — SettingsStore, a tiny JSON file that persists an API key entered through the UI (data/settings.json) so it survives a restart. Optional: a monolith normally wouldn't pass one in, since it manages its own secrets.
  • converter.py — calls convert_file() from the vendored helpers.
  • models.py — Job, TrainingParams, JobStatus — plain dataclasses, no pydantic/FastAPI dependency.
• app/api/ — the HTTP layer:
  • routes.py::build_router(service) returns a plain APIRouter. A monolith does monolith_app.include_router(build_router(shared_service), prefix="/krea2-trainer") and gets the same endpoints without a second FastAPI app.
  • app.py::create_app(service, poll_interval_seconds=...) — the standalone app: same routes + the web UI's static files + the background poller.
• app/worker/poller.py — a simple asyncio loop that calls service.advance_all() every KREA2_POLL_INTERVAL_SECONDS seconds. Only needed standalone; a monolith would typically advance jobs from its own task queue (Celery/RQ/cron/webhook) by calling service.advance(job_id) itself.
• app/vendor/krea_lora_convert_helpers.py — an exact copy of ComfyUI-Pixaroma's _krea_lora_convert_helpers.py (stdlib-only, no torch/comfy imports), so this app does not depend on a ComfyUI install being present. Set KREA2_HELPERS_PATH to point at a live ComfyUI checkout instead, if you want a single source of truth for the mapping tables.
• app/web/static/ — a plain vanilla HTML/CSS/JS UI for standalone mode.

Job state machine

PENDING → PACKAGING → QUEUED → IN_PROGRESS → DOWNLOADING → CONVERTING → DONE
                                                                      ↘ FAILED / CANCELLED (from any state)

TrainingService.advance(job_id) performs exactly one bounded step and persists progress through JobStore, so it is safe to call repeatedly (on a timer, or from an external task queue).

Setup (standalone)

python -m venv .venv
./.venv/Scripts/pip install -r requirements.txt   # Windows
# .venv/bin/pip install -r requirements.txt        # macOS/Linux

cp .env.example .env
# fill in FAL_KEY in .env (https://fal.ai/dashboard/keys)

Run:

./.venv/Scripts/python -m app.main

Open http://127.0.0.1:8001 for the web UI. The REST API lives under /api/*, with interactive docs at /docs.

Connecting to fal.ai

The top of the web UI shows a live connection indicator:

• 🟢 green — the key works; shows the account's username and credit balance when available.
• 🟡 yellow — the key is valid but fal.ai returned an error reading extra info (see below).
• 🔴 red — no working key (missing/invalid, or a network error).

Note: the username/balance readout comes from fal.ai's GET /v1/account/billing endpoint, which the docs mark as requiring an admin-scoped key. A regular API key still trains LoRAs fine, but the status bar may show "connected" without a balance and a permission warning - that's expected, not a bug.

You can paste a different key straight into the UI ("Set API key" → "Save & connect") instead of editing .env; it's persisted to data/settings.json and takes precedence over FAL_KEY on the next restart. "Reconnect" just re-runs the same check against whichever key is currently active (useful after a transient network error, or to refresh the balance).

Usage

1. Drop images into input/<set_name>/ (PNG/JPG/JPEG/WebP), optionally with a same-stem image.txt caption next to each image.png. Or upload files straight from the web UI ("Image set" section) — it creates/fills the folder under input/ for you.
2. Enter a trigger phrase and parameters (auto_captioning, steps, learning_rate, resolution) and start the training job.
3. The background poller zips the images, uploads them to fal.ai, polls the job status by request_id, and once it's done, downloads the .safetensors file and converts it.
4. Results land in outputs/<trigger_phrase>/:
   • <name>.safetensors — as returned by fal.ai (for diffusers-style loaders);
   • <name>_converted.safetensors — with renamed keys for ComfyUI.

REST API

Method	Path	Description
GET	/api/fal/status	Connection indicator: verifies the active key, returns username + credit balance
POST	/api/fal/reconnect	Re-run the same check (no body)
POST	/api/fal/key	Set/replace the active fal.ai API key ({"api_key": "..."}), persisted for standalone mode
GET	/api/input-sets	List folders under input/ with image counts
POST	/api/input-sets/{name}/images	Upload files (multipart) into input/{name}/
POST	/api/trainings	Create a training job
GET	/api/trainings	List jobs
GET	/api/trainings/{id}	Get one job's status
POST	/api/trainings/{id}/cancel	Cancel a job (and on fal.ai's side, if not finished yet)
POST	/api/trainings/{id}/advance	Manually step the pipeline once (when nothing is running the background poller)
GET	/api/trainings/{id}/download?variant=raw|converted	Download a .safetensors file

POST /api/trainings body:

{
  "input_set": "my_character",
  "trigger_phrase": "ohwx_character",
  "auto_captioning": "Object/Character",
  "steps": 100,
  "learning_rate": 0.0005,
  "resolution": 768
}

Integrating into a monolith

Don't use app.api.app.create_app — build a TrainingService with your own JobStore/Storage implementations (e.g. backed by your DB and object storage) and mount the router:

from app.api.routes import build_router
from app.core.training_service import TrainingService

service = TrainingService(my_job_store, my_storage, my_fal_gateway, ...)
monolith_app.include_router(build_router(service), prefix="/krea2-trainer")

The monolith then decides what advances the pipeline — a periodic task calling service.advance_all(), a fal.ai webhook handler (pass fal_webhook to submit, then call _submit/_poll-equivalent logic when the webhook fires instead of polling), or a manual service.advance(job_id) call from a task queue.

Tests

./.venv/Scripts/pip install -r requirements-dev.txt
./.venv/Scripts/python -m pytest

Tests cover the converter (app/core/converter.py, against synthetic .safetensors files), the filesystem storage (app/core/storage.py), and the full job lifecycle in TrainingService with a fake FalGateway (no real calls to fal.ai).

Known limitations / follow-ups

• JsonFileJobStore is sized for standalone use (human/poll-rate writes), not high write concurrency — swap in a DB-backed JobStore for a monolith.
• Cancellation on fal.ai's side is best-effort: if the request is already running, cancelling may not stop the actual computation on fal.ai, but the job is still marked cancelled locally and stops being advanced.
• fal.ai's auto_captioning="Custom" supports a custom captioning prompt; the current UI/schema doesn't expose that field separately — add auto_captioning_prompt to TrainingParams.to_fal_arguments if you need it.

License

MIT - see LICENSE. app/vendor/krea_lora_convert_helpers.py is vendored from the separately MIT-licensed ComfyUI-Pixaroma project; its original copyright notice is kept intact in THIRD_PARTY_NOTICES.md.