From 6c163508e18cd11bc4a4d9a18fd17b40db3a5772 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 13 Jul 2026 13:14:22 +0200 Subject: [PATCH] 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 = []