diffusers/tools

1#!/usr/bin/env python32from diffusers import UNet2DConditionModel3import torch4 5torch.cuda.set_per_process_memory_fraction(0.5, device="cuda:1")6 7unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", variant="fp16", torch_dtype=torch.float16)8unet.train()9unet.enable_gradient_checkpointing()10unet = unet.to("cuda:1")11 12batch_size = 213 14sample = torch.randn((1, 4, 128, 128)).half().to(unet.device).repeat(batch_size, 1, 1, 1)15time_ids = (torch.arange(6) / 6)[None, :].half().to(unet.device).repeat(batch_size, 1)16encoder_hidden_states = torch.randn((1, 77, 2048)).half().to(unet.device).repeat(batch_size, 1, 1)17text_embeds = torch.randn((1, 1280)).half().to(unet.device).repeat(batch_size, 1)18 19out = unet(sample, 1.0, added_cond_kwargs={"time_ids": time_ids, "text_embeds": text_embeds}, encoder_hidden_states=encoder_hidden_states).sample20 21loss = ((out - sample) ** 2).mean()22loss.backward()23 24print(torch.cuda.max_memory_allocated(device=unet.device))25