Spaces:
Runtime error
Runtime error
File size: 55,664 Bytes
6fedad3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 |
# Copyright (c) 2024 Jaerin Lee
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from transformers import Blip2Processor, Blip2ForConditionalGeneration
from diffusers import DiffusionPipeline, LCMScheduler, DDIMScheduler, AutoencoderTiny
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
from einops import rearrange
from typing import Tuple, List, Literal, Optional, Union
from tqdm import tqdm
from PIL import Image
from util import gaussian_lowpass, blend, get_panorama_views, shift_to_mask_bbox_center
class StableMultiDiffusionPipeline(nn.Module):
def __init__(
self,
device: torch.device,
dtype: torch.dtype = torch.float16,
sd_version: Literal['1.5', '2.0', '2.1', 'xl'] = '1.5',
hf_key: Optional[str] = None,
lora_key: Optional[str] = None,
load_from_local: bool = False, # Turn on if you have already downloaed LoRA & Hugging Face hub is down.
default_mask_std: float = 1.0, # 8.0
default_mask_strength: float = 1.0,
default_prompt_strength: float = 1.0, # 8.0
default_bootstrap_steps: int = 1,
default_boostrap_mix_steps: float = 1.0,
default_bootstrap_leak_sensitivity: float = 0.2,
default_preprocess_mask_cover_alpha: float = 0.3,
t_index_list: List[int] = [0, 4, 12, 25, 37], # [0, 5, 16, 18, 20, 37], # [0, 12, 25, 37], # Magic number.
mask_type: Literal['discrete', 'semi-continuous', 'continuous'] = 'discrete',
) -> None:
r"""Stabilized MultiDiffusion for fast sampling.
Accelrated region-based text-to-image synthesis with Latent Consistency
Model while preserving mask fidelity and quality.
Args:
device (torch.device): Specify CUDA device.
dtype (torch.dtype): Default precision used in the sampling
process. By default, it is FP16.
sd_version (Literal['1.5', '2.0', '2.1', 'xl']): StableDiffusion
version. Currently, only 1.5 is supported.
hf_key (Optional[str]): Custom StableDiffusion checkpoint for
stylized generation.
lora_key (Optional[str]): Custom LCM LoRA for acceleration.
load_from_local (bool): Turn on if you have already downloaed LoRA
& Hugging Face hub is down.
default_mask_std (float): Preprocess mask with Gaussian blur with
specified standard deviation.
default_mask_strength (float): Preprocess mask by multiplying it
globally with the specified variable. Caution: extremely
sensitive. Recommended range: 0.98-1.
default_prompt_strength (float): Preprocess foreground prompts
globally by linearly interpolating its embedding with the
background prompt embeddint with specified mix ratio. Useful
control handle for foreground blending. Recommended range:
0.5-1.
default_bootstrap_steps (int): Bootstrapping stage steps to
encourage region separation. Recommended range: 1-3.
default_boostrap_mix_steps (float): Bootstrapping background is a
linear interpolation between background latent and the white
image latent. This handle controls the mix ratio. Available
range: 0-(number of bootstrapping inference steps). For
example, 2.3 means that for the first two steps, white image
is used as a bootstrapping background and in the third step,
mixture of white (0.3) and registered background (0.7) is used
as a bootstrapping background.
default_bootstrap_leak_sensitivity (float): Postprocessing at each
inference step by masking away the remaining bootstrap
backgrounds t Recommended range: 0-1.
default_preprocess_mask_cover_alpha (float): Optional preprocessing
where each mask covered by other masks is reduced in its alpha
value by this specified factor.
t_index_list (List[int]): The default scheduling for LCM scheduler.
mask_type (Literal['discrete', 'semi-continuous', 'continuous']):
defines the mask quantization modes. Details in the codes of
`self.process_mask`. Basically, this (subtly) controls the
smoothness of foreground-background blending. More continuous
means more blending, but smaller generated patch depending on
the mask standard deviation.
"""
super().__init__()
self.device = device
self.dtype = dtype
self.sd_version = sd_version
self.default_mask_std = default_mask_std
self.default_mask_strength = default_mask_strength
self.default_prompt_strength = default_prompt_strength
self.default_t_list = t_index_list
self.default_bootstrap_steps = default_bootstrap_steps
self.default_boostrap_mix_steps = default_boostrap_mix_steps
self.default_bootstrap_leak_sensitivity = default_bootstrap_leak_sensitivity
self.default_preprocess_mask_cover_alpha = default_preprocess_mask_cover_alpha
self.mask_type = mask_type
print(f'[INFO] Loading Stable Diffusion...')
variant = None
lora_weight_name = None
if self.sd_version == '1.5':
if hf_key is not None:
print(f'[INFO] Using Hugging Face custom model key: {hf_key}')
model_key = hf_key
else:
model_key = 'runwayml/stable-diffusion-v1-5'
variant = 'fp16'
lora_key = 'latent-consistency/lcm-lora-sdv1-5'
lora_weight_name = 'pytorch_lora_weights.safetensors'
# elif self.sd_version == 'xl':
# model_key = 'stabilityai/stable-diffusion-xl-base-1.0'
# lora_key = 'latent-consistency/lcm-lora-sdxl'
# variant = 'fp16'
# lora_weight_name = 'pytorch_lora_weights.safetensors'
else:
raise ValueError(f'Stable Diffusion version {self.sd_version} not supported.')
# Create model
self.i2t_processor = Blip2Processor.from_pretrained('Salesforce/blip2-opt-2.7b')
self.i2t_model = Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-opt-2.7b')
self.pipe = DiffusionPipeline.from_pretrained(model_key, variant=variant, torch_dtype=dtype).to(self.device)
if lora_key is None:
print(f'[INFO] LCM LoRA is not available for SD version {sd_version}. Using DDIM Scheduler instead...')
self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
self.scheduler = self.pipe.scheduler
self.default_num_inference_steps = 50
self.default_guidance_scale = 7.5
else:
self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
self.scheduler = self.pipe.scheduler
self.pipe.load_lora_weights(lora_key, weight_name=lora_weight_name, adapter_name='lcm')
self.default_num_inference_steps = 4
self.default_guidance_scale = 1.0
self.prepare_lcm_schedule(t_index_list, 50)
self.vae = self.pipe.vae
self.tokenizer = self.pipe.tokenizer
self.text_encoder = self.pipe.text_encoder
self.unet = self.pipe.unet
self.vae_scale_factor = self.pipe.vae_scale_factor
# Prepare white background for bootstrapping.
self.get_white_background(768, 768)
print(f'[INFO] Model is loaded!')
def prepare_lcm_schedule(
self,
t_index_list: Optional[List[int]] = None,
num_inference_steps: Optional[int] = None,
) -> None:
r"""Set up different inference schedule for the diffusion model.
You do not have to run this explicitly if you want to use the default
setting, but if you want other time schedules, run this function
between the module initialization and the main call.
Note:
- Recommended t_index_lists for LCMs:
- [0, 12, 25, 37]: Default schedule for 4 steps. Best for
panorama. Not recommended if you want to use bootstrapping.
Because bootstrapping stage affects the initial structuring
of the generated image & in this four step LCM, this is done
with only at the first step, the structure may be distorted.
- [0, 4, 12, 25, 37]: Recommended if you would use 1-step boot-
strapping. Default initialization in this implementation.
- [0, 5, 16, 18, 20, 37]: Recommended if you would use 2-step
bootstrapping.
- Due to the characteristic of SD1.5 LCM LoRA, setting
`num_inference_steps` larger than 20 may results in overly blurry
and unrealistic images. Beware!
Args:
t_index_list (Optional[List[int]]): The specified scheduling step
regarding the maximum timestep as `num_inference_steps`, which
is by default, 50. That means that
`t_index_list=[0, 12, 25, 37]` is a relative time indices basd
on the full scale of 50. If None, reinitialize the module with
the default value.
num_inference_steps (Optional[int]): The maximum timestep of the
sampler. Defines relative scale of the `t_index_list`. Rarely
used in practice. If None, reinitialize the module with the
default value.
"""
if t_index_list is None:
t_index_list = self.default_t_list
if num_inference_steps is None:
num_inference_steps = self.default_num_inference_steps
self.scheduler.set_timesteps(num_inference_steps)
self.timesteps = torch.as_tensor([
self.scheduler.timesteps[t] for t in t_index_list
], dtype=torch.long)
shape = (len(t_index_list), 1, 1, 1)
c_skip_list = []
c_out_list = []
for timestep in self.timesteps:
c_skip, c_out = self.scheduler.get_scalings_for_boundary_condition_discrete(timestep)
c_skip_list.append(c_skip)
c_out_list.append(c_out)
self.c_skip = torch.stack(c_skip_list).view(*shape).to(dtype=self.dtype, device=self.device)
self.c_out = torch.stack(c_out_list).view(*shape).to(dtype=self.dtype, device=self.device)
alpha_prod_t_sqrt_list = []
beta_prod_t_sqrt_list = []
for timestep in self.timesteps:
alpha_prod_t_sqrt = self.scheduler.alphas_cumprod[timestep].sqrt()
beta_prod_t_sqrt = (1 - self.scheduler.alphas_cumprod[timestep]).sqrt()
alpha_prod_t_sqrt_list.append(alpha_prod_t_sqrt)
beta_prod_t_sqrt_list.append(beta_prod_t_sqrt)
alpha_prod_t_sqrt = (torch.stack(alpha_prod_t_sqrt_list).view(*shape)
.to(dtype=self.dtype, device=self.device))
beta_prod_t_sqrt = (torch.stack(beta_prod_t_sqrt_list).view(*shape)
.to(dtype=self.dtype, device=self.device))
self.alpha_prod_t_sqrt = alpha_prod_t_sqrt
self.beta_prod_t_sqrt = beta_prod_t_sqrt
noise_lvs = (1 - self.scheduler.alphas_cumprod[self.timesteps].to(self.device)) ** 0.5
self.noise_lvs = noise_lvs[None, :, None, None, None]
self.next_noise_lvs = torch.cat([noise_lvs[1:], noise_lvs.new_zeros(1)])[None, :, None, None, None]
@torch.no_grad()
def get_text_embeds(self, prompt: str, negative_prompt: str) -> Tuple[torch.Tensor]:
r"""Text embeddings from string text prompts.
Args:
prompt (str): A text prompt string.
negative_prompt: An optional negative text prompt string. Good for
high-quality generation.
Returns:
A tuple of (negative, positive) prompt embeddings of (1, 77, 768).
"""
kwargs = dict(padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt')
# Tokenize text and get embeddings.
text_input = self.tokenizer(prompt, truncation=True, **kwargs)
text_embeds = self.text_encoder(text_input.input_ids.to(self.device))[0]
uncond_input = self.tokenizer(negative_prompt, **kwargs)
uncond_embeds = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
return uncond_embeds, text_embeds
@torch.no_grad()
def get_text_prompts(self, image: Image.Image) -> str:
r"""A convenient method to extract text prompt from an image.
This is called if the user does not provide background prompt but only
the background image. We use BLIP-2 to automatically generate prompts.
Args:
image (Image.Image): A PIL image.
Returns:
A single string of text prompt.
"""
question = 'Question: What are in the image? Answer:'
inputs = self.i2t_processor(image, question, return_tensors='pt')
out = self.i2t_model.generate(**inputs, max_new_tokens=77)
prompt = self.i2t_processor.decode(out[0], skip_special_tokens=True).strip()
return prompt
@torch.no_grad()
def encode_imgs(
self,
imgs: torch.Tensor,
generator: Optional[torch.Generator] = None,
vae: Optional[nn.Module] = None,
) -> torch.Tensor:
r"""A wrapper function for VAE encoder of the latent diffusion model.
Args:
imgs (torch.Tensor): An image to get StableDiffusion latents.
Expected shape: (B, 3, H, W). Expected pixel scale: [0, 1].
generator (Optional[torch.Generator]): Seed for KL-Autoencoder.
vae (Optional[nn.Module]): Explicitly specify VAE (used for
the demo application with TinyVAE).
Returns:
An image latent embedding with 1/8 size (depending on the auto-
encoder. Shape: (B, 4, H//8, W//8).
"""
def _retrieve_latents(
encoder_output: torch.Tensor,
generator: Optional[torch.Generator] = None,
sample_mode: str = 'sample',
):
if hasattr(encoder_output, 'latent_dist') and sample_mode == 'sample':
return encoder_output.latent_dist.sample(generator)
elif hasattr(encoder_output, 'latent_dist') and sample_mode == 'argmax':
return encoder_output.latent_dist.mode()
elif hasattr(encoder_output, 'latents'):
return encoder_output.latents
else:
raise AttributeError('Could not access latents of provided encoder_output')
vae = self.vae if vae is None else vae
imgs = 2 * imgs - 1
latents = vae.config.scaling_factor * _retrieve_latents(vae.encode(imgs), generator=generator)
return latents
@torch.no_grad()
def decode_latents(self, latents: torch.Tensor, vae: Optional[nn.Module] = None) -> torch.Tensor:
r"""A wrapper function for VAE decoder of the latent diffusion model.
Args:
latents (torch.Tensor): An image latent to get associated images.
Expected shape: (B, 4, H//8, W//8).
vae (Optional[nn.Module]): Explicitly specify VAE (used for
the demo application with TinyVAE).
Returns:
An image latent embedding with 1/8 size (depending on the auto-
encoder. Shape: (B, 3, H, W).
"""
vae = self.vae if vae is None else vae
latents = 1 / vae.config.scaling_factor * latents
imgs = vae.decode(latents).sample
imgs = (imgs / 2 + 0.5).clip_(0, 1)
return imgs
@torch.no_grad()
def get_white_background(self, height: int, width: int) -> torch.Tensor:
r"""White background image latent for bootstrapping or in case of
absent background.
Additionally stores the maximally-sized white latent for fast retrieval
in the future. By default, we initially call this with 768x768 sized
white image, so the function is rarely visited twice.
Args:
height (int): The height of the white *image*, not its latent.
width (int): The width of the white *image*, not its latent.
Returns:
A white image latent of size (1, 4, height//8, width//8). A cropped
version of the stored white latent is returned if the requested
size is smaller than what we already have created.
"""
if not hasattr(self, 'white') or self.white.shape[-2] < height or self.white.shape[-1] < width:
white = torch.ones(1, 3, height, width, dtype=self.dtype, device=self.device)
self.white = self.encode_imgs(white)
return self.white
return self.white[..., :(height // self.vae_scale_factor), :(width // self.vae_scale_factor)]
@torch.no_grad()
def process_mask(
self,
masks: Union[torch.Tensor, Image.Image, List[Image.Image]],
strength: Optional[Union[torch.Tensor, float]] = None,
std: Optional[Union[torch.Tensor, float]] = None,
height: int = 512,
width: int = 512,
use_boolean_mask: bool = True,
timesteps: Optional[torch.Tensor] = None,
preprocess_mask_cover_alpha: Optional[float] = None,
) -> Tuple[torch.Tensor]:
r"""Fast preprocess of masks for region-based generation with fine-
grained controls.
Mask preprocessing is done in four steps:
1. Resizing: Resize the masks into the specified width and height by
nearest neighbor interpolation.
2. (Optional) Ordering: Masks with higher indices are considered to
cover the masks with smaller indices. Covered masks are decayed
in its alpha value by the specified factor of
`preprocess_mask_cover_alpha`.
3. Blurring: Gaussian blur is applied to the mask with the specified
standard deviation (isotropic). This results in gradual increase of
masked region as the timesteps evolve, naturally blending fore-
ground and the predesignated background. Not strictly required if
you want to produce images from scratch withoout background.
4. Quantization: Split the real-numbered masks of value between [0, 1]
into predefined noise levels for each quantized scheduling step of
the diffusion sampler. For example, if the diffusion model sampler
has noise level of [0.9977, 0.9912, 0.9735, 0.8499, 0.5840], which
is the default noise level of this module with schedule [0, 4, 12,
25, 37], the masks are split into binary masks whose values are
greater than these levels. This results in tradual increase of mask
region as the timesteps increase. Details are described in our
paper at https://arxiv.org/pdf/2403.09055.pdf.
On the Three Modes of `mask_type`:
`self.mask_type` is predefined at the initialization stage of this
pipeline. Three possible modes are available: 'discrete', 'semi-
continuous', and 'continuous'. These define the mask quantization
modes we use. Basically, this (subtly) controls the smoothness of
foreground-background blending. Continuous modes produces nonbinary
masks to further blend foreground and background latents by linear-
ly interpolating between them. Semi-continuous masks only applies
continuous mask at the last step of the LCM sampler. Due to the
large step size of the LCM scheduler, we find that our continuous
blending helps generating seamless inpainting and editing results.
Args:
masks (Union[torch.Tensor, Image.Image, List[Image.Image]]): Masks.
strength (Optional[Union[torch.Tensor, float]]): Mask strength that
overrides the default value. A globally multiplied factor to
the mask at the initial stage of processing. Can be applied
seperately for each mask.
std (Optional[Union[torch.Tensor, float]]): Mask blurring Gaussian
kernel's standard deviation. Overrides the default value. Can
be applied seperately for each mask.
height (int): The height of the expected generation. Mask is
resized to (height//8, width//8) with nearest neighbor inter-
polation.
width (int): The width of the expected generation. Mask is resized
to (height//8, width//8) with nearest neighbor interpolation.
use_boolean_mask (bool): Specify this to treat the mask image as
a boolean tensor. The retion with dark part darker than 0.5 of
the maximal pixel value (that is, 127.5) is considered as the
designated mask.
timesteps (Optional[torch.Tensor]): Defines the scheduler noise
levels that acts as bins of mask quantization.
preprocess_mask_cover_alpha (Optional[float]): Optional pre-
processing where each mask covered by other masks is reduced in
its alpha value by this specified factor. Overrides the default
value.
Returns: A tuple of tensors.
- masks: Preprocessed (ordered, blurred, and quantized) binary/non-
binary masks (see the explanation on `mask_type` above) for
region-based image synthesis.
- masks_blurred: Gaussian blurred masks. Used for optionally
specified foreground-background blending after image
generation.
- std: Mask blur standard deviation. Used for optionally specified
foreground-background blending after image generation.
"""
if isinstance(masks, Image.Image):
masks = [masks]
if isinstance(masks, (tuple, list)):
# Assumes white background for Image.Image;
# inverted boolean masks with shape (1, 1, H, W) for torch.Tensor.
if use_boolean_mask:
proc = lambda m: T.ToTensor()(m)[None, -1:] < 0.5
else:
proc = lambda m: 1.0 - T.ToTensor()(m)[None, -1:]
masks = torch.cat([proc(mask) for mask in masks], dim=0).float().clip_(0, 1)
masks = F.interpolate(masks.float(), size=(height, width), mode='bilinear', align_corners=False)
masks = masks.to(self.device)
# Background mask alpha is decayed by the specified factor where foreground masks covers it.
if preprocess_mask_cover_alpha is None:
preprocess_mask_cover_alpha = self.default_preprocess_mask_cover_alpha
if preprocess_mask_cover_alpha > 0:
masks = torch.stack([
torch.where(
masks[i + 1:].sum(dim=0) > 0,
mask * preprocess_mask_cover_alpha,
mask,
) if i < len(masks) - 1 else mask
for i, mask in enumerate(masks)
], dim=0)
# Scheduler noise levels for mask quantization.
if timesteps is None:
noise_lvs = self.noise_lvs
next_noise_lvs = self.next_noise_lvs
else:
noise_lvs_ = (1 - self.scheduler.alphas_cumprod[timesteps].to(self.device)) ** 0.5
noise_lvs = noise_lvs_[None, :, None, None, None]
next_noise_lvs = torch.cat([noise_lvs_[1:], noise_lvs_.new_zeros(1)])[None, :, None, None, None]
# Mask preprocessing parameters are fetched from the default settings.
if std is None:
std = self.default_mask_std
if isinstance(std, (int, float)):
std = [std] * len(masks)
if isinstance(std, (list, tuple)):
std = torch.as_tensor(std, dtype=torch.float, device=self.device)
if strength is None:
strength = self.default_mask_strength
if isinstance(strength, (int, float)):
strength = [strength] * len(masks)
if isinstance(strength, (list, tuple)):
strength = torch.as_tensor(strength, dtype=torch.float, device=self.device)
if (std > 0).any():
std = torch.where(std > 0, std, 1e-5)
masks = gaussian_lowpass(masks, std)
masks_blurred = masks
# NOTE: This `strength` aligns with `denoising strength`. However, with LCM, using strength < 0.96
# gives unpleasant results.
masks = masks * strength[:, None, None, None]
masks = masks.unsqueeze(1).repeat(1, noise_lvs.shape[1], 1, 1, 1)
# Mask is quantized according to the current noise levels specified by the scheduler.
if self.mask_type == 'discrete':
# Discrete mode.
masks = masks > noise_lvs
elif self.mask_type == 'semi-continuous':
# Semi-continuous mode (continuous at the last step only).
masks = torch.cat((
masks[:, :-1] > noise_lvs[:, :-1],
(
(masks[:, -1:] - next_noise_lvs[:, -1:]) / (noise_lvs[:, -1:] - next_noise_lvs[:, -1:])
).clip_(0, 1),
), dim=1)
elif self.mask_type == 'continuous':
# Continuous mode: Have the exact same `1` coverage with discrete mode, but the mask gradually
# decreases continuously after the discrete mode boundary to become `0` at the
# next lower threshold.
masks = ((masks - next_noise_lvs) / (noise_lvs - next_noise_lvs)).clip_(0, 1)
# NOTE: Post processing mask strength does not align with conventional 'denoising_strength'. However,
# fine-grained mask alpha channel tuning is available with this form.
# masks = masks * strength[None, :, None, None, None]
h = height // self.vae_scale_factor
w = width // self.vae_scale_factor
masks = rearrange(masks.float(), 'p t () h w -> (p t) () h w')
masks = F.interpolate(masks, size=(h, w), mode='nearest')
masks = rearrange(masks.to(self.dtype), '(p t) () h w -> p t () h w', p=len(std))
return masks, masks_blurred, std
def scheduler_step(
self,
noise_pred: torch.Tensor,
idx: int,
latent: torch.Tensor,
) -> torch.Tensor:
r"""Denoise-only step for reverse diffusion scheduler.
Designed to match the interface of the original `pipe.scheduler.step`,
which is a combination of this method and the following
`scheduler_add_noise`.
Args:
noise_pred (torch.Tensor): Noise prediction results from the U-Net.
idx (int): Instead of timesteps (in [0, 1000]-scale) use indices
for the timesteps tensor (ranged in [0, len(timesteps)-1]).
latent (torch.Tensor): Noisy latent.
Returns:
A denoised tensor with the same size as latent.
"""
F_theta = (latent - self.beta_prod_t_sqrt[idx] * noise_pred) / self.alpha_prod_t_sqrt[idx]
return self.c_out[idx] * F_theta + self.c_skip[idx] * latent
def scheduler_add_noise(
self,
latent: torch.Tensor,
noise: Optional[torch.Tensor],
idx: int,
) -> torch.Tensor:
r"""Separated noise-add step for the reverse diffusion scheduler.
Designed to match the interface of the original
`pipe.scheduler.add_noise`.
Args:
latent (torch.Tensor): Denoised latent.
noise (torch.Tensor): Added noise. Can be None. If None, a random
noise is newly sampled for addition.
idx (int): Instead of timesteps (in [0, 1000]-scale) use indices
for the timesteps tensor (ranged in [0, len(timesteps)-1]).
Returns:
A noisy tensor with the same size as latent.
"""
if idx >= len(self.alpha_prod_t_sqrt) or idx < 0:
# The last step does not require noise addition.
return latent
noise = torch.randn_like(latent) if noise is None else noise
return self.alpha_prod_t_sqrt[idx] * latent + self.beta_prod_t_sqrt[idx] * noise
@torch.no_grad()
def sample(
self,
prompts: Union[str, List[str]],
negative_prompts: Union[str, List[str]] = '',
height: int = 512,
width: int = 512,
num_inference_steps: Optional[int] = None,
guidance_scale: Optional[float] = None,
batch_size: int = 1,
) -> Image.Image:
r"""StableDiffusionPipeline for single-prompt single-tile generation.
Minimal Example:
>>> device = torch.device('cuda:0')
>>> smd = StableMultiDiffusionPipeline(device)
>>> image = smd.sample('A photo of the dolomites')
>>> image.save('my_creation.png')
Args:
prompts (Union[str, List[str]]): A text prompt.
negative_prompts (Union[str, List[str]]): A negative text prompt.
height (int): Height of a generated image.
width (int): Width of a generated image.
num_inference_steps (Optional[int]): Number of inference steps.
Default inference scheduling is used if none is specified.
guidance_scale (Optional[float]): Classifier guidance scale.
Default value is used if none is specified.
batch_size (int): Number of images to generate.
Returns: A PIL.Image image.
"""
if num_inference_steps is None:
num_inference_steps = self.default_num_inference_steps
if guidance_scale is None:
guidance_scale = self.default_guidance_scale
self.scheduler.set_timesteps(num_inference_steps)
if isinstance(prompts, str):
prompts = [prompts]
if isinstance(negative_prompts, str):
negative_prompts = [negative_prompts]
# Calculate text embeddings.
uncond_embeds, text_embeds = self.get_text_embeds(prompts, negative_prompts) # [2, 77, 768]
text_embeds = torch.cat([uncond_embeds.mean(dim=0, keepdim=True), text_embeds.mean(dim=0, keepdim=True)])
h = height // self.vae_scale_factor
w = width // self.vae_scale_factor
latent = torch.randn((batch_size, self.unet.config.in_channels, h, w), dtype=self.dtype, device=self.device)
with torch.autocast('cuda'):
for i, t in enumerate(tqdm(self.scheduler.timesteps)):
# Expand the latents if we are doing classifier-free guidance.
latent_model_input = torch.cat([latent] * 2)
# Perform one step of the reverse diffusion.
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeds)['sample']
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
latent = self.scheduler.step(noise_pred, t, latent)['prev_sample']
# Return PIL Image.
latent = latent.to(dtype=self.dtype)
imgs = [T.ToPILImage()(self.decode_latents(l[None])[0]) for l in latent]
return imgs
@torch.no_grad()
def sample_panorama(
self,
prompts: Union[str, List[str]],
negative_prompts: Union[str, List[str]] = '',
height: int = 512,
width: int = 2048,
num_inference_steps: Optional[int] = None,
guidance_scale: Optional[float] = None,
tile_size: Optional[int] = None,
) -> Image.Image:
r"""Large size image generation from a single set of prompts.
Minimal Example:
>>> device = torch.device('cuda:0')
>>> smd = StableMultiDiffusionPipeline(device)
>>> image = smd.sample_panorama(
>>> 'A photo of Alps', height=512, width=3072)
>>> image.save('my_panorama_creation.png')
Args:
prompts (Union[str, List[str]]): A text prompt.
negative_prompts (Union[str, List[str]]): A negative text prompt.
height (int): Height of a generated image. It is tiled if larger
than `tile_size`.
width (int): Width of a generated image. It is tiled if larger
than `tile_size`.
num_inference_steps (Optional[int]): Number of inference steps.
Default inference scheduling is used if none is specified.
guidance_scale (Optional[float]): Classifier guidance scale.
Default value is used if none is specified.
tile_size (Optional[int]): Tile size of the panorama generation.
Works best with the default training size of the Stable-
Diffusion model, i.e., 512 or 768 for SD1.5 and 1024 for SDXL.
Returns: A PIL.Image image of a panorama (large-size) image.
"""
if num_inference_steps is None:
num_inference_steps = self.default_num_inference_steps
self.scheduler.set_timesteps(num_inference_steps)
timesteps = self.timesteps
use_custom_timesteps = False
else:
self.scheduler.set_timesteps(num_inference_steps)
timesteps = self.scheduler.timesteps
use_custom_timesteps = True
if guidance_scale is None:
guidance_scale = self.default_guidance_scale
if isinstance(prompts, str):
prompts = [prompts]
if isinstance(negative_prompts, str):
negative_prompts = [negative_prompts]
# Calculate text embeddings.
uncond_embeds, text_embeds = self.get_text_embeds(prompts, negative_prompts) # [2, 77, 768]
text_embeds = torch.cat([uncond_embeds.mean(dim=0, keepdim=True), text_embeds.mean(dim=0, keepdim=True)])
# Define panorama grid and get views
h = height // self.vae_scale_factor
w = width // self.vae_scale_factor
latent = torch.randn((1, self.unet.config.in_channels, h, w), dtype=self.dtype, device=self.device)
if tile_size is None:
tile_size = min(min(height, width), 512)
views, masks = get_panorama_views(h, w, tile_size // self.vae_scale_factor)
masks = masks.to(dtype=self.dtype, device=self.device)
value = torch.zeros_like(latent)
with torch.autocast('cuda'):
for i, t in enumerate(tqdm(timesteps)):
value.zero_()
for j, (h_start, h_end, w_start, w_end) in enumerate(views):
# TODO we can support batches, and pass multiple views at once to the unet
latent_view = latent[:, :, h_start:h_end, w_start:w_end]
# Expand the latents if we are doing classifier-free guidance.
latent_model_input = torch.cat([latent_view] * 2)
# Perform one step of the reverse diffusion.
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeds)['sample']
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
# Compute the denoising step.
latents_view_denoised = self.scheduler_step(noise_pred, i, latent_view) # (1, 4, h, w)
mask = masks[..., j:j + 1, h_start:h_end, w_start:w_end] # (1, 1, h, w)
value[..., h_start:h_end, w_start:w_end] += mask * latents_view_denoised # (1, 1, h, w)
# Update denoised latent.
latent = value.clone()
if i < len(timesteps) - 1:
latent = self.scheduler_add_noise(latent, None, i + 1)
# Return PIL Image.
imgs = self.decode_latents(latent)
img = T.ToPILImage()(imgs[0].cpu())
return img
@torch.no_grad()
def __call__(
self,
prompts: Optional[Union[str, List[str]]] = None,
negative_prompts: Union[str, List[str]] = '',
suffix: Optional[str] = None, #', background is ',
background: Optional[Union[torch.Tensor, Image.Image]] = None,
background_prompt: Optional[str] = None,
background_negative_prompt: str = '',
height: int = 512,
width: int = 512,
num_inference_steps: Optional[int] = None,
guidance_scale: Optional[float] = None,
prompt_strengths: Optional[Union[torch.Tensor, float, List[float]]] = None,
masks: Optional[Union[Image.Image, List[Image.Image]]] = None,
mask_strengths: Optional[Union[torch.Tensor, float, List[float]]] = None,
mask_stds: Optional[Union[torch.Tensor, float, List[float]]] = None,
use_boolean_mask: bool = True,
do_blend: bool = True,
tile_size: int = 768,
bootstrap_steps: Optional[int] = None,
boostrap_mix_steps: Optional[float] = None,
bootstrap_leak_sensitivity: Optional[float] = None,
preprocess_mask_cover_alpha: Optional[float] = None,
) -> Image.Image:
r"""Arbitrary-size image generation from multiple pairs of (regional)
text prompt-mask pairs.
This is a main routine for this pipeline.
Example:
>>> device = torch.device('cuda:0')
>>> smd = StableMultiDiffusionPipeline(device)
>>> prompts = {... specify prompts}
>>> masks = {... specify mask tensors}
>>> height, width = masks.shape[-2:]
>>> image = smd(
>>> prompts, masks=masks.float(), height=height, width=width)
>>> image.save('my_beautiful_creation.png')
Args:
prompts (Union[str, List[str]]): A text prompt.
negative_prompts (Union[str, List[str]]): A negative text prompt.
suffix (Optional[str]): One option for blending foreground prompts
with background prompts by simply appending background prompt
to the end of each foreground prompt with this `middle word` in
between. For example, if you set this as `, background is`,
then the foreground prompt will be changed into
`(fg), background is (bg)` before conditional generation.
background (Optional[Union[torch.Tensor, Image.Image]]): a
background image, if the user wants to draw in front of the
specified image. Background prompt will automatically generated
with a BLIP-2 model.
background_prompt (Optional[str]): The background prompt is used
for preprocessing foreground prompt embeddings to blend
foreground and background.
background_negative_prompt (Optional[str]): The negative background
prompt.
height (int): Height of a generated image. It is tiled if larger
than `tile_size`.
width (int): Width of a generated image. It is tiled if larger
than `tile_size`.
num_inference_steps (Optional[int]): Number of inference steps.
Default inference scheduling is used if none is specified.
guidance_scale (Optional[float]): Classifier guidance scale.
Default value is used if none is specified.
prompt_strength (float): Overrides default value. Preprocess
foreground prompts globally by linearly interpolating its
embedding with the background prompt embeddint with specified
mix ratio. Useful control handle for foreground blending.
Recommended range: 0.5-1.
masks (Optional[Union[Image.Image, List[Image.Image]]]): a list of
mask images. Each mask associates with each of the text prompts
and each of the negative prompts. If specified as an image, it
regards the image as a boolean mask. Also accepts torch.Tensor
masks, which can have nonbinary values for fine-grained
controls in mixing regional generations.
mask_strengths (Optional[Union[torch.Tensor, float, List[float]]]):
Overrides the default value. an be assigned for each mask
separately. Preprocess mask by multiplying it globally with the
specified variable. Caution: extremely sensitive. Recommended
range: 0.98-1.
mask_stds (Optional[Union[torch.Tensor, float, List[float]]]):
Overrides the default value. Can be assigned for each mask
separately. Preprocess mask with Gaussian blur with specified
standard deviation. Recommended range: 0-64.
use_boolean_mask (bool): Turn this off if you want to treat the
mask image as nonbinary one. The module will use the last
channel of the given image in `masks` as the mask value.
do_blend (bool): Blend the generated foreground and the optionally
predefined background by smooth boundary obtained from Gaussian
blurs of the foreground `masks` with the given `mask_stds`.
tile_size (Optional[int]): Tile size of the panorama generation.
Works best with the default training size of the Stable-
Diffusion model, i.e., 512 or 768 for SD1.5 and 1024 for SDXL.
bootstrap_steps (int): Overrides the default value. Bootstrapping
stage steps to encourage region separation. Recommended range:
1-3.
boostrap_mix_steps (float): Overrides the default value.
Bootstrapping background is a linear interpolation between
background latent and the white image latent. This handle
controls the mix ratio. Available range: 0-(number of
bootstrapping inference steps). For example, 2.3 means that for
the first two steps, white image is used as a bootstrapping
background and in the third step, mixture of white (0.3) and
registered background (0.7) is used as a bootstrapping
background.
bootstrap_leak_sensitivity (float): Overrides the default value.
Postprocessing at each inference step by masking away the
remaining bootstrap backgrounds t Recommended range: 0-1.
preprocess_mask_cover_alpha (float): Overrides the default value.
Optional preprocessing where each mask covered by other masks
is reduced in its alpha value by this specified factor.
Returns: A PIL.Image image of a panorama (large-size) image.
"""
### Simplest cases
# prompts is None: return background.
# masks is None but prompts is not None: return prompts
# masks is not None and prompts is not None: Do StableMultiDiffusion.
if prompts is None or (isinstance(prompts, (list, tuple, str)) and len(prompts) == 0):
if background is None and background_prompt is not None:
return sample(background_prompt, background_negative_prompt, height, width, num_inference_steps, guidance_scale)
return background
elif masks is None or (isinstance(masks, (list, tuple)) and len(masks) == 0):
return sample(prompts, negative_prompts, height, width, num_inference_steps, guidance_scale)
### Prepare generation
if num_inference_steps is not None:
self.prepare_lcm_schedule(list(range(num_inference_steps)), num_inference_steps)
if guidance_scale is None:
guidance_scale = self.default_guidance_scale
### Prompts & Masks
# asserts #m > 0 and #p > 0.
# #m == #p == #n > 0: We happily generate according to the prompts & masks.
# #m != #p: #p should be 1 and we will broadcast text embeds of p through m masks.
# #p != #n: #n should be 1 and we will broadcast negative embeds n through p prompts.
if isinstance(masks, Image.Image):
masks = [masks]
if isinstance(prompts, str):
prompts = [prompts]
if isinstance(negative_prompts, str):
negative_prompts = [negative_prompts]
num_masks = len(masks)
num_prompts = len(prompts)
num_nprompts = len(negative_prompts)
assert num_prompts in (num_masks, 1), \
f'The number of prompts {num_prompts} should match the number of masks {num_masks}!'
assert num_nprompts in (num_prompts, 1), \
f'The number of negative prompts {num_nprompts} should match the number of prompts {num_prompts}!'
fg_masks, masks_g, std = self.process_mask(
masks,
mask_strengths,
mask_stds,
height=height,
width=width,
use_boolean_mask=use_boolean_mask,
timesteps=self.timesteps,
preprocess_mask_cover_alpha=preprocess_mask_cover_alpha,
) # (p, t, 1, H, W)
bg_masks = (1 - fg_masks.sum(dim=0)).clip_(0, 1) # (T, 1, h, w)
has_background = bg_masks.sum() > 0
h = (height + self.vae_scale_factor - 1) // self.vae_scale_factor
w = (width + self.vae_scale_factor - 1) // self.vae_scale_factor
### Background
# background == None && background_prompt == None: Initialize with white background.
# background == None && background_prompt != None: Generate background *along with other prompts*.
# background != None && background_prompt == None: Retrieve text prompt using BLIP.
# background != None && background_prompt != None: Use the given arguments.
# not has_background: no effect of prompt_strength (the mix ratio between fg prompt & bg prompt)
# has_background && prompt_strength != 1: mix only for this case.
bg_latent = None
if has_background:
if background is None and background_prompt is not None:
fg_masks = torch.cat((bg_masks[None], fg_masks), dim=0)
if suffix is not None:
prompts = [p + suffix + background_prompt for p in prompts]
prompts = [background_prompt] + prompts
negative_prompts = [background_negative_prompt] + negative_prompts
has_background = False # Regard that background does not exist.
else:
if background is None and background_prompt is None:
background = torch.ones(1, 3, height, width, dtype=self.dtype, device=self.device)
background_prompt = 'simple white background image'
elif background is not None and background_prompt is None:
background_prompt = self.get_text_prompts(background)
if suffix is not None:
prompts = [p + suffix + background_prompt for p in prompts]
prompts = [background_prompt] + prompts
negative_prompts = [background_negative_prompt] + negative_prompts
if isinstance(background, Image.Image):
background = T.ToTensor()(background).to(dtype=self.dtype, device=self.device)[None]
background = F.interpolate(background, size=(height, width), mode='bicubic', align_corners=False)
bg_latent = self.encode_imgs(background)
# Bootstrapping stage preparation.
if bootstrap_steps is None:
bootstrap_steps = self.default_bootstrap_steps
if boostrap_mix_steps is None:
boostrap_mix_steps = self.default_boostrap_mix_steps
if bootstrap_leak_sensitivity is None:
bootstrap_leak_sensitivity = self.default_bootstrap_leak_sensitivity
if bootstrap_steps > 0:
height_ = min(height, tile_size)
width_ = min(width, tile_size)
white = self.get_white_background(height, width) # (1, 4, h, w)
### Prepare text embeddings (optimized for the minimal encoder batch size)
uncond_embeds, text_embeds = self.get_text_embeds(prompts, negative_prompts) # [2 * len(prompts), 77, 768]
if has_background:
# First channel is background prompt text embeds. Background prompt itself is not used for generation.
s = prompt_strengths
if prompt_strengths is None:
s = self.default_prompt_strength
if isinstance(s, (int, float)):
s = [s] * num_prompts
if isinstance(s, (list, tuple)):
assert len(s) == num_prompts, \
f'The number of prompt strengths {len(s)} should match the number of prompts {num_prompts}!'
s = torch.as_tensor(s, dtype=self.dtype, device=self.device)
s = s[:, None, None]
be = text_embeds[:1]
bu = uncond_embeds[:1]
fe = text_embeds[1:]
fu = uncond_embeds[1:]
if num_prompts > num_nprompts:
# # negative prompts = 1; # prompts > 1.
assert fu.shape[0] == 1 and fe.shape == num_prompts
fu = fu.repeat(num_prompts, 1, 1)
text_embeds = torch.lerp(be, fe, s) # (p, 77, 768)
uncond_embeds = torch.lerp(bu, fu, s) # (n, 77, 768)
elif num_prompts > num_nprompts:
# # negative prompts = 1; # prompts > 1.
assert uncond_embeds.shape[0] == 1 and text_embeds.shape[0] == num_prompts
uncond_embeds = uncond_embeds.repeat(num_prompts, 1, 1)
assert uncond_embeds.shape[0] == text_embeds.shape[0] == num_prompts
if num_masks > num_prompts:
assert masks.shape[0] == num_masks and num_prompts == 1
text_embeds = text_embeds.repeat(num_masks, 1, 1)
uncond_embeds = uncond_embeds.repeat(num_masks, 1, 1)
text_embeds = torch.cat([uncond_embeds, text_embeds])
### Run
# Latent initialization.
if self.timesteps[0] < 999 and has_background:
latent = self.scheduler_add_noise(bg_latent, None, 0)
else:
latent = torch.randn((1, self.unet.config.in_channels, h, w), dtype=self.dtype, device=self.device)
# Tiling (if needed).
if height > tile_size or width > tile_size:
t = (tile_size + self.vae_scale_factor - 1) // self.vae_scale_factor
views, tile_masks = get_panorama_views(h, w, t)
tile_masks = tile_masks.to(self.device)
else:
views = [(0, h, 0, w)]
tile_masks = latent.new_ones((1, 1, h, w))
value = torch.zeros_like(latent)
count_all = torch.zeros_like(latent)
with torch.autocast('cuda'):
for i, t in enumerate(tqdm(self.timesteps)):
fg_mask = fg_masks[:, i]
bg_mask = bg_masks[i:i + 1]
value.zero_()
count_all.zero_()
for j, (h_start, h_end, w_start, w_end) in enumerate(views):
fg_mask_ = fg_mask[..., h_start:h_end, w_start:w_end]
latent_ = latent[..., h_start:h_end, w_start:w_end].repeat(num_masks, 1, 1, 1)
# Bootstrap for tight background.
if i < bootstrap_steps:
mix_ratio = min(1, max(0, boostrap_mix_steps - i))
# Treat the first foreground latent as the background latent if one does not exist.
bg_latent_ = bg_latent[..., h_start:h_end, w_start:w_end] if has_background else latent_[:1]
white_ = white[..., h_start:h_end, w_start:w_end]
bg_latent_ = mix_ratio * white_ + (1.0 - mix_ratio) * bg_latent_
bg_latent_ = self.scheduler_add_noise(bg_latent_, None, i)
latent_ = (1.0 - fg_mask_) * bg_latent_ + fg_mask_ * latent_
# Centering.
latent_ = shift_to_mask_bbox_center(latent_, fg_mask_, reverse=True)
# Perform one step of the reverse diffusion.
noise_pred = self.unet(torch.cat([latent_] * 2), t, encoder_hidden_states=text_embeds)['sample']
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
latent_ = self.scheduler_step(noise_pred, i, latent_)
if i < bootstrap_steps:
# Uncentering.
latent_ = shift_to_mask_bbox_center(latent_, fg_mask_)
# Remove leakage (optional).
leak = (latent_ - bg_latent_).pow(2).mean(dim=1, keepdim=True)
leak_sigmoid = torch.sigmoid(leak / bootstrap_leak_sensitivity) * 2 - 1
fg_mask_ = fg_mask_ * leak_sigmoid
# Mix the latents.
fg_mask_ = fg_mask_ * tile_masks[:, j:j + 1, h_start:h_end, w_start:w_end]
value[..., h_start:h_end, w_start:w_end] += (fg_mask_ * latent_).sum(dim=0, keepdim=True)
count_all[..., h_start:h_end, w_start:w_end] += fg_mask_.sum(dim=0, keepdim=True)
latent = torch.where(count_all > 0, value / count_all, value)
bg_mask = (1 - count_all).clip_(0, 1) # (T, 1, h, w)
if has_background:
latent = (1 - bg_mask) * latent + bg_mask * bg_latent
# Noise is added after mixing.
if i < len(self.timesteps) - 1:
latent = self.scheduler_add_noise(latent, None, i + 1)
# Return PIL Image.
image = self.decode_latents(latent.to(dtype=self.dtype))[0]
if has_background and do_blend:
fg_mask = torch.sum(masks_g, dim=0).clip_(0, 1)
image = blend(image, background[0], fg_mask)
else:
image = T.ToPILImage()(image)
return image
|