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):