[Chore] allow auraflow latest to be torch compile compatible.

[sayakpaul]

What does this PR do?

Otherwise, we run into a shape broadcasting problem during torch.compile(). Note that this doesn't happen when the internal AuraDiffusion/auradiffusion-v0.1a0 variant.

Benchmarking script

Benchmarking script

importtorchtorch.set_float32_matmul_precision("high")
fromdiffusersimportAuraFlowPipelineimporttimeid="fal/AuraFlow"# "AuraDiffusion/auradiffusion-v0.1a0"pipeline=AuraFlowPipeline.from_pretrained(
id, torch_dtype=torch.float16
).to("cuda")
pipeline.set_progress_bar_config(disable=True)
torch._inductor.config.conv_1x1_as_mm=Truetorch._inductor.config.coordinate_descent_tuning=Truetorch._inductor.config.epilogue_fusion=Falsetorch._inductor.config.coordinate_descent_check_all_directions=Truepipeline.transformer.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.transformer=torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
pipeline.vae.decode=torch.compile(pipeline.vae.decode, mode="max-autotune", fullgraph=True)
prompt="a photo of a cat"for_inrange(3):
_=pipeline(
prompt=prompt,
num_inference_steps=50, guidance_scale=5.0,
generator=torch.manual_seed(1),
)
start=time.time()
for_inrange(10):
_=pipeline(
prompt=prompt,
num_inference_steps=50, guidance_scale=5.0,
generator=torch.manual_seed(1),
)
end=time.time()
avg_inference_time= (end-start) /10print(f"Average inference time: {avg_inference_time:.3f} seconds.")

Results

Results

| With `torch.compile()` | Without `torch.compile()` | |----------|----------| | 7.293 seconds | 8.801 seconds |

@apolinario FYI.

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

[yiyixuxu]

I don't think there is anything wrong with torch.compile
maybe you should update the default height and width for aura pipelines now the checkpoint only supports 1024 https://github.com/huggingface/diffusers/blob/bbd2f9d4e9ae70b04fedf65903fd1fb035437db4/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py#L394C9-L394C15

[bghira]

it's a shame that the height and width aren't just None, and then default to the sample size in the model config. but i guess it's actually hard to concretely calculate that. we can hardcode it based on max_pos_embed value. i'm not sure if we're going to make any 512px models available. they are very fast and nice to work with.

[sayakpaul]

@yiyixuxu done.

@bghira

height=heightorself.transformer.config.sample_size*self.vae_scale_factorwidth=widthorself.transformer.config.sample_size*self.vae_scale_factor

is our automatic resolution inference.

[apolinario]

@sayakpaul , by changing the default, you un-did the initial change, is that intended? f6c975a

[sayakpaul]

yeah that is right.