






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



				"""simple docstring"""



				_lowercase				:   Union[str, Any]      =  UnCLIPImageVariationPipeline
				_lowercase				:   Tuple      =  IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''guidance_scale'''}
				_lowercase				:   Dict      =  IMAGE_VARIATION_BATCH_PARAMS

				_lowercase				:   Dict      =  [
				    '''generator''',
				    '''return_dict''',
				    '''decoder_num_inference_steps''',
				    '''super_res_num_inference_steps''',
				]
				_lowercase				:   Optional[int]      =  False



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






											'''simple docstring'''







											return 3_2



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






											'''simple docstring'''







											return 3_2



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






											'''simple docstring'''







											return self.time_input_dim



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






											'''simple docstring'''







											return self.time_input_dim * 4



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






											'''simple docstring'''







											return 1_0_0



				@property
				def 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







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



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






											'''simple docstring'''







											torch.manual_seed(0)
											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=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(UpperCamelCase__)



				@property
				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											torch.manual_seed(0)
											snake_case__	  =       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=3_2			,		intermediate_size=3_7			,		patch_size=1			,		)
											return CLIPVisionModelWithProjection(UpperCamelCase__)



				@property
				def 					__magic_name__							(      self    :    Any):






											'''simple docstring'''







											torch.manual_seed(0)

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

											snake_case__	  =       UnCLIPTextProjModel(**UpperCamelCase__)
											return model



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






											'''simple docstring'''







											torch.manual_seed(0)

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

											snake_case__	  =       UNetaDConditionModel(**UpperCamelCase__)
											return model



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






											'''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 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







											torch.manual_seed(0)

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



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






											'''simple docstring'''







											torch.manual_seed(1)

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



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											snake_case__	  =       self.dummy_decoder
											snake_case__	  =       self.dummy_text_proj
											snake_case__	  =       self.dummy_text_encoder
											snake_case__	  =       self.dummy_tokenizer
											snake_case__	  =       self.dummy_super_res_first
											snake_case__	  =       self.dummy_super_res_last

											snake_case__	  =       UnCLIPScheduler(
											    variance_type="""learned_range"""			,		prediction_type="""epsilon"""			,		num_train_timesteps=1_0_0_0			,		)

											snake_case__	  =       UnCLIPScheduler(
											    variance_type="""fixed_small_log"""			,		prediction_type="""epsilon"""			,		num_train_timesteps=1_0_0_0			,		)

											snake_case__	  =       CLIPImageProcessor(crop_size=3_2			,		size=3_2)

											snake_case__	  =       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 					__magic_name__							(      self    :    Any			,		UpperCamelCase__    :    Dict			,		UpperCamelCase__    :    str=0			,		UpperCamelCase__    :    Dict=True):






											'''simple docstring'''







											snake_case__	  =       floats_tensor((1, 3, 3_2, 3_2)			,		rng=random.Random(UpperCamelCase__)).to(UpperCamelCase__)
											if str(UpperCamelCase__).startswith("""mps"""):
																		snake_case__	  =       torch.manual_seed(UpperCamelCase__)
											else:
																		snake_case__	  =       torch.Generator(device=UpperCamelCase__).manual_seed(UpperCamelCase__)

											if pil_image:
																		snake_case__	  =       input_image * 0.5 + 0.5
																		snake_case__	  =       input_image.clamp(0			,		1)
																		snake_case__	  =       input_image.cpu().permute(0			,		2			,		3			,		1).float().numpy()
																		snake_case__	  =       DiffusionPipeline.numpy_to_pil(UpperCamelCase__)[0]

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



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											snake_case__	  =       """cpu"""

											snake_case__	  =       self.get_dummy_components()

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

											pipe.set_progress_bar_config(disable=UpperCamelCase__)

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)

											snake_case__	  =       pipe(**UpperCamelCase__)
											snake_case__	  =       output.images

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)

											snake_case__	  =       pipe(
											    **UpperCamelCase__			,		return_dict=UpperCamelCase__			,		)[0]

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

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

											snake_case__	  =       np.array(
											    [
											        0.99_97,
											        0.00_02,
											        0.99_97,
											        0.99_97,
											        0.99_69,
											        0.00_23,
											        0.99_97,
											        0.99_69,
											        0.99_70,
											    ])

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






											'''simple docstring'''







											snake_case__	  =       """cpu"""

											snake_case__	  =       self.get_dummy_components()

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

											pipe.set_progress_bar_config(disable=UpperCamelCase__)

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)

											snake_case__	  =       pipe(**UpperCamelCase__)
											snake_case__	  =       output.images

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)

											snake_case__	  =       pipe(
											    **UpperCamelCase__			,		return_dict=UpperCamelCase__			,		)[0]

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

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

											snake_case__	  =       np.array([0.99_97, 0.00_03, 0.99_97, 0.99_97, 0.99_70, 0.00_24, 0.99_97, 0.99_71, 0.99_71])

											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 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											snake_case__	  =       """cpu"""

											snake_case__	  =       self.get_dummy_components()

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

											pipe.set_progress_bar_config(disable=UpperCamelCase__)

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)
											snake_case__	  =       [
											    pipeline_inputs["""image"""],
											    pipeline_inputs["""image"""],
											]

											snake_case__	  =       pipe(**UpperCamelCase__)
											snake_case__	  =       output.images

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)
											snake_case__	  =       [
											    tuple_pipeline_inputs["""image"""],
											    tuple_pipeline_inputs["""image"""],
											]

											snake_case__	  =       pipe(
											    **UpperCamelCase__			,		return_dict=UpperCamelCase__			,		)[0]

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

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

											snake_case__	  =       np.array(
											    [
											        0.99_97,
											        0.99_89,
											        0.00_08,
											        0.00_21,
											        0.99_60,
											        0.00_18,
											        0.00_14,
											        0.00_02,
											        0.99_33,
											    ])

											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 					__magic_name__							(      self    :    int):






											'''simple docstring'''







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


											class    _lowerCAmelCase			:



															"""simple docstring"""



															_lowercase				:   str      =  1


											snake_case__	  =       self.get_dummy_components()

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

											pipe.set_progress_bar_config(disable=UpperCamelCase__)

											snake_case__	  =       torch.Generator(device=UpperCamelCase__).manual_seed(0)
											snake_case__	  =       pipe.decoder.dtype
											snake_case__	  =       1

											snake_case__	  =       (
											    batch_size,
											    pipe.decoder.config.in_channels,
											    pipe.decoder.config.sample_size,
											    pipe.decoder.config.sample_size,
											)
											snake_case__	  =       pipe.prepare_latents(
											    UpperCamelCase__			,		dtype=UpperCamelCase__			,		device=UpperCamelCase__			,		generator=UpperCamelCase__			,		latents=UpperCamelCase__			,		scheduler=DummyScheduler())

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

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)

											snake_case__	  =       pipe(
											    **UpperCamelCase__			,		decoder_latents=UpperCamelCase__			,		super_res_latents=UpperCamelCase__).images

											snake_case__	  =       self.get_dummy_inputs(UpperCamelCase__			,		pil_image=UpperCamelCase__)
											# Don't pass image, instead pass embedding
											snake_case__	  =       pipeline_inputs.pop("""image""")
											snake_case__	  =       pipe.image_encoder(UpperCamelCase__).image_embeds

											snake_case__	  =       pipe(
											    **UpperCamelCase__			,		decoder_latents=UpperCamelCase__			,		super_res_latents=UpperCamelCase__			,		image_embeddings=UpperCamelCase__			,		).images

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



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






											'''simple docstring'''







											snake_case__	  =       torch_device == """cpu"""

											# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
											snake_case__	  =       1E-2

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



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






											'''simple docstring'''







											snake_case__	  =       torch_device == """cpu"""
											snake_case__	  =       True
											snake_case__	  =       [
											    """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 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											snake_case__	  =       [
											    """decoder_num_inference_steps""",
											    """super_res_num_inference_steps""",
											]

											if torch_device == "mps":
																		# TODO: MPS errors with larger batch sizes
																		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 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											return super().test_dict_tuple_outputs_equivalent()



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






											'''simple docstring'''







											return super().test_save_load_local()



				@skip_mps
				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											return super().test_save_load_optional_components()


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



				"""simple docstring"""



				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







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



				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											snake_case__	  =       load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png""")
											snake_case__	  =       load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
											    """/unclip/karlo_v1_alpha_cat_variation_fp16.npy""")

											snake_case__	  =       UnCLIPImageVariationPipeline.from_pretrained(
											    """kakaobrain/karlo-v1-alpha-image-variations"""			,		torch_dtype=torch.floataa)
											snake_case__	  =       pipeline.to(UpperCamelCase__)
											pipeline.set_progress_bar_config(disable=UpperCamelCase__)

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

											snake_case__	  =       output.images[0]

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

											assert_mean_pixel_difference(UpperCamelCase__			,		UpperCamelCase__			,		1_5)








import os
import time

import pytest

from datasets.utils.filelock import FileLock, Timeout

def _UpperCAmelCase   (	a						:	List[Any] ):
							snake_case__	  =       FileLock(str(tmpdir / """foo.lock""" ) )
							snake_case__	  =       FileLock(str(tmpdir / """foo.lock""" ) )
							snake_case__	  =       0.01
							with locka.acquire():
														with pytest.raises(a ):
																					snake_case__	  =       time.time()
																					locka.acquire(a )
																					assert time.time() - _start > timeout



def _UpperCAmelCase   (	a						:	str ):
							snake_case__	  =       """a""" * 1000 + """.lock"""
							snake_case__	  =       FileLock(str(tmpdir / filename ) )
							assert locka._lock_file.endswith(""".lock""" )
							assert not locka._lock_file.endswith(a )
							assert len(os.path.basename(locka._lock_file ) ) <= 255
							snake_case__	  =       FileLock(tmpdir / filename )
							with locka.acquire():
														with pytest.raises(a ):
																					locka.acquire(0 )
