diffusers/community-pipelines-mirror

1import re2from copy import deepcopy3from dataclasses import asdict, dataclass4from enum import Enum5from typing import List, Optional, Union6 7import numpy as np8import torch9from numpy import exp, pi, sqrt10from torchvision.transforms.functional import resize11from tqdm.auto import tqdm12from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer13 14from diffusers.models import AutoencoderKL, UNet2DConditionModel15from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin16from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker17from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler18 19 20def preprocess_image(image):21    from PIL import Image22 23    """Preprocess an input image24 25    Same as26    https://github.com/huggingface/diffusers/blob/1138d63b519e37f0ce04e027b9f4a3261d27c628/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py#L4427    """28    w, h = image.size29    w, h = (x - x % 32 for x in (w, h))  # resize to integer multiple of 3230    image = image.resize((w, h), resample=Image.LANCZOS)31    image = np.array(image).astype(np.float32) / 255.032    image = image[None].transpose(0, 3, 1, 2)33    image = torch.from_numpy(image)34    return 2.0 * image - 1.035 36 37@dataclass38class CanvasRegion:39    """Class defining a rectangular region in the canvas"""40 41    row_init: int  # Region starting row in pixel space (included)42    row_end: int  # Region end row in pixel space (not included)43    col_init: int  # Region starting column in pixel space (included)44    col_end: int  # Region end column in pixel space (not included)45    region_seed: int = None  # Seed for random operations in this region46    noise_eps: float = 0.0  # Deviation of a zero-mean gaussian noise to be applied over the latents in this region. Useful for slightly "rerolling" latents47 48    def __post_init__(self):49        # Initialize arguments if not specified50        if self.region_seed is None:51            self.region_seed = np.random.randint(9999999999)52        # Check coordinates are non-negative53        for coord in [self.row_init, self.row_end, self.col_init, self.col_end]:54            if coord < 0:55                raise ValueError(56                    f"A CanvasRegion must be defined with non-negative indices, found ({self.row_init}, {self.row_end}, {self.col_init}, {self.col_end})"57                )58        # Check coordinates are divisible by 8, else we end up with nasty rounding error when mapping to latent space59        for coord in [self.row_init, self.row_end, self.col_init, self.col_end]:60            if coord // 8 != coord / 8:61                raise ValueError(62                    f"A CanvasRegion must be defined with locations divisible by 8, found ({self.row_init}-{self.row_end}, {self.col_init}-{self.col_end})"63                )64        # Check noise eps is non-negative65        if self.noise_eps < 0:66            raise ValueError(f"A CanvasRegion must be defined noises eps non-negative, found {self.noise_eps}")67        # Compute coordinates for this region in latent space68        self.latent_row_init = self.row_init // 869        self.latent_row_end = self.row_end // 870        self.latent_col_init = self.col_init // 871        self.latent_col_end = self.col_end // 872 73    @property74    def width(self):75        return self.col_end - self.col_init76 77    @property78    def height(self):79        return self.row_end - self.row_init80 81    def get_region_generator(self, device="cpu"):82        """Creates a torch.Generator based on the random seed of this region"""83        # Initialize region generator84        return torch.Generator(device).manual_seed(self.region_seed)85 86    @property87    def __dict__(self):88        return asdict(self)89 90 91class MaskModes(Enum):92    """Modes in which the influence of diffuser is masked"""93 94    CONSTANT = "constant"95    GAUSSIAN = "gaussian"96    QUARTIC = "quartic"  # See https://en.wikipedia.org/wiki/Kernel_(statistics)97 98 99@dataclass100class DiffusionRegion(CanvasRegion):101    """Abstract class defining a region where some class of diffusion process is acting"""102 103    pass104 105 106@dataclass107class Text2ImageRegion(DiffusionRegion):108    """Class defining a region where a text guided diffusion process is acting"""109 110    prompt: str = ""  # Text prompt guiding the diffuser in this region111    guidance_scale: float = 7.5  # Guidance scale of the diffuser in this region. If None, randomize112    mask_type: MaskModes = MaskModes.GAUSSIAN.value  # Kind of weight mask applied to this region113    mask_weight: float = 1.0  # Global weights multiplier of the mask114    tokenized_prompt = None  # Tokenized prompt115    encoded_prompt = None  # Encoded prompt116 117    def __post_init__(self):118        super().__post_init__()119        # Mask weight cannot be negative120        if self.mask_weight < 0:121            raise ValueError(122                f"A Text2ImageRegion must be defined with non-negative mask weight, found {self.mask_weight}"123            )124        # Mask type must be an actual known mask125        if self.mask_type not in [e.value for e in MaskModes]:126            raise ValueError(127                f"A Text2ImageRegion was defined with mask {self.mask_type}, which is not an accepted mask ({[e.value for e in MaskModes]})"128            )129        # Randomize arguments if given as None130        if self.guidance_scale is None:131            self.guidance_scale = np.random.randint(5, 30)132        # Clean prompt133        self.prompt = re.sub(" +", " ", self.prompt).replace("\n", " ")134 135    def tokenize_prompt(self, tokenizer):136        """Tokenizes the prompt for this diffusion region using a given tokenizer"""137        self.tokenized_prompt = tokenizer(138            self.prompt,139            padding="max_length",140            max_length=tokenizer.model_max_length,141            truncation=True,142            return_tensors="pt",143        )144 145    def encode_prompt(self, text_encoder, device):146        """Encodes the previously tokenized prompt for this diffusion region using a given encoder"""147        assert self.tokenized_prompt is not None, ValueError(148            "Prompt in diffusion region must be tokenized before encoding"149        )150        self.encoded_prompt = text_encoder(self.tokenized_prompt.input_ids.to(device))[0]151 152 153@dataclass154class Image2ImageRegion(DiffusionRegion):155    """Class defining a region where an image guided diffusion process is acting"""156 157    reference_image: torch.Tensor = None158    strength: float = 0.8  # Strength of the image159 160    def __post_init__(self):161        super().__post_init__()162        if self.reference_image is None:163            raise ValueError("Must provide a reference image when creating an Image2ImageRegion")164        if self.strength < 0 or self.strength > 1:165            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {self.strength}")166        # Rescale image to region shape167        self.reference_image = resize(self.reference_image, size=[self.height, self.width])168 169    def encode_reference_image(self, encoder, device, generator, cpu_vae=False):170        """Encodes the reference image for this Image2Image region into the latent space"""171        # Place encoder in CPU or not following the parameter cpu_vae172        if cpu_vae:173            # Note here we use mean instead of sample, to avoid moving also generator to CPU, which is troublesome174            self.reference_latents = encoder.cpu().encode(self.reference_image).latent_dist.mean.to(device)175        else:176            self.reference_latents = encoder.encode(self.reference_image.to(device)).latent_dist.sample(177                generator=generator178            )179        self.reference_latents = 0.18215 * self.reference_latents180 181    @property182    def __dict__(self):183        # This class requires special casting to dict because of the reference_image tensor. Otherwise it cannot be casted to JSON184 185        # Get all basic fields from parent class186        super_fields = {key: getattr(self, key) for key in DiffusionRegion.__dataclass_fields__.keys()}187        # Pack other fields188        return {**super_fields, "reference_image": self.reference_image.cpu().tolist(), "strength": self.strength}189 190 191class RerollModes(Enum):192    """Modes in which the reroll regions operate"""193 194    RESET = "reset"  # Completely reset the random noise in the region195    EPSILON = "epsilon"  # Alter slightly the latents in the region196 197 198@dataclass199class RerollRegion(CanvasRegion):200    """Class defining a rectangular canvas region in which initial latent noise will be rerolled"""201 202    reroll_mode: RerollModes = RerollModes.RESET.value203 204 205@dataclass206class MaskWeightsBuilder:207    """Auxiliary class to compute a tensor of weights for a given diffusion region"""208 209    latent_space_dim: int  # Size of the U-net latent space210    nbatch: int = 1  # Batch size in the U-net211 212    def compute_mask_weights(self, region: DiffusionRegion) -> torch.tensor:213        """Computes a tensor of weights for a given diffusion region"""214        MASK_BUILDERS = {215            MaskModes.CONSTANT.value: self._constant_weights,216            MaskModes.GAUSSIAN.value: self._gaussian_weights,217            MaskModes.QUARTIC.value: self._quartic_weights,218        }219        return MASK_BUILDERS[region.mask_type](region)220 221    def _constant_weights(self, region: DiffusionRegion) -> torch.tensor:222        """Computes a tensor of constant for a given diffusion region"""223        latent_width = region.latent_col_end - region.latent_col_init224        latent_height = region.latent_row_end - region.latent_row_init225        return torch.ones(self.nbatch, self.latent_space_dim, latent_height, latent_width) * region.mask_weight226 227    def _gaussian_weights(self, region: DiffusionRegion) -> torch.tensor:228        """Generates a gaussian mask of weights for tile contributions"""229        latent_width = region.latent_col_end - region.latent_col_init230        latent_height = region.latent_row_end - region.latent_row_init231 232        var = 0.01233        midpoint = (latent_width - 1) / 2  # -1 because index goes from 0 to latent_width - 1234        x_probs = [235            exp(-(x - midpoint) * (x - midpoint) / (latent_width * latent_width) / (2 * var)) / sqrt(2 * pi * var)236            for x in range(latent_width)237        ]238        midpoint = (latent_height - 1) / 2239        y_probs = [240            exp(-(y - midpoint) * (y - midpoint) / (latent_height * latent_height) / (2 * var)) / sqrt(2 * pi * var)241            for y in range(latent_height)242        ]243 244        weights = np.outer(y_probs, x_probs) * region.mask_weight245        return torch.tile(torch.tensor(weights), (self.nbatch, self.latent_space_dim, 1, 1))246 247    def _quartic_weights(self, region: DiffusionRegion) -> torch.tensor:248        """Generates a quartic mask of weights for tile contributions249 250        The quartic kernel has bounded support over the diffusion region, and a smooth decay to the region limits.251        """252        quartic_constant = 15.0 / 16.0253 254        support = (np.array(range(region.latent_col_init, region.latent_col_end)) - region.latent_col_init) / (255            region.latent_col_end - region.latent_col_init - 1256        ) * 1.99 - (1.99 / 2.0)257        x_probs = quartic_constant * np.square(1 - np.square(support))258        support = (np.array(range(region.latent_row_init, region.latent_row_end)) - region.latent_row_init) / (259            region.latent_row_end - region.latent_row_init - 1260        ) * 1.99 - (1.99 / 2.0)261        y_probs = quartic_constant * np.square(1 - np.square(support))262 263        weights = np.outer(y_probs, x_probs) * region.mask_weight264        return torch.tile(torch.tensor(weights), (self.nbatch, self.latent_space_dim, 1, 1))265 266 267class StableDiffusionCanvasPipeline(DiffusionPipeline, StableDiffusionMixin):268    """Stable Diffusion pipeline that mixes several diffusers in the same canvas"""269 270    def __init__(271        self,272        vae: AutoencoderKL,273        text_encoder: CLIPTextModel,274        tokenizer: CLIPTokenizer,275        unet: UNet2DConditionModel,276        scheduler: Union[DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler],277        safety_checker: StableDiffusionSafetyChecker,278        feature_extractor: CLIPFeatureExtractor,279    ):280        super().__init__()281        self.register_modules(282            vae=vae,283            text_encoder=text_encoder,284            tokenizer=tokenizer,285            unet=unet,286            scheduler=scheduler,287            safety_checker=safety_checker,288            feature_extractor=feature_extractor,289        )290 291    def decode_latents(self, latents, cpu_vae=False):292        """Decodes a given array of latents into pixel space"""293        # scale and decode the image latents with vae294        if cpu_vae:295            lat = deepcopy(latents).cpu()296            vae = deepcopy(self.vae).cpu()297        else:298            lat = latents299            vae = self.vae300 301        lat = 1 / 0.18215 * lat302        image = vae.decode(lat).sample303 304        image = (image / 2 + 0.5).clamp(0, 1)305        image = image.cpu().permute(0, 2, 3, 1).numpy()306 307        return self.numpy_to_pil(image)308 309    def get_latest_timestep_img2img(self, num_inference_steps, strength):310        """Finds the latest timesteps where an img2img strength does not impose latents anymore"""311        # get the original timestep using init_timestep312        offset = self.scheduler.config.get("steps_offset", 0)313        init_timestep = int(num_inference_steps * (1 - strength)) + offset314        init_timestep = min(init_timestep, num_inference_steps)315 316        t_start = min(max(num_inference_steps - init_timestep + offset, 0), num_inference_steps - 1)317        latest_timestep = self.scheduler.timesteps[t_start]318 319        return latest_timestep320 321    @torch.no_grad()322    def __call__(323        self,324        canvas_height: int,325        canvas_width: int,326        regions: List[DiffusionRegion],327        num_inference_steps: Optional[int] = 50,328        seed: Optional[int] = 12345,329        reroll_regions: Optional[List[RerollRegion]] = None,330        cpu_vae: Optional[bool] = False,331        decode_steps: Optional[bool] = False,332    ):333        if reroll_regions is None:334            reroll_regions = []335        batch_size = 1336 337        if decode_steps:338            steps_images = []339 340        # Prepare scheduler341        self.scheduler.set_timesteps(num_inference_steps, device=self.device)342 343        # Split diffusion regions by their kind344        text2image_regions = [region for region in regions if isinstance(region, Text2ImageRegion)]345        image2image_regions = [region for region in regions if isinstance(region, Image2ImageRegion)]346 347        # Prepare text embeddings348        for region in text2image_regions:349            region.tokenize_prompt(self.tokenizer)350            region.encode_prompt(self.text_encoder, self.device)351 352        # Create original noisy latents using the timesteps353        latents_shape = (batch_size, self.unet.config.in_channels, canvas_height // 8, canvas_width // 8)354        generator = torch.Generator(self.device).manual_seed(seed)355        init_noise = torch.randn(latents_shape, generator=generator, device=self.device)356 357        # Reset latents in seed reroll regions, if requested358        for region in reroll_regions:359            if region.reroll_mode == RerollModes.RESET.value:360                region_shape = (361                    latents_shape[0],362                    latents_shape[1],363                    region.latent_row_end - region.latent_row_init,364                    region.latent_col_end - region.latent_col_init,365                )366                init_noise[367                    :,368                    :,369                    region.latent_row_init : region.latent_row_end,370                    region.latent_col_init : region.latent_col_end,371                ] = torch.randn(region_shape, generator=region.get_region_generator(self.device), device=self.device)372 373        # Apply epsilon noise to regions: first diffusion regions, then reroll regions374        all_eps_rerolls = regions + [r for r in reroll_regions if r.reroll_mode == RerollModes.EPSILON.value]375        for region in all_eps_rerolls:376            if region.noise_eps > 0:377                region_noise = init_noise[378                    :,379                    :,380                    region.latent_row_init : region.latent_row_end,381                    region.latent_col_init : region.latent_col_end,382                ]383                eps_noise = (384                    torch.randn(385                        region_noise.shape, generator=region.get_region_generator(self.device), device=self.device386                    )387                    * region.noise_eps388                )389                init_noise[390                    :,391                    :,392                    region.latent_row_init : region.latent_row_end,393                    region.latent_col_init : region.latent_col_end,394                ] += eps_noise395 396        # scale the initial noise by the standard deviation required by the scheduler397        latents = init_noise * self.scheduler.init_noise_sigma398 399        # Get unconditional embeddings for classifier free guidance in text2image regions400        for region in text2image_regions:401            max_length = region.tokenized_prompt.input_ids.shape[-1]402            uncond_input = self.tokenizer(403                [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"404            )405            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]406 407            # For classifier free guidance, we need to do two forward passes.408            # Here we concatenate the unconditional and text embeddings into a single batch409            # to avoid doing two forward passes410            region.encoded_prompt = torch.cat([uncond_embeddings, region.encoded_prompt])411 412        # Prepare image latents413        for region in image2image_regions:414            region.encode_reference_image(self.vae, device=self.device, generator=generator)415 416        # Prepare mask of weights for each region417        mask_builder = MaskWeightsBuilder(latent_space_dim=self.unet.config.in_channels, nbatch=batch_size)418        mask_weights = [mask_builder.compute_mask_weights(region).to(self.device) for region in text2image_regions]419 420        # Diffusion timesteps421        for i, t in tqdm(enumerate(self.scheduler.timesteps)):422            # Diffuse each region423            noise_preds_regions = []424 425            # text2image regions426            for region in text2image_regions:427                region_latents = latents[428                    :,429                    :,430                    region.latent_row_init : region.latent_row_end,431                    region.latent_col_init : region.latent_col_end,432                ]433                # expand the latents if we are doing classifier free guidance434                latent_model_input = torch.cat([region_latents] * 2)435                # scale model input following scheduler rules436                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)437                # predict the noise residual438                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=region.encoded_prompt)["sample"]439                # perform guidance440                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)441                noise_pred_region = noise_pred_uncond + region.guidance_scale * (noise_pred_text - noise_pred_uncond)442                noise_preds_regions.append(noise_pred_region)443 444            # Merge noise predictions for all tiles445            noise_pred = torch.zeros(latents.shape, device=self.device)446            contributors = torch.zeros(latents.shape, device=self.device)447            # Add each tile contribution to overall latents448            for region, noise_pred_region, mask_weights_region in zip(449                text2image_regions, noise_preds_regions, mask_weights450            ):451                noise_pred[452                    :,453                    :,454                    region.latent_row_init : region.latent_row_end,455                    region.latent_col_init : region.latent_col_end,456                ] += noise_pred_region * mask_weights_region457                contributors[458                    :,459                    :,460                    region.latent_row_init : region.latent_row_end,461                    region.latent_col_init : region.latent_col_end,462                ] += mask_weights_region463            # Average overlapping areas with more than 1 contributor464            noise_pred /= contributors465            noise_pred = torch.nan_to_num(466                noise_pred467            )  # Replace NaNs by zeros: NaN can appear if a position is not covered by any DiffusionRegion468 469            # compute the previous noisy sample x_t -> x_t-1470            latents = self.scheduler.step(noise_pred, t, latents).prev_sample471 472            # Image2Image regions: override latents generated by the scheduler473            for region in image2image_regions:474                influence_step = self.get_latest_timestep_img2img(num_inference_steps, region.strength)475                # Only override in the timesteps before the last influence step of the image (given by its strength)476                if t > influence_step:477                    timestep = t.repeat(batch_size)478                    region_init_noise = init_noise[479                        :,480                        :,481                        region.latent_row_init : region.latent_row_end,482                        region.latent_col_init : region.latent_col_end,483                    ]484                    region_latents = self.scheduler.add_noise(region.reference_latents, region_init_noise, timestep)485                    latents[486                        :,487                        :,488                        region.latent_row_init : region.latent_row_end,489                        region.latent_col_init : region.latent_col_end,490                    ] = region_latents491 492            if decode_steps:493                steps_images.append(self.decode_latents(latents, cpu_vae))494 495        # scale and decode the image latents with vae496        image = self.decode_latents(latents, cpu_vae)497 498        output = {"images": image}499        if decode_steps:500            output = {**output, "steps_images": steps_images}501        return output502