21 Commits
Author SHA1 Message Date
CthomasdesignandClaude Opus 4.8 e9e8b52d6f fix: use cached base model without re-verifying against HF (Xet 403)
krea/Krea-2-Turbo moved to HuggingFace Xet storage, whose CDN 403s public
downloads. On every model load the pipeline re-verified the base assets
(VAE + Qwen3-VL encoder + tokenizer) via a HEAD size-check; the Xet
redirect returns no usable content-length, so a fully-cached file looked
stale and triggered a re-download that 403'd — breaking the model loader
even though the files were present.

_base_dir() now short-circuits to the cache when the assets are all
present (no network at all), and _http_download trusts a cached file when
the size is unknowable. Verified: load with base_dir=None (the ComfyUI
path) resolves the cache in 0.0s and the pipeline loads with no 403. v0.7.3.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-13 19:27:20 -04:00
CthomasdesignandClaude Opus 4.8 8881d961c0 fix: model loader honors extra_model_paths (shared model stores)
The Krea2 model loader listed builds by raw os.listdir on the install
ComfyUI/models/krea2 only, so a transformer in an extra_model_paths
location (e.g. a Comfy Desktop shared store) was invisible — the loader
showed no model even though the LoRA node (which uses folder_paths)
found its files there. Register the 'krea2' folder and resolve builds
through folder_paths.get_folder_paths / get_full_path so every registered
search path is checked. v0.7.2.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 17:36:44 -04:00
CthomasdesignandClaude Opus 4.8 2a7bb53800 tools: discover weights in the ComfyUI-Shared model store
Weights now live in the Comfy Desktop shared store
(~/ComfyUI-Shared/models/{krea2,loras}) rather than a project-local Models/
folder bridged by per-file symlinks. The test harness searches there first,
falls back to the legacy Models/ location, and honors the KREA2_TEST_* env
overrides.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 17:30:46 -04:00
CthomasdesignandClaude Opus 4.8 c9f8f0ace1 docs: README overhaul (architecture, editing guide, env vars, tools)
Architecture overview (the MLX island), editing guide (grounding_px
semantics, aspect-ratio rule, two-ref order and caveats, Turbo vs
Raw-removals limitation), example-workflow index, environment-variable
reference, tools/harness section.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 17:16:08 -04:00
CthomasdesignandClaude Fable 5 4d8db6bdf8 chore: declare macOS-only via classifiers; v0.7.1
The registry/Manager read OS constraints from [project] classifiers; this
node is Apple-Silicon-only (MLX/Metal).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 13:59:11 -04:00
CthomasdesignandClaude Fable 5 5fcf343a4d refactor(tools): machine-portable test scripts, shared asset helpers
_shared gains default_model()/default_lora() (env overrides
KREA2_TEST_MODEL/KREA2_TEST_LORA, else <repo>/../Models discovery with
clear errors) and fox_source()/ref_uint8() for the reference source image
boilerplate repeated across scripts. test_img2img/test_edit/
test_edit_grounded no longer hardcode machine paths; all tools compile
and resolve on a fresh checkout.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 13:59:11 -04:00
CthomasdesignandClaude Fable 5 1fc866c144 refactor(nodes): shared _run helper + widget factories, DESCRIPTION + tooltips
Extract the triplicated LoRA/progress/Cancel/safety/tensor plumbing into
one _run() helper; common widgets come from factories returning fresh
specs. Every node gains a DESCRIPTION and every widget a tooltip
(documented ComfyUI features). Module docstring now lists all 6 nodes.

Workflow-compatibility guarded: INPUT_TYPES key order and widget types
verified structurally identical to a pre-refactor snapshot (ComfyUI
stores widget values positionally).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 13:59:11 -04:00
CthomasdesignandClaude Fable 5 b4ebb360e3 refactor(engine): tidy grounded additions — no behavior change
text_encoder: hoist os/json imports, drop redundant inner glob import,
remove unused VISION_START_ID, single VISION_BLOCK constant feeding both
the template and encode_grounded, grounded constants grouped at top,
module docstring covers both conditioning paths, type hints.
sampling: shared _decode() tail for sample/sample_edit; sample_edit gains
minres/maxres params (same defaults) instead of hardcoded literals.
pipeline: real OrderedDict annotation, explicit grounding_px validation.

Verified: full parity bit-exact (all 5 cases incl. LoRA), grounded
regression vs the saved torch golden numerically identical
(cos 1.00012243), bench step~ unchanged at 7.79s.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 13:59:11 -04:00
CthomasdesignandClaude Opus 4.8 3ec843606c feat: two-reference grounded edits + preprocess robustness fix
Generalize the grounded encode to N reference images (scene + subject):
both go through the vision tower with a per-image <vision> block, spliced
and deepstack-injected at their own token positions, with multi-image
get_rope_index. Validated vs full torch Qwen3-VL — two-image grounded tap
matches to cosine 1.0, mean|Δ| 3.1e-5 (single-image unchanged). generate_edit
now feeds both image + image_b through grounding when grounding_px > 0.

Fix a ZeroDivisionError in preprocess: max_pixels is now clamped so a tiny
grounding_px (e.g. from a stale workflow whose widgets predate the
grounding_px control) can't zero the area ceiling. The example edit
workflow gains the grounding_px widget so its values align with the node.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 20:37:02 -04:00
CthomasdesignandClaude Opus 4.8 af579ffb6a feat: image-grounded edit conditioning via Qwen3-VL vision tower (v0.7.0)
Completes rung B — the semantic path of krea2_edit. Krea2Edit gains a
grounding_px control: >0 runs the source image through the (now wired-up)
Qwen3-VL vision tower and reads the instruction *while looking at the
image*, matching what the identity LoRA was trained with; 0 keeps plain
text conditioning.

Integration (krea2_engine/text_encoder.py): splice merged vision tokens
at <image_pad>, inject the 3 deepstack features at decoder layers 0/1/2,
apply interleaved multimodal RoPE (3D t/h/w positions, mrope_section
24/20/20), and tap the same 12 layers. Torch-free at runtime; the vision
tower is loaded lazily on first grounded use.

Validated against the full PyTorch Qwen3-VL model end-to-end: the grounded
12-layer tap matches to cosine 1.0, mean|Δ| 2.7e-5 (a deepstack recording-
order bug found and fixed via the per-layer diff). txt2img stays
pixel-identical to the parity reference; grounded edit runs end-to-end and
preserves the source scene better than the plain path (see tools/).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 19:58:43 -04:00
CthomasdesignandClaude Opus 4.8 1db2179abe feat(wip): vendor Qwen3-VL vision tower (MLX) + torch-free preprocessing
Groundwork for image-grounded conditioning (rung B). Vendors mlx-vlm's
pure-MLX Qwen3-VL vision tower into krea2_engine/vision/ (self-contained:
inlined SDPA helper, local VisionConfig) and adds torch-free image
preprocessing (smart-resize + CLIP normalize + patch flatten).

Both validated against the PyTorch reference (tools/test_vision_tower.py):
- tower vs transformers Qwen3VLVisionModel: merged embeds + all 3 deepstack
  features match to max|Δ|~2e-4, cosine 1.0 (bf16 rounding).
- preprocessing vs transformers image processor: max|Δ|~2e-7, grid exact.

Loads the visual.* weights already present in the Krea-2-Turbo text
encoder (315 tensors, 0 missing). Not yet wired into any node — the
grounded-encode integration (M-RoPE, deepstack injection, 12-layer tap)
is the next step. mlx-vlm is MIT; attribution added to NOTICE.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 19:31:17 -04:00
CthomasdesignandClaude Opus 4.8 984aabc0d1 feat: Krea2Edit node — in-context image editing on Turbo/MLX (v0.6.0)
New additive node for identity-preserving edits: the clean VAE-encoded
source rides the sequence as a frame≠0 reference token block while a
fresh-noise target is denoised, and only the target is decoded — the
appearance path from lbouaraba/comfyui-krea2edit, mapped onto our engine
(sequence [text | src(frame=1) | target(frame=0)], mirroring the 3-axis
RoPE frame scheme that was already present but unused).

Engine additions are isolated: transformer gains prepare_edit +
denoise_step_edit, sampling gains sample_edit + build_edit_positions,
pipeline gains generate_edit; the txt2img/img2img sample() path is
untouched (verified txt2img stays pixel-identical to the parity ref).

Pairs with the krea2_edit identity LoRA (conradlocke/krea2-identity-edit),
which loads through the existing Krea2LoRA node unchanged (ai-toolkit
reference naming, 28 blocks, 0 unmatched keys). Validated: source fox +
"red scarf" keeps the fox's identity while applying the edit; the LoRA
sharpens edit coherence markedly (see tools/test_edit.py). Second
reference image supported for multi-ref edit LoRAs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 18:52:53 -04:00
CthomasdesignandClaude Opus 4.8 0f46e19cbc feat: Krea2Img2Img node — image-to-image on Turbo/MLX (v0.5.0)
New additive node (existing four untouched): VAE-encodes a source image
and denoises from a noised version of it. Rectified-flow interpolant
x_t=(1-t)·z0+t·noise, entering the schedule at t=denoise; denoise=1.0
takes the exact txt2img code path (verified pixel-identical to the
parity reference — the hard guard that generation is undisturbed).

Engine changes are additive: sample() gains init_latent/strength
(strength>=1.0 is a byte-identical no-op), pipeline gains
generate_img2img() reusing the prompt cache + LoRA path, nodes.py gains
_from_image_tensor. VAE encode input range verified [-1,1] (round-trip
45 dB). Denoise sweep 0.3→0.9 is monotonic and coherent; contact sheet
and strength=1.0 guard in tools/test_img2img.py.

First rung of the image-conditioning ladder (img2img → in-context edit
→ grounded vision); reuses the VAE-encode plumbing the edit path needs.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-08 16:47:06 -04:00
CthomasdesignandClaude Fable 5 98cb8b5380 docs: v0.4.0 — performance section, NOTICE amendment, version bump
README documents the measured speedups (512² ~35% / 1024² ~15% on
M1 Max), the prompt-embedding cache, and the KREA2_EXACT_LEGACY seed-
reproducibility flag. NOTICE records the performance modifications to
the vendored engine, as its attribution chain previously stated the
engine was otherwise unmodified.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
v0.4.0
2026-07-08 06:27:33 -04:00
CthomasdesignandClaude Fable 5 151033af99 engine: prompt-embedding cache + encode-once for batches
Krea2Pipeline now routes generate() through _encode_cached: an 8-entry
LRU keyed by prompt string. Seed sweeps with a fixed prompt (the dominant
ComfyUI iteration pattern) skip the text encoder entirely, and
num_images>1 encodes its prompt once instead of once per image.

Cache hits return the same arrays a fresh encode produced, so this is
bit-exact on the single-image path: verified repeat generate() calls are
pixel-identical and match the engine-path parity reference exactly; quick
parity (incl. LoRA) stays bit-exact. Batch img0 was never seed-compatible
with a single-image run (verified pre-existing on the legacy path).
LoRA-safe (LoRAs wrap only the DiT). KREA2_EXACT_LEGACY=1 bypasses.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 23:12:17 -04:00
CthomasdesignandClaude Fable 5 bfc60c11c5 engine: dynamic-length text encoding (541 -> ~prompt-length tokens)
Pad prompts only to the longest in the batch instead of always to 541
tokens; the 5 suffix tokens keep their absolute RoPE positions from the
padded layout (positions are index-assigned and only trailing pad columns
are dropped, so every real token's rotation is unchanged — verified:
embeddings layer bit-exact, divergence enters only via kernel
reassociation at ~1e-6 f32 and grows smoothly, no structural jump).
Encoder work drops ~12x for short prompts. KREA2_EXACT_LEGACY=1 restores
always-max padding.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 22:28:15 -04:00
CthomasdesignandClaude Fable 5 bf7b9cb7cf engine: trim padded text tokens before the DiT (~30% fewer tokens at 512²)
Prompts are padded to 512 text tokens; ~470 of them are dead padding that
previously rode through all 28 attention blocks on every denoising step
and forced a dense (B,1,L,L) additive mask. Gather the union of valid
columns instead (layout is [prompt, padding, suffix], so this is a gather,
not a slice); single-prompt batches then run fully mask-free.

Mathematically exact (f32 check: velocity diff collapses to ~3e-5,
1-cos = 0.0), but the shorter sequence tiles bf16 kernels differently, so
fixed seeds render equal-quality, non-bit-identical images vs 0.3.x —
side-by-side reviewed and accepted. KREA2_EXACT_LEGACY=1 restores the
padded path for pre-upgrade seed reproducibility.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 22:08:01 -04:00
CthomasdesignandClaude Fable 5 98ebf5b0bf engine: hoist step-invariant conditioning out of the denoise loop
Split SingleStreamDiT.__call__ into prepare_conditioning() (text fusion,
RoPE tables, attention masks — identical every step) and denoise_step();
sample() now prepares once and steps 8 times instead of recomputing the
whole text path per step. Also eliminates the two .item() GPU syncs that
_additive_mask forced on every step. __call__ still composes both, so the
external API is unchanged. Parity: bit-exact incl. LoRA case.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 21:49:00 -04:00
CthomasdesignandClaude Fable 5 80d0bf8f71 engine: use native GQA in fused attention instead of repeating K/V heads
mx.fast.scaled_dot_product_attention handles grouped KV heads (48Q/12KV)
natively; passing them directly is bit-identical to the mx.repeat
workaround (parity: max|d|=0.0 incl. LoRA case) while skipping the 4x
K/V materialization in both the DiT and the Qwen3 text encoder.
Wall-clock neutral at 512² (quantized matmuls dominate); saves memory
traffic and peak footprint, which matters on 32 GB machines.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 21:41:50 -04:00
CthomasdesignandClaude Fable 5 59deb1e3e5 tools: benchmark + output-parity harness for engine optimization work
bench.py attributes generation time to encode / per-step / decode stages
via the pipeline's own step_callback; parity.py snapshots decoded float32
pixels for a fixed prompt/seed/size (+LoRA) matrix and asserts later runs
stay uint8-identical — the quality gate for every optimization commit.
Reference outputs and results land in tools/reference/ (gitignored).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 21:07:07 -04:00
CthomasdesignandClaude Sonnet 5 2ff0827d08 Initial commit: self-contained Krea-2-Turbo MLX nodes for ComfyUI
Krea2ModelLoader / Krea2LoRA / Krea2Generate / Krea2Unload nodes running
Krea-2-Turbo on Apple MLX in-process, returning ComfyUI IMAGE tensors.
The krea2_engine/ pipeline is vendored directly in this repo (no external
git dependency) — see NOTICE for the full attribution chain back to the
original Krea 2 model and the MLX port it builds on.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-07 20:32:37 -04:00