* sync

* sync
This commit is contained in:
lllyasviel
2023-10-12 03:08:12 -07:00
committed by GitHub
parent 42bb1cda2f
commit 4c867c1b8b
9 changed files with 208 additions and 405 deletions
+2 -31
View File
@@ -3,47 +3,18 @@ import comfy.clip_vision
import safetensors.torch as sf
import comfy.model_management as model_management
import contextlib
import comfy.ldm.modules.attention as attention
from fooocus_extras.resampler import Resampler
from comfy.model_patcher import ModelPatcher
if model_management.xformers_enabled():
import xformers
import xformers.ops
SD_V12_CHANNELS = [320] * 4 + [640] * 4 + [1280] * 4 + [1280] * 6 + [640] * 6 + [320] * 6 + [1280] * 2
SD_XL_CHANNELS = [640] * 8 + [1280] * 40 + [1280] * 60 + [640] * 12 + [1280] * 20
def sdp(q, k, v, extra_options):
if model_management.xformers_enabled():
b, _, _ = q.shape
q, k, v = map(
lambda t: t.unsqueeze(3)
.reshape(b, t.shape[1], extra_options["n_heads"], extra_options["dim_head"])
.permute(0, 2, 1, 3)
.reshape(b * extra_options["n_heads"], t.shape[1], extra_options["dim_head"])
.contiguous(),
(q, k, v),
)
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=None)
out = (
out.unsqueeze(0)
.reshape(b, extra_options["n_heads"], out.shape[1], extra_options["dim_head"])
.permute(0, 2, 1, 3)
.reshape(b, out.shape[1], extra_options["n_heads"] * extra_options["dim_head"])
)
else:
b, _, _ = q.shape
q, k, v = map(
lambda t: t.view(b, -1, extra_options["n_heads"], extra_options["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)
out = out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"])
return out
return attention.optimized_attention(q, k, v, heads=extra_options["n_heads"], mask=None)
class ImageProjModel(torch.nn.Module):