diff --git a/comfy/ldm/krea2/model.py b/comfy/ldm/krea2/model.py index ecb16254f..8001812d7 100644 --- a/comfy/ldm/krea2/model.py +++ b/comfy/ldm/krea2/model.py @@ -15,6 +15,7 @@ from einops import rearrange import comfy.model_management import comfy.patcher_extension import comfy.ldm.common_dit +import comfy.utils from comfy.ldm.flux.layers import EmbedND, timestep_embedding from comfy.ldm.flux.math import apply_rope from comfy.ldm.modules.attention import optimized_attention_masked @@ -73,11 +74,20 @@ class Attention(nn.Module): self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype) def forward(self, x, freqs=None, mask=None, transformer_options={}): + transformer_patches = transformer_options.get("patches", {}) + extra_options = transformer_options.copy() q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x) q = rearrange(q, "B L (H D) -> B H L D", H=self.heads) k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads) v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads) q, k = self.qknorm(q, k) + + if "block_index" in transformer_options and "attn1_patch" in transformer_patches: + for p in transformer_patches["attn1_patch"]: + out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options) + q, k, v = out.get("q", q), out.get("k", k), out.get("v", v) + freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask) + if freqs is not None: q, k = apply_rope(q, k, freqs) if self.kvheads != self.heads: @@ -86,6 +96,11 @@ class Attention(nn.Module): v = v.repeat_interleave(rep, dim=1) out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True, transformer_options=transformer_options) + + if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches: + for p in transformer_patches["attn1_output_patch"]: + out = p(out, extra_options) + return self.wo(out * F.sigmoid(gate)) @@ -158,8 +173,44 @@ 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, timestep_zero_index=None, transformer_options={}): prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec) + if timestep_zero_index is not None: + bs = x.shape[0] + ref_prescale = prescale[bs:] + ref_preshift = preshift[bs:] + ref_pregate = pregate[bs:] + ref_postscale = postscale[bs:] + ref_postshift = postshift[bs:] + ref_postgate = postgate[bs:] + prescale = prescale[:bs] + preshift = preshift[:bs] + pregate = pregate[:bs] + postscale = postscale[:bs] + postshift = postshift[:bs] + postgate = postgate[:bs] + + pre = self.prenorm(x) + pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift) + pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift) + attn = self.attn(pre, freqs, mask, transformer_options=transformer_options) + del pre + attn[:, :timestep_zero_index].mul_(pregate) + attn[:, timestep_zero_index:].mul_(ref_pregate) + x = x + attn + del attn + + post = self.postnorm(x) + post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift) + post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift) + mlp = self.mlp(post) + del post + mlp[:, :timestep_zero_index].mul_(postgate) + mlp[:, timestep_zero_index:].mul_(ref_postgate) + x = x + mlp + del mlp + return x + 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 @@ -181,7 +232,7 @@ class LastLayer(nn.Module): class SingleStreamDiT(nn.Module): def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4, layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12, - txtheads=20, txtkvheads=20, image_model=None, + txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None, device=None, dtype=None, operations=None, **kwargs): super().__init__() self.dtype = dtype @@ -191,6 +242,7 @@ class SingleStreamDiT(nn.Module): self.heads = heads self.txtdim = txtdim self.txtlayers = txtlayers + self.default_ref_method = default_ref_method headdim = features // heads axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)] @@ -221,61 +273,110 @@ class SingleStreamDiT(nn.Module): operations.Linear(features, features * 6, device=device, dtype=dtype), ) - def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs): + def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs): return comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, self, comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options), - ).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs) + ).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs) - def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs): + def process_img(self, x, index=0): + patch = self.patch + x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch)) + h, w = x.shape[-2] // patch, x.shape[-1] // patch + img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) + + img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32) + img_ids[..., 0] = index + img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None] + img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :] + return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w + + def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs): + transformer_options = transformer_options.copy() temporal = x.ndim == 5 if temporal: b5, c5, t5, h5, w5 = x.shape x = x.reshape(b5 * t5, c5, h5, w5) - bs, c, H_orig, W_orig = x.shape + bs, _, h_orig, w_orig = x.shape patch = self.patch - # Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end. - x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch)) - H, W = x.shape[-2], x.shape[-1] - h_, w_ = H // patch, W // patch # context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim). context = self._unpack_context(context) - img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch) + img, imgpos, h_, w_ = self.process_img(x) + img_tokens = img.shape[1] + timestep_zero_index = None + ref_method = kwargs.get("ref_latents_method", self.default_ref_method) + if ref_method is not None and ref_latents is not None and len(ref_latents) > 0: + ref_tokens = [] + ref_pos = [] + ref_num_tokens = [] + for index, ref in enumerate(ref_latents, 1): + if ref.ndim == 5: + rb, rc, rt, rh5, rw5 = ref.shape + ref = ref.reshape(rb * rt, rc, rh5, rw5) + ref = comfy.utils.repeat_to_batch_size(ref, bs) + kontext, kontext_ids, _, _ = self.process_img(ref, index=index) + ref_tokens.append(kontext) + ref_pos.append(kontext_ids) + ref_num_tokens.append(kontext.shape[1]) + img = torch.cat([img] + ref_tokens, dim=1) + imgpos = torch.cat([imgpos] + ref_pos, dim=1) + del ref_tokens, ref_pos + if ref_method == "index_timestep_zero": + timestep_zero_index = img_tokens + transformer_options["reference_image_num_tokens"] = ref_num_tokens + img = self.first(img) t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype)) tvec = self.tproj(t) + if timestep_zero_index is not None: + t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype)) + tvec = torch.cat((tvec, self.tproj(t0)), dim=0) context = self.txtfusion(context, mask=None, transformer_options=transformer_options) context = self.txtmlp(context) - txtlen, imglen = context.shape[1], img.shape[1] + txtlen = context.shape[1] + device = context.device + txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32) + + patches = transformer_options.get("patches", {}) + if "post_input" in patches: + for p in patches["post_input"]: + out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options}) + img, context = out["img"], out["txt"] + imgpos, txtpos = out["img_ids"], out["txt_ids"] + combined = torch.cat((context, img), dim=1) + del context, img + if timestep_zero_index is not None: + timestep_zero_index += txtlen # 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) + del txtpos, imgpos freqs = self.pe_embedder(pos) + del pos - for block in self.blocks: - combined = block(combined, tvec, freqs, None, transformer_options=transformer_options) + transformer_options["total_blocks"] = len(self.blocks) + transformer_options["block_type"] = "single" + transformer_options["img_slice"] = [txtlen, combined.shape[1]] + for i, block in enumerate(self.blocks): + transformer_options["block_index"] = i + combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options) final = self.last(combined, t) - out = final[:, txtlen:txtlen + imglen, :] + del combined + out = final[:, txtlen:txtlen + img_tokens, :] out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_, w=w_, ph=patch, pw=patch, c=self.channels) - out = out[:, :, :H_orig, :W_orig] # crop padding back off + out = out[:, :, :h_orig, :w_orig] # crop padding back off if temporal: - out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2) + out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2) return out def _unpack_context(self, context): diff --git a/comfy/model_base.py b/comfy/model_base.py index 98f5ba48b..0f705316c 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2227,10 +2227,7 @@ class Omnigen2(BaseModel): 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) + out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents]) return out def extra_conds_shapes(self, **kwargs): @@ -2317,12 +2314,30 @@ 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()), ref_latents)) // 16]) return out class HunyuanImage21(BaseModel):