





'''simple docstring'''



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

@skip_mps
class     A_ (  lowerCAmelCase_    ,			lowerCAmelCase_    ,			lowerCAmelCase_    ,			unittest.TestCase  ):

			'''simple docstring'''





			_lowerCAmelCase						=    StableDiffusionModelEditingPipeline
			_lowerCAmelCase						=    TEXT_TO_IMAGE_PARAMS
			_lowerCAmelCase						=    TEXT_TO_IMAGE_BATCH_PARAMS
			_lowerCAmelCase						=    TEXT_TO_IMAGE_IMAGE_PARAMS
			_lowerCAmelCase						=    TEXT_TO_IMAGE_IMAGE_PARAMS




			def       a	(					self   ):
								torch.manual_seed(0   )
								_UpperCamelCase											=   UNetaDConditionModel(
								    block_out_channels=(32, 64)			,     layers_per_block=2			,     sample_size=32			,     in_channels=4			,     out_channels=4			,     down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,     up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,     cross_attention_dim=32			,     )
								_UpperCamelCase											=   DDIMScheduler()
								torch.manual_seed(0   )
								_UpperCamelCase											=   AutoencoderKL(
								    block_out_channels=[32, 64]			,     in_channels=3			,     out_channels=3			,     down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,     up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,     latent_channels=4			,     )
								torch.manual_seed(0   )
								_UpperCamelCase											=   CLIPTextConfig(
								    bos_token_id=0			,     eos_token_id=2			,     hidden_size=32			,     intermediate_size=37			,     layer_norm_eps=1e-0_5			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=10_00			,     )
								_UpperCamelCase											=   CLIPTextModel(A_   )
								_UpperCamelCase											=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

								_UpperCamelCase											=   {
								    "unet": unet,
								    "scheduler": scheduler,
								    "vae": vae,
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "safety_checker": None,
								    "feature_extractor": None,
								}
								return components




			def       a	(					self			,     A_			,     A_=0   ):
								_UpperCamelCase											=   torch.manual_seed(A_   )
								_UpperCamelCase											=   {
								    "prompt": "A field of roses",
								    "generator": generator,
								    # Setting height and width to None to prevent OOMs on CPU.
								    "height": None,
								    "width": None,
								    "num_inference_steps": 2,
								    "guidance_scale": 6.0,
								    "output_type": "numpy",
								}
								return inputs




			def       a	(					self   ):
								_UpperCamelCase											=   "cpu"  # ensure determinism for the device-dependent torch.Generator
								_UpperCamelCase											=   self.get_dummy_components()
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline(**A_   )
								_UpperCamelCase											=   sd_pipe.to(A_   )
								sd_pipe.set_progress_bar_config(disable=A_   )

								_UpperCamelCase											=   self.get_dummy_inputs(A_   )
								_UpperCamelCase											=   sd_pipe(**A_   ).images
								_UpperCamelCase											=   image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								_UpperCamelCase											=   np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889]   )

								assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2




			def       a	(					self   ):
								_UpperCamelCase											=   "cpu"  # ensure determinism for the device-dependent torch.Generator
								_UpperCamelCase											=   self.get_dummy_components()
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline(**A_   )
								_UpperCamelCase											=   sd_pipe.to(A_   )
								sd_pipe.set_progress_bar_config(disable=A_   )

								_UpperCamelCase											=   self.get_dummy_inputs(A_   )
								_UpperCamelCase											=   "french fries"
								_UpperCamelCase											=   sd_pipe(**A_			,     negative_prompt=A_   )
								_UpperCamelCase											=   output.images
								_UpperCamelCase											=   image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								_UpperCamelCase											=   np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913]   )

								assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2




			def       a	(					self   ):
								_UpperCamelCase											=   "cpu"  # ensure determinism for the device-dependent torch.Generator
								_UpperCamelCase											=   self.get_dummy_components()
								_UpperCamelCase											=   EulerAncestralDiscreteScheduler(
								    beta_start=0.0_0085			,     beta_end=0.012			,     beta_schedule="scaled_linear"   )
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline(**A_   )
								_UpperCamelCase											=   sd_pipe.to(A_   )
								sd_pipe.set_progress_bar_config(disable=A_   )

								_UpperCamelCase											=   self.get_dummy_inputs(A_   )
								_UpperCamelCase											=   sd_pipe(**A_   ).images
								_UpperCamelCase											=   image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								_UpperCamelCase											=   np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027]   )

								assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2




			def       a	(					self   ):
								_UpperCamelCase											=   "cpu"  # ensure determinism for the device-dependent torch.Generator
								_UpperCamelCase											=   self.get_dummy_components()
								_UpperCamelCase											=   PNDMScheduler()
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline(**A_   )
								_UpperCamelCase											=   sd_pipe.to(A_   )
								sd_pipe.set_progress_bar_config(disable=A_   )

								_UpperCamelCase											=   self.get_dummy_inputs(A_   )
								# the pipeline does not expect pndm so test if it raises error.
								with self.assertRaises(A_   ):
													_UpperCamelCase											=   sd_pipe(**A_   ).images




			def       a	(					self   ):
								super().test_inference_batch_single_identical(expected_max_diff=5e-3   )




			def       a	(					self   ):
								super().test_attention_slicing_forward_pass(expected_max_diff=5e-3   )




@slow
@require_torch_gpu
class     A_ (  unittest.TestCase  ):

			'''simple docstring'''





			def       a	(					self   ):
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




			def       a	(					self			,     A_=0   ):
								_UpperCamelCase											=   torch.manual_seed(A_   )
								_UpperCamelCase											=   {
								    "prompt": "A field of roses",
								    "generator": generator,
								    "num_inference_steps": 3,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs




			def       a	(					self   ):
								_UpperCamelCase											=   "CompVis/stable-diffusion-v1-4"
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline.from_pretrained(A_			,     safety_checker=A_   )
								pipe.to(A_   )
								pipe.set_progress_bar_config(disable=A_   )
								pipe.enable_attention_slicing()

								_UpperCamelCase											=   self.get_inputs()
								_UpperCamelCase											=   pipe(**A_   ).images
								_UpperCamelCase											=   image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)

								_UpperCamelCase											=   np.array(
								    [0.674_9496, 0.638_6453, 0.5144_3267, 0.6609_4905, 0.6192_1215, 0.549_1332, 0.574_4417, 0.5807_5106, 0.517_4658]   )

								assert np.abs(expected_slice - image_slice   ).max() < 1e-2

								# make sure image changes after editing
								pipe.edit_model("A pack of roses"			,     "A pack of blue roses"   )

								_UpperCamelCase											=   pipe(**A_   ).images
								_UpperCamelCase											=   image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 5_12, 5_12, 3)

								assert np.abs(expected_slice - image_slice   ).max() > 1e-1




			def       a	(					self   ):
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								_UpperCamelCase											=   "CompVis/stable-diffusion-v1-4"
								_UpperCamelCase											=   DDIMScheduler.from_pretrained(A_			,     subfolder="scheduler"   )
								_UpperCamelCase											=   StableDiffusionModelEditingPipeline.from_pretrained(
								    A_			,     scheduler=A_			,     safety_checker=A_   )
								_UpperCamelCase											=   pipe.to(A_   )
								pipe.set_progress_bar_config(disable=A_   )
								pipe.enable_attention_slicing(1   )
								pipe.enable_sequential_cpu_offload()

								_UpperCamelCase											=   self.get_inputs()
								_UpperCamelCase											=   pipe(**A_   )

								_UpperCamelCase											=   torch.cuda.max_memory_allocated()
								# make sure that less than 4.4 GB is allocated
								assert mem_bytes < 4.4 * 10**9










'''simple docstring'''



import os
from typing import List, Union

import tensorflow as tf
from tensorflow_text import BertTokenizer as BertTokenizerLayer
from tensorflow_text import FastBertTokenizer, ShrinkLongestTrimmer, case_fold_utfa, combine_segments, pad_model_inputs

from .tokenization_bert import BertTokenizer

class     A_ (  tf.keras.layers.Layer  ):

			'''simple docstring'''





			def __init__(					self			,     A_			,     A_			,     A_ = None			,     A_ = None			,     A_ = None			,     A_ = "longest"			,     A_ = True			,     A_ = 5_12			,     A_ = None			,     A_ = True			,     A_ = True			,     A_ = True			,     **A_			,     ):
								super().__init__()
								if use_fast_bert_tokenizer:
													_UpperCamelCase											=   FastBertTokenizer(
													    A_			,     token_out_type=tf.intaa			,     lower_case_nfd_strip_accents=A_			,     **A_   )
								else:
													_UpperCamelCase											=   tf.lookup.StaticVocabularyTable(
													    tf.lookup.KeyValueTensorInitializer(
													        keys=A_			,     key_dtype=tf.string			,     values=tf.range(tf.size(A_			,     out_type=tf.intaa   )			,     dtype=tf.intaa   )			,     value_dtype=tf.intaa			,     )			,     num_oov_buckets=1			,     )
													_UpperCamelCase											=   BertTokenizerLayer(
													    A_			,     token_out_type=tf.intaa			,     lower_case=A_			,     **A_   )

								_UpperCamelCase											=   vocab_list
								_UpperCamelCase											=   do_lower_case
								_UpperCamelCase											=   cls_token_id or vocab_list.index("[CLS]"   )
								_UpperCamelCase											=   sep_token_id or vocab_list.index("[SEP]"   )
								_UpperCamelCase											=   pad_token_id or vocab_list.index("[PAD]"   )
								_UpperCamelCase											=   ShrinkLongestTrimmer(max_length - 3			,     axis=1   )  # Allow room for special tokens
								_UpperCamelCase											=   max_length
								_UpperCamelCase											=   padding
								_UpperCamelCase											=   truncation
								_UpperCamelCase											=   pad_to_multiple_of
								_UpperCamelCase											=   return_token_type_ids
								_UpperCamelCase											=   return_attention_mask




			@classmethod
			def       a	(					cls			,     A_			,     **A_   ):  # noqa: F821
								_UpperCamelCase											=   kwargs.pop("do_lower_case"			,     A_   )
								_UpperCamelCase											=   tokenizer.do_lower_case if do_lower_case is None else do_lower_case
								_UpperCamelCase											=   kwargs.pop("cls_token_id"			,     A_   )
								_UpperCamelCase											=   tokenizer.cls_token_id if cls_token_id is None else cls_token_id
								_UpperCamelCase											=   kwargs.pop("sep_token_id"			,     A_   )
								_UpperCamelCase											=   tokenizer.sep_token_id if sep_token_id is None else sep_token_id
								_UpperCamelCase											=   kwargs.pop("pad_token_id"			,     A_   )
								_UpperCamelCase											=   tokenizer.pad_token_id if pad_token_id is None else pad_token_id

								_UpperCamelCase											=   tokenizer.get_vocab()
								_UpperCamelCase											=   sorted([(wordpiece, idx) for wordpiece, idx in vocab.items()]			,     key=lambda A_   : x[1]   )
								_UpperCamelCase											=   [entry[0] for entry in vocab]
								return cls(
								    vocab_list=A_			,     do_lower_case=A_			,     cls_token_id=A_			,     sep_token_id=A_			,     pad_token_id=A_			,     **A_			,     )




			@classmethod
			def       a	(					cls			,     A_			,     *A_			,     **A_   ):
								try:
													_UpperCamelCase											=   BertTokenizer.from_pretrained(A_			,     *A_			,     **A_   )
								except:  # noqa: E722
													from .tokenization_bert_fast import BertTokenizerFast

													_UpperCamelCase											=   BertTokenizerFast.from_pretrained(A_			,     *A_			,     **A_   )
								return cls.from_tokenizer(A_			,     **A_   )




			def       a	(					self			,     A_   ):
								if self.do_lower_case:
													_UpperCamelCase											=   case_fold_utfa(A_   )
								_UpperCamelCase											=   self.tf_tokenizer.tokenize(A_   )
								return tokens.merge_dims(1			,     -1   )




			def       a	(					self			,     A_			,     A_=None			,     A_=None			,     A_=None			,     A_=None			,     A_=None			,     A_=None			,     A_=None			,     ):
								if padding is None:
													_UpperCamelCase											=   self.padding
								if padding not in ("longest", "max_length"):
													raise ValueError("Padding must be either 'longest' or 'max_length'!"   )
								if max_length is not None and text_pair is not None:
													# Because we have to instantiate a Trimmer to do it properly
													raise ValueError("max_length cannot be overridden at call time when truncating paired texts!"   )
								if max_length is None:
													_UpperCamelCase											=   self.max_length
								if truncation is None:
													_UpperCamelCase											=   self.truncation
								if pad_to_multiple_of is None:
													_UpperCamelCase											=   self.pad_to_multiple_of
								if return_token_type_ids is None:
													_UpperCamelCase											=   self.return_token_type_ids
								if return_attention_mask is None:
													_UpperCamelCase											=   self.return_attention_mask
								if not isinstance(A_			,     tf.Tensor   ):
													_UpperCamelCase											=   tf.convert_to_tensor(A_   )
								if text_pair is not None and not isinstance(A_			,     tf.Tensor   ):
													_UpperCamelCase											=   tf.convert_to_tensor(A_   )
								if text_pair is not None:
													if text.shape.rank > 1:
																		raise ValueError("text argument should not be multidimensional when a text pair is supplied!"   )
													if text_pair.shape.rank > 1:
																		raise ValueError("text_pair should not be multidimensional!"   )
								if text.shape.rank == 2:
													_UpperCamelCase				,	_UpperCamelCase											=   text[:, 0], text[:, 1]
								_UpperCamelCase											=   self.unpaired_tokenize(A_   )
								if text_pair is None:  # Unpaired text
													if truncation:
																		_UpperCamelCase											=   text[:, : max_length - 2]  # Allow room for special tokens
													_UpperCamelCase				,	_UpperCamelCase											=   combine_segments(
													    (text,)			,     start_of_sequence_id=self.cls_token_id			,     end_of_segment_id=self.sep_token_id   )
								else:  # Paired text
													_UpperCamelCase											=   self.unpaired_tokenize(A_   )
													if truncation:
																		_UpperCamelCase				,	_UpperCamelCase											=   self.paired_trimmer.trim([text, text_pair]   )
													_UpperCamelCase				,	_UpperCamelCase											=   combine_segments(
													    (text, text_pair)			,     start_of_sequence_id=self.cls_token_id			,     end_of_segment_id=self.sep_token_id   )
								if padding == "longest":
													_UpperCamelCase											=   input_ids.bounding_shape(axis=1   )
													if pad_to_multiple_of is not None:
																		# No ceiling division in tensorflow, so we negate floordiv instead
																		_UpperCamelCase											=   pad_to_multiple_of * (-tf.math.floordiv(-pad_length			,     A_   ))
								else:
													_UpperCamelCase											=   max_length

								_UpperCamelCase				,	_UpperCamelCase											=   pad_model_inputs(A_			,     max_seq_length=A_			,     pad_value=self.pad_token_id   )
								_UpperCamelCase											=   {"input_ids": input_ids}
								if return_attention_mask:
													_UpperCamelCase											=   attention_mask
								if return_token_type_ids:
													_UpperCamelCase				,	_UpperCamelCase											=   pad_model_inputs(
													    A_			,     max_seq_length=A_			,     pad_value=self.pad_token_id   )
													_UpperCamelCase											=   token_type_ids
								return output




			def       a	(					self   ):
								return {
								    "vocab_list": self.vocab_list,
								    "do_lower_case": self.do_lower_case,
								    "cls_token_id": self.cls_token_id,
								    "sep_token_id": self.sep_token_id,
								    "pad_token_id": self.pad_token_id,
								}



