From e4cb8309b48d4f10d48085b491ed5413e8f45ed5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Jul 2026 22:15:51 -0400 Subject: [PATCH] Implement timestep zero reference images to krea 2 for ostris ref loras. --- comfy/ldm/krea2/model.py | 110 ++++++++++++++++++++++++++++++++------- comfy/model_base.py | 14 +++++ 2 files changed, 105 insertions(+), 19 deletions(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index ecb16254f..594d99789 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -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 @@ -158,8 +159,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 @@ -221,61 +258,96 @@ 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): 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 + if 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 + timestep_zero_index = img_tokens + transformer_options = transformer_options.copy() + 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] 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) + 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): diff --git a/comfy/model_base.py b/comfy/model_base.py index dcfa555dc..5f00f19c9 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2282,12 +2282,26 @@ 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) + 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):