



"""simple docstring"""

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class     lowerCAmelCase__						(	_lowerCamelCase							,							unittest.TestCase ):


	'''simple docstring'''







	_SCREAMING_SNAKE_CASE	:		str  						= KandinskyVaaImgaImgPipeline
	_SCREAMING_SNAKE_CASE	:		int  						= ['''image_embeds''', '''negative_image_embeds''', '''image''']
	_SCREAMING_SNAKE_CASE	:		Optional[int]  						= [
	    '''image_embeds''',
	    '''negative_image_embeds''',
	    '''image''',
	]
	_SCREAMING_SNAKE_CASE	:		Optional[int]  						= [
	    '''generator''',
	    '''height''',
	    '''width''',
	    '''strength''',
	    '''guidance_scale''',
	    '''num_inference_steps''',
	    '''return_dict''',
	    '''guidance_scale''',
	    '''num_images_per_prompt''',
	    '''output_type''',
	    '''return_dict''',
	]
	_SCREAMING_SNAKE_CASE	:		List[Any]  						= False





	@property
	def 			_lowerCAmelCase     (      self				:       Optional[int]    )						->	Any:





					"""simple docstring"""


					return 32





	@property
	def 			_lowerCAmelCase     (      self				:       List[str]    )						->	int:





					"""simple docstring"""


					return 32





	@property
	def 			_lowerCAmelCase     (      self				:       List[str]    )						->	int:





					"""simple docstring"""


					return self.time_input_dim





	@property
	def 			_lowerCAmelCase     (      self				:       Tuple    )						->	List[Any]:





					"""simple docstring"""


					return self.time_input_dim * 4





	@property
	def 			_lowerCAmelCase     (      self				:       int    )						->	Tuple:





					"""simple docstring"""


					return 100





	@property
	def 			_lowerCAmelCase     (      self				:       List[Any]    )						->	List[str]:





					"""simple docstring"""


					torch.manual_seed(0    )

					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     {
					    'in_channels': 4,
					    # Out channels is double in channels because predicts mean and variance
					    'out_channels': 8,
					    'addition_embed_type': 'image',
					    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
					    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
					    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
					    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
					    'layers_per_block': 1,
					    'encoder_hid_dim': self.text_embedder_hidden_size,
					    'encoder_hid_dim_type': 'image_proj',
					    'cross_attention_dim': self.cross_attention_dim,
					    'attention_head_dim': 4,
					    'resnet_time_scale_shift': 'scale_shift',
					    'class_embed_type': None,
					}

					SCREAMING_SNAKE_CASE						:							int							=     UNetaDConditionModel(**_SCREAMING_SNAKE_CASE    )
					return model





	@property
	def 			_lowerCAmelCase     (      self				:       Optional[int]    )						->	Optional[Any]:





					"""simple docstring"""


					return {
					    "block_out_channels": [32, 64],
					    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
					    "in_channels": 3,
					    "latent_channels": 4,
					    "layers_per_block": 1,
					    "norm_num_groups": 8,
					    "norm_type": "spatial",
					    "num_vq_embeddings": 12,
					    "out_channels": 3,
					    "up_block_types": [
					        "AttnUpDecoderBlock2D",
					        "UpDecoderBlock2D",
					    ],
					    "vq_embed_dim": 4,
					}





	@property
	def 			_lowerCAmelCase     (      self				:       Any    )						->	List[str]:





					"""simple docstring"""


					torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							List[str]							=     VQModel(**self.dummy_movq_kwargs    )
					return model





	def 			_lowerCAmelCase     (      self				:       int    )						->	int:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[int]							=     self.dummy_unet
					SCREAMING_SNAKE_CASE						:							List[str]							=     self.dummy_movq

					SCREAMING_SNAKE_CASE						:							List[Any]							=     {
					    'num_train_timesteps': 1_000,
					    'beta_schedule': 'linear',
					    'beta_start': 0.0_0_0_8_5,
					    'beta_end': 0.0_1_2,
					    'clip_sample': False,
					    'set_alpha_to_one': False,
					    'steps_offset': 0,
					    'prediction_type': 'epsilon',
					    'thresholding': False,
					}

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     DDIMScheduler(**_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     {
					    'unet': unet,
					    'scheduler': scheduler,
					    'movq': movq,
					}

					return components





	def 			_lowerCAmelCase     (      self				:       Optional[int]     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=0    )						->	Union[str, Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							int							=     floats_tensor((1, self.text_embedder_hidden_size)     ,       rng=random.Random(_SCREAMING_SNAKE_CASE    )    ).to(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     floats_tensor((1, self.text_embedder_hidden_size)     ,       rng=random.Random(seed + 1    )    ).to(
					    _SCREAMING_SNAKE_CASE    )
					# create init_image
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     floats_tensor((1, 3, 64, 64)     ,       rng=random.Random(_SCREAMING_SNAKE_CASE    )    ).to(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     image.cpu().permute(0     ,       2     ,       3     ,       1    )[0]
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE    )    ).convert('RGB'    ).resize((256, 256)    )

					if str(_SCREAMING_SNAKE_CASE    ).startswith('mps'    ):
									SCREAMING_SNAKE_CASE						:							str							=     torch.manual_seed(_SCREAMING_SNAKE_CASE    )
					else:
									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     torch.Generator(device=_SCREAMING_SNAKE_CASE    ).manual_seed(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     {
					    'image': init_image,
					    'image_embeds': image_embeds,
					    'negative_image_embeds': negative_image_embeds,
					    'generator': generator,
					    'height': 64,
					    'width': 64,
					    'num_inference_steps': 10,
					    'guidance_scale': 7.0,
					    'strength': 0.2,
					    'output_type': 'np',
					}
					return inputs







	def 			_lowerCAmelCase     (      self				:       Tuple    )						->	int:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							str							=     'cpu'

					SCREAMING_SNAKE_CASE						:							Tuple							=     self.get_dummy_components()

					SCREAMING_SNAKE_CASE						:							Tuple							=     self.pipeline_class(**_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							Tuple							=     pipe.to(_SCREAMING_SNAKE_CASE    )

					pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE    )    )
					SCREAMING_SNAKE_CASE						:							Dict							=     output.images

					SCREAMING_SNAKE_CASE						:							List[str]							=     pipe(
					    **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE    )     ,       return_dict=_SCREAMING_SNAKE_CASE     ,       )[0]

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     image[0, -3:, -3:, -1]
					SCREAMING_SNAKE_CASE						:							List[Any]							=     image_from_tuple[0, -3:, -3:, -1]

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

					SCREAMING_SNAKE_CASE						:							Tuple							=     np.array(
					    [0.6_1_9_9_7_7_8, 0.6_3_9_8_4_4_0_6, 0.4_6_1_4_5_7_8_5, 0.6_2_9_4_4_9_8_4, 0.5_6_2_2_2_1_5, 0.4_7_3_0_6_1_3_2, 0.4_7_4_4_1_4_5_6, 0.4_6_0_7_6_0_6, 0.4_8_7_1_9_2_6_3]    )
					assert (
					    np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
					), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
					assert (
					    np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2
					), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




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


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       Dict    )						->	int:





					"""simple docstring"""


					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()







	def 			_lowerCAmelCase     (      self				:       Union[str, Any]    )						->	Union[str, Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
					    '/kandinskyv22/kandinskyv22_img2img_frog.npy'    )

					SCREAMING_SNAKE_CASE						:							str							=     load_image(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'    )
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     'A red cartoon frog, 4k'

					SCREAMING_SNAKE_CASE						:							Tuple							=     KandinskyVaaPriorPipeline.from_pretrained(
					    'kandinsky-community/kandinsky-2-2-prior'     ,       torch_dtype=torch.floataa    )
					pipe_prior.to(_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Tuple							=     KandinskyVaaImgaImgPipeline.from_pretrained(
					    'kandinsky-community/kandinsky-2-2-decoder'     ,       torch_dtype=torch.floataa    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     pipeline.to(_SCREAMING_SNAKE_CASE    )

					pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							int							=     torch.Generator(device='cpu'    ).manual_seed(0    )
					SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							List[str]							=     pipe_prior(
					    _SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       num_inference_steps=5     ,       negative_prompt=''     ,       ).to_tuple()

					SCREAMING_SNAKE_CASE						:							Tuple							=     pipeline(
					    image=_SCREAMING_SNAKE_CASE     ,       image_embeds=_SCREAMING_SNAKE_CASE     ,       negative_image_embeds=_SCREAMING_SNAKE_CASE     ,       generator=_SCREAMING_SNAKE_CASE     ,       num_inference_steps=100     ,       height=768     ,       width=768     ,       strength=0.2     ,       output_type='np'     ,       )

					SCREAMING_SNAKE_CASE						:							Dict							=     output.images[0]

					assert image.shape == (768, 768, 3)

					assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )










"""simple docstring"""

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


A_						:    Optional[Any]				=      {
    'configuration_squeezebert': [
        'SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'SqueezeBertConfig',
        'SqueezeBertOnnxConfig',
    ],
    'tokenization_squeezebert': ['SqueezeBertTokenizer'],
}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							A_						:    Tuple				=      ['SqueezeBertTokenizerFast']

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							A_						:    int				=      [
							    'SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'SqueezeBertForMaskedLM',
							    'SqueezeBertForMultipleChoice',
							    'SqueezeBertForQuestionAnswering',
							    'SqueezeBertForSequenceClassification',
							    'SqueezeBertForTokenClassification',
							    'SqueezeBertModel',
							    'SqueezeBertModule',
							    'SqueezeBertPreTrainedModel',
							]


if TYPE_CHECKING:
							from .configuration_squeezebert import (
							    SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    SqueezeBertConfig,
							    SqueezeBertOnnxConfig,
							)
							from .tokenization_squeezebert import SqueezeBertTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_squeezebert_fast import SqueezeBertTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_squeezebert import (
														    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    SqueezeBertForMaskedLM,
														    SqueezeBertForMultipleChoice,
														    SqueezeBertForQuestionAnswering,
														    SqueezeBertForSequenceClassification,
														    SqueezeBertForTokenClassification,
														    SqueezeBertModel,
														    SqueezeBertModule,
														    SqueezeBertPreTrainedModel,
														)

else:
							import sys

							A_						:    Optional[Any]				=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





