


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 						UpperCAmelCase__   (      A_      ,						A_      ,						unittest.TestCase		):



     '''simple docstring'''
     UpperCAmelCase_         =  StableDiffusionParadigmsPipeline
     UpperCAmelCase_         =  TEXT_TO_IMAGE_PARAMS
     UpperCAmelCase_         =  TEXT_TO_IMAGE_BATCH_PARAMS
     UpperCAmelCase_         =  TEXT_TO_IMAGE_IMAGE_PARAMS
     UpperCAmelCase_         =  TEXT_TO_IMAGE_IMAGE_PARAMS




     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



            torch.manual_seed(0		)
            _lowercase					:				Dict           =			UNetaDConditionModel(
                block_out_channels=(32, 64)	,      layers_per_block=2	,      sample_size=32	,      in_channels=4	,      out_channels=4	,      down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,      up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,      cross_attention_dim=32	,      attention_head_dim=(2, 4)	,      use_linear_projection=UpperCamelCase	,      )
            _lowercase					:				Optional[Any]           =			DDIMParallelScheduler(
                beta_start=0.00085	,      beta_end=0.012	,      beta_schedule='''scaled_linear'''	,      clip_sample=UpperCamelCase	,      set_alpha_to_one=UpperCamelCase	,      )
            torch.manual_seed(0		)
            _lowercase					:				str           =			AutoencoderKL(
                block_out_channels=[32, 64]	,      in_channels=3	,      out_channels=3	,      down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,      up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,      latent_channels=4	,      sample_size=1_28	,      )
            torch.manual_seed(0		)
            _lowercase					:				Tuple           =			CLIPTextConfig(
                bos_token_id=0	,      eos_token_id=2	,      hidden_size=32	,      intermediate_size=37	,      layer_norm_eps=1E-05	,      num_attention_heads=4	,      num_hidden_layers=5	,      pad_token_id=1	,      vocab_size=10_00	,      hidden_act='''gelu'''	,      projection_dim=5_12	,      )
            _lowercase					:				Union[str, Any]           =			CLIPTextModel(UpperCamelCase		)
            _lowercase					:				List[str]           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

            _lowercase					:				Optional[int]           =			{
                '''unet''': unet,
                '''scheduler''': scheduler,
                '''vae''': vae,
                '''text_encoder''': text_encoder,
                '''tokenizer''': tokenizer,
                '''safety_checker''': None,
                '''feature_extractor''': None,
            }
            return components




     def   lowerCAmelCase_   ( self       :       List[str]	,      UpperCamelCase       :       List[Any]	,      UpperCamelCase       :       Any=0		):





            """simple docstring"""



            if str(UpperCamelCase		).startswith('''mps'''		):
                   _lowercase					:				int           =			torch.manual_seed(UpperCamelCase		)
            else:
                   _lowercase					:				Union[str, Any]           =			torch.Generator(device=UpperCamelCase		).manual_seed(UpperCamelCase		)
            _lowercase					:				Any           =			{
                '''prompt''': '''a photograph of an astronaut riding a horse''',
                '''generator''': generator,
                '''num_inference_steps''': 10,
                '''guidance_scale''': 6.0,
                '''output_type''': '''numpy''',
                '''parallel''': 3,
                '''debug''': True,
            }
            return inputs




     def   lowerCAmelCase_   ( self       :       List[str]		):





            """simple docstring"""



            _lowercase					:				str           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
            _lowercase					:				List[str]           =			self.get_dummy_components()
            _lowercase					:				str           =			StableDiffusionParadigmsPipeline(**UpperCamelCase		)
            _lowercase					:				int           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Dict           =			self.get_dummy_inputs(UpperCamelCase		)
            _lowercase					:				Union[str, Any]           =			sd_pipe(**UpperCamelCase		).images
            _lowercase					:				Optional[Any]           =			image[0, -3:, -3:, -1]
            assert image.shape == (1, 64, 64, 3)

            _lowercase					:				Tuple           =			np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]		)

            assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2




     def   lowerCAmelCase_   ( self       :       Optional[int]		):





            """simple docstring"""



            _lowercase					:				List[Any]           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
            _lowercase					:				int           =			self.get_dummy_components()
            torch.manual_seed(0		)
            _lowercase					:				Union[str, Any]           =			DDPMParallelScheduler()
            torch.manual_seed(0		)
            _lowercase					:				Any           =			StableDiffusionParadigmsPipeline(**UpperCamelCase		)
            _lowercase					:				Tuple           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Union[str, Any]           =			self.get_dummy_inputs(UpperCamelCase		)
            _lowercase					:				Optional[int]           =			sd_pipe(**UpperCamelCase		).images
            _lowercase					:				Tuple           =			image[0, -3:, -3:, -1]
            assert image.shape == (1, 64, 64, 3)

            _lowercase					:				Optional[Any]           =			np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]		)

            assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2




     def   lowerCAmelCase_   ( self       :       Union[str, Any]		):





            """simple docstring"""



            super().test_inference_batch_consistent(batch_sizes=[1, 2]		)




     def   lowerCAmelCase_   ( self       :       Any		):





            """simple docstring"""



            super().test_inference_batch_single_identical(batch_size=2	,      expected_max_diff=3E-3		)




     def   lowerCAmelCase_   ( self       :       Optional[Any]		):





            """simple docstring"""



            _lowercase					:				str           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
            _lowercase					:				Tuple           =			self.get_dummy_components()
            _lowercase					:				Tuple           =			StableDiffusionParadigmsPipeline(**UpperCamelCase		)
            _lowercase					:				List[Any]           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Any           =			self.get_dummy_inputs(UpperCamelCase		)
            _lowercase					:				Dict           =			'''french fries'''
            _lowercase					:				List[Any]           =			sd_pipe(**UpperCamelCase	,      negative_prompt=UpperCamelCase		)
            _lowercase					:				Any           =			output.images
            _lowercase					:				Tuple           =			image[0, -3:, -3:, -1]

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

            _lowercase					:				List[Any]           =			np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]		)

            assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2




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



     '''simple docstring'''




     def   lowerCAmelCase_   ( self       :       Tuple		):





            """simple docstring"""



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




     def   lowerCAmelCase_   ( self       :       List[Any]	,      UpperCamelCase       :       Union[str, Any]=0		):





            """simple docstring"""



            _lowercase					:				int           =			torch.Generator(device=UpperCamelCase		).manual_seed(UpperCamelCase		)
            _lowercase					:				List[Any]           =			{
                '''prompt''': '''a photograph of an astronaut riding a horse''',
                '''generator''': generator,
                '''num_inference_steps''': 10,
                '''guidance_scale''': 7.5,
                '''output_type''': '''numpy''',
                '''parallel''': 3,
                '''debug''': True,
            }
            return inputs




     def   lowerCAmelCase_   ( self       :       Any		):





            """simple docstring"""



            _lowercase					:				Dict           =			'''stabilityai/stable-diffusion-2-base'''
            _lowercase					:				str           =			DDIMParallelScheduler.from_pretrained(UpperCamelCase	,      subfolder='''scheduler'''		)
            _lowercase					:				List[str]           =			StableDiffusionParadigmsPipeline.from_pretrained(UpperCamelCase	,      scheduler=UpperCamelCase	,      safety_checker=UpperCamelCase		)
            pipe.to(UpperCamelCase		)
            pipe.set_progress_bar_config(disable=UpperCamelCase		)
            pipe.enable_attention_slicing()

            _lowercase					:				Tuple           =			self.get_inputs()
            _lowercase					:				Dict           =			pipe(**UpperCamelCase		).images
            _lowercase					:				Tuple           =			image[0, -3:, -3:, -1].flatten()

            assert image.shape == (1, 5_12, 5_12, 3)

            _lowercase					:				Tuple           =			np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]		)

            assert np.abs(expected_slice - image_slice		).max() < 1E-2



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__				       =			logging.get_logger(__name__)

UpperCamelCase__				       =			{
    'facebook/sam-vit-huge': 'https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json',
    'facebook/sam-vit-large': 'https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json',
    'facebook/sam-vit-base': 'https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json',
}

class 						UpperCAmelCase__   (      A_		):



     '''simple docstring'''




     def __init__( self       :       List[Any]	,      UpperCamelCase       :       Optional[Any]=2_56	,      UpperCamelCase       :       str=10_24	,      UpperCamelCase       :       Any=16	,      UpperCamelCase       :       Optional[Any]=16	,      UpperCamelCase       :       Union[str, Any]=4	,      UpperCamelCase       :       str="gelu"	,      UpperCamelCase       :       Dict=1E-6	,      **UpperCamelCase       :       Tuple	,      ):





            """simple docstring"""



            super().__init__(**UpperCamelCase		)
            _lowercase					:				Union[str, Any]           =			hidden_size
            _lowercase					:				Optional[Any]           =			image_size
            _lowercase					:				Tuple           =			patch_size
            _lowercase					:				List[str]           =			image_size // patch_size
            _lowercase					:				Optional[Any]           =			mask_input_channels
            _lowercase					:				str           =			num_point_embeddings
            _lowercase					:				List[Any]           =			hidden_act
            _lowercase					:				Union[str, Any]           =			layer_norm_eps




class 						UpperCAmelCase__   (      A_		):



     '''simple docstring'''




     def __init__( self       :       Optional[int]	,      UpperCamelCase       :       str=2_56	,      UpperCamelCase       :       int="relu"	,      UpperCamelCase       :       Any=20_48	,      UpperCamelCase       :       List[Any]=2	,      UpperCamelCase       :       Union[str, Any]=8	,      UpperCamelCase       :       List[str]=2	,      UpperCamelCase       :       Any=3	,      UpperCamelCase       :       List[Any]=3	,      UpperCamelCase       :       int=2_56	,      UpperCamelCase       :       List[str]=1E-6	,      **UpperCamelCase       :       Tuple	,      ):





            """simple docstring"""



            super().__init__(**UpperCamelCase		)
            _lowercase					:				List[Any]           =			hidden_size
            _lowercase					:				Optional[int]           =			hidden_act
            _lowercase					:				Union[str, Any]           =			mlp_dim
            _lowercase					:				Optional[int]           =			num_hidden_layers
            _lowercase					:				List[str]           =			num_attention_heads
            _lowercase					:				Dict           =			attention_downsample_rate
            _lowercase					:				Optional[Any]           =			num_multimask_outputs
            _lowercase					:				Optional[Any]           =			iou_head_depth
            _lowercase					:				Dict           =			iou_head_hidden_dim
            _lowercase					:				List[str]           =			layer_norm_eps




class 						UpperCAmelCase__   (      A_		):



     '''simple docstring'''




     def __init__( self       :       List[str]	,      UpperCamelCase       :       Union[str, Any]=7_68	,      UpperCamelCase       :       Optional[Any]=2_56	,      UpperCamelCase       :       Optional[Any]=12	,      UpperCamelCase       :       Dict=12	,      UpperCamelCase       :       Dict=3	,      UpperCamelCase       :       Optional[int]=10_24	,      UpperCamelCase       :       Any=16	,      UpperCamelCase       :       Union[str, Any]="gelu"	,      UpperCamelCase       :       int=1E-06	,      UpperCamelCase       :       Tuple=0.0	,      UpperCamelCase       :       List[str]=1E-10	,      UpperCamelCase       :       List[Any]=True	,      UpperCamelCase       :       Any=4.0	,      UpperCamelCase       :       str=True	,      UpperCamelCase       :       Optional[int]=True	,      UpperCamelCase       :       Optional[int]=14	,      UpperCamelCase       :       Any=[2, 5, 8, 11]	,      UpperCamelCase       :       Tuple=1_28	,      UpperCamelCase       :       str=None	,      **UpperCamelCase       :       int	,      ):





            """simple docstring"""



            super().__init__(**UpperCamelCase		)

            _lowercase					:				Union[str, Any]           =			hidden_size
            _lowercase					:				Optional[int]           =			output_channels
            _lowercase					:				Dict           =			num_hidden_layers
            _lowercase					:				str           =			num_attention_heads
            _lowercase					:				Dict           =			num_channels
            _lowercase					:				List[Any]           =			image_size
            _lowercase					:				int           =			patch_size
            _lowercase					:				Optional[Any]           =			hidden_act
            _lowercase					:				Any           =			layer_norm_eps
            _lowercase					:				Any           =			attention_dropout
            _lowercase					:				Optional[int]           =			initializer_range
            _lowercase					:				List[str]           =			qkv_bias
            _lowercase					:				Optional[Any]           =			mlp_ratio
            _lowercase					:				Optional[Any]           =			use_abs_pos
            _lowercase					:				List[Any]           =			use_rel_pos
            _lowercase					:				str           =			window_size
            _lowercase					:				Tuple           =			global_attn_indexes
            _lowercase					:				List[Any]           =			num_pos_feats
            _lowercase					:				int           =			int(hidden_size * mlp_ratio		) if mlp_dim is None else mlp_dim




class 						UpperCAmelCase__   (      A_		):



     '''simple docstring'''
     UpperCAmelCase_         =  '''sam'''
     UpperCAmelCase_         =  True




     def __init__( self       :       Tuple	,      UpperCamelCase       :       Tuple=None	,      UpperCamelCase       :       Any=None	,      UpperCamelCase       :       Tuple=None	,      UpperCamelCase       :       Dict=0.02	,      **UpperCamelCase       :       Union[str, Any]	,      ):





            """simple docstring"""



            super().__init__(**UpperCamelCase		)
            _lowercase					:				Optional[int]           =			vision_config if vision_config is not None else {}
            _lowercase					:				Any           =			prompt_encoder_config if prompt_encoder_config is not None else {}
            _lowercase					:				Optional[Any]           =			mask_decoder_config if mask_decoder_config is not None else {}

            if isinstance(UpperCamelCase	,      UpperCamelCase		):
                   _lowercase					:				Dict           =			vision_config.to_dict()
            if isinstance(UpperCamelCase	,      UpperCamelCase		):
                   _lowercase					:				List[str]           =			prompt_encoder_config.to_dict()
            if isinstance(UpperCamelCase	,      UpperCamelCase		):
                   _lowercase					:				Optional[int]           =			mask_decoder_config.to_dict()

            _lowercase					:				Union[str, Any]           =			SamVisionConfig(**UpperCamelCase		)
            _lowercase					:				Optional[Any]           =			SamPromptEncoderConfig(**UpperCamelCase		)
            _lowercase					:				List[Any]           =			SamMaskDecoderConfig(**UpperCamelCase		)
            _lowercase					:				Union[str, Any]           =			initializer_range




     def   lowerCAmelCase_   ( self       :       Union[str, Any]		):





            """simple docstring"""



            _lowercase					:				List[Any]           =			copy.deepcopy(self.__dict__		)
            _lowercase					:				int           =			self.vision_config.to_dict()
            _lowercase					:				Optional[int]           =			self.prompt_encoder_config.to_dict()
            _lowercase					:				List[Any]           =			self.mask_decoder_config.to_dict()
            _lowercase					:				List[Any]           =			self.__class__.model_type
            return output