tsi-org/tango

1import inspect2from typing import List, Optional, Union3 4import torch5from torch import nn6from torch.nn import functional as F7from torchvision import transforms8from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer9 10from diffusers import (11    AutoencoderKL,12    DDIMScheduler,13    DiffusionPipeline,14    DPMSolverMultistepScheduler,15    LMSDiscreteScheduler,16    PNDMScheduler,17    UNet2DConditionModel,18)19from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput20 21 22class MakeCutouts(nn.Module):23    def __init__(self, cut_size, cut_power=1.0):24        super().__init__()25 26        self.cut_size = cut_size27        self.cut_power = cut_power28 29    def forward(self, pixel_values, num_cutouts):30        sideY, sideX = pixel_values.shape[2:4]31        max_size = min(sideX, sideY)32        min_size = min(sideX, sideY, self.cut_size)33        cutouts = []34        for _ in range(num_cutouts):35            size = int(torch.rand([]) ** self.cut_power * (max_size - min_size) + min_size)36            offsetx = torch.randint(0, sideX - size + 1, ())37            offsety = torch.randint(0, sideY - size + 1, ())38            cutout = pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]39            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))40        return torch.cat(cutouts)41 42 43def spherical_dist_loss(x, y):44    x = F.normalize(x, dim=-1)45    y = F.normalize(y, dim=-1)46    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)47 48 49def set_requires_grad(model, value):50    for param in model.parameters():51        param.requires_grad = value52 53 54class CLIPGuidedStableDiffusion(DiffusionPipeline):55    """CLIP guided stable diffusion based on the amazing repo by @crowsonkb and @Jack00056    - https://github.com/Jack000/glid-3-xl57    - https://github.dev/crowsonkb/k-diffusion58    """59 60    def __init__(61        self,62        vae: AutoencoderKL,63        text_encoder: CLIPTextModel,64        clip_model: CLIPModel,65        tokenizer: CLIPTokenizer,66        unet: UNet2DConditionModel,67        scheduler: Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler],68        feature_extractor: CLIPImageProcessor,69    ):70        super().__init__()71        self.register_modules(72            vae=vae,73            text_encoder=text_encoder,74            clip_model=clip_model,75            tokenizer=tokenizer,76            unet=unet,77            scheduler=scheduler,78            feature_extractor=feature_extractor,79        )80 81        self.normalize = transforms.Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std)82        self.cut_out_size = (83            feature_extractor.size84            if isinstance(feature_extractor.size, int)85            else feature_extractor.size["shortest_edge"]86        )87        self.make_cutouts = MakeCutouts(self.cut_out_size)88 89        set_requires_grad(self.text_encoder, False)90        set_requires_grad(self.clip_model, False)91 92    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):93        if slice_size == "auto":94            # half the attention head size is usually a good trade-off between95            # speed and memory96            slice_size = self.unet.config.attention_head_dim // 297        self.unet.set_attention_slice(slice_size)98 99    def disable_attention_slicing(self):100        self.enable_attention_slicing(None)101 102    def freeze_vae(self):103        set_requires_grad(self.vae, False)104 105    def unfreeze_vae(self):106        set_requires_grad(self.vae, True)107 108    def freeze_unet(self):109        set_requires_grad(self.unet, False)110 111    def unfreeze_unet(self):112        set_requires_grad(self.unet, True)113 114    @torch.enable_grad()115    def cond_fn(116        self,117        latents,118        timestep,119        index,120        text_embeddings,121        noise_pred_original,122        text_embeddings_clip,123        clip_guidance_scale,124        num_cutouts,125        use_cutouts=True,126    ):127        latents = latents.detach().requires_grad_()128 129        latent_model_input = self.scheduler.scale_model_input(latents, timestep)130 131        # predict the noise residual132        noise_pred = self.unet(latent_model_input, timestep, encoder_hidden_states=text_embeddings).sample133 134        if isinstance(self.scheduler, (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)):135            alpha_prod_t = self.scheduler.alphas_cumprod[timestep]136            beta_prod_t = 1 - alpha_prod_t137            # compute predicted original sample from predicted noise also called138            # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf139            pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)140 141            fac = torch.sqrt(beta_prod_t)142            sample = pred_original_sample * (fac) + latents * (1 - fac)143        elif isinstance(self.scheduler, LMSDiscreteScheduler):144            sigma = self.scheduler.sigmas[index]145            sample = latents - sigma * noise_pred146        else:147            raise ValueError(f"scheduler type {type(self.scheduler)} not supported")148 149        sample = 1 / self.vae.config.scaling_factor * sample150        image = self.vae.decode(sample).sample151        image = (image / 2 + 0.5).clamp(0, 1)152 153        if use_cutouts:154            image = self.make_cutouts(image, num_cutouts)155        else:156            image = transforms.Resize(self.cut_out_size)(image)157        image = self.normalize(image).to(latents.dtype)158 159        image_embeddings_clip = self.clip_model.get_image_features(image)160        image_embeddings_clip = image_embeddings_clip / image_embeddings_clip.norm(p=2, dim=-1, keepdim=True)161 162        if use_cutouts:163            dists = spherical_dist_loss(image_embeddings_clip, text_embeddings_clip)164            dists = dists.view([num_cutouts, sample.shape[0], -1])165            loss = dists.sum(2).mean(0).sum() * clip_guidance_scale166        else:167            loss = spherical_dist_loss(image_embeddings_clip, text_embeddings_clip).mean() * clip_guidance_scale168 169        grads = -torch.autograd.grad(loss, latents)[0]170 171        if isinstance(self.scheduler, LMSDiscreteScheduler):172            latents = latents.detach() + grads * (sigma**2)173            noise_pred = noise_pred_original174        else:175            noise_pred = noise_pred_original - torch.sqrt(beta_prod_t) * grads176        return noise_pred, latents177 178    @torch.no_grad()179    def __call__(180        self,181        prompt: Union[str, List[str]],182        height: Optional[int] = 512,183        width: Optional[int] = 512,184        num_inference_steps: Optional[int] = 50,185        guidance_scale: Optional[float] = 7.5,186        num_images_per_prompt: Optional[int] = 1,187        eta: float = 0.0,188        clip_guidance_scale: Optional[float] = 100,189        clip_prompt: Optional[Union[str, List[str]]] = None,190        num_cutouts: Optional[int] = 4,191        use_cutouts: Optional[bool] = True,192        generator: Optional[torch.Generator] = None,193        latents: Optional[torch.FloatTensor] = None,194        output_type: Optional[str] = "pil",195        return_dict: bool = True,196    ):197        if isinstance(prompt, str):198            batch_size = 1199        elif isinstance(prompt, list):200            batch_size = len(prompt)201        else:202            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")203 204        if height % 8 != 0 or width % 8 != 0:205            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")206 207        # get prompt text embeddings208        text_input = self.tokenizer(209            prompt,210            padding="max_length",211            max_length=self.tokenizer.model_max_length,212            truncation=True,213            return_tensors="pt",214        )215        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]216        # duplicate text embeddings for each generation per prompt217        text_embeddings = text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)218 219        if clip_guidance_scale > 0:220            if clip_prompt is not None:221                clip_text_input = self.tokenizer(222                    clip_prompt,223                    padding="max_length",224                    max_length=self.tokenizer.model_max_length,225                    truncation=True,226                    return_tensors="pt",227                ).input_ids.to(self.device)228            else:229                clip_text_input = text_input.input_ids.to(self.device)230            text_embeddings_clip = self.clip_model.get_text_features(clip_text_input)231            text_embeddings_clip = text_embeddings_clip / text_embeddings_clip.norm(p=2, dim=-1, keepdim=True)232            # duplicate text embeddings clip for each generation per prompt233            text_embeddings_clip = text_embeddings_clip.repeat_interleave(num_images_per_prompt, dim=0)234 235        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)236        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`237        # corresponds to doing no classifier free guidance.238        do_classifier_free_guidance = guidance_scale > 1.0239        # get unconditional embeddings for classifier free guidance240        if do_classifier_free_guidance:241            max_length = text_input.input_ids.shape[-1]242            uncond_input = self.tokenizer([""], padding="max_length", max_length=max_length, return_tensors="pt")243            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]244            # duplicate unconditional embeddings for each generation per prompt245            uncond_embeddings = uncond_embeddings.repeat_interleave(num_images_per_prompt, dim=0)246 247            # For classifier free guidance, we need to do two forward passes.248            # Here we concatenate the unconditional and text embeddings into a single batch249            # to avoid doing two forward passes250            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])251 252        # get the initial random noise unless the user supplied it253 254        # Unlike in other pipelines, latents need to be generated in the target device255        # for 1-to-1 results reproducibility with the CompVis implementation.256        # However this currently doesn't work in `mps`.257        latents_shape = (batch_size * num_images_per_prompt, self.unet.in_channels, height // 8, width // 8)258        latents_dtype = text_embeddings.dtype259        if latents is None:260            if self.device.type == "mps":261                # randn does not work reproducibly on mps262                latents = torch.randn(latents_shape, generator=generator, device="cpu", dtype=latents_dtype).to(263                    self.device264                )265            else:266                latents = torch.randn(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)267        else:268            if latents.shape != latents_shape:269                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")270            latents = latents.to(self.device)271 272        # set timesteps273        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())274        extra_set_kwargs = {}275        if accepts_offset:276            extra_set_kwargs["offset"] = 1277 278        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)279 280        # Some schedulers like PNDM have timesteps as arrays281        # It's more optimized to move all timesteps to correct device beforehand282        timesteps_tensor = self.scheduler.timesteps.to(self.device)283 284        # scale the initial noise by the standard deviation required by the scheduler285        latents = latents * self.scheduler.init_noise_sigma286 287        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature288        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.289        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502290        # and should be between [0, 1]291        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())292        extra_step_kwargs = {}293        if accepts_eta:294            extra_step_kwargs["eta"] = eta295 296        # check if the scheduler accepts generator297        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())298        if accepts_generator:299            extra_step_kwargs["generator"] = generator300 301        for i, t in enumerate(self.progress_bar(timesteps_tensor)):302            # expand the latents if we are doing classifier free guidance303            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents304            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)305 306            # predict the noise residual307            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample308 309            # perform classifier free guidance310            if do_classifier_free_guidance:311                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)312                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)313 314            # perform clip guidance315            if clip_guidance_scale > 0:316                text_embeddings_for_guidance = (317                    text_embeddings.chunk(2)[1] if do_classifier_free_guidance else text_embeddings318                )319                noise_pred, latents = self.cond_fn(320                    latents,321                    t,322                    i,323                    text_embeddings_for_guidance,324                    noise_pred,325                    text_embeddings_clip,326                    clip_guidance_scale,327                    num_cutouts,328                    use_cutouts,329                )330 331            # compute the previous noisy sample x_t -> x_t-1332            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample333 334        # scale and decode the image latents with vae335        latents = 1 / self.vae.config.scaling_factor * latents336        image = self.vae.decode(latents).sample337 338        image = (image / 2 + 0.5).clamp(0, 1)339        image = image.cpu().permute(0, 2, 3, 1).numpy()340 341        if output_type == "pil":342            image = self.numpy_to_pil(image)343 344        if not return_dict:345            return (image, None)346 347        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None)348