




import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase	,  _UpperCamelCase	,  _UpperCamelCase	,  unittest.TestCase   ):



	"""simple docstring"""


	__magic_name__   :       str      	=		StableUnCLIPPipeline
	__magic_name__   :       str      	=		TEXT_TO_IMAGE_PARAMS
	__magic_name__   :       str      	=		TEXT_TO_IMAGE_BATCH_PARAMS
	__magic_name__   :       Union[str, Any]      	=		TEXT_TO_IMAGE_IMAGE_PARAMS
	__magic_name__   :       Any      	=		TEXT_TO_IMAGE_IMAGE_PARAMS

	# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
	__magic_name__   :       Any      	=		False






	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       int:







						"""simple docstring"""







						__UpperCamelCase					:	Tuple             =      32
						__UpperCamelCase					:	Dict             =      embedder_hidden_size

						# prior components

						torch.manual_seed(0	)
						__UpperCamelCase					:	Tuple             =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

						torch.manual_seed(0	)
						__UpperCamelCase					:	str             =      CLIPTextModelWithProjection(
						    CLIPTextConfig(
						        bos_token_id=0					,     eos_token_id=2					,     hidden_size=lowerCAmelCase					,     projection_dim=lowerCAmelCase					,     intermediate_size=37					,     layer_norm_eps=1E-05					,     num_attention_heads=4					,     num_hidden_layers=5					,     pad_token_id=1					,     vocab_size=1000					,     )	)

						torch.manual_seed(0	)
						__UpperCamelCase					:	List[str]             =      PriorTransformer(
						    num_attention_heads=2					,     attention_head_dim=12					,     embedding_dim=lowerCAmelCase					,     num_layers=1					,     )

						torch.manual_seed(0	)
						__UpperCamelCase					:	Dict             =      DDPMScheduler(
						    variance_type="""fixed_small_log"""					,     prediction_type="""sample"""					,     num_train_timesteps=1000					,     clip_sample=lowerCAmelCase					,     clip_sample_range=5.0					,     beta_schedule="""squaredcos_cap_v2"""					,     )

						# regular denoising components

						torch.manual_seed(0	)
						__UpperCamelCase					:	Union[str, Any]             =      StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase	)
						__UpperCamelCase					:	List[Any]             =      DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""	)

						torch.manual_seed(0	)
						__UpperCamelCase					:	List[str]             =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

						torch.manual_seed(0	)
						__UpperCamelCase					:	Any             =      CLIPTextModel(
						    CLIPTextConfig(
						        bos_token_id=0					,     eos_token_id=2					,     hidden_size=lowerCAmelCase					,     projection_dim=32					,     intermediate_size=37					,     layer_norm_eps=1E-05					,     num_attention_heads=4					,     num_hidden_layers=5					,     pad_token_id=1					,     vocab_size=1000					,     )	)

						torch.manual_seed(0	)
						__UpperCamelCase					:	Optional[Any]             =      UNetaDConditionModel(
						    sample_size=32					,     in_channels=4					,     out_channels=4					,     down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")					,     up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")					,     block_out_channels=(32, 64)					,     attention_head_dim=(2, 4)					,     class_embed_type="""projection"""					,     projection_class_embeddings_input_dim=embedder_projection_dim * 2					,     cross_attention_dim=lowerCAmelCase					,     layers_per_block=1					,     upcast_attention=lowerCAmelCase					,     use_linear_projection=lowerCAmelCase					,     )

						torch.manual_seed(0	)
						__UpperCamelCase					:	Optional[int]             =      DDIMScheduler(
						    beta_schedule="""scaled_linear"""					,     beta_start=0.0_00_85					,     beta_end=0.0_12					,     prediction_type="""v_prediction"""					,     set_alpha_to_one=lowerCAmelCase					,     steps_offset=1					,     )

						torch.manual_seed(0	)
						__UpperCamelCase					:	Optional[int]             =      AutoencoderKL()

						__UpperCamelCase					:	List[str]             =      {
						    # prior components
						    """prior_tokenizer""": prior_tokenizer,
						    """prior_text_encoder""": prior_text_encoder,
						    """prior""": prior,
						    """prior_scheduler""": prior_scheduler,
						    # image noising components
						    """image_normalizer""": image_normalizer,
						    """image_noising_scheduler""": image_noising_scheduler,
						    # regular denoising components
						    """tokenizer""": tokenizer,
						    """text_encoder""": text_encoder,
						    """unet""": unet,
						    """scheduler""": scheduler,
						    """vae""": vae,
						}

						return components






	def 					lowerCamelCase__     (					self		:						str					,     lowerCAmelCase		:						List[str]					,     lowerCAmelCase		:						Dict=0	)	->       Union[str, Any]:







						"""simple docstring"""







						if str(lowerCAmelCase	).startswith("""mps"""	):
											__UpperCamelCase					:	Tuple             =      torch.manual_seed(lowerCAmelCase	)
						else:
											__UpperCamelCase					:	List[str]             =      torch.Generator(device=lowerCAmelCase	).manual_seed(lowerCAmelCase	)
						__UpperCamelCase					:	Dict             =      {
						    """prompt""": """A painting of a squirrel eating a burger""",
						    """generator""": generator,
						    """num_inference_steps""": 2,
						    """prior_num_inference_steps""": 2,
						    """output_type""": """numpy""",
						}
						return inputs






	def 					lowerCamelCase__     (					self		:						List[str]	)	->       Optional[Any]:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[int]             =      torch_device == """cpu"""

						self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       Tuple:







						"""simple docstring"""







						__UpperCamelCase					:	int             =      torch_device in ["""cpu""", """mps"""]

						self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase	)






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



	"""simple docstring"""


	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       int:







						"""simple docstring"""







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






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       Tuple:







						"""simple docstring"""







						__UpperCamelCase					:	List[Any]             =      load_numpy(
						    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""	)

						__UpperCamelCase					:	Tuple             =      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""					,     torch_dtype=torch.floataa	)
						pipe.to(lowerCAmelCase	)
						pipe.set_progress_bar_config(disable=lowerCAmelCase	)
						# stable unclip will oom when integration tests are run on a V100,
						# so turn on memory savings
						pipe.enable_attention_slicing()
						pipe.enable_sequential_cpu_offload()

						__UpperCamelCase					:	List[str]             =      torch.Generator(device="""cpu"""	).manual_seed(0	)
						__UpperCamelCase					:	str             =      pipe("""anime turle"""					,     generator=lowerCAmelCase					,     output_type="""np"""	)

						__UpperCamelCase					:	Any             =      output.images[0]

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

						assert_mean_pixel_difference(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       str:







						"""simple docstring"""







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

						__UpperCamelCase					:	Any             =      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""					,     torch_dtype=torch.floataa	)
						__UpperCamelCase					:	Optional[int]             =      pipe.to(lowerCAmelCase	)
						pipe.set_progress_bar_config(disable=lowerCAmelCase	)
						pipe.enable_attention_slicing()
						pipe.enable_sequential_cpu_offload()

						__UpperCamelCase					:	int             =      pipe(
						    """anime turtle"""					,     prior_num_inference_steps=2					,     num_inference_steps=2					,     output_type="""np"""					,     )

						__UpperCamelCase					:	Optional[Any]             =      torch.cuda.max_memory_allocated()
						# make sure that less than 7 GB is allocated
						assert mem_bytes < 7 * 10**9












import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


a__		     =     {
    '''debug''': logging.DEBUG,
    '''info''': logging.INFO,
    '''warning''': logging.WARNING,
    '''error''': logging.ERROR,
    '''critical''': logging.CRITICAL,
}

a__		     =     logging.WARNING





def 				A__ ()     ->						Any:
					__UpperCamelCase					:	Union[str, Any]             =      os.getenv("""DATASETS_VERBOSITY"""		,	snake_case				)
					if env_level_str:
										if env_level_str in log_levels:
															return log_levels[env_level_str]
										else:
															logging.getLogger().warning(
															    F'''Unknown option DATASETS_VERBOSITY={env_level_str}, '''
															    F'''has to be one of: { ", ".join(log_levels.keys()				) }'''				)
					return _default_log_level





def 				A__ ()     ->						str:
					return __name__.split("""."""				)[0]





def 				A__ ()     ->						logging.Logger:
					return logging.getLogger(_get_library_name()				)





def 				A__ ()     ->						None:
					# Apply our default configuration to the library root logger.
					__UpperCamelCase					:	Any             =      _get_library_root_logger()
					library_root_logger.setLevel(_get_default_logging_level()				)





def 				A__ ()     ->						None:
					__UpperCamelCase					:	List[Any]             =      _get_library_root_logger()
					library_root_logger.setLevel(logging.NOTSET				)





def 				A__ (snake_case   :		Optional[str] = None				)     ->						logging.Logger:
					if name is None:
										__UpperCamelCase					:	Optional[int]             =      _get_library_name()
					return logging.getLogger(snake_case				)





def 				A__ ()     ->						int:
					return _get_library_root_logger().getEffectiveLevel()





def 				A__ (snake_case   :		int				)     ->						None:
					_get_library_root_logger().setLevel(snake_case				)





def 				A__ ()     ->						Dict:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						Any:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						Dict:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						str:
					return set_verbosity(snake_case				)





def 				A__ ()     ->						None:
					__UpperCamelCase					:	str             =      False





def 				A__ ()     ->						None:
					__UpperCamelCase					:	List[str]             =      True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()


class    SCREAMING_SNAKE_CASE_				:



	"""simple docstring"""


	def __init__(					self		:						Optional[int]					,     *lowerCAmelCase		:						str					,     **lowerCAmelCase		:						List[str]	)	->       Union[str, Any]:  # pylint: disable=unused-argument







						"""simple docstring"""







						__UpperCamelCase					:	List[str]             =      args[0] if args else None






	def __iter__(					self		:						List[Any]	)	->       List[Any]:







						"""simple docstring"""







						return iter(self._iterator	)






	def __getattr__(					self		:						int					,     lowerCAmelCase		:						List[str]	)	->       Dict:







						"""simple docstring"""







						def empty_fn(*lowerCAmelCase		:						List[Any]					,     **lowerCAmelCase		:						Dict	):  # pylint: disable=unused-argument
											return

						return empty_fn






	def __enter__(					self		:						Tuple	)	->       Any:







						"""simple docstring"""







						return self






	def __exit__(					self		:						str					,     lowerCAmelCase		:						Dict					,     lowerCAmelCase		:						Dict					,     lowerCAmelCase		:						List[str]	)	->       Tuple:







						"""simple docstring"""







						return






a__		     =     True


class    SCREAMING_SNAKE_CASE_				:



	"""simple docstring"""


	def __call__(					self		:						str					,     *lowerCAmelCase		:						List[Any]					,     lowerCAmelCase		:						Dict=False					,     **lowerCAmelCase		:						Tuple	)	->       Optional[Any]:







						"""simple docstring"""







						if _tqdm_active and not disable:
											return tqdm_lib.tqdm(*lowerCAmelCase					,     **lowerCAmelCase	)
						else:
											return EmptyTqdm(*lowerCAmelCase					,     **lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Any					,     *lowerCAmelCase		:						List[str]					,     **lowerCAmelCase		:						str	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	Any             =      None
						if _tqdm_active:
											return tqdm_lib.tqdm.set_lock(*lowerCAmelCase					,     **lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Union[str, Any]:







						"""simple docstring"""







						if _tqdm_active:
											return tqdm_lib.tqdm.get_lock()






a__		     =     _tqdm_cls()





def 				A__ ()     ->						bool:
					global _tqdm_active
					return bool(_tqdm_active				)





def 				A__ ()     ->						List[str]:
					global _tqdm_active
					__UpperCamelCase					:	Tuple             =      True







def 				A__ ()     ->						int:
					global _tqdm_active
					__UpperCamelCase					:	Any             =      False






