JaehaerysNL/StableDiffusion
1
0
Fork 0
mirror of https://github.com/Stability-AI/stablediffusion.git synced 2024-12-22 15:44:58 +00:00
Code Issues Projects Releases Packages Wiki Activity

create timestep embedding tensor on device

Browse source
This commit is contained in:
drhead 2024-05-17 09:54:59 -04:00 committed by GitHub
parent cf1d67a6fd
commit e5fb3e6a5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
ldm/modules/diffusionmodules/util.py
View file

@ -170,8 +170,8 @@ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
if not repeat_only: if not repeat_only:
half = dim // 2 half = dim // 2
freqs = torch.exp( freqs = torch.exp(
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
).to(device=timesteps.device) )
args = timesteps[:, None].float() * freqs[None] args = timesteps[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2: if dim % 2: