


'''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  UpperCAmelCase (       A__		,			unittest.TestCase):
	_lowerCamelCase :					Optional[Any]    					=					UnCLIPPipeline
	_lowerCamelCase :					List[str]    					=					TEXT_TO_IMAGE_PARAMS - {
	    'negative_prompt',
	    'height',
	    'width',
	    'negative_prompt_embeds',
	    'guidance_scale',
	    'prompt_embeds',
	    'cross_attention_kwargs',
	}
	_lowerCamelCase :					Dict    					=					TEXT_TO_IMAGE_BATCH_PARAMS
	_lowerCamelCase :					Dict    					=					[
	    'generator',
	    'return_dict',
	    'prior_num_inference_steps',
	    'decoder_num_inference_steps',
	    'super_res_num_inference_steps',
	]
	_lowerCamelCase :					str    					=					False

	@property
	def 		lowercase_						( self :	str					):





				"""simple docstring"""



				return 32

	@property
	def 		lowercase_						( self :	List[Any]					):





				"""simple docstring"""



				return 32

	@property
	def 		lowercase_						( self :	Any					):





				"""simple docstring"""



				return self.time_input_dim

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





				"""simple docstring"""



				return self.time_input_dim * 4

	@property
	def 		lowercase_						( self :	Any					):





				"""simple docstring"""



				return 100

	@property
	def 		lowercase_						( self :	int					):





				"""simple docstring"""



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

	@property
	def 		lowercase_						( self :	int					):





				"""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=37,					layer_norm_eps=1e-0_5,					num_attention_heads=4,					num_hidden_layers=5,					pad_token_id=1,					vocab_size=1000,					)
				return CLIPTextModelWithProjection(_lowerCAmelCase					)

	@property
	def 		lowercase_						( self :	List[str]					):





				"""simple docstring"""



				torch.manual_seed(0					)

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

				UpperCamelCase__ 						=						PriorTransformer(**_lowerCAmelCase					)
				return model

	@property
	def 		lowercase_						( self :	List[Any]					):





				"""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(**_lowerCAmelCase					)
				return model

	@property
	def 		lowercase_						( self :	List[str]					):





				"""simple docstring"""



				torch.manual_seed(0					)

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

				UpperCamelCase__ 						=						UNetaDConditionModel(**_lowerCAmelCase					)
				return model

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





				"""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 		lowercase_						( self :	Optional[int]					):





				"""simple docstring"""



				torch.manual_seed(0					)

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

	@property
	def 		lowercase_						( self :	str					):





				"""simple docstring"""



				torch.manual_seed(1					)

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

	def 		lowercase_						( self :	Optional[int]					):





				"""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=1000,					clip_sample_range=5.0,					)

				UpperCamelCase__ 						=						UnCLIPScheduler(
				    variance_type="learned_range",					prediction_type="epsilon",					num_train_timesteps=1000,					)

				UpperCamelCase__ 						=						UnCLIPScheduler(
				    variance_type="fixed_small_log",					prediction_type="epsilon",					num_train_timesteps=1000,					)

				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 :	Tuple,					a_ :	str,					a_ :	List[str]=0					):





				"""simple docstring"""



				if str(_lowerCAmelCase					).startswith("mps"					):
							UpperCamelCase__ 						=						torch.manual_seed(_lowerCAmelCase					)
				else:
							UpperCamelCase__ 						=						torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)
				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 :	Any					):





				"""simple docstring"""



				UpperCamelCase__ 						=						"cpu"

				UpperCamelCase__ 						=						self.get_dummy_components()

				UpperCamelCase__ 						=						self.pipeline_class(**_lowerCAmelCase					)
				UpperCamelCase__ 						=						pipe.to(_lowerCAmelCase					)

				pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

				UpperCamelCase__ 						=						pipe(**self.get_dummy_inputs(_lowerCAmelCase					)					)
				UpperCamelCase__ 						=						output.images

				UpperCamelCase__ 						=						pipe(
				    **self.get_dummy_inputs(_lowerCAmelCase					),					return_dict=_lowerCAmelCase,					)[0]

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

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

				UpperCamelCase__ 						=						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 		lowercase_						( self :	str					):





				"""simple docstring"""



				UpperCamelCase__ 						=						torch.device("cpu"					)







				class  UpperCAmelCase :
					_lowerCamelCase :					Optional[int]    					=					1







				UpperCamelCase__ 						=						self.get_dummy_components()

				UpperCamelCase__ 						=						self.pipeline_class(**_lowerCAmelCase					)
				UpperCamelCase__ 						=						pipe.to(_lowerCAmelCase					)

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

				UpperCamelCase__ 						=						torch.Generator(device=_lowerCAmelCase					).manual_seed(0					)
				UpperCamelCase__ 						=						prior.dtype
				UpperCamelCase__ 						=						1

				UpperCamelCase__ 						=						(batch_size, prior.config.embedding_dim)
				UpperCamelCase__ 						=						pipe.prepare_latents(
				    _lowerCAmelCase,					dtype=_lowerCAmelCase,					device=_lowerCAmelCase,					generator=_lowerCAmelCase,					latents=_lowerCAmelCase,					scheduler=DummyScheduler()					)
				UpperCamelCase__ 						=						(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
				UpperCamelCase__ 						=						pipe.prepare_latents(
				    _lowerCAmelCase,					dtype=_lowerCAmelCase,					device=_lowerCAmelCase,					generator=_lowerCAmelCase,					latents=_lowerCAmelCase,					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(
				    _lowerCAmelCase,					dtype=_lowerCAmelCase,					device=_lowerCAmelCase,					generator=_lowerCAmelCase,					latents=_lowerCAmelCase,					scheduler=DummyScheduler()					)

				pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

				UpperCamelCase__ 						=						"this is a prompt example"

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

				UpperCamelCase__ 						=						tokenizer(
				    _lowerCAmelCase,					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=_lowerCAmelCase					).manual_seed(0					)
				UpperCamelCase__ 						=						pipe(
				    generator=_lowerCAmelCase,					prior_num_inference_steps=2,					decoder_num_inference_steps=2,					super_res_num_inference_steps=2,					prior_latents=_lowerCAmelCase,					decoder_latents=_lowerCAmelCase,					super_res_latents=_lowerCAmelCase,					text_model_output=_lowerCAmelCase,					text_attention_mask=_lowerCAmelCase,					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 :	Any					):





				"""simple docstring"""



				UpperCamelCase__ 						=						torch_device == "cpu"

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

	@skip_mps
	def 		lowercase_						( self :	Tuple					):





				"""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=_lowerCAmelCase,					relax_max_difference=_lowerCAmelCase,					additional_params_copy_to_batched_inputs=_lowerCAmelCase,					)

	def 		lowercase_						( self :	Optional[int]					):





				"""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=_lowerCAmelCase,					additional_params_copy_to_batched_inputs=_lowerCAmelCase,					)
				else:
							self._test_inference_batch_consistent(
							    additional_params_copy_to_batched_inputs=_lowerCAmelCase					)

	@skip_mps
	def 		lowercase_						( self :	List[str]					):





				"""simple docstring"""



				return super().test_dict_tuple_outputs_equivalent()

	@skip_mps
	def 		lowercase_						( self :	List[Any]					):





				"""simple docstring"""



				return super().test_save_load_local()

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





				"""simple docstring"""



				return super().test_save_load_optional_components()







@nightly
class  UpperCAmelCase (       unittest.TestCase):

	def 		lowercase_						( self :	int					):





				"""simple docstring"""



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

	def 		lowercase_						( self :	Dict					):





				"""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=_lowerCAmelCase					)

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

				UpperCamelCase__ 						=						output.images[0]

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







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

	def 		lowercase_						( self :	Tuple					):





				"""simple docstring"""



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

	def 		lowercase_						( self :	int					):





				"""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(_lowerCAmelCase					)
				pipeline.set_progress_bar_config(disable=_lowerCAmelCase					)

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

				UpperCamelCase__ 						=						output.images[0]

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

				assert_mean_pixel_difference(_lowerCAmelCase,					_lowerCAmelCase					)

	def 		lowercase_						( self :	List[Any]					):





				"""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(_lowerCAmelCase					)
				pipe.set_progress_bar_config(disable=_lowerCAmelCase					)
				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 * 10**9



'''simple docstring'''

from typing import List

from ...processing_utils import ProcessorMixin
from ...utils import is_torch_available


if is_torch_available():
       import torch







class  UpperCAmelCase (       SCREAMING_SNAKE_CASE__):
 _lowerCamelCase :					Tuple    					=					['image_processor', 'tokenizer']
 _lowerCamelCase :					Tuple    					=					'OneFormerImageProcessor'
 _lowerCamelCase :					List[Any]    					=					('CLIPTokenizer', 'CLIPTokenizerFast')

 def __init__( self :	Any,					a_ :	Dict=None,					a_ :	int=None,					a_ :	int = 77,					a_ :	int = 77,					**a_ :	int					):





    """simple docstring"""



    if image_processor is None:
       raise ValueError("You need to specify an `image_processor`."					)
    if tokenizer is None:
       raise ValueError("You need to specify a `tokenizer`."					)

    UpperCamelCase__ 						=						max_seq_length
    UpperCamelCase__ 						=						task_seq_length

    super().__init__(a_,					a_					)

 def 		lowercase_						( self :	Any,					a_ :	Optional[Any]=None,					a_ :	Optional[int]=77					):





    """simple docstring"""



    if text_list is None:
       raise ValueError("tokens cannot be None."					)

    UpperCamelCase__ 						=						self.tokenizer(a_,					padding="max_length",					max_length=a_,					truncation=a_					)

    UpperCamelCase__						,		UpperCamelCase__ 						=						tokens["attention_mask"], tokens["input_ids"]

    UpperCamelCase__ 						=						[]
    for attn_mask, input_id in zip(a_,					a_					):
       UpperCamelCase__ 						=						torch.tensor(a_					) * torch.tensor(a_					)
       token_inputs.append(token.unsqueeze(0					)					)

    UpperCamelCase__ 						=						torch.cat(a_,					dim=0					)
    return token_inputs

 def __call__( self :	str,					a_ :	Dict=None,					a_ :	str=None,					a_ :	Tuple=None,					**a_ :	Any					):





    """simple docstring"""



    if task_inputs is None:
       raise ValueError("You have to specify the task_input. Found None."					)
    elif images is None:
       raise ValueError("You have to specify the image. Found None."					)

    if not all(task in ["semantic", "instance", "panoptic"] for task in task_inputs					):
       raise ValueError("task_inputs must be semantic, instance, or panoptic."					)

    UpperCamelCase__ 						=						self.image_processor(a_,					a_,					a_,					**a_					)

    if isinstance(a_,					a_					):
       UpperCamelCase__ 						=						[task_inputs]

    if isinstance(a_,					a_					) and all(isinstance(a_,					a_					) for task_input in task_inputs					):
       UpperCamelCase__ 						=						[]
       for task in task_inputs:
          UpperCamelCase__ 						=						f'the task is {task}'
          task_token_inputs.append(a_					)
       UpperCamelCase__ 						=						self._preprocess_text(a_,					max_length=self.task_seq_length					)
    else:
       raise TypeError("Task Inputs should be a string or a list of strings."					)

    if hasattr(a_,					"text_inputs"					):
       UpperCamelCase__ 						=						encoded_inputs.text_inputs

       UpperCamelCase__ 						=						[]
       for texts in texts_list:
          UpperCamelCase__ 						=						self._preprocess_text(a_,					max_length=self.max_seq_length					)
          text_inputs.append(text_input_list.unsqueeze(0					)					)

       UpperCamelCase__ 						=						torch.cat(a_,					dim=0					)

    return encoded_inputs

 def 		lowercase_						( self :	Union[str, Any],					a_ :	Tuple=None,					a_ :	List[str]=None,					a_ :	Optional[Any]=None,					**a_ :	Tuple					):





    """simple docstring"""



    if task_inputs is None:
       raise ValueError("You have to specify the task_input. Found None."					)
    elif images is None:
       raise ValueError("You have to specify the image. Found None."					)

    if not all(task in ["semantic", "instance", "panoptic"] for task in task_inputs					):
       raise ValueError("task_inputs must be semantic, instance, or panoptic."					)

    UpperCamelCase__ 						=						self.image_processor.encode_inputs(a_,					a_,					a_,					**a_					)

    if isinstance(a_,					a_					):
       UpperCamelCase__ 						=						[task_inputs]

    if isinstance(a_,					a_					) and all(isinstance(a_,					a_					) for task_input in task_inputs					):
       UpperCamelCase__ 						=						[]
       for task in task_inputs:
          UpperCamelCase__ 						=						f'the task is {task}'
          task_token_inputs.append(a_					)
       UpperCamelCase__ 						=						self._preprocess_text(a_,					max_length=self.task_seq_length					)
    else:
       raise TypeError("Task Inputs should be a string or a list of strings."					)

    if hasattr(a_,					"text_inputs"					):
       UpperCamelCase__ 						=						encoded_inputs.text_inputs

       UpperCamelCase__ 						=						[]
       for texts in texts_list:
          UpperCamelCase__ 						=						self._preprocess_text(a_,					max_length=self.max_seq_length					)
          text_inputs.append(text_input_list.unsqueeze(0					)					)

       UpperCamelCase__ 						=						torch.cat(a_,					dim=0					)

    return encoded_inputs

 def 		lowercase_						( self :	Any,					*a_ :	int,					**a_ :	Optional[int]					):





    """simple docstring"""



    return self.image_processor.post_process_semantic_segmentation(*a_,					**a_					)

 def 		lowercase_						( self :	str,					*a_ :	Tuple,					**a_ :	Optional[int]					):





    """simple docstring"""



    return self.image_processor.post_process_instance_segmentation(*a_,					**a_					)

 def 		lowercase_						( self :	Dict,					*a_ :	str,					**a_ :	int					):





    """simple docstring"""



    return self.image_processor.post_process_panoptic_segmentation(*a_,					**a_					)