Spaces:
Running
on
Zero
Running
on
Zero
| from typing import Union, List | |
| import PIL | |
| import torch | |
| import torchvision.transforms as T | |
| from einops import repeat | |
| from kandinsky3.model.unet import UNet | |
| from kandinsky3.movq import MoVQ | |
| from kandinsky3.condition_encoders import T5TextConditionEncoder | |
| from kandinsky3.condition_processors import T5TextConditionProcessor | |
| from kandinsky3.model.diffusion import BaseDiffusion, get_named_beta_schedule | |
| class Kandinsky3T2IPipeline: | |
| def __init__( | |
| self, | |
| device_map: Union[str, torch.device, dict], | |
| dtype_map: Union[str, torch.dtype, dict], | |
| unet: UNet, | |
| null_embedding: torch.Tensor, | |
| t5_processor: T5TextConditionProcessor, | |
| t5_encoder: T5TextConditionEncoder, | |
| movq: MoVQ, | |
| gan: bool, | |
| ): | |
| self.device_map = device_map | |
| self.dtype_map = dtype_map | |
| self.to_pil = T.ToPILImage() | |
| self.unet = unet | |
| self.null_embedding = null_embedding | |
| self.t5_processor = t5_processor | |
| self.t5_encoder = t5_encoder | |
| self.movq = movq | |
| self.gan = gan | |
| def __call__( | |
| self, | |
| text: str, | |
| negative_text: str = None, | |
| images_num: int = 1, | |
| bs: int = 1, | |
| width: int = 1024, | |
| height: int = 1024, | |
| guidance_scale: float = 3.0, | |
| steps: int = 50, | |
| eta: float = 1.0 | |
| ) -> List[PIL.Image.Image]: | |
| betas = get_named_beta_schedule('cosine', 1000) | |
| base_diffusion = BaseDiffusion(betas, 0.99) | |
| times = list(range(999, 0, -1000 // steps)) | |
| if self.gan: | |
| times = list(range(979, 0, -250)) | |
| condition_model_input, negative_condition_model_input = self.t5_processor.encode(text, negative_text) | |
| for input_type in condition_model_input: | |
| condition_model_input[input_type] = condition_model_input[input_type][None].to( | |
| self.device_map['text_encoder'] | |
| ) | |
| if negative_condition_model_input is not None: | |
| for input_type in negative_condition_model_input: | |
| negative_condition_model_input[input_type] = negative_condition_model_input[input_type][None].to( | |
| self.device_map['text_encoder'] | |
| ) | |
| pil_images = [] | |
| with torch.no_grad(): | |
| with torch.cuda.amp.autocast(dtype=self.dtype_map['text_encoder']): | |
| context, context_mask = self.t5_encoder(condition_model_input) | |
| if negative_condition_model_input is not None: | |
| negative_context, negative_context_mask = self.t5_encoder(negative_condition_model_input) | |
| else: | |
| negative_context, negative_context_mask = None, None | |
| k, m = images_num // bs, images_num % bs | |
| for minibatch in [bs] * k + [m]: | |
| if minibatch == 0: | |
| continue | |
| bs_context = repeat(context, '1 n d -> b n d', b=minibatch) | |
| bs_context_mask = repeat(context_mask, '1 n -> b n', b=minibatch) | |
| if negative_context is not None: | |
| bs_negative_context = repeat(negative_context, '1 n d -> b n d', b=minibatch) | |
| bs_negative_context_mask = repeat(negative_context_mask, '1 n -> b n', b=minibatch) | |
| else: | |
| bs_negative_context, bs_negative_context_mask = None, None | |
| with torch.cuda.amp.autocast(dtype=self.dtype_map['unet']): | |
| images = base_diffusion.p_sample_loop( | |
| self.unet, (minibatch, 4, height // 8, width // 8), times, self.device_map['unet'], | |
| bs_context, bs_context_mask, self.null_embedding, guidance_scale, eta, | |
| negative_context=bs_negative_context, negative_context_mask=bs_negative_context_mask, | |
| gan=self.gan | |
| ) | |
| with torch.cuda.amp.autocast(dtype=self.dtype_map['movq']): | |
| images = torch.cat([self.movq.decode(image) for image in images.chunk(2)]) | |
| images = torch.clip((images + 1.) / 2., 0., 1.) | |
| for images_chunk in images.chunk(1): | |
| pil_images += [self.to_pil(image) for image in images_chunk] | |
| return pil_images | |