Merge 641ce082a9 into 0f42ba5146

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={}):
+    def forward(self, x, vec, freqs, mask=None, transformer_options={}, vec_ref=None, split=None):
-        prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+        if vec_ref is not None and split is not None:
-        x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
+            m = self.mod(vec)
-        x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
+            r = self.mod(vec_ref)
-        return x
+
+            # 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):
@@ -253,21 +296,47 @@ 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", "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, device, bs, h_, w_, img.dtype
+        )
+
+        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)
 
-        for block in self.blocks:
+        if len(ref_num_tokens) > 0:
-            combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
+            # 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, :]
@@ -288,3 +357,34 @@ 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, 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:
+            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.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)
+                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)
+                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] = 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)
+
+        return ref_tokens_list, ref_pos_ids_list, ref_num_tokens

comfy/model_base.py

@@ -2317,12 +2317,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):