lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-04-30 03:01:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

better cast

Browse Source
This commit is contained in:
layerdiffusion
2024-08-07 22:19:17 -07:00
parent 78c65708ea
commit c00b45aa71
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/nn/flux.py
View File

@@ -105,7 +105,7 @@ class RMSNorm(nn.Module):
to_args = dict(device=x.device, dtype=x.dtype) to_args = dict(device=x.device, dtype=x.dtype)
x = x.float() x = x.float()
rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6) rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
return (x * rrms.to(x) * self.scale.to(x)).to(**to_args) return (x * rrms).to(**to_args) * self.scale.to(**to_args)
class QKNorm(nn.Module): class QKNorm(nn.Module):