






import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device





def 						A_					(						lowercase_			)				->List[str]:







 """simple docstring"""

 SCREAMING_SNAKE_CASE				    =		{}
 for name in unet.attn_processors.keys():
  SCREAMING_SNAKE_CASE				    =		None if name.endswith('attn1.processor'			) else unet.config.cross_attention_dim
  if name.startswith('mid_block'			):
   SCREAMING_SNAKE_CASE				    =		unet.config.block_out_channels[-1]
  elif name.startswith('up_blocks'			):
   SCREAMING_SNAKE_CASE				    =		int(name[len('up_blocks.'			)]			)
   SCREAMING_SNAKE_CASE				    =		list(reversed(unet.config.block_out_channels			)			)[block_id]
  elif name.startswith('down_blocks'			):
   SCREAMING_SNAKE_CASE				    =		int(name[len('down_blocks.'			)]			)
   SCREAMING_SNAKE_CASE				    =		unet.config.block_out_channels[block_id]
  SCREAMING_SNAKE_CASE				    =		(
      LoRAAttnProcessor2_0 if hasattr(lowercase_	,  'scaled_dot_product_attention'			) else LoRAAttnProcessor
  )
  SCREAMING_SNAKE_CASE				    =		lora_attn_processor_class(
      hidden_size=lowercase_	,  cross_attention_dim=lowercase_			)
 SCREAMING_SNAKE_CASE				    =		AttnProcsLayers(lowercase_			)
 return lora_attn_procs, unet_lora_layers





def 						A_					(						lowercase_			)				->Union[str, Any]:







 """simple docstring"""

 SCREAMING_SNAKE_CASE				    =		{}
 SCREAMING_SNAKE_CASE				    =		(
     LoRAAttnProcessor2_0 if hasattr(lowercase_	,  'scaled_dot_product_attention'			) else LoRAAttnProcessor
 )
 for name, module in text_encoder.named_modules():
  if name.endswith(lowercase_			):
   SCREAMING_SNAKE_CASE				    =		lora_attn_processor_class(
       hidden_size=module.out_proj.out_features	,  cross_attention_dim=lowercase_			)
 return text_lora_attn_procs





def 						A_					(						lowercase_			)				->int:







 """simple docstring"""

 SCREAMING_SNAKE_CASE				    =		create_text_encoder_lora_attn_procs(lowercase_			)
 SCREAMING_SNAKE_CASE				    =		AttnProcsLayers(lowercase_			)
 return text_encoder_lora_layers





def 						A_					(						lowercase_	,  lowercase_=False			)				->Optional[Any]:







 """simple docstring"""

 for _, attn_proc in text_lora_attn_procs.items():
  # set up.weights
  for layer_name, layer_module in attn_proc.named_modules():
   if layer_name.endswith('_lora'			):
    SCREAMING_SNAKE_CASE				    =		(
        torch.randn_like(layer_module.up.weight			)
        if randn_weight
        else torch.zeros_like(layer_module.up.weight			)
    )
    SCREAMING_SNAKE_CASE				    =		torch.nn.Parameter(lowercase_			)
class 	a_(		unittest.TestCase	):



     """simple docstring"""



     def     __UpperCamelCase						(		self					:	Dict)							->	str:

      """simple docstring"""
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)  ,  layers_per_block=2  ,  sample_size=3_2  ,  in_channels=4  ,  out_channels=4  ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')  ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')  ,  cross_attention_dim=3_2  ,  )
      SCREAMING_SNAKE_CASE				    =		DDIMScheduler(
          beta_start=0.0_00_85  ,  beta_end=0.0_12  ,  beta_schedule='scaled_linear'  ,  clip_sample=lowerCAmelCase__  ,  set_alpha_to_one=lowerCAmelCase__  ,  steps_offset=1  ,  )
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		AutoencoderKL(
          block_out_channels=[3_2, 6_4]  ,  in_channels=3  ,  out_channels=3  ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,  latent_channels=4  ,  )
      SCREAMING_SNAKE_CASE				    =		CLIPTextConfig(
          bos_token_id=0  ,  eos_token_id=2  ,  hidden_size=3_2  ,  intermediate_size=3_7  ,  layer_norm_eps=1e-05  ,  num_attention_heads=4  ,  num_hidden_layers=5  ,  pad_token_id=1  ,  vocab_size=1_0_0_0  ,  )
      SCREAMING_SNAKE_CASE				    =		CLIPTextModel(lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		create_unet_lora_layers(lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		create_text_encoder_lora_layers(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		{
          'unet': unet,
          'scheduler': scheduler,
          'vae': vae,
          'text_encoder': text_encoder,
          'tokenizer': tokenizer,
          'safety_checker': None,
          'feature_extractor': None,
      }
      SCREAMING_SNAKE_CASE				    =		{
          'unet_lora_layers': unet_lora_layers,
          'text_encoder_lora_layers': text_encoder_lora_layers,
          'unet_lora_attn_procs': unet_lora_attn_procs,
      }
      return pipeline_components, lora_components





     def     __UpperCamelCase						(		self					:	List[str])							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		1
      SCREAMING_SNAKE_CASE				    =		1_0
      SCREAMING_SNAKE_CASE				    =		4
      SCREAMING_SNAKE_CASE				    =		(3_2, 3_2)

      SCREAMING_SNAKE_CASE				    =		torch.manual_seed(0)
      SCREAMING_SNAKE_CASE				    =		floats_tensor((batch_size, num_channels) + sizes)
      SCREAMING_SNAKE_CASE				    =		torch.randint(1  ,  lowerCAmelCase__  ,  size=(batch_size, sequence_length)  ,  generator=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		{
          'prompt': 'A painting of a squirrel eating a burger',
          'generator': generator,
          'num_inference_steps': 2,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }

      return noise, input_ids, pipeline_inputs

      # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb





     def     __UpperCamelCase						(		self					:	List[str])							->	List[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		7_7

      SCREAMING_SNAKE_CASE				    =		torch.randint(2  ,  5_6  ,  size=(1, max_seq_length)  ,  generator=torch.manual_seed(0))

      SCREAMING_SNAKE_CASE				    =		{}
      SCREAMING_SNAKE_CASE				    =		inputs
      return prepared_inputs





     def     __UpperCamelCase						(		self					:	int  ,  lowerCAmelCase__					:	List[Any])							->	Union[str, Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      LoraLoaderMixin.save_lora_weights(
          save_directory=lowerCAmelCase__  ,  unet_lora_layers=lora_components['unet_lora_layers']  ,  text_encoder_lora_layers=lora_components['text_encoder_lora_layers']  ,  )
      self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.bin')))





     def     __UpperCamelCase						(		self					:	Any)							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_inputs()

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		original_images[0, -3:, -3:, -1]

      with tempfile.TemporaryDirectory() as tmpdirname:
       LoraLoaderMixin.save_lora_weights(
           save_directory=lowerCAmelCase__  ,  unet_lora_layers=lora_components['unet_lora_layers']  ,  text_encoder_lora_layers=lora_components['text_encoder_lora_layers']  ,  )
       self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.bin')))
       sd_pipe.load_lora_weights(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		lora_images[0, -3:, -3:, -1]

      # Outputs shouldn't match.
      self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)  ,  torch.from_numpy(lowerCAmelCase__)))





     def     __UpperCamelCase						(		self					:	Tuple)							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_inputs()

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		original_images[0, -3:, -3:, -1]

      with tempfile.TemporaryDirectory() as tmpdirname:
       LoraLoaderMixin.save_lora_weights(
           save_directory=lowerCAmelCase__  ,  unet_lora_layers=lora_components['unet_lora_layers']  ,  text_encoder_lora_layers=lora_components['text_encoder_lora_layers']  ,  safe_serialization=lowerCAmelCase__  ,  )
       self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.safetensors')))
       sd_pipe.load_lora_weights(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		lora_images[0, -3:, -3:, -1]

      # Outputs shouldn't match.
      self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)  ,  torch.from_numpy(lowerCAmelCase__)))





     def     __UpperCamelCase						(		self					:	List[str])							->	List[str]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		lora_components['unet_lora_attn_procs']
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_inputs()

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		original_images[0, -3:, -3:, -1]

      with tempfile.TemporaryDirectory() as tmpdirname:
       SCREAMING_SNAKE_CASE				    =		sd_pipe.unet
       unet.set_attn_processor(lowerCAmelCase__)
       unet.save_attn_procs(lowerCAmelCase__)
       self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.bin')))
       sd_pipe.load_lora_weights(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		lora_images[0, -3:, -3:, -1]

      # Outputs shouldn't match.
      self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)  ,  torch.from_numpy(lowerCAmelCase__)))





     def     __UpperCamelCase						(		self					:	Dict)							->	Optional[int]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		self.get_dummy_tokens()

      # inference without lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_without_lora.shape == (1, 7_7, 3_2)

      # create lora_attn_procs with zeroed out up.weights
      SCREAMING_SNAKE_CASE				    =		create_text_encoder_lora_attn_procs(pipe.text_encoder)
      set_lora_up_weights(lowerCAmelCase__  ,  randn_weight=lowerCAmelCase__)

      # monkey patch
      pipe._modify_text_encoder(lowerCAmelCase__)

      # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
      del text_attn_procs
      gc.collect()

      # inference with lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_with_lora.shape == (1, 7_7, 3_2)

      assert torch.allclose(
          lowerCAmelCase__  ,  lowerCAmelCase__), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

      # create lora_attn_procs with randn up.weights
      SCREAMING_SNAKE_CASE				    =		create_text_encoder_lora_attn_procs(pipe.text_encoder)
      set_lora_up_weights(lowerCAmelCase__  ,  randn_weight=lowerCAmelCase__)

      # monkey patch
      pipe._modify_text_encoder(lowerCAmelCase__)

      # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
      del text_attn_procs
      gc.collect()

      # inference with lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_with_lora.shape == (1, 7_7, 3_2)

      assert not torch.allclose(
          lowerCAmelCase__  ,  lowerCAmelCase__), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"





     def     __UpperCamelCase						(		self					:	Optional[int])							->	Optional[int]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		self.get_dummy_tokens()

      # inference without lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_without_lora.shape == (1, 7_7, 3_2)

      # create lora_attn_procs with randn up.weights
      SCREAMING_SNAKE_CASE				    =		create_text_encoder_lora_attn_procs(pipe.text_encoder)
      set_lora_up_weights(lowerCAmelCase__  ,  randn_weight=lowerCAmelCase__)

      # monkey patch
      pipe._modify_text_encoder(lowerCAmelCase__)

      # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
      del text_attn_procs
      gc.collect()

      # inference with lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_with_lora.shape == (1, 7_7, 3_2)

      assert not torch.allclose(
          lowerCAmelCase__  ,  lowerCAmelCase__), "lora outputs should be different to without lora outputs"

      # remove monkey patch
      pipe._remove_text_encoder_monkey_patch()

      # inference with removed lora
      SCREAMING_SNAKE_CASE				    =		pipe.text_encoder(**lowerCAmelCase__)[0]
      assert outputs_without_lora_removed.shape == (1, 7_7, 3_2)

      assert torch.allclose(
          lowerCAmelCase__  ,  lowerCAmelCase__), "remove lora monkey patch should restore the original outputs"





     def     __UpperCamelCase						(		self					:	Dict)							->	Any:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_inputs()

      with tempfile.TemporaryDirectory() as tmpdirname:
       LoraLoaderMixin.save_lora_weights(
           save_directory=lowerCAmelCase__  ,  unet_lora_layers=lora_components['unet_lora_layers']  ,  text_encoder_lora_layers=lora_components['text_encoder_lora_layers']  ,  )
       self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.bin')))
       sd_pipe.load_lora_weights(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		lora_images[0, -3:, -3:, -1]

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__  ,  cross_attention_kwargs={'scale': 0.5}).images
      SCREAMING_SNAKE_CASE				    =		lora_images_with_scale[0, -3:, -3:, -1]

      # Outputs shouldn't match.
      self.assertFalse(
          torch.allclose(torch.from_numpy(lowerCAmelCase__)  ,  torch.from_numpy(lowerCAmelCase__)))





     def     __UpperCamelCase						(		self					:	str)							->	Any:

      """simple docstring"""
      with tempfile.TemporaryDirectory() as tmpdirname:
       self.create_lora_weight_file(lowerCAmelCase__)

       SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
       SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
       SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

       # check if vanilla attention processors are used
       for _, module in sd_pipe.unet.named_modules():
        if isinstance(lowerCAmelCase__  ,  lowerCAmelCase__):
         self.assertIsInstance(module.processor  ,  (AttnProcessor, AttnProcessor2_0))

            # load LoRA weight file
       sd_pipe.load_lora_weights(lowerCAmelCase__)

       # check if lora attention processors are used
       for _, module in sd_pipe.unet.named_modules():
        if isinstance(lowerCAmelCase__  ,  lowerCAmelCase__):
         SCREAMING_SNAKE_CASE				    =		(
             LoRAAttnProcessor2_0 if hasattr(lowerCAmelCase__  ,  'scaled_dot_product_attention') else LoRAAttnProcessor
         )
         self.assertIsInstance(module.processor  ,  lowerCAmelCase__)





     @unittest.skipIf(torch_device != 'cuda'  ,  'This test is supposed to run on GPU')
     def     __UpperCamelCase						(		self					:	str)							->	int:

      """simple docstring"""
      with tempfile.TemporaryDirectory() as tmpdirname:
       self.create_lora_weight_file(lowerCAmelCase__)

       SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
       SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
       SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

       # enable XFormers
       sd_pipe.enable_xformers_memory_efficient_attention()

       # check if xFormers attention processors are used
       for _, module in sd_pipe.unet.named_modules():
        if isinstance(lowerCAmelCase__  ,  lowerCAmelCase__):
         self.assertIsInstance(module.processor  ,  lowerCAmelCase__)

            # load LoRA weight file
       sd_pipe.load_lora_weights(lowerCAmelCase__)

       # check if lora attention processors are used
       for _, module in sd_pipe.unet.named_modules():
        if isinstance(lowerCAmelCase__  ,  lowerCAmelCase__):
         self.assertIsInstance(module.processor  ,  lowerCAmelCase__)





     @unittest.skipIf(torch_device != 'cuda'  ,  'This test is supposed to run on GPU')
     def     __UpperCamelCase						(		self					:	int)							->	Optional[int]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_components()
      SCREAMING_SNAKE_CASE				    =		StableDiffusionPipeline(**lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		sd_pipe.to(lowerCAmelCase__)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__)

      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.get_dummy_inputs()

      # enable XFormers
      sd_pipe.enable_xformers_memory_efficient_attention()

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		original_images[0, -3:, -3:, -1]

      with tempfile.TemporaryDirectory() as tmpdirname:
       LoraLoaderMixin.save_lora_weights(
           save_directory=lowerCAmelCase__  ,  unet_lora_layers=lora_components['unet_lora_layers']  ,  text_encoder_lora_layers=lora_components['text_encoder_lora_layers']  ,  )
       self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__  ,  'pytorch_lora_weights.bin')))
       sd_pipe.load_lora_weights(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE				    =		sd_pipe(**lowerCAmelCase__).images
      SCREAMING_SNAKE_CASE				    =		lora_images[0, -3:, -3:, -1]

      # Outputs shouldn't match.
      self.assertFalse(torch.allclose(torch.from_numpy(lowerCAmelCase__)  ,  torch.from_numpy(lowerCAmelCase__)))







"""simple docstring"""




import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
     from transformers.pipelines.conversational import Conversation

SCREAMING_SNAKE_CASE_  :      Optional[int]							=      logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_  :      Optional[int]							=      {
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
}

SCREAMING_SNAKE_CASE_  :      Dict							=      {
    '''vocab_file''': {
        '''gpt2''': '''https://huggingface.co/gpt2/resolve/main/vocab.json''',
        '''gpt2-medium''': '''https://huggingface.co/gpt2-medium/resolve/main/vocab.json''',
        '''gpt2-large''': '''https://huggingface.co/gpt2-large/resolve/main/vocab.json''',
        '''gpt2-xl''': '''https://huggingface.co/gpt2-xl/resolve/main/vocab.json''',
        '''distilgpt2''': '''https://huggingface.co/distilgpt2/resolve/main/vocab.json''',
    },
    '''merges_file''': {
        '''gpt2''': '''https://huggingface.co/gpt2/resolve/main/merges.txt''',
        '''gpt2-medium''': '''https://huggingface.co/gpt2-medium/resolve/main/merges.txt''',
        '''gpt2-large''': '''https://huggingface.co/gpt2-large/resolve/main/merges.txt''',
        '''gpt2-xl''': '''https://huggingface.co/gpt2-xl/resolve/main/merges.txt''',
        '''distilgpt2''': '''https://huggingface.co/distilgpt2/resolve/main/merges.txt''',
    },
}

SCREAMING_SNAKE_CASE_  :      Optional[Any]							=      {
    '''gpt2''': 1024,
    '''gpt2-medium''': 1024,
    '''gpt2-large''': 1024,
    '''gpt2-xl''': 1024,
    '''distilgpt2''': 1024,
}



@lru_cache()
def 						UpperCAmelCase__ (					)   ->   Optional[int]:




   """simple docstring"""
   lowerCamelCase__						       =      (
       list(range(ord("!" )						,     ord("~" ) + 1 ) ) + list(range(ord("¡" )						,     ord("¬" ) + 1 ) ) + list(range(ord("®" )						,     ord("ÿ" ) + 1 ) )
   )
   lowerCamelCase__						       =      bs[:]
   lowerCamelCase__						       =      0
   for b in range(2**8 ):
      if b not in bs:
         bs.append(A__ )
         cs.append(2**8 + n )
         n += 1
   lowerCamelCase__						       =      [chr(A__ ) for n in cs]
   return dict(zip(A__						,     A__ ) )



def 						UpperCAmelCase__ (					A__ )   ->   int:




   """simple docstring"""
   lowerCamelCase__						       =      set()
   lowerCamelCase__						       =      word[0]
   for char in word[1:]:
      pairs.add((prev_char, char) )
      lowerCamelCase__						       =      char
   return pairs


class 		_A							(      __a       ):
      __a             =						VOCAB_FILES_NAMES
      __a             =						PRETRAINED_VOCAB_FILES_MAP
      __a             =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __a             =						['input_ids', 'attention_mask']





      def __init__(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__="replace"     ,    SCREAMING_SNAKE_CASE__="<|endoftext|>"     ,    SCREAMING_SNAKE_CASE__="<|endoftext|>"     ,    SCREAMING_SNAKE_CASE__="<|endoftext|>"     ,    SCREAMING_SNAKE_CASE__=None     ,    SCREAMING_SNAKE_CASE__=False     ,    SCREAMING_SNAKE_CASE__=False     ,    **SCREAMING_SNAKE_CASE__     ,    ) ->       List[str]:
         lowerCamelCase__						       =      AddedToken(SCREAMING_SNAKE_CASE__     ,    lstrip=SCREAMING_SNAKE_CASE__     ,    rstrip=SCREAMING_SNAKE_CASE__		) if isinstance(SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__		) else bos_token
         lowerCamelCase__						       =      AddedToken(SCREAMING_SNAKE_CASE__     ,    lstrip=SCREAMING_SNAKE_CASE__     ,    rstrip=SCREAMING_SNAKE_CASE__		) if isinstance(SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__		) else eos_token
         lowerCamelCase__						       =      AddedToken(SCREAMING_SNAKE_CASE__     ,    lstrip=SCREAMING_SNAKE_CASE__     ,    rstrip=SCREAMING_SNAKE_CASE__		) if isinstance(SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__		) else unk_token
         lowerCamelCase__						       =      AddedToken(SCREAMING_SNAKE_CASE__     ,    lstrip=SCREAMING_SNAKE_CASE__     ,    rstrip=SCREAMING_SNAKE_CASE__		) if isinstance(SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__		) else pad_token
         super().__init__(
             errors=SCREAMING_SNAKE_CASE__     ,    unk_token=SCREAMING_SNAKE_CASE__     ,    bos_token=SCREAMING_SNAKE_CASE__     ,    eos_token=SCREAMING_SNAKE_CASE__     ,    pad_token=SCREAMING_SNAKE_CASE__     ,    add_prefix_space=SCREAMING_SNAKE_CASE__     ,    add_bos_token=SCREAMING_SNAKE_CASE__     ,    **SCREAMING_SNAKE_CASE__     ,    )
         lowerCamelCase__						       =      add_bos_token

         with open(SCREAMING_SNAKE_CASE__     ,    encoding="utf-8"		) as vocab_handle:
            lowerCamelCase__						       =      json.load(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						       =      {v: k for k, v in self.encoder.items()}
         lowerCamelCase__						       =      errors  # how to handle errors in decoding
         lowerCamelCase__						       =      bytes_to_unicode()
         lowerCamelCase__						       =      {v: k for k, v in self.byte_encoder.items()}
         with open(SCREAMING_SNAKE_CASE__     ,    encoding="utf-8"		) as merges_handle:
            lowerCamelCase__						       =      merges_handle.read().split("\n"		)[1:-1]
         lowerCamelCase__						       =      [tuple(merge.split()		) for merge in bpe_merges]
         lowerCamelCase__						       =      dict(zip(SCREAMING_SNAKE_CASE__     ,    range(len(SCREAMING_SNAKE_CASE__		)		)		)		)
         lowerCamelCase__						       =      {}
         lowerCamelCase__						       =      add_prefix_space

         # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
         lowerCamelCase__						       =      re.compile(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"		)





      @property
      def        _lowerCamelCase	(	self		) ->       Optional[int]:
         return len(self.encoder		)





      def        _lowerCamelCase	(	self		) ->       str:
         return dict(self.encoder     ,    **self.added_tokens_encoder		)





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       Any:
         if token in self.cache:
            return self.cache[token]
         lowerCamelCase__						       =      tuple(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						       =      get_pairs(SCREAMING_SNAKE_CASE__		)

         if not pairs:
            return token

         while True:
            lowerCamelCase__						       =      min(SCREAMING_SNAKE_CASE__     ,    key=lambda SCREAMING_SNAKE_CASE__		: self.bpe_ranks.get(SCREAMING_SNAKE_CASE__     ,    float("inf"		)		)		)
            if bigram not in self.bpe_ranks:
               break
            lowerCamelCase__       ,							lowerCamelCase__						       =      bigram
            lowerCamelCase__						       =      []
            lowerCamelCase__						       =      0
            while i < len(SCREAMING_SNAKE_CASE__		):
               try:
                  lowerCamelCase__						       =      word.index(SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__		)
               except ValueError:
                  new_word.extend(word[i:]		)
                  break
               else:
                  new_word.extend(word[i:j]		)
                  lowerCamelCase__						       =      j

               if word[i] == first and i < len(SCREAMING_SNAKE_CASE__		) - 1 and word[i + 1] == second:
                  new_word.append(first + second		)
                  i += 2
               else:
                  new_word.append(word[i]		)
                  i += 1
            lowerCamelCase__						       =      tuple(SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						       =      new_word
            if len(SCREAMING_SNAKE_CASE__		) == 1:
               break
            else:
               lowerCamelCase__						       =      get_pairs(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						       =      " ".join(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						       =      word
         return word





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__=None		) ->       Optional[int]:
         if self.add_bos_token:
            lowerCamelCase__						       =      [self.bos_token_id]
         else:
            lowerCamelCase__						       =      []

         lowerCamelCase__						       =      bos_token_ids + token_ids_a

         if token_ids_a is None:
            return output

         return output + bos_token_ids + token_ids_a





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__ = None     ,    SCREAMING_SNAKE_CASE__ = False		) ->       List[int]:
         if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_a=SCREAMING_SNAKE_CASE__     ,    token_ids_a=SCREAMING_SNAKE_CASE__     ,    already_has_special_tokens=SCREAMING_SNAKE_CASE__		)

         if not self.add_bos_token:
            return super().get_special_tokens_mask(
                token_ids_a=SCREAMING_SNAKE_CASE__     ,    token_ids_a=SCREAMING_SNAKE_CASE__     ,    already_has_special_tokens=SCREAMING_SNAKE_CASE__		)

         if token_ids_a is None:
            return [1] + ([0] * len(SCREAMING_SNAKE_CASE__		))
         return [1] + ([0] * len(SCREAMING_SNAKE_CASE__		)) + [1] + ([0] * len(SCREAMING_SNAKE_CASE__		))





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       Tuple:
         lowerCamelCase__						       =      []
         for token in re.findall(self.pat     ,    SCREAMING_SNAKE_CASE__		):
            lowerCamelCase__						       =      "".join(
                self.byte_encoder[b] for b in token.encode("utf-8"		)		)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(SCREAMING_SNAKE_CASE__		).split(" "		)		)
         return bpe_tokens





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       Tuple:
         return self.encoder.get(SCREAMING_SNAKE_CASE__     ,    self.encoder.get(self.unk_token		)		)





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       Any:
         return self.decoder.get(SCREAMING_SNAKE_CASE__		)





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       str:
         lowerCamelCase__						       =      "".join(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						       =      bytearray([self.byte_decoder[c] for c in text]		).decode("utf-8"     ,    errors=self.errors		)
         return text





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__ = None		) ->       Tuple[str]:
         if not os.path.isdir(SCREAMING_SNAKE_CASE__		):
            logger.error(f'Vocabulary path ({save_directory}) should be a directory'		)
            return
         lowerCamelCase__						       =      os.path.join(
             SCREAMING_SNAKE_CASE__     ,    (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]		)
         lowerCamelCase__						       =      os.path.join(
             SCREAMING_SNAKE_CASE__     ,    (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]		)

         with open(SCREAMING_SNAKE_CASE__     ,    "w"     ,    encoding="utf-8"		) as f:
            f.write(json.dumps(self.encoder     ,    indent=2     ,    sort_keys=SCREAMING_SNAKE_CASE__     ,    ensure_ascii=SCREAMING_SNAKE_CASE__		) + "\n"		)

         lowerCamelCase__						       =      0
         with open(SCREAMING_SNAKE_CASE__     ,    "w"     ,    encoding="utf-8"		) as writer:
            writer.write("#version: 0.2\n"		)
            for bpe_tokens, token_index in sorted(self.bpe_ranks.items()     ,    key=lambda SCREAMING_SNAKE_CASE__		: kv[1]		):
               if index != token_index:
                  logger.warning(
                      f'Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'
                      " Please check that the tokenizer is not corrupted!"		)
                  lowerCamelCase__						       =      token_index
               writer.write(" ".join(SCREAMING_SNAKE_CASE__		) + "\n"		)
               index += 1

         return vocab_file, merge_file





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__=False     ,    **SCREAMING_SNAKE_CASE__		) ->       int:
         lowerCamelCase__						       =      kwargs.pop("add_prefix_space"     ,    self.add_prefix_space		)
         if is_split_into_words or add_prefix_space:
            lowerCamelCase__						       =      " " + text
         return (text, kwargs)





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__		) ->       List[int]:
         lowerCamelCase__						       =      []
         for is_user, text in conversation.iter_texts():
            input_ids.extend(self.encode(SCREAMING_SNAKE_CASE__     ,    add_special_tokens=SCREAMING_SNAKE_CASE__		) + [self.eos_token_id]		)
         if len(SCREAMING_SNAKE_CASE__		) > self.model_max_length:
            lowerCamelCase__						       =      input_ids[-self.model_max_length :]
         return input_ids




