ControlNets for Multi-GPU

[ishan-modi]

Currently most of our documentation has the following way of using ControlNets

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

controlnet = SanaControlNetModel.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNet_diffusers",
    torch_dtype=torch.float16
)

pipe = SanaControlNetPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    variant="fp16",
    controlnet=controlnet,
    torch_dtype=torch.float16,
)

pipe.to('cuda')
pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

This does work for a single GPU use case, but it is not appropriate for multi gpu, because controlnet and pipe both will have a separate device map, which causes issue related to device alignment during forward pass.

Following is how it should be documented

Single GPU

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

pipe = SanaControlNetPipeline.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers",
    variant="fp16",
    torch_dtype=torch.float16
).to('cuda')

pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

Multi GPU

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

pipe = SanaControlNetPipeline.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers",
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="balanced"
)

pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

Let me know if my thinking is correct on this one and if we need to change Docstring
cc - @a-r-r-o-w