






"""simple docstring"""




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  __lowercase (		lowercase_     , unittest.TestCase					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			UnCLIPPipeline
     SCREAMING_SNAKE_CASE				  =			TEXT_TO_IMAGE_PARAMS - {
         "negative_prompt",
         "height",
         "width",
         "negative_prompt_embeds",
         "guidance_scale",
         "prompt_embeds",
         "cross_attention_kwargs",
     }
     SCREAMING_SNAKE_CASE				  =			TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE				  =			[
         "generator",
         "return_dict",
         "prior_num_inference_steps",
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]
     SCREAMING_SNAKE_CASE				  =			False




     @property
     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         return 32




     @property
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         return 32




     @property
     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



         return self.time_input_dim




     @property
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         return self.time_input_dim * 4




     @property
     def 						lowerCAmelCase_       (			self    :     Union[str, Any]			):


         """simple docstring"""



         return 100




     @property
     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



         __A             =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)
         return tokenizer




     @property
     def 						lowerCAmelCase_       (			self    :     Dict			):


         """simple docstring"""



         torch.manual_seed(0			)
         __A             =		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=1_000	,		)
         return CLIPTextModelWithProjection(UpperCamelCase_			)




     @property
     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         torch.manual_seed(0			)

         __A             =		{
             """num_attention_heads""": 2,
             """attention_head_dim""": 12,
             """embedding_dim""": self.text_embedder_hidden_size,
             """num_layers""": 1,
         }

         __A             =		PriorTransformer(**UpperCamelCase_			)
         return model




     @property
     def 						lowerCAmelCase_       (			self    :     Tuple			):


         """simple docstring"""



         torch.manual_seed(0			)

         __A             =		{
             """clip_embeddings_dim""": self.text_embedder_hidden_size,
             """time_embed_dim""": self.time_embed_dim,
             """cross_attention_dim""": self.cross_attention_dim,
         }

         __A             =		UnCLIPTextProjModel(**UpperCamelCase_			)
         return model




     @property
     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



         torch.manual_seed(0			)

         __A             =		{
             """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""",
         }

         __A             =		UNetaDConditionModel(**UpperCamelCase_			)
         return model




     @property
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         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    :     List[str]			):


         """simple docstring"""



         torch.manual_seed(0			)

         __A             =		UNetaDModel(**self.dummy_super_res_kwargs			)
         return model




     @property
     def 						lowerCAmelCase_       (			self    :     Optional[Any]			):


         """simple docstring"""



         torch.manual_seed(1			)

         __A             =		UNetaDModel(**self.dummy_super_res_kwargs			)
         return model




     def 						lowerCAmelCase_       (			self    :     List[str]			):


         """simple docstring"""



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

         __A             =		UnCLIPScheduler(
             variance_type="""fixed_small_log"""	,		prediction_type="""sample"""	,		num_train_timesteps=1_000	,		clip_sample_range=5.0	,		)

         __A             =		UnCLIPScheduler(
             variance_type="""learned_range"""	,		prediction_type="""epsilon"""	,		num_train_timesteps=1_000	,		)

         __A             =		UnCLIPScheduler(
             variance_type="""fixed_small_log"""	,		prediction_type="""epsilon"""	,		num_train_timesteps=1_000	,		)

         __A             =		{
             """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    :     int	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Optional[int]=0			):


         """simple docstring"""



         if str(UpperCamelCase_			).startswith("""mps"""			):
             __A             =		torch.manual_seed(UpperCamelCase_			)
         else:
             __A             =		torch.Generator(device=UpperCamelCase_			).manual_seed(UpperCamelCase_			)
         __A             =		{
             """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    :     str			):


         """simple docstring"""



         __A             =		"""cpu"""

         __A             =		self.get_dummy_components()

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

         pipe.set_progress_bar_config(disable=UpperCamelCase_			)

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

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

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

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

         __A             =		np.array(
             [
                 0.9997,
                 0.9988,
                 0.0028,
                 0.9997,
                 0.9984,
                 0.9965,
                 0.0029,
                 0.9986,
                 0.0025,
             ]			)

         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			):


         """simple docstring"""



         __A             =		torch.device("""cpu"""			)


         class  __lowercase :







              '''simple docstring'''
              SCREAMING_SNAKE_CASE				  =			1
         __A             =		self.get_dummy_components()

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

         __A             =		components["""prior"""]
         __A             =		components["""decoder"""]
         __A             =		components["""super_res_first"""]
         __A             =		components["""tokenizer"""]
         __A             =		components["""text_encoder"""]

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

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

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

         pipe.set_progress_bar_config(disable=UpperCamelCase_			)

         __A             =		"""this is a prompt example"""

         __A             =		torch.Generator(device=UpperCamelCase_			).manual_seed(0			)
         __A             =		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"""	,		)
         __A             =		output.images

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

         __A             =		torch.Generator(device=UpperCamelCase_			).manual_seed(0			)
         __A             =		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    :     Optional[Any]			):


         """simple docstring"""



         __A             =		torch_device == """cpu"""

         self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase_	,		expected_max_diff=0.01			)




     @skip_mps
     def 						lowerCAmelCase_       (			self    :     Dict			):


         """simple docstring"""



         __A             =		torch_device == """cpu"""
         __A             =		True
         __A             =		[
             """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[int]			):


         """simple docstring"""



         __A             =		[
             """prior_num_inference_steps""",
             """decoder_num_inference_steps""",
             """super_res_num_inference_steps""",
         ]

         if torch_device == "mps":
             # TODO: MPS errors with larger batch sizes
             __A             =		[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    :     Dict			):


         """simple docstring"""



         return super().test_dict_tuple_outputs_equivalent()




     @skip_mps
     def 						lowerCAmelCase_       (			self    :     Union[str, Any]			):


         """simple docstring"""



         return super().test_save_load_local()





     @skip_mps
     def 						lowerCAmelCase_       (			self    :     Union[str, Any]			):


         """simple docstring"""



         return super().test_save_load_optional_components()


@nightly
class  __lowercase (		unittest.TestCase					):







     '''simple docstring'''




     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



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





     def 						lowerCAmelCase_       (			self    :     str			):


         """simple docstring"""



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

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

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

         __A             =		output.images[0]

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


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







     '''simple docstring'''




     def 						lowerCAmelCase_       (			self    :     Tuple			):


         """simple docstring"""



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




     def 						lowerCAmelCase_       (			self    :     Dict			):


         """simple docstring"""



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

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

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

         __A             =		output.images[0]

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

         assert_mean_pixel_difference(UpperCamelCase_	,		UpperCamelCase_			)





     def 						lowerCAmelCase_       (			self    :     Optional[Any]			):


         """simple docstring"""



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

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

         __A             =		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"""	,		)

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












import copy
from typing import Dict

from ...generation.configuration_utils import GenerationConfig
from ...utils import logging


__a :    str								=				logging.get_logger(__name__)


class  __lowercase (		lowercase_					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			"semantic"




     def __init__(			self    :     Optional[int]	,		UpperCamelCase_    :     List[Any]=10_000	,		UpperCamelCase_    :     int=True	,		UpperCamelCase_    :     Dict=768	,		UpperCamelCase_    :     str=False	,		UpperCamelCase_    :     Dict=False	,		UpperCamelCase_    :     List[Any]=False	,		UpperCamelCase_    :     Any=False	,		UpperCamelCase_    :     Any=0.7	,		UpperCamelCase_    :     List[Any]=True	,		UpperCamelCase_    :     Dict=10_048	,		UpperCamelCase_    :     Tuple=129_595	,		UpperCamelCase_    :     List[str]=129_599	,		UpperCamelCase_    :     Optional[int]=10_000	,		UpperCamelCase_    :     Dict=256	,		UpperCamelCase_    :     int=49.9	,		**UpperCamelCase_    :     Any	,		):


         """simple docstring"""



         super().__init__(
             temperature=UpperCamelCase_	,		do_sample=UpperCamelCase_	,		eos_token_id=UpperCamelCase_	,		renormalize_logits=UpperCamelCase_	,		max_new_tokens=UpperCamelCase_	,		output_scores=UpperCamelCase_	,		return_dict_in_generate=UpperCamelCase_	,		output_hidden_states=UpperCamelCase_	,		output_attentions=UpperCamelCase_	,		**UpperCamelCase_	,		)

         __A             =		text_encoding_offset
         __A             =		text_pad_token
         __A             =		eos_token_id
         __A             =		semantic_infer_token
         __A             =		semantic_vocab_size
         __A             =		max_input_semantic_length
         __A             =		semantic_rate_hz


class  __lowercase (		lowercase_					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			"coarse_acoustics"




     def __init__(			self    :     int	,		UpperCamelCase_    :     List[Any]=True	,		UpperCamelCase_    :     int=False	,		UpperCamelCase_    :     Dict=False	,		UpperCamelCase_    :     Dict=False	,		UpperCamelCase_    :     Union[str, Any]=False	,		UpperCamelCase_    :     Any=0.7	,		UpperCamelCase_    :     Optional[Any]=True	,		UpperCamelCase_    :     str=12_048	,		UpperCamelCase_    :     int=75	,		UpperCamelCase_    :     Dict=2	,		UpperCamelCase_    :     List[str]=12_050	,		UpperCamelCase_    :     str=256	,		UpperCamelCase_    :     int = 630	,		UpperCamelCase_    :     int = 60	,		**UpperCamelCase_    :     Union[str, Any]	,		):


         """simple docstring"""



         super().__init__(
             temperature=UpperCamelCase_	,		do_sample=UpperCamelCase_	,		renormalize_logits=UpperCamelCase_	,		output_scores=UpperCamelCase_	,		return_dict_in_generate=UpperCamelCase_	,		output_hidden_states=UpperCamelCase_	,		output_attentions=UpperCamelCase_	,		**UpperCamelCase_	,		)

         __A             =		coarse_semantic_pad_token
         __A             =		coarse_rate_hz
         __A             =		n_coarse_codebooks
         __A             =		coarse_infer_token
         __A             =		max_coarse_input_length
         __A             =		max_coarse_history
         __A             =		sliding_window_len


class  __lowercase (		lowercase_					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			"fine_acoustics"




     def __init__(			self    :     Any	,		UpperCamelCase_    :     Dict=0.5	,		UpperCamelCase_    :     List[Any]=512	,		UpperCamelCase_    :     int=1_024	,		UpperCamelCase_    :     Union[str, Any]=8	,		**UpperCamelCase_    :     int	,		):


         """simple docstring"""



         super().__init__(temperature=UpperCamelCase_			)

         __A             =		max_fine_history_length
         __A             =		max_fine_input_length
         __A             =		n_fine_codebooks


class  __lowercase (		lowercase_					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			"bark"
     SCREAMING_SNAKE_CASE				  =			True




     def __init__(			self    :     Dict	,		UpperCamelCase_    :     Dict = None	,		UpperCamelCase_    :     Dict = None	,		UpperCamelCase_    :     Dict = None	,		UpperCamelCase_    :     List[str]=24_000	,		UpperCamelCase_    :     Union[str, Any]=1_024	,		**UpperCamelCase_    :     List[str]	,		):


         """simple docstring"""



         if semantic_config is None:
             __A             =		{}
             logger.info("""semantic_config is None. initializing the semantic model with default values."""			)

         if coarse_acoustics_config is None:
             __A             =		{}
             logger.info("""coarse_acoustics_config is None. initializing the coarse model with default values."""			)

         if fine_acoustics_config is None:
             __A             =		{}
             logger.info("""fine_acoustics_config is None. initializing the fine model with default values."""			)

         __A             =		BarkSemanticGenerationConfig(**UpperCamelCase_			)
         __A             =		BarkCoarseGenerationConfig(**UpperCamelCase_			)
         __A             =		BarkFineGenerationConfig(**UpperCamelCase_			)

         __A             =		sample_rate
         __A             =		codebook_size




     @classmethod
     def 						lowerCAmelCase_       (			cls    :     int	,		UpperCamelCase_    :     BarkSemanticGenerationConfig	,		UpperCamelCase_    :     BarkCoarseGenerationConfig	,		UpperCamelCase_    :     BarkFineGenerationConfig	,		**UpperCamelCase_    :     Tuple	,		):


         """simple docstring"""



         return cls(
             semantic_config=semantic_config.to_dict()	,		coarse_acoustics_config=coarse_acoustics_config.to_dict()	,		fine_acoustics_config=fine_acoustics_config.to_dict()	,		**UpperCamelCase_	,		)





     def 						lowerCAmelCase_       (			self    :     List[str]			):


         """simple docstring"""



         __A             =		copy.deepcopy(self.__dict__			)

         __A             =		self.semantic_config.to_dict()
         __A             =		self.coarse_acoustics_config.to_dict()
         __A             =		self.fine_acoustics_config.to_dict()

         __A             =		self.__class__.model_type
         return output




