Cthomasdesign-ComfyUI-Krea2-MLX/README.md

ComfyUI-Krea2-MLX

Run Krea-2-Turbo locally on Apple Silicon inside ComfyUI, using Apple's MLX framework — text-to-image, img2img, and identity-preserving image editing with dual conditioning (in-context reference frames + Qwen3-VL vision grounding). Quantized to fit a 32 GB Mac; supports stackable LoRAs (official Krea / diffusers and ai-toolkit / PEFT formats, including the krea2_edit identity LoRA).

macOS on Apple Silicon (M-series) only. The MLX dependencies do not run on Windows/Linux. ~24 GB unified memory recommended for 1024²; editing with grounding wants ~32 GB.

Architecture

The pure-MLX pipeline (krea2_engine/) is vendored directly in this repo — no external git dependency — and runs in-process, returning standard ComfyUI IMAGE tensors:

Krea2 Model Loader ──KREA2_PIPE──┬── Krea2 Generate ──┐
Krea2 LoRA ──KREA2_LORASTACK─────┼── Krea2 Img2Img ───┼──IMAGE──► any ComfyUI node
LoadImage ──IMAGE────────────────┴── Krea2 Edit ──────┘   (save, upscale, …)

Inside the pipe, everything is MLX on Metal: the Qwen3-VL-4B text encoder (with its vision tower for grounded editing), the 12.9B Krea-2 MMDiT transformer, the flow-matching sampler, and the Qwen-Image VAE. No PyTorch in the generation path. It's a self-contained island: it does not expose a torch MODEL, so it does not plug into ComfyUI's native KSampler/ControlNet — by design.

Install

Via ComfyUI-Manager (recommended): search for Krea2-MLX and install; Manager runs requirements.txt, which installs the vendored pipeline's MLX deps into ComfyUI's Python.

Manual:

cd ComfyUI/custom_nodes
git clone https://github.com/cthomasdesign/ComfyUI-Krea2-MLX
<ComfyUI-python> -m pip install -r ComfyUI-Krea2-MLX/requirements.txt

Restart ComfyUI.

Models

Put a transformer build in ComfyUI/models/krea2/ — the Krea2 Model Loader lists what's there and infers precision from the filename (mixed_4_8 / 8bit / bf16):

Build	Size	Repo → file
mixed-4/8	9.8 GB	avlp12/Krea-2-Turbo-Alis-MLX-mixed-4-8 → transformer_mixed_4_8.safetensors
8-bit	14 GB	avlp12/Krea-2-Turbo-Alis-MLX-8bit → transformer_8bit.safetensors

On first generation the VAE + Qwen3-VL-4B text encoder (~9 GB) download automatically from krea/Krea-2-Turbo (cached in ~/.cache/krea2_alis_mlx). Put LoRAs in ComfyUI/models/loras/ — e.g. conradlocke/krea2-identity-edit for editing.

Nodes (category Krea2 MLX)

Node	Purpose
Krea2 Model Loader (MLX)	Pick a build from models/krea2 → KREA2_PIPE (cached; stays resident).
Krea2 LoRA (MLX)	Pick a LoRA from models/loras + strength → KREA2_LORASTACK. Chain several to stack.
Krea2 Generate (MLX)	Text→image. Prompt/size/steps/seed (+ LoRA stack) → IMAGE. Progress bar + Cancel supported; NSFW safety_filter on by default.
Krea2 Img2Img (MLX)	Source image + prompt + denoise → IMAGE. Reinterprets the whole frame: denoise 0→1 goes from near-copy to full reinterpretation (1.0 == plain text→image).
Krea2 Edit (MLX)	Source image + instruction (+ optional 2nd reference, + grounding_px) → IMAGE. Identity-preserving edit: the source rides as a clean reference frame while a fresh target is generated.
Krea2 Unload (MLX)	Frees the cached model and clears the Metal cache to reclaim memory.

Ready-made graphs in example_workflows/: text-to-image (Krea2-MLX-example.json), img2img (Krea2-MLX-img2img.json), identity edit with the LoRA chained in (Krea2-MLX-edit.json).

Editing guide (Krea2 Edit)

The edit node implements the dual conditioning the krea2_edit identity LoRA was trained with: the source enters (1) as clean VAE reference tokens the transformer attends to while generating, and (2) through the Qwen3-VL vision tower, so the instruction is read while looking at the image. Stack the identity LoRA at strength 1.0 for best results.

• Match the output aspect ratio to the source — training pairs are same-size; mismatches degrade preservation.
• grounding_px (default 768): lower = stronger edit adherence and more uniform scene changes; higher = stronger identity/likeness. 0 = plain-text conditioning (no vision tower).
• Two references (image_b): training order is scene first, subject second. Two-person inputs tend to drift faces toward each other — for best identity, chain single-reference edits instead.
• What works on Turbo (this island): add, recolor, restyle, re-stage — at 8 steps, no CFG. Large removals/deletions were tuned for Krea-2-Raw at 20 steps with CFG, which this Turbo-only island doesn't run; expect them to be weaker.
• Img2Img vs Edit: img2img reinterprets everything (style transfers, variations); Edit preserves the subject and applies the instruction (targeted changes).

Performance

The vendored engine is heavily optimized for Apple Silicon — measured on an M1 Max 32 GB (mixed-4/8 build, 8 steps):

	before (v0.3)	after (v0.4+)	faster
512², short prompt	~103s	~67s	~35%
512², long prompt	~104s	~72s	~31%
1024²	~310s	~263s	~15%

Repeat generations with the same prompt (seed sweeps) additionally skip the text encoder via a per-prompt embedding cache. What changed: native grouped-query attention; step-invariant conditioning hoisted out of the denoising loop; padded text tokens trimmed before the DiT (~30% of the sequence at 512²); dynamic-length text encoding; prompt-embedding cache. All validated mathematically faithful to the reference computation (see tools/).

Seed reproducibility: the sequence-length optimizations alter bf16 kernel rounding order, so a fixed seed renders an equal-quality but not pixel-identical image vs v0.3 — the same class of change as an MLX or hardware upgrade. Within a version, generation is fully deterministic (same seed → same image, bit-for-bit).

Environment variables

Variable	Effect
KREA2_EXACT_LEGACY=1	Bit-exact pre-0.4 outputs: restores padded encoding/attention, bypasses the prompt cache (slower). For reproducing old seeds.
KREA2_BASE_DIR	Local krea/Krea-2-Turbo snapshot (skips the first-run download).
KREA2_DISABLE_SAFETY=1	Globally disables the NSFW filter (the license requires filtering in deployments — see Notes).
KREA2_SAFETY_THRESHOLD	NSFW flag threshold (default 0.85).

Development: bench & parity harness (tools/)

tools/bench.py attributes generation time to encode / per-step / decode stages; tools/parity.py snapshots reference outputs and asserts later runs stay pixel-identical — the quality gate every engine change passes through. tools/test_*.py validate each feature (img2img guard, edit forward, vision tower vs the PyTorch reference, grounded encode, two-ref grounding).

Notes

• One 12.9B model (~18–20 GB with encoder/VAE; +~2.7 GB with the vision tower and identity LoRA) stays resident. On 32 GB, avoid loading large torch models simultaneously; use Krea2 Unload to free memory.
• The safety_filter runs the pure-MLX Falconsai/nsfw_image_detection classifier; the Krea 2 Community License (§4.2) requires reasonable content-filtering in deployments — keep it on there.
• Where required by law, disclose that outputs are AI-generated (License §4.3).

License & attribution

Independent, unofficial — not an official Krea product or endorsed by Krea. The model is a Derivative of krea/Krea-2-Turbo under the Krea 2 Community License (see LICENSE); the MLX port is by avlp12 (Alis MLX), with LoRA support added in krea2-mlx-studio. The engine (krea2_engine/) is vendored here from that lineage; the Qwen3-VL vision tower is vendored from mlx-vlm (MIT) — see NOTICE for the full attribution chain and what changed at each step. Model weights are downloaded by the user, not redistributed here. Node/engine code here: MIT (per upstream). NSFW classifier: Falconsai/nsfw_image_detection (Apache-2.0). By using the weights you agree to the Krea 2 Community License (commercial use requires < $1M annual revenue, or an enterprise license).