From 1a4bd9e9a6fc2b364ebb547dbd80736548cf9f5c Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Oct 2023 15:47:53 -0400 Subject: [PATCH 01/15] Refactor the attention functions. There's no reason for the whole CrossAttention object to be repeated when only the operation in the middle changes. --- comfy/ldm/modules/attention.py | 538 +++++++------------- comfy/ldm/modules/diffusionmodules/model.py | 13 - 2 files changed, 186 insertions(+), 365 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index fcae6b66a..3230cfaf5 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -94,253 +94,220 @@ def zero_module(module): def Normalize(in_channels, dtype=None, device=None): return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device) +def attention_basic(q, k, v, heads, mask=None): + h = heads + scale = (q.shape[-1] // heads) ** -0.5 + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) -class SpatialSelfAttention(nn.Module): - def __init__(self, in_channels): - super().__init__() - self.in_channels = in_channels + # force cast to fp32 to avoid overflowing + if _ATTN_PRECISION =="fp32": + with torch.autocast(enabled=False, device_type = 'cuda'): + q, k = q.float(), k.float() + sim = einsum('b i d, b j d -> b i j', q, k) * scale + else: + sim = einsum('b i d, b j d -> b i j', q, k) * scale - self.norm = Normalize(in_channels) - self.q = torch.nn.Conv2d(in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0) - self.k = torch.nn.Conv2d(in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0) - self.v = torch.nn.Conv2d(in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0) - self.proj_out = torch.nn.Conv2d(in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0) + del q, k - def forward(self, x): - h_ = x - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) + if exists(mask): + mask = rearrange(mask, 'b ... -> b (...)') + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, 'b j -> (b h) () j', h=h) + sim.masked_fill_(~mask, max_neg_value) - # compute attention - b,c,h,w = q.shape - q = rearrange(q, 'b c h w -> b (h w) c') - k = rearrange(k, 'b c h w -> b c (h w)') - w_ = torch.einsum('bij,bjk->bik', q, k) + # attention, what we cannot get enough of + sim = sim.softmax(dim=-1) - w_ = w_ * (int(c)**(-0.5)) - w_ = torch.nn.functional.softmax(w_, dim=2) - - # attend to values - v = rearrange(v, 'b c h w -> b c (h w)') - w_ = rearrange(w_, 'b i j -> b j i') - h_ = torch.einsum('bij,bjk->bik', v, w_) - h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h) - h_ = self.proj_out(h_) - - return x+h_ + out = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v) + out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + return out -class CrossAttentionBirchSan(nn.Module): - def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops): - super().__init__() - inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) +def attention_sub_quad(query, key, value, heads, mask=None): + scale = (query.shape[-1] // heads) ** -0.5 + query = query.unflatten(-1, (heads, -1)).transpose(1,2).flatten(end_dim=1) + key_t = key.transpose(1,2).unflatten(1, (heads, -1)).flatten(end_dim=1) + del key + value = value.unflatten(-1, (heads, -1)).transpose(1,2).flatten(end_dim=1) - self.scale = dim_head ** -0.5 - self.heads = heads + dtype = query.dtype + upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32 + if upcast_attention: + bytes_per_token = torch.finfo(torch.float32).bits//8 + else: + bytes_per_token = torch.finfo(query.dtype).bits//8 + batch_x_heads, q_tokens, _ = query.shape + _, _, k_tokens = key_t.shape + qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens - self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) + mem_free_total, mem_free_torch = model_management.get_free_memory(query.device, True) - self.to_out = nn.Sequential( - operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), - nn.Dropout(dropout) - ) + chunk_threshold_bytes = mem_free_torch * 0.5 #Using only this seems to work better on AMD - def forward(self, x, context=None, value=None, mask=None): - h = self.heads + kv_chunk_size_min = None - query = self.to_q(x) - context = default(context, x) - key = self.to_k(context) - if value is not None: - value = self.to_v(value) - else: - value = self.to_v(context) + #not sure at all about the math here + #TODO: tweak this + if mem_free_total > 8192 * 1024 * 1024 * 1.3: + query_chunk_size_x = 1024 * 4 + elif mem_free_total > 4096 * 1024 * 1024 * 1.3: + query_chunk_size_x = 1024 * 2 + else: + query_chunk_size_x = 1024 + kv_chunk_size_min_x = None + kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024 + if kv_chunk_size_x < 1024: + kv_chunk_size_x = None - del context, x + if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes: + # the big matmul fits into our memory limit; do everything in 1 chunk, + # i.e. send it down the unchunked fast-path + query_chunk_size = q_tokens + kv_chunk_size = k_tokens + else: + query_chunk_size = query_chunk_size_x + kv_chunk_size = kv_chunk_size_x + kv_chunk_size_min = kv_chunk_size_min_x - query = query.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1) - key_t = key.transpose(1,2).unflatten(1, (self.heads, -1)).flatten(end_dim=1) - del key - value = value.unflatten(-1, (self.heads, -1)).transpose(1,2).flatten(end_dim=1) + hidden_states = efficient_dot_product_attention( + query, + key_t, + value, + query_chunk_size=query_chunk_size, + kv_chunk_size=kv_chunk_size, + kv_chunk_size_min=kv_chunk_size_min, + use_checkpoint=False, + upcast_attention=upcast_attention, + ) - dtype = query.dtype - upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32 - if upcast_attention: - bytes_per_token = torch.finfo(torch.float32).bits//8 - else: - bytes_per_token = torch.finfo(query.dtype).bits//8 - batch_x_heads, q_tokens, _ = query.shape - _, _, k_tokens = key_t.shape - qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens + hidden_states = hidden_states.to(dtype) - mem_free_total, mem_free_torch = model_management.get_free_memory(query.device, True) + hidden_states = hidden_states.unflatten(0, (-1, heads)).transpose(1,2).flatten(start_dim=2) + return hidden_states - chunk_threshold_bytes = mem_free_torch * 0.5 #Using only this seems to work better on AMD +def attention_split(q, k, v, heads, mask=None): + scale = (q.shape[-1] // heads) ** -0.5 + h = heads + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) - kv_chunk_size_min = None + r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) - #not sure at all about the math here - #TODO: tweak this - if mem_free_total > 8192 * 1024 * 1024 * 1.3: - query_chunk_size_x = 1024 * 4 - elif mem_free_total > 4096 * 1024 * 1024 * 1.3: - query_chunk_size_x = 1024 * 2 - else: - query_chunk_size_x = 1024 - kv_chunk_size_min_x = None - kv_chunk_size_x = (int((chunk_threshold_bytes // (batch_x_heads * bytes_per_token * query_chunk_size_x)) * 2.0) // 1024) * 1024 - if kv_chunk_size_x < 1024: - kv_chunk_size_x = None + mem_free_total = model_management.get_free_memory(q.device) - if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes: - # the big matmul fits into our memory limit; do everything in 1 chunk, - # i.e. send it down the unchunked fast-path - query_chunk_size = q_tokens - kv_chunk_size = k_tokens - else: - query_chunk_size = query_chunk_size_x - kv_chunk_size = kv_chunk_size_x - kv_chunk_size_min = kv_chunk_size_min_x - - hidden_states = efficient_dot_product_attention( - query, - key_t, - value, - query_chunk_size=query_chunk_size, - kv_chunk_size=kv_chunk_size, - kv_chunk_size_min=kv_chunk_size_min, - use_checkpoint=self.training, - upcast_attention=upcast_attention, - ) - - hidden_states = hidden_states.to(dtype) - - hidden_states = hidden_states.unflatten(0, (-1, self.heads)).transpose(1,2).flatten(start_dim=2) - - out_proj, dropout = self.to_out - hidden_states = out_proj(hidden_states) - hidden_states = dropout(hidden_states) - - return hidden_states + gb = 1024 ** 3 + tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() + modifier = 3 if q.element_size() == 2 else 2.5 + mem_required = tensor_size * modifier + steps = 1 -class CrossAttentionDoggettx(nn.Module): - def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops): - super().__init__() - inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) + if mem_required > mem_free_total: + steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2))) + # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " + # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") - self.scale = dim_head ** -0.5 - self.heads = heads + if steps > 64: + max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 + raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). ' + f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free') - self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - - self.to_out = nn.Sequential( - operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), - nn.Dropout(dropout) - ) - - def forward(self, x, context=None, value=None, mask=None): - h = self.heads - - q_in = self.to_q(x) - context = default(context, x) - k_in = self.to_k(context) - if value is not None: - v_in = self.to_v(value) - del value - else: - v_in = self.to_v(context) - del context, x - - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in)) - del q_in, k_in, v_in - - r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) - - mem_free_total = model_management.get_free_memory(q.device) - - gb = 1024 ** 3 - tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() - modifier = 3 if q.element_size() == 2 else 2.5 - mem_required = tensor_size * modifier - steps = 1 - - - if mem_required > mem_free_total: - steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2))) - # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " - # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") - - if steps > 64: - max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 - raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). ' - f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free') - - # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size) - first_op_done = False - cleared_cache = False - while True: - try: - slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1] - for i in range(0, q.shape[1], slice_size): - end = i + slice_size - if _ATTN_PRECISION =="fp32": - with torch.autocast(enabled=False, device_type = 'cuda'): - s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * self.scale - else: - s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale - first_op_done = True - - s2 = s1.softmax(dim=-1).to(v.dtype) - del s1 - - r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v) - del s2 - break - except model_management.OOM_EXCEPTION as e: - if first_op_done == False: - model_management.soft_empty_cache(True) - if cleared_cache == False: - cleared_cache = True - print("out of memory error, emptying cache and trying again") - continue - steps *= 2 - if steps > 64: - raise e - print("out of memory error, increasing steps and trying again", steps) + # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size) + first_op_done = False + cleared_cache = False + while True: + try: + slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1] + for i in range(0, q.shape[1], slice_size): + end = i + slice_size + if _ATTN_PRECISION =="fp32": + with torch.autocast(enabled=False, device_type = 'cuda'): + s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale else: + s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale + first_op_done = True + + s2 = s1.softmax(dim=-1).to(v.dtype) + del s1 + + r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v) + del s2 + break + except model_management.OOM_EXCEPTION as e: + if first_op_done == False: + model_management.soft_empty_cache(True) + if cleared_cache == False: + cleared_cache = True + print("out of memory error, emptying cache and trying again") + continue + steps *= 2 + if steps > 64: raise e + print("out of memory error, increasing steps and trying again", steps) + else: + raise e - del q, k, v + del q, k, v - r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h) - del r1 + r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h) + del r1 + return r2 - return self.to_out(r2) +def attention_xformers(q, k, v, heads, mask=None): + b, _, _ = q.shape + q, k, v = map( + lambda t: t.unsqueeze(3) + .reshape(b, t.shape[1], heads, -1) + .permute(0, 2, 1, 3) + .reshape(b * heads, t.shape[1], -1) + .contiguous(), + (q, k, v), + ) + + # actually compute the attention, what we cannot get enough of + out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None) + + if exists(mask): + raise NotImplementedError + out = ( + out.unsqueeze(0) + .reshape(b, heads, out.shape[1], -1) + .permute(0, 2, 1, 3) + .reshape(b, out.shape[1], -1) + ) + return out + +def attention_pytorch(q, k, v, heads, mask=None): + b, _, dim_head = q.shape + dim_head //= heads + q, k, v = map( + lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), + (q, k, v), + ) + + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + + if exists(mask): + raise NotImplementedError + out = ( + out.transpose(1, 2).reshape(b, -1, heads * dim_head) + ) + return out + +optimized_attention = attention_basic + +if model_management.xformers_enabled(): + print("Using xformers cross attention") + optimized_attention = attention_xformers +elif model_management.pytorch_attention_enabled(): + print("Using pytorch cross attention") + optimized_attention = attention_pytorch +else: + if args.use_split_cross_attention: + print("Using split optimization for cross attention") + optimized_attention = attention_split + else: + print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention") + optimized_attention = attention_sub_quad class CrossAttention(nn.Module): def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops): @@ -348,62 +315,6 @@ class CrossAttention(nn.Module): inner_dim = dim_head * heads context_dim = default(context_dim, query_dim) - self.scale = dim_head ** -0.5 - self.heads = heads - - self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - - self.to_out = nn.Sequential( - operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), - nn.Dropout(dropout) - ) - - def forward(self, x, context=None, value=None, mask=None): - h = self.heads - - q = self.to_q(x) - context = default(context, x) - k = self.to_k(context) - if value is not None: - v = self.to_v(value) - del value - else: - v = self.to_v(context) - - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) - - # force cast to fp32 to avoid overflowing - if _ATTN_PRECISION =="fp32": - with torch.autocast(enabled=False, device_type = 'cuda'): - q, k = q.float(), k.float() - sim = einsum('b i d, b j d -> b i j', q, k) * self.scale - else: - sim = einsum('b i d, b j d -> b i j', q, k) * self.scale - - del q, k - - if exists(mask): - mask = rearrange(mask, 'b ... -> b (...)') - max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b j -> (b h) () j', h=h) - sim.masked_fill_(~mask, max_neg_value) - - # attention, what we cannot get enough of - sim = sim.softmax(dim=-1) - - out = einsum('b i j, b j d -> b i d', sim, v) - out = rearrange(out, '(b h) n d -> b n (h d)', h=h) - return self.to_out(out) - -class MemoryEfficientCrossAttention(nn.Module): - # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 - def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, dtype=None, device=None, operations=comfy.ops): - super().__init__() - inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) - self.heads = heads self.dim_head = dim_head @@ -412,7 +323,6 @@ class MemoryEfficientCrossAttention(nn.Module): self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout)) - self.attention_op: Optional[Any] = None def forward(self, x, context=None, value=None, mask=None): q = self.to_q(x) @@ -424,85 +334,9 @@ class MemoryEfficientCrossAttention(nn.Module): else: v = self.to_v(context) - b, _, _ = q.shape - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, t.shape[1], self.heads, self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b * self.heads, t.shape[1], self.dim_head) - .contiguous(), - (q, k, v), - ) - - # actually compute the attention, what we cannot get enough of - out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) - - if exists(mask): - raise NotImplementedError - out = ( - out.unsqueeze(0) - .reshape(b, self.heads, out.shape[1], self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b, out.shape[1], self.heads * self.dim_head) - ) + out = optimized_attention(q, k, v, self.heads, mask) return self.to_out(out) -class CrossAttentionPytorch(nn.Module): - def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=comfy.ops): - super().__init__() - inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) - - self.heads = heads - self.dim_head = dim_head - - self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device) - - self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout)) - self.attention_op: Optional[Any] = None - - def forward(self, x, context=None, value=None, mask=None): - q = self.to_q(x) - context = default(context, x) - k = self.to_k(context) - if value is not None: - v = self.to_v(value) - del value - else: - v = self.to_v(context) - - b, _, _ = q.shape - q, k, v = map( - lambda t: t.view(b, -1, self.heads, self.dim_head).transpose(1, 2), - (q, k, v), - ) - - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) - - if exists(mask): - raise NotImplementedError - out = ( - out.transpose(1, 2).reshape(b, -1, self.heads * self.dim_head) - ) - - return self.to_out(out) - -if model_management.xformers_enabled(): - print("Using xformers cross attention") - CrossAttention = MemoryEfficientCrossAttention -elif model_management.pytorch_attention_enabled(): - print("Using pytorch cross attention") - CrossAttention = CrossAttentionPytorch -else: - if args.use_split_cross_attention: - print("Using split optimization for cross attention") - CrossAttention = CrossAttentionDoggettx - else: - print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention") - CrossAttention = CrossAttentionBirchSan - class BasicTransformerBlock(nn.Module): def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py index 5f38640c3..e6cf954ff 100644 --- a/comfy/ldm/modules/diffusionmodules/model.py +++ b/comfy/ldm/modules/diffusionmodules/model.py @@ -6,7 +6,6 @@ import numpy as np from einops import rearrange from typing import Optional, Any -from ..attention import MemoryEfficientCrossAttention from comfy import model_management import comfy.ops @@ -352,15 +351,6 @@ class MemoryEfficientAttnBlockPytorch(nn.Module): out = self.proj_out(out) return x+out -class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention): - def forward(self, x, context=None, mask=None): - b, c, h, w = x.shape - x = rearrange(x, 'b c h w -> b (h w) c') - out = super().forward(x, context=context, mask=mask) - out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w, c=c) - return x + out - - def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown' if model_management.xformers_enabled_vae() and attn_type == "vanilla": @@ -376,9 +366,6 @@ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): return MemoryEfficientAttnBlock(in_channels) elif attn_type == "vanilla-pytorch": return MemoryEfficientAttnBlockPytorch(in_channels) - elif type == "memory-efficient-cross-attn": - attn_kwargs["query_dim"] = in_channels - return MemoryEfficientCrossAttentionWrapper(**attn_kwargs) elif attn_type == "none": return nn.Identity(in_channels) else: From ac7d8cfa875e08623993da8109cc73a68df42379 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Oct 2023 20:24:17 -0400 Subject: [PATCH 02/15] Allow attn_mask in attention_pytorch. --- comfy/ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 3230cfaf5..ac0d9c8c5 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -284,7 +284,7 @@ def attention_pytorch(q, k, v, heads, mask=None): (q, k, v), ) - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) if exists(mask): raise NotImplementedError From 20d3852aa1a23db24b288531ef07c26b49308629 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Oct 2023 20:35:50 -0400 Subject: [PATCH 03/15] Pull some small changes from the other repo. --- comfy/model_management.py | 5 +++-- comfy/utils.py | 4 ++++ comfy_extras/nodes_custom_sampler.py | 3 ++- execution.py | 19 ++++++++++--------- folder_paths.py | 2 ++ nodes.py | 2 +- 6 files changed, 22 insertions(+), 13 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 8b8963726..3b43b21ac 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -354,6 +354,8 @@ def load_models_gpu(models, memory_required=0): current_loaded_models.insert(0, current_loaded_models.pop(index)) models_already_loaded.append(loaded_model) else: + if hasattr(x, "model"): + print(f"Requested to load {x.model.__class__.__name__}") models_to_load.append(loaded_model) if len(models_to_load) == 0: @@ -363,7 +365,7 @@ def load_models_gpu(models, memory_required=0): free_memory(extra_mem, d, models_already_loaded) return - print("loading new") + print(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}") total_memory_required = {} for loaded_model in models_to_load: @@ -405,7 +407,6 @@ def load_model_gpu(model): def cleanup_models(): to_delete = [] for i in range(len(current_loaded_models)): - print(sys.getrefcount(current_loaded_models[i].model)) if sys.getrefcount(current_loaded_models[i].model) <= 2: to_delete = [i] + to_delete diff --git a/comfy/utils.py b/comfy/utils.py index 7843b58cc..df016ef9e 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -408,6 +408,10 @@ def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_am output[b:b+1] = out/out_div return output +PROGRESS_BAR_ENABLED = True +def set_progress_bar_enabled(enabled): + global PROGRESS_BAR_ENABLED + PROGRESS_BAR_ENABLED = enabled PROGRESS_BAR_HOOK = None def set_progress_bar_global_hook(function): diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 9391c7147..b52ad8fbd 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -3,6 +3,7 @@ import comfy.sample from comfy.k_diffusion import sampling as k_diffusion_sampling import latent_preview import torch +import comfy.utils class BasicScheduler: @@ -219,7 +220,7 @@ class SamplerCustom: x0_output = {} callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output) - disable_pbar = False + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed) out = latent.copy() diff --git a/execution.py b/execution.py index 5f5d6c738..918c2bc5c 100644 --- a/execution.py +++ b/execution.py @@ -2,6 +2,7 @@ import os import sys import copy import json +import logging import threading import heapq import traceback @@ -156,7 +157,7 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute if server.client_id is not None: server.send_sync("executed", { "node": unique_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id) except comfy.model_management.InterruptProcessingException as iex: - print("Processing interrupted") + logging.info("Processing interrupted") # skip formatting inputs/outputs error_details = { @@ -177,8 +178,8 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute for node_id, node_outputs in outputs.items(): output_data_formatted[node_id] = [[format_value(x) for x in l] for l in node_outputs] - print("!!! Exception during processing !!!") - print(traceback.format_exc()) + logging.error("!!! Exception during processing !!!") + logging.error(traceback.format_exc()) error_details = { "node_id": unique_id, @@ -636,11 +637,11 @@ def validate_prompt(prompt): if valid is True: good_outputs.add(o) else: - print(f"Failed to validate prompt for output {o}:") + logging.error(f"Failed to validate prompt for output {o}:") if len(reasons) > 0: - print("* (prompt):") + logging.error("* (prompt):") for reason in reasons: - print(f" - {reason['message']}: {reason['details']}") + logging.error(f" - {reason['message']}: {reason['details']}") errors += [(o, reasons)] for node_id, result in validated.items(): valid = result[0] @@ -656,11 +657,11 @@ def validate_prompt(prompt): "dependent_outputs": [], "class_type": class_type } - print(f"* {class_type} {node_id}:") + logging.error(f"* {class_type} {node_id}:") for reason in reasons: - print(f" - {reason['message']}: {reason['details']}") + logging.error(f" - {reason['message']}: {reason['details']}") node_errors[node_id]["dependent_outputs"].append(o) - print("Output will be ignored") + logging.error("Output will be ignored") if len(good_outputs) == 0: errors_list = [] diff --git a/folder_paths.py b/folder_paths.py index 898513b0e..5d121b443 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -29,6 +29,8 @@ folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions) +folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""}) + output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") diff --git a/nodes.py b/nodes.py index 16bf07cca..208cbc84f 100644 --- a/nodes.py +++ b/nodes.py @@ -1202,7 +1202,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, noise_mask = latent["noise_mask"] callback = latent_preview.prepare_callback(model, steps) - disable_pbar = False + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) From 88733c997fd807a572d4a214d2c15fc5dd17b3c6 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Oct 2023 21:29:03 -0400 Subject: [PATCH 04/15] pytorch_attention_enabled can now return True when xformers is enabled. --- comfy/ldm/modules/diffusionmodules/model.py | 2 +- comfy/model_management.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py index e6cf954ff..6576df4b6 100644 --- a/comfy/ldm/modules/diffusionmodules/model.py +++ b/comfy/ldm/modules/diffusionmodules/model.py @@ -355,7 +355,7 @@ def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown' if model_management.xformers_enabled_vae() and attn_type == "vanilla": attn_type = "vanilla-xformers" - if model_management.pytorch_attention_enabled() and attn_type == "vanilla": + elif model_management.pytorch_attention_enabled() and attn_type == "vanilla": attn_type = "vanilla-pytorch" print(f"making attention of type '{attn_type}' with {in_channels} in_channels") if attn_type == "vanilla": diff --git a/comfy/model_management.py b/comfy/model_management.py index 3b43b21ac..3c390d9ca 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -154,14 +154,18 @@ def is_nvidia(): return True return False -ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention +ENABLE_PYTORCH_ATTENTION = False +if args.use_pytorch_cross_attention: + ENABLE_PYTORCH_ATTENTION = True + XFORMERS_IS_AVAILABLE = False + VAE_DTYPE = torch.float32 try: if is_nvidia(): torch_version = torch.version.__version__ if int(torch_version[0]) >= 2: - if ENABLE_PYTORCH_ATTENTION == False and XFORMERS_IS_AVAILABLE == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: ENABLE_PYTORCH_ATTENTION = True if torch.cuda.is_bf16_supported(): VAE_DTYPE = torch.bfloat16 @@ -186,7 +190,6 @@ if ENABLE_PYTORCH_ATTENTION: torch.backends.cuda.enable_math_sdp(True) torch.backends.cuda.enable_flash_sdp(True) torch.backends.cuda.enable_mem_efficient_sdp(True) - XFORMERS_IS_AVAILABLE = False if args.lowvram: set_vram_to = VRAMState.LOW_VRAM From 41d2c5660dab2404c7db78ff14c7faecd6c57e1f Mon Sep 17 00:00:00 2001 From: Chris Date: Thu, 12 Oct 2023 14:26:53 +1100 Subject: [PATCH 05/15] add query --- web/scripts/app.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/scripts/app.js b/web/scripts/app.js index 7698d0f11..3cf3585d2 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -1592,7 +1592,7 @@ export class ComfyApp { all_inputs = all_inputs.concat(Object.keys(parent.inputs)) for (let parent_input in all_inputs) { parent_input = all_inputs[parent_input]; - if (parent.inputs[parent_input].type === node.inputs[i].type) { + if (parent.inputs[parent_input]?.type === node.inputs[i].type) { link = parent.getInputLink(parent_input); if (link) { parent = parent.getInputNode(parent_input); From 851a4bdb803871d62ceca4249036da34449891ed Mon Sep 17 00:00:00 2001 From: Nick Teeple <64399276+nickteeple@users.noreply.github.com> Date: Thu, 12 Oct 2023 21:26:27 +0800 Subject: [PATCH 06/15] Update extra_model_paths.yaml.example with comfy specific example --- extra_model_paths.yaml.example | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index 36078fffc..0870811f3 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -1,5 +1,20 @@ #Rename this to extra_model_paths.yaml and ComfyUI will load it +#config for comfyui +#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc. +#in this example, we make a central folder called c:/machine-learning/ and replicate the ComfyUI models folder structure inside +comfyui: + base_path: c:/machine-learning/ + checkpoints: models/checkpoints/ + clip: models/clip/ + clip_vision: models/clip_vision/ + configs: models/configs/ + controlnet: models/controlnet/ + embeddings: models/embeddings/ + loras: models/loras/ + upscale_models: models/upscale_models/ + vae: models/vae/ + #config for a1111 ui #all you have to do is change the base_path to where yours is installed a111: From fee3b0c0700aedaafacd72ef90d49c7be8c1a003 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 12 Oct 2023 20:54:43 -0400 Subject: [PATCH 07/15] Move and comment out. --- extra_model_paths.yaml.example | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index 0870811f3..846d04dbe 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -1,19 +1,5 @@ #Rename this to extra_model_paths.yaml and ComfyUI will load it -#config for comfyui -#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc. -#in this example, we make a central folder called c:/machine-learning/ and replicate the ComfyUI models folder structure inside -comfyui: - base_path: c:/machine-learning/ - checkpoints: models/checkpoints/ - clip: models/clip/ - clip_vision: models/clip_vision/ - configs: models/configs/ - controlnet: models/controlnet/ - embeddings: models/embeddings/ - loras: models/loras/ - upscale_models: models/upscale_models/ - vae: models/vae/ #config for a1111 ui #all you have to do is change the base_path to where yours is installed @@ -34,6 +20,21 @@ a111: hypernetworks: models/hypernetworks controlnet: models/ControlNet +#config for comfyui +#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc. + +#comfyui: +# base_path: path/to/comfyui/ +# checkpoints: models/checkpoints/ +# clip: models/clip/ +# clip_vision: models/clip_vision/ +# configs: models/configs/ +# controlnet: models/controlnet/ +# embeddings: models/embeddings/ +# loras: models/loras/ +# upscale_models: models/upscale_models/ +# vae: models/vae/ + #other_ui: # base_path: path/to/ui # checkpoints: models/checkpoints From 87097a11c32ac7393093055c4ed1432553995560 Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Fri, 13 Oct 2023 12:26:54 -0300 Subject: [PATCH 08/15] Fix FeatherMask --- comfy_extras/nodes_mask.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py index 9b0b289c1..3cc4cfd76 100644 --- a/comfy_extras/nodes_mask.py +++ b/comfy_extras/nodes_mask.py @@ -282,10 +282,10 @@ class FeatherMask: def feather(self, mask, left, top, right, bottom): output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone() - left = min(left, output.shape[1]) - right = min(right, output.shape[1]) - top = min(top, output.shape[0]) - bottom = min(bottom, output.shape[0]) + left = min(left, output.shape[-1]) + right = min(right, output.shape[-1]) + top = min(top, output.shape[-2]) + bottom = min(bottom, output.shape[-2]) for x in range(left): feather_rate = (x + 1.0) / left From b5fa3d28d7a9242a73c40b3b37e526267cbd67fd Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Fri, 13 Oct 2023 13:40:53 -0300 Subject: [PATCH 09/15] Fix MaskComposite --- comfy_extras/nodes_mask.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py index 3cc4cfd76..d8c65c2b6 100644 --- a/comfy_extras/nodes_mask.py +++ b/comfy_extras/nodes_mask.py @@ -240,8 +240,8 @@ class MaskComposite: right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2])) visible_width, visible_height = (right - left, bottom - top,) - source_portion = source[:visible_height, :visible_width] - destination_portion = destination[top:bottom, left:right] + source_portion = source[:, :visible_height, :visible_width] + destination_portion = destination[:, top:bottom, left:right] if operation == "multiply": output[:, top:bottom, left:right] = destination_portion * source_portion From 9a55dadb4c9c82d772d934125112b8b0613e4f28 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 13 Oct 2023 14:35:21 -0400 Subject: [PATCH 10/15] Refactor code so model can be a dtype other than fp32 or fp16. --- comfy/cldm/cldm.py | 6 ++-- comfy/controlnet.py | 11 +++---- .../modules/diffusionmodules/openaimodel.py | 6 ++-- comfy/model_detection.py | 32 +++++++++---------- comfy/model_management.py | 5 +++ comfy/sd.py | 20 ++++++------ 6 files changed, 39 insertions(+), 41 deletions(-) diff --git a/comfy/cldm/cldm.py b/comfy/cldm/cldm.py index 251483131..f982d648c 100644 --- a/comfy/cldm/cldm.py +++ b/comfy/cldm/cldm.py @@ -34,8 +34,7 @@ class ControlNet(nn.Module): dims=2, num_classes=None, use_checkpoint=False, - use_fp16=False, - use_bf16=False, + dtype=torch.float32, num_heads=-1, num_head_channels=-1, num_heads_upsample=-1, @@ -108,8 +107,7 @@ class ControlNet(nn.Module): self.conv_resample = conv_resample self.num_classes = num_classes self.use_checkpoint = use_checkpoint - self.dtype = th.float16 if use_fp16 else th.float32 - self.dtype = th.bfloat16 if use_bf16 else self.dtype + self.dtype = dtype self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample diff --git a/comfy/controlnet.py b/comfy/controlnet.py index ea219c7e5..73a40acfa 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -292,8 +292,8 @@ def load_controlnet(ckpt_path, model=None): controlnet_config = None if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format - use_fp16 = comfy.model_management.should_use_fp16() - controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data, use_fp16) + unet_dtype = comfy.model_management.unet_dtype() + controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data, unet_dtype) diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config) diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight" diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias" @@ -353,8 +353,8 @@ def load_controlnet(ckpt_path, model=None): return net if controlnet_config is None: - use_fp16 = comfy.model_management.should_use_fp16() - controlnet_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, use_fp16, True).unet_config + unet_dtype = comfy.model_management.unet_dtype() + controlnet_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, unet_dtype, True).unet_config controlnet_config.pop("out_channels") controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1] control_model = comfy.cldm.cldm.ControlNet(**controlnet_config) @@ -383,8 +383,7 @@ def load_controlnet(ckpt_path, model=None): missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False) print(missing, unexpected) - if use_fp16: - control_model = control_model.half() + control_model = control_model.to(unet_dtype) global_average_pooling = False filename = os.path.splitext(ckpt_path)[0] diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index b42637c82..bf58a4045 100644 --- a/comfy/ldm/modules/diffusionmodules/openaimodel.py +++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py @@ -296,8 +296,7 @@ class UNetModel(nn.Module): dims=2, num_classes=None, use_checkpoint=False, - use_fp16=False, - use_bf16=False, + dtype=th.float32, num_heads=-1, num_head_channels=-1, num_heads_upsample=-1, @@ -370,8 +369,7 @@ class UNetModel(nn.Module): self.conv_resample = conv_resample self.num_classes = num_classes self.use_checkpoint = use_checkpoint - self.dtype = th.float16 if use_fp16 else th.float32 - self.dtype = th.bfloat16 if use_bf16 else self.dtype + self.dtype = dtype self.num_heads = num_heads self.num_head_channels = num_head_channels self.num_heads_upsample = num_heads_upsample diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 787c78575..0ff2e7fb5 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -14,7 +14,7 @@ def count_blocks(state_dict_keys, prefix_string): count += 1 return count -def detect_unet_config(state_dict, key_prefix, use_fp16): +def detect_unet_config(state_dict, key_prefix, dtype): state_dict_keys = list(state_dict.keys()) unet_config = { @@ -32,7 +32,7 @@ def detect_unet_config(state_dict, key_prefix, use_fp16): else: unet_config["adm_in_channels"] = None - unet_config["use_fp16"] = use_fp16 + unet_config["dtype"] = dtype model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0] in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1] @@ -116,15 +116,15 @@ def model_config_from_unet_config(unet_config): print("no match", unet_config) return None -def model_config_from_unet(state_dict, unet_key_prefix, use_fp16, use_base_if_no_match=False): - unet_config = detect_unet_config(state_dict, unet_key_prefix, use_fp16) +def model_config_from_unet(state_dict, unet_key_prefix, dtype, use_base_if_no_match=False): + unet_config = detect_unet_config(state_dict, unet_key_prefix, dtype) model_config = model_config_from_unet_config(unet_config) if model_config is None and use_base_if_no_match: return comfy.supported_models_base.BASE(unet_config) else: return model_config -def unet_config_from_diffusers_unet(state_dict, use_fp16): +def unet_config_from_diffusers_unet(state_dict, dtype): match = {} attention_resolutions = [] @@ -147,47 +147,47 @@ def unet_config_from_diffusers_unet(state_dict, use_fp16): match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1] SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 2, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64} SDXL_refiner = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2560, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 384, + 'num_classes': 'sequential', 'adm_in_channels': 2560, 'dtype': dtype, 'in_channels': 4, 'model_channels': 384, 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 4, 4, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 4, 'use_linear_in_transformer': True, 'context_dim': 1280, "num_head_channels": 64} SD21 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'adm_in_channels': None, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, + 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024, "num_head_channels": 64} SD21_uncliph = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2048, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 2048, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024, "num_head_channels": 64} SD21_unclipl = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 1536, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 1536, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 1024} SD15 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'adm_in_channels': None, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, + 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [1, 2, 4], 'transformer_depth': [1, 1, 1, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': False, 'context_dim': 768, "num_heads": 8} SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [4], 'transformer_depth': [0, 0, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64} SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 4, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [], 'transformer_depth': [0, 0, 0], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 0, 'use_linear_in_transformer': True, "num_head_channels": 64, 'context_dim': 1} SDXL_diffusers_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, - 'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': use_fp16, 'in_channels': 9, 'model_channels': 320, + 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 9, 'model_channels': 320, 'num_res_blocks': 2, 'attention_resolutions': [2, 4], 'transformer_depth': [0, 2, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10, 'use_linear_in_transformer': True, 'context_dim': 2048, "num_head_channels": 64} @@ -203,8 +203,8 @@ def unet_config_from_diffusers_unet(state_dict, use_fp16): return unet_config return None -def model_config_from_diffusers_unet(state_dict, use_fp16): - unet_config = unet_config_from_diffusers_unet(state_dict, use_fp16) +def model_config_from_diffusers_unet(state_dict, dtype): + unet_config = unet_config_from_diffusers_unet(state_dict, dtype) if unet_config is not None: return model_config_from_unet_config(unet_config) return None diff --git a/comfy/model_management.py b/comfy/model_management.py index 3c390d9ca..1161c2447 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -448,6 +448,11 @@ def unet_inital_load_device(parameters, dtype): else: return cpu_dev +def unet_dtype(device=None, model_params=0): + if should_use_fp16(device=device, model_params=model_params): + return torch.float16 + return torch.float32 + def text_encoder_offload_device(): if args.gpu_only: return get_torch_device() diff --git a/comfy/sd.py b/comfy/sd.py index cfd6fb3cb..fd8b94df8 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -327,7 +327,9 @@ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_cl if "params" in model_config_params["unet_config"]: unet_config = model_config_params["unet_config"]["params"] if "use_fp16" in unet_config: - fp16 = unet_config["use_fp16"] + fp16 = unet_config.pop("use_fp16") + if fp16: + unet_config["dtype"] = torch.float16 noise_aug_config = None if "noise_aug_config" in model_config_params: @@ -405,12 +407,12 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o clip_target = None parameters = comfy.utils.calculate_parameters(sd, "model.diffusion_model.") - fp16 = model_management.should_use_fp16(model_params=parameters) + unet_dtype = model_management.unet_dtype(model_params=parameters) class WeightsLoader(torch.nn.Module): pass - model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", fp16) + model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", unet_dtype) if model_config is None: raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path)) @@ -418,12 +420,8 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o if output_clipvision: clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True) - dtype = torch.float32 - if fp16: - dtype = torch.float16 - if output_model: - inital_load_device = model_management.unet_inital_load_device(parameters, dtype) + inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype) offload_device = model_management.unet_offload_device() model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device) model.load_model_weights(sd, "model.diffusion_model.") @@ -458,15 +456,15 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o def load_unet(unet_path): #load unet in diffusers format sd = comfy.utils.load_torch_file(unet_path) parameters = comfy.utils.calculate_parameters(sd) - fp16 = model_management.should_use_fp16(model_params=parameters) + unet_dtype = model_management.unet_dtype(model_params=parameters) if "input_blocks.0.0.weight" in sd: #ldm - model_config = model_detection.model_config_from_unet(sd, "", fp16) + model_config = model_detection.model_config_from_unet(sd, "", unet_dtype) if model_config is None: raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path)) new_sd = sd else: #diffusers - model_config = model_detection.model_config_from_diffusers_unet(sd, fp16) + model_config = model_detection.model_config_from_diffusers_unet(sd, unet_dtype) if model_config is None: print("ERROR UNSUPPORTED UNET", unet_path) return None From fd4c5f07e7b4d034c7f9fa96d6853f4aeec7b2de Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 13 Oct 2023 14:51:10 -0400 Subject: [PATCH 11/15] Add a --bf16-unet to test running the unet in bf16. --- comfy/cli_args.py | 2 ++ comfy/model_management.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 35d44164f..d86557646 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -53,6 +53,8 @@ fp_group = parser.add_mutually_exclusive_group() fp_group.add_argument("--force-fp32", action="store_true", help="Force fp32 (If this makes your GPU work better please report it).") fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.") +parser.add_argument("--bf16-unet", action="store_true", help="Run the UNET in bf16. This should only be used for testing stuff.") + fpvae_group = parser.add_mutually_exclusive_group() fpvae_group.add_argument("--fp16-vae", action="store_true", help="Run the VAE in fp16, might cause black images.") fpvae_group.add_argument("--fp32-vae", action="store_true", help="Run the VAE in full precision fp32.") diff --git a/comfy/model_management.py b/comfy/model_management.py index 1161c2447..c24c7b27e 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -449,6 +449,8 @@ def unet_inital_load_device(parameters, dtype): return cpu_dev def unet_dtype(device=None, model_params=0): + if args.bf16_unet: + return torch.bfloat16 if should_use_fp16(device=device, model_params=model_params): return torch.float16 return torch.float32 From 25f0f4e9c8fdbad43311b8e5161b3e6efa9d58d9 Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Sat, 14 Oct 2023 11:54:33 -0300 Subject: [PATCH 12/15] Shortcut Alt + C to collapse/uncollapse selected nodes --- README.md | 1 + web/scripts/app.js | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/README.md b/README.md index 6bef25cee..d622c9072 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ Workflow examples can be found on the [Examples page](https://comfyanonymous.git | Ctrl + S | Save workflow | | Ctrl + O | Load workflow | | Ctrl + A | Select all nodes | +| Alt + C | Collapse/uncollapse selected nodes | | Ctrl + M | Mute/unmute selected nodes | | Ctrl + B | Bypass selected nodes (acts like the node was removed from the graph and the wires reconnected through) | | Delete/Backspace | Delete selected nodes | diff --git a/web/scripts/app.js b/web/scripts/app.js index 3cf3585d2..1a07d69bc 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -928,6 +928,16 @@ export class ComfyApp { block_default = true; } + // Alt + C collapse/uncollapse + if (e.key === 'c' && e.altKey) { + if (this.selected_nodes) { + for (var i in this.selected_nodes) { + this.selected_nodes[i].collapse() + } + } + block_default = true; + } + // Ctrl+C Copy if ((e.key === 'c') && (e.metaKey || e.ctrlKey)) { // Trigger onCopy From 2e6270e328d4b43f9d76172f21f6d7636bf4a452 Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Sat, 14 Oct 2023 11:56:44 -0300 Subject: [PATCH 13/15] Stop auto queue on error --- web/scripts/ui.js | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/web/scripts/ui.js b/web/scripts/ui.js index 1e7920167..c3b3fbda1 100644 --- a/web/scripts/ui.js +++ b/web/scripts/ui.js @@ -809,7 +809,8 @@ export class ComfyUI { if ( this.lastQueueSize != 0 && status.exec_info.queue_remaining == 0 && - document.getElementById("autoQueueCheckbox").checked + document.getElementById("autoQueueCheckbox").checked && + ! app.lastExecutionError ) { app.queuePrompt(0, this.batchCount); } From 8d04978298d8bb2bcb969608a994fe4ef39be3ee Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Sat, 14 Oct 2023 11:59:35 -0300 Subject: [PATCH 14/15] Allow all extensions if extension list is empty --- folder_paths.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/folder_paths.py b/folder_paths.py index 5d121b443..4a38deec0 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -146,7 +146,7 @@ def recursive_search(directory, excluded_dir_names=None): return result, dirs def filter_files_extensions(files, extensions): - return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files))) + return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files))) From a7b65b9505b0504fefc7b57a5e80a243bd6630cf Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Sat, 14 Oct 2023 12:11:49 -0300 Subject: [PATCH 15/15] Group menu option select nodes --- web/extensions/core/groupOptions.js | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/web/extensions/core/groupOptions.js b/web/extensions/core/groupOptions.js index 1d935e90a..d523737dc 100644 --- a/web/extensions/core/groupOptions.js +++ b/web/extensions/core/groupOptions.js @@ -38,6 +38,15 @@ app.registerExtension({ } } + options.push({ + content: "Select Nodes", + callback: () => { + this.selectNodes(nodesInGroup); + this.graph.change(); + this.canvas.focus(); + } + }); + // Modes // 0: Always // 1: On Event