
"""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_						:


				"""simple docstring"""



				def 	__magic_name__    (self    )					->		Optional[Any]:



											"""simple docstring"""




											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Any           =   TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""    )

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

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[str]           =   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__						:       List[str]           =   DDPMScheduler(
											    num_train_timesteps=10_00   ,       beta_schedule="""squaredcos_cap_v2"""   ,       beta_start=0.0001   ,       beta_end=0.02   ,       thresholding=SCREAMING_SNAKE_CASE__   ,       dynamic_thresholding_ratio=0.95   ,       sample_max_value=1.0   ,       prediction_type="""epsilon"""   ,       variance_type="""learned_range"""   ,       )

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[str]           =   IFWatermarker()

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



				def 	__magic_name__    (self    )					->		Tuple:



											"""simple docstring"""




											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   TaEncoderModel.from_pretrained("""hf-internal-testing/tiny-random-t5"""    )

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

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   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.414   ,       time_embedding_act_fn="""gelu"""   ,       time_embedding_dim=32   ,       )
											unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   DDPMScheduler(
											    num_train_timesteps=10_00   ,       beta_schedule="""squaredcos_cap_v2"""   ,       beta_start=0.0001   ,       beta_end=0.02   ,       thresholding=SCREAMING_SNAKE_CASE__   ,       dynamic_thresholding_ratio=0.95   ,       sample_max_value=1.0   ,       prediction_type="""epsilon"""   ,       variance_type="""learned_range"""   ,       )

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   DDPMScheduler(
											    num_train_timesteps=10_00   ,       beta_schedule="""squaredcos_cap_v2"""   ,       beta_start=0.0001   ,       beta_end=0.02   ,       )

											torch.manual_seed(0    )
											SCREAMING_SNAKE_CASE__						:       List[str]           =   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 	__magic_name__    (self    )					->		Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Any           =   self.get_dummy_components()
											SCREAMING_SNAKE_CASE__						:       List[str]           =   self.pipeline_class(**SCREAMING_SNAKE_CASE__    )
											pipe.to(SCREAMING_SNAKE_CASE__    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       int           =   inputs["""prompt"""]
											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   inputs["""generator"""]
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   inputs["""num_inference_steps"""]
											SCREAMING_SNAKE_CASE__						:       List[Any]           =   inputs["""output_type"""]

											if "image" in inputs:
																		SCREAMING_SNAKE_CASE__						:       Any           =   inputs["""image"""]
											else:
																		SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   None

											if "mask_image" in inputs:
																		SCREAMING_SNAKE_CASE__						:       str           =   inputs["""mask_image"""]
											else:
																		SCREAMING_SNAKE_CASE__						:       str           =   None

											if "original_image" in inputs:
																		SCREAMING_SNAKE_CASE__						:       str           =   inputs["""original_image"""]
											else:
																		SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   None

											SCREAMING_SNAKE_CASE__      ,      SCREAMING_SNAKE_CASE__						:       int           =   pipe.encode_prompt(SCREAMING_SNAKE_CASE__    )

											# inputs with prompt converted to embeddings
											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   {
											    """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__						:       Optional[Any]           =   image

											if mask_image is not None:
																		SCREAMING_SNAKE_CASE__						:       Optional[int]           =   mask_image

											if original_image is not None:
																		SCREAMING_SNAKE_CASE__						:       str           =   original_image

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

											SCREAMING_SNAKE_CASE__						:       Dict           =   pipe(**SCREAMING_SNAKE_CASE__    )[0]

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

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

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

											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       int           =   inputs["""generator"""]
											SCREAMING_SNAKE_CASE__						:       Any           =   inputs["""num_inference_steps"""]
											SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   inputs["""output_type"""]

											# inputs with prompt converted to embeddings
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   {
											    """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__						:       Any           =   image

											if mask_image is not None:
																		SCREAMING_SNAKE_CASE__						:       int           =   mask_image

											if original_image is not None:
																		SCREAMING_SNAKE_CASE__						:       str           =   original_image

											SCREAMING_SNAKE_CASE__						:       List[str]           =   pipe_loaded(**SCREAMING_SNAKE_CASE__    )[0]

											SCREAMING_SNAKE_CASE__						:       Dict           =   np.abs(to_np(SCREAMING_SNAKE_CASE__    ) - to_np(SCREAMING_SNAKE_CASE__    )    ).max()
											self.assertLess(SCREAMING_SNAKE_CASE__   ,       1E-4    )



				def 	__magic_name__    (self    )					->		str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:       Optional[Any]           =   self.get_dummy_components()
											SCREAMING_SNAKE_CASE__						:       Optional[int]           =   self.pipeline_class(**SCREAMING_SNAKE_CASE__    )
											pipe.to(SCREAMING_SNAKE_CASE__    )
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__    )

											SCREAMING_SNAKE_CASE__						:       List[Any]           =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       Tuple           =   pipe(**SCREAMING_SNAKE_CASE__    )[0]

											with tempfile.TemporaryDirectory() as tmpdir:
																		pipe.save_pretrained(SCREAMING_SNAKE_CASE__    )
																		SCREAMING_SNAKE_CASE__						:       Union[str, Any]           =   self.pipeline_class.from_pretrained(SCREAMING_SNAKE_CASE__    )
																		pipe_loaded.to(SCREAMING_SNAKE_CASE__    )
																		pipe_loaded.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__    )

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

											SCREAMING_SNAKE_CASE__						:       List[str]           =   self.get_dummy_inputs(SCREAMING_SNAKE_CASE__    )
											SCREAMING_SNAKE_CASE__						:       str           =   pipe_loaded(**SCREAMING_SNAKE_CASE__    )[0]

											SCREAMING_SNAKE_CASE__						:       Tuple           =   np.abs(to_np(SCREAMING_SNAKE_CASE__    ) - to_np(SCREAMING_SNAKE_CASE__    )    ).max()
											self.assertLess(SCREAMING_SNAKE_CASE__   ,       1E-4    )












import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
	import torch

	from transformers import AutoModelForImageClassification

if is_vision_available():
	from transformers import AutoImageProcessor




@require_torch
@require_vision
class  UpperCamelCase_      (			unittest.TestCase				):



	'''simple docstring'''




	@slow
	def 							_UpperCamelCase			(       self					)       ->    int:
								snake_case_				=							AutoImageProcessor.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'					)
								snake_case_				=							AutoModelForImageClassification.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'					)
								model.to(a					)

								from datasets import load_dataset

								snake_case_				=							load_dataset('nielsr/rvlcdip-demo'					)

								snake_case_				=							dataset['train'][0]['image'].convert('RGB'					)

								snake_case_				=							image_processor(a       ,			return_tensors='pt'					).to(a					)

								# forward pass
								with torch.no_grad():
															snake_case_				=							model(**a					)
															snake_case_				=							outputs.logits

								snake_case_				=							torch.Size((1, 16)					)
								self.assertEqual(logits.shape       ,			a					)

								snake_case_				=							torch.tensor(
								    [-0.4_158, -0.4_092, -0.4_347]       ,			device=a       ,			dtype=torch.float       ,			)
								self.assertTrue(torch.allclose(logits[0, :3]       ,			a       ,			atol=1E-4					)					)





