

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_			(				lowerCamelCase_					,     unittest.TestCase	):



      """simple docstring"""







      __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 			_lowerCAmelCase	( self	:				Union[str, Any]							):
           return 32
      @property
      def 			_lowerCAmelCase	( self	:				Tuple							):
           return 32
      @property
      def 			_lowerCAmelCase	( self	:				Union[str, Any]							):
           return self.time_input_dim
      @property
      def 			_lowerCAmelCase	( self	:				Optional[int]							):
           return self.time_input_dim * 4
      @property
      def 			_lowerCAmelCase	( self	:				str							):
           return 100
      @property
      def 			_lowerCAmelCase	( self	:				Tuple							):
           SCREAMING_SNAKE_CASE 		=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)
           return tokenizer
      @property
      def 			_lowerCAmelCase	( self	:				str							):
           torch.manual_seed(0							)
           SCREAMING_SNAKE_CASE 		=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=37     ,layer_norm_eps=1e-05     ,num_attention_heads=4     ,num_hidden_layers=5     ,pad_token_id=1     ,vocab_size=1000     ,)
           return CLIPTextModelWithProjection(UpperCamelCase_							)
      @property
      def 			_lowerCAmelCase	( self	:				Dict							):
           torch.manual_seed(0							)

           SCREAMING_SNAKE_CASE 		={
               'num_attention_heads': 2,
               'attention_head_dim': 12,
               'embedding_dim': self.text_embedder_hidden_size,
               'num_layers': 1,
           }

           SCREAMING_SNAKE_CASE 		=PriorTransformer(**UpperCamelCase_							)
           return model
      @property
      def 			_lowerCAmelCase	( self	:				Optional[int]							):
           torch.manual_seed(0							)

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

           SCREAMING_SNAKE_CASE 		=UnCLIPTextProjModel(**UpperCamelCase_							)
           return model
      @property
      def 			_lowerCAmelCase	( self	:				List[str]							):
           torch.manual_seed(0							)

           SCREAMING_SNAKE_CASE 		={
               'sample_size': 32,
               # 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',
           }

           SCREAMING_SNAKE_CASE 		=UNetaDConditionModel(**UpperCamelCase_							)
           return model
      @property
      def 			_lowerCAmelCase	( self	:				Optional[int]							):
           return {
               "sample_size": 64,
               "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 			_lowerCAmelCase	( self	:				str							):
           torch.manual_seed(0							)

           SCREAMING_SNAKE_CASE 		=UNetaDModel(**self.dummy_super_res_kwargs							)
           return model
      @property
      def 			_lowerCAmelCase	( self	:				int							):
           torch.manual_seed(1							)

           SCREAMING_SNAKE_CASE 		=UNetaDModel(**self.dummy_super_res_kwargs							)
           return model
      def 			_lowerCAmelCase	( self	:				List[Any]							):
           SCREAMING_SNAKE_CASE 		=self.dummy_prior
           SCREAMING_SNAKE_CASE 		=self.dummy_decoder
           SCREAMING_SNAKE_CASE 		=self.dummy_text_proj
           SCREAMING_SNAKE_CASE 		=self.dummy_text_encoder
           SCREAMING_SNAKE_CASE 		=self.dummy_tokenizer
           SCREAMING_SNAKE_CASE 		=self.dummy_super_res_first
           SCREAMING_SNAKE_CASE 		=self.dummy_super_res_last

           SCREAMING_SNAKE_CASE 		=UnCLIPScheduler(
               variance_type='fixed_small_log'     ,prediction_type='sample'     ,num_train_timesteps=1000     ,clip_sample_range=5.0     ,)

           SCREAMING_SNAKE_CASE 		=UnCLIPScheduler(
               variance_type='learned_range'     ,prediction_type='epsilon'     ,num_train_timesteps=1000     ,)

           SCREAMING_SNAKE_CASE 		=UnCLIPScheduler(
               variance_type='fixed_small_log'     ,prediction_type='epsilon'     ,num_train_timesteps=1000     ,)

           SCREAMING_SNAKE_CASE 		={
               '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 			_lowerCAmelCase	( self	:				Any     ,snake_case	:				List[str]     ,snake_case	:				List[str]=0							):





           if str(UpperCamelCase_							).startswith('mps'							):
                SCREAMING_SNAKE_CASE 		=torch.manual_seed(UpperCamelCase_							)
           else:
                SCREAMING_SNAKE_CASE 		=torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
           SCREAMING_SNAKE_CASE 		={
               '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 			_lowerCAmelCase	( self	:				int							):
           SCREAMING_SNAKE_CASE 		='cpu'

           SCREAMING_SNAKE_CASE 		=self.get_dummy_components()

           SCREAMING_SNAKE_CASE 		=self.pipeline_class(**UpperCamelCase_							)
           SCREAMING_SNAKE_CASE 		=pipe.to(UpperCamelCase_							)

           pipe.set_progress_bar_config(disable=UpperCamelCase_							)

           SCREAMING_SNAKE_CASE 		=pipe(**self.get_dummy_inputs(UpperCamelCase_							)							)
           SCREAMING_SNAKE_CASE 		=output.images

           SCREAMING_SNAKE_CASE 		=pipe(
               **self.get_dummy_inputs(UpperCamelCase_							)     ,return_dict=UpperCamelCase_     ,)[0]

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

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

           SCREAMING_SNAKE_CASE 		=np.array(
               [
                   0.9_997,
                   0.9_988,
                   0.0_028,
                   0.9_997,
                   0.9_984,
                   0.9_965,
                   0.0_029,
                   0.9_986,
                   0.0_025,
               ]							)

           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 			_lowerCAmelCase	( self	:				Dict							):
           SCREAMING_SNAKE_CASE 		=torch.device('cpu'							)



           class 						a_			:



                 """simple docstring"""







                 __UpperCAmelCase     						= 1







           SCREAMING_SNAKE_CASE 		=self.get_dummy_components()

           SCREAMING_SNAKE_CASE 		=self.pipeline_class(**UpperCamelCase_							)
           SCREAMING_SNAKE_CASE 		=pipe.to(UpperCamelCase_							)

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

           SCREAMING_SNAKE_CASE 		=torch.Generator(device=UpperCamelCase_							).manual_seed(0							)
           SCREAMING_SNAKE_CASE 		=prior.dtype
           SCREAMING_SNAKE_CASE 		=1

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

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

           pipe.set_progress_bar_config(disable=UpperCamelCase_							)

           SCREAMING_SNAKE_CASE 		='this is a prompt example'

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

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

           SCREAMING_SNAKE_CASE 		=torch.Generator(device=UpperCamelCase_							).manual_seed(0							)
           SCREAMING_SNAKE_CASE 		=pipe(
               generator=UpperCamelCase_     ,prior_num_inference_steps=2     ,decoder_num_inference_steps=2     ,super_res_num_inference_steps=2     ,prior_latents=UpperCamelCase_     ,decoder_latents=UpperCamelCase_     ,super_res_latents=UpperCamelCase_     ,text_model_output=UpperCamelCase_     ,text_attention_mask=UpperCamelCase_     ,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 			_lowerCAmelCase	( self	:				List[Any]							):
           SCREAMING_SNAKE_CASE 		=torch_device == 'cpu'

           self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase_     ,expected_max_diff=0.01							)
      @skip_mps
      def 			_lowerCAmelCase	( self	:				Tuple							):
           SCREAMING_SNAKE_CASE 		=torch_device == 'cpu'
           SCREAMING_SNAKE_CASE 		=True
           SCREAMING_SNAKE_CASE 		=[
               'prior_num_inference_steps',
               'decoder_num_inference_steps',
               'super_res_num_inference_steps',
           ]

           self._test_inference_batch_single_identical(
               test_max_difference=UpperCamelCase_     ,relax_max_difference=UpperCamelCase_     ,additional_params_copy_to_batched_inputs=UpperCamelCase_     ,)
      def 			_lowerCAmelCase	( self	:				Optional[Any]							):
           SCREAMING_SNAKE_CASE 		=[
               'prior_num_inference_steps',
               'decoder_num_inference_steps',
               'super_res_num_inference_steps',
           ]

           if torch_device == "mps":
                # TODO: MPS errors with larger batch sizes
                SCREAMING_SNAKE_CASE 		=[2, 3]
                self._test_inference_batch_consistent(
                    batch_sizes=UpperCamelCase_     ,additional_params_copy_to_batched_inputs=UpperCamelCase_     ,)
           else:
                self._test_inference_batch_consistent(
                    additional_params_copy_to_batched_inputs=UpperCamelCase_							)
      @skip_mps
      def 			_lowerCAmelCase	( self	:				int							):
           return super().test_dict_tuple_outputs_equivalent()
      @skip_mps
      def 			_lowerCAmelCase	( self	:				Any							):
           return super().test_save_load_local()


      @skip_mps
      def 			_lowerCAmelCase	( self	:				Optional[int]							):
           return super().test_save_load_optional_components()







@nightly
class 						a_			(				unittest.TestCase	):



      """simple docstring"""







      def 			_lowerCAmelCase	( self	:				Optional[Any]							):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


      def 			_lowerCAmelCase	( self	:				Dict							):
           SCREAMING_SNAKE_CASE 		=load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/unclip/karlo_v1_alpha_horse_cpu.npy'							)

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

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

           SCREAMING_SNAKE_CASE 		=output.images[0]

           assert image.shape == (256, 256, 3)
           assert np.abs(expected_image - image							).max() < 1e-1







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



      """simple docstring"""







      def 			_lowerCAmelCase	( self	:				List[Any]							):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()
      def 			_lowerCAmelCase	( self	:				Dict							):
           SCREAMING_SNAKE_CASE 		=load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/unclip/karlo_v1_alpha_horse_fp16.npy'							)

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

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

           SCREAMING_SNAKE_CASE 		=output.images[0]

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

           assert_mean_pixel_difference(UpperCamelCase_     ,UpperCamelCase_							)


      def 			_lowerCAmelCase	( self	:				List[str]							):
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

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

           SCREAMING_SNAKE_CASE 		=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'     ,)

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






'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_A					:     List[Any]      ={'''configuration_umt5''': ['''UMT5Config''', '''UMT5OnnxConfig''']}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_A					:     Optional[Any]      =[
		    '''UMT5EncoderModel''',
		    '''UMT5ForConditionalGeneration''',
		    '''UMT5ForQuestionAnswering''',
		    '''UMT5Model''',
		    '''UMT5PreTrainedModel''',
		]

if TYPE_CHECKING:
		from .configuration_umta import UMTaConfig, UMTaOnnxConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_umta import (
				    UMTaEncoderModel,
				    UMTaForConditionalGeneration,
				    UMTaForQuestionAnswering,
				    UMTaModel,
				    UMTaPreTrainedModel,
				)
else:
		import sys

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