xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-14 10:57:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'comfyanonymous:master' into add-xla-support

Browse Source
This commit is contained in:
HGB 2024-11-18 03:59:41 +07:00 committed by GitHub
commit d7c88efa58
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 55 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
app/user_manager.py
View File

@ -1,3 +1,4 @@
from __future__ import annotations
import json import json
import os import os
import re import re

28
comfy/ldm/aura/mmdit.py
View File

@ -437,7 +437,8 @@ class MMDiT(nn.Module):
pos_encoding = pos_encoding[:,from_h:from_h+h,from_w:from_w+w] 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]) 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 # patchify x, add PE
b, c, h, w = x.shape b, c, h, w = x.shape
@ -458,14 +459,35 @@ class MMDiT(nn.Module):
global_cond = self.t_embedder(t, x.dtype) # B, D global_cond = self.t_embedder(t, x.dtype) # B, D
blocks_replace = patches_replace.get("dit", {})
if len(self.double_layers) > 0: if len(self.double_layers) > 0:
for layer in self.double_layers: 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) c, x = layer(c, x, global_cond, **kwargs)
if len(self.single_layers) > 0: if len(self.single_layers) > 0:
c_len = c.size(1) c_len = c.size(1)
cx = torch.cat([c, x], dim=1) cx = torch.cat([c, x], dim=1)
for layer in self.single_layers: 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) cx = layer(cx, global_cond, **kwargs)
x = cx[:, c_len:] x = cx[:, c_len:]

20
comfy/ldm/genmo/joint_model/asymm_models_joint.py
View File

@ -494,8 +494,9 @@ class AsymmDiTJoint(nn.Module):
packed_indices: Dict[str, torch.Tensor] = None, packed_indices: Dict[str, torch.Tensor] = None,
rope_cos: torch.Tensor = None, rope_cos: torch.Tensor = None,
rope_sin: 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_feat = context
y_mask = attention_mask y_mask = attention_mask
sigma = timestep sigma = timestep
@ -515,7 +516,24 @@ class AsymmDiTJoint(nn.Module):
) )
del y_mask del y_mask
blocks_replace = patches_replace.get("dit", {})
for i, block in enumerate(self.blocks): for i, block in enumerate(self.blocks):
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, y_feat = block(
x, x,
c, c,