






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 						__snake_case					(   self     :   int	):
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''	)

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

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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	)
											lowerCAmelCase__      			=						DDPMScheduler(
											    num_train_timesteps=1_000 ,							beta_schedule='''squaredcos_cap_v2''' ,							beta_start=0.0_001 ,							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	)
											lowerCAmelCase__      			=						IFWatermarker()

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




							def 						__snake_case					(   self     :   Any	):
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''	)

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

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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	)
											lowerCAmelCase__      			=						DDPMScheduler(
											    num_train_timesteps=1_000 ,							beta_schedule='''squaredcos_cap_v2''' ,							beta_start=0.0_001 ,							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	)
											lowerCAmelCase__      			=						DDPMScheduler(
											    num_train_timesteps=1_000 ,							beta_schedule='''squaredcos_cap_v2''' ,							beta_start=0.0_001 ,							beta_end=0.02 ,							)

											torch.manual_seed(0	)
											lowerCAmelCase__      			=						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 						__snake_case					(   self     :   Any	):
											lowerCAmelCase__      			=						self.get_dummy_components()
											lowerCAmelCase__      			=						self.pipeline_class(**SCREAMING_SNAKE_CASE_	)
											pipe.to(SCREAMING_SNAKE_CASE_	)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						inputs['''prompt''']
											lowerCAmelCase__      			=						inputs['''generator''']
											lowerCAmelCase__      			=						inputs['''num_inference_steps''']
											lowerCAmelCase__      			=						inputs['''output_type''']

											if "image" in inputs:
															lowerCAmelCase__      			=						inputs['''image''']
											else:
															lowerCAmelCase__      			=						None

											if "mask_image" in inputs:
															lowerCAmelCase__      			=						inputs['''mask_image''']
											else:
															lowerCAmelCase__      			=						None

											if "original_image" in inputs:
															lowerCAmelCase__      			=						inputs['''original_image''']
											else:
															lowerCAmelCase__      			=						None

											lowerCAmelCase__			,		lowerCAmelCase__      			=						pipe.encode_prompt(SCREAMING_SNAKE_CASE_	)

											# inputs with prompt converted to embeddings
											lowerCAmelCase__      			=						{
											    '''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:
															lowerCAmelCase__      			=						image

											if mask_image is not None:
															lowerCAmelCase__      			=						mask_image

											if original_image is not None:
															lowerCAmelCase__      			=						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_	)

											lowerCAmelCase__      			=						pipe(**SCREAMING_SNAKE_CASE_	)[0]

											with tempfile.TemporaryDirectory() as tmpdir:
															pipe.save_pretrained(SCREAMING_SNAKE_CASE_	)
															lowerCAmelCase__      			=						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.' ,							)

											lowerCAmelCase__      			=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						inputs['''generator''']
											lowerCAmelCase__      			=						inputs['''num_inference_steps''']
											lowerCAmelCase__      			=						inputs['''output_type''']

											# inputs with prompt converted to embeddings
											lowerCAmelCase__      			=						{
											    '''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:
															lowerCAmelCase__      			=						image

											if mask_image is not None:
															lowerCAmelCase__      			=						mask_image

											if original_image is not None:
															lowerCAmelCase__      			=						original_image

											lowerCAmelCase__      			=						pipe_loaded(**SCREAMING_SNAKE_CASE_	)[0]

											lowerCAmelCase__      			=						np.abs(to_np(SCREAMING_SNAKE_CASE_	) - to_np(SCREAMING_SNAKE_CASE_	)	).max()
											self.assertLess(SCREAMING_SNAKE_CASE_ ,							1e-4	)




							def 						__snake_case					(   self     :   Tuple	):
											lowerCAmelCase__      			=						self.get_dummy_components()
											lowerCAmelCase__      			=						self.pipeline_class(**SCREAMING_SNAKE_CASE_	)
											pipe.to(SCREAMING_SNAKE_CASE_	)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						pipe(**SCREAMING_SNAKE_CASE_	)[0]

											with tempfile.TemporaryDirectory() as tmpdir:
															pipe.save_pretrained(SCREAMING_SNAKE_CASE_	)
															lowerCAmelCase__      			=						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

											lowerCAmelCase__      			=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						pipe_loaded(**SCREAMING_SNAKE_CASE_	)[0]

											lowerCAmelCase__      			=						np.abs(to_np(SCREAMING_SNAKE_CASE_	) - to_np(SCREAMING_SNAKE_CASE_	)	).max()
											self.assertLess(SCREAMING_SNAKE_CASE_ ,							1e-4	)













import pytest

import datasets


# Import fixture modules as plugins
_UpperCAmelCase			:		Tuple   				=    ["tests.fixtures.files", "tests.fixtures.hub", "tests.fixtures.fsspec"]





def 	lowerCAmelCase_		(lowercase__						:       Optional[int]     ,					lowercase__						:       int       )		->    Any:


				'''simple docstring'''




				for item in items:
								if any(marker in item.keywords for marker in ['''integration''', '''unit''']       ):
												continue
								item.add_marker(pytest.mark.unit       )





def 	lowerCAmelCase_		(lowercase__						:       str       )		->    List[str]:


				'''simple docstring'''




				config.addinivalue_line('''markers'''     ,					'''torchaudio_latest: mark test to run with torchaudio>=0.12'''       )





@pytest.fixture(autouse=lowercase__       )
def 	lowerCAmelCase_		(lowercase__						:       int     ,					lowercase__						:       Dict       )		->    Optional[Any]:


				'''simple docstring'''




				lowerCAmelCase__      			=						tmp_path_factory.getbasetemp() / '''cache'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''datasets'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''metrics'''
				lowerCAmelCase__      			=						test_hf_cache_home / '''modules'''
				monkeypatch.setattr('''datasets.config.HF_DATASETS_CACHE'''     ,					str(lowercase__       )       )
				monkeypatch.setattr('''datasets.config.HF_METRICS_CACHE'''     ,					str(lowercase__       )       )
				monkeypatch.setattr('''datasets.config.HF_MODULES_CACHE'''     ,					str(lowercase__       )       )
				lowerCAmelCase__      			=						test_hf_datasets_cache / '''downloads'''
				monkeypatch.setattr('''datasets.config.DOWNLOADED_DATASETS_PATH'''     ,					str(lowercase__       )       )
				lowerCAmelCase__      			=						test_hf_datasets_cache / '''downloads''' / '''extracted'''
				monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''     ,					str(lowercase__       )       )





@pytest.fixture(autouse=lowercase__     ,					scope='''session'''       )
def 	lowerCAmelCase_		()		->    int:


				'''simple docstring'''




				datasets.disable_progress_bar()





@pytest.fixture(autouse=lowercase__       )
def 	lowerCAmelCase_		(lowercase__						:       Dict       )		->    List[str]:


				'''simple docstring'''




				monkeypatch.setattr('''datasets.config.HF_UPDATE_DOWNLOAD_COUNTS'''     ,					lowercase__       )





@pytest.fixture
def 	lowerCAmelCase_		(lowercase__						:       List[str]       )		->    List[Any]:


				'''simple docstring'''




				monkeypatch.setattr('''sqlalchemy.util.deprecations.SILENCE_UBER_WARNING'''     ,					lowercase__       )





