From 22a1d7ce7896d09da375febe6279c9f89c0de933 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sat, 16 Nov 2024 20:42:21 -0500 Subject: [PATCH 1/3] Fix 3.8 compatibility in user_manager.py (#5645) --- app/user_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/app/user_manager.py b/app/user_manager.py index 4e545c0db..bb6217995 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -1,3 +1,4 @@ +from __future__ import annotations import json import os import re From 41886af138939719b86f30bacbafc76a142d89de Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 16 Nov 2024 14:47:06 -0500 Subject: [PATCH 2/3] Add transformer options blocks replace patch to mochi. --- .../genmo/joint_model/asymm_models_joint.py | 36 ++++++++++++++----- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/comfy/ldm/genmo/joint_model/asymm_models_joint.py b/comfy/ldm/genmo/joint_model/asymm_models_joint.py index c36a00068..45c938966 100644 --- a/comfy/ldm/genmo/joint_model/asymm_models_joint.py +++ b/comfy/ldm/genmo/joint_model/asymm_models_joint.py @@ -494,8 +494,9 @@ class AsymmDiTJoint(nn.Module): packed_indices: Dict[str, torch.Tensor] = None, rope_cos: torch.Tensor = None, rope_sin: torch.Tensor = None, - control=None, **kwargs + control=None, transformer_options={}, **kwargs ): + patches_replace = transformer_options.get("patches_replace", {}) y_feat = context y_mask = attention_mask sigma = timestep @@ -515,15 +516,32 @@ class AsymmDiTJoint(nn.Module): ) del y_mask + blocks_replace = patches_replace.get("dit", {}) for i, block in enumerate(self.blocks): - x, y_feat = block( - x, - c, - y_feat, - rope_cos=rope_cos, - rope_sin=rope_sin, - crop_y=num_tokens, - ) # (B, M, D), (B, L, D) + if ("double_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["img"], out["txt"] = block( + args["img"], + args["vec"], + args["txt"], + rope_cos=args["rope_cos"], + rope_sin=args["rope_sin"], + crop_y=args["num_tokens"] + ) + return out + out = blocks_replace[("double_block", i)]({"img": x, "txt": y_feat, "vec": c, "rope_cos": rope_cos, "rope_sin": rope_sin, "num_tokens": num_tokens}, {"original_block": block_wrap}) + y_feat = out["txt"] + x = out["img"] + else: + x, y_feat = block( + x, + c, + y_feat, + rope_cos=rope_cos, + rope_sin=rope_sin, + crop_y=num_tokens, + ) # (B, M, D), (B, L, D) del y_feat # Final layers don't use dense text features. x = self.final_layer(x, c) # (B, M, patch_size ** 2 * out_channels) From d9f90965c8601b671faca9c8784b529e389c49aa Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 17 Nov 2024 08:19:59 -0500 Subject: [PATCH 3/3] Support block replace patches in auraflow. --- comfy/ldm/aura/mmdit.py | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/comfy/ldm/aura/mmdit.py b/comfy/ldm/aura/mmdit.py index cd9a42185..77090372e 100644 --- a/comfy/ldm/aura/mmdit.py +++ b/comfy/ldm/aura/mmdit.py @@ -437,7 +437,8 @@ class MMDiT(nn.Module): pos_encoding = pos_encoding[:,from_h:from_h+h,from_w:from_w+w] return x + pos_encoding.reshape(1, -1, self.positional_encoding.shape[-1]) - def forward(self, x, timestep, context, **kwargs): + def forward(self, x, timestep, context, transformer_options={}, **kwargs): + patches_replace = transformer_options.get("patches_replace", {}) # patchify x, add PE b, c, h, w = x.shape @@ -458,15 +459,36 @@ class MMDiT(nn.Module): global_cond = self.t_embedder(t, x.dtype) # B, D + blocks_replace = patches_replace.get("dit", {}) if len(self.double_layers) > 0: - for layer in self.double_layers: - c, x = layer(c, x, global_cond, **kwargs) + for i, layer in enumerate(self.double_layers): + if ("double_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["txt"], out["img"] = layer(args["txt"], + args["img"], + args["vec"]) + return out + out = blocks_replace[("double_block", i)]({"img": x, "txt": c, "vec": global_cond}, {"original_block": block_wrap}) + c = out["txt"] + x = out["img"] + else: + c, x = layer(c, x, global_cond, **kwargs) if len(self.single_layers) > 0: c_len = c.size(1) cx = torch.cat([c, x], dim=1) - for layer in self.single_layers: - cx = layer(cx, global_cond, **kwargs) + for i, layer in enumerate(self.single_layers): + if ("single_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["img"] = layer(args["img"], args["vec"]) + return out + + out = blocks_replace[("single_block", i)]({"img": cx, "vec": global_cond}, {"original_block": block_wrap}) + cx = out["img"] + else: + cx = layer(cx, global_cond, **kwargs) x = cx[:, c_len:]