

import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class   __lowercase      (_UpperCAmelCase		,   _UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  StableDiffusionXLImgaImgPipeline
 _UpperCamelCase			  =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 _UpperCamelCase			  =  PipelineTesterMixin.required_optional_params - {"""latents"""}
 _UpperCamelCase			  =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 _UpperCamelCase			  =  IMAGE_TO_IMAGE_IMAGE_PARAMS
 _UpperCamelCase			  =  IMAGE_TO_IMAGE_IMAGE_PARAMS
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       torch.manual_seed(0					)
       __lowerCAmelCase					:						Any					=					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''')      ,							attention_head_dim=(2, 4)      ,							use_linear_projection=A_      ,							addition_embed_type='''text_time'''      ,							addition_time_embed_dim=8      ,							transformer_layers_per_block=(1, 2)      ,							projection_class_embeddings_input_dim=80      ,							cross_attention_dim=64      ,							)
       __lowerCAmelCase					:						Optional[Any]					=					EulerDiscreteScheduler(
           beta_start=0.00_085      ,							beta_end=0.012      ,							steps_offset=1      ,							beta_schedule='''scaled_linear'''      ,							timestep_spacing='''leading'''      ,							)
       torch.manual_seed(0					)
       __lowerCAmelCase					:						str					=					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      ,							sample_size=128      ,							)
       torch.manual_seed(0					)
       __lowerCAmelCase					:						Any					=					CLIPTextConfig(
           bos_token_id=0      ,							eos_token_id=2      ,							hidden_size=32      ,							intermediate_size=37      ,							layer_norm_eps=1e-05      ,							num_attention_heads=4      ,							num_hidden_layers=5      ,							pad_token_id=1      ,							vocab_size=1000      ,							hidden_act='''gelu'''      ,							projection_dim=32      ,							)
       __lowerCAmelCase					:						Any					=					CLIPTextModel(A_					)
       __lowerCAmelCase					:						Dict					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,							local_files_only=A_					)

       __lowerCAmelCase					:						int					=					CLIPTextModelWithProjection(A_					)
       __lowerCAmelCase					:						Optional[int]					=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''      ,							local_files_only=A_					)

       __lowerCAmelCase					:						int					=					{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''text_encoder_2''': text_encoder_a,
           '''tokenizer_2''': tokenizer_a,
           # "safety_checker": None,
           # "feature_extractor": None,
       }
       return components
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_=0					)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					floats_tensor((1, 3, 32, 32)      ,							rng=random.Random(A_					)					).to(A_					)
       __lowerCAmelCase					:						int					=					image / 2 + 0.5
       if str(A_					).startswith('''mps'''					):
             __lowerCAmelCase					:						Union[str, Any]					=					torch.manual_seed(A_					)
       else:
             __lowerCAmelCase					:						int					=					torch.Generator(device=A_					).manual_seed(A_					)
       __lowerCAmelCase					:						Dict					=					{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 5.0,
           '''output_type''': '''numpy''',
           '''strength''': 0.75,
       }
       return inputs
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __lowerCAmelCase					:						Optional[int]					=					self.get_dummy_components()
       __lowerCAmelCase					:						Dict					=					StableDiffusionXLImgaImgPipeline(**A_					)
       __lowerCAmelCase					:						Optional[int]					=					sd_pipe.to(A_					)
       sd_pipe.set_progress_bar_config(disable=A_					)

       __lowerCAmelCase					:						Union[str, Any]					=					self.get_dummy_inputs(A_					)
       __lowerCAmelCase					:						Any					=					sd_pipe(**A_					).images
       __lowerCAmelCase					:						Tuple					=					image[0, -3:, -3:, -1]

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

       __lowerCAmelCase					:						Dict					=					np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]					)

       assert np.abs(image_slice.flatten() - expected_slice					).max() < 1e-2
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3					)
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       super().test_inference_batch_single_identical(expected_max_diff=3e-3					)
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       pass





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



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					self.get_dummy_components()
       __lowerCAmelCase					:						str					=					StableDiffusionXLImgaImgPipeline(**A_					)
       __lowerCAmelCase					:						List[str]					=					sd_pipe.to(A_					)
       __lowerCAmelCase					:						int					=					sd_pipe.to(A_					)
       sd_pipe.set_progress_bar_config(disable=A_					)

       # forward without prompt embeds
       __lowerCAmelCase					:						Optional[Any]					=					self.get_dummy_inputs(A_					)
       __lowerCAmelCase					:						List[str]					=					3 * ['''this is a negative prompt''']
       __lowerCAmelCase					:						List[Any]					=					negative_prompt
       __lowerCAmelCase					:						Dict					=					3 * [inputs['''prompt''']]

       __lowerCAmelCase					:						Any					=					sd_pipe(**A_					)
       __lowerCAmelCase					:						List[str]					=					output.images[0, -3:, -3:, -1]

       # forward with prompt embeds
       __lowerCAmelCase					:						Optional[Any]					=					self.get_dummy_inputs(A_					)
       __lowerCAmelCase					:						Tuple					=					3 * ['''this is a negative prompt''']
       __lowerCAmelCase					:						List[str]					=					3 * [inputs.pop('''prompt'''					)]

       (
           (
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							
       )					:						Tuple					=					sd_pipe.encode_prompt(A_      ,							negative_prompt=A_					)

       __lowerCAmelCase					:						Dict					=					sd_pipe(
           **A_      ,							prompt_embeds=A_      ,							negative_prompt_embeds=A_      ,							pooled_prompt_embeds=A_      ,							negative_pooled_prompt_embeds=A_      ,							)
       __lowerCAmelCase					:						Optional[int]					=					output.images[0, -3:, -3:, -1]

       # make sure that it's equal
       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1e-4







@slow
@require_torch_gpu
class   __lowercase      (unittest.TestCase ):
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_="cpu"      ,							A_=torch.floataa      ,							A_=0					)					->List[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					torch.Generator(device=A_					).manual_seed(A_					)
       __lowerCAmelCase					:						Dict					=					np.random.RandomState(A_					).standard_normal((1, 4, 64, 64)					)
       __lowerCAmelCase					:						str					=					torch.from_numpy(A_					).to(device=A_      ,							dtype=A_					)
       __lowerCAmelCase					:						Union[str, Any]					=					{
           '''prompt''': '''a photograph of an astronaut riding a horse''',
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs





 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''					)
       pipe.to(A_					)
       pipe.set_progress_bar_config(disable=A_					)

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

       assert image.shape == (1, 512, 512, 3)
       __lowerCAmelCase					:						Dict					=					np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]					)
       assert np.abs(image_slice - expected_slice					).max() < 7e-3





import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_torch, require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_torch_available, is_vision_available


if is_torch_available():
 import torch


if is_vision_available():
 from PIL import Image

 from transformers import MgpstrProcessor, ViTImageProcessor







@require_torch
@require_vision
class   __lowercase      (unittest.TestCase ):
 _UpperCamelCase			  =  ViTImageProcessor if is_vision_available() else None
 @property
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       return self.image_processor_tester.prepare_image_processor_dict()
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						Tuple					=					(3, 32, 128)
       __lowerCAmelCase					:						List[str]					=					tempfile.mkdtemp()

       # fmt: off
       __lowerCAmelCase					:						List[str]					=					['''[GO]''', '''[s]''', '''0''', '''1''', '''2''', '''3''', '''4''', '''5''', '''6''', '''7''', '''8''', '''9''', '''a''', '''b''', '''c''', '''d''', '''e''', '''f''', '''g''', '''h''', '''i''', '''j''', '''k''', '''l''', '''m''', '''n''', '''o''', '''p''', '''q''', '''r''', '''s''', '''t''', '''u''', '''v''', '''w''', '''x''', '''y''', '''z''']
       # fmt: on
       __lowerCAmelCase					:						Optional[int]					=					dict(zip(A_      ,							range(len(A_					)					)					)					)

       __lowerCAmelCase					:						Tuple					=					os.path.join(self.tmpdirname      ,							VOCAB_FILES_NAMES['''vocab_file''']					)
       with open(self.vocab_file      ,							'''w'''      ,							encoding='''utf-8'''					) as fp:
             fp.write(json.dumps(A_					) + '''\n'''					)

       __lowerCAmelCase					:						Union[str, Any]					=					{
           '''do_normalize''': False,
           '''do_resize''': True,
           '''image_processor_type''': '''ViTImageProcessor''',
           '''resample''': 3,
           '''size''': {'''height''': 32, '''width''': 128},
       }
       __lowerCAmelCase					:						Optional[Any]					=					os.path.join(self.tmpdirname      ,							A_					)
       with open(self.image_processor_file      ,							'''w'''      ,							encoding='''utf-8'''					) as fp:
             json.dump(A_      ,							A_					)
 def 	UpperCamelCase__  (							self      ,							**A_					)					->Tuple:



       '''simple docstring'''





       return MgpstrTokenizer.from_pretrained(self.tmpdirname      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							**A_					)					->Tuple:



       '''simple docstring'''





       return ViTImageProcessor.from_pretrained(self.tmpdirname      ,							**A_					)
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       shutil.rmtree(self.tmpdirname					)
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						Tuple					=					np.random.randint(255      ,							size=(3, 30, 400)      ,							dtype=np.uinta					)

       __lowerCAmelCase					:						str					=					Image.fromarray(np.moveaxis(A_      ,							0      ,							-1					)					)

       return image_input
 def 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					self.get_tokenizer()
       __lowerCAmelCase					:						List[Any]					=					self.get_image_processor()

       __lowerCAmelCase					:						List[Any]					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)
       processor.save_pretrained(self.tmpdirname					)
       __lowerCAmelCase					:						Union[str, Any]					=					MgpstrProcessor.from_pretrained(self.tmpdirname      ,							use_fast=A_					)

       self.assertEqual(processor.char_tokenizer.get_vocab()      ,							tokenizer.get_vocab()					)
       self.assertIsInstance(processor.char_tokenizer      ,							A_					)

       self.assertEqual(processor.image_processor.to_json_string()      ,							image_processor.to_json_string()					)
       self.assertIsInstance(processor.image_processor      ,							A_					)
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					self.get_tokenizer()
       __lowerCAmelCase					:						Union[str, Any]					=					self.get_image_processor()

       __lowerCAmelCase					:						List[Any]					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)
       processor.save_pretrained(self.tmpdirname					)

       __lowerCAmelCase					:						List[Any]					=					self.get_tokenizer(bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''					)
       __lowerCAmelCase					:						int					=					self.get_image_processor(do_normalize=A_      ,							padding_value=1.0					)

       __lowerCAmelCase					:						int					=					MgpstrProcessor.from_pretrained(
           self.tmpdirname      ,							bos_token='''(BOS)'''      ,							eos_token='''(EOS)'''      ,							do_normalize=A_      ,							padding_value=1.0					)

       self.assertEqual(processor.char_tokenizer.get_vocab()      ,							tokenizer_add_kwargs.get_vocab()					)
       self.assertIsInstance(processor.char_tokenizer      ,							A_					)

       self.assertEqual(processor.image_processor.to_json_string()      ,							image_processor_add_kwargs.to_json_string()					)
       self.assertIsInstance(processor.image_processor      ,							A_					)
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					self.get_image_processor()
       __lowerCAmelCase					:						Optional[Any]					=					self.get_tokenizer()

       __lowerCAmelCase					:						int					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						Optional[int]					=					self.prepare_image_inputs()

       __lowerCAmelCase					:						Optional[Any]					=					image_processor(A_      ,							return_tensors='''np'''					)
       __lowerCAmelCase					:						Tuple					=					processor(images=A_      ,							return_tensors='''np'''					)

       for key in input_image_proc.keys():
             self.assertAlmostEqual(input_image_proc[key].sum()      ,							input_processor[key].sum()      ,							delta=1e-2					)
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					self.get_image_processor()
       __lowerCAmelCase					:						Union[str, Any]					=					self.get_tokenizer()

       __lowerCAmelCase					:						Optional[Any]					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						Any					=					'''test'''

       __lowerCAmelCase					:						Dict					=					processor(text=A_					)

       __lowerCAmelCase					:						str					=					tokenizer(A_					)
       for key in encoded_tok.keys():
             self.assertListEqual(encoded_tok[key]      ,							encoded_processor[key]					)
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					self.get_image_processor()
       __lowerCAmelCase					:						Any					=					self.get_tokenizer()

       __lowerCAmelCase					:						str					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						List[Any]					=					'''test'''
       __lowerCAmelCase					:						int					=					self.prepare_image_inputs()

       __lowerCAmelCase					:						int					=					processor(text=A_      ,							images=A_					)

       self.assertListEqual(list(inputs.keys()					)      ,							['''pixel_values''', '''labels''']					)

       # test if it raises when no input is passed
       with pytest.raises(A_					):
             processor()
 def 	UpperCamelCase__  (							self					)					->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					self.get_image_processor()
       __lowerCAmelCase					:						int					=					self.get_tokenizer()

       __lowerCAmelCase					:						Any					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						List[str]					=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9], [3, 4, 3, 1, 1, 8, 9]]

       __lowerCAmelCase					:						Optional[int]					=					processor.char_decode(A_					)
       __lowerCAmelCase					:						Tuple					=					tokenizer.batch_decode(A_					)
       __lowerCAmelCase					:						Any					=					[seq.replace(''' '''      ,							''''''					) for seq in decoded_tok]

       self.assertListEqual(A_      ,							A_					)
 def 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					self.get_image_processor()
       __lowerCAmelCase					:						Any					=					self.get_tokenizer()

       __lowerCAmelCase					:						int					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						Union[str, Any]					=					None
       __lowerCAmelCase					:						Optional[Any]					=					self.prepare_image_inputs()

       __lowerCAmelCase					:						List[Any]					=					processor(text=A_      ,							images=A_					)

       self.assertListEqual(list(inputs.keys()					)      ,							processor.model_input_names					)





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



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					self.get_image_processor()
       __lowerCAmelCase					:						List[str]					=					self.get_tokenizer()

       __lowerCAmelCase					:						Any					=					MgpstrProcessor(tokenizer=A_      ,							image_processor=A_					)

       __lowerCAmelCase					:						List[Any]					=					torch.randn(1      ,							27      ,							38					)
       __lowerCAmelCase					:						Optional[int]					=					torch.randn(1      ,							27      ,							5_0257					)
       __lowerCAmelCase					:						Optional[Any]					=					torch.randn(1      ,							27      ,							3_0522					)

       __lowerCAmelCase					:						List[str]					=					processor.batch_decode([char_input, bpe_input, wp_input]					)

       self.assertListEqual(list(results.keys()					)      ,							['''generated_text''', '''scores''', '''char_preds''', '''bpe_preds''', '''wp_preds''']					)


