diffusers/community-pipelines-mirror

1import inspect2from copy import deepcopy3from enum import Enum4from typing import List, Optional, Tuple, Union5 6import torch7from tqdm.auto import tqdm8 9from diffusers.models import AutoencoderKL, UNet2DConditionModel10from diffusers.pipelines.pipeline_utils import DiffusionPipeline11from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker12from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler13from diffusers.utils import logging14 15 16try:17    from ligo.segments import segment18    from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer19except ImportError:20    raise ImportError("Please install transformers and ligo-segments to use the mixture pipeline")21 22logger = logging.get_logger(__name__)  # pylint: disable=invalid-name23 24EXAMPLE_DOC_STRING = """25    Examples:26        ```py27        >>> from diffusers import LMSDiscreteScheduler, DiffusionPipeline28 29        >>> scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)30        >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", scheduler=scheduler, custom_pipeline="mixture_tiling")31        >>> pipeline.to("cuda")32 33        >>> image = pipeline(34        >>>     prompt=[[35        >>>         "A charming house in the countryside, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",36        >>>         "A dirt road in the countryside crossing pastures, by jakub rozalski, sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece",37        >>>         "An old and rusty giant robot lying on a dirt road, by jakub rozalski, dark sunset lighting, elegant, highly detailed, smooth, sharp focus, artstation, stunning masterpiece"38        >>>     ]],39        >>>     tile_height=640,40        >>>     tile_width=640,41        >>>     tile_row_overlap=0,42        >>>     tile_col_overlap=256,43        >>>     guidance_scale=8,44        >>>     seed=7178915308,45        >>>     num_inference_steps=50,46    >>> )["images"][0]47        ```48"""49 50 51def _tile2pixel_indices(tile_row, tile_col, tile_width, tile_height, tile_row_overlap, tile_col_overlap):52    """Given a tile row and column numbers returns the range of pixels affected by that tiles in the overall image53 54    Returns a tuple with:55        - Starting coordinates of rows in pixel space56        - Ending coordinates of rows in pixel space57        - Starting coordinates of columns in pixel space58        - Ending coordinates of columns in pixel space59    """60    px_row_init = 0 if tile_row == 0 else tile_row * (tile_height - tile_row_overlap)61    px_row_end = px_row_init + tile_height62    px_col_init = 0 if tile_col == 0 else tile_col * (tile_width - tile_col_overlap)63    px_col_end = px_col_init + tile_width64    return px_row_init, px_row_end, px_col_init, px_col_end65 66 67def _pixel2latent_indices(px_row_init, px_row_end, px_col_init, px_col_end):68    """Translates coordinates in pixel space to coordinates in latent space"""69    return px_row_init // 8, px_row_end // 8, px_col_init // 8, px_col_end // 870 71 72def _tile2latent_indices(tile_row, tile_col, tile_width, tile_height, tile_row_overlap, tile_col_overlap):73    """Given a tile row and column numbers returns the range of latents affected by that tiles in the overall image74 75    Returns a tuple with:76        - Starting coordinates of rows in latent space77        - Ending coordinates of rows in latent space78        - Starting coordinates of columns in latent space79        - Ending coordinates of columns in latent space80    """81    px_row_init, px_row_end, px_col_init, px_col_end = _tile2pixel_indices(82        tile_row, tile_col, tile_width, tile_height, tile_row_overlap, tile_col_overlap83    )84    return _pixel2latent_indices(px_row_init, px_row_end, px_col_init, px_col_end)85 86 87def _tile2latent_exclusive_indices(88    tile_row, tile_col, tile_width, tile_height, tile_row_overlap, tile_col_overlap, rows, columns89):90    """Given a tile row and column numbers returns the range of latents affected only by that tile in the overall image91 92    Returns a tuple with:93        - Starting coordinates of rows in latent space94        - Ending coordinates of rows in latent space95        - Starting coordinates of columns in latent space96        - Ending coordinates of columns in latent space97    """98    row_init, row_end, col_init, col_end = _tile2latent_indices(99        tile_row, tile_col, tile_width, tile_height, tile_row_overlap, tile_col_overlap100    )101    row_segment = segment(row_init, row_end)102    col_segment = segment(col_init, col_end)103    # Iterate over the rest of tiles, clipping the region for the current tile104    for row in range(rows):105        for column in range(columns):106            if row != tile_row and column != tile_col:107                clip_row_init, clip_row_end, clip_col_init, clip_col_end = _tile2latent_indices(108                    row, column, tile_width, tile_height, tile_row_overlap, tile_col_overlap109                )110                row_segment = row_segment - segment(clip_row_init, clip_row_end)111                col_segment = col_segment - segment(clip_col_init, clip_col_end)112    # return row_init, row_end, col_init, col_end113    return row_segment[0], row_segment[1], col_segment[0], col_segment[1]114 115 116class StableDiffusionExtrasMixin:117    """Mixin providing additional convenience method to Stable Diffusion pipelines"""118 119    def decode_latents(self, latents, cpu_vae=False):120        """Decodes a given array of latents into pixel space"""121        # scale and decode the image latents with vae122        if cpu_vae:123            lat = deepcopy(latents).cpu()124            vae = deepcopy(self.vae).cpu()125        else:126            lat = latents127            vae = self.vae128 129        lat = 1 / 0.18215 * lat130        image = vae.decode(lat).sample131 132        image = (image / 2 + 0.5).clamp(0, 1)133        image = image.cpu().permute(0, 2, 3, 1).numpy()134 135        return self.numpy_to_pil(image)136 137 138class StableDiffusionTilingPipeline(DiffusionPipeline, StableDiffusionExtrasMixin):139    def __init__(140        self,141        vae: AutoencoderKL,142        text_encoder: CLIPTextModel,143        tokenizer: CLIPTokenizer,144        unet: UNet2DConditionModel,145        scheduler: Union[DDIMScheduler, PNDMScheduler],146        safety_checker: StableDiffusionSafetyChecker,147        feature_extractor: CLIPFeatureExtractor,148    ):149        super().__init__()150        self.register_modules(151            vae=vae,152            text_encoder=text_encoder,153            tokenizer=tokenizer,154            unet=unet,155            scheduler=scheduler,156            safety_checker=safety_checker,157            feature_extractor=feature_extractor,158        )159 160    class SeedTilesMode(Enum):161        """Modes in which the latents of a particular tile can be re-seeded"""162 163        FULL = "full"164        EXCLUSIVE = "exclusive"165 166    @torch.no_grad()167    def __call__(168        self,169        prompt: Union[str, List[List[str]]],170        num_inference_steps: Optional[int] = 50,171        guidance_scale: Optional[float] = 7.5,172        eta: Optional[float] = 0.0,173        seed: Optional[int] = None,174        tile_height: Optional[int] = 512,175        tile_width: Optional[int] = 512,176        tile_row_overlap: Optional[int] = 256,177        tile_col_overlap: Optional[int] = 256,178        guidance_scale_tiles: Optional[List[List[float]]] = None,179        seed_tiles: Optional[List[List[int]]] = None,180        seed_tiles_mode: Optional[Union[str, List[List[str]]]] = "full",181        seed_reroll_regions: Optional[List[Tuple[int, int, int, int, int]]] = None,182        cpu_vae: Optional[bool] = False,183    ):184        r"""185        Function to run the diffusion pipeline with tiling support.186 187        Args:188            prompt: either a single string (no tiling) or a list of lists with all the prompts to use (one list for each row of tiles). This will also define the tiling structure.189            num_inference_steps: number of diffusions steps.190            guidance_scale: classifier-free guidance.191            seed: general random seed to initialize latents.192            tile_height: height in pixels of each grid tile.193            tile_width: width in pixels of each grid tile.194            tile_row_overlap: number of overlap pixels between tiles in consecutive rows.195            tile_col_overlap: number of overlap pixels between tiles in consecutive columns.196            guidance_scale_tiles: specific weights for classifier-free guidance in each tile.197            guidance_scale_tiles: specific weights for classifier-free guidance in each tile. If None, the value provided in guidance_scale will be used.198            seed_tiles: specific seeds for the initialization latents in each tile. These will override the latents generated for the whole canvas using the standard seed parameter.199            seed_tiles_mode: either "full" "exclusive". If "full", all the latents affected by the tile be overriden. If "exclusive", only the latents that are affected exclusively by this tile (and no other tiles) will be overriden.200            seed_reroll_regions: a list of tuples in the form (start row, end row, start column, end column, seed) defining regions in pixel space for which the latents will be overriden using the given seed. Takes priority over seed_tiles.201            cpu_vae: the decoder from latent space to pixel space can require too mucho GPU RAM for large images. If you find out of memory errors at the end of the generation process, try setting this parameter to True to run the decoder in CPU. Slower, but should run without memory issues.202 203        Examples:204 205        Returns:206            A PIL image with the generated image.207 208        """209        if not isinstance(prompt, list) or not all(isinstance(row, list) for row in prompt):210            raise ValueError(f"`prompt` has to be a list of lists but is {type(prompt)}")211        grid_rows = len(prompt)212        grid_cols = len(prompt[0])213        if not all(len(row) == grid_cols for row in prompt):214            raise ValueError("All prompt rows must have the same number of prompt columns")215        if not isinstance(seed_tiles_mode, str) and (216            not isinstance(seed_tiles_mode, list) or not all(isinstance(row, list) for row in seed_tiles_mode)217        ):218            raise ValueError(f"`seed_tiles_mode` has to be a string or list of lists but is {type(prompt)}")219        if isinstance(seed_tiles_mode, str):220            seed_tiles_mode = [[seed_tiles_mode for _ in range(len(row))] for row in prompt]221 222        modes = [mode.value for mode in self.SeedTilesMode]223        if any(mode not in modes for row in seed_tiles_mode for mode in row):224            raise ValueError(f"Seed tiles mode must be one of {modes}")225        if seed_reroll_regions is None:226            seed_reroll_regions = []227        batch_size = 1228 229        # create original noisy latents using the timesteps230        height = tile_height + (grid_rows - 1) * (tile_height - tile_row_overlap)231        width = tile_width + (grid_cols - 1) * (tile_width - tile_col_overlap)232        latents_shape = (batch_size, self.unet.config.in_channels, height // 8, width // 8)233        generator = torch.Generator("cuda").manual_seed(seed)234        latents = torch.randn(latents_shape, generator=generator, device=self.device)235 236        # overwrite latents for specific tiles if provided237        if seed_tiles is not None:238            for row in range(grid_rows):239                for col in range(grid_cols):240                    if (seed_tile := seed_tiles[row][col]) is not None:241                        mode = seed_tiles_mode[row][col]242                        if mode == self.SeedTilesMode.FULL.value:243                            row_init, row_end, col_init, col_end = _tile2latent_indices(244                                row, col, tile_width, tile_height, tile_row_overlap, tile_col_overlap245                            )246                        else:247                            row_init, row_end, col_init, col_end = _tile2latent_exclusive_indices(248                                row,249                                col,250                                tile_width,251                                tile_height,252                                tile_row_overlap,253                                tile_col_overlap,254                                grid_rows,255                                grid_cols,256                            )257                        tile_generator = torch.Generator("cuda").manual_seed(seed_tile)258                        tile_shape = (latents_shape[0], latents_shape[1], row_end - row_init, col_end - col_init)259                        latents[:, :, row_init:row_end, col_init:col_end] = torch.randn(260                            tile_shape, generator=tile_generator, device=self.device261                        )262 263        # overwrite again for seed reroll regions264        for row_init, row_end, col_init, col_end, seed_reroll in seed_reroll_regions:265            row_init, row_end, col_init, col_end = _pixel2latent_indices(266                row_init, row_end, col_init, col_end267            )  # to latent space coordinates268            reroll_generator = torch.Generator("cuda").manual_seed(seed_reroll)269            region_shape = (latents_shape[0], latents_shape[1], row_end - row_init, col_end - col_init)270            latents[:, :, row_init:row_end, col_init:col_end] = torch.randn(271                region_shape, generator=reroll_generator, device=self.device272            )273 274        # Prepare scheduler275        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())276        extra_set_kwargs = {}277        if accepts_offset:278            extra_set_kwargs["offset"] = 1279        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)280        # if we use LMSDiscreteScheduler, let's make sure latents are multiplied by sigmas281        if isinstance(self.scheduler, LMSDiscreteScheduler):282            latents = latents * self.scheduler.sigmas[0]283 284        # get prompts text embeddings285        text_input = [286            [287                self.tokenizer(288                    col,289                    padding="max_length",290                    max_length=self.tokenizer.model_max_length,291                    truncation=True,292                    return_tensors="pt",293                )294                for col in row295            ]296            for row in prompt297        ]298        text_embeddings = [[self.text_encoder(col.input_ids.to(self.device))[0] for col in row] for row in text_input]299 300        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)301        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`302        # corresponds to doing no classifier free guidance.303        do_classifier_free_guidance = guidance_scale > 1.0  # TODO: also active if any tile has guidance scale304        # get unconditional embeddings for classifier free guidance305        if do_classifier_free_guidance:306            for i in range(grid_rows):307                for j in range(grid_cols):308                    max_length = text_input[i][j].input_ids.shape[-1]309                    uncond_input = self.tokenizer(310                        [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"311                    )312                    uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]313 314                    # For classifier free guidance, we need to do two forward passes.315                    # Here we concatenate the unconditional and text embeddings into a single batch316                    # to avoid doing two forward passes317                    text_embeddings[i][j] = torch.cat([uncond_embeddings, text_embeddings[i][j]])318 319        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature320        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.321        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502322        # and should be between [0, 1]323        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())324        extra_step_kwargs = {}325        if accepts_eta:326            extra_step_kwargs["eta"] = eta327 328        # Mask for tile weights strength329        tile_weights = self._gaussian_weights(tile_width, tile_height, batch_size)330 331        # Diffusion timesteps332        for i, t in tqdm(enumerate(self.scheduler.timesteps)):333            # Diffuse each tile334            noise_preds = []335            for row in range(grid_rows):336                noise_preds_row = []337                for col in range(grid_cols):338                    px_row_init, px_row_end, px_col_init, px_col_end = _tile2latent_indices(339                        row, col, tile_width, tile_height, tile_row_overlap, tile_col_overlap340                    )341                    tile_latents = latents[:, :, px_row_init:px_row_end, px_col_init:px_col_end]342                    # expand the latents if we are doing classifier free guidance343                    latent_model_input = torch.cat([tile_latents] * 2) if do_classifier_free_guidance else tile_latents344                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)345                    # predict the noise residual346                    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings[row][col])[347                        "sample"348                    ]349                    # perform guidance350                    if do_classifier_free_guidance:351                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)352                        guidance = (353                            guidance_scale354                            if guidance_scale_tiles is None or guidance_scale_tiles[row][col] is None355                            else guidance_scale_tiles[row][col]356                        )357                        noise_pred_tile = noise_pred_uncond + guidance * (noise_pred_text - noise_pred_uncond)358                        noise_preds_row.append(noise_pred_tile)359                noise_preds.append(noise_preds_row)360            # Stitch noise predictions for all tiles361            noise_pred = torch.zeros(latents.shape, device=self.device)362            contributors = torch.zeros(latents.shape, device=self.device)363            # Add each tile contribution to overall latents364            for row in range(grid_rows):365                for col in range(grid_cols):366                    px_row_init, px_row_end, px_col_init, px_col_end = _tile2latent_indices(367                        row, col, tile_width, tile_height, tile_row_overlap, tile_col_overlap368                    )369                    noise_pred[:, :, px_row_init:px_row_end, px_col_init:px_col_end] += (370                        noise_preds[row][col] * tile_weights371                    )372                    contributors[:, :, px_row_init:px_row_end, px_col_init:px_col_end] += tile_weights373            # Average overlapping areas with more than 1 contributor374            noise_pred /= contributors375 376            # compute the previous noisy sample x_t -> x_t-1377            latents = self.scheduler.step(noise_pred, t, latents).prev_sample378 379        # scale and decode the image latents with vae380        image = self.decode_latents(latents, cpu_vae)381 382        return {"images": image}383 384    def _gaussian_weights(self, tile_width, tile_height, nbatches):385        """Generates a gaussian mask of weights for tile contributions"""386        import numpy as np387        from numpy import exp, pi, sqrt388 389        latent_width = tile_width // 8390        latent_height = tile_height // 8391 392        var = 0.01393        midpoint = (latent_width - 1) / 2  # -1 because index goes from 0 to latent_width - 1394        x_probs = [395            exp(-(x - midpoint) * (x - midpoint) / (latent_width * latent_width) / (2 * var)) / sqrt(2 * pi * var)396            for x in range(latent_width)397        ]398        midpoint = latent_height / 2399        y_probs = [400            exp(-(y - midpoint) * (y - midpoint) / (latent_height * latent_height) / (2 * var)) / sqrt(2 * pi * var)401            for y in range(latent_height)402        ]403 404        weights = np.outer(y_probs, x_probs)405        return torch.tile(torch.tensor(weights, device=self.device), (nbatches, self.unet.config.in_channels, 1, 1))406