





import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, 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_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class      A_	(a_					,							unittest.TestCase      ):
	UpperCAmelCase__ =      UnCLIPPipeline
	UpperCAmelCase__ =      TEXT_TO_IMAGE_PARAMS - {
	    '''negative_prompt''',
	    '''height''',
	    '''width''',
	    '''negative_prompt_embeds''',
	    '''guidance_scale''',
	    '''prompt_embeds''',
	    '''cross_attention_kwargs''',
	}
	UpperCAmelCase__ =      TEXT_TO_IMAGE_BATCH_PARAMS
	UpperCAmelCase__ =      [
	    '''generator''',
	    '''return_dict''',
	    '''prior_num_inference_steps''',
	    '''decoder_num_inference_steps''',
	    '''super_res_num_inference_steps''',
	]
	UpperCAmelCase__ =      False
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return 3_2
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return 3_2
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return self.time_input_dim
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return self.time_input_dim * 4
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return 1_0_0
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)
				return tokenizer
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase       		=    CLIPTextConfig(
				    bos_token_id=0							, eos_token_id=2							, hidden_size=self.text_embedder_hidden_size							, projection_dim=self.text_embedder_hidden_size							, 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							, )
				return CLIPTextModelWithProjection(_A		)
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)

				UpperCAmelCase       		=    {
				    '''num_attention_heads''': 2,
				    '''attention_head_dim''': 1_2,
				    '''embedding_dim''': self.text_embedder_hidden_size,
				    '''num_layers''': 1,
				}

				UpperCAmelCase       		=    PriorTransformer(**_A		)
				return model
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)

				UpperCAmelCase       		=    {
				    '''clip_embeddings_dim''': self.text_embedder_hidden_size,
				    '''time_embed_dim''': self.time_embed_dim,
				    '''cross_attention_dim''': self.cross_attention_dim,
				}

				UpperCAmelCase       		=    UnCLIPTextProjModel(**_A		)
				return model
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)

				UpperCAmelCase       		=    {
				    '''sample_size''': 3_2,
				    # RGB in channels
				    '''in_channels''': 3,
				    # Out channels is double in channels because predicts mean and variance
				    '''out_channels''': 6,
				    '''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,
				    '''cross_attention_dim''': self.cross_attention_dim,
				    '''attention_head_dim''': 4,
				    '''resnet_time_scale_shift''': '''scale_shift''',
				    '''class_embed_type''': '''identity''',
				}

				UpperCAmelCase       		=    UNetaDConditionModel(**_A		)
				return model
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return {
				    "sample_size": 6_4,
				    "layers_per_block": 1,
				    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
				    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
				    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
				    "in_channels": 6,
				    "out_channels": 3,
				}
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(0		)

				UpperCAmelCase       		=    UNetaDModel(**self.dummy_super_res_kwargs		)
				return model
	@property
	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.manual_seed(1		)

				UpperCAmelCase       		=    UNetaDModel(**self.dummy_super_res_kwargs		)
				return model
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    self.dummy_prior
				UpperCAmelCase       		=    self.dummy_decoder
				UpperCAmelCase       		=    self.dummy_text_proj
				UpperCAmelCase       		=    self.dummy_text_encoder
				UpperCAmelCase       		=    self.dummy_tokenizer
				UpperCAmelCase       		=    self.dummy_super_res_first
				UpperCAmelCase       		=    self.dummy_super_res_last

				UpperCAmelCase       		=    UnCLIPScheduler(
				    variance_type='''fixed_small_log'''							, prediction_type='''sample'''							, num_train_timesteps=1_0_0_0							, clip_sample_range=5.0							, )

				UpperCAmelCase       		=    UnCLIPScheduler(
				    variance_type='''learned_range'''							, prediction_type='''epsilon'''							, num_train_timesteps=1_0_0_0							, )

				UpperCAmelCase       		=    UnCLIPScheduler(
				    variance_type='''fixed_small_log'''							, prediction_type='''epsilon'''							, num_train_timesteps=1_0_0_0							, )

				UpperCAmelCase       		=    {
				    '''prior''': prior,
				    '''decoder''': decoder,
				    '''text_proj''': text_proj,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''super_res_first''': super_res_first,
				    '''super_res_last''': super_res_last,
				    '''prior_scheduler''': prior_scheduler,
				    '''decoder_scheduler''': decoder_scheduler,
				    '''super_res_scheduler''': super_res_scheduler,
				}

				return components
	def 					_lowercase    (   self							, _A							, _A=0		):






				'''simple docstring'''







				if str(_A		).startswith('''mps'''		):
							UpperCAmelCase       		=    torch.manual_seed(_A		)
				else:
							UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(_A		)
				UpperCAmelCase       		=    {
				    '''prompt''': '''horse''',
				    '''generator''': generator,
				    '''prior_num_inference_steps''': 2,
				    '''decoder_num_inference_steps''': 2,
				    '''super_res_num_inference_steps''': 2,
				    '''output_type''': '''numpy''',
				}
				return inputs
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    '''cpu'''

				UpperCAmelCase       		=    self.get_dummy_components()

				UpperCAmelCase       		=    self.pipeline_class(**_A		)
				UpperCAmelCase       		=    pipe.to(_A		)

				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase       		=    pipe(**self.get_dummy_inputs(_A		)		)
				UpperCAmelCase       		=    output.images

				UpperCAmelCase       		=    pipe(
				    **self.get_dummy_inputs(_A		)							, return_dict=_A							, )[0]

				UpperCAmelCase       		=    image[0, -3:, -3:, -1]
				UpperCAmelCase       		=    image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 6_4, 6_4, 3)

				UpperCAmelCase       		=    np.array(
				    [
				        0.99_97,
				        0.99_88,
				        0.00_28,
				        0.99_97,
				        0.99_84,
				        0.99_65,
				        0.00_29,
				        0.99_86,
				        0.00_25,
				    ]		)

				assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    torch.device('''cpu'''		)

				class      A_	:
					UpperCAmelCase__ =      1
				UpperCAmelCase       		=    self.get_dummy_components()

				UpperCAmelCase       		=    self.pipeline_class(**_A		)
				UpperCAmelCase       		=    pipe.to(_A		)

				UpperCAmelCase       		=    components['''prior''']
				UpperCAmelCase       		=    components['''decoder''']
				UpperCAmelCase       		=    components['''super_res_first''']
				UpperCAmelCase       		=    components['''tokenizer''']
				UpperCAmelCase       		=    components['''text_encoder''']

				UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(0		)
				UpperCAmelCase       		=    prior.dtype
				UpperCAmelCase       		=    1

				UpperCAmelCase       		=    (batch_size, prior.config.embedding_dim)
				UpperCAmelCase       		=    pipe.prepare_latents(
				    _A							, dtype=_A							, device=_A							, generator=_A							, latents=_A							, scheduler=DummyScheduler()		)
				UpperCAmelCase       		=    (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
				UpperCAmelCase       		=    pipe.prepare_latents(
				    _A							, dtype=_A							, device=_A							, generator=_A							, latents=_A							, scheduler=DummyScheduler()		)

				UpperCAmelCase       		=    (
				    batch_size,
				    super_res_first.config.in_channels // 2,
				    super_res_first.config.sample_size,
				    super_res_first.config.sample_size,
				)
				UpperCAmelCase       		=    pipe.prepare_latents(
				    _A							, dtype=_A							, device=_A							, generator=_A							, latents=_A							, scheduler=DummyScheduler()		)

				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase       		=    '''this is a prompt example'''

				UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(0		)
				UpperCAmelCase       		=    pipe(
				    [prompt]							, generator=_A							, prior_num_inference_steps=2							, decoder_num_inference_steps=2							, super_res_num_inference_steps=2							, prior_latents=_A							, decoder_latents=_A							, super_res_latents=_A							, output_type='''np'''							, )
				UpperCAmelCase       		=    output.images

				UpperCAmelCase       		=    tokenizer(
				    _A							, padding='''max_length'''							, max_length=tokenizer.model_max_length							, return_tensors='''pt'''							, )
				UpperCAmelCase       		=    text_encoder(text_inputs.input_ids		)
				UpperCAmelCase       		=    text_inputs.attention_mask

				UpperCAmelCase       		=    torch.Generator(device=_A		).manual_seed(0		)
				UpperCAmelCase       		=    pipe(
				    generator=_A							, prior_num_inference_steps=2							, decoder_num_inference_steps=2							, super_res_num_inference_steps=2							, prior_latents=_A							, decoder_latents=_A							, super_res_latents=_A							, text_model_output=_A							, text_attention_mask=_A							, output_type='''np'''							, )[0]

				# make sure passing text embeddings manually is identical
				assert np.abs(image - image_from_text		).max() < 1E-4
	@skip_mps
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    torch_device == '''cpu'''

				self._test_attention_slicing_forward_pass(test_max_difference=_A							, expected_max_diff=0.01		)
	@skip_mps
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    torch_device == '''cpu'''
				UpperCAmelCase       		=    True
				UpperCAmelCase       		=    [
				    '''prior_num_inference_steps''',
				    '''decoder_num_inference_steps''',
				    '''super_res_num_inference_steps''',
				]

				self._test_inference_batch_single_identical(
				    test_max_difference=_A							, relax_max_difference=_A							, additional_params_copy_to_batched_inputs=_A							, )
	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    [
				    '''prior_num_inference_steps''',
				    '''decoder_num_inference_steps''',
				    '''super_res_num_inference_steps''',
				]

				if torch_device == "mps":
							# TODO: MPS errors with larger batch sizes
							UpperCAmelCase       		=    [2, 3]
							self._test_inference_batch_consistent(
							    batch_sizes=_A							, additional_params_copy_to_batched_inputs=_A							, )
				else:
							self._test_inference_batch_consistent(
							    additional_params_copy_to_batched_inputs=_A		)
	@skip_mps
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return super().test_dict_tuple_outputs_equivalent()
	@skip_mps
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return super().test_save_load_local()







	@skip_mps
	def 					_lowercase    (   self		):






				'''simple docstring'''







				return super().test_save_load_optional_components()

@nightly
class      A_	(unittest.TestCase      ):
	def 					_lowercase    (   self		):






				'''simple docstring'''







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







	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/unclip/karlo_v1_alpha_horse_cpu.npy'''		)

				UpperCAmelCase       		=    UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''		)
				pipeline.set_progress_bar_config(disable=_A		)

				UpperCAmelCase       		=    torch.manual_seed(0		)
				UpperCAmelCase       		=    pipeline(
				    '''horse'''							, num_images_per_prompt=1							, generator=_A							, output_type='''np'''							, )

				UpperCAmelCase       		=    output.images[0]

				assert image.shape == (2_5_6, 2_5_6, 3)
				assert np.abs(expected_image - image		).max() < 1E-1

@slow
@require_torch_gpu
class      A_	(unittest.TestCase      ):
	def 					_lowercase    (   self		):






				'''simple docstring'''







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






				'''simple docstring'''







				UpperCAmelCase       		=    load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/unclip/karlo_v1_alpha_horse_fp16.npy'''		)

				UpperCAmelCase       		=    UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''							, torch_dtype=torch.floataa		)
				UpperCAmelCase       		=    pipeline.to(_A		)
				pipeline.set_progress_bar_config(disable=_A		)

				UpperCAmelCase       		=    torch.Generator(device='''cpu'''		).manual_seed(0		)
				UpperCAmelCase       		=    pipeline(
				    '''horse'''							, generator=_A							, output_type='''np'''							, )

				UpperCAmelCase       		=    output.images[0]

				assert image.shape == (2_5_6, 2_5_6, 3)

				assert_mean_pixel_difference(_A							, _A		)







	def 					_lowercase    (   self		):






				'''simple docstring'''







				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

				UpperCAmelCase       		=    UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''							, torch_dtype=torch.floataa		)
				UpperCAmelCase       		=    pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)
				pipe.enable_attention_slicing()
				pipe.enable_sequential_cpu_offload()

				UpperCAmelCase       		=    pipe(
				    '''horse'''							, num_images_per_prompt=1							, prior_num_inference_steps=2							, decoder_num_inference_steps=2							, super_res_num_inference_steps=2							, output_type='''np'''							, )

				UpperCAmelCase       		=    torch.cuda.max_memory_allocated()
				# make sure that less than 7 GB is allocated
				assert mem_bytes < 7 * 1_0**9










import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


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

class      A_	(a_      ):
	UpperCAmelCase__ =      '''speech-encoder-decoder'''
	UpperCAmelCase__ =      True
	def __init__(   self							, **_A		):






				'''simple docstring'''







				super().__init__(**_A		)
				if "encoder" not in kwargs or "decoder" not in kwargs:
							raise ValueError(
							    F"""A configuraton of type {self.model_type} cannot be instantiated because not both `encoder` and"""
							    F""" `decoder` sub-configurations are passed, but only {kwargs}"""		)

				UpperCAmelCase       		=    kwargs.pop('''encoder'''		)
				UpperCAmelCase       		=    encoder_config.pop('''model_type'''		)
				UpperCAmelCase       		=    kwargs.pop('''decoder'''		)
				UpperCAmelCase       		=    decoder_config.pop('''model_type'''		)

				UpperCAmelCase       		=    AutoConfig.for_model(_A							, **_A		)
				UpperCAmelCase       		=    AutoConfig.for_model(_A							, **_A		)
				UpperCAmelCase       		=    True
	@classmethod
	def 					_lowercase    (   cls							, _A							, _A							, **_A		):






				'''simple docstring'''







				logger.info('''Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''		)
				UpperCAmelCase       		=    True
				UpperCAmelCase       		=    True

				return cls(encoder=encoder_config.to_dict()							, decoder=decoder_config.to_dict()							, **_A		)







	def 					_lowercase    (   self		):






				'''simple docstring'''







				UpperCAmelCase       		=    copy.deepcopy(self.__dict__		)
				UpperCAmelCase       		=    self.encoder.to_dict()
				UpperCAmelCase       		=    self.decoder.to_dict()
				UpperCAmelCase       		=    self.__class__.model_type
				return output



