github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added support to train Krea2 as an edit model

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-04 09:12:39 -06:00
parent f63221e577
commit b1e1a834d4
7 changed files with 519 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
extensions_built_in/diffusion_models/krea2/src/mmdit.py
View File

@@ -328,6 +328,33 @@ class SingleStreamBlock(nn.Module):
def forward(
self, x: Tensor, vec: Tensor, freqs: Tensor, mask: Tensor | None = None
) -> Tensor:
# ``vec`` is the (B, 1, 6*features) modulation input, or a tuple
# ``(vec, refvec, split)`` for reference-image conditioning: tokens
# ``[:split]`` (text + noisy image) are modulated with ``vec`` while
# tokens ``[split:]`` (clean reference tokens) use ``refvec`` built from
# t=0 (ComfyUI Kontext "index_timestep_zero"). Applied per span rather
# than materializing a per-token (B, L, 6*features) tensor.
if isinstance(vec, tuple):
vec, refvec, split = vec
m = self.mod(vec)
r = self.mod(refvec)
def mod(h, scale, shift):
return torch.cat(
(
(1 + m[scale]) * h[:, :split] + m[shift],
(1 + r[scale]) * h[:, split:] + r[shift],
),
dim=1,
)
def gate(h, g):
return torch.cat((m[g] * h[:, :split], r[g] * h[:, split:]), dim=1)
x = x + gate(self.attn(mod(self.prenorm(x), 0, 1), freqs, mask), 2)
x = x + gate(self.mlp(mod(self.postnorm(x), 3, 4)), 5)
return x
prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
x = x + pregate * self.attn(
(1 + prescale) * self.prenorm(x) + preshift, freqs, mask
@@ -417,6 +444,7 @@ class SingleStreamDiT(nn.Module):
t: Tensor,
pos: Tensor,
mask: Tensor | None = None,
reflen: int = 0,
) -> Tensor:
img = self.first(img)
t = self.tmlp(temb(t, self.config.tdim, device=img.device, dtype=img.dtype))
@@ -438,6 +466,23 @@ class SingleStreamDiT(nn.Module):
mask = F.pad(mask, (0, _padlen), value=False)
pos = F.pad(pos, (0, 0, 0, _padlen))
blockvec = tvec
if reflen > 0:
# The last ``reflen`` image tokens are clean reference tokens: they
# get t=0 modulation (ComfyUI Kontext "index_timestep_zero") while
# text + noisy image tokens keep the real t. Padding tokens fall in
# the t=0 span, but they are masked from attention and sliced off
# the output, so their values never matter.
t0 = self.tmlp(
temb(
torch.zeros_like(t[:, 0, 0]),
self.config.tdim,
device=img.device,
dtype=img.dtype,
)
)
blockvec = (tvec, self.tproj(t0), txtlen + imglen - reflen)
mask = _mask(mask)
freqs = self.posemb(pos)
@@ -447,15 +492,15 @@ class SingleStreamDiT(nn.Module):
combined = checkpoint(
block,
combined,
tvec,
blockvec,
freqs,
mask,
use_reentrant=False,
)
else:
combined = block(combined, tvec, freqs, mask)
combined = block(combined, blockvec, freqs, mask)
final = self.last(combined, t)
output = final[:, txtlen : txtlen + imglen, :]
output = final[:, txtlen : txtlen + imglen - reflen, :]
return output

91
extensions_built_in/diffusion_models/krea2/src/pipeline.py
View File

@@ -90,20 +90,78 @@ def prepare(
return img, pos, mask
def pack_ref_latents(
ref_latents: List[List[torch.Tensor]], patch: int, device, dtype
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Patchify per-sample reference latents into padded ref tokens / pos / mask.
``ref_latents`` is a list (one entry per batch item) of lists of ``(C, h, w)``
reference latents. The i-th reference of a sample is placed on RoPE axis 0 at
index ``i + 1`` with its own y/x grid starting at 0 -- the ComfyUI Kontext
"index" placement (axis 0 is otherwise always 0, so the base weights see the
references as a new "frame" axis). Samples with fewer reference tokens are
right-padded and masked out. Returns ``(tokens (B, Lr, C*p*p),
pos (B, Lr, 3), mask (B, Lr))``.
"""
token_dim = None
seqs, ids = [], []
for refs in ref_latents:
toks, rpos = [], []
for i, ref in enumerate(refs):
ref = ref.to(device, dtype)
_, h, w = ref.shape
h_, w_ = h // patch, w // patch
toks.append(
rearrange(ref, "c (h ph) (w pw) -> (h w) (c ph pw)", ph=patch, pw=patch)
)
token_dim = toks[-1].shape[-1]
refids = torch.zeros((h_, w_, 3), device=device)
refids[..., 0] = i + 1
refids[..., 1] = torch.arange(h_, device=device)[:, None]
refids[..., 2] = torch.arange(w_, device=device)[None, :]
rpos.append(refids.reshape(-1, 3))
seqs.append(toks)
ids.append(rpos)
b = len(seqs)
# a sample may have no references (its span is fully masked/padded)
seqs = [
torch.cat(t, dim=0)
if t
else torch.zeros(0, token_dim, device=device, dtype=dtype)
for t in seqs
]
ids = [torch.cat(p, dim=0) if p else torch.zeros(0, 3, device=device) for p in ids]
max_len = max(s.shape[0] for s in seqs)
tokens = torch.zeros(b, max_len, token_dim, device=device, dtype=dtype)
pos = torch.zeros(b, max_len, 3, device=device)
mask = torch.zeros(b, max_len, device=device, dtype=torch.bool)
for i, (s, p) in enumerate(zip(seqs, ids)):
ln = s.shape[0]
tokens[i, :ln] = s
pos[i, :ln] = p
mask[i, :ln] = True
return tokens, pos, mask
def predict_velocity(
model: SingleStreamDiT,
latents: torch.Tensor, # (B, C, h, w)
t: torch.Tensor, # (B,) flow time in [0, 1] (1 = pure noise)
context: torch.Tensor, # (B, Lt, n*d) flattened stacked Qwen3-VL features
text_mask: torch.Tensor, # (B, Lt) 1 for real text tokens
ref_latents: Optional[List[List[torch.Tensor]]] = None, # per-sample (C, h, w) refs
) -> torch.Tensor:
"""Run the MMDiT on the packed [text | image] sequence.
"""Run the MMDiT on the packed [text | image | refs] sequence.
``latents`` stay in the unpacked ``(B, C, h, w)`` latent layout; image-token
packing is internal to this function. ``context`` arrives 2D-per-sample
flattened ``(B, Lt, n*d)`` and is restored to ``(B, Lt, n, d)`` for the MMDiT.
Returns the velocity ``noise - clean`` reshaped back to ``(B, C, h, w)``. No
time flip / negation: Krea's convention matches toolkit's.
``ref_latents`` (optional) are clean reference latents appended after the
image tokens and conditioned at t=0 ("index_timestep_zero"); the prediction
only ever covers the noisy target tokens. Returns the velocity
``noise - clean`` reshaped back to ``(B, C, h, w)``. No time flip / negation:
Krea's convention matches toolkit's.
"""
patch = model.config.patch
b, c, h, w = latents.shape
@@ -116,7 +174,17 @@ def predict_velocity(
img_tokens, pos, mask = prepare(latents, context.shape[1], patch, text_mask)
out = model(img=img_tokens, context=context, t=t, pos=pos, mask=mask)
reflen = 0
if ref_latents is not None and any(len(r) > 0 for r in ref_latents):
ref_tokens, ref_pos, ref_mask = pack_ref_latents(
ref_latents, patch, img_tokens.device, img_tokens.dtype
)
reflen = ref_tokens.shape[1]
img_tokens = torch.cat((img_tokens, ref_tokens), dim=1)
pos = torch.cat((pos, ref_pos), dim=1)
mask = torch.cat((mask, ref_mask), dim=1)
out = model(img=img_tokens, context=context, t=t, pos=pos, mask=mask, reflen=reflen)
# (B, imglen, c*p*p) -> (B, c, h, w)
velocity = rearrange(
@@ -193,6 +261,7 @@ class Krea2Pipeline:
guidance_scale: float = 4.5,
latents: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
ref_latents: Optional[List[List[torch.Tensor]]] = None,
**kwargs,
) -> List[Image.Image]:
model = self.model
@@ -242,11 +311,21 @@ class Krea2Pipeline:
for tcurr, tprev in zip(ts[:-1], ts[1:]):
t = torch.full((latents.shape[0],), tcurr, dtype=dtype, device=device)
v_cond = predict_velocity(
transformer, latents.to(dtype), t, cond_feats, cond_mask
transformer,
latents.to(dtype),
t,
cond_feats,
cond_mask,
ref_latents=ref_latents,
)
if do_cfg:
v_uncond = predict_velocity(
transformer, latents.to(dtype), t, uncond_feats, uncond_mask
transformer,
latents.to(dtype),
t,
uncond_feats,
uncond_mask,
ref_latents=ref_latents,
)
v = v_cond + guidance_scale * (v_cond - v_uncond)
else:

74
extensions_built_in/diffusion_models/krea2/src/text_encoder.py
View File

@@ -13,6 +13,8 @@ its hidden states are sliced off the returned features, exactly like the
reference.
"""
from typing import List, Optional
import torch
from torch import Tensor
@@ -43,6 +45,9 @@ def encode_krea_prompt(
max_length: int = 512,
select_layers: tuple[int, ...] = SELECT_LAYERS,
prefix_idx: int = PROMPT_TEMPLATE_ENCODE_START_IDX,
images: Optional[List[Tensor]] = None,
vl_processor=None,
dtype: Optional[torch.dtype] = None,
) -> Tensor:
"""Encode a single prompt into stacked Qwen3-VL hidden states.
@@ -50,6 +55,15 @@ def encode_krea_prompt(
dtype) holding the prompt + suffix token features -- the system prefix has
been sliced off. ``L`` is the natural (unpadded) length so the caller stores
one tensor per prompt and pads to the batch max at the model call.
``images`` (optional) are reference images -- ``(C, H, W)`` tensors in
``[0, 1]`` -- embedded in the user message ahead of the prompt via named
vision placeholders (``Picture 1: <|vision_start|><|image_pad|><|vision_end|>``,
the ComfyUI ``TextEncodeQwenImageEditPlus`` layout). The ``vl_processor``
(Qwen3-VL AutoProcessor) expands each ``<|image_pad|>`` to the image's token
grid, so the returned features carry the vision tokens as extra conditioning.
The system prefix is unchanged, so ``prefix_idx`` slicing stays valid and the
image + prompt tokens all survive the slice.
"""
device = qwen.device
@@ -61,21 +75,59 @@ def encode_krea_prompt(
suffix_ids = suffix_inputs["input_ids"]
suffix_mask = suffix_inputs["attention_mask"].bool()
# Prefix + prompt at natural length (no padding); truncate very long prompts.
text = PROMPT_TEMPLATE_ENCODE_PREFIX + prompt
inputs = tokenizer(
[text],
truncation=True,
return_length=False,
return_overflowing_tokens=False,
max_length=max_length + prefix_idx,
return_tensors="pt",
).to(device, non_blocking=True)
extra_inputs = {}
if images is not None and len(images) > 0:
image_prompt = "".join(
f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>"
for i in range(len(images))
)
text = PROMPT_TEMPLATE_ENCODE_PREFIX + image_prompt + prompt
# No truncation here: the expanded image-pad runs must stay intact.
inputs = vl_processor(
text=[text],
images=list(images),
return_tensors="pt",
do_rescale=False,
).to(device)
for k, v in inputs.items():
if k in ("input_ids", "attention_mask"):
continue
if (
isinstance(v, torch.Tensor)
and v.is_floating_point()
and dtype is not None
):
v = v.to(dtype)
extra_inputs[k] = v
else:
# Prefix + prompt at natural length (no padding); truncate very long prompts.
text = PROMPT_TEMPLATE_ENCODE_PREFIX + prompt
inputs = tokenizer(
[text],
truncation=True,
return_length=False,
return_overflowing_tokens=False,
max_length=max_length + prefix_idx,
return_tensors="pt",
).to(device, non_blocking=True)
input_ids = torch.cat([inputs["input_ids"], suffix_ids], dim=1)
mask = torch.cat([inputs["attention_mask"].bool(), suffix_mask], dim=1)
states = qwen(input_ids=input_ids, attention_mask=mask, output_hidden_states=True)
# mm_token_type_ids (used for M-RoPE) must cover the appended suffix tokens
# too; they are plain text -> type 0.
if "mm_token_type_ids" in extra_inputs:
tt = extra_inputs["mm_token_type_ids"]
extra_inputs["mm_token_type_ids"] = torch.cat(
[tt, torch.zeros_like(suffix_ids, dtype=tt.dtype)], dim=1
)
states = qwen(
input_ids=input_ids,
attention_mask=mask,
output_hidden_states=True,
**extra_inputs,
)
# (1, L, num_layers, hidden)
hiddens = torch.stack([states.hidden_states[i] for i in select_layers], dim=2)