Spaces:
Sleeping
Sleeping
Fix timestep_embedding for torch.compile compatibility
Browse filesRevert 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)
|
| 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:
|