diff --git a/comfy/gligen.py b/comfy/gligen.py index 592522767..acb1fa67d 100644 --- a/comfy/gligen.py +++ b/comfy/gligen.py @@ -5,19 +5,6 @@ from inspect import isfunction import comfy.ops ops = comfy.ops.manual_cast -def exists(val): - return val is not None - - -def uniq(arr): - return{el: True for el in arr}.keys() - - -def default(val, d): - if exists(val): - return val - return d() if isfunction(d) else d - # feedforward class GEGLU(nn.Module): @@ -34,7 +21,8 @@ class FeedForward(nn.Module): def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.): super().__init__() inner_dim = int(dim * mult) - dim_out = default(dim_out, dim) + if not dim_out: + dim_out = dim() if isfunction(dim) else dim project_in = nn.Sequential( ops.Linear(dim, inner_dim), nn.GELU() diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 65a8bcf42..4750a3a9a 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -24,33 +24,9 @@ FORCE_UPCAST_ATTENTION_DTYPE = model_management.force_upcast_attention_dtype() def get_attn_precision(attn_precision): if args.dont_upcast_attention: return None - if FORCE_UPCAST_ATTENTION_DTYPE is not None: - return FORCE_UPCAST_ATTENTION_DTYPE - return attn_precision - -def exists(val): - return val is not None - - -def uniq(arr): - return{el: True for el in arr}.keys() - - -def default(val, d): - if exists(val): - return val - return d - - -def max_neg_value(t): - return -torch.finfo(t.dtype).max - - -def init_(tensor): - dim = tensor.shape[-1] - std = 1 / math.sqrt(dim) - tensor.uniform_(-std, std) - return tensor + if FORCE_UPCAST_ATTENTION_DTYPE is None: + return attn_precision + return FORCE_UPCAST_ATTENTION_DTYPE # feedforward @@ -68,7 +44,8 @@ class FeedForward(nn.Module): def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., dtype=None, device=None, operations=ops): super().__init__() inner_dim = int(dim * mult) - dim_out = default(dim_out, dim) + if dim_out is None: + dim_out = dim project_in = nn.Sequential( operations.Linear(dim, inner_dim, dtype=dtype, device=device), nn.GELU() @@ -121,7 +98,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape del q, k - if exists(mask): + if mask is not None: if mask.dtype == torch.bool: mask = rearrange(mask, 'b ... -> b (...)') #TODO: check if this bool part matches pytorch attention max_neg_value = -torch.finfo(sim.dtype).max @@ -449,7 +426,8 @@ class CrossAttention(nn.Module): def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., attn_precision=None, dtype=None, device=None, operations=ops): super().__init__() inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) + if context_dim is None: + context_dim = query_dim self.attn_precision = attn_precision self.heads = heads @@ -463,7 +441,8 @@ class CrossAttention(nn.Module): def forward(self, x, context=None, value=None, mask=None): q = self.to_q(x) - context = default(context, x) + if context is None: + context = x k = self.to_k(context) if value is not None: v = self.to_v(value) @@ -649,7 +628,7 @@ class SpatialTransformer(nn.Module): disable_self_attn=False, use_linear=False, use_checkpoint=True, attn_precision=None, dtype=None, device=None, operations=ops): super().__init__() - if exists(context_dim) and not isinstance(context_dim, list): + if context_dim is not None and not isinstance(context_dim, list): context_dim = [context_dim] * depth self.in_channels = in_channels inner_dim = n_heads * d_head @@ -799,7 +778,7 @@ class SpatialVideoTransformer(SpatialTransformer): _, _, h, w = x.shape x_in = x spatial_context = None - if exists(context): + if context is not None: spatial_context = context if self.use_spatial_context: @@ -861,5 +840,3 @@ class SpatialVideoTransformer(SpatialTransformer): x = self.proj_out(x) out = x + x_in return out - - diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py index 491a58a20..ee3ae0be5 100644 --- a/comfy/ldm/modules/diffusionmodules/mmdit.py +++ b/comfy/ldm/modules/diffusionmodules/mmdit.py @@ -11,10 +11,6 @@ from .util import timestep_embedding import comfy.ops import comfy.ldm.common_dit -def default(x, y): - if x is not None: - return x - return y class Mlp(nn.Module): """ MLP as used in Vision Transformer, MLP-Mixer and related networks @@ -733,14 +729,14 @@ class MMDiT(nn.Module): self.learn_sigma = learn_sigma self.in_channels = in_channels default_out_channels = in_channels * 2 if learn_sigma else in_channels - self.out_channels = default(out_channels, default_out_channels) + self.out_channels = default_out_channels if out_channels is None else out_channels self.patch_size = patch_size self.pos_embed_scaling_factor = pos_embed_scaling_factor self.pos_embed_offset = pos_embed_offset self.pos_embed_max_size = pos_embed_max_size - # hidden_size = default(hidden_size, 64 * depth) - # num_heads = default(num_heads, hidden_size // 64) + # hidden_size = 64 * depth if hidden_size is None else hidden_size + # num_heads = hidden_size // 64 if num_heads is None else num_heads # apply magic --> this defines a head_size of 64 self.hidden_size = 64 * depth @@ -882,7 +878,7 @@ class MMDiT(nn.Module): context = torch.cat( ( repeat(self.register, "1 ... -> b ...", b=x.shape[0]), - default(context, torch.Tensor([]).type_as(x)), + torch.Tensor([]).type_as(x) is context is None else context, ), 1, ) @@ -952,4 +948,3 @@ class OpenAISignatureMMDITWrapper(MMDiT): **kwargs, ) -> torch.Tensor: return super().forward(x, timesteps, context=context, y=y, control=control) - diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index 6535e899c..ac08bf547 100644 --- a/comfy/ldm/modules/diffusionmodules/openaimodel.py +++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py @@ -13,8 +13,7 @@ from .util import ( timestep_embedding, AlphaBlender, ) -from ..attention import SpatialTransformer, SpatialVideoTransformer, default -from comfy.ldm.util import exists +from ..attention import SpatialTransformer, SpatialVideoTransformer import comfy.ops ops = comfy.ops.disable_weight_init @@ -301,11 +300,11 @@ class VideoResBlock(ResBlock): ) self.time_stack = ResBlock( - default(out_channels, channels), + channels if out_channels is None else out_channels, emb_channels, dropout=dropout, dims=3, - out_channels=default(out_channels, channels), + out_channels=channels if out_channels is None else out_channels, use_scale_shift_norm=False, use_conv=False, up=False, @@ -642,12 +641,12 @@ class UNetModel(nn.Module): if legacy: #num_heads = 1 dim_head = ch // num_heads if use_spatial_transformer else num_head_channels - if exists(disable_self_attentions): + if disable_self_attentions is not None: disabled_sa = disable_self_attentions[level] else: disabled_sa = False - if not exists(num_attention_blocks) or nr < num_attention_blocks[level]: + if not num_attention_blocks is None or nr < num_attention_blocks[level]: layers.append(get_attention_layer( ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, disable_self_attn=disabled_sa, use_checkpoint=use_checkpoint) @@ -768,12 +767,12 @@ class UNetModel(nn.Module): if legacy: #num_heads = 1 dim_head = ch // num_heads if use_spatial_transformer else num_head_channels - if exists(disable_self_attentions): + if disable_self_attentions is not None: disabled_sa = disable_self_attentions[level] else: disabled_sa = False - if not exists(num_attention_blocks) or i < num_attention_blocks[level]: + if num_attention_blocks is None or i < num_attention_blocks[level]: layers.append( get_attention_layer( ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, diff --git a/comfy/ldm/util.py b/comfy/ldm/util.py index 8c09ca1c7..aeb40876b 100644 --- a/comfy/ldm/util.py +++ b/comfy/ldm/util.py @@ -44,16 +44,6 @@ def isimage(x): return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1) -def exists(x): - return x is not None - - -def default(val, d): - if exists(val): - return val - return d() if isfunction(d) else d - - def mean_flat(tensor): """ https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86 @@ -194,4 +184,4 @@ class AdamWwithEMAandWings(optim.Optimizer): for param, ema_param in zip(params_with_grad, ema_params_with_grad): ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay) - return loss \ No newline at end of file + return loss