LayerStream: Krea2 Chunking

Here's a failure mode that confuses people more than a clean OOM: your Krea 2 render doesn't crash, it just never finishes. The VRAM readout shows a healthy 24 GB in use, the progress bar creeps, and you wonder whether the sampler is stuck.

That's the case LayerStream: Krea2 Chunking was built for. In the author's production report - a 24.5 GB bf16 Krea 2 on a 24 GB card, no special launch flags - the workflow with this node finished in about a minute. Without it, it "runs forever". Same VRAM reading both times. The difference is where the memory goes.

Weights stream; activations don't

ComfyUI's Dynamic VRAM streams weights layer by layer, so a model larger than your card isn't automatically a dead end. But it fills the card with weights while it does that, and activations have nowhere left to live - so they spill into shared system memory, which is orders of magnitude slower. Nothing errors. Everything just gets glacial.

Activations scale with token count, and Krea 2's is generous: at 4 MP, roughly 16k tokens through a hidden width of 6144, with SwiGLU MLPs at 4× expansion and 48-head K/V. It's a deliberately conventional 12B single-stream DiT, which is exactly why this works on it.

What the node does to the block

Same two-pronged approach as the H3 sibling. Anything token-wise - the pre/post norms, modulation, the SwiGLU MLP, the projections - runs on slices and accumulates into the residual in place.

Attention runs in two passes. Pass one computes only K and V for every token, GQA-expanded the way the stock block does it. Pass two recomputes q and the attention gate per slice, attends against the full K/V, projects, and accumulates. The final layer holds an fp32 copy of the whole sequence, so it runs stock first and only slices when it runs out of memory.

It's honest about the edge cases: the block falls back to the stock implementation for reference-image editing (the timestep_zero_index path), for attention masks, and when another node has patched attention. No savings there, but no surprises either.

Be clear about exactness

This is where it differs from the H3 node, and it matters if you pin seeds. With a single chunk - bigger than the sequence - the output is bit-identical to stock. With several chunks on a bf16 model, the matmuls round differently at different slice sizes, so you land around 38.5 dB PSNR at 4 MP against a stock render. Visually the same image. Numerically not the same file. Whether int8 or fp8 Krea 2 files behave like H3 (exact) or like bf16 is still untested.

So: treat it as "numerically equivalent", and don't file a bug over a few changed pixels at eight chunks.

The inputs

model in, model out. Place it after the LoRA loader, before the sampler/guider - it's a model patch, so it has to sit in the model chain.

• mlp_chunk_tokens - default 4096, tokens per MLP/norm slice. Smaller is lower peak VRAM, slightly slower.
• attn_chunk_tokens - default 4096, query tokens per attention slice. 0 disables the two-pass attention and gives you stock attention, leaving only the MLP chunked.

Install

One pack, two nodes - install osiworx-pack and you get both (plus Megapixel Size):

cd ComfyUI/custom_nodes
git clone https://github.com/osi1880vr/osiworx-pack

Or search "osiworx-pack" in ComfyUI Manager. No pip dependencies; it's pure Python against torch and einops. Restart and look under osiworx/layer. The node only registers if your ComfyUI ships comfy.ldm.krea2; developed against 0.39.0.

When you actually need it

Not at 1 MP. The author measured Krea 2 at 1024×1024 under hard caps of 4 GB, 2 GB and even 1 GB with no node at all - weight streaming alone handles that, and the images came out identical. At 1 MP this node is dead weight in your graph.

At 4 MP it buys roughly 1 GB of lower floor: stock failed at a 2 GB cap (an RMSNorm fp32 copy was what blew), the node held. A real but modest saving, and you pay between 2.7× and 4× in wall time for it, because layer weights get fetched again for every slice.

Two traps worth knowing

You need tiled VAE decode. At 4 MP the plain VAEDecode ran out of memory at a 3 GB cap even though sampling had finished cleanly. Swap in VAEDecodeTiled (tile 512) or you'll blame the wrong node.

Use the largest chunk that fits. Dropping to 256 tokens doesn't buy you a smaller card, just a longer wait. And if the node never appears in the list, check your console for osiworx-pack - the pack skips modules it can't import rather than breaking your ComfyUI, which is nice right up until you can't find the node.