Skip to content

GitLab

"doc/git@developer.sourcefind.cn:wangsen/paddle_dbnet.git" did not exist on "1c2c26989d12841ceaf5ef37c5fc994b6ad731c8"
Commit 6971646b authored by comfyanonymous's avatar comfyanonymous
Browse files

Speed up model loading a bit. 

Default pytorch Linear initializes the weights which is useless and slow.
parent 84f13f82
Hide whitespace changes
Inline Side-by-side
Showing with 43 additions and 25 deletions
comfy/ldm/modules/attention.py
View file @ 6971646b
......@@ -10,6 +10,7 @@ from .diffusionmodules.util import checkpoint
from .sub_quadratic_attention import efficient_dot_product_attention
from comfy import model_management
import comfy.ops
from . import tomesd
......@@ -52,7 +53,7 @@ def init_(tensor):
class GEGLU(nn.Module):
def __init__(self, dim_in, dim_out):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out * 2)
self.proj = comfy.ops.Linear(dim_in, dim_out * 2)
def forward(self, x):
x, gate = self.proj(x).chunk(2, dim=-1)
......@@ -65,14 +66,14 @@ class FeedForward(nn.Module):
inner_dim = int(dim * mult)
dim_out = default(dim_out, dim)
project_in = nn.Sequential(
nn.Linear(dim, inner_dim),
comfy.ops.Linear(dim, inner_dim),
nn.GELU()
) if not glu else GEGLU(dim, inner_dim)
self.net = nn.Sequential(
project_in,
nn.Dropout(dropout),
nn.Linear(inner_dim, dim_out)
comfy.ops.Linear(inner_dim, dim_out)
)
def forward(self, x):
......@@ -154,12 +155,12 @@ class CrossAttentionBirchSan(nn.Module):
self.scale = dim_head ** -0.5
self.heads = heads
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_q = comfy.ops.Linear(query_dim, inner_dim, bias=False)
self.to_k = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_v = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, query_dim),
comfy.ops.Linear(inner_dim, query_dim),
nn.Dropout(dropout)
)
......@@ -251,12 +252,12 @@ class CrossAttentionDoggettx(nn.Module):
self.scale = dim_head ** -0.5
self.heads = heads
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_q = comfy.ops.Linear(query_dim, inner_dim, bias=False)
self.to_k = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_v = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, query_dim),
comfy.ops.Linear(inner_dim, query_dim),
nn.Dropout(dropout)
)
......@@ -349,12 +350,12 @@ class CrossAttention(nn.Module):
self.scale = dim_head ** -0.5
self.heads = heads
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_q = comfy.ops.Linear(query_dim, inner_dim, bias=False)
self.to_k = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_v = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, query_dim),
comfy.ops.Linear(inner_dim, query_dim),
nn.Dropout(dropout)
)
......@@ -407,11 +408,11 @@ class MemoryEfficientCrossAttention(nn.Module):
self.heads = heads
self.dim_head = dim_head
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_q = comfy.ops.Linear(query_dim, inner_dim, bias=False)
self.to_k = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_v = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
self.to_out = nn.Sequential(comfy.ops.Linear(inner_dim, query_dim), nn.Dropout(dropout))
self.attention_op: Optional[Any] = None
def forward(self, x, context=None, value=None, mask=None):
......@@ -456,11 +457,11 @@ class CrossAttentionPytorch(nn.Module):
self.heads = heads
self.dim_head = dim_head
self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
self.to_q = comfy.ops.Linear(query_dim, inner_dim, bias=False)
self.to_k = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_v = comfy.ops.Linear(context_dim, inner_dim, bias=False)
self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
self.to_out = nn.Sequential(comfy.ops.Linear(inner_dim, query_dim), nn.Dropout(dropout))
self.attention_op: Optional[Any] = None
def forward(self, x, context=None, value=None, mask=None):
......@@ -601,7 +602,7 @@ class SpatialTransformer(nn.Module):
stride=1,
padding=0)
else:
self.proj_in = nn.Linear(in_channels, inner_dim)
self.proj_in = comfy.ops.Linear(in_channels, inner_dim)
self.transformer_blocks = nn.ModuleList(
[BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
......@@ -615,7 +616,7 @@ class SpatialTransformer(nn.Module):
stride=1,
padding=0))
else:
self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
self.proj_out = zero_module(comfy.ops.Linear(in_channels, inner_dim))
self.use_linear = use_linear
def forward(self, x, context=None, transformer_options={}):
......
comfy/ops.py 0 → 100644
View file @ 6971646b
import torch
class Linear(torch.nn.Module):
def __init__(self, in_features: int, out_features: int, bias: bool = True,
device=None, dtype=None) -> None:
factory_kwargs = {'device': device, 'dtype': dtype}
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = torch.nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
if bias:
self.bias = torch.nn.Parameter(torch.empty(out_features, **factory_kwargs))
else:
self.register_parameter('bias', None)
def forward(self, input):
return torch.nn.functional.linear(input, self.weight, self.bias)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment