From 64ed140c89ce85506cc2536a5b64f40fcbb2dc40 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 23 Jun 2026 22:03:06 +0200 Subject: [PATCH 1/6] Implement reference latent support for Krea 2 model --- comfy/ldm/krea2/model.py | 72 ++++++++++++++++++++++++++++++++++++++-- comfy/model_base.py | 19 +++++++++++ 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index ecb16254f..91a089153 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -253,16 +253,30 @@ class SingleStreamDiT(nn.Module): context = self.txtmlp(context) txtlen, imglen = context.shape[1], img.shape[1] - combined = torch.cat((context, img), dim=1) + + ref_latents = kwargs.get("ref_latents", None) + ref_latents_method = kwargs.get("ref_latents_method", "offset") + device = img.device + + ref_tokens_list, ref_pos_ids_list, ref_num_tokens = self._process_ref_latents( + ref_latents, ref_latents_method, device, bs + ) + + if len(ref_num_tokens) > 0: + transformer_options = transformer_options.copy() + if "reference_image_num_tokens" not in transformer_options: + transformer_options["reference_image_num_tokens"] = [] + transformer_options["reference_image_num_tokens"].extend(ref_num_tokens) + + combined = torch.cat([context, img] + ref_tokens_list, dim=1) # 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) + pos = torch.cat([txtpos, imgpos] + ref_pos_ids_list, dim=1) freqs = self.pe_embedder(pos) @@ -288,3 +302,55 @@ class SingleStreamDiT(nn.Module): f"Load the text encoder with CLIPLoader type 'krea2'." ) return context.reshape(b, seq, self.txtlayers, self.txtdim) + + def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs): + ref_tokens_list = [] + ref_pos_ids_list = [] + ref_num_tokens = [] + patch = self.patch + + if ref_latents is not None: + h = 0 + w = 0 + index = 0 + index_ref_method = (ref_latents_method == "index") or (ref_latents_method == "index_timestep_zero") + negative_ref_method = ref_latents_method == "negative_index" + + for ref in ref_latents: + ref_pad = comfy.ldm.common_dit.pad_to_patch_size(ref, (patch, patch)) + ref_b, ref_c, ref_h, ref_w = ref_pad.shape + ref_gh = ref_h // patch + ref_gw = ref_w // patch + + if index_ref_method: + index += 1 + gh_offset = 0 + gw_offset = 0 + elif negative_ref_method: + index -= 1 + gh_offset = 0 + gw_offset = 0 + else: # offset/default + index = 1 + gh_offset = 0 + gw_offset = 0 + if ref_gh + h > ref_gw + w: + gw_offset = w + else: + gh_offset = h + h = max(h, ref_gh + gh_offset) + w = max(w, ref_gw + gw_offset) + + ref_tokens = rearrange(ref_pad, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) + ref_tokens = self.first(ref_tokens) + ref_tokens_list.append(ref_tokens) + ref_num_tokens.append(ref_tokens.shape[1]) + + ref_pos_ids = torch.zeros(ref_gh, ref_gw, 3, device=device, dtype=torch.float32) + ref_pos_ids[..., 0] = index + ref_pos_ids[..., 1] = torch.arange(ref_gh, device=device, dtype=torch.float32)[:, None] + gh_offset + ref_pos_ids[..., 2] = torch.arange(ref_gw, device=device, dtype=torch.float32)[None, :] + gw_offset + ref_pos_ids = ref_pos_ids.reshape(1, ref_gh * ref_gw, 3).repeat(bs, 1, 1) + ref_pos_ids_list.append(ref_pos_ids) + + return ref_tokens_list, ref_pos_ids_list, ref_num_tokens diff --git a/comfy/model_base.py b/comfy/model_base.py index dcfa555dc..b85e11f91 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2282,12 +2282,31 @@ 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()[2:]), ref_latents))]) return out class HunyuanImage21(BaseModel): From 9fecaeb2d5ebbba2bcc370c699d58e63e9be54e8 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 23 Jun 2026 22:21:39 +0200 Subject: [PATCH 2/6] Implement 5D-to-4D temporal latent reshape fix --- comfy/ldm/krea2/model.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index 91a089153..255119cbd 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -317,6 +317,9 @@ class SingleStreamDiT(nn.Module): negative_ref_method = ref_latents_method == "negative_index" for ref in ref_latents: + if ref.ndim == 5: + ref_b5, ref_c5, ref_t5, ref_h5, ref_w5 = ref.shape + ref = ref.reshape(ref_b5 * ref_t5, ref_c5, ref_h5, ref_w5) ref_pad = comfy.ldm.common_dit.pad_to_patch_size(ref, (patch, patch)) ref_b, ref_c, ref_h, ref_w = ref_pad.shape ref_gh = ref_h // patch From 0cfc1568b10caa191f231f8d33a1b255567a2bcd Mon Sep 17 00:00:00 2001 From: silveroxides Date: Tue, 30 Jun 2026 02:44:31 +0200 Subject: [PATCH 3/6] fix --- comfy/ldm/krea2/model.py | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index 255119cbd..5f3c801c3 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -259,7 +259,7 @@ class SingleStreamDiT(nn.Module): device = img.device ref_tokens_list, ref_pos_ids_list, ref_num_tokens = self._process_ref_latents( - ref_latents, ref_latents_method, device, bs + ref_latents, ref_latents_method, device, bs, h_, w_ ) if len(ref_num_tokens) > 0: @@ -303,15 +303,15 @@ class SingleStreamDiT(nn.Module): ) return context.reshape(b, seq, self.txtlayers, self.txtdim) - def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs): + def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs, h_main, w_main): ref_tokens_list = [] ref_pos_ids_list = [] ref_num_tokens = [] patch = self.patch if ref_latents is not None: - h = 0 - w = 0 + h = h_main + w = w_main index = 0 index_ref_method = (ref_latents_method == "index") or (ref_latents_method == "index_timestep_zero") negative_ref_method = ref_latents_method == "negative_index" @@ -327,22 +327,19 @@ class SingleStreamDiT(nn.Module): if index_ref_method: index += 1 - gh_offset = 0 - gw_offset = 0 elif negative_ref_method: index -= 1 - gh_offset = 0 - gw_offset = 0 else: # offset/default - index = 1 - gh_offset = 0 - gw_offset = 0 - if ref_gh + h > ref_gw + w: - gw_offset = w - else: - gh_offset = h - h = max(h, ref_gh + gh_offset) - w = max(w, ref_gw + gw_offset) + index = 0 # Stay in t_idx = 0 plane to remain 100% in-distribution + + gh_offset = 0 + gw_offset = 0 + if ref_gh + h > ref_gw + w: + gw_offset = w + else: + gh_offset = h + h = max(h, ref_gh + gh_offset) + w = max(w, ref_gw + gw_offset) ref_tokens = rearrange(ref_pad, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) ref_tokens = self.first(ref_tokens) From 7252db85560e94c46632c17fddb1435a9e24bd92 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 11 Jul 2026 16:17:01 +0200 Subject: [PATCH 4/6] Correct implementation issues with the help of Ostris custom node as reference --- comfy/ldm/krea2/model.py | 106 +++++++++++++++++++++++++-------------- 1 file changed, 68 insertions(+), 38 deletions(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index 5f3c801c3..78e7e1952 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -158,11 +158,54 @@ 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={}): - prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec) - 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 + def forward(self, x, vec, freqs, mask=None, transformer_options={}, vec_ref=None, split=None): + if vec_ref is not None and split is not None: + m = self.mod(vec) + r = self.mod(vec_ref) + + # prenorm and attention + h = self.prenorm(x) + h_mod = torch.cat( + ( + (1 + m[0]) * h[:, :split] + m[1], + (1 + r[0]) * h[:, split:] + r[1] + ), + dim=1 + ) + attn_out = self.attn(h_mod, freqs, mask, transformer_options=transformer_options) + attn_gate = torch.cat( + ( + m[2] * attn_out[:, :split], + r[2] * attn_out[:, split:] + ), + dim=1 + ) + x = x + attn_gate + + # postnorm and mlp + h = self.postnorm(x) + h_mod = torch.cat( + ( + (1 + m[3]) * h[:, :split] + m[4], + (1 + r[3]) * h[:, split:] + r[4] + ), + dim=1 + ) + mlp_out = self.mlp(h_mod) + mlp_gate = torch.cat( + ( + m[5] * mlp_out[:, :split], + r[5] * mlp_out[:, split:] + ), + dim=1 + ) + x = x + mlp_gate + return x + else: + prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec) + 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 class LastLayer(nn.Module): @@ -259,7 +302,7 @@ class SingleStreamDiT(nn.Module): device = img.device ref_tokens_list, ref_pos_ids_list, ref_num_tokens = self._process_ref_latents( - ref_latents, ref_latents_method, device, bs, h_, w_ + ref_latents, ref_latents_method, device, bs, h_, w_, img.dtype ) if len(ref_num_tokens) > 0: @@ -280,8 +323,16 @@ class SingleStreamDiT(nn.Module): freqs = self.pe_embedder(pos) - for block in self.blocks: - combined = block(combined, tvec, freqs, None, transformer_options=transformer_options) + if len(ref_num_tokens) > 0: + # Compute tvec0 for timestep=0 (reference) + t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype)) + tvec0 = self.tproj(t0) + split = txtlen + imglen + for block in self.blocks: + combined = block(combined, tvec, freqs, None, transformer_options=transformer_options, vec_ref=tvec0, split=split) + else: + for block in self.blocks: + combined = block(combined, tvec, freqs, None, transformer_options=transformer_options) final = self.last(combined, t) out = final[:, txtlen:txtlen + imglen, :] @@ -303,43 +354,22 @@ class SingleStreamDiT(nn.Module): ) return context.reshape(b, seq, self.txtlayers, self.txtdim) - def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs, h_main, w_main): + def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs, h_main, w_main, dtype): ref_tokens_list = [] ref_pos_ids_list = [] ref_num_tokens = [] patch = self.patch if ref_latents is not None: - h = h_main - w = w_main - index = 0 - index_ref_method = (ref_latents_method == "index") or (ref_latents_method == "index_timestep_zero") - negative_ref_method = ref_latents_method == "negative_index" - - for ref in ref_latents: + for i, ref in enumerate(ref_latents): if ref.ndim == 5: ref_b5, ref_c5, ref_t5, ref_h5, ref_w5 = ref.shape ref = ref.reshape(ref_b5 * ref_t5, ref_c5, ref_h5, ref_w5) ref_pad = comfy.ldm.common_dit.pad_to_patch_size(ref, (patch, patch)) - ref_b, ref_c, ref_h, ref_w = ref_pad.shape - ref_gh = ref_h // patch - ref_gw = ref_w // patch - - if index_ref_method: - index += 1 - elif negative_ref_method: - index -= 1 - else: # offset/default - index = 0 # Stay in t_idx = 0 plane to remain 100% in-distribution - - gh_offset = 0 - gw_offset = 0 - if ref_gh + h > ref_gw + w: - gw_offset = w - else: - gh_offset = h - h = max(h, ref_gh + gh_offset) - w = max(w, ref_gw + gw_offset) + ref_pad = comfy.utils.repeat_to_batch_size(ref_pad, bs) + ref_pad = ref_pad.to(device=device, dtype=dtype) + ref_gh = ref_pad.shape[-2] // patch + ref_gw = ref_pad.shape[-1] // patch ref_tokens = rearrange(ref_pad, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) ref_tokens = self.first(ref_tokens) @@ -347,9 +377,9 @@ class SingleStreamDiT(nn.Module): ref_num_tokens.append(ref_tokens.shape[1]) ref_pos_ids = torch.zeros(ref_gh, ref_gw, 3, device=device, dtype=torch.float32) - ref_pos_ids[..., 0] = index - ref_pos_ids[..., 1] = torch.arange(ref_gh, device=device, dtype=torch.float32)[:, None] + gh_offset - ref_pos_ids[..., 2] = torch.arange(ref_gw, device=device, dtype=torch.float32)[None, :] + gw_offset + ref_pos_ids[..., 0] = i + 1.0 + ref_pos_ids[..., 1] = torch.arange(ref_gh, device=device, dtype=torch.float32)[:, None] + ref_pos_ids[..., 2] = torch.arange(ref_gw, device=device, dtype=torch.float32)[None, :] ref_pos_ids = ref_pos_ids.reshape(1, ref_gh * ref_gw, 3).repeat(bs, 1, 1) ref_pos_ids_list.append(ref_pos_ids) From 6c163508e18cd11bc4a4d9a18fd17b40db3a5772 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 13:14:22 +0200 Subject: [PATCH 5/6] Fix reference latent method validation for Krea2 --- comfy/ldm/krea2/model.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index 78e7e1952..f9bc74b1c 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -298,11 +298,15 @@ class SingleStreamDiT(nn.Module): txtlen, imglen = context.shape[1], img.shape[1] ref_latents = kwargs.get("ref_latents", None) - ref_latents_method = kwargs.get("ref_latents_method", "offset") + ref_latents_method = kwargs.get("ref_latents_method", "index_timestep_zero") + if ref_latents_method is None: + ref_latents_method = "index_timestep_zero" + if ref_latents_method != "index_timestep_zero": + raise ValueError(f"Unsupported Krea2 reference latent method: {ref_latents_method}") device = img.device ref_tokens_list, ref_pos_ids_list, ref_num_tokens = self._process_ref_latents( - ref_latents, ref_latents_method, device, bs, h_, w_, img.dtype + ref_latents, device, bs, h_, w_, img.dtype ) if len(ref_num_tokens) > 0: @@ -354,7 +358,7 @@ class SingleStreamDiT(nn.Module): ) return context.reshape(b, seq, self.txtlayers, self.txtdim) - def _process_ref_latents(self, ref_latents, ref_latents_method, device, bs, h_main, w_main, dtype): + def _process_ref_latents(self, ref_latents, device, bs, h_main, w_main, dtype): ref_tokens_list = [] ref_pos_ids_list = [] ref_num_tokens = [] From 641ce082a93d1a502dede8f61984318cca6045ef Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 13:17:41 +0200 Subject: [PATCH 6/6] Fix 5D reference latent reshaping block --- comfy/ldm/krea2/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index f9bc74b1c..d4081eb7c 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -368,7 +368,7 @@ class SingleStreamDiT(nn.Module): for i, ref in enumerate(ref_latents): if ref.ndim == 5: ref_b5, ref_c5, ref_t5, ref_h5, ref_w5 = ref.shape - ref = ref.reshape(ref_b5 * ref_t5, ref_c5, ref_h5, ref_w5) + ref = ref.movedim(2, 1).reshape(ref_b5 * ref_t5, ref_c5, ref_h5, ref_w5) ref_pad = comfy.ldm.common_dit.pad_to_patch_size(ref, (patch, patch)) ref_pad = comfy.utils.repeat_to_batch_size(ref_pad, bs) ref_pad = ref_pad.to(device=device, dtype=dtype)