multimodalart/pix2pix-zero

1import sys2import numpy as np3import torch4import torch.nn.functional as F5from random import randrange6from typing import Any, Callable, Dict, List, Optional, Union, Tuple7from diffusers import DDIMScheduler8from diffusers.schedulers.scheduling_ddim import DDIMSchedulerOutput9from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput10sys.path.insert(0, "src/utils")11from base_pipeline import BasePipeline12from cross_attention import prep_unet13 14 15class DDIMInversion(BasePipeline):16 17    def auto_corr_loss(self, x, random_shift=True):18        B,C,H,W = x.shape19        assert B==120        x = x.squeeze(0)21        # x must be shape [C,H,W] now22        reg_loss = 0.023        for ch_idx in range(x.shape[0]):24            noise = x[ch_idx][None, None,:,:]25            while True:26                if random_shift: roll_amount = randrange(noise.shape[2]//2)27                else: roll_amount = 128                reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=2)).mean()**229                reg_loss += (noise*torch.roll(noise, shifts=roll_amount, dims=3)).mean()**230                if noise.shape[2] <= 8:31                    break32                noise = F.avg_pool2d(noise, kernel_size=2)33        return reg_loss34    35    def kl_divergence(self, x):36        _mu = x.mean()37        _var = x.var()38        return _var + _mu**2 - 1 - torch.log(_var+1e-7)39 40 41    def __call__(42        self,43        prompt: Union[str, List[str]] = None,44        num_inversion_steps: int = 50,45        guidance_scale: float = 7.5,46        negative_prompt: Optional[Union[str, List[str]]] = None,47        num_images_per_prompt: Optional[int] = 1,48        eta: float = 0.0,49        output_type: Optional[str] = "pil",50        return_dict: bool = True,51        cross_attention_kwargs: Optional[Dict[str, Any]] = None,52        img=None, # the input image as a PIL image53        torch_dtype=torch.float32,54 55        # inversion regularization parameters56        lambda_ac: float = 20.0,57        lambda_kl: float = 20.0,58        num_reg_steps: int = 5,59        num_ac_rolls: int = 5,60    ):61        62        # 0. modify the unet to be useful :D63        self.unet = prep_unet(self.unet)64 65        # set the scheduler to be the Inverse DDIM scheduler66        # self.scheduler = MyDDIMScheduler.from_config(self.scheduler.config)67 68        device = self._execution_device69        do_classifier_free_guidance = guidance_scale > 1.070        self.scheduler.set_timesteps(num_inversion_steps, device=device)71        timesteps = self.scheduler.timesteps72 73        # Encode the input image with the first stage model74        x0 = np.array(img)/25575        x0 = torch.from_numpy(x0).type(torch_dtype).permute(2, 0, 1).unsqueeze(dim=0).repeat(1, 1, 1, 1).cuda()76        x0 = (x0 - 0.5) * 2.77        with torch.no_grad():78            x0_enc = self.vae.encode(x0).latent_dist.sample().to(device, torch_dtype)79        latents = x0_enc = 0.18215 * x0_enc80 81        # Decode and return the image82        with torch.no_grad():83            x0_dec = self.decode_latents(x0_enc.detach())84        image_x0_dec = self.numpy_to_pil(x0_dec)85 86        with torch.no_grad():87            prompt_embeds = self._encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt).to(device)88        extra_step_kwargs = self.prepare_extra_step_kwargs(None, eta)89 90        # Do the inversion91        num_warmup_steps = len(timesteps) - num_inversion_steps * self.scheduler.order # should be 0?92        with self.progress_bar(total=num_inversion_steps) as progress_bar:93            for i, t in enumerate(timesteps.flip(0)[1:-1]):94                # expand the latents if we are doing classifier free guidance95                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents96                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)97 98                # predict the noise residual99                with torch.no_grad():100                    noise_pred = self.unet(latent_model_input,t,encoder_hidden_states=prompt_embeds,cross_attention_kwargs=cross_attention_kwargs,).sample101 102                # perform guidance103                if do_classifier_free_guidance:104                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)105                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)106 107                # regularization of the noise prediction108                e_t = noise_pred109                for _outer in range(num_reg_steps):110                    if lambda_ac>0:111                        for _inner in range(num_ac_rolls):112                            _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)113                            l_ac = self.auto_corr_loss(_var)114                            l_ac.backward()115                            _grad = _var.grad.detach()/num_ac_rolls116                            e_t = e_t - lambda_ac*_grad117                    if lambda_kl>0:118                        _var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)119                        l_kld = self.kl_divergence(_var)120                        l_kld.backward()121                        _grad = _var.grad.detach()122                        e_t = e_t - lambda_kl*_grad123                    e_t = e_t.detach()124                noise_pred = e_t125 126                # compute the previous noisy sample x_t -> x_t-1127                latents = self.scheduler.step(noise_pred, t, latents, reverse=True, **extra_step_kwargs).prev_sample128 129                # call the callback, if provided130                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):131                    progress_bar.update()132        133 134        x_inv = latents.detach().clone()135        # reconstruct the image136 137        # 8. Post-processing138        image = self.decode_latents(latents.detach())139        image = self.numpy_to_pil(image)140        return x_inv, image, image_x0_dec