| from diffusers import DiffusionPipeline | |
| import torch | |
| import pdb | |
| if __name__ == "__main__": | |
| 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 | |
| torch._inductor.config.force_fuse_int_mm_with_mul = True | |
| torch._inductor.config.use_mixed_mm = True | |
| prompt = "3 happy foxes play together" | |
| pipe = DiffusionPipeline.from_pretrained("models/newdream-sdxl-20", torch_dtype=torch.float32).to("cuda:0") | |
| pipe.unet.to(memory_format=torch.channels_last) | |
| # pipe.vae.to(memory_format=torch.channels_last) | |
| pipe.fuse_qkv_projections() | |
| compiled_unet = torch.compile(pipe.unet, | |
| backend="inductor", | |
| fullgraph=True, | |
| mode="max-autotune", | |
| ) | |
| pipe.unet = compiled_unet | |
| pipe(prompt) | |
| pipe(prompt) | |
| # compiled_decode = torch.compile(pipe.vae.decode, backend="tensorrt", dynamic=False) | |
| pdb.set_trace() | |