diffusers/community-pipelines-mirror

1import inspect2import os3 4import numpy as np5import torch6import torch.nn.functional as nnf7from PIL import Image8from torch.optim.adam import Adam9from tqdm import tqdm10 11from diffusers import StableDiffusionPipeline12from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput13 14 15def retrieve_timesteps(16    scheduler,17    num_inference_steps=None,18    device=None,19    timesteps=None,20    **kwargs,21):22    """23    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles24    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.25    Args:26        scheduler (`SchedulerMixin`):27            The scheduler to get timesteps from.28        num_inference_steps (`int`):29            The number of diffusion steps used when generating samples with a pre-trained model. If used,30            `timesteps` must be `None`.31        device (`str` or `torch.device`, *optional*):32            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.33        timesteps (`List[int]`, *optional*):34                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default35                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`36                must be `None`.37 38    Returns:39        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the40        second element is the number of inference steps.41    """42    if timesteps is not None:43        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())44        if not accepts_timesteps:45            raise ValueError(46                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"47                f" timestep schedules. Please check whether you are using the correct scheduler."48            )49        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)50        timesteps = scheduler.timesteps51        num_inference_steps = len(timesteps)52    else:53        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)54        timesteps = scheduler.timesteps55    return timesteps, num_inference_steps56 57 58class NullTextPipeline(StableDiffusionPipeline):59    def get_noise_pred(self, latents, t, context):60        latents_input = torch.cat([latents] * 2)61        guidance_scale = 7.562        noise_pred = self.unet(latents_input, t, encoder_hidden_states=context)["sample"]63        noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)64        noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)65        latents = self.prev_step(noise_pred, t, latents)66        return latents67 68    def get_noise_pred_single(self, latents, t, context):69        noise_pred = self.unet(latents, t, encoder_hidden_states=context)["sample"]70        return noise_pred71 72    @torch.no_grad()73    def image2latent(self, image_path):74        image = Image.open(image_path).convert("RGB")75        image = np.array(image)76        image = torch.from_numpy(image).float() / 127.5 - 177        image = image.permute(2, 0, 1).unsqueeze(0).to(self.device)78        latents = self.vae.encode(image)["latent_dist"].mean79        latents = latents * 0.1821580        return latents81 82    @torch.no_grad()83    def latent2image(self, latents):84        latents = 1 / 0.18215 * latents.detach()85        image = self.vae.decode(latents)["sample"].detach()86        image = self.processor.postprocess(image, output_type="pil")[0]87        return image88 89    def prev_step(self, model_output, timestep, sample):90        prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps91        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]92        alpha_prod_t_prev = (93            self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod94        )95        beta_prod_t = 1 - alpha_prod_t96        pred_original_sample = (sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.597        pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output98        prev_sample = alpha_prod_t_prev**0.5 * pred_original_sample + pred_sample_direction99        return prev_sample100 101    def next_step(self, model_output, timestep, sample):102        timestep, next_timestep = (103            min(timestep - self.scheduler.config.num_train_timesteps // self.num_inference_steps, 999),104            timestep,105        )106        alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod107        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]108        beta_prod_t = 1 - alpha_prod_t109        next_original_sample = (sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5110        next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output111        next_sample = alpha_prod_t_next**0.5 * next_original_sample + next_sample_direction112        return next_sample113 114    def null_optimization(self, latents, context, num_inner_steps, epsilon):115        uncond_embeddings, cond_embeddings = context.chunk(2)116        uncond_embeddings_list = []117        latent_cur = latents[-1]118        bar = tqdm(total=num_inner_steps * self.num_inference_steps)119        for i in range(self.num_inference_steps):120            uncond_embeddings = uncond_embeddings.clone().detach()121            uncond_embeddings.requires_grad = True122            optimizer = Adam([uncond_embeddings], lr=1e-2 * (1.0 - i / 100.0))123            latent_prev = latents[len(latents) - i - 2]124            t = self.scheduler.timesteps[i]125            with torch.no_grad():126                noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings)127            for j in range(num_inner_steps):128                noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings)129                noise_pred = noise_pred_uncond + 7.5 * (noise_pred_cond - noise_pred_uncond)130                latents_prev_rec = self.prev_step(noise_pred, t, latent_cur)131                loss = nnf.mse_loss(latents_prev_rec, latent_prev)132                optimizer.zero_grad()133                loss.backward()134                optimizer.step()135                loss_item = loss.item()136                bar.update()137                if loss_item < epsilon + i * 2e-5:138                    break139            for j in range(j + 1, num_inner_steps):140                bar.update()141            uncond_embeddings_list.append(uncond_embeddings[:1].detach())142            with torch.no_grad():143                context = torch.cat([uncond_embeddings, cond_embeddings])144                latent_cur = self.get_noise_pred(latent_cur, t, context)145        bar.close()146        return uncond_embeddings_list147 148    @torch.no_grad()149    def ddim_inversion_loop(self, latent, context):150        self.scheduler.set_timesteps(self.num_inference_steps)151        _, cond_embeddings = context.chunk(2)152        all_latent = [latent]153        latent = latent.clone().detach()154        with torch.no_grad():155            for i in range(0, self.num_inference_steps):156                t = self.scheduler.timesteps[len(self.scheduler.timesteps) - i - 1]157                noise_pred = self.unet(latent, t, encoder_hidden_states=cond_embeddings)["sample"]158                latent = self.next_step(noise_pred, t, latent)159                all_latent.append(latent)160        return all_latent161 162    def get_context(self, prompt):163        uncond_input = self.tokenizer(164            [""], padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt"165        )166        uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]167        text_input = self.tokenizer(168            [prompt],169            padding="max_length",170            max_length=self.tokenizer.model_max_length,171            truncation=True,172            return_tensors="pt",173        )174        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]175        context = torch.cat([uncond_embeddings, text_embeddings])176        return context177 178    def invert(179        self, image_path: str, prompt: str, num_inner_steps=10, early_stop_epsilon=1e-6, num_inference_steps=50180    ):181        self.num_inference_steps = num_inference_steps182        context = self.get_context(prompt)183        latent = self.image2latent(image_path)184        ddim_latents = self.ddim_inversion_loop(latent, context)185        if os.path.exists(image_path + ".pt"):186            uncond_embeddings = torch.load(image_path + ".pt")187        else:188            uncond_embeddings = self.null_optimization(ddim_latents, context, num_inner_steps, early_stop_epsilon)189            uncond_embeddings = torch.stack(uncond_embeddings, 0)190            torch.save(uncond_embeddings, image_path + ".pt")191        return ddim_latents[-1], uncond_embeddings192 193    @torch.no_grad()194    def __call__(195        self,196        prompt,197        uncond_embeddings,198        inverted_latent,199        num_inference_steps: int = 50,200        timesteps=None,201        guidance_scale=7.5,202        negative_prompt=None,203        num_images_per_prompt=1,204        generator=None,205        latents=None,206        prompt_embeds=None,207        negative_prompt_embeds=None,208        output_type="pil",209    ):210        self._guidance_scale = guidance_scale211        # 0. Default height and width to unet212        height = self.unet.config.sample_size * self.vae_scale_factor213        width = self.unet.config.sample_size * self.vae_scale_factor214        # to deal with lora scaling and other possible forward hook215        callback_steps = None216        # 1. Check inputs. Raise error if not correct217        self.check_inputs(218            prompt,219            height,220            width,221            callback_steps,222            negative_prompt,223            prompt_embeds,224            negative_prompt_embeds,225        )226        # 2. Define call parameter227        device = self._execution_device228        # 3. Encode input prompt229        prompt_embeds, _ = self.encode_prompt(230            prompt,231            device,232            num_images_per_prompt,233            self.do_classifier_free_guidance,234            negative_prompt,235            prompt_embeds=prompt_embeds,236            negative_prompt_embeds=negative_prompt_embeds,237        )238        # 4. Prepare timesteps239        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)240        latents = inverted_latent241        with self.progress_bar(total=num_inference_steps) as progress_bar:242            for i, t in enumerate(timesteps):243                noise_pred_uncond = self.unet(latents, t, encoder_hidden_states=uncond_embeddings[i])["sample"]244                noise_pred = self.unet(latents, t, encoder_hidden_states=prompt_embeds)["sample"]245                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)246                # compute the previous noisy sample x_t -> x_t-1247                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]248                progress_bar.update()249        if not output_type == "latent":250            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[251                0252            ]253        else:254            image = latents255        image = self.image_processor.postprocess(256            image, output_type=output_type, do_denormalize=[True] * image.shape[0]257        )258        # Offload all models259        self.maybe_free_model_hooks()260        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=False)261