Created
June 5, 2024 05:56
-
-
Save sayakpaul/29d3a14905cfcbf611fe71ebd22e9b23 to your computer and use it in GitHub Desktop.
Benchmarking script for running Hunyuan DiT with `torch.compile()`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment