lora_layers_reference.md

LoRA Layer Inventory — Edit Anything checkpoints

Inventory of every tensor in two builds of the edit_anything_reference_v0.1_r128 LoRA.

Both builds share the same canonical basename (edit_anything_reference_v0.1_r128) and are distinguished by the extras suffix that scripts/split_editanything_lora.py appends to the output filenames:

• edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding.{standard,module}.safetensors — the original build. Only ships ref_adaln_proj + role_embedding.
• edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj.{standard,module}.safetensors — the continuation, fine-tuned with the video_to_video_ref_visual_adaln strategy. Adds the ref_attn LoRA branch and the ref_visual_proj projector on top of the original extras.

In the rest of this doc the two are referred to by their suffix only:

• ref_adaln_proj-role_embedding
• ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj

Rank is 128 in both (encoded in the LoRA tensor shapes; no alpha keys saved). Dtype is bfloat16 throughout. All LoRA modules cover 48 transformer blocks.

---

1. Summary

	ref_adaln_proj-role_embedding	ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj
Total tensors	965	1356
LoRA-target modules	10	14
LoRA tensors (A+B)	960	1344
Extra (non-LoRA) tensors	5	12
ref_attn LoRA branch	❌ absent	✅ trained on 48 blocks
ref_visual_proj (visual cross-attn projector)	❌ absent	✅ present (7 tensors)
ref_adaln_proj (global appearance AdaLN)	✅ (fc1 input dim 256)	✅ (fc1 input dim 768)
role_embedding	✅ shape (1, 128)	✅ shape (1, 128)

---

2. LoRA adapters

Each row = one target module type. Each entry = (lora_A.weight, lora_B.weight) duplicated across the 48 blocks of diffusion_model.transformer_blocks.*.

Module	ref_adaln_proj-role_embedding	ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj	Notes
attn1.to_q	✅	✅	self-attention query
attn1.to_k	✅	✅	self-attention key
attn1.to_v	✅	✅	self-attention value
attn1.to_out.0	✅	✅	self-attention output proj
attn2.to_q	✅	✅	cross-attention to text (Gemma)
attn2.to_k	✅	✅
attn2.to_v	✅	✅
attn2.to_out.0	✅	✅
ff.net.0.proj	✅	✅	feed-forward up-projection
ff.net.2	✅	✅	feed-forward down-projection
ref_attn.to_q	—	✅	new — visual reference cross-attention
ref_attn.to_k	—	✅	new
ref_attn.to_v	—	✅	new
ref_attn.to_out.0	—	✅	new

Key naming: diffusion_model.transformer_blocks.{0..47}.{module}.{lora_A|lora_B}.weight

Training freeze policy for the ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj build (per stage2_ref_visual_adaln_crossattn_from_v01_r128.yaml):

• attn1.* adapters loaded from the ref_adaln_proj-role_embedding build but frozen (trainable_include_patterns excludes them).
• attn2.*, ff.*, ref_attn.* are trainable.

---

3. Non-LoRA modules (the module sidecar)

These tensors live at the top of the state dict (no transformer_blocks.* prefix) and are consumed by the custom inference path (LTXVEditAnythingModuleLoader + LTXVEditAnythingLoopingSampler), not by the standard ComfyUI LoRA loader.

3.1. role_embedding — appearance role bias

Key	Shape	Notes
role_embedding.embedding.weight	(1, 128)	1 slot (appearance). Padded to (3, 128) at inference; entry stored at slot 1 (ref_img role).

Present in both builds with the same shape. In the ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj build it is frozen (use_visual_ref_role_embedding: false); wandb shows its norm stays flat at ~0.125 throughout training.

3.2. ref_adaln_proj — global AdaLN appearance anchor

Two-layer MLP that pools the reference latent into a vector added to every block's AdaLN timestep bias.

Key	ref_adaln_proj-role_embedding shape	ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj shape
ref_adaln_proj.fc1.weight	(512, 256)	(512, 768)
ref_adaln_proj.fc1.bias	(512,)	(512,)
ref_adaln_proj.proj.weight	(36864, 512)	(36864, 512)
ref_adaln_proj.proj.bias	(36864,)	(36864,)

⚠️ Shape mismatch on fc1.weight. The ref_adaln_proj-role_embedding build was trained with a 2-scale pool (avg_1x1 ‖ max_1x1 → 256-dim input). The ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj build was trained with a 3-scale pool (avg_1x1 ‖ avg_2x2 ‖ max_1x1 → 768-dim). Because of this incompatibility the trainer reinitializes ref_adaln_proj from scratch when continuing from ref_adaln_proj-role_embedding; the AdaLN projector in the continuation is not a fine-tune of the original one. The output dim 36864 = AdaLN param count for the LTX-2 transformer (read at runtime via preprocessor.adaln.linear.out_features).

3.3. ref_visual_proj — visual cross-attention memory tokens

Present in ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj only. SafeVisualRefProjector (training file video_to_video_ref_visual.py). Produces 32 visual memory tokens consumed by the new ref_attn branch.

Key	Shape	Notes
ref_visual_proj.fc1.weight	(1024, 384)	input 384 = 128 (local pooled) + 128 (global mean) + 128 (global std)
ref_visual_proj.fc1.bias	(1024,)	xavier init gain 0.1
ref_visual_proj.proj.weight	(4096, 1024)	maps to context_dim 4096; xavier init gain 0.05
ref_visual_proj.proj.bias	(4096,)
ref_visual_proj.norm.weight	(4096,)	LayerNorm γ
ref_visual_proj.norm.bias	(4096,)	LayerNorm β
ref_visual_proj.pos_embed	(1, 32, 4096)	per-token learned positional bias

Forward (matches SafeVisualRefProjector.forward):

tokens = local ‖ global_mean ‖ global_std          # [B, 32, 384]
tokens = proj(silu(fc1(tokens)))                   # → [B, 32, 4096]
tokens = LayerNorm(tokens)
tokens = tokens + pos_embed[:, :32]
return tokens * token_scale                        # training default 0.25

Not present in ref_adaln_proj-role_embedding — this entire branch is new.

---

4. Total tensor counts (sanity check)

ref_adaln_proj-role_embedding

LoRA: 10 modules × 48 blocks × 2 (A,B)            = 960
ref_adaln_proj: 4 (fc1.{w,b}, proj.{w,b})         =   4
role_embedding: 1                                 =   1
                                              total= 965 ✓

ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj

LoRA: 14 modules × 48 blocks × 2 (A,B)            = 1344
ref_adaln_proj: 4                                 =    4
ref_visual_proj: 7                                =    7
role_embedding: 1                                 =    1
                                              total= 1356 ✓

---

5. Loading checkpoint at inference

Use scripts/split_editanything_lora.py to split each raw training checkpoint into:

• *.standard.safetensors — LoRA on attn1/attn2/ff only; safe to feed to ComfyUI's standard LoraLoader.
• *.module.safetensors — everything else (role_embedding, ref_adaln_proj, ref_visual_proj, ref_attn LoRA adapters); feed to LTXVEditAnythingModuleLoader.

The filename suffix lists every extra that ended up in the module sidecar, so it is obvious at a glance which mechanisms a given pair carries. Order is fixed: ref_adaln_proj, role_embedding, ref_attn, ref_visual_proj.

Canonical output names

edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding.standard.safetensors
edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding.module.safetensors

edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj.standard.safetensors
edit_anything_reference_v0.1_r128_ref_adaln_proj-role_embedding-ref_attn-ref_visual_proj.module.safetensors

Command

python3 /data/training/ltx-edit-trainer/scripts/split_editanything_lora.py \
  <raw-checkpoint>.safetensors --output-dir <dir> [--overwrite]