pipeline_leditspp_stable_diffusion.py

importinspect
importmath
fromitertoolsimportrepeat
fromtypingimportAny, Callable, Dict, List, Optional, Tuple, Union

importtorch
importtorch.nn.functionalasF
frompackagingimportversion
fromtransformersimportCLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from ...configuration_utilsimportFrozenDict
from ...image_processorimportPipelineImageInput, VaeImageProcessor
from ...loadersimportFromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
from ...modelsimportAutoencoderKL, UNet2DConditionModel
from ...models.attention_processorimportAttention, AttnProcessor
from ...models.loraimportadjust_lora_scale_text_encoder
from ...pipelines.stable_diffusion.safety_checkerimportStableDiffusionSafetyChecker
from ...schedulersimportDDIMScheduler, DPMSolverMultistepScheduler
from ...utilsimport (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
scale_lora_layers,
unscale_lora_layers,
)
from ...utils.torch_utilsimportrandn_tensor
from ..pipeline_utilsimportDiffusionPipeline
from .pipeline_outputimportLEditsPPDiffusionPipelineOutput, LEditsPPInversionPipelineOutput


ifis_torch_xla_available():
importtorch_xla.core.xla_modelasxm

XLA_AVAILABLE=True
else:
XLA_AVAILABLE=False

logger=logging.get_logger(__name__) # pylint: disable=invalid-name


EXAMPLE_DOC_STRING="""
 Examples:
 ```py
 >>> import torch

 >>> from diffusers import LEditsPPPipelineStableDiffusion
 >>> from diffusers.utils import load_image

 >>> pipe = LEditsPPPipelineStableDiffusion.from_pretrained(
 ... "runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16
 ... )
 >>> pipe.enable_vae_tiling()
 >>> pipe = pipe.to("cuda")

 >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/cherry_blossom.png"
 >>> image = load_image(img_url).resize((512, 512))

 >>> _ = pipe.invert(image=image, num_inversion_steps=50, skip=0.1)

 >>> edited_image = pipe(
 ... editing_prompt=["cherry blossom"], edit_guidance_scale=10.0, edit_threshold=0.75
 ... ).images[0]
 ```
"""


# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline.AttentionStore
classLeditsAttentionStore:
@staticmethod
defget_empty_store():
return {"down_cross": [], "mid_cross": [], "up_cross": [], "down_self": [], "mid_self": [], "up_self": []}

def__call__(self, attn, is_cross: bool, place_in_unet: str, editing_prompts, PnP=False):
# attn.shape = batch_size * head_size, seq_len query, seq_len_key
ifattn.shape[1] <=self.max_size:
bs=1+int(PnP) +editing_prompts
skip=2ifPnPelse1# skip PnP & unconditional
attn=torch.stack(attn.split(self.batch_size)).permute(1, 0, 2, 3)
source_batch_size=int(attn.shape[1] //bs)
self.forward(attn[:, skip*source_batch_size :], is_cross, place_in_unet)

defforward(self, attn, is_cross: bool, place_in_unet: str):
key=f"{place_in_unet}_{'cross'ifis_crosselse'self'}"

self.step_store[key].append(attn)

defbetween_steps(self, store_step=True):
ifstore_step:
ifself.average:
iflen(self.attention_store) ==0:
self.attention_store=self.step_store
else:
forkeyinself.attention_store:
foriinrange(len(self.attention_store[key])):
self.attention_store[key][i] +=self.step_store[key][i]
else:
iflen(self.attention_store) ==0:
self.attention_store= [self.step_store]
else:
self.attention_store.append(self.step_store)

self.cur_step+=1
self.step_store=self.get_empty_store()

defget_attention(self, step: int):
ifself.average:
attention= {
key: [item/self.cur_stepforiteminself.attention_store[key]] forkeyinself.attention_store
 }
else:
assertstepisnotNone
attention=self.attention_store[step]
returnattention

defaggregate_attention(
self, attention_maps, prompts, res: Union[int, Tuple[int]], from_where: List[str], is_cross: bool, select: int
 ):
out= [[] forxinrange(self.batch_size)]
ifisinstance(res, int):
num_pixels=res**2
resolution= (res, res)
else:
num_pixels=res[0] *res[1]
resolution=res[:2]

forlocationinfrom_where:
forbs_iteminattention_maps[f"{location}_{'cross'ifis_crosselse'self'}"]:
forbatch, iteminenumerate(bs_item):
ifitem.shape[1] ==num_pixels:
cross_maps=item.reshape(len(prompts), -1, *resolution, item.shape[-1])[select]
out[batch].append(cross_maps)

out=torch.stack([torch.cat(x, dim=0) forxinout])
# average over heads
out=out.sum(1) /out.shape[1]
returnout

def__init__(self, average: bool, batch_size=1, max_resolution=16, max_size: int=None):
self.step_store=self.get_empty_store()
self.attention_store= []
self.cur_step=0
self.average=average
self.batch_size=batch_size
ifmax_sizeisNone:
self.max_size=max_resolution**2
elifmax_sizeisnotNoneandmax_resolutionisNone:
self.max_size=max_size
else:
raiseValueError("Only allowed to set one of max_resolution or max_size")


# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionAttendAndExcitePipeline.GaussianSmoothing
classLeditsGaussianSmoothing:
def__init__(self, device):
kernel_size= [3, 3]
sigma= [0.5, 0.5]

# The gaussian kernel is the product of the gaussian function of each dimension.
kernel=1
meshgrids=torch.meshgrid([torch.arange(size, dtype=torch.float32) forsizeinkernel_size], indexing="ij")
forsize, std, mgridinzip(kernel_size, sigma, meshgrids):
mean= (size-1) /2
kernel*=1/ (std*math.sqrt(2*math.pi)) *torch.exp(-(((mgrid-mean) / (2*std)) **2))

# Make sure sum of values in gaussian kernel equals 1.
kernel=kernel/torch.sum(kernel)

# Reshape to depthwise convolutional weight
kernel=kernel.view(1, 1, *kernel.size())
kernel=kernel.repeat(1, *[1] * (kernel.dim() -1))

self.weight=kernel.to(device)

def__call__(self, input):
"""
 Arguments:
 Apply gaussian filter to input.
 input (torch.Tensor): Input to apply gaussian filter on.
 Returns:
 filtered (torch.Tensor): Filtered output.
 """
returnF.conv2d(input, weight=self.weight.to(input.dtype))


classLEDITSCrossAttnProcessor:
def__init__(self, attention_store, place_in_unet, pnp, editing_prompts):
self.attnstore=attention_store
self.place_in_unet=place_in_unet
self.editing_prompts=editing_prompts
self.pnp=pnp

def__call__(
self,
attn: Attention,
hidden_states,
encoder_hidden_states,
attention_mask=None,
temb=None,
 ):
batch_size, sequence_length, _= (
hidden_states.shapeifencoder_hidden_statesisNoneelseencoder_hidden_states.shape
 )
attention_mask=attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)

query=attn.to_q(hidden_states)

ifencoder_hidden_statesisNone:
encoder_hidden_states=hidden_states
elifattn.norm_cross:
encoder_hidden_states=attn.norm_encoder_hidden_states(encoder_hidden_states)

key=attn.to_k(encoder_hidden_states)
value=attn.to_v(encoder_hidden_states)

query=attn.head_to_batch_dim(query)
key=attn.head_to_batch_dim(key)
value=attn.head_to_batch_dim(value)

attention_probs=attn.get_attention_scores(query, key, attention_mask)
self.attnstore(
attention_probs,
is_cross=True,
place_in_unet=self.place_in_unet,
editing_prompts=self.editing_prompts,
PnP=self.pnp,
 )

hidden_states=torch.bmm(attention_probs, value)
hidden_states=attn.batch_to_head_dim(hidden_states)

# linear proj
hidden_states=attn.to_out[0](hidden_states)
# dropout
hidden_states=attn.to_out[1](hidden_states)

hidden_states=hidden_states/attn.rescale_output_factor
returnhidden_states


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
defrescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
r"""
 Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
 Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
 Flawed](https://arxiv.org/pdf/2305.08891.pdf).

 Args:
 noise_cfg (`torch.Tensor`):
 The predicted noise tensor for the guided diffusion process.
 noise_pred_text (`torch.Tensor`):
 The predicted noise tensor for the text-guided diffusion process.
 guidance_rescale (`float`, *optional*, defaults to 0.0):
 A rescale factor applied to the noise predictions.

 Returns:
 noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor.
 """
std_text=noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg=noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# rescale the results from guidance (fixes overexposure)
noise_pred_rescaled=noise_cfg* (std_text/std_cfg)
# mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
noise_cfg=guidance_rescale*noise_pred_rescaled+ (1-guidance_rescale) *noise_cfg
returnnoise_cfg


classLEditsPPPipelineStableDiffusion(
DiffusionPipeline, TextualInversionLoaderMixin, StableDiffusionLoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin
):
"""
 Pipeline for textual image editing using LEDits++ with Stable Diffusion.

 This model inherits from [`DiffusionPipeline`] and builds on the [`StableDiffusionPipeline`]. Check the superclass
 documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular
 device, etc.).

 Args:
 vae ([`AutoencoderKL`]):
 Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
 text_encoder ([`~transformers.CLIPTextModel`]):
 Frozen text-encoder. Stable Diffusion uses the text portion of
 [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
 the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
 tokenizer ([`~transformers.CLIPTokenizer`]):
 Tokenizer of class
 [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
 unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
 scheduler ([`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]):
 A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of
 [`DPMSolverMultistepScheduler`] or [`DDIMScheduler`]. If any other scheduler is passed it will
 automatically be set to [`DPMSolverMultistepScheduler`].
 safety_checker ([`StableDiffusionSafetyChecker`]):
 Classification module that estimates whether generated images could be considered offensive or harmful.
 Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.
 feature_extractor ([`~transformers.CLIPImageProcessor`]):
 Model that extracts features from generated images to be used as inputs for the `safety_checker`.
 """

model_cpu_offload_seq="text_encoder->unet->vae"
_exclude_from_cpu_offload= ["safety_checker"]
_callback_tensor_inputs= ["latents", "prompt_embeds", "negative_prompt_embeds"]
_optional_components= ["safety_checker", "feature_extractor", "image_encoder"]

def__init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: Union[DDIMScheduler, DPMSolverMultistepScheduler],
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool=True,
 ):
super().__init__()

ifnotisinstance(scheduler, DDIMScheduler) andnotisinstance(scheduler, DPMSolverMultistepScheduler):
scheduler=DPMSolverMultistepScheduler.from_config(
scheduler.config, algorithm_type="sde-dpmsolver++", solver_order=2
 )
logger.warning(
"This pipeline only supports DDIMScheduler and DPMSolverMultistepScheduler. "
"The scheduler has been changed to DPMSolverMultistepScheduler."
 )

ifschedulerisnotNoneandgetattr(scheduler.config, "steps_offset", 1) !=1:
deprecation_message= (
f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
"to update the config accordingly as leaving `steps_offset` might led to incorrect results"
" in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
" it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
" file"
 )
deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
new_config=dict(scheduler.config)
new_config["steps_offset"] =1
scheduler._internal_dict=FrozenDict(new_config)

ifschedulerisnotNoneandgetattr(scheduler.config, "clip_sample", False) isTrue:
deprecation_message= (
f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
" `clip_sample` should be set to False in the configuration file. Please make sure to update the"
" config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
" future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
" nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
 )
deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
new_config=dict(scheduler.config)
new_config["clip_sample"] =False
scheduler._internal_dict=FrozenDict(new_config)

ifsafety_checkerisNoneandrequires_safety_checker:
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
 )

ifsafety_checkerisnotNoneandfeature_extractorisNone:
raiseValueError(
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
 )

is_unet_version_less_0_9_0= (
unetisnotNone
andhasattr(unet.config, "_diffusers_version")
andversion.parse(version.parse(unet.config._diffusers_version).base_version) <version.parse("0.9.0.dev0")
 )
is_unet_sample_size_less_64= (
unetisnotNoneandhasattr(unet.config, "sample_size") andunet.config.sample_size<64
 )
ifis_unet_version_less_0_9_0andis_unet_sample_size_less_64:
deprecation_message= (
"The configuration file of the unet has set the default `sample_size` to smaller than"
" 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
" following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
" CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
" \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
" configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
" in the config might lead to incorrect results in future versions. If you have downloaded this"
" checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
" the `unet/config.json` file"
 )
deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
new_config=dict(unet.config)
new_config["sample_size"] =64
unet._internal_dict=FrozenDict(new_config)

self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
 )
self.vae_scale_factor=2** (len(self.vae.config.block_out_channels) -1) ifgetattr(self, "vae", None) else8
self.image_processor=VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
self.register_to_config(requires_safety_checker=requires_safety_checker)

self.inversion_steps=None

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
defrun_safety_checker(self, image, device, dtype):
ifself.safety_checkerisNone:
has_nsfw_concept=None
else:
iftorch.is_tensor(image):
feature_extractor_input=self.image_processor.postprocess(image, output_type="pil")
else:
feature_extractor_input=self.image_processor.numpy_to_pil(image)
safety_checker_input=self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
image, has_nsfw_concept=self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
 )
returnimage, has_nsfw_concept

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
defdecode_latents(self, latents):
deprecation_message="The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)

latents=1/self.vae.config.scaling_factor*latents
image=self.vae.decode(latents, return_dict=False)[0]
image= (image/2+0.5).clamp(0, 1)
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
image=image.cpu().permute(0, 2, 3, 1).float().numpy()
returnimage

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
defprepare_extra_step_kwargs(self, eta, generator=None):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
# and should be between [0, 1]

accepts_eta="eta"inset(inspect.signature(self.scheduler.step).parameters.keys())
extra_step_kwargs= {}
ifaccepts_eta:
extra_step_kwargs["eta"] =eta

# check if the scheduler accepts generator
accepts_generator="generator"inset(inspect.signature(self.scheduler.step).parameters.keys())
ifaccepts_generator:
extra_step_kwargs["generator"] =generator
returnextra_step_kwargs

# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs
defcheck_inputs(
self,
negative_prompt=None,
editing_prompt_embeddings=None,
negative_prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
 ):
ifcallback_on_step_end_tensor_inputsisnotNoneandnotall(
kinself._callback_tensor_inputsforkincallback_on_step_end_tensor_inputs
 ):
raiseValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[kforkincallback_on_step_end_tensor_inputsifknotinself._callback_tensor_inputs]}"
 )
ifnegative_promptisnotNoneandnegative_prompt_embedsisnotNone:
raiseValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
 )

ifediting_prompt_embeddingsisnotNoneandnegative_prompt_embedsisnotNone:
ifediting_prompt_embeddings.shape!=negative_prompt_embeds.shape:
raiseValueError(
"`editing_prompt_embeddings` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `editing_prompt_embeddings` {editing_prompt_embeddings.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
 )

# Modified from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
defprepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, latents):
# shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)

# if latents.shape != shape:
# raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")

latents=latents.to(device)

# scale the initial noise by the standard deviation required by the scheduler
latents=latents*self.scheduler.init_noise_sigma
returnlatents

defprepare_unet(self, attention_store, PnP: bool=False):
attn_procs= {}
fornameinself.unet.attn_processors.keys():
ifname.startswith("mid_block"):
place_in_unet="mid"
elifname.startswith("up_blocks"):
place_in_unet="up"
elifname.startswith("down_blocks"):
place_in_unet="down"
else:
continue

if"attn2"innameandplace_in_unet!="mid":
attn_procs[name] =LEDITSCrossAttnProcessor(
attention_store=attention_store,
place_in_unet=place_in_unet,
pnp=PnP,
editing_prompts=self.enabled_editing_prompts,
 )
else:
attn_procs[name] =AttnProcessor()

self.unet.set_attn_processor(attn_procs)

defencode_prompt(
self,
device,
num_images_per_prompt,
enable_edit_guidance,
negative_prompt=None,
editing_prompt=None,
negative_prompt_embeds: Optional[torch.Tensor] =None,
editing_prompt_embeds: Optional[torch.Tensor] =None,
lora_scale: Optional[float] =None,
clip_skip: Optional[int] =None,
 ):
r"""
 Encodes the prompt into text encoder hidden states.

 Args:
 device: (`torch.device`):
 torch device
 num_images_per_prompt (`int`):
 number of images that should be generated per prompt
 enable_edit_guidance (`bool`):
 whether to perform any editing or reconstruct the input image instead
 negative_prompt (`str` or `List[str]`, *optional*):
 The prompt or prompts not to guide the image generation. If not defined, one has to pass
 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
 less than `1`).
 editing_prompt (`str` or `List[str]`, *optional*):
 Editing prompt(s) to be encoded. If not defined, one has to pass `editing_prompt_embeds` instead.
 editing_prompt_embeds (`torch.Tensor`, *optional*):
 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
 provided, text embeddings will be generated from `prompt` input argument.
 negative_prompt_embeds (`torch.Tensor`, *optional*):
 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
 argument.
 lora_scale (`float`, *optional*):
 A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
 clip_skip (`int`, *optional*):
 Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
 the output of the pre-final layer will be used for computing the prompt embeddings.
 """
# set lora scale so that monkey patched LoRA
# function of text encoder can correctly access it
iflora_scaleisnotNoneandisinstance(self, StableDiffusionLoraLoaderMixin):
self._lora_scale=lora_scale

# dynamically adjust the LoRA scale
ifnotUSE_PEFT_BACKEND:
adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
else:
scale_lora_layers(self.text_encoder, lora_scale)

batch_size=self.batch_size
num_edit_tokens=None

ifnegative_prompt_embedsisNone:
uncond_tokens: List[str]
ifnegative_promptisNone:
uncond_tokens= [""] *batch_size
elifisinstance(negative_prompt, str):
uncond_tokens= [negative_prompt]
elifbatch_size!=len(negative_prompt):
raiseValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but exoected"
f"{batch_size} based on the input images. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
 )
else:
uncond_tokens=negative_prompt

# textual inversion: procecss multi-vector tokens if necessary
ifisinstance(self, TextualInversionLoaderMixin):
uncond_tokens=self.maybe_convert_prompt(uncond_tokens, self.tokenizer)

uncond_input=self.tokenizer(
uncond_tokens,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
 )

ifhasattr(self.text_encoder.config, "use_attention_mask") andself.text_encoder.config.use_attention_mask:
attention_mask=uncond_input.attention_mask.to(device)
else:
attention_mask=None

negative_prompt_embeds=self.text_encoder(
uncond_input.input_ids.to(device),
attention_mask=attention_mask,
 )
negative_prompt_embeds=negative_prompt_embeds[0]

ifself.text_encoderisnotNone:
prompt_embeds_dtype=self.text_encoder.dtype
elifself.unetisnotNone:
prompt_embeds_dtype=self.unet.dtype
else:
prompt_embeds_dtype=negative_prompt_embeds.dtype

negative_prompt_embeds=negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)

ifenable_edit_guidance:
ifediting_prompt_embedsisNone:
# textual inversion: procecss multi-vector tokens if necessary
# if isinstance(self, TextualInversionLoaderMixin):
# prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
ifisinstance(editing_prompt, str):
editing_prompt= [editing_prompt]

max_length=negative_prompt_embeds.shape[1]
text_inputs=self.tokenizer(
 [xforiteminediting_promptforxinrepeat(item, batch_size)],
padding="max_length",
max_length=max_length,
truncation=True,
return_tensors="pt",
return_length=True,
 )

num_edit_tokens=text_inputs.length-2# not counting startoftext and endoftext
text_input_ids=text_inputs.input_ids
untruncated_ids=self.tokenizer(
 [xforiteminediting_promptforxinrepeat(item, batch_size)],
padding="longest",
return_tensors="pt",
 ).input_ids

ifuntruncated_ids.shape[-1] >=text_input_ids.shape[-1] andnottorch.equal(
text_input_ids, untruncated_ids
 ):
removed_text=self.tokenizer.batch_decode(
untruncated_ids[:, self.tokenizer.model_max_length-1 : -1]
 )
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
 )

if (
hasattr(self.text_encoder.config, "use_attention_mask")
andself.text_encoder.config.use_attention_mask
 ):
attention_mask=text_inputs.attention_mask.to(device)
else:
attention_mask=None

ifclip_skipisNone:
editing_prompt_embeds=self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
editing_prompt_embeds=editing_prompt_embeds[0]
else:
editing_prompt_embeds=self.text_encoder(
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
 )
# Access the `hidden_states` first, that contains a tuple of
# all the hidden states from the encoder layers. Then index into
# the tuple to access the hidden states from the desired layer.
editing_prompt_embeds=editing_prompt_embeds[-1][-(clip_skip+1)]
# We also need to apply the final LayerNorm here to not mess with the
# representations. The `last_hidden_states` that we typically use for
# obtaining the final prompt representations passes through the LayerNorm
# layer.
editing_prompt_embeds=self.text_encoder.text_model.final_layer_norm(editing_prompt_embeds)

editing_prompt_embeds=editing_prompt_embeds.to(dtype=negative_prompt_embeds.dtype, device=device)

bs_embed_edit, seq_len, _=editing_prompt_embeds.shape
editing_prompt_embeds=editing_prompt_embeds.to(dtype=negative_prompt_embeds.dtype, device=device)
editing_prompt_embeds=editing_prompt_embeds.repeat(1, num_images_per_prompt, 1)
editing_prompt_embeds=editing_prompt_embeds.view(bs_embed_edit*num_images_per_prompt, seq_len, -1)

# get unconditional embeddings for classifier free guidance

# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
seq_len=negative_prompt_embeds.shape[1]

negative_prompt_embeds=negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)

negative_prompt_embeds=negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
negative_prompt_embeds=negative_prompt_embeds.view(batch_size*num_images_per_prompt, seq_len, -1)

ifisinstance(self, StableDiffusionLoraLoaderMixin) andUSE_PEFT_BACKEND:
# Retrieve the original scale by scaling back the LoRA layers
unscale_lora_layers(self.text_encoder, lora_scale)

returnediting_prompt_embeds, negative_prompt_embeds, num_edit_tokens

@property
defguidance_rescale(self):
returnself._guidance_rescale

@property
defclip_skip(self):
returnself._clip_skip

@property
defcross_attention_kwargs(self):
returnself._cross_attention_kwargs

defenable_vae_slicing(self):
r"""
 Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
 compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
 """
self.vae.enable_slicing()

defdisable_vae_slicing(self):
r"""
 Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
 computing decoding in one step.
 """
self.vae.disable_slicing()

defenable_vae_tiling(self):
r"""
 Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
 compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
 processing larger images.
 """
self.vae.enable_tiling()

defdisable_vae_tiling(self):
r"""
 Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
 computing decoding in one step.
 """
self.vae.disable_tiling()

@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def__call__(
self,
negative_prompt: Optional[Union[str, List[str]]] =None,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] =None,
output_type: Optional[str] ="pil",
return_dict: bool=True,
editing_prompt: Optional[Union[str, List[str]]] =None,
editing_prompt_embeds: Optional[torch.Tensor] =None,
negative_prompt_embeds: Optional[torch.Tensor] =None,
reverse_editing_direction: Optional[Union[bool, List[bool]]] =False,
edit_guidance_scale: Optional[Union[float, List[float]]] =5,
edit_warmup_steps: Optional[Union[int, List[int]]] =0,
edit_cooldown_steps: Optional[Union[int, List[int]]] =None,
edit_threshold: Optional[Union[float, List[float]]] =0.9,
user_mask: Optional[torch.Tensor] =None,
sem_guidance: Optional[List[torch.Tensor]] =None,
use_cross_attn_mask: bool=False,
use_intersect_mask: bool=True,
attn_store_steps: Optional[List[int]] = [],
store_averaged_over_steps: bool=True,
cross_attention_kwargs: Optional[Dict[str, Any]] =None,
guidance_rescale: float=0.0,
clip_skip: Optional[int] =None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] =None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
**kwargs,
 ):
r"""
 The call function to the pipeline for editing. The
 [`~pipelines.ledits_pp.LEditsPPPipelineStableDiffusion.invert`] method has to be called beforehand. Edits will
 always be performed for the last inverted image(s).

 Args:
 negative_prompt (`str` or `List[str]`, *optional*):
 The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
 if `guidance_scale` is less than `1`).
 generator (`torch.Generator`, *optional*):
 One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
 to make generation deterministic.
 output_type (`str`, *optional*, defaults to `"pil"`):
 The output format of the generate image. Choose between
 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
 return_dict (`bool`, *optional*, defaults to `True`):
 Whether or not to return a [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] instead of a plain
 tuple.
 editing_prompt (`str` or `List[str]`, *optional*):
 The prompt or prompts to guide the image generation. The image is reconstructed by setting
 `editing_prompt = None`. Guidance direction of prompt should be specified via
 `reverse_editing_direction`.
 editing_prompt_embeds (`torch.Tensor>`, *optional*):
 Pre-computed embeddings to use for guiding the image generation. Guidance direction of embedding should
 be specified via `reverse_editing_direction`.
 negative_prompt_embeds (`torch.Tensor`, *optional*):
 Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
 not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
 reverse_editing_direction (`bool` or `List[bool]`, *optional*, defaults to `False`):
 Whether the corresponding prompt in `editing_prompt` should be increased or decreased.
 edit_guidance_scale (`float` or `List[float]`, *optional*, defaults to 5):
 Guidance scale for guiding the image generation. If provided as list values should correspond to
 `editing_prompt`. `edit_guidance_scale` is defined as `s_e` of equation 12 of [LEDITS++
 Paper](https://arxiv.org/abs/2301.12247).
 edit_warmup_steps (`float` or `List[float]`, *optional*, defaults to 10):
 Number of diffusion steps (for each prompt) for which guidance will not be applied.
 edit_cooldown_steps (`float` or `List[float]`, *optional*, defaults to `None`):
 Number of diffusion steps (for each prompt) after which guidance will no longer be applied.
 edit_threshold (`float` or `List[float]`, *optional*, defaults to 0.9):
 Masking threshold of guidance. Threshold should be proportional to the image region that is modified.
 'edit_threshold' is defined as 'λ' of equation 12 of [LEDITS++
 Paper](https://arxiv.org/abs/2301.12247).
 user_mask (`torch.Tensor`, *optional*):
 User-provided mask for even better control over the editing process. This is helpful when LEDITS++'s
 implicit masks do not meet user preferences.
 sem_guidance (`List[torch.Tensor]`, *optional*):
 List of pre-generated guidance vectors to be applied at generation. Length of the list has to
 correspond to `num_inference_steps`.
 use_cross_attn_mask (`bool`, defaults to `False`):
 Whether cross-attention masks are used. Cross-attention masks are always used when use_intersect_mask
 is set to true. Cross-attention masks are defined as 'M^1' of equation 12 of [LEDITS++
 paper](https://arxiv.org/pdf/2311.16711.pdf).
 use_intersect_mask (`bool`, defaults to `True`):
 Whether the masking term is calculated as intersection of cross-attention masks and masks derived from
 the noise estimate. Cross-attention mask are defined as 'M^1' and masks derived from the noise estimate
 are defined as 'M^2' of equation 12 of [LEDITS++ paper](https://arxiv.org/pdf/2311.16711.pdf).
 attn_store_steps (`List[int]`, *optional*):
 Steps for which the attention maps are stored in the AttentionStore. Just for visualization purposes.
 store_averaged_over_steps (`bool`, defaults to `True`):
 Whether the attention maps for the 'attn_store_steps' are stored averaged over the diffusion steps. If
 False, attention maps for each step are stores separately. Just for visualization purposes.
 cross_attention_kwargs (`dict`, *optional*):
 A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
 [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
 guidance_rescale (`float`, *optional*, defaults to 0.0):
 Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
 Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
 using zero terminal SNR.
 clip_skip (`int`, *optional*):
 Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
 the output of the pre-final layer will be used for computing the prompt embeddings.
 callback_on_step_end (`Callable`, *optional*):
 A function that calls at the end of each denoising steps during the inference. The function is called
 with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
 callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
 `callback_on_step_end_tensor_inputs`.
 callback_on_step_end_tensor_inputs (`List`, *optional*):
 The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
 `._callback_tensor_inputs` attribute of your pipeline class.

 Examples:

 Returns:
 [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] or `tuple`:
 [`~pipelines.ledits_pp.LEditsPPDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. When
 returning a tuple, the first element is a list with the generated images, and the second element is a list
 of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" (nsfw)
 content, according to the `safety_checker`.
 """

ifself.inversion_stepsisNone:
raiseValueError(
"You need to invert an input image first before calling the pipeline. The `invert` method has to be called beforehand. Edits will always be performed for the last inverted image(s)."
 )

eta=self.eta
num_images_per_prompt=1
latents=self.init_latents

zs=self.zs
self.scheduler.set_timesteps(len(self.scheduler.timesteps))

ifuse_intersect_mask:
use_cross_attn_mask=True

ifuse_cross_attn_mask:
self.smoothing=LeditsGaussianSmoothing(self.device)

ifuser_maskisnotNone:
user_mask=user_mask.to(self.device)

org_prompt=""

# 1. Check inputs. Raise error if not correct
self.check_inputs(
negative_prompt,
editing_prompt_embeds,
negative_prompt_embeds,
callback_on_step_end_tensor_inputs,
 )

self._guidance_rescale=guidance_rescale
self._clip_skip=clip_skip
self._cross_attention_kwargs=cross_attention_kwargs

# 2. Define call parameters
batch_size=self.batch_size

ifediting_prompt:
enable_edit_guidance=True
ifisinstance(editing_prompt, str):
editing_prompt= [editing_prompt]
self.enabled_editing_prompts=len(editing_prompt)
elifediting_prompt_embedsisnotNone:
enable_edit_guidance=True
self.enabled_editing_prompts=editing_prompt_embeds.shape[0]
else:
self.enabled_editing_prompts=0
enable_edit_guidance=False

# 3. Encode input prompt
lora_scale= (
self.cross_attention_kwargs.get("scale", None) ifself.cross_attention_kwargsisnotNoneelseNone
 )

edit_concepts, uncond_embeddings, num_edit_tokens=self.encode_prompt(
editing_prompt=editing_prompt,
device=self.device,
num_images_per_prompt=num_images_per_prompt,
enable_edit_guidance=enable_edit_guidance,
negative_prompt=negative_prompt,
editing_prompt_embeds=editing_prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
 )

# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
ifenable_edit_guidance:
text_embeddings=torch.cat([uncond_embeddings, edit_concepts])
self.text_cross_attention_maps= [editing_prompt] ifisinstance(editing_prompt, str) elseediting_prompt
else:
text_embeddings=torch.cat([uncond_embeddings])

# 4. Prepare timesteps
# self.scheduler.set_timesteps(num_inference_steps, device=self.device)
timesteps=self.inversion_steps
t_to_idx= {int(v): kfork, vinenumerate(timesteps[-zs.shape[0] :])}

ifuse_cross_attn_mask:
self.attention_store=LeditsAttentionStore(
average=store_averaged_over_steps,
batch_size=batch_size,
max_size=(latents.shape[-2] /4.0) * (latents.shape[-1] /4.0),
max_resolution=None,
 )
self.prepare_unet(self.attention_store, PnP=False)
resolution=latents.shape[-2:]
att_res= (int(resolution[0] /4), int(resolution[1] /4))

# 5. Prepare latent variables
num_channels_latents=self.unet.config.in_channels
latents=self.prepare_latents(
batch_size*num_images_per_prompt,
num_channels_latents,
None,
None,
text_embeddings.dtype,
self.device,
latents,
 )

# 6. Prepare extra step kwargs.
extra_step_kwargs=self.prepare_extra_step_kwargs(eta)

self.sem_guidance=None
self.activation_mask=None

# 7. Denoising loop
num_warmup_steps=0
withself.progress_bar(total=len(timesteps)) asprogress_bar:
fori, tinenumerate(timesteps):
# expand the latents if we are doing classifier free guidance

ifenable_edit_guidance:
latent_model_input=torch.cat([latents] * (1+self.enabled_editing_prompts))
else:
latent_model_input=latents

latent_model_input=self.scheduler.scale_model_input(latent_model_input, t)

text_embed_input=text_embeddings

# predict the noise residual
noise_pred=self.unet(latent_model_input, t, encoder_hidden_states=text_embed_input).sample

noise_pred_out=noise_pred.chunk(1+self.enabled_editing_prompts) # [b,4, 64, 64]
noise_pred_uncond=noise_pred_out[0]
noise_pred_edit_concepts=noise_pred_out[1:]