Not an expert in machine learning, but AFAIK diffusion models use a completely different architecture, therefore you can't use the same code to run optimized versions of both. But maybe the core ideas can be adapted to diffusion somehow.

Thanks! I don't have much experience with diffusion models, but technically any multi-layer model could benefit from loading weights one by one