File size: 2,582 Bytes
e53e4fc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import torch
from diffusers import AutoencoderKL, DiffusionPipeline
from transformers import CLIPTextModel, CLIPTokenizer
from mv_unet import SPADUnetModel
from diffusers.schedulers import DPMSolverMultistepScheduler
class SPADPipeline(DiffusionPipeline):
def __init__(
self,
vae: AutoencoderKL,
unet: SPADUnetModel,
tokenizer: CLIPTokenizer,
text_encoder: CLIPTextModel,
scheduler: DPMSolverMultistepScheduler,
):
super().__init__()
self.vae = vae
self.unet = unet
self.tokenizer = tokenizer
self.text_encoder = text_encoder
self.scheduler = scheduler
# make sure all our models are on the same device
self.vae.to(self.device)
self.unet.to(self.device)
self.text_encoder.to(self.device)
def encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt=None):
text_input = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
return_tensors="pt"
)
text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
# we duplicate the text embeddings for each generation, just to save time :)
bs_embed, seq_len, _ = text_embeddings.shape
text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
return text_embeddings
def __call__(self, prompt, num_inference_steps=50, guidance_scale=7.5):
# encide the prompt into the text embeddings
text_embeddings = self.encode_prompt(prompt, self.device, 1, do_classifier_free_guidance=False)
# this is the initial noise sample
latents = torch.randn(
(text_embeddings.shape[0], self.unet.in_channels, self.unet.image_size, self.unet.image_size),
device=self.device
)
# setting up the scheduler
self.scheduler.set_timesteps(num_inference_steps)
# iterate and generate
for t in self.scheduler.timesteps:
latents = self.scheduler.scale_model_input(latents, t)
latents = self.unet(latents, t, text_embeddings)["sample"]
latents = self.scheduler.step(latents, t, latents, guidance_scale=guidance_scale)["prev_sample"]
# decode latents into images
images = self.vae.decode(latents)
images = (images / 2 + 0.5).clamp(0, 1)
return images |