Oysiyl Claude Sonnet 4.5 commited on
Commit
ecce6e7
·
1 Parent(s): 576412f

Fix timestep_embedding for torch.compile compatibility

Browse files

Revert to CPU-first tensor creation then .to(device) instead of
device= parameter. The device= parameter causes torch.compile to fail
with ConstantVariable assertion error for torch.device type.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

comfy/ldm/modules/diffusionmodules/util.py CHANGED
@@ -268,8 +268,10 @@ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
268
  if not repeat_only:
269
  half = dim // 2
270
  freqs = torch.exp(
271
- -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
272
- )
 
 
273
  args = timesteps[:, None].float() * freqs[None]
274
  embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
275
  if dim % 2:
 
268
  if not repeat_only:
269
  half = dim // 2
270
  freqs = torch.exp(
271
+ -math.log(max_period)
272
+ * torch.arange(start=0, end=half, dtype=torch.float32)
273
+ / half
274
+ ).to(timesteps.device)
275
  args = timesteps[:, None].float() * freqs[None]
276
  embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
277
  if dim % 2: