junnyu/webui_ppdiffusers
0
1# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.2# Copyright 2023 The HuggingFace Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15#16# modified from https://github.com/AUTOMATIC1111/stable-diffusion-webui17# Here is the AGPL-3.0 license https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/LICENSE.txt18 19import inspect20import shutil21from pathlib import Path22from typing import Any, Callable, Dict, List, Optional, Union23 24import paddle25import paddle.nn as nn26 27from paddlenlp.transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer28from ppdiffusers.models import AutoencoderKL, UNet2DConditionModel29from ppdiffusers.pipelines.pipeline_utils import DiffusionPipeline30from ppdiffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput31from ppdiffusers.pipelines.stable_diffusion.safety_checker import (32 StableDiffusionSafetyChecker,33)34from ppdiffusers.schedulers import KarrasDiffusionSchedulers35from ppdiffusers.utils import (36 PPDIFFUSERS_CACHE,37 logging,38 ppdiffusers_url_download,39 randn_tensor,40 safetensors_load,41 smart_load,42 torch_load,43)44 45logger = logging.get_logger(__name__) # pylint: disable=invalid-name46 47 48import copy49import os50import os.path51 52from huggingface_hub.file_download import _request_wrapper, hf_raise_for_status53 54# lark omegaconf55 56 57def get_civitai_download_url(display_url, url_prefix="https://civitai.com"):58 if "api/download" in display_url:59 return display_url60 import bs461 import requests62 63 headers = {64 "User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36 QIHU 360SE"65 }66 r = requests.get(display_url, headers=headers)67 soup = bs4.BeautifulSoup(r.text, "lxml")68 download_url = None69 for a in soup.find_all("a", href=True):70 if "Download" in str(a):71 download_url = url_prefix + a["href"].split("?")[0]72 break73 return download_url74 75 76def http_file_name(77 url: str,78 *,79 proxies=None,80 headers: Optional[Dict[str, str]] = None,81 timeout=10.0,82 max_retries=0,83):84 """85 Get a remote file name.86 """87 headers = copy.deepcopy(headers) or {}88 r = _request_wrapper(89 method="GET",90 url=url,91 stream=True,92 proxies=proxies,93 headers=headers,94 timeout=timeout,95 max_retries=max_retries,96 )97 hf_raise_for_status(r)98 displayed_name = url.split("/")[-1]99 content_disposition = r.headers.get("Content-Disposition")100 if content_disposition is not None and "filename=" in content_disposition:101 # Means file is on CDN102 displayed_name = content_disposition.split("filename=")[-1]103 return displayed_name104 105 106@paddle.no_grad()107def load_lora(108 pipeline,109 state_dict: dict,110 LORA_PREFIX_UNET: str = "lora_unet",111 LORA_PREFIX_TEXT_ENCODER: str = "lora_te",112 ratio: float = 1.0,113):114 ratio = float(ratio)115 visited = []116 for key in state_dict:117 if ".alpha" in key or ".lora_up" in key or key in visited:118 continue119 120 if "text" in key:121 tmp_layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")122 hf_to_ppnlp = {123 "encoder": "transformer",124 "fc1": "linear1",125 "fc2": "linear2",126 }127 layer_infos = []128 for layer_info in tmp_layer_infos:129 if layer_info == "mlp":130 continue131 layer_infos.append(hf_to_ppnlp.get(layer_info, layer_info))132 curr_layer: paddle.nn.Linear = pipeline.text_encoder133 else:134 layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")135 curr_layer: paddle.nn.Linear = pipeline.unet136 137 temp_name = layer_infos.pop(0)138 while len(layer_infos) > -1:139 try:140 if temp_name == "to":141 raise ValueError()142 curr_layer = curr_layer.__getattr__(temp_name)143 if len(layer_infos) > 0:144 temp_name = layer_infos.pop(0)145 elif len(layer_infos) == 0:146 break147 except Exception:148 if len(temp_name) > 0:149 temp_name += "_" + layer_infos.pop(0)150 else:151 temp_name = layer_infos.pop(0)152 153 triplet_keys = [key, key.replace("lora_down", "lora_up"), key.replace("lora_down.weight", "alpha")]154 dtype: paddle.dtype = curr_layer.weight.dtype155 weight_down: paddle.Tensor = state_dict[triplet_keys[0]].cast(dtype)156 weight_up: paddle.Tensor = state_dict[triplet_keys[1]].cast(dtype)157 rank: float = float(weight_down.shape[0])158 if triplet_keys[2] in state_dict:159 alpha: float = state_dict[triplet_keys[2]].cast(dtype).item()160 scale: float = alpha / rank161 else:162 scale = 1.0163 164 if not hasattr(curr_layer, "backup_weights"):165 curr_layer.backup_weights = curr_layer.weight.clone()166 167 if len(weight_down.shape) == 4:168 if weight_down.shape[2:4] == [1, 1]:169 # conv2d 1x1170 curr_layer.weight.copy_(171 curr_layer.weight172 + ratio173 * paddle.matmul(weight_up.squeeze([-1, -2]), weight_down.squeeze([-1, -2])).unsqueeze([-1, -2])174 * scale,175 True,176 )177 else:178 # conv2d 3x3179 curr_layer.weight.copy_(180 curr_layer.weight181 + ratio182 * paddle.nn.functional.conv2d(weight_down.transpose([1, 0, 2, 3]), weight_up).transpose(183 [1, 0, 2, 3]184 )185 * scale,186 True,187 )188 else:189 # linear190 curr_layer.weight.copy_(curr_layer.weight + ratio * paddle.matmul(weight_up, weight_down).T * scale, True)191 192 # update visited list193 visited.extend(triplet_keys)194 return pipeline195 196 197class WebUIStableDiffusionPipeline(DiffusionPipeline):198 r"""199 Pipeline for text-to-image generation using Stable Diffusion.200 201 This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the202 library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)203 204 Args:205 vae ([`AutoencoderKL`]):206 Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.207 text_encoder ([`CLIPTextModel`]):208 Frozen text-encoder. Stable Diffusion uses the text portion of209 [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically210 the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.211 tokenizer (`CLIPTokenizer`):212 Tokenizer of class213 [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).214 unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.215 scheduler ([`SchedulerMixin`]):216 A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of217 [`DDIMScheduler`], [`LMSDiscreteScheduler`], [`PNDMScheduler`], [`EulerDiscreteScheduler`], [`EulerAncestralDiscreteScheduler`]218 or [`DPMSolverMultistepScheduler`].219 safety_checker ([`StableDiffusionSafetyChecker`]):220 Classification module that estimates whether generated images could be considered offensive or harmful.221 Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.222 feature_extractor ([`CLIPFeatureExtractor`]):223 Model that extracts features from generated images to be used as inputs for the `safety_checker`.224 """225 _optional_components = ["safety_checker", "feature_extractor"]226 enable_emphasis = True227 comma_padding_backtrack = 20228 LORA_DIR = os.path.join(PPDIFFUSERS_CACHE, "lora")229 TI_DIR = os.path.join(PPDIFFUSERS_CACHE, "textual_inversion")230 231 def __init__(232 self,233 vae: AutoencoderKL,234 text_encoder: CLIPTextModel,235 tokenizer: CLIPTokenizer,236 unet: UNet2DConditionModel,237 scheduler: KarrasDiffusionSchedulers,238 safety_checker: StableDiffusionSafetyChecker,239 feature_extractor: CLIPFeatureExtractor,240 requires_safety_checker: bool = True,241 ):242 super().__init__()243 244 if safety_checker is None and requires_safety_checker:245 logger.warning(246 f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"247 " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"248 " results in services or applications open to the public. PaddleNLP team, diffusers team and Hugging Face"249 " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"250 " it only for use-cases that involve analyzing network behavior or auditing its results. For more"251 " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."252 )253 254 if safety_checker is not None and feature_extractor is None:255 raise ValueError(256 f"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"257 " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."258 )259 260 self.register_modules(261 vae=vae,262 text_encoder=text_encoder,263 tokenizer=tokenizer,264 unet=unet,265 scheduler=scheduler,266 safety_checker=safety_checker,267 feature_extractor=feature_extractor,268 )269 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)270 self.register_to_config(requires_safety_checker=requires_safety_checker)271 272 # custom data273 clip_model = FrozenCLIPEmbedder(text_encoder, tokenizer)274 self.sj = StableDiffusionModelHijack(clip_model)275 self.orginal_scheduler_config = self.scheduler.config276 self.supported_scheduler = [277 "pndm",278 "lms",279 "euler",280 "euler-ancestral",281 "dpm-multi",282 "dpm-single",283 "unipc-multi",284 "ddim",285 "ddpm",286 "deis-multi",287 "heun",288 "kdpm2-ancestral",289 "kdpm2",290 ]291 self.weights_has_changed = False292 293 # register_state_dict_hook to fix text_encoder, when we save_pretrained text model.294 def map_to(state_dict, *args, **kwargs):295 if "text_model.token_embedding.wrapped.weight" in state_dict:296 state_dict["text_model.token_embedding.weight"] = state_dict.pop(297 "text_model.token_embedding.wrapped.weight"298 )299 return state_dict300 301 self.text_encoder.register_state_dict_hook(map_to)302 303 def add_ti_embedding_dir(self, embeddings_dir=None):304 self.sj.embedding_db.add_embedding_dir(embeddings_dir)305 self.sj.embedding_db.load_textual_inversion_embeddings()306 307 def clear_ti_embedding(self):308 self.sj.embedding_db.clear_embedding_dirs()309 self.sj.embedding_db.load_textual_inversion_embeddings(True)310 311 def download_civitai_lora_file(self, url):312 if os.path.isfile(url):313 dst = os.path.join(self.LORA_DIR, os.path.basename(url))314 shutil.copyfile(url, dst)315 return dst316 317 download_url = get_civitai_download_url(url) or url318 file_path = ppdiffusers_url_download(319 download_url, cache_dir=self.LORA_DIR, filename=http_file_name(download_url).strip('"')320 )321 return file_path322 323 def download_civitai_ti_file(self, url):324 if os.path.isfile(url):325 dst = os.path.join(self.TI_DIR, os.path.basename(url))326 shutil.copyfile(url, dst)327 return dst328 329 download_url = get_civitai_download_url(url) or url330 file_path = ppdiffusers_url_download(331 download_url, cache_dir=self.TI_DIR, filename=http_file_name(download_url).strip('"')332 )333 return file_path334 335 def change_scheduler(self, scheduler_type="ddim"):336 self.switch_scheduler(scheduler_type)337 338 def switch_scheduler(self, scheduler_type="ddim"):339 scheduler_type = scheduler_type.lower()340 from ppdiffusers import (341 DDIMScheduler,342 DDPMScheduler,343 DEISMultistepScheduler,344 DPMSolverMultistepScheduler,345 DPMSolverSinglestepScheduler,346 EulerAncestralDiscreteScheduler,347 EulerDiscreteScheduler,348 HeunDiscreteScheduler,349 KDPM2AncestralDiscreteScheduler,350 KDPM2DiscreteScheduler,351 LMSDiscreteScheduler,352 PNDMScheduler,353 UniPCMultistepScheduler,354 )355 356 if scheduler_type == "pndm":357 scheduler = PNDMScheduler.from_config(self.orginal_scheduler_config, skip_prk_steps=True)358 elif scheduler_type == "lms":359 scheduler = LMSDiscreteScheduler.from_config(self.orginal_scheduler_config)360 elif scheduler_type == "heun":361 scheduler = HeunDiscreteScheduler.from_config(self.orginal_scheduler_config)362 elif scheduler_type == "euler":363 scheduler = EulerDiscreteScheduler.from_config(self.orginal_scheduler_config)364 elif scheduler_type == "euler-ancestral":365 scheduler = EulerAncestralDiscreteScheduler.from_config(self.orginal_scheduler_config)366 elif scheduler_type == "dpm-multi":367 scheduler = DPMSolverMultistepScheduler.from_config(self.orginal_scheduler_config)368 elif scheduler_type == "dpm-single":369 scheduler = DPMSolverSinglestepScheduler.from_config(self.orginal_scheduler_config)370 elif scheduler_type == "kdpm2-ancestral":371 scheduler = KDPM2AncestralDiscreteScheduler.from_config(self.orginal_scheduler_config)372 elif scheduler_type == "kdpm2":373 scheduler = KDPM2DiscreteScheduler.from_config(self.orginal_scheduler_config)374 elif scheduler_type == "unipc-multi":375 scheduler = UniPCMultistepScheduler.from_config(self.orginal_scheduler_config)376 elif scheduler_type == "ddim":377 scheduler = DDIMScheduler.from_config(378 self.orginal_scheduler_config,379 steps_offset=1,380 clip_sample=False,381 set_alpha_to_one=False,382 )383 elif scheduler_type == "ddpm":384 scheduler = DDPMScheduler.from_config(385 self.orginal_scheduler_config,386 )387 elif scheduler_type == "deis-multi":388 scheduler = DEISMultistepScheduler.from_config(389 self.orginal_scheduler_config,390 )391 else:392 raise ValueError(393 f"Scheduler of type {scheduler_type} doesn't exist! Please choose in {self.supported_scheduler}!"394 )395 self.scheduler = scheduler396 397 @paddle.no_grad()398 def _encode_prompt(399 self,400 prompt: str,401 do_classifier_free_guidance: float = 7.5,402 negative_prompt: str = None,403 num_inference_steps: int = 50,404 ):405 if do_classifier_free_guidance:406 assert isinstance(negative_prompt, str)407 negative_prompt = [negative_prompt]408 uc = get_learned_conditioning(self.sj.clip, negative_prompt, num_inference_steps)409 else:410 uc = None411 412 c = get_multicond_learned_conditioning(self.sj.clip, prompt, num_inference_steps)413 return c, uc414 415 def run_safety_checker(self, image, dtype):416 if self.safety_checker is not None:417 safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pd")418 image, has_nsfw_concept = self.safety_checker(419 images=image, clip_input=safety_checker_input.pixel_values.cast(dtype)420 )421 else:422 has_nsfw_concept = None423 return image, has_nsfw_concept424 425 def decode_latents(self, latents):426 latents = 1 / self.vae.config.scaling_factor * latents427 image = self.vae.decode(latents).sample428 image = (image / 2 + 0.5).clip(0, 1)429 # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16430 image = image.transpose([0, 2, 3, 1]).cast("float32").numpy()431 return image432 433 def prepare_extra_step_kwargs(self, generator, eta):434 # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature435 # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.436 # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502437 # and should be between [0, 1]438 439 accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())440 extra_step_kwargs = {}441 if accepts_eta:442 extra_step_kwargs["eta"] = eta443 444 # check if the scheduler accepts generator445 accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())446 if accepts_generator:447 extra_step_kwargs["generator"] = generator448 return extra_step_kwargs449 450 def check_inputs(451 self,452 prompt,453 height,454 width,455 callback_steps,456 negative_prompt=None,457 ):458 if height % 8 != 0 or width % 8 != 0:459 raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")460 461 if (callback_steps is None) or (462 callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)463 ):464 raise ValueError(465 f"`callback_steps` has to be a positive integer but is {callback_steps} of type"466 f" {type(callback_steps)}."467 )468 469 if prompt is not None and not isinstance(prompt, str):470 raise ValueError(f"`prompt` has to be of type `str` but is {type(prompt)}")471 472 if negative_prompt is not None and not isinstance(negative_prompt, str):473 raise ValueError(f"`negative_prompt` has to be of type `str` but is {type(negative_prompt)}")474 475 def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, generator, latents=None):476 shape = [batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor]477 if isinstance(generator, list) and len(generator) != batch_size:478 raise ValueError(479 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"480 f" size of {batch_size}. Make sure the batch size matches the length of the generators."481 )482 483 if latents is None:484 latents = randn_tensor(shape, generator=generator, dtype=dtype)485 486 # scale the initial noise by the standard deviation required by the scheduler487 latents = latents * self.scheduler.init_noise_sigma488 return latents489 490 @paddle.no_grad()491 def __call__(492 self,493 prompt: str = None,494 height: Optional[int] = None,495 width: Optional[int] = None,496 num_inference_steps: int = 50,497 guidance_scale: float = 7.5,498 negative_prompt: str = None,499 eta: float = 0.0,500 generator: Optional[Union[paddle.Generator, List[paddle.Generator]]] = None,501 latents: Optional[paddle.Tensor] = None,502 output_type: Optional[str] = "pil",503 return_dict: bool = True,504 callback: Optional[Callable[[int, int, paddle.Tensor], None]] = None,505 callback_steps: Optional[int] = 1,506 cross_attention_kwargs: Optional[Dict[str, Any]] = None,507 clip_skip: int = 1,508 enable_lora: bool = True,509 ):510 r"""511 Function invoked when calling the pipeline for generation.512 513 Args:514 prompt (`str`, *optional*):515 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.516 instead.517 height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):518 The height in pixels of the generated image.519 width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):520 The width in pixels of the generated image.521 num_inference_steps (`int`, *optional*, defaults to 50):522 The number of denoising steps. More denoising steps usually lead to a higher quality image at the523 expense of slower inference.524 guidance_scale (`float`, *optional*, defaults to 7.5):525 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).526 `guidance_scale` is defined as `w` of equation 2. of [Imagen527 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >528 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,529 usually at the expense of lower image quality.530 negative_prompt (`str`, *optional*):531 The prompt or prompts not to guide the image generation. If not defined, one has to pass532 `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.533 Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).534 eta (`float`, *optional*, defaults to 0.0):535 Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to536 [`schedulers.DDIMScheduler`], will be ignored for others.537 generator (`paddle.Generator` or `List[paddle.Generator]`, *optional*):538 One or a list of paddle generator(s) to make generation deterministic.539 latents (`paddle.Tensor`, *optional*):540 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image541 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents542 tensor will ge generated by sampling using the supplied random `generator`.543 output_type (`str`, *optional*, defaults to `"pil"`):544 The output format of the generate image. Choose between545 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.546 return_dict (`bool`, *optional*, defaults to `True`):547 Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a548 plain tuple.549 callback (`Callable`, *optional*):550 A function that will be called every `callback_steps` steps during inference. The function will be551 called with the following arguments: `callback(step: int, timestep: int, latents: paddle.Tensor)`.552 callback_steps (`int`, *optional*, defaults to 1):553 The frequency at which the `callback` function will be called. If not specified, the callback will be554 called at every step.555 cross_attention_kwargs (`dict`, *optional*):556 A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under557 `self.processor` in558 [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).559 clip_skip (`int`, *optional*, defaults to 1):560 CLIP_stop_at_last_layers, if clip_skip <= 1, we will use the last_hidden_state from text_encoder.561 Examples:562 563 Returns:564 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:565 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.566 When returning a tuple, the first element is a list with the generated images, and the second element is a567 list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"568 (nsfw) content, according to the `safety_checker`.569 """570 self.add_ti_embedding_dir(self.TI_DIR)571 572 try:573 # 0. Default height and width to unet574 height = height or max(self.unet.config.sample_size * self.vae_scale_factor, 512)575 width = width or max(self.unet.config.sample_size * self.vae_scale_factor, 512)576 577 # 1. Check inputs. Raise error if not correct578 self.check_inputs(579 prompt,580 height,581 width,582 callback_steps,583 negative_prompt,584 )585 586 batch_size = 1587 588 # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)589 # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`590 # corresponds to doing no classifier free guidance.591 do_classifier_free_guidance = guidance_scale > 1.0592 593 prompts, extra_network_data = parse_prompts([prompt])594 595 if enable_lora and self.LORA_DIR is not None:596 if os.path.exists(self.LORA_DIR):597 lora_mapping = {p.stem: p.absolute() for p in Path(self.LORA_DIR).glob("*.safetensors")}598 for params in extra_network_data["lora"]:599 assert len(params.items) > 0600 name = params.items[0]601 if name in lora_mapping:602 ratio = float(params.items[1]) if len(params.items) > 1 else 1.0603 lora_state_dict = smart_load(lora_mapping[name], map_location=paddle.get_device())604 self.weights_has_changed = True605 load_lora(self, state_dict=lora_state_dict, ratio=ratio)606 del lora_state_dict607 else:608 print(f"We can't find lora weight: {name}! Please make sure that exists!")609 else:610 if len(extra_network_data["lora"]) > 0:611 print(f"{self.LORA_DIR} not exists, so we cant load loras!")612 613 self.sj.clip.CLIP_stop_at_last_layers = clip_skip614 # 3. Encode input prompt615 prompt_embeds, negative_prompt_embeds = self._encode_prompt(616 prompts,617 do_classifier_free_guidance,618 negative_prompt,619 num_inference_steps=num_inference_steps,620 )621 622 # 4. Prepare timesteps623 self.scheduler.set_timesteps(num_inference_steps)624 timesteps = self.scheduler.timesteps625 626 # 5. Prepare latent variables627 num_channels_latents = self.unet.in_channels628 latents = self.prepare_latents(629 batch_size,630 num_channels_latents,631 height,632 width,633 self.unet.dtype,634 generator,635 latents,636 )637 638 # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline639 extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)640 641 # 7. Denoising loop642 num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order643 with self.progress_bar(total=num_inference_steps) as progress_bar:644 for i, t in enumerate(timesteps):645 step = i // self.scheduler.order646 do_batch = False647 conds_list, cond_tensor = reconstruct_multicond_batch(prompt_embeds, step)648 try:649 weight = conds_list[0][0][1]650 except Exception:651 weight = 1.0652 if do_classifier_free_guidance:653 uncond_tensor = reconstruct_cond_batch(negative_prompt_embeds, step)654 do_batch = cond_tensor.shape[1] == uncond_tensor.shape[1]655 656 # expand the latents if we are doing classifier free guidance657 latent_model_input = paddle.concat([latents] * 2) if do_batch else latents658 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)659 660 if do_batch:661 noise_pred = self.unet(662 latent_model_input,663 t,664 encoder_hidden_states=paddle.concat([uncond_tensor, cond_tensor]),665 cross_attention_kwargs=cross_attention_kwargs,666 ).sample667 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)668 noise_pred = noise_pred_uncond + weight * guidance_scale * (669 noise_pred_text - noise_pred_uncond670 )671 else:672 noise_pred = self.unet(673 latent_model_input,674 t,675 encoder_hidden_states=cond_tensor,676 cross_attention_kwargs=cross_attention_kwargs,677 ).sample678 679 if do_classifier_free_guidance:680 noise_pred_uncond = self.unet(681 latent_model_input,682 t,683 encoder_hidden_states=uncond_tensor,684 cross_attention_kwargs=cross_attention_kwargs,685 ).sample686 noise_pred = noise_pred_uncond + weight * guidance_scale * (noise_pred - noise_pred_uncond)687 688 # compute the previous noisy sample x_t -> x_t-1689 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample690 691 # call the callback, if provided692 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):693 progress_bar.update()694 if callback is not None and i % callback_steps == 0:695 callback(i, t, latents)696 697 if output_type == "latent":698 image = latents699 has_nsfw_concept = None700 elif output_type == "pil":701 # 8. Post-processing702 image = self.decode_latents(latents)703 704 # 9. Run safety checker705 image, has_nsfw_concept = self.run_safety_checker(image, self.unet.dtype)706 707 # 10. Convert to PIL708 image = self.numpy_to_pil(image)709 else:710 # 8. Post-processing711 image = self.decode_latents(latents)712 713 # 9. Run safety checker714 image, has_nsfw_concept = self.run_safety_checker(image, self.unet.dtype)715 716 if not return_dict:717 return (image, has_nsfw_concept)718 719 return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)720 except Exception as e:721 raise ValueError(e)722 finally:723 if enable_lora and self.weights_has_changed:724 for sub_layer in self.text_encoder.sublayers(include_self=True):725 if hasattr(sub_layer, "backup_weights"):726 sub_layer.weight.copy_(sub_layer.backup_weights, True)727 for sub_layer in self.unet.sublayers(include_self=True):728 if hasattr(sub_layer, "backup_weights"):729 sub_layer.weight.copy_(sub_layer.backup_weights, True)730 self.weights_has_changed = False731 732 733# clip.py734import math735from collections import namedtuple736 737 738class PromptChunk:739 """740 This object contains token ids, weight (multipliers:1.4) and textual inversion embedding info for a chunk of prompt.741 If a prompt is short, it is represented by one PromptChunk, otherwise, multiple are necessary.742 Each PromptChunk contains an exact amount of tokens - 77, which includes one for start and end token,743 so just 75 tokens from prompt.744 """745 746 def __init__(self):747 self.tokens = []748 self.multipliers = []749 self.fixes = []750 751 752PromptChunkFix = namedtuple("PromptChunkFix", ["offset", "embedding"])753"""An object of this type is a marker showing that textual inversion embedding's vectors have to placed at offset in the prompt754chunk. Thos objects are found in PromptChunk.fixes and, are placed into FrozenCLIPEmbedderWithCustomWordsBase.hijack.fixes, and finally755are applied by sd_hijack.EmbeddingsWithFixes's forward function."""756 757 758class FrozenCLIPEmbedder(nn.Layer):759 """Uses the CLIP transformer encoder for text (from huggingface)"""760 761 LAYERS = ["last", "pooled", "hidden"]762 763 def __init__(self, text_encoder, tokenizer, freeze=True, layer="last", layer_idx=None):764 super().__init__()765 assert layer in self.LAYERS766 self.tokenizer = tokenizer767 self.text_encoder = text_encoder768 if freeze:769 self.freeze()770 self.layer = layer771 self.layer_idx = layer_idx772 if layer == "hidden":773 assert layer_idx is not None774 assert 0 <= abs(layer_idx) <= 12775 776 def freeze(self):777 self.text_encoder.eval()778 for param in self.parameters():779 param.stop_gradient = False780 781 def forward(self, text):782 batch_encoding = self.tokenizer(783 text,784 truncation=True,785 max_length=self.tokenizer.model_max_length,786 padding="max_length",787 return_tensors="pd",788 )789 tokens = batch_encoding["input_ids"]790 outputs = self.text_encoder(input_ids=tokens, output_hidden_states=self.layer == "hidden", return_dict=True)791 if self.layer == "last":792 z = outputs.last_hidden_state793 elif self.layer == "pooled":794 z = outputs.pooler_output[:, None, :]795 else:796 z = outputs.hidden_states[self.layer_idx]797 return z798 799 def encode(self, text):800 return self(text)801 802 803class FrozenCLIPEmbedderWithCustomWordsBase(nn.Layer):804 """A pytorch module that is a wrapper for FrozenCLIPEmbedder module. it enhances FrozenCLIPEmbedder, making it possible to805 have unlimited prompt length and assign weights to tokens in prompt.806 """807 808 def __init__(self, wrapped, hijack):809 super().__init__()810 811 self.wrapped = wrapped812 """Original FrozenCLIPEmbedder module; can also be FrozenOpenCLIPEmbedder or xlmr.BertSeriesModelWithTransformation,813 depending on model."""814 815 self.hijack = hijack816 self.chunk_length = 75817 818 def empty_chunk(self):819 """creates an empty PromptChunk and returns it"""820 821 chunk = PromptChunk()822 chunk.tokens = [self.id_start] + [self.id_end] * (self.chunk_length + 1)823 chunk.multipliers = [1.0] * (self.chunk_length + 2)824 return chunk825 826 def get_target_prompt_token_count(self, token_count):827 """returns the maximum number of tokens a prompt of a known length can have before it requires one more PromptChunk to be represented"""828 829 return math.ceil(max(token_count, 1) / self.chunk_length) * self.chunk_length830 831 def tokenize(self, texts):832 """Converts a batch of texts into a batch of token ids"""833 834 raise NotImplementedError835 836 def encode_with_text_encoder(self, tokens):837 """838 converts a batch of token ids (in python lists) into a single tensor with numeric respresentation of those tokens;839 All python lists with tokens are assumed to have same length, usually 77.840 if input is a list with B elements and each element has T tokens, expected output shape is (B, T, C), where C depends on841 model - can be 768 and 1024.842 Among other things, this call will read self.hijack.fixes, apply it to its inputs, and clear it (setting it to None).843 """844 845 raise NotImplementedError846 847 def encode_embedding_init_text(self, init_text, nvpt):848 """Converts text into a tensor with this text's tokens' embeddings. Note that those are embeddings before they are passed through849 transformers. nvpt is used as a maximum length in tokens. If text produces less teokens than nvpt, only this many is returned."""850 851 raise NotImplementedError852 853 def tokenize_line(self, line):854 """855 this transforms a single prompt into a list of PromptChunk objects - as many as needed to856 represent the prompt.857 Returns the list and the total number of tokens in the prompt.858 """859 860 if WebUIStableDiffusionPipeline.enable_emphasis:861 parsed = parse_prompt_attention(line)862 else:863 parsed = [[line, 1.0]]864 865 tokenized = self.tokenize([text for text, _ in parsed])866 867 chunks = []868 chunk = PromptChunk()869 token_count = 0870 last_comma = -1871 872 def next_chunk(is_last=False):873 """puts current chunk into the list of results and produces the next one - empty;874 if is_last is true, tokens <end-of-text> tokens at the end won't add to token_count"""875 nonlocal token_count876 nonlocal last_comma877 nonlocal chunk878 879 if is_last:880 token_count += len(chunk.tokens)881 else:882 token_count += self.chunk_length883 884 to_add = self.chunk_length - len(chunk.tokens)885 if to_add > 0:886 chunk.tokens += [self.id_end] * to_add887 chunk.multipliers += [1.0] * to_add888 889 chunk.tokens = [self.id_start] + chunk.tokens + [self.id_end]890 chunk.multipliers = [1.0] + chunk.multipliers + [1.0]891 892 last_comma = -1893 chunks.append(chunk)894 chunk = PromptChunk()895 896 for tokens, (text, weight) in zip(tokenized, parsed):897 if text == "BREAK" and weight == -1:898 next_chunk()899 continue900 901 position = 0902 while position < len(tokens):903 token = tokens[position]904 905 if token == self.comma_token:906 last_comma = len(chunk.tokens)907 908 # this is when we are at the end of alloted 75 tokens for the current chunk, and the current token is not a comma. opts.comma_padding_backtrack909 # is a setting that specifies that if there is a comma nearby, the text after the comma should be moved out of this chunk and into the next.910 elif (911 WebUIStableDiffusionPipeline.comma_padding_backtrack != 0912 and len(chunk.tokens) == self.chunk_length913 and last_comma != -1914 and len(chunk.tokens) - last_comma <= WebUIStableDiffusionPipeline.comma_padding_backtrack915 ):916 break_location = last_comma + 1917 918 reloc_tokens = chunk.tokens[break_location:]919 reloc_mults = chunk.multipliers[break_location:]920 921 chunk.tokens = chunk.tokens[:break_location]922 chunk.multipliers = chunk.multipliers[:break_location]923 924 next_chunk()925 chunk.tokens = reloc_tokens926 chunk.multipliers = reloc_mults927 928 if len(chunk.tokens) == self.chunk_length:929 next_chunk()930 931 embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(932 tokens, position933 )934 if embedding is None:935 chunk.tokens.append(token)936 chunk.multipliers.append(weight)937 position += 1938 continue939 940 emb_len = int(embedding.vec.shape[0])941 if len(chunk.tokens) + emb_len > self.chunk_length:942 next_chunk()943 944 chunk.fixes.append(PromptChunkFix(len(chunk.tokens), embedding))945 946 chunk.tokens += [0] * emb_len947 chunk.multipliers += [weight] * emb_len948 position += embedding_length_in_tokens949 950 if len(chunk.tokens) > 0 or len(chunks) == 0:951 next_chunk(is_last=True)952 953 return chunks, token_count954 955 def process_texts(self, texts):956 """957 Accepts a list of texts and calls tokenize_line() on each, with cache. Returns the list of results and maximum958 length, in tokens, of all texts.959 """960 961 token_count = 0962 963 cache = {}964 batch_chunks = []965 for line in texts:966 if line in cache:967 chunks = cache[line]968 else:969 chunks, current_token_count = self.tokenize_line(line)970 token_count = max(current_token_count, token_count)971 972 cache[line] = chunks973 974 batch_chunks.append(chunks)975 976 return batch_chunks, token_count977 978 def forward(self, texts):979 """980 Accepts an array of texts; Passes texts through transformers network to create a tensor with numerical representation of those texts.981 Returns a tensor with shape of (B, T, C), where B is length of the array; T is length, in tokens, of texts (including padding) - T will982 be a multiple of 77; and C is dimensionality of each token - for SD1 it's 768, and for SD2 it's 1024.983 An example shape returned by this function can be: (2, 77, 768).984 Webui usually sends just one text at a time through this function - the only time when texts is an array with more than one elemenet985 is when you do prompt editing: "a picture of a [cat:dog:0.4] eating ice cream"986 """987 988 batch_chunks, token_count = self.process_texts(texts)989 990 used_embeddings = {}991 chunk_count = max([len(x) for x in batch_chunks])992 993 zs = []994 for i in range(chunk_count):995 batch_chunk = [chunks[i] if i < len(chunks) else self.empty_chunk() for chunks in batch_chunks]996 997 tokens = [x.tokens for x in batch_chunk]998 multipliers = [x.multipliers for x in batch_chunk]999 self.hijack.fixes = [x.fixes for x in batch_chunk]1000 1001 for fixes in self.hijack.fixes:1002 for position, embedding in fixes:1003 used_embeddings[embedding.name] = embedding1004 1005 z = self.process_tokens(tokens, multipliers)1006 zs.append(z)1007 1008 if len(used_embeddings) > 0:1009 embeddings_list = ", ".join(1010 [f"{name} [{embedding.checksum()}]" for name, embedding in used_embeddings.items()]1011 )1012 self.hijack.comments.append(f"Used embeddings: {embeddings_list}")1013 1014 return paddle.concat(zs, axis=1)1015 1016 def process_tokens(self, remade_batch_tokens, batch_multipliers):1017 """1018 sends one single prompt chunk to be encoded by transformers neural network.1019 remade_batch_tokens is a batch of tokens - a list, where every element is a list of tokens; usually1020 there are exactly 77 tokens in the list. batch_multipliers is the same but for multipliers instead of tokens.1021 Multipliers are used to give more or less weight to the outputs of transformers network. Each multiplier1022 corresponds to one token.1023 """1024 tokens = paddle.to_tensor(remade_batch_tokens)1025 1026 # this is for SD2: SD1 uses the same token for padding and end of text, while SD2 uses different ones.1027 if self.id_end != self.id_pad:1028 for batch_pos in range(len(remade_batch_tokens)):1029 index = remade_batch_tokens[batch_pos].index(self.id_end)1030 tokens[batch_pos, index + 1 : tokens.shape[1]] = self.id_pad1031 1032 z = self.encode_with_text_encoder(tokens)1033 1034 # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise1035 batch_multipliers = paddle.to_tensor(batch_multipliers)1036 original_mean = z.mean()1037 z = z * batch_multipliers.reshape(1038 batch_multipliers.shape1039 + [1040 1,1041 ]1042 ).expand(z.shape)1043 new_mean = z.mean()1044 z = z * (original_mean / new_mean)1045 1046 return z1047 1048 1049class FrozenCLIPEmbedderWithCustomWords(FrozenCLIPEmbedderWithCustomWordsBase):1050 def __init__(self, wrapped, hijack, CLIP_stop_at_last_layers=-1):1051 super().__init__(wrapped, hijack)1052 self.CLIP_stop_at_last_layers = CLIP_stop_at_last_layers1053 self.tokenizer = wrapped.tokenizer1054 1055 vocab = self.tokenizer.get_vocab()1056 1057 self.comma_token = vocab.get(",</w>", None)1058 1059 self.token_mults = {}1060 tokens_with_parens = [(k, v) for k, v in vocab.items() if "(" in k or ")" in k or "[" in k or "]" in k]1061 for text, ident in tokens_with_parens:1062 mult = 1.01063 for c in text:1064 if c == "[":1065 mult /= 1.11066 if c == "]":1067 mult *= 1.11068 if c == "(":1069 mult *= 1.11070 if c == ")":1071 mult /= 1.11072 1073 if mult != 1.0:1074 self.token_mults[ident] = mult1075 1076 self.id_start = self.wrapped.tokenizer.bos_token_id1077 self.id_end = self.wrapped.tokenizer.eos_token_id1078 self.id_pad = self.id_end1079 1080 def tokenize(self, texts):1081 tokenized = self.wrapped.tokenizer(texts, truncation=False, add_special_tokens=False)["input_ids"]1082 1083 return tokenized1084 1085 def encode_with_text_encoder(self, tokens):1086 output_hidden_states = self.CLIP_stop_at_last_layers > 11087 outputs = self.wrapped.text_encoder(1088 input_ids=tokens, output_hidden_states=output_hidden_states, return_dict=True1089 )1090 1091 if output_hidden_states:1092 z = outputs.hidden_states[-self.CLIP_stop_at_last_layers]1093 z = self.wrapped.text_encoder.text_model.ln_final(z)1094 else:1095 z = outputs.last_hidden_state1096 1097 return z1098 1099 def encode_embedding_init_text(self, init_text, nvpt):1100 embedding_layer = self.wrapped.text_encoder.text_model1101 ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pd", add_special_tokens=False)[1102 "input_ids"1103 ]1104 embedded = embedding_layer.token_embedding.wrapped(ids).squeeze(0)1105 1106 return embedded1107 1108 1109# extra_networks.py1110import re1111from collections import defaultdict1112 1113 1114class ExtraNetworkParams:1115 def __init__(self, items=None):1116 self.items = items or []1117 1118 1119re_extra_net = re.compile(r"<(\w+):([^>]+)>")1120 1121 1122def parse_prompt(prompt):1123 res = defaultdict(list)1124 1125 def found(m):1126 name = m.group(1)1127 args = m.group(2)1128 1129 res[name].append(ExtraNetworkParams(items=args.split(":")))1130 1131 return ""1132 1133 prompt = re.sub(re_extra_net, found, prompt)1134 1135 return prompt, res1136 1137 1138def parse_prompts(prompts):1139 res = []1140 extra_data = None1141 1142 for prompt in prompts:1143 updated_prompt, parsed_extra_data = parse_prompt(prompt)1144 1145 if extra_data is None:1146 extra_data = parsed_extra_data1147 1148 res.append(updated_prompt)1149 1150 return res, extra_data1151 1152 1153# image_embeddings.py1154 1155import base641156import json1157import zlib1158 1159import numpy as np1160from PIL import Image1161 1162 1163class EmbeddingDecoder(json.JSONDecoder):1164 def __init__(self, *args, **kwargs):1165 json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)1166 1167 def object_hook(self, d):1168 if "TORCHTENSOR" in d:1169 return paddle.to_tensor(np.array(d["TORCHTENSOR"]))1170 return d1171 1172 1173def embedding_from_b64(data):1174 d = base64.b64decode(data)1175 return json.loads(d, cls=EmbeddingDecoder)1176 1177 1178def lcg(m=2**32, a=1664525, c=1013904223, seed=0):1179 while True:1180 seed = (a * seed + c) % m1181 yield seed % 2551182 1183 1184def xor_block(block):1185 g = lcg()1186 randblock = np.array([next(g) for _ in range(np.product(block.shape))]).astype(np.uint8).reshape(block.shape)1187 return np.bitwise_xor(block.astype(np.uint8), randblock & 0x0F)1188 1189 1190def crop_black(img, tol=0):1191 mask = (img > tol).all(2)1192 mask0, mask1 = mask.any(0), mask.any(1)1193 col_start, col_end = mask0.argmax(), mask.shape[1] - mask0[::-1].argmax()1194 row_start, row_end = mask1.argmax(), mask.shape[0] - mask1[::-1].argmax()1195 return img[row_start:row_end, col_start:col_end]1196 1197 1198def extract_image_data_embed(image):1199 d = 31200 outarr = (