Benchmarking script for running Hunyuan DiT with `torch.compile()`.

run_hunyuan_dit_compile.py

import torch

torch.set_float32_matmul_precision("high")

from diffusers import HunyuanDiTPipeline
import argparse
import time


def load_pipeline(args):
    id = "XCLiu/HunyuanDiT-0523"
    pipeline = HunyuanDiTPipeline.from_pretrained(
        id,
        torch_dtype=torch.float16
    ).to("cuda")
    pipeline.set_progress_bar_config(disable=True)

    if args.use_compile:
        torch._inductor.config.conv_1x1_as_mm = True
        torch._inductor.config.coordinate_descent_tuning = True
        torch._inductor.config.epilogue_fusion = False
        torch._inductor.config.coordinate_descent_check_all_directions = True

        pipeline.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)
    return pipeline


def run_benchmark(args):
    pipeline = load_pipeline(args)

    prompt = "一个宇航员在骑马"
    for _ in range(3):
        _ = pipeline(
            prompt=prompt,
            generator=torch.manual_seed(1),
        )

    start = time.time()
    for _ in range(10):
        _ = pipeline(
            prompt=prompt,
            generator=torch.manual_seed(1),
        )
    end = time.time()
    avg_inference_time = (end - start) / 10
    print(f"Average inference time: {avg_inference_time:.3f} seconds.")

    image = pipeline(
        prompt=prompt,
        generator=torch.manual_seed(1),
    ).images[0]
    filename = "_".join(prompt.split(" "))
    image.save(f"diffusers_{filename}_compile@{args.use_compile}.png")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_compile", type=bool, default=True)
    args = parser.parse_args()
    run_benchmark(args)