enable compilation in qwen image.

[sayakpaul]

What does this PR do?

• Adds tests for the Qwen transformer model tests.
• Enable compilation with triggering recompilations.

Timing (compilation) gathered from an H100:

• PR branch: timings.mean()=tensor(42.8954)
• tests/qwen-image: timings.mean()=tensor(75.9385)

Code

from diffusers import DiffusionPipeline
import torch
import time

pipe = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16).to("cuda")
# pipe.load_lora_weights("trained-qwen-image-lora")
pipe.transformer.compile(fullgraph=True)

timings = []
for _ in range(3):
    start = time.time()
    image = pipe(
        "a 3dicon, a llama with a signboard saying 'Qwen is awesome'", 
        guidance_scale=1.0, 
        num_inference_steps=50,

    ).images[0]
    end = time.time()
    timings.append(end - start)

timings = torch.tensor(timings)
print(f"{timings.mean()=}")
image.save("llama_pretrained.png")

[sayakpaul]

Recompilation trigger one.

[sayakpaul]

Recompilation trigger two.

[sayakpaul]

Will investigate in a follow-up.

[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.

[a-r-r-o-w]

@sayakpaul Could you rebase with main? Sorry I didn't see this before stacked over tests/qwen-image

[sayakpaul]

Done!

[a-r-r-o-w]

TODO: we need to remove frame (can be done in future PR)

[a-r-r-o-w]

Let's remove the self.rope_cache and just use the lru_cache implementation? WDYT @yiyixuxu?

WDYT about maybe putting maxsize=128 or something here so that long running services that use diffusers don't accidentally die with OOM (probably very unlikely though) @sayakpaul?

[sayakpaul]

maxsize=128 sounds reasonable to me.

[a-r-r-o-w]

This is most likely not equivalent. When registered as buffer, if the model is loaded in bf16, the precision of these will bf16 instead of fp32. Doing RoPE in bf16 may harm image quality, so we need to be careful here. Not sure what's best to do here -- maybe for now we can put the rope layer in _keep_modules_in_fp32?

This recompilation related problem seems to have become too common with RoPE. Maybe we need to rethink the design a bit.

[sayakpaul]

Just for the record, sharing the recompilation error we get without the buffer implementation:

>               raise exc.RecompileError(message)

E               torch._dynamo.exc.RecompileError: Recompiling function forward in /fsx/sayak/diffusers/src/diffusers/models/transformers/transformer_qwenimage.py:529

E                   triggered by the following guard failure(s):

E                   - 0/0: tensor 'self._modules['pos_embed'].neg_freqs' dispatch key set mismatch. expected DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), actual DispatchKeySet(CUDA, BackendSelect, ADInplaceOrView, AutogradCUDA)


../miniconda3/envs/diffusers/lib/python3.10/site-packages/torch/_dynamo/guards.py:3822: RecompileError

But I agree with your first point.