DMD: Inference After Training Completes

[kgmills]

Hello,

Thanks for providing the step distillation code PixArt-alpha using DMD. When I run the code, it saves FSDP checkpoints periodically, and then one more FSDP checkpoint when training is done, in ".bin" format.

Say I want to take any of these checkpoint files and use them for inference using the existing diffusers pipeline ("PixArt-alpha/PixArt-Alpha-DMD-XL-2-512x512", just replacing the provided DiT weights with the ones I trained), how do I accomplish this correctly? I've tried loading the state_dict from the .bin file and even using 'tools/convert_pixart_to_diffusers.py', but to no avail.

Any help/advice with this matter will be greatly appeciated.

Thanks.