RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Implement regular and timestep zero reference images to krea 2 for ostris and identity edit ref loras. (#14843)

This commit is contained in:
comfyanonymous 2026-07-18 17:12:18 -07:00 committed by GitHub
parent 83082a51c4
commit c9602625e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 143 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

147
comfy/ldm/krea2/model.py
View File

@ -15,6 +15,7 @@ from einops import rearrange
import comfy.model_management
import comfy.patcher_extension
import comfy.ldm.common_dit
import comfy.utils
from comfy.ldm.flux.layers import EmbedND, timestep_embedding
from comfy.ldm.flux.math import apply_rope
from comfy.ldm.modules.attention import optimized_attention_masked
@ -73,11 +74,20 @@ class Attention(nn.Module):
self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
def forward(self, x, freqs=None, mask=None, transformer_options={}):
transformer_patches = transformer_options.get("patches", {})
extra_options = transformer_options.copy()
q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
q, k = self.qknorm(q, k)
if "block_index" in transformer_options and "attn1_patch" in transformer_patches:
for p in transformer_patches["attn1_patch"]:
out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options)
q, k, v = out.get("q", q), out.get("k", k), out.get("v", v)
freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask)
if freqs is not None:
q, k = apply_rope(q, k, freqs)
if self.kvheads != self.heads:
@ -86,6 +96,11 @@ class Attention(nn.Module):
v = v.repeat_interleave(rep, dim=1)
out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True,
transformer_options=transformer_options)
if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches:
for p in transformer_patches["attn1_output_patch"]:
out = p(out, extra_options)
return self.wo(out * F.sigmoid(gate))
@ -158,8 +173,44 @@ class SingleStreamBlock(nn.Module):
self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
def forward(self, x, vec, freqs, mask=None, transformer_options={}):
def forward(self, x, vec, freqs, mask=None, timestep_zero_index=None, transformer_options={}):
prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
if timestep_zero_index is not None:
bs = x.shape[0]
ref_prescale = prescale[bs:]
ref_preshift = preshift[bs:]
ref_pregate = pregate[bs:]
ref_postscale = postscale[bs:]
ref_postshift = postshift[bs:]
ref_postgate = postgate[bs:]
prescale = prescale[:bs]
preshift = preshift[:bs]
pregate = pregate[:bs]
postscale = postscale[:bs]
postshift = postshift[:bs]
postgate = postgate[:bs]
pre = self.prenorm(x)
pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift)
pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift)
attn = self.attn(pre, freqs, mask, transformer_options=transformer_options)
del pre
attn[:, :timestep_zero_index].mul_(pregate)
attn[:, timestep_zero_index:].mul_(ref_pregate)
x = x + attn
del attn
post = self.postnorm(x)
post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift)
post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift)
mlp = self.mlp(post)
del post
mlp[:, :timestep_zero_index].mul_(postgate)
mlp[:, timestep_zero_index:].mul_(ref_postgate)
x = x + mlp
del mlp
return x
x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
return x
@ -181,7 +232,7 @@ class LastLayer(nn.Module):
class SingleStreamDiT(nn.Module):
def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12,
txtheads=20, txtkvheads=20, image_model=None,
txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None,
device=None, dtype=None, operations=None, **kwargs):
super().__init__()
self.dtype = dtype
@ -191,6 +242,7 @@ class SingleStreamDiT(nn.Module):
self.heads = heads
self.txtdim = txtdim
self.txtlayers = txtlayers
self.default_ref_method = default_ref_method
headdim = features // heads
axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
@ -221,61 +273,110 @@ class SingleStreamDiT(nn.Module):
operations.Linear(features, features * 6, device=device, dtype=dtype),
)
def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options),
).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs)
).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs)
def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
def process_img(self, x, index=0):
patch = self.patch
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
h, w = x.shape[-2] // patch, x.shape[-1] // patch
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32)
img_ids[..., 0] = index
img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None]
img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :]
return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w
def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
transformer_options = transformer_options.copy()
temporal = x.ndim == 5
if temporal:
b5, c5, t5, h5, w5 = x.shape
x = x.reshape(b5 * t5, c5, h5, w5)
bs, c, H_orig, W_orig = x.shape
bs, _, h_orig, w_orig = x.shape
patch = self.patch
# Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end.
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
H, W = x.shape[-2], x.shape[-1]
h_, w_ = H // patch, W // patch
# context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim).
context = self._unpack_context(context)
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
img, imgpos, h_, w_ = self.process_img(x)
img_tokens = img.shape[1]
timestep_zero_index = None
ref_method = kwargs.get("ref_latents_method", self.default_ref_method)
if ref_method is not None and ref_latents is not None and len(ref_latents) > 0:
ref_tokens = []
ref_pos = []
ref_num_tokens = []
for index, ref in enumerate(ref_latents, 1):
if ref.ndim == 5:
rb, rc, rt, rh5, rw5 = ref.shape
ref = ref.reshape(rb * rt, rc, rh5, rw5)
ref = comfy.utils.repeat_to_batch_size(ref, bs)
kontext, kontext_ids, _, _ = self.process_img(ref, index=index)
ref_tokens.append(kontext)
ref_pos.append(kontext_ids)
ref_num_tokens.append(kontext.shape[1])
img = torch.cat([img] + ref_tokens, dim=1)
imgpos = torch.cat([imgpos] + ref_pos, dim=1)
del ref_tokens, ref_pos
if ref_method == "index_timestep_zero":
timestep_zero_index = img_tokens
transformer_options["reference_image_num_tokens"] = ref_num_tokens
img = self.first(img)
t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype))
tvec = self.tproj(t)
if timestep_zero_index is not None:
t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype))
tvec = torch.cat((tvec, self.tproj(t0)), dim=0)
context = self.txtfusion(context, mask=None, transformer_options=transformer_options)
context = self.txtmlp(context)
txtlen, imglen = context.shape[1], img.shape[1]
txtlen = context.shape[1]
device = context.device
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
patches = transformer_options.get("patches", {})
if "post_input" in patches:
for p in patches["post_input"]:
out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options})
img, context = out["img"], out["txt"]
imgpos, txtpos = out["img_ids"], out["txt_ids"]
combined = torch.cat((context, img), dim=1)
del context, img
if timestep_zero_index is not None:
timestep_zero_index += txtlen
# Position ids: text at 0, image at (0, h_idx, w_idx).
device = combined.device
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
pos = torch.cat((txtpos, imgpos), dim=1)
del txtpos, imgpos
freqs = self.pe_embedder(pos)
del pos
for block in self.blocks:
combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "single"
transformer_options["img_slice"] = [txtlen, combined.shape[1]]
for i, block in enumerate(self.blocks):
transformer_options["block_index"] = i
combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options)
final = self.last(combined, t)
out = final[:, txtlen:txtlen + imglen, :]
del combined
out = final[:, txtlen:txtlen + img_tokens, :]
out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
h=h_, w=w_, ph=patch, pw=patch, c=self.channels)
out = out[:, :, :H_orig, :W_orig] # crop padding back off
out = out[:, :, :h_orig, :w_orig] # crop padding back off
if temporal:
out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2)
out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2)
return out
def _unpack_context(self, context):

23
comfy/model_base.py
View File

@ -2227,10 +2227,7 @@ class Omnigen2(BaseModel):
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
latents = []
for lat in ref_latents:
latents.append(self.process_latent_in(lat))
out['ref_latents'] = comfy.conds.CONDList(latents)
out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
return out
def extra_conds_shapes(self, **kwargs):
@ -2317,12 +2314,30 @@ class Ideogram4(BaseModel):
class Krea2(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT)
self.memory_usage_factor_conds = ("ref_latents",)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
cross_attn = kwargs.get("cross_attn", None)
if cross_attn is not None:
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
latents = []
for lat in ref_latents:
latents.append(self.process_latent_in(lat))
out['ref_latents'] = comfy.conds.CONDList(latents)
ref_latents_method = kwargs.get("reference_latents_method", None)
if ref_latents_method is not None:
out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
return out
def extra_conds_shapes(self, **kwargs):
out = {}
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
return out
class HunyuanImage21(BaseModel):