Support for PyTorch 2.1 torch.export To Load / Save Compiled Stable Diffusion Models

[djj0s3]

Diffusers supports PyTorch 2.0's torch.compile which yields significant performance gains. However, the torch2.0 implementation did not include support for saving and loading the compiled result, which unfortunately limits its use in some environments. An example would be a production environment where new GPUs need to spin up to meet demand but require a "warm" up with compilation happening each time. With the new PyTorch 2.1 support, we should be able to save the compilation once and load it dynamically. PyTorch2.1 is now in RC so seems like a good time for this to be officially supported by Diffusers.

[djj0s3]

Upon reading the torch 2.1 docs, I'd expect to do something like the pseudocode below to utilize torch.export.* with Diffusers:

pipe = StableDiffusionPipeline.from_pretrained(
            <MODEL>,
            torch_dtype=torch.float16,
        )
example_inputs = {
            'prompt_embeds': promptEmbeds,
            'negative_prompt_embeds': negPromptEmbeds,
            'num_images_per_prompt': numImages,
            'guidance_scale': guidanceScale,
            'generator': generator,
            'num_inference_steps': numSteps,
        }
fileName = "exported_pipe";
if not os.path.exists(fileName):
	exported_pipe = torch.export.export(pipe, args=(), kwargs=example_inputs)
        torch.export.save(exported_program, fileName)
else:
        pipe = torch.export.load(fileName)

[patrickvonplaten]

@djj0s3 we usually only compile the unet, so I guess we would only have to export/load the unet no?

[djj0s3]

That makes sense to me. Export seems to trace the whole model though sorta like tracing if I follow the docs correctly. What inputs would we send through unet?

[patrickvonplaten]

The unet takes all the inputs as defined here:

diffusers/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

Line 831 in 04d696d

noise_pred = self.unet(

[djj0s3]

Very helpful, thank you. I got reasonably far with that approach but got a bit stuck (it would fail on serialization, and fail on use downstream. So, it is not quite as straightforward to use as torch.compile in my experience so far. :( Are any maintainers going to take a spin at documenting/supporting torch.export.export for torch2.1? I'd love to follow that issue/PR if so.

[yuananf]

PyTorch 2.1 was released. https://pytorch.org/blog/pytorch-2-1/

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[AbhinavGopal]

Wondering about this right now. Was anyone able to figure out how to save + load compiled unets?