flux torch cudagraphs compilation broken again on 2.6.0
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Expected Behavior
I should be able to torch.compile (TorchCompileNode) a Flux model with cudagraphs enabled
### Actual Behavior
cudagraphs phase of inductor "reduce-overhead" call is skipped, but can be reproduced with anything that creates cudagraphs compilations.
### Steps to Reproduce
add a torch compile node on a cuda machine and test a flux workflow
[flux-compile-0.json](https://github.com/user-attachments/files/18854941/flux-compile-0.json)
### Debug Logs
```powershell
skipping cudagraphs due to skipping cudagraphs due to cpu device (arg7_1). Found from :
File "comfy\ldm\flux\model.py", in forward
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
File "comfy\ldm\flux\model.py", in forward_orig
vec = self.time_in(timestep_embedding(timesteps, 256).to(img.dtype))
File "comfy\ldm\flux\layers.py", in forward
return self.out_layer(self.silu(self.in_layer(x)))
File "comfy\ops.py", in forward
return self.forward_comfy_cast_weights(*args, **kwargs)
File "comfy\ops.py", in forward_comfy_cast_weights
weight, bias = cast_bias_weight(self, input)
File "comfy\ops.py", in cast_bias_weight
bias = model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function)
File "comfy\model_management.py", in cast_to
r.copy_(weight, non_blocking=non_blocking)
```
### Other
```
torch==2.6.0+cu124
torchaudio==2.6.0+cu124
```
Contributor guide
Assessment
This issue has not been assessed yet.