CthomasdesignandClaude Opus 4.8
984aabc0d1
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>
334 lines
13 KiB
Python
334 lines
13 KiB
Python
"""MLX port of Krea-2 SingleStreamDiT (krea/Krea-2-Turbo transformer).
|
|
|
|
Faithful line-by-line port of krea-2-official/mmdit.py. Module attribute names
|
|
match the `turbo.safetensors` tensor names exactly, so weights load with no
|
|
remapping. Reference casts the whole model to bf16 at inference; RMSNorm upcasts
|
|
to f32 internally (weight = stored `scale` + 1.0).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
import mlx.core as mx
|
|
from mlx import nn
|
|
from mlx.core.fast import scaled_dot_product_attention
|
|
|
|
SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi)
|
|
|
|
|
|
def gelu_tanh(x: mx.array) -> mx.array:
|
|
# Matches torch.nn.GELU(approximate="tanh").
|
|
return 0.5 * x * (1.0 + mx.tanh(SQRT_2_OVER_PI * (x + 0.044715 * x * x * x)))
|
|
|
|
|
|
class GELUTanh(nn.Module):
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return gelu_tanh(x)
|
|
|
|
|
|
@dataclass
|
|
class Krea2Config:
|
|
features: int = 6144
|
|
tdim: int = 256
|
|
txtdim: int = 2560
|
|
heads: int = 48
|
|
kvheads: int = 12
|
|
multiplier: int = 4
|
|
layers: int = 28
|
|
patch: int = 2
|
|
channels: int = 16
|
|
theta: float = 1000.0
|
|
txtheads: int = 20
|
|
txtkvheads: int = 20
|
|
txtlayers: int = 12 # number of selected encoder hidden-state layers fed in
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
"""Krea2 RMSNorm: effective weight = (scale + 1), math in float32, eps=1e-5."""
|
|
|
|
def __init__(self, features: int, eps: float = 1e-5):
|
|
super().__init__()
|
|
self.scale = mx.zeros((features,))
|
|
self.eps = eps
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
dt = x.dtype
|
|
t = x.astype(mx.float32)
|
|
t = t * mx.rsqrt(mx.mean(t * t, axis=-1, keepdims=True) + self.eps)
|
|
t = t * (1.0 + self.scale.astype(mx.float32))
|
|
return t.astype(dt)
|
|
|
|
|
|
class QKNorm(nn.Module):
|
|
def __init__(self, dim: int):
|
|
super().__init__()
|
|
self.qnorm = RMSNorm(dim)
|
|
self.knorm = RMSNorm(dim)
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
def __init__(self, features: int, multiplier: int, multiple: int = 128):
|
|
super().__init__()
|
|
mlpdim = int(2 * features / 3) * multiplier
|
|
mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
|
|
self.gate = nn.Linear(features, mlpdim, bias=False)
|
|
self.up = nn.Linear(features, mlpdim, bias=False)
|
|
self.down = nn.Linear(mlpdim, features, bias=False)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.down(nn.silu(self.gate(x)) * self.up(x))
|
|
|
|
|
|
def _make_rope(pos: mx.array, axes: list[int], theta: float) -> tuple[mx.array, mx.array]:
|
|
"""pos: (L, 3) float positions -> (cos, sin) each (L, sum(axes)//2)."""
|
|
cos_parts, sin_parts = [], []
|
|
for i, d in enumerate(axes):
|
|
scale = mx.arange(0, d, 2).astype(mx.float32) / d # (d/2,)
|
|
omega = 1.0 / (theta**scale)
|
|
freqs = pos[:, i : i + 1] * omega[None, :] # (L, d/2)
|
|
cos_parts.append(mx.cos(freqs))
|
|
sin_parts.append(mx.sin(freqs))
|
|
return mx.concatenate(cos_parts, axis=-1), mx.concatenate(sin_parts, axis=-1)
|
|
|
|
|
|
def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
|
|
"""x: (B, H, L, D); cos/sin: (L, D/2). Interleaved (adjacent-pair) rotation."""
|
|
b, h, l, d = x.shape
|
|
xf = x.astype(mx.float32).reshape(b, h, l, d // 2, 2)
|
|
x0, x1 = xf[..., 0], xf[..., 1]
|
|
c, s = cos[None, None], sin[None, None] # (1,1,L,D/2)
|
|
o0 = x0 * c - x1 * s
|
|
o1 = x0 * s + x1 * c
|
|
out = mx.stack([o0, o1], axis=-1).reshape(b, h, l, d)
|
|
return out.astype(x.dtype)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, dim: int, heads: int, kvheads: int | None = None):
|
|
super().__init__()
|
|
self.heads = heads
|
|
self.kvheads = kvheads if kvheads is not None else heads
|
|
self.headdim = dim // heads
|
|
self.scale = self.headdim**-0.5
|
|
self.wq = nn.Linear(dim, self.headdim * self.heads, bias=False)
|
|
self.wk = nn.Linear(dim, self.headdim * self.kvheads, bias=False)
|
|
self.wv = nn.Linear(dim, self.headdim * self.kvheads, bias=False)
|
|
self.gate = nn.Linear(dim, dim, bias=False)
|
|
self.qknorm = QKNorm(self.headdim)
|
|
self.wo = nn.Linear(dim, dim, bias=False)
|
|
|
|
def __call__(
|
|
self,
|
|
qkv: mx.array,
|
|
cos: mx.array | None = None,
|
|
sin: mx.array | None = None,
|
|
mask: mx.array | None = None,
|
|
) -> mx.array:
|
|
b, l, _ = qkv.shape
|
|
q = self.wq(qkv).reshape(b, l, self.heads, self.headdim).transpose(0, 2, 1, 3)
|
|
k = self.wk(qkv).reshape(b, l, self.kvheads, self.headdim).transpose(0, 2, 1, 3)
|
|
v = self.wv(qkv).reshape(b, l, self.kvheads, self.headdim).transpose(0, 2, 1, 3)
|
|
gate = self.gate(qkv)
|
|
|
|
q = self.qknorm.qnorm(q)
|
|
k = self.qknorm.knorm(k)
|
|
|
|
if cos is not None:
|
|
q = _apply_rope(q, cos, sin)
|
|
k = _apply_rope(k, cos, sin)
|
|
|
|
# GQA: mx.fast SDPA groups kv heads natively (bit-identical to repeating them).
|
|
out = scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
|
|
out = out.transpose(0, 2, 1, 3).reshape(b, l, self.heads * self.headdim)
|
|
return self.wo(out * mx.sigmoid(gate))
|
|
|
|
|
|
class DoubleSharedModulation(nn.Module):
|
|
def __init__(self, dim: int):
|
|
super().__init__()
|
|
self.lin = mx.zeros((6 * dim,))
|
|
|
|
def __call__(self, vec: mx.array) -> tuple[mx.array, ...]:
|
|
out = vec + self.lin
|
|
return mx.split(out, 6, axis=-1)
|
|
|
|
|
|
class SimpleModulation(nn.Module):
|
|
def __init__(self, dim: int):
|
|
super().__init__()
|
|
self.lin = mx.zeros((2, dim))
|
|
|
|
def __call__(self, vec: mx.array) -> tuple[mx.array, mx.array]:
|
|
out = vec + self.lin[None] # (B,1,d)+(1,2,d) -> (B,2,d)
|
|
scale, shift = mx.split(out, 2, axis=1)
|
|
return scale, shift
|
|
|
|
|
|
class SingleStreamBlock(nn.Module):
|
|
def __init__(self, features: int, heads: int, multiplier: int, kvheads: int):
|
|
super().__init__()
|
|
self.mod = DoubleSharedModulation(features)
|
|
self.prenorm = RMSNorm(features)
|
|
self.postnorm = RMSNorm(features)
|
|
self.attn = Attention(features, heads, kvheads)
|
|
self.mlp = SwiGLU(features, multiplier)
|
|
|
|
def __call__(self, x, vec, cos, sin, mask):
|
|
prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
|
|
x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, cos, sin, mask)
|
|
x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
|
|
return x
|
|
|
|
|
|
class TextFusionBlock(nn.Module):
|
|
def __init__(self, features: int, heads: int, multiplier: int, kvheads: int):
|
|
super().__init__()
|
|
self.prenorm = RMSNorm(features)
|
|
self.postnorm = RMSNorm(features)
|
|
self.attn = Attention(features, heads, kvheads)
|
|
self.mlp = SwiGLU(features, multiplier)
|
|
|
|
def __call__(self, x, mask=None):
|
|
x = x + self.attn(self.prenorm(x), mask=mask)
|
|
x = x + self.mlp(self.postnorm(x))
|
|
return x
|
|
|
|
|
|
class TextFusionTransformer(nn.Module):
|
|
def __init__(self, num_txt_layers, txt_dim, heads, multiplier, kvheads):
|
|
super().__init__()
|
|
self.layerwise_blocks = [TextFusionBlock(txt_dim, heads, multiplier, kvheads) for _ in range(2)]
|
|
self.projector = nn.Linear(num_txt_layers, 1, bias=False)
|
|
self.refiner_blocks = [TextFusionBlock(txt_dim, heads, multiplier, kvheads) for _ in range(2)]
|
|
|
|
def __call__(self, x: mx.array, mask: mx.array | None = None) -> mx.array:
|
|
# x: (B, L, n_layers, D)
|
|
b, l, n, d = x.shape
|
|
x = x.reshape(b * l, n, d)
|
|
for block in self.layerwise_blocks:
|
|
x = block(x, mask=None)
|
|
# (b l) n d -> b l d n
|
|
x = x.reshape(b, l, n, d).transpose(0, 1, 3, 2)
|
|
x = self.projector(x) # (b, l, d, 1)
|
|
x = x[..., 0] # (b, l, d)
|
|
for block in self.refiner_blocks:
|
|
x = block(x, mask=mask)
|
|
return x
|
|
|
|
|
|
class LastLayer(nn.Module):
|
|
def __init__(self, features: int, patch: int, channels: int):
|
|
super().__init__()
|
|
self.norm = RMSNorm(features)
|
|
self.linear = nn.Linear(features, patch * patch * channels, bias=True)
|
|
self.modulation = SimpleModulation(features)
|
|
|
|
def __call__(self, x: mx.array, tvec: mx.array) -> mx.array:
|
|
scale, shift = self.modulation(tvec)
|
|
x = (1 + scale) * self.norm(x) + shift
|
|
return self.linear(x)
|
|
|
|
|
|
def _timestep_embed(t: mx.array, dim: int, period: float = 1e4, tfactor: float = 1e3) -> mx.array:
|
|
half = dim // 2
|
|
freqs = mx.exp(-math.log(period) * mx.arange(half).astype(mx.float32) / half)
|
|
args = (t.astype(mx.float32) * tfactor)[:, None, None] * freqs # (B,1,half)
|
|
return mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) # (B,1,dim)
|
|
|
|
|
|
def _additive_mask(valid: mx.array, dtype: mx.Dtype) -> mx.array | None:
|
|
"""valid: (B, L) {0,1}. Returns additive (B,1,L,L) mask, or None if all valid."""
|
|
if bool(mx.all(valid >= 0.5).item()):
|
|
return None
|
|
m = valid.astype(mx.float32)
|
|
full = m[:, :, None] * m[:, None, :] # (B,L,L)
|
|
add = (1.0 - full) * -1e9
|
|
return add[:, None].astype(dtype)
|
|
|
|
|
|
class SingleStreamDiT(nn.Module):
|
|
def __init__(self, cfg: Krea2Config = Krea2Config()):
|
|
super().__init__()
|
|
self.cfg = cfg
|
|
headdim = cfg.features // cfg.heads
|
|
self.axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
|
|
assert sum(self.axes) == headdim
|
|
|
|
self.first = nn.Linear(cfg.channels * cfg.patch**2, cfg.features, bias=True)
|
|
self.blocks = [SingleStreamBlock(cfg.features, cfg.heads, cfg.multiplier, cfg.kvheads) for _ in range(cfg.layers)]
|
|
self.tmlp = [nn.Linear(cfg.tdim, cfg.features), GELUTanh(), nn.Linear(cfg.features, cfg.features)]
|
|
self.txtfusion = TextFusionTransformer(cfg.txtlayers, cfg.txtdim, cfg.txtheads, cfg.multiplier, cfg.txtkvheads)
|
|
self.txtmlp = [RMSNorm(cfg.txtdim), nn.Linear(cfg.txtdim, cfg.features), GELUTanh(), nn.Linear(cfg.features, cfg.features)]
|
|
self.last = LastLayer(cfg.features, cfg.patch, cfg.channels)
|
|
self.tproj = [GELUTanh(), nn.Linear(cfg.features, cfg.features * 6)]
|
|
|
|
def _run_seq(self, layers, x):
|
|
for layer in layers:
|
|
x = layer(x)
|
|
return x
|
|
|
|
def prepare_conditioning(self, context, pos, mask, dtype):
|
|
"""Step-invariant work (text fusion, RoPE tables, attention masks), computed once per
|
|
generation and reused across denoising steps. Bit-identical to running it per step.
|
|
|
|
context: (B, seq, n_layers, txtdim); pos: (L, 3); mask: (B, L) validity {0,1}.
|
|
Returns (fused_ctx, cos, sin, full_mask)."""
|
|
txtlen = context.shape[1]
|
|
txtmask = _additive_mask(mask[:, :txtlen], dtype)
|
|
context = self.txtfusion(context, mask=txtmask)
|
|
context = self._run_seq(self.txtmlp, context)
|
|
cos, sin = _make_rope(pos.astype(mx.float32), self.axes, self.cfg.theta)
|
|
full_mask = _additive_mask(mask, dtype)
|
|
return context, cos, sin, full_mask
|
|
|
|
def denoise_step(self, img, fused_ctx, t, cos, sin, full_mask):
|
|
"""One velocity evaluation. img: (B, Limg, channels*patch^2); t: (B,)."""
|
|
img = self.first(img)
|
|
t_emb = self._run_seq(self.tmlp, _timestep_embed(t, self.cfg.tdim).astype(img.dtype)) # (B,1,feat)
|
|
tvec = self._run_seq(self.tproj, t_emb) # (B,1,6*feat)
|
|
|
|
txtlen = fused_ctx.shape[1]
|
|
combined = mx.concatenate([fused_ctx, img], axis=1)
|
|
for block in self.blocks:
|
|
combined = block(combined, tvec, cos, sin, full_mask)
|
|
|
|
final = self.last(combined, t_emb)
|
|
return final[:, txtlen : txtlen + img.shape[1], :]
|
|
|
|
def prepare_edit(self, context, src_patched, pos, mask, dtype):
|
|
"""In-context edit conditioning. Same as prepare_conditioning, plus projects the CLEAN
|
|
source token block(s) through `first` once (fixed across steps). `pos`/`mask` must span the
|
|
full [text | src(s) | target] sequence. src_patched: (B, srclen, channels*patch^2).
|
|
Returns (fused_ctx, src_tokens, cos, sin, full_mask)."""
|
|
fused_ctx, cos, sin, full_mask = self.prepare_conditioning(context, pos, mask, dtype)
|
|
src_tokens = self.first(src_patched)
|
|
return fused_ctx, src_tokens, cos, sin, full_mask
|
|
|
|
def denoise_step_edit(self, img, fused_ctx, src_tokens, t, cos, sin, full_mask):
|
|
"""Edit velocity: sequence is [text | src(clean) | target]; only the target is returned.
|
|
`src_tokens` are pre-projected (frame≠0 via RoPE in cos/sin) and fixed across steps."""
|
|
img = self.first(img)
|
|
t_emb = self._run_seq(self.tmlp, _timestep_embed(t, self.cfg.tdim).astype(img.dtype))
|
|
tvec = self._run_seq(self.tproj, t_emb)
|
|
|
|
txtlen, srclen = fused_ctx.shape[1], src_tokens.shape[1]
|
|
combined = mx.concatenate([fused_ctx, src_tokens, img], axis=1)
|
|
for block in self.blocks:
|
|
combined = block(combined, tvec, cos, sin, full_mask)
|
|
|
|
final = self.last(combined, t_emb)
|
|
return final[:, txtlen + srclen : txtlen + srclen + img.shape[1], :]
|
|
|
|
def __call__(
|
|
self,
|
|
img: mx.array, # (B, Limg, channels*patch^2)
|
|
context: mx.array, # (B, seq, n_layers, txtdim)
|
|
t: mx.array, # (B,) current timestep in [0,1]
|
|
pos: mx.array, # (L, 3) positions for [txt; img]
|
|
mask: mx.array, # (B, L) validity {0,1}
|
|
) -> mx.array:
|
|
fused_ctx, cos, sin, full_mask = self.prepare_conditioning(context, pos, mask, img.dtype)
|
|
return self.denoise_step(img, fused_ctx, t, cos, sin, full_mask)
|