[torch.compile] Make HiDream torch.compile ready

[sayakpaul]

What does this PR do?

Trying to make the HiDream model fully compatible with torch.compile() but it fails with:
https://pastebin.com/EbCFqBvw

To reproduce run the following from a GPU machine:

RUN_COMPILE=1 RUN_SLOW=1 pytest tests/models/transformers/test_models_transformer_hidream.py -k "test_torch_compile_recompilation_and_graph_break"

I am on the following env:

- 🤗 Diffusers version: 0.34.0.dev0
- Platform: Linux-6.8.0-55-generic-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.10.12
- PyTorch version (GPU?): 2.7.0+cu126 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.30.2
- Transformers version: 4.51.3
- Accelerate version: 1.6.0.dev0
- PEFT version: 0.15.2.dev0
- Bitsandbytes version: 0.45.3
- Safetensors version: 0.5.3
- xFormers version: not installed
- Accelerator: NVIDIA GeForce RTX 4090, 24564 MiB
NVIDIA GeForce RTX 4090, 24564 MiB
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

@anijain2305 @StrongerXi would you have any pointers?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Just reimplemented it to eliminate the numpy() dependency.

[sayakpaul]

Relevant test for this PR.

[StrongerXi]

The graph break seems to be induced by @torch.no_grad:

diffusers/src/diffusers/models/transformers/transformer_hidream_image.py

Line 388 in d0c0239

@torch.no_grad()

@anijain2305 is this known?

[sayakpaul]

The graph break seems to be induced by @torch.no_grad:

diffusers/src/diffusers/models/transformers/transformer_hidream_image.py

Line 388 in d0c0239

@torch.no_grad()

@anijain2305 is this known?

Even if we remove the decorator, it still fails with the same error.

[anijain2305]

LGTM

Edit - Checked the messages, missed that there is still a graph break. I can take a look today.

[sayakpaul]

Thanks! Appreciate it.

…

On Thu, 8 May 2025 at 7:06 PM, Animesh Jain ***@***.***> wrote: ***@***.**** approved this pull request. LGTM — Reply to this email directly, view it on GitHub <#11477 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFPE2TCL5CXY6KDTOMUTXM325NMWRAVCNFSM6AAAAAB4H4U5SOVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDQMRVGE3DSNZQHE> . You are receiving this because you authored the thread.Message ID: ***@***.***>