| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import inspect |
| import json |
| import os |
| import re |
| from typing import Any, Callable, Dict, List, Optional, Union |
|
|
| import safetensors.torch |
| import torch |
| from huggingface_hub import hf_hub_download |
| from huggingface_hub.utils import EntryNotFoundError |
| from packaging import version |
| from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer |
|
|
| from diffusers.configuration_utils import FrozenDict |
| from diffusers.image_processor import VaeImageProcessor |
| from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin |
| from diffusers.models import AutoencoderKL, UNet2DConditionModel |
| from diffusers.models.lora import adjust_lora_scale_text_encoder |
| from diffusers.schedulers import KarrasDiffusionSchedulers |
| from diffusers.utils import deprecate, logging, replace_example_docstring, scale_lora_layers, unscale_lora_layers |
| from diffusers.utils.torch_utils import randn_tensor |
| from diffusers.pipelines.pipeline_utils import DiffusionPipeline |
| from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput |
| from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker |
|
|
|
|
| logger = logging.get_logger(__name__) |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, input_size, hidden_size, num_heads, dropout=0.1): |
| super(TransformerBlock, self).__init__() |
| self.attention = nn.MultiheadAttention(input_size, num_heads, dropout=dropout) |
| self.norm1 = nn.LayerNorm(input_size) |
| self.ffn = nn.Sequential( |
| nn.Linear(input_size, hidden_size), |
| nn.ReLU(), |
| nn.Linear(hidden_size, input_size) |
| ) |
| self.norm2 = nn.LayerNorm(input_size) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| |
| attn_output, _ = self.attention(x, x, x) |
| |
| x = x + self.dropout(self.norm1(attn_output)) |
| |
| |
| ffn_output = self.ffn(x) |
| |
| x = x + self.dropout(self.norm2(ffn_output)) |
| |
| return x |
|
|
| class AdaptiveTokenMapping_v1(nn.Module): |
| NUM_HEADS = 6 |
| NUM_LAYERS = 4 |
|
|
| def __init__(self, input_size, hidden_size, output_size, dropout=0.1): |
| super(AdaptiveTokenMapping_v1, self).__init__() |
| self.embedding = nn.Linear(input_size, hidden_size) |
| self.transformer_blocks = nn.ModuleList([ |
| TransformerBlock(hidden_size, hidden_size, self.NUM_HEADS, dropout) |
| for _ in range(self.NUM_LAYERS) |
| ]) |
| self.output_layer = nn.Linear(hidden_size, output_size) |
|
|
| @property |
| def device(self): |
| return self.output_layer.weight.device |
|
|
| @property |
| def dtype(self): |
| return self.output_layer.weight.dtype |
|
|
| def forward(self, x): |
| x = self.embedding(x) |
| for transformer_block in self.transformer_blocks: |
| x = transformer_block(x) |
| x = self.output_layer(x) |
| return x |
|
|
|
|
| AITTI_REPO_ID = "itsmag11/AITTI" |
|
|
|
|
| def load_adaptive_mapping_state_dict(path): |
| try: |
| return safetensors.torch.load_file(path) |
| except Exception: |
| |
| with open(path, "rb") as f: |
| return torch.load(f, map_location="cpu", weights_only=True) |
|
|
|
|
| EXAMPLE_DOC_STRING = """ |
| Examples: |
| ```py |
| >>> import torch |
| >>> from diffusers import DiffusionPipeline |
| |
| >>> pipe = DiffusionPipeline.from_pretrained( |
| ... "stable-diffusion-v1-5/stable-diffusion-v1-5", |
| ... custom_pipeline="itsmag11/AITTI", |
| ... torch_dtype=torch.float16, |
| ... ).to("cuda") |
| >>> pipe.load_aitti("gender") |
| |
| >>> image = pipe("A photo of a <gender-inclusive> doctor").images[0] |
| ``` |
| """ |
|
|
|
|
| def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): |
| """ |
| Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and |
| Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 |
| """ |
| 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) |
| |
| noise_pred_rescaled = noise_cfg * (std_text / std_cfg) |
| |
| noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg |
| return noise_cfg |
|
|
|
|
| class StableDiffusionAdaptiveTokenPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin): |
| r""" |
| Pipeline for text-to-image generation using Stable Diffusion. |
| |
| This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods |
| implemented for all pipelines (downloading, saving, running on a particular device, etc.). |
| |
| The pipeline also inherits the following loading methods: |
| - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings |
| - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights |
| - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights |
| - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files |
| |
| 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 ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). |
| tokenizer ([`~transformers.CLIPTokenizer`]): |
| A `CLIPTokenizer` to tokenize text. |
| unet ([`UNet2DConditionModel`]): |
| A `UNet2DConditionModel` to denoise the encoded image latents. |
| scheduler ([`SchedulerMixin`]): |
| A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of |
| [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. |
| 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/runwayml/stable-diffusion-v1-5) for more details |
| about a model's potential harms. |
| feature_extractor ([`~transformers.CLIPImageProcessor`]): |
| A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. |
| """ |
| model_cpu_offload_seq = "text_encoder->unet->vae" |
| _optional_components = ["safety_checker", "feature_extractor", "adaptive_mapping"] |
| _exclude_from_cpu_offload = ["safety_checker"] |
|
|
| def __init__( |
| self, |
| vae: AutoencoderKL, |
| text_encoder: CLIPTextModel, |
| tokenizer: CLIPTokenizer, |
| unet: UNet2DConditionModel, |
| scheduler: KarrasDiffusionSchedulers, |
| safety_checker: StableDiffusionSafetyChecker, |
| feature_extractor: CLIPImageProcessor, |
| adaptive_mapping: Optional[AdaptiveTokenMapping_v1] = None, |
| requires_safety_checker: bool = True, |
| ): |
| super().__init__() |
|
|
| if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 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) |
|
|
| if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: |
| 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) |
|
|
| if safety_checker is None and requires_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 ." |
| ) |
|
|
| if safety_checker is not None and feature_extractor is None: |
| raise ValueError( |
| "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 = hasattr(unet.config, "_diffusers_version") and version.parse( |
| version.parse(unet.config._diffusers_version).base_version |
| ) < version.parse("0.9.0.dev0") |
| is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 |
| if is_unet_version_less_0_9_0 and is_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, |
| adaptive_mapping=adaptive_mapping, |
| ) |
| self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) |
| self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) |
| self.use_peft_backend = False |
| self.register_to_config(requires_safety_checker=requires_safety_checker) |
| self.aitti_token = None |
| self.aitti_change_step = 0 |
|
|
| def load_aitti( |
| self, |
| attribute: Optional[str] = None, |
| pretrained_model_name_or_path: str = AITTI_REPO_ID, |
| token: Optional[str] = None, |
| adaptive_mapping_weight_name: str = "adaptive_mapping.safetensors", |
| learned_embeds_weight_name: str = "learned_embeds.safetensors", |
| **hub_kwargs, |
| ) -> str: |
| r""" |
| Load an AITTI inclusive token (adaptive token mapping + learned token embedding) and return the token string. |
| |
| Examples: |
| >>> pipe.load_aitti("gender") # "<gender-inclusive>" from https://huggingface.co/itsmag11/AITTI |
| >>> pipe.load_aitti(pretrained_model_name_or_path="./gender-inclusive", token="<gender-diverse>") # local training output |
| |
| Args: |
| attribute (`str`, *optional*): |
| Sub-folder holding the weights, e.g. `"gender"`, `"race"` or `"age"` for the released tokens. |
| pretrained_model_name_or_path (`str`, defaults to `"itsmag11/AITTI"`): |
| Hugging Face Hub repo id or local directory. |
| token (`str`, *optional*): |
| Token string to register. Defaults to the one in `config.json` or in the learned embedding file. |
| hub_kwargs: |
| Extra arguments forwarded to `huggingface_hub.hf_hub_download` (e.g. `revision`, `cache_dir`). |
| """ |
|
|
| def resolve(filename): |
| if os.path.isdir(pretrained_model_name_or_path): |
| folder = pretrained_model_name_or_path |
| if attribute is not None and os.path.isdir(os.path.join(folder, attribute)): |
| folder = os.path.join(folder, attribute) |
| path = os.path.join(folder, filename) |
| return path if os.path.isfile(path) else None |
| subpath = filename if attribute is None else f"{attribute}/{filename}" |
| try: |
| return hf_hub_download(pretrained_model_name_or_path, subpath, **hub_kwargs) |
| except EntryNotFoundError: |
| return None |
|
|
| config = {} |
| config_path = resolve("config.json") |
| if config_path is not None: |
| with open(config_path) as f: |
| config = json.load(f) |
|
|
| adaptive_mapping_path = resolve(adaptive_mapping_weight_name) |
| learned_embeds_path = resolve(learned_embeds_weight_name) |
| if adaptive_mapping_path is None or learned_embeds_path is None: |
| raise FileNotFoundError( |
| f"Could not find {adaptive_mapping_weight_name} and {learned_embeds_weight_name} in " |
| f"{pretrained_model_name_or_path}" + (f"/{attribute}" if attribute else "") |
| ) |
|
|
| embed_dim = self.text_encoder.config.hidden_size |
| adaptive_mapping = AdaptiveTokenMapping_v1(embed_dim, embed_dim, embed_dim) |
| adaptive_mapping.load_state_dict(load_adaptive_mapping_state_dict(adaptive_mapping_path), strict=True) |
| adaptive_mapping.requires_grad_(False).eval() |
| adaptive_mapping.to(device=self.text_encoder.device, dtype=self.text_encoder.dtype) |
| self.register_modules(adaptive_mapping=adaptive_mapping) |
|
|
| learned_embeds = safetensors.torch.load_file(learned_embeds_path) |
| token = token or config.get("token") or next(iter(learned_embeds)) |
| if token in self.tokenizer.get_vocab(): |
| self.unload_textual_inversion(token) |
| self.load_textual_inversion(learned_embeds, token=token) |
| self.aitti_token = token |
| self.aitti_change_step = config.get("change_step", 0) |
| return token |
|
|
| @staticmethod |
| def _infer_profession_name(prompt: str, token_name: str) -> str: |
| words = prompt.split(" ") |
| start = words.index(token_name) + 1 |
| profession_words = [] |
| for word in words[start:]: |
| if not re.fullmatch(r"[A-Za-z\-]+", word): |
| break |
| profession_words.append(word) |
| if not profession_words: |
| raise ValueError( |
| f"Could not infer `profession_name` from prompt {prompt!r}; please pass it explicitly, " |
| "e.g. pipe('A photo of a <gender-inclusive> doctor', profession_name='doctor')." |
| ) |
| return " ".join(profession_words) |
|
|
| def enable_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() |
|
|
| def disable_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() |
|
|
| def enable_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() |
|
|
| def disable_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() |
|
|
| def _encode_prompt( |
| self, |
| prompt, |
| device, |
| num_images_per_prompt, |
| do_classifier_free_guidance, |
| negative_prompt=None, |
| prompt_embeds: Optional[torch.FloatTensor] = None, |
| negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
| lora_scale: Optional[float] = None, |
| **kwargs, |
| ): |
| deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." |
| deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False) |
|
|
| prompt_embeds_tuple = self.encode_prompt( |
| prompt=prompt, |
| device=device, |
| num_images_per_prompt=num_images_per_prompt, |
| do_classifier_free_guidance=do_classifier_free_guidance, |
| negative_prompt=negative_prompt, |
| prompt_embeds=prompt_embeds, |
| negative_prompt_embeds=negative_prompt_embeds, |
| lora_scale=lora_scale, |
| **kwargs, |
| ) |
|
|
| |
| prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]]) |
|
|
| return prompt_embeds |
| |
| @property |
| def clip_text_model(self): |
| |
| return getattr(self.text_encoder, "text_model", self.text_encoder) |
|
|
| def get_lookup_embeddings(self, input_ids, inputs_embeds=None): |
| if input_ids is None: |
| raise ValueError("You have to specify input_ids") |
|
|
| input_shape = input_ids.size() |
| input_ids = input_ids.view(-1, input_shape[-1]) |
|
|
| if inputs_embeds is None: |
| inputs_embeds = self.clip_text_model.embeddings.token_embedding(input_ids) |
|
|
| return inputs_embeds |
| |
| def add_position_embeddings(self, input_ids, inputs_embeds, position_ids=None): |
| seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] |
| if position_ids is None: |
| position_ids = self.clip_text_model.embeddings.position_ids[:, :seq_length] |
|
|
| position_embeddings = self.clip_text_model.embeddings.position_embedding(position_ids) |
| embeddings = inputs_embeds + position_embeddings |
| return embeddings |
| |
| def get_text_embeddings(self, input_ids, hidden_states, |
| attention_mask = None, |
| output_attentions = None, |
| output_hidden_states = None, |
| return_dict = None): |
| input_shape = input_ids.size() |
| text_model = self.clip_text_model |
|
|
| if "causal_attention_mask" in inspect.signature(text_model.encoder.forward).parameters: |
| from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask |
| |
| |
| causal_attention_mask = _create_4d_causal_attention_mask( |
| input_shape, hidden_states.dtype, device=hidden_states.device |
| ) |
| |
| if attention_mask is not None: |
| |
| attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) |
|
|
| encoder_outputs = text_model.encoder( |
| inputs_embeds=hidden_states, |
| attention_mask=attention_mask, |
| causal_attention_mask=causal_attention_mask, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| ) |
| else: |
| from transformers.masking_utils import create_causal_mask |
|
|
| attention_mask = create_causal_mask( |
| config=text_model.config, |
| inputs_embeds=hidden_states, |
| attention_mask=attention_mask, |
| past_key_values=None, |
| ) |
| encoder_outputs = text_model.encoder( |
| inputs_embeds=hidden_states, |
| attention_mask=attention_mask, |
| is_causal=True, |
| ) |
|
|
| last_hidden_state = encoder_outputs[0] |
| last_hidden_state = text_model.final_layer_norm(last_hidden_state) |
|
|
| return last_hidden_state |
|
|
| def encode_prompt( |
| self, |
| prompt, |
| device, |
| num_images_per_prompt, |
| do_classifier_free_guidance, |
| negative_prompt=None, |
| prompt_embeds: Optional[torch.FloatTensor] = None, |
| negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
| lora_scale: Optional[float] = None, |
| clip_skip: Optional[int] = None, |
| profession_name: Optional[str] = None, |
| token_name: Optional[str] = None, |
| ): |
| r""" |
| Encodes the prompt into text encoder hidden states. |
| |
| Args: |
| prompt (`str` or `List[str]`, *optional*): |
| prompt to be encoded |
| device: (`torch.device`): |
| torch device |
| num_images_per_prompt (`int`): |
| number of images that should be generated per prompt |
| do_classifier_free_guidance (`bool`): |
| whether to use classifier free guidance or not |
| 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`). |
| prompt_embeds (`torch.FloatTensor`, *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.FloatTensor`, *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. |
| """ |
| |
| |
| ori_prompt = prompt.replace(f'{token_name} ', '') |
| inclusive_prompt = prompt |
| prompt = [inclusive_prompt, ori_prompt] |
| |
| print(f'prompt pair: {prompt}', flush=True) |
| |
| |
|
|
| |
| |
| if lora_scale is not None and isinstance(self, LoraLoaderMixin): |
| self._lora_scale = lora_scale |
|
|
| |
| if not self.use_peft_backend: |
| adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) |
| else: |
| scale_lora_layers(self.text_encoder, lora_scale) |
|
|
| if prompt is not None and isinstance(prompt, str): |
| batch_size = 1 |
| elif prompt is not None and isinstance(prompt, list): |
| batch_size = len(prompt) |
| else: |
| batch_size = prompt_embeds.shape[0] |
| |
| if prompt_embeds is None: |
| print(f'prompt before: {prompt}') |
| |
| if isinstance(self, TextualInversionLoaderMixin): |
| prompt = self.maybe_convert_prompt(prompt, self.tokenizer) |
|
|
| text_inputs = self.tokenizer( |
| prompt, |
| padding="max_length", |
| max_length=self.tokenizer.model_max_length, |
| truncation=True, |
| return_tensors="pt", |
| ) |
| text_input_ids = text_inputs.input_ids |
| untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids |
|
|
| |
| |
| placeholders_mask = torch.zeros_like(text_input_ids[0]).bool() |
| placeholders_mask[prompt[0].split(' ').index(token_name) + 1] = True |
|
|
| concept_mask = torch.zeros_like(text_input_ids[0]).bool() |
| for profession_word in profession_name.split(' '): |
| concept_mask[prompt[0].split(' ').index(profession_word) + 1] = True |
| |
|
|
| if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.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") and self.text_encoder.config.use_attention_mask: |
| attention_mask = text_inputs.attention_mask.to(device) |
| else: |
| attention_mask = None |
|
|
| if self.text_encoder is not None: |
| prompt_embeds_dtype = self.text_encoder.dtype |
| elif self.unet is not None: |
| prompt_embeds_dtype = self.unet.dtype |
| else: |
| prompt_embeds_dtype = prompt_embeds.dtype |
|
|
| if clip_skip is None: |
| prompt_embeds_ori = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask) |
| prompt_embeds_ori = prompt_embeds_ori[0] |
|
|
| token_embeddings = self.get_lookup_embeddings(text_input_ids.to(device)) |
|
|
| print(f'token_embeddings: {token_embeddings.shape}') |
| |
|
|
| |
| |
|
|
| token_embeddings_ori = token_embeddings.clone() |
|
|
| print(f'=== adaptive token mapping ===', flush=True) |
|
|
| placeholders_token_embeddings = token_embeddings[0][placeholders_mask] |
| concept_token_embeddings = token_embeddings[0][concept_mask] |
| input_token_embeddings = torch.cat([placeholders_token_embeddings, concept_token_embeddings], dim=0) |
| adaptive_output = self.adaptive_mapping(input_token_embeddings) |
| placeholders_num = placeholders_token_embeddings.shape[0] |
| adaptive_tokens = adaptive_output[:placeholders_num].to(dtype=prompt_embeds_dtype, device=device) |
| token_embeddings[0][placeholders_mask] = adaptive_tokens.unsqueeze(0) |
|
|
| assert (token_embeddings[0][placeholders_mask] == token_embeddings_ori[0][placeholders_mask]).all() == False |
| assert (token_embeddings[0][~placeholders_mask] == token_embeddings_ori[0][~placeholders_mask]).all() == True |
| assert (token_embeddings[1] == token_embeddings_ori[1]).all() == True |
|
|
| |
| |
| token_embeddings_w_pos = self.add_position_embeddings(text_input_ids.to(device), token_embeddings) |
| prompt_embeds = self.get_text_embeddings(text_input_ids.to(device), token_embeddings_w_pos, attention_mask=attention_mask) |
|
|
| else: |
| prompt_embeds = self.text_encoder( |
| text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True |
| ) |
| |
| |
| |
| prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)] |
| |
| |
| |
| |
| prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds) |
|
|
| prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) |
|
|
| bs_embed, seq_len, _ = prompt_embeds.shape |
| |
| prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) |
| prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) |
|
|
| |
| if do_classifier_free_guidance and negative_prompt_embeds is None: |
| uncond_tokens: List[str] |
| if negative_prompt is None: |
| uncond_tokens = [""] * batch_size |
| elif prompt is not None and type(prompt) is not type(negative_prompt): |
| raise TypeError( |
| f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" |
| f" {type(prompt)}." |
| ) |
| elif isinstance(negative_prompt, str): |
| uncond_tokens = [negative_prompt] |
| elif batch_size != len(negative_prompt): |
| raise ValueError( |
| f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" |
| f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" |
| " the batch size of `prompt`." |
| ) |
| else: |
| uncond_tokens = negative_prompt |
|
|
| |
| if isinstance(self, TextualInversionLoaderMixin): |
| uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer) |
|
|
| max_length = prompt_embeds.shape[1] |
| uncond_input = self.tokenizer( |
| uncond_tokens, |
| padding="max_length", |
| max_length=max_length, |
| truncation=True, |
| return_tensors="pt", |
| ) |
|
|
| if hasattr(self.text_encoder.config, "use_attention_mask") and self.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] |
|
|
| if do_classifier_free_guidance: |
| |
| 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) |
|
|
| if isinstance(self, LoraLoaderMixin) and self.use_peft_backend: |
| |
| unscale_lora_layers(self.text_encoder) |
|
|
| return prompt_embeds, negative_prompt_embeds |
|
|
| def run_safety_checker(self, image, device, dtype): |
| if self.safety_checker is None: |
| has_nsfw_concept = None |
| else: |
| if torch.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) |
| ) |
| return image, has_nsfw_concept |
|
|
| def decode_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) |
| |
| image = image.cpu().permute(0, 2, 3, 1).float().numpy() |
| return image |
|
|
| def prepare_extra_step_kwargs(self, generator, eta): |
| |
| |
| |
| |
|
|
| accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) |
| extra_step_kwargs = {} |
| if accepts_eta: |
| extra_step_kwargs["eta"] = eta |
|
|
| |
| accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) |
| if accepts_generator: |
| extra_step_kwargs["generator"] = generator |
| return extra_step_kwargs |
|
|
| def check_inputs( |
| self, |
| prompt, |
| height, |
| width, |
| callback_steps, |
| negative_prompt=None, |
| prompt_embeds=None, |
| negative_prompt_embeds=None, |
| ): |
| if height % 8 != 0 or width % 8 != 0: |
| raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") |
|
|
| if (callback_steps is None) or ( |
| callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) |
| ): |
| raise ValueError( |
| f"`callback_steps` has to be a positive integer but is {callback_steps} of type" |
| f" {type(callback_steps)}." |
| ) |
|
|
| if prompt is not None and prompt_embeds is not None: |
| raise ValueError( |
| f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" |
| " only forward one of the two." |
| ) |
| elif prompt is None and prompt_embeds is None: |
| raise ValueError( |
| "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." |
| ) |
| elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): |
| raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") |
|
|
| if negative_prompt is not None and negative_prompt_embeds is not None: |
| raise ValueError( |
| 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." |
| ) |
|
|
| if prompt_embeds is not None and negative_prompt_embeds is not None: |
| if prompt_embeds.shape != negative_prompt_embeds.shape: |
| raise ValueError( |
| "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" |
| f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" |
| f" {negative_prompt_embeds.shape}." |
| ) |
|
|
| def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): |
| shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) |
| if isinstance(generator, list) and len(generator) != batch_size: |
| raise ValueError( |
| f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" |
| f" size of {batch_size}. Make sure the batch size matches the length of the generators." |
| ) |
|
|
| if latents is None: |
| latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) |
| else: |
| latents = latents.to(device) |
|
|
| |
| latents = latents * self.scheduler.init_noise_sigma |
| return latents |
|
|
| def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): |
| r"""Enables the FreeU mechanism as in https://arxiv.org/abs/2309.11497. |
| |
| The suffixes after the scaling factors represent the stages where they are being applied. |
| |
| Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of the values |
| that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. |
| |
| Args: |
| s1 (`float`): |
| Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to |
| mitigate "oversmoothing effect" in the enhanced denoising process. |
| s2 (`float`): |
| Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to |
| mitigate "oversmoothing effect" in the enhanced denoising process. |
| b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. |
| b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. |
| """ |
| if not hasattr(self, "unet"): |
| raise ValueError("The pipeline must have `unet` for using FreeU.") |
| self.unet.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2) |
|
|
| def disable_freeu(self): |
| """Disables the FreeU mechanism if enabled.""" |
| self.unet.disable_freeu() |
|
|
| @torch.no_grad() |
| @replace_example_docstring(EXAMPLE_DOC_STRING) |
| def __call__( |
| self, |
| prompt: Union[str, List[str]] = None, |
| height: Optional[int] = None, |
| width: Optional[int] = None, |
| num_inference_steps: int = 50, |
| guidance_scale: float = 7.5, |
| negative_prompt: Optional[Union[str, List[str]]] = None, |
| num_images_per_prompt: Optional[int] = 1, |
| eta: float = 0.0, |
| generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, |
| latents: Optional[torch.FloatTensor] = None, |
| prompt_embeds: Optional[torch.FloatTensor] = None, |
| negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
| output_type: Optional[str] = "pil", |
| return_dict: bool = True, |
| callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, |
| callback_steps: int = 1, |
| cross_attention_kwargs: Optional[Dict[str, Any]] = None, |
| guidance_rescale: float = 0.0, |
| clip_skip: Optional[int] = None, |
| profession_name: Optional[str] = None, |
| token_name: Optional[str] = None, |
| change_step: Optional[int] = None, |
| ): |
| r""" |
| The call function to the pipeline for generation. |
| |
| Args: |
| prompt (`str` or `List[str]`, *optional*): |
| The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. |
| height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): |
| The height in pixels of the generated image. |
| width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): |
| The width in pixels of the generated image. |
| num_inference_steps (`int`, *optional*, defaults to 50): |
| The number of denoising steps. More denoising steps usually lead to a higher quality image at the |
| expense of slower inference. |
| guidance_scale (`float`, *optional*, defaults to 7.5): |
| A higher guidance scale value encourages the model to generate images closely linked to the text |
| `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. |
| negative_prompt (`str` or `List[str]`, *optional*): |
| The prompt or prompts to guide what to not include in image generation. If not defined, you need to |
| pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). |
| num_images_per_prompt (`int`, *optional*, defaults to 1): |
| The number of images to generate per prompt. |
| eta (`float`, *optional*, defaults to 0.0): |
| Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies |
| to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. |
| generator (`torch.Generator` or `List[torch.Generator]`, *optional*): |
| A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make |
| generation deterministic. |
| latents (`torch.FloatTensor`, *optional*): |
| Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image |
| generation. Can be used to tweak the same generation with different prompts. If not provided, a latents |
| tensor is generated by sampling using the supplied random `generator`. |
| prompt_embeds (`torch.FloatTensor`, *optional*): |
| Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not |
| provided, text embeddings are generated from the `prompt` input argument. |
| negative_prompt_embeds (`torch.FloatTensor`, *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. |
| output_type (`str`, *optional*, defaults to `"pil"`): |
| The output format of the generated image. Choose between `PIL.Image` or `np.array`. |
| return_dict (`bool`, *optional*, defaults to `True`): |
| Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a |
| plain tuple. |
| callback (`Callable`, *optional*): |
| A function that calls every `callback_steps` steps during inference. The function is called with the |
| following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. |
| callback_steps (`int`, *optional*, defaults to 1): |
| The frequency at which the `callback` function is called. If not specified, the callback is called at |
| every step. |
| 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. |
| profession_name (`str`, *optional*): |
| The concept word(s) following the inclusive token, e.g. `"doctor"`. Inferred from the prompt if omitted. |
| token_name (`str`, *optional*): |
| The inclusive token in the prompt. Defaults to the token loaded by `load_aitti`. |
| change_step (`int`, *optional*): |
| Number of initial denoising steps that use the original prompt (without the inclusive token). |
| Defaults to the value in the loaded token's `config.json`, otherwise 0. |
| |
| Examples: |
| |
| Returns: |
| [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: |
| If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, |
| otherwise a `tuple` is returned where the first element is a list with the generated images and the |
| second element is a list of `bool`s indicating whether the corresponding generated image contains |
| "not-safe-for-work" (nsfw) content. |
| """ |
| |
| height = height or self.unet.config.sample_size * self.vae_scale_factor |
| width = width or self.unet.config.sample_size * self.vae_scale_factor |
|
|
| |
| self.check_inputs( |
| prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds |
| ) |
|
|
| |
| if prompt is not None and isinstance(prompt, str): |
| batch_size = 1 |
| elif prompt is not None and isinstance(prompt, list): |
| batch_size = len(prompt) |
| else: |
| batch_size = prompt_embeds.shape[0] |
|
|
| token_name = token_name or self.aitti_token |
| if token_name is None: |
| raise ValueError("Call `pipe.load_aitti(...)` first or pass `token_name`.") |
| if self.adaptive_mapping is None: |
| raise ValueError("No adaptive mapping loaded. Call `pipe.load_aitti(...)` first.") |
| if profession_name is None: |
| profession_name = self._infer_profession_name(prompt, token_name) |
| if change_step is None: |
| change_step = self.aitti_change_step |
|
|
| device = self._execution_device |
| |
| |
| |
| do_classifier_free_guidance = guidance_scale > 1.0 |
|
|
| |
| text_encoder_lora_scale = ( |
| cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None |
| ) |
| prompt_embeds, negative_prompt_embeds = self.encode_prompt( |
| prompt, |
| device, |
| num_images_per_prompt, |
| do_classifier_free_guidance, |
| negative_prompt, |
| prompt_embeds=prompt_embeds, |
| negative_prompt_embeds=negative_prompt_embeds, |
| lora_scale=text_encoder_lora_scale, |
| clip_skip=clip_skip, |
| profession_name=profession_name, |
| token_name=token_name |
| ) |
| |
| |
| |
| if do_classifier_free_guidance: |
| |
| inclusive_prompt_embeds = torch.cat([negative_prompt_embeds[0].unsqueeze(0), prompt_embeds[0].unsqueeze(0)]) |
| ori_prompt_embeds = torch.cat([negative_prompt_embeds[1].unsqueeze(0), prompt_embeds[1].unsqueeze(0)]) |
| print(f'inclusive_prompt_embeds: {inclusive_prompt_embeds.shape}') |
| print(f'ori_prompt_embeds: {ori_prompt_embeds.shape}') |
| |
|
|
| |
| self.scheduler.set_timesteps(num_inference_steps, device=device) |
| timesteps = self.scheduler.timesteps |
|
|
| |
| num_channels_latents = self.unet.config.in_channels |
| latents = self.prepare_latents( |
| batch_size * num_images_per_prompt, |
| num_channels_latents, |
| height, |
| width, |
| inclusive_prompt_embeds.dtype, |
| device, |
| generator, |
| latents, |
| ) |
|
|
| |
| extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) |
|
|
| |
| num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order |
| with self.progress_bar(total=num_inference_steps) as progress_bar: |
| for i, t in enumerate(timesteps): |
| prompt_embeds = ori_prompt_embeds if i < change_step else inclusive_prompt_embeds |
| |
| latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents |
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) |
|
|
| |
| noise_pred = self.unet( |
| latent_model_input, |
| t, |
| encoder_hidden_states=prompt_embeds, |
| cross_attention_kwargs=cross_attention_kwargs, |
| return_dict=False, |
| )[0] |
|
|
| |
| if do_classifier_free_guidance: |
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) |
| noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) |
|
|
| if do_classifier_free_guidance and guidance_rescale > 0.0: |
| |
| noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) |
|
|
| |
| latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] |
|
|
| |
| if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): |
| progress_bar.update() |
| if callback is not None and i % callback_steps == 0: |
| step_idx = i // getattr(self.scheduler, "order", 1) |
| callback(step_idx, t, latents) |
|
|
| if not output_type == "latent": |
| image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] |
| image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) |
| else: |
| image = latents |
| has_nsfw_concept = None |
|
|
| if has_nsfw_concept is None: |
| do_denormalize = [True] * image.shape[0] |
| else: |
| do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] |
|
|
| image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) |
|
|
| |
| self.maybe_free_model_hooks() |
|
|
| if not return_dict: |
| return (image, has_nsfw_concept) |
|
|
| return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) |
|
|
