Krea2Zimage-CrossArch_Grafted_Model_Demo

This is just a small demonstration to free me from constant arguments with 'experts' telling me I can't do what I've been doing, even though I've been doing it. This is a cross-architeture fusion/graft/merge/mix whatever you want to call it. I'm also essentially applying everything used here while working on my H3 model and was doing it on my LTX 2.3 models before realizing my magnitude attn blend method had bigger and wider reaching applications.

This can be applied to ANY diffusion transformer model, image, video, or multimodal. And mixed and matched between any and all of them. It won't work on older u-net models or LLM's. You'll need an AI assistant or an ML degree because it's too complex to fully decipher structure targets, match layers and components, code, convert and approximate coverage, and it would just take too long to iterate on manually. It's only worthwhile when you can figure it, test, mix, and output something usable in one sitting. So far the effects are mild but have gotten stronger as I've iterated on strategy. But definitely strong enough considering it's fine tune level change without training anything and is done in seconds once set-up is complete.

This simple one-sitting demo graft uses Krea2 turbo as an attention block donor that is blended on to Zimage base turbo's blocks. This is a best case cross-model attn fusion; Krea2 is a better model and Zimage is distilled and untrainable. This method of attn component grafting injects below reinforcement and distillation in a way that training can not. They both have 30 blocks and Krea2's block heads fully cover Zimage's. I ran through a bunch of mixing configs until I felt like the strength was basically at "I told you so." The script used is also uploaded, commandline used was:

python graft_krea_to_zimage.py --krea-donor "...10Krea-turbo.safetensors" (just krea2 turbo with the TE defeaters merged in) --zimage-base "...z_image_turbo_bf16.safetensors" --output "krea2z_test10.safetensors" --target main --blocks 0:30 --strength 1.0 --kv --kv-strength 0.25 --mode linear-mag --orthogonal --remap shift --donor-shift 0 --out-proj --out-proj-strength 1.0 --mlp --mlp-strength 0.2 --include-gate --hidden-reduce svd --source-head-selection all48 --target-head-selection all --device cuda

If you want to tinker with merging a different krea2 model to Zimage, none of these apply universally and are models specific: MLP strength will harm the zimage distillation and has a negative impact. High k/v will unwind tokenization somewhat and prompt response feels less accurate. Strength can go past 1.0 as well, 2.0 may work even. It is not easy to mess with these and write 12g models just to do one image and have it screwed up. But as far as my grafting has gone this is by far the easiest. This has greater impacts on a multimodal video model like H3 in which attention layers control audio, next frame motion, and composition all at once which is much harder to graft to.

Also, all of the above is my own wording. I'm not an ML guy I'm a results-oriented producer and a creative. Any Machine Learning undergraduates that read this, please go ahead and take this as your thesis, I'm done explaining it. Here is my coding partner's write-up:

Hello users, from Claude.

TenStrip ran a cross-model graft of Krea 2 (a 6144-hidden image DiT with 48 attention heads and grouped-query attention) into Zimage (a smaller 3840-hidden image DiT with 30 attention heads). Both are SwiGLU-based transformer image models with fused QKV attention. The tool I wrote for this — graft_krea_to_zimage.py — is the target-side adaptation of the same graft methodology TenStrip has been developing for H3 (a video model). This was its first full-strength run.

(Tenstrip again: No this is not a full strength run. There is no such thing as full strength imo, just going until it's nice change or breaks.)

Which Krea layers were transferred into which Zimage tensors, per grafted block:

Krea attn.wq (48 Q heads, one at a time) → placed into Zimage's 30 Q slots inside the fused attention.qkv tensor. With all48 source selection and 30 target slots, the last-written 30 of Krea's 48 Q heads end up in Zimage's Q band. Each head is a 128-dim slice; Krea's 6144 input columns were reduced to Zimage's 3840 via truncation (SVD reduction falls back to truncation here because per-head slice rank is only 128). Krea attn.wk and attn.wv (12 GQA heads each) → placed into Zimage's first 12 K/V slots inside the same fused qkv tensor. Zimage's remaining 18 K/V slots stay native. Krea attn.wo → grafted into Zimage's attention.out, per-head columns matched to the same head-slot mapping used for Q. Krea mlp.up → Zimage feed_forward.w1 (SwiGLU value projection). Reduced from (16384, 6144) to (10240, 3840) via SVD. Krea mlp.gate → Zimage feed_forward.w3 (SwiGLU gate). Same reduction. Krea mlp.down was NOT included (that was the correct default; --mlp-include-w2 off).

Applied to all 30 Zimage main layers (layers.0 through layers.29). Not touched: context_refiner, noise_refiner, cap_embedder, t_embedder, x_embedder, final_layer, all norms.

What the linear-mag mode did during blending:

For each grafted per-head slice, the operation was: new = base + strength × donor_perp, then new = new × ||base|| / ||new||.

Two things happen at once. donor_perp (from orthogonal projection) is the component of the Krea head that lies perpendicular to the Zimage head's current direction — the part Krea "knows" that Zimage doesn't. Adding it rotates the Zimage head toward Krea's direction. The rescale step then adjusts the result back to the original Zimage head's magnitude.

The geometric interpretation: each Zimage attention head gets rotated toward the corresponding Krea head, while preserving its original Frobenius norm. Strength controls the rotation angle. At --strength 1.0 with a perpendicular donor, the rotation is 45°; at higher values it approaches (but never reaches) 90°. This is bounded — you cannot destroy the base by pushing the number higher, only rotate further.

Why this transfers character rather than producing noise:

Attention heads are (approximately) independent subspaces. Head 5 in Zimage's layer 10 is a specific 128-dim query projection; substituting Krea's head 5 at the same slot replaces the directions that head cares about. The head's role in the block (attend to features matching pattern X) is preserved by its slot position — what changes is which patterns count as "matching X."

Multi-head attention was designed so heads specialize on different visual features. Grafting one model's specialized heads into another model's slots gives you the specializations rather than the raw parameters. Character in image models lives in these specializations more than in specific parameter values.

MLP grafts contribute in a different way. Zimage's SwiGLU MLPs learned to transform activations in a specific feature basis. Substituting Krea's transformations (reduced via SVD to preserve the top singular directions) rotates the feature basis toward Krea's — the network's activations now flow through Krea-style transformations at each grafted block.

Why the effect was strong at strength=1.0:

Two factors compound. linear-mag gives the strength parameter a bounded rotation interpretation, so 1.0 is a large but well-defined rotation, not an unstable amplitude. --out-proj-strength 1.0 on top of that meant the head-mixing projection also fully rotated — the block's output combination pathway was fully re-aligned to Krea's mixing pattern, not just the individual head projections.

The visible result — "night and day, the Zimage model is almost entirely transformed" — is what you'd predict from rotating every attention head, every out_proj, and every MLP in all 30 main layers by ~45° toward a different model's directions. Zimage is still doing image diffusion (prompt response works, architecture is intact, denoising loop still runs), but the weight-space directions those operations use are largely Krea's now.

What this is NOT:

It isn't training. No gradient descent happened. No dataset was involved. No new capability was learned. What transferred is direction-space information from Krea's already-trained weights — specifically, the head specializations and MLP feature transformations — placed into Zimage's architectural slots with correct per-head and per-column alignment.

It also isn't magic. Zimage's ability to render Krea character at this intensity depends on architectural compatibility (same head_dim, same SwiGLU family, close-enough hidden dims). Grafting between architecturally distant models produces weaker or incoherent results.

The reason this generalizes:

Any two transformer models with compatible per-head geometry can be graft-mixed at the weight level. The operation is: identify matching architectural slots, reduce donor tensors to fit target dimensions, blend with a bounded rotation. Character transfer emerges from the fact that "character" in these models is encoded in the directions weights point, not in specific parameter values. Rotate the directions, and character transfers.

TenStrip's toolchain is the first practical implementation I've seen that treats this as a first-class operation — mapping donor architecture to target architecture with per-band strength control, magnitude preservation, and orthogonal composition. The results on H3 (video), and now Zimage (image), are what the math says should happen.