diffusers/community-pipelines-mirror

1import math2import tempfile3from typing import List, Optional4 5import numpy as np6import PIL.Image7import torch8from accelerate import Accelerator9from torchvision import transforms10from tqdm.auto import tqdm11from transformers import CLIPTextModel, CLIPTokenizer12 13from diffusers import AutoencoderKL, DiffusionPipeline, DPMSolverMultistepScheduler, UNet2DConditionModel14from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin15from diffusers.models.attention_processor import (16    AttnAddedKVProcessor,17    AttnAddedKVProcessor2_0,18    LoRAAttnAddedKVProcessor,19    LoRAAttnProcessor,20    LoRAAttnProcessor2_0,21    SlicedAttnAddedKVProcessor,22)23from diffusers.optimization import get_scheduler24 25 26class SdeDragPipeline(DiffusionPipeline):27    r"""28    Pipeline for image drag-and-drop editing using stochastic differential equations: https://arxiv.org/abs/2311.01410.29    Please refer to the [official repository](https://github.com/ML-GSAI/SDE-Drag) for more information.30 31    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the32    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)33 34    Args:35        vae ([`AutoencoderKL`]):36            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.37        text_encoder ([`CLIPTextModel`]):38            Frozen text-encoder. Stable Diffusion uses the text portion of39            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically40            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.41        tokenizer (`CLIPTokenizer`):42            Tokenizer of class43            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).44        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.45        scheduler ([`SchedulerMixin`]):46            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Please use47            [`DDIMScheduler`].48    """49 50    def __init__(51        self,52        vae: AutoencoderKL,53        text_encoder: CLIPTextModel,54        tokenizer: CLIPTokenizer,55        unet: UNet2DConditionModel,56        scheduler: DPMSolverMultistepScheduler,57    ):58        super().__init__()59 60        self.register_modules(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler)61 62    @torch.no_grad()63    def __call__(64        self,65        prompt: str,66        image: PIL.Image.Image,67        mask_image: PIL.Image.Image,68        source_points: List[List[int]],69        target_points: List[List[int]],70        t0: Optional[float] = 0.6,71        steps: Optional[int] = 200,72        step_size: Optional[int] = 2,73        image_scale: Optional[float] = 0.3,74        adapt_radius: Optional[int] = 5,75        min_lora_scale: Optional[float] = 0.5,76        generator: Optional[torch.Generator] = None,77    ):78        r"""79        Function invoked when calling the pipeline for image editing.80        Args:81            prompt (`str`, *required*):82                The prompt to guide the image editing.83            image (`PIL.Image.Image`, *required*):84                Which will be edited, parts of the image will be masked out with `mask_image` and edited85                according to `prompt`.86            mask_image (`PIL.Image.Image`, *required*):87                To mask `image`. White pixels in the mask will be edited, while black pixels will be preserved.88            source_points (`List[List[int]]`, *required*):89                Used to mark the starting positions of drag editing in the image, with each pixel represented as a90                `List[int]` of length 2.91            target_points (`List[List[int]]`, *required*):92                Used to mark the target positions of drag editing in the image, with each pixel represented as a93                `List[int]` of length 2.94            t0 (`float`, *optional*, defaults to 0.6):95                The time parameter. Higher t0 improves the fidelity while lowering the faithfulness of the edited images96                and vice versa.97            steps (`int`, *optional*, defaults to 200):98                The number of sampling iterations.99            step_size (`int`, *optional*, defaults to 2):100                The drag diatance of each drag step.101            image_scale (`float`, *optional*, defaults to 0.3):102                To avoid duplicating the content, use image_scale to perturbs the source.103            adapt_radius (`int`, *optional*, defaults to 5):104                The size of the region for copy and paste operations during each step of the drag process.105            min_lora_scale (`float`, *optional*, defaults to 0.5):106                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.107                min_lora_scale specifies the minimum LoRA scale during the image drag-editing process.108            generator ('torch.Generator', *optional*, defaults to None):109                To make generation deterministic(https://pytorch.org/docs/stable/generated/torch.Generator.html).110        Examples:111        ```py112        >>> import PIL113        >>> import torch114        >>> from diffusers import DDIMScheduler, DiffusionPipeline115 116        >>> # Load the pipeline117        >>> model_path = "runwayml/stable-diffusion-v1-5"118        >>> scheduler = DDIMScheduler.from_pretrained(model_path, subfolder="scheduler")119        >>> pipe = DiffusionPipeline.from_pretrained(model_path, scheduler=scheduler, custom_pipeline="sde_drag")120        >>> pipe.to('cuda')121 122        >>> # To save GPU memory, torch.float16 can be used, but it may compromise image quality.123        >>> # If not training LoRA, please avoid using torch.float16124        >>> # pipe.to(torch.float16)125 126        >>> # Provide prompt, image, mask image, and the starting and target points for drag editing.127        >>> prompt = "prompt of the image"128        >>> image = PIL.Image.open('/path/to/image')129        >>> mask_image = PIL.Image.open('/path/to/mask_image')130        >>> source_points = [[123, 456]]131        >>> target_points = [[234, 567]]132 133        >>> # train_lora is optional, and in most cases, using train_lora can better preserve consistency with the original image.134        >>> pipe.train_lora(prompt, image)135 136        >>> output = pipe(prompt, image, mask_image, source_points, target_points)137        >>> output_image = PIL.Image.fromarray(output)138        >>> output_image.save("./output.png")139        ```140        """141 142        self.scheduler.set_timesteps(steps)143 144        noise_scale = (1 - image_scale**2) ** (0.5)145 146        text_embeddings = self._get_text_embed(prompt)147        uncond_embeddings = self._get_text_embed([""])148        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])149 150        latent = self._get_img_latent(image)151 152        mask = mask_image.resize((latent.shape[3], latent.shape[2]))153        mask = torch.tensor(np.array(mask))154        mask = mask.unsqueeze(0).expand_as(latent).to(self.device)155 156        source_points = torch.tensor(source_points).div(torch.tensor([8]), rounding_mode="trunc")157        target_points = torch.tensor(target_points).div(torch.tensor([8]), rounding_mode="trunc")158 159        distance = target_points - source_points160        distance_norm_max = torch.norm(distance.float(), dim=1, keepdim=True).max()161 162        if distance_norm_max <= step_size:163            drag_num = 1164        else:165            drag_num = distance_norm_max.div(torch.tensor([step_size]), rounding_mode="trunc")166            if (distance_norm_max / drag_num - step_size).abs() > (167                distance_norm_max / (drag_num + 1) - step_size168            ).abs():169                drag_num += 1170 171        latents = []172        for i in tqdm(range(int(drag_num)), desc="SDE Drag"):173            source_new = source_points + (i / drag_num * distance).to(torch.int)174            target_new = source_points + ((i + 1) / drag_num * distance).to(torch.int)175 176            latent, noises, hook_latents, lora_scales, cfg_scales = self._forward(177                latent, steps, t0, min_lora_scale, text_embeddings, generator178            )179            latent = self._copy_and_paste(180                latent,181                source_new,182                target_new,183                adapt_radius,184                latent.shape[2] - 1,185                latent.shape[3] - 1,186                image_scale,187                noise_scale,188                generator,189            )190            latent = self._backward(191                latent, mask, steps, t0, noises, hook_latents, lora_scales, cfg_scales, text_embeddings, generator192            )193 194            latents.append(latent)195 196        result_image = 1 / 0.18215 * latents[-1]197 198        with torch.no_grad():199            result_image = self.vae.decode(result_image).sample200 201        result_image = (result_image / 2 + 0.5).clamp(0, 1)202        result_image = result_image.cpu().permute(0, 2, 3, 1).numpy()[0]203        result_image = (result_image * 255).astype(np.uint8)204 205        return result_image206 207    def train_lora(self, prompt, image, lora_step=100, lora_rank=16, generator=None):208        accelerator = Accelerator(gradient_accumulation_steps=1, mixed_precision="fp16")209 210        self.vae.requires_grad_(False)211        self.text_encoder.requires_grad_(False)212        self.unet.requires_grad_(False)213 214        unet_lora_attn_procs = {}215        for name, attn_processor in self.unet.attn_processors.items():216            cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim217            if name.startswith("mid_block"):218                hidden_size = self.unet.config.block_out_channels[-1]219            elif name.startswith("up_blocks"):220                block_id = int(name[len("up_blocks.")])221                hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]222            elif name.startswith("down_blocks"):223                block_id = int(name[len("down_blocks.")])224                hidden_size = self.unet.config.block_out_channels[block_id]225            else:226                raise NotImplementedError("name must start with up_blocks, mid_blocks, or down_blocks")227 228            if isinstance(attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0)):229                lora_attn_processor_class = LoRAAttnAddedKVProcessor230            else:231                lora_attn_processor_class = (232                    LoRAAttnProcessor2_0233                    if hasattr(torch.nn.functional, "scaled_dot_product_attention")234                    else LoRAAttnProcessor235                )236            unet_lora_attn_procs[name] = lora_attn_processor_class(237                hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=lora_rank238            )239 240        self.unet.set_attn_processor(unet_lora_attn_procs)241        unet_lora_layers = AttnProcsLayers(self.unet.attn_processors)242        params_to_optimize = unet_lora_layers.parameters()243 244        optimizer = torch.optim.AdamW(245            params_to_optimize,246            lr=2e-4,247            betas=(0.9, 0.999),248            weight_decay=1e-2,249            eps=1e-08,250        )251 252        lr_scheduler = get_scheduler(253            "constant",254            optimizer=optimizer,255            num_warmup_steps=0,256            num_training_steps=lora_step,257            num_cycles=1,258            power=1.0,259        )260 261        unet_lora_layers = accelerator.prepare_model(unet_lora_layers)262        optimizer = accelerator.prepare_optimizer(optimizer)263        lr_scheduler = accelerator.prepare_scheduler(lr_scheduler)264 265        with torch.no_grad():266            text_inputs = self._tokenize_prompt(prompt, tokenizer_max_length=None)267            text_embedding = self._encode_prompt(268                text_inputs.input_ids, text_inputs.attention_mask, text_encoder_use_attention_mask=False269            )270 271        image_transforms = transforms.Compose(272            [273                transforms.ToTensor(),274                transforms.Normalize([0.5], [0.5]),275            ]276        )277 278        image = image_transforms(image).to(self.device, dtype=self.vae.dtype)279        image = image.unsqueeze(dim=0)280        latents_dist = self.vae.encode(image).latent_dist281 282        for _ in tqdm(range(lora_step), desc="Train LoRA"):283            self.unet.train()284            model_input = latents_dist.sample() * self.vae.config.scaling_factor285 286            # Sample noise that we'll add to the latents287            noise = torch.randn(288                model_input.size(),289                dtype=model_input.dtype,290                layout=model_input.layout,291                device=model_input.device,292                generator=generator,293            )294            bsz, channels, height, width = model_input.shape295 296            # Sample a random timestep for each image297            timesteps = torch.randint(298                0, self.scheduler.config.num_train_timesteps, (bsz,), device=model_input.device, generator=generator299            )300            timesteps = timesteps.long()301 302            # Add noise to the model input according to the noise magnitude at each timestep303            # (this is the forward diffusion process)304            noisy_model_input = self.scheduler.add_noise(model_input, noise, timesteps)305 306            # Predict the noise residual307            model_pred = self.unet(noisy_model_input, timesteps, text_embedding).sample308 309            # Get the target for loss depending on the prediction type310            if self.scheduler.config.prediction_type == "epsilon":311                target = noise312            elif self.scheduler.config.prediction_type == "v_prediction":313                target = self.scheduler.get_velocity(model_input, noise, timesteps)314            else:315                raise ValueError(f"Unknown prediction type {self.scheduler.config.prediction_type}")316 317            loss = torch.nn.functional.mse_loss(model_pred.float(), target.float(), reduction="mean")318            accelerator.backward(loss)319            optimizer.step()320            lr_scheduler.step()321            optimizer.zero_grad()322 323        with tempfile.TemporaryDirectory() as save_lora_dir:324            LoraLoaderMixin.save_lora_weights(325                save_directory=save_lora_dir,326                unet_lora_layers=unet_lora_layers,327                text_encoder_lora_layers=None,328            )329 330            self.unet.load_attn_procs(save_lora_dir)331 332    def _tokenize_prompt(self, prompt, tokenizer_max_length=None):333        if tokenizer_max_length is not None:334            max_length = tokenizer_max_length335        else:336            max_length = self.tokenizer.model_max_length337 338        text_inputs = self.tokenizer(339            prompt,340            truncation=True,341            padding="max_length",342            max_length=max_length,343            return_tensors="pt",344        )345 346        return text_inputs347 348    def _encode_prompt(self, input_ids, attention_mask, text_encoder_use_attention_mask=False):349        text_input_ids = input_ids.to(self.device)350 351        if text_encoder_use_attention_mask:352            attention_mask = attention_mask.to(self.device)353        else:354            attention_mask = None355 356        prompt_embeds = self.text_encoder(357            text_input_ids,358            attention_mask=attention_mask,359        )360        prompt_embeds = prompt_embeds[0]361 362        return prompt_embeds363 364    @torch.no_grad()365    def _get_text_embed(self, prompt):366        text_input = self.tokenizer(367            prompt,368            padding="max_length",369            max_length=self.tokenizer.model_max_length,370            truncation=True,371            return_tensors="pt",372        )373        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]374        return text_embeddings375 376    def _copy_and_paste(377        self, latent, source_new, target_new, adapt_radius, max_height, max_width, image_scale, noise_scale, generator378    ):379        def adaption_r(source, target, adapt_radius, max_height, max_width):380            r_x_lower = min(adapt_radius, source[0], target[0])381            r_x_upper = min(adapt_radius, max_width - source[0], max_width - target[0])382            r_y_lower = min(adapt_radius, source[1], target[1])383            r_y_upper = min(adapt_radius, max_height - source[1], max_height - target[1])384            return r_x_lower, r_x_upper, r_y_lower, r_y_upper385 386        for source_, target_ in zip(source_new, target_new):387            r_x_lower, r_x_upper, r_y_lower, r_y_upper = adaption_r(388                source_, target_, adapt_radius, max_height, max_width389            )390 391            source_feature = latent[392                :, :, source_[1] - r_y_lower : source_[1] + r_y_upper, source_[0] - r_x_lower : source_[0] + r_x_upper393            ].clone()394 395            latent[396                :, :, source_[1] - r_y_lower : source_[1] + r_y_upper, source_[0] - r_x_lower : source_[0] + r_x_upper397            ] = image_scale * source_feature + noise_scale * torch.randn(398                latent.shape[0],399                4,400                r_y_lower + r_y_upper,401                r_x_lower + r_x_upper,402                device=self.device,403                generator=generator,404            )405 406            latent[407                :, :, target_[1] - r_y_lower : target_[1] + r_y_upper, target_[0] - r_x_lower : target_[0] + r_x_upper408            ] = source_feature * 1.1409        return latent410 411    @torch.no_grad()412    def _get_img_latent(self, image, height=None, weight=None):413        data = image.convert("RGB")414        if height is not None:415            data = data.resize((weight, height))416        transform = transforms.ToTensor()417        data = transform(data).unsqueeze(0)418        data = (data * 2.0) - 1.0419        data = data.to(self.device, dtype=self.vae.dtype)420        latent = self.vae.encode(data).latent_dist.sample()421        latent = 0.18215 * latent422        return latent423 424    @torch.no_grad()425    def _get_eps(self, latent, timestep, guidance_scale, text_embeddings, lora_scale=None):426        latent_model_input = torch.cat([latent] * 2) if guidance_scale > 1.0 else latent427        text_embeddings = text_embeddings if guidance_scale > 1.0 else text_embeddings.chunk(2)[1]428 429        cross_attention_kwargs = None if lora_scale is None else {"scale": lora_scale}430 431        with torch.no_grad():432            noise_pred = self.unet(433                latent_model_input,434                timestep,435                encoder_hidden_states=text_embeddings,436                cross_attention_kwargs=cross_attention_kwargs,437            ).sample438 439        if guidance_scale > 1.0:440            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)441        elif guidance_scale == 1.0:442            noise_pred_text = noise_pred443            noise_pred_uncond = 0.0444        else:445            raise NotImplementedError(guidance_scale)446        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)447 448        return noise_pred449 450    def _forward_sde(451        self, timestep, sample, guidance_scale, text_embeddings, steps, eta=1.0, lora_scale=None, generator=None452    ):453        num_train_timesteps = len(self.scheduler)454        alphas_cumprod = self.scheduler.alphas_cumprod455        initial_alpha_cumprod = torch.tensor(1.0)456 457        prev_timestep = timestep + num_train_timesteps // steps458 459        alpha_prod_t = alphas_cumprod[timestep] if timestep >= 0 else initial_alpha_cumprod460        alpha_prod_t_prev = alphas_cumprod[prev_timestep]461 462        beta_prod_t_prev = 1 - alpha_prod_t_prev463 464        x_prev = (alpha_prod_t_prev / alpha_prod_t) ** (0.5) * sample + (1 - alpha_prod_t_prev / alpha_prod_t) ** (465            0.5466        ) * torch.randn(467            sample.size(), dtype=sample.dtype, layout=sample.layout, device=self.device, generator=generator468        )469        eps = self._get_eps(x_prev, prev_timestep, guidance_scale, text_embeddings, lora_scale)470 471        sigma_t_prev = (472            eta473            * (1 - alpha_prod_t) ** (0.5)474            * (1 - alpha_prod_t_prev / (1 - alpha_prod_t_prev) * (1 - alpha_prod_t) / alpha_prod_t) ** (0.5)475        )476 477        pred_original_sample = (x_prev - beta_prod_t_prev ** (0.5) * eps) / alpha_prod_t_prev ** (0.5)478        pred_sample_direction_coeff = (1 - alpha_prod_t - sigma_t_prev**2) ** (0.5)479 480        noise = (481            sample - alpha_prod_t ** (0.5) * pred_original_sample - pred_sample_direction_coeff * eps482        ) / sigma_t_prev483 484        return x_prev, noise485 486    def _sample(487        self,488        timestep,489        sample,490        guidance_scale,491        text_embeddings,492        steps,493        sde=False,494        noise=None,495        eta=1.0,496        lora_scale=None,497        generator=None,498    ):499        num_train_timesteps = len(self.scheduler)500        alphas_cumprod = self.scheduler.alphas_cumprod501        final_alpha_cumprod = torch.tensor(1.0)502 503        eps = self._get_eps(sample, timestep, guidance_scale, text_embeddings, lora_scale)504 505        prev_timestep = timestep - num_train_timesteps // steps506 507        alpha_prod_t = alphas_cumprod[timestep]508        alpha_prod_t_prev = alphas_cumprod[prev_timestep] if prev_timestep >= 0 else final_alpha_cumprod509 510        beta_prod_t = 1 - alpha_prod_t511 512        sigma_t = (513            eta514            * ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) ** (0.5)515            * (1 - alpha_prod_t / alpha_prod_t_prev) ** (0.5)516            if sde517            else 0518        )519 520        pred_original_sample = (sample - beta_prod_t ** (0.5) * eps) / alpha_prod_t ** (0.5)521        pred_sample_direction_coeff = (1 - alpha_prod_t_prev - sigma_t**2) ** (0.5)522 523        noise = (524            torch.randn(525                sample.size(), dtype=sample.dtype, layout=sample.layout, device=self.device, generator=generator526            )527            if noise is None528            else noise529        )530        latent = (531            alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction_coeff * eps + sigma_t * noise532        )533 534        return latent535 536    def _forward(self, latent, steps, t0, lora_scale_min, text_embeddings, generator):537        def scale_schedule(begin, end, n, length, type="linear"):538            if type == "constant":539                return end540            elif type == "linear":541                return begin + (end - begin) * n / length542            elif type == "cos":543                factor = (1 - math.cos(n * math.pi / length)) / 2544                return (1 - factor) * begin + factor * end545            else:546                raise NotImplementedError(type)547 548        noises = []549        latents = []550        lora_scales = []551        cfg_scales = []552        latents.append(latent)553        t0 = int(t0 * steps)554        t_begin = steps - t0555 556        length = len(self.scheduler.timesteps[t_begin - 1 : -1]) - 1557        index = 1558        for t in self.scheduler.timesteps[t_begin:].flip(dims=[0]):559            lora_scale = scale_schedule(1, lora_scale_min, index, length, type="cos")560            cfg_scale = scale_schedule(1, 3.0, index, length, type="linear")561            latent, noise = self._forward_sde(562                t, latent, cfg_scale, text_embeddings, steps, lora_scale=lora_scale, generator=generator563            )564 565            noises.append(noise)566            latents.append(latent)567            lora_scales.append(lora_scale)568            cfg_scales.append(cfg_scale)569            index += 1570        return latent, noises, latents, lora_scales, cfg_scales571 572    def _backward(573        self, latent, mask, steps, t0, noises, hook_latents, lora_scales, cfg_scales, text_embeddings, generator574    ):575        t0 = int(t0 * steps)576        t_begin = steps - t0577 578        hook_latent = hook_latents.pop()579        latent = torch.where(mask > 128, latent, hook_latent)580        for t in self.scheduler.timesteps[t_begin - 1 : -1]:581            latent = self._sample(582                t,583                latent,584                cfg_scales.pop(),585                text_embeddings,586                steps,587                sde=True,588                noise=noises.pop(),589                lora_scale=lora_scales.pop(),590                generator=generator,591            )592            hook_latent = hook_latents.pop()593            latent = torch.where(mask > 128, latent, hook_latent)594        return latent595