






"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE       ,       unittest.TestCase						):





	"""simple docstring"""




	_A       :				Union[str, Any]    =				UnCLIPImageVariationPipeline
	_A       :				Optional[int]    =				IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
	_A       :				List[Any]    =				IMAGE_VARIATION_BATCH_PARAMS

	_A       :				Dict    =				[
	    """generator""",
	    """return_dict""",
	    """decoder_num_inference_steps""",
	    """super_res_num_inference_steps""",
	]
	_A       :				Optional[Any]    =				False





	@property
	def    lowerCamelCase(self							):
		return 32





	@property
	def    lowerCamelCase(self							):
		return 32





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





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





	@property
	def    lowerCamelCase(self							):
		return 100





	@property
	def    lowerCamelCase(self							):
		A_		:						int	    =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							)
		return tokenizer





	@property
	def    lowerCamelCase(self							):
		torch.manual_seed(0							)
		A_		:						Optional[int]	    =   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    lowerCamelCase(self							):
		torch.manual_seed(0							)
		A_		:						Optional[int]	    =   CLIPVisionConfig(
		    hidden_size=self.text_embedder_hidden_size     ,     projection_dim=self.text_embedder_hidden_size     ,     num_hidden_layers=5     ,     num_attention_heads=4     ,     image_size=32     ,     intermediate_size=37     ,     patch_size=1     ,     )
		return CLIPVisionModelWithProjection(lowerCAmelCase_							)





	@property
	def    lowerCamelCase(self							):
		torch.manual_seed(0							)

		A_		:						List[Any]	    =   {
		    """clip_embeddings_dim""": self.text_embedder_hidden_size,
		    """time_embed_dim""": self.time_embed_dim,
		    """cross_attention_dim""": self.cross_attention_dim,
		}

		A_		:						Optional[Any]	    =   UnCLIPTextProjModel(**lowerCAmelCase_							)
		return model





	@property
	def    lowerCamelCase(self							):
		torch.manual_seed(0							)

		A_		:						Optional[int]	    =   {
		    """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_		:						Union[str, Any]	    =   UNetaDConditionModel(**lowerCAmelCase_							)
		return model





	@property
	def    lowerCamelCase(self							):
		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							):
		torch.manual_seed(0							)

		A_		:						int	    =   UNetaDModel(**self.dummy_super_res_kwargs							)
		return model





	@property
	def    lowerCamelCase(self							):
		# seeded differently to get different unet than `self.dummy_super_res_first`
		torch.manual_seed(1							)

		A_		:						List[str]	    =   UNetaDModel(**self.dummy_super_res_kwargs							)
		return model





	def    lowerCamelCase(self							):
		A_		:						List[str]	    =   self.dummy_decoder
		A_		:						Union[str, Any]	    =   self.dummy_text_proj
		A_		:						List[Any]	    =   self.dummy_text_encoder
		A_		:						Dict	    =   self.dummy_tokenizer
		A_		:						Tuple	    =   self.dummy_super_res_first
		A_		:						Optional[int]	    =   self.dummy_super_res_last

		A_		:						Optional[int]	    =   UnCLIPScheduler(
		    variance_type="""learned_range"""     ,     prediction_type="""epsilon"""     ,     num_train_timesteps=1000     ,     )

		A_		:						Tuple	    =   UnCLIPScheduler(
		    variance_type="""fixed_small_log"""     ,     prediction_type="""epsilon"""     ,     num_train_timesteps=1000     ,     )

		A_		:						Union[str, Any]	    =   CLIPImageProcessor(crop_size=32     ,     size=32							)

		A_		:						Union[str, Any]	    =   self.dummy_image_encoder

		return {
		    "decoder": decoder,
		    "text_encoder": text_encoder,
		    "tokenizer": tokenizer,
		    "text_proj": text_proj,
		    "feature_extractor": feature_extractor,
		    "image_encoder": image_encoder,
		    "super_res_first": super_res_first,
		    "super_res_last": super_res_last,
		    "decoder_scheduler": decoder_scheduler,
		    "super_res_scheduler": super_res_scheduler,
		}





	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_=0     ,     lowerCAmelCase_=True							):
		A_		:						str	    =   floats_tensor((1, 3, 32, 32)     ,     rng=random.Random(lowerCAmelCase_							)							).to(lowerCAmelCase_							)
		if str(lowerCAmelCase_							).startswith("""mps"""							):
			A_		:						Dict	    =   torch.manual_seed(lowerCAmelCase_							)
		else:
			A_		:						Union[str, Any]	    =   torch.Generator(device=lowerCAmelCase_							).manual_seed(lowerCAmelCase_							)

		if pil_image:
			A_		:						str	    =   input_image * 0.5 + 0.5
			A_		:						Optional[int]	    =   input_image.clamp(0     ,     1							)
			A_		:						Union[str, Any]	    =   input_image.cpu().permute(0     ,     2     ,     3     ,     1							).float().numpy()
			A_		:						Union[str, Any]	    =   DiffusionPipeline.numpy_to_pil(lowerCAmelCase_							)[0]

		return {
		    "image": input_image,
		    "generator": generator,
		    "decoder_num_inference_steps": 2,
		    "super_res_num_inference_steps": 2,
		    "output_type": "np",
		}





	def    lowerCamelCase(self							):
		A_		:						Tuple	    =   """cpu"""

		A_		:						Union[str, Any]	    =   self.get_dummy_components()

		A_		:						int	    =   self.pipeline_class(**lowerCAmelCase_							)
		A_		:						Dict	    =   pipe.to(lowerCAmelCase_							)

		pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

		A_		:						List[str]	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)

		A_		:						int	    =   pipe(**lowerCAmelCase_							)
		A_		:						Union[str, Any]	    =   output.images

		A_		:						int	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)

		A_		:						List[str]	    =   pipe(
		    **lowerCAmelCase_     ,     return_dict=lowerCAmelCase_     ,     )[0]

		A_		:						Tuple	    =   image[0, -3:, -3:, -1]
		A_		:						int	    =   image_from_tuple[0, -3:, -3:, -1]

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

		A_		:						List[str]	    =   np.array(
		    [
		        0.9997,
		        0.0002,
		        0.9997,
		        0.9997,
		        0.9969,
		        0.0023,
		        0.9997,
		        0.9969,
		        0.9970,
		    ]							)

		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							):
		A_		:						Optional[Any]	    =   """cpu"""

		A_		:						Union[str, Any]	    =   self.get_dummy_components()

		A_		:						Optional[int]	    =   self.pipeline_class(**lowerCAmelCase_							)
		A_		:						str	    =   pipe.to(lowerCAmelCase_							)

		pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

		A_		:						Dict	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)

		A_		:						Tuple	    =   pipe(**lowerCAmelCase_							)
		A_		:						List[Any]	    =   output.images

		A_		:						Tuple	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)

		A_		:						Optional[Any]	    =   pipe(
		    **lowerCAmelCase_     ,     return_dict=lowerCAmelCase_     ,     )[0]

		A_		:						str	    =   image[0, -3:, -3:, -1]
		A_		:						Optional[Any]	    =   image_from_tuple[0, -3:, -3:, -1]

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

		A_		:						Optional[int]	    =   np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]							)

		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							):
		A_		:						Tuple	    =   """cpu"""

		A_		:						Union[str, Any]	    =   self.get_dummy_components()

		A_		:						Union[str, Any]	    =   self.pipeline_class(**lowerCAmelCase_							)
		A_		:						int	    =   pipe.to(lowerCAmelCase_							)

		pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

		A_		:						List[str]	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)
		A_		:						Optional[int]	    =   [
		    pipeline_inputs["""image"""],
		    pipeline_inputs["""image"""],
		]

		A_		:						str	    =   pipe(**lowerCAmelCase_							)
		A_		:						Optional[int]	    =   output.images

		A_		:						List[str]	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)
		A_		:						str	    =   [
		    tuple_pipeline_inputs["""image"""],
		    tuple_pipeline_inputs["""image"""],
		]

		A_		:						Optional[Any]	    =   pipe(
		    **lowerCAmelCase_     ,     return_dict=lowerCAmelCase_     ,     )[0]

		A_		:						str	    =   image[0, -3:, -3:, -1]
		A_		:						List[str]	    =   image_from_tuple[0, -3:, -3:, -1]

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

		A_		:						List[Any]	    =   np.array(
		    [
		        0.9997,
		        0.9989,
		        0.0008,
		        0.0021,
		        0.9960,
		        0.0018,
		        0.0014,
		        0.0002,
		        0.9933,
		    ]							)

		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							):
		A_		:						List[str]	    =   torch.device("""cpu"""							)



		class      SCREAMING_SNAKE_CASE  :





			"""simple docstring"""




			_A       :				Optional[int]    =				1






		A_		:						Any	    =   self.get_dummy_components()

		A_		:						int	    =   self.pipeline_class(**lowerCAmelCase_							)
		A_		:						int	    =   pipe.to(lowerCAmelCase_							)

		pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

		A_		:						Dict	    =   torch.Generator(device=lowerCAmelCase_							).manual_seed(0							)
		A_		:						str	    =   pipe.decoder.dtype
		A_		:						List[str]	    =   1

		A_		:						Tuple	    =   (
		    batch_size,
		    pipe.decoder.config.in_channels,
		    pipe.decoder.config.sample_size,
		    pipe.decoder.config.sample_size,
		)
		A_		:						Union[str, Any]	    =   pipe.prepare_latents(
		    lowerCAmelCase_     ,     dtype=lowerCAmelCase_     ,     device=lowerCAmelCase_     ,     generator=lowerCAmelCase_     ,     latents=lowerCAmelCase_     ,     scheduler=DummyScheduler()							)

		A_		:						List[str]	    =   (
		    batch_size,
		    pipe.super_res_first.config.in_channels // 2,
		    pipe.super_res_first.config.sample_size,
		    pipe.super_res_first.config.sample_size,
		)
		A_		:						Optional[Any]	    =   pipe.prepare_latents(
		    lowerCAmelCase_     ,     dtype=lowerCAmelCase_     ,     device=lowerCAmelCase_     ,     generator=lowerCAmelCase_     ,     latents=lowerCAmelCase_     ,     scheduler=DummyScheduler()							)

		A_		:						Union[str, Any]	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)

		A_		:						Any	    =   pipe(
		    **lowerCAmelCase_     ,     decoder_latents=lowerCAmelCase_     ,     super_res_latents=lowerCAmelCase_							).images

		A_		:						int	    =   self.get_dummy_inputs(lowerCAmelCase_     ,     pil_image=lowerCAmelCase_							)
		# Don't pass image, instead pass embedding
		A_		:						str	    =   pipeline_inputs.pop("""image"""							)
		A_		:						List[str]	    =   pipe.image_encoder(lowerCAmelCase_							).image_embeds

		A_		:						Optional[int]	    =   pipe(
		    **lowerCAmelCase_     ,     decoder_latents=lowerCAmelCase_     ,     super_res_latents=lowerCAmelCase_     ,     image_embeddings=lowerCAmelCase_     ,     ).images

		# make sure passing text embeddings manually is identical
		assert np.abs(img_out_a - img_out_a							).max() < 1e-4





	@skip_mps
	def    lowerCamelCase(self							):
		A_		:						int	    =   torch_device == """cpu"""

		# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
		A_		:						Dict	    =   1e-2

		self._test_attention_slicing_forward_pass(
		    test_max_difference=lowerCAmelCase_     ,     expected_max_diff=lowerCAmelCase_							)





	@skip_mps
	def    lowerCamelCase(self							):
		A_		:						str	    =   torch_device == """cpu"""
		A_		:						Any	    =   True
		A_		:						List[Any]	    =   [
		    """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    lowerCamelCase(self							):
		A_		:						Dict	    =   [
		    """decoder_num_inference_steps""",
		    """super_res_num_inference_steps""",
		]

		if torch_device == "mps":
			# TODO: MPS errors with larger batch sizes
			A_		:						Dict	    =   [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    lowerCamelCase(self							):
		return super().test_dict_tuple_outputs_equivalent()





	@skip_mps
	def    lowerCamelCase(self							):
		return super().test_save_load_local()





	@skip_mps
	def    lowerCamelCase(self							):
		return super().test_save_load_optional_components()






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





	"""simple docstring"""




	def    lowerCamelCase(self							):
		# clean up the VRAM after each test
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def    lowerCamelCase(self							):
		A_		:						List[Any]	    =   load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"""							)
		A_		:						str	    =   load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/unclip/karlo_v1_alpha_cat_variation_fp16.npy"""							)

		A_		:						Any	    =   UnCLIPImageVariationPipeline.from_pretrained(
		    """kakaobrain/karlo-v1-alpha-image-variations"""     ,     torch_dtype=torch.floataa							)
		A_		:						List[Any]	    =   pipeline.to(lowerCAmelCase_							)
		pipeline.set_progress_bar_config(disable=lowerCAmelCase_							)

		A_		:						Optional[int]	    =   torch.Generator(device="""cpu"""							).manual_seed(0							)
		A_		:						str	    =   pipeline(
		    lowerCAmelCase_     ,     generator=lowerCAmelCase_     ,     output_type="""np"""     ,     )

		A_		:						Dict	    =   output.images[0]

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

		assert_mean_pixel_difference(lowerCAmelCase_     ,     lowerCAmelCase_     ,     15							)









"""simple docstring"""
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.auto import AutoModelForSequenceClassification, AutoTokenizer
from .base import PipelineTool



class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE						):





	"""simple docstring"""




	_A       :				Optional[int]    =				"""facebook/bart-large-mnli"""
	_A       :				str    =				(
	    """This is a tool that classifies an English text using provided labels. It takes two inputs: `text`, which """
	    """should be the text to classify, and `labels`, which should be the list of labels to use for classification. """
	    """It returns the most likely label in the list of provided `labels` for the input text."""
	)
	_A       :				List[str]    =				"""text_classifier"""
	_A       :				Optional[int]    =				AutoTokenizer
	_A       :				Optional[Any]    =				AutoModelForSequenceClassification

	_A       :				List[str]    =				["""text""", ["""text"""]]
	_A       :				Dict    =				["""text"""]





	def    lowerCamelCase(self							):
		super().setup()
		A_		:						int	    =   self.model.config
		A_		:						List[str]	    =   -1
		for idx, label in config.idalabel.items():
			if label.lower().startswith("""entail"""							):
				A_		:						List[Any]	    =   int(lowerCAmelCase_							)
		if self.entailment_id == -1:
			raise ValueError("""Could not determine the entailment ID from the model config, please pass it at init."""							)





	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						List[Any]	    =   labels
		return self.pre_processor(
		    [text] * len(lowerCAmelCase_							)     ,     [f"""This example is {label}""" for label in labels]     ,     return_tensors="""pt"""     ,     padding="""max_length"""     ,     )





	def    lowerCamelCase(self     ,     lowerCAmelCase_							):
		A_		:						str	    =   outputs.logits
		A_		:						Optional[int]	    =   torch.argmax(logits[:, 2]							).item()
		return self._labels[label_id]

