wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 20:38:15 +08:00
Code Issues Actions 27 Packages Projects Releases Wiki Activity

Fix reference latent method validation for Krea2

Browse Source
This commit is contained in:
silveroxides 2026-07-13 13:14:22 +02:00
parent 7252db8556
commit 6c163508e1
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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 = []