CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

prep

Browse Source
This commit is contained in:
Haoming 2026-02-21 15:10:45 +08:00
parent 24279f0894
commit 2f59cd1a31
2 changed files with 19 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
backend/nn/wan.py
View File

@ -15,6 +15,10 @@ from backend.nn.flux import EmbedND, apply_rope1
from backend.utils import pad_to_patch_size
def attention(*args, **kwargs): # for Radial Attention
return attention_function(*args, **kwargs)
def sinusoidal_embedding_1d(dim, position):
# preprocess
assert dim % 2 == 0
@ -47,7 +51,7 @@ class WanSelfAttention(nn.Module):
self.norm_q = nn.RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
self.norm_k = nn.RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
def forward(self, x, freqs):
def forward(self, x, freqs, transformer_options={}):
r"""
Args:
x(Tensor): Shape [B, L, num_heads, C / num_heads]
@ -66,11 +70,12 @@ class WanSelfAttention(nn.Module):
q = qkv_fn_q(x)
k = qkv_fn_k(x)
x = attention_function(
x = attention(
q.view(b, s, n * d),
k.view(b, s, n * d),
self.v(x).view(b, s, n * d),
heads=self.num_heads,
transformer_options=transformer_options,
)
x = self.o(x)
@ -177,6 +182,7 @@ class WanAttentionBlock(nn.Module):
freqs,
context,
context_img_len=257,
transformer_options={},
):
r"""
Args:
@ -193,7 +199,7 @@ class WanAttentionBlock(nn.Module):
# assert e[0].dtype == torch.float32
# self-attention
y = self.self_attn(self.norm1(x) * (1 + repeat_e(e[1], x)) + repeat_e(e[0], x), freqs)
y = self.self_attn(self.norm1(x) * (1 + repeat_e(e[1], x)) + repeat_e(e[0], x), freqs, transformer_options=transformer_options)
x = x + y * repeat_e(e[2], x)
@ -357,18 +363,21 @@ class WanModel(nn.Module):
patches_replace = transformer_options.get("patches_replace", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
for i, block in enumerate(self.blocks):
transformer_options["block_index"] = i
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len)
out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
return out
out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
# head
x = self.head(x, e)

4
modules/sd_samplers_kdiffusion.py
View File

@ -175,6 +175,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
"s_min_uncond": self.s_min_uncond,
}
p.sd_model.forge_objects.unet.model_options["transformer_options"]["sampling_sigmas"] = sigmas
samples = self.launch_sampling(
t_enc + 1,
lambda: self.func(self.model_wrap_cfg, xi, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs),
@ -228,6 +230,8 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
"s_min_uncond": self.s_min_uncond,
}
p.sd_model.forge_objects.unet.model_options["transformer_options"]["sampling_sigmas"] = sigmas
samples = self.launch_sampling(
steps,
lambda: self.func(self.model_wrap_cfg, x, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs),