|
|
|
from diffusers import UNet2DConditionModel |
|
import torch |
|
|
|
unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", variant="fp16", torch_dtype=torch.float16) |
|
unet.train() |
|
unet.enable_gradient_checkpointing() |
|
unet = unet.to("cuda:1") |
|
|
|
batch_size = 8 |
|
|
|
sample = torch.randn((1, 4, 128, 128)).half().to(unet.device).repeat(batch_size, 1, 1, 1) |
|
time_ids = (torch.arange(6) / 6)[None, :].half().to(unet.device).repeat(batch_size, 1) |
|
encoder_hidden_states = torch.randn((1, 77, 2048)).half().to(unet.device).repeat(batch_size, 1, 1) |
|
text_embeds = torch.randn((1, 1280)).half().to(unet.device).repeat(batch_size, 1) |
|
|
|
out = unet(sample, 1.0, added_cond_kwargs={"time_ids": time_ids, "text_embeds": text_embeds}, encoder_hidden_states=encoder_hidden_states).sample |
|
|
|
loss = ((out - sample) ** 2).mean() |
|
loss.backward() |
|
|
|
print(torch.cuda.max_memory_allocated(device=unet.device)) |
|
|
|
|
|
|
|
|
|
|
