
'''simple docstring'''







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_snake_case				:   List[Any]     =				False


class 			lowerCAmelCase (				unittest.TestCase			):




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




		@property
		def     lowercase       (				self      ):
								return 12




		@property
		def     lowercase       (				self      ):
								return 12




		@property
		def     lowercase       (				self      ):
								return 32




		@property
		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= VQModel(
								    block_out_channels=[32, 64]					,						in_channels=3					,						out_channels=3					,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]					,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]					,						latent_channels=3					,						num_vq_embeddings=self.num_embed					,						vq_embed_dim=3					,						)
								return model




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




		@property
		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= CLIPTextConfig(
								    bos_token_id=0					,						eos_token_id=2					,						hidden_size=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 CLIPTextModel(UpperCamelCase      )




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

								_SCREAMING_SNAKE_CASE				= 12
								_SCREAMING_SNAKE_CASE				= 12

								_SCREAMING_SNAKE_CASE				= {
								    "attention_bias": True,
								    "cross_attention_dim": 32,
								    "attention_head_dim": height * width,
								    "num_attention_heads": 1,
								    "num_vector_embeds": self.num_embed,
								    "num_embeds_ada_norm": self.num_embeds_ada_norm,
								    "norm_num_groups": 32,
								    "sample_size": width,
								    "activation_fn": "geglu-approximate",
								}

								_SCREAMING_SNAKE_CASE				= TransformeraDModel(**UpperCamelCase      )
								return model




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= "cpu"

								_SCREAMING_SNAKE_CASE				= self.dummy_vqvae
								_SCREAMING_SNAKE_CASE				= self.dummy_text_encoder
								_SCREAMING_SNAKE_CASE				= self.dummy_tokenizer
								_SCREAMING_SNAKE_CASE				= self.dummy_transformer
								_SCREAMING_SNAKE_CASE				= VQDiffusionScheduler(self.num_embed      )
								_SCREAMING_SNAKE_CASE				= LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= VQDiffusionPipeline(
								    vqvae=UpperCamelCase					,						text_encoder=UpperCamelCase					,						tokenizer=UpperCamelCase					,						transformer=UpperCamelCase					,						scheduler=UpperCamelCase					,						learned_classifier_free_sampling_embeddings=UpperCamelCase					,						)
								_SCREAMING_SNAKE_CASE				= pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= "teddy bear playing in the pool"

								_SCREAMING_SNAKE_CASE				= torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= pipe([prompt]					,						generator=UpperCamelCase					,						num_inference_steps=2					,						output_type="np"      )
								_SCREAMING_SNAKE_CASE				= output.images

								_SCREAMING_SNAKE_CASE				= torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= pipe(
								    [prompt]					,						generator=UpperCamelCase					,						output_type="np"					,						return_dict=UpperCamelCase					,						num_inference_steps=2      )[0]

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

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

								_SCREAMING_SNAKE_CASE				= np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]      )

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

								_SCREAMING_SNAKE_CASE				= self.dummy_vqvae
								_SCREAMING_SNAKE_CASE				= self.dummy_text_encoder
								_SCREAMING_SNAKE_CASE				= self.dummy_tokenizer
								_SCREAMING_SNAKE_CASE				= self.dummy_transformer
								_SCREAMING_SNAKE_CASE				= VQDiffusionScheduler(self.num_embed      )
								_SCREAMING_SNAKE_CASE				= LearnedClassifierFreeSamplingEmbeddings(
								    learnable=UpperCamelCase					,						hidden_size=self.text_embedder_hidden_size					,						length=tokenizer.model_max_length      )

								_SCREAMING_SNAKE_CASE				= VQDiffusionPipeline(
								    vqvae=UpperCamelCase					,						text_encoder=UpperCamelCase					,						tokenizer=UpperCamelCase					,						transformer=UpperCamelCase					,						scheduler=UpperCamelCase					,						learned_classifier_free_sampling_embeddings=UpperCamelCase					,						)
								_SCREAMING_SNAKE_CASE				= pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= "teddy bear playing in the pool"

								_SCREAMING_SNAKE_CASE				= torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= pipe([prompt]					,						generator=UpperCamelCase					,						num_inference_steps=2					,						output_type="np"      )
								_SCREAMING_SNAKE_CASE				= output.images

								_SCREAMING_SNAKE_CASE				= torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= pipe(
								    [prompt]					,						generator=UpperCamelCase					,						output_type="np"					,						return_dict=UpperCamelCase					,						num_inference_steps=2      )[0]

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

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

								_SCREAMING_SNAKE_CASE				= np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]      )

								assert np.abs(image_slice.flatten() - expected_slice      ).max() < 2.0
								assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2


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




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




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"      )

								_SCREAMING_SNAKE_CASE				= VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"      )
								_SCREAMING_SNAKE_CASE				= pipeline.to(UpperCamelCase      )
								pipeline.set_progress_bar_config(disable=UpperCamelCase      )

								# requires GPU generator for gumbel softmax
								# don't use GPU generator in tests though
								_SCREAMING_SNAKE_CASE				= torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= pipeline(
								    "teddy bear playing in the pool"					,						num_images_per_prompt=1					,						generator=UpperCamelCase					,						output_type="np"					,						)

								_SCREAMING_SNAKE_CASE				= output.images[0]

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

'''simple docstring'''







import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np


class 			lowerCAmelCase :




		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= UNetaDConditionModel(
								    sample_size=32					,						layers_per_block=1					,						block_out_channels=[32, 64]					,						down_block_types=[
								        "ResnetDownsampleBlock2D",
								        "SimpleCrossAttnDownBlock2D",
								    ]					,						mid_block_type="UNetMidBlock2DSimpleCrossAttn"					,						up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]					,						in_channels=3					,						out_channels=6					,						cross_attention_dim=32					,						encoder_hid_dim=32					,						attention_head_dim=8					,						addition_embed_type="text"					,						addition_embed_type_num_heads=2					,						cross_attention_norm="group_norm"					,						resnet_time_scale_shift="scale_shift"					,						act_fn="gelu"					,						)
								unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						thresholding=UpperCamelCase					,						dynamic_thresholding_ratio=0.95					,						sample_max_value=1.0					,						prediction_type="epsilon"					,						variance_type="learned_range"					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}




		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= UNetaDConditionModel(
								    sample_size=32					,						layers_per_block=[1, 2]					,						block_out_channels=[32, 64]					,						down_block_types=[
								        "ResnetDownsampleBlock2D",
								        "SimpleCrossAttnDownBlock2D",
								    ]					,						mid_block_type="UNetMidBlock2DSimpleCrossAttn"					,						up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]					,						in_channels=6					,						out_channels=6					,						cross_attention_dim=32					,						encoder_hid_dim=32					,						attention_head_dim=8					,						addition_embed_type="text"					,						addition_embed_type_num_heads=2					,						cross_attention_norm="group_norm"					,						resnet_time_scale_shift="scale_shift"					,						act_fn="gelu"					,						class_embed_type="timestep"					,						mid_block_scale_factor=1.4_14					,						time_embedding_act_fn="gelu"					,						time_embedding_dim=32					,						)
								unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						thresholding=UpperCamelCase					,						dynamic_thresholding_ratio=0.95					,						sample_max_value=1.0					,						prediction_type="epsilon"					,						variance_type="learned_range"					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "image_noising_scheduler": image_noising_scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.get_dummy_components()
								_SCREAMING_SNAKE_CASE				= self.pipeline_class(**UpperCamelCase      )
								pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= inputs["prompt"]
								_SCREAMING_SNAKE_CASE				= inputs["generator"]
								_SCREAMING_SNAKE_CASE				= inputs["num_inference_steps"]
								_SCREAMING_SNAKE_CASE				= inputs["output_type"]

								if "image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								if "mask_image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["mask_image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								if "original_image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["original_image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= pipe.encode_prompt(UpperCamelCase      )

								# inputs with prompt converted to embeddings
								_SCREAMING_SNAKE_CASE				= {
								    "prompt_embeds": prompt_embeds,
								    "negative_prompt_embeds": negative_prompt_embeds,
								    "generator": generator,
								    "num_inference_steps": num_inference_steps,
								    "output_type": output_type,
								}

								if image is not None:
														_SCREAMING_SNAKE_CASE				= image

								if mask_image is not None:
														_SCREAMING_SNAKE_CASE				= mask_image

								if original_image is not None:
														_SCREAMING_SNAKE_CASE				= original_image

								# set all optional components to None
								for optional_component in pipe._optional_components:
														setattr(UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= pipe(**UpperCamelCase      )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(UpperCamelCase      )
														_SCREAMING_SNAKE_CASE				= self.pipeline_class.from_pretrained(UpperCamelCase      )
														pipe_loaded.to(UpperCamelCase      )
														pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								for optional_component in pipe._optional_components:
														self.assertTrue(
														    getattr(UpperCamelCase					,						UpperCamelCase      ) is None					,						F'`{optional_component}` did not stay set to None after loading.'					,						)

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= inputs["generator"]
								_SCREAMING_SNAKE_CASE				= inputs["num_inference_steps"]
								_SCREAMING_SNAKE_CASE				= inputs["output_type"]

								# inputs with prompt converted to embeddings
								_SCREAMING_SNAKE_CASE				= {
								    "prompt_embeds": prompt_embeds,
								    "negative_prompt_embeds": negative_prompt_embeds,
								    "generator": generator,
								    "num_inference_steps": num_inference_steps,
								    "output_type": output_type,
								}

								if image is not None:
														_SCREAMING_SNAKE_CASE				= image

								if mask_image is not None:
														_SCREAMING_SNAKE_CASE				= mask_image

								if original_image is not None:
														_SCREAMING_SNAKE_CASE				= original_image

								_SCREAMING_SNAKE_CASE				= pipe_loaded(**UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
								self.assertLess(UpperCamelCase					,						1e-4      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.get_dummy_components()
								_SCREAMING_SNAKE_CASE				= self.pipeline_class(**UpperCamelCase      )
								pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= pipe(**UpperCamelCase      )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(UpperCamelCase      )
														_SCREAMING_SNAKE_CASE				= self.pipeline_class.from_pretrained(UpperCamelCase      )
														pipe_loaded.to(UpperCamelCase      )
														pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= pipe_loaded(**UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
								self.assertLess(UpperCamelCase					,						1e-4      )