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

Cleanup unnecessary & unused helper functions

Browse Source 
Cleanup unnecessary & unused helper functions
This commit is contained in:
John Slegers 2024-08-12 01:28:30 +02:00
parent 1de69fe4d5
commit 87e3f6eb35
5 changed files with 26 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
comfy/gligen.py
View File

@ -5,19 +5,6 @@ from inspect import isfunction
import comfy.ops import comfy.ops
ops = comfy.ops.manual_cast 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 # feedforward
class GEGLU(nn.Module): 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.): def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
super().__init__() super().__init__()
inner_dim = int(dim * mult) 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( project_in = nn.Sequential(
ops.Linear(dim, inner_dim), ops.Linear(dim, inner_dim),
nn.GELU() nn.GELU()

47
comfy/ldm/modules/attention.py
View File

@ -24,33 +24,9 @@ FORCE_UPCAST_ATTENTION_DTYPE = model_management.force_upcast_attention_dtype()
def get_attn_precision(attn_precision): def get_attn_precision(attn_precision):
if args.dont_upcast_attention: if args.dont_upcast_attention:
return None return None
if FORCE_UPCAST_ATTENTION_DTYPE is not None: if FORCE_UPCAST_ATTENTION_DTYPE is None:
return FORCE_UPCAST_ATTENTION_DTYPE return attn_precision
return attn_precision return FORCE_UPCAST_ATTENTION_DTYPE
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
# feedforward # 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): def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., dtype=None, device=None, operations=ops):
super().__init__() super().__init__()
inner_dim = int(dim * mult) inner_dim = int(dim * mult)
dim_out = default(dim_out, dim) if dim_out is None:
dim_out = dim
project_in = nn.Sequential( project_in = nn.Sequential(
operations.Linear(dim, inner_dim, dtype=dtype, device=device), operations.Linear(dim, inner_dim, dtype=dtype, device=device),
nn.GELU() nn.GELU()
@ -121,7 +98,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
del q, k del q, k
if exists(mask): if mask is not None:
if mask.dtype == torch.bool: if mask.dtype == torch.bool:
mask = rearrange(mask, 'b ... -> b (...)') #TODO: check if this bool part matches pytorch attention mask = rearrange(mask, 'b ... -> b (...)') #TODO: check if this bool part matches pytorch attention
max_neg_value = -torch.finfo(sim.dtype).max 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): 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__() super().__init__()
inner_dim = dim_head * heads 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.attn_precision = attn_precision
self.heads = heads self.heads = heads
@ -463,7 +441,8 @@ class CrossAttention(nn.Module):
def forward(self, x, context=None, value=None, mask=None): def forward(self, x, context=None, value=None, mask=None):
q = self.to_q(x) q = self.to_q(x)
context = default(context, x) if context is None:
context = x
k = self.to_k(context) k = self.to_k(context)
if value is not None: if value is not None:
v = self.to_v(value) v = self.to_v(value)
@ -649,7 +628,7 @@ class SpatialTransformer(nn.Module):
disable_self_attn=False, use_linear=False, disable_self_attn=False, use_linear=False,
use_checkpoint=True, attn_precision=None, dtype=None, device=None, operations=ops): use_checkpoint=True, attn_precision=None, dtype=None, device=None, operations=ops):
super().__init__() 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 context_dim = [context_dim] * depth
self.in_channels = in_channels self.in_channels = in_channels
inner_dim = n_heads * d_head inner_dim = n_heads * d_head
@ -799,7 +778,7 @@ class SpatialVideoTransformer(SpatialTransformer):
_, _, h, w = x.shape _, _, h, w = x.shape
x_in = x x_in = x
spatial_context = None spatial_context = None
if exists(context): if context is not None:
spatial_context = context spatial_context = context
if self.use_spatial_context: if self.use_spatial_context:
@ -861,5 +840,3 @@ class SpatialVideoTransformer(SpatialTransformer):
x = self.proj_out(x) x = self.proj_out(x)
out = x + x_in out = x + x_in
return out return out

13
comfy/ldm/modules/diffusionmodules/mmdit.py
View File

@ -11,10 +11,6 @@ from .util import timestep_embedding
import comfy.ops import comfy.ops
import comfy.ldm.common_dit import comfy.ldm.common_dit
def default(x, y):
if x is not None:
return x
return y
class Mlp(nn.Module): class Mlp(nn.Module):
""" MLP as used in Vision Transformer, MLP-Mixer and related networks """ 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.learn_sigma = learn_sigma
self.in_channels = in_channels self.in_channels = in_channels
default_out_channels = in_channels * 2 if learn_sigma else 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.patch_size = patch_size
self.pos_embed_scaling_factor = pos_embed_scaling_factor self.pos_embed_scaling_factor = pos_embed_scaling_factor
self.pos_embed_offset = pos_embed_offset self.pos_embed_offset = pos_embed_offset
self.pos_embed_max_size = pos_embed_max_size self.pos_embed_max_size = pos_embed_max_size
# hidden_size = default(hidden_size, 64 * depth) # hidden_size = 64 * depth if hidden_size is None else hidden_size
# num_heads = default(num_heads, hidden_size // 64) # num_heads = hidden_size // 64 if num_heads is None else num_heads
# apply magic --> this defines a head_size of 64 # apply magic --> this defines a head_size of 64
self.hidden_size = 64 * depth self.hidden_size = 64 * depth
@ -882,7 +878,7 @@ class MMDiT(nn.Module):
context = torch.cat( context = torch.cat(
( (
repeat(self.register, "1 ... -> b ...", b=x.shape[0]), 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, 1,
) )
@ -952,4 +948,3 @@ class OpenAISignatureMMDITWrapper(MMDiT):
**kwargs, **kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
return super().forward(x, timesteps, context=context, y=y, control=control) return super().forward(x, timesteps, context=context, y=y, control=control)

15
comfy/ldm/modules/diffusionmodules/openaimodel.py
View File

@ -13,8 +13,7 @@ from .util import (
timestep_embedding, timestep_embedding,
AlphaBlender, AlphaBlender,
) )
from ..attention import SpatialTransformer, SpatialVideoTransformer, default from ..attention import SpatialTransformer, SpatialVideoTransformer
from comfy.ldm.util import exists
import comfy.ops import comfy.ops
ops = comfy.ops.disable_weight_init ops = comfy.ops.disable_weight_init
@ -301,11 +300,11 @@ class VideoResBlock(ResBlock):
) )
self.time_stack = ResBlock( self.time_stack = ResBlock(
default(out_channels, channels), channels if out_channels is None else out_channels,
emb_channels, emb_channels,
dropout=dropout, dropout=dropout,
dims=3, 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_scale_shift_norm=False,
use_conv=False, use_conv=False,
up=False, up=False,
@ -642,12 +641,12 @@ class UNetModel(nn.Module):
if legacy: if legacy:
#num_heads = 1 #num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels 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] disabled_sa = disable_self_attentions[level]
else: else:
disabled_sa = False 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( layers.append(get_attention_layer(
ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
disable_self_attn=disabled_sa, use_checkpoint=use_checkpoint) disable_self_attn=disabled_sa, use_checkpoint=use_checkpoint)
@ -768,12 +767,12 @@ class UNetModel(nn.Module):
if legacy: if legacy:
#num_heads = 1 #num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels 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] disabled_sa = disable_self_attentions[level]
else: else:
disabled_sa = False 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( layers.append(
get_attention_layer( get_attention_layer(
ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim, ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,

12
comfy/ldm/util.py
View File

@ -44,16 +44,6 @@ def isimage(x):
return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1) 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): def mean_flat(tensor):
""" """
https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86 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): 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) ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
return loss return loss