| |
| 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)) |
|
|
|
|
| |
| |
|
|
