



import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__     , unittest.TestCase   ):
       __lowerCamelCase		:		Optional[int]      =RePaintPipeline
       __lowerCamelCase		:		Union[str, Any]      =IMAGE_INPAINTING_PARAMS - {'width', 'height', 'guidance_scale'}
       __lowerCamelCase		:		int      =PipelineTesterMixin.required_optional_params - {
           'latents',
           'num_images_per_prompt',
           'callback',
           'callback_steps',
       }
       __lowerCamelCase		:		List[Any]      =IMAGE_INPAINTING_BATCH_PARAMS
       __lowerCamelCase		:		List[Any]      =False





       def   UpperCamelCase_   (				self						:						Any		):

           '''simple docstring'''
           torch.manual_seed(0		)
           torch.manual_seed(0		)
           __a					       =						UNetaDModel(
               block_out_channels=(32, 64) ,			layers_per_block=2 ,			sample_size=32 ,			in_channels=3 ,			out_channels=3 ,			down_block_types=("""DownBlock2D""", """AttnDownBlock2D""") ,			up_block_types=("""AttnUpBlock2D""", """UpBlock2D""") ,			)
           __a					       =						RePaintScheduler()
           __a					       =						{"""unet""": unet, """scheduler""": scheduler}
           return components





       def   UpperCamelCase_   (				self						:						Optional[int] ,			__lowercase						:						Dict ,			__lowercase						:						List[str]=0		):

           '''simple docstring'''
           if str(__lowercase		).startswith("""mps"""		):
               __a					       =						torch.manual_seed(__lowercase		)
           else:
               __a					       =						torch.Generator(device=__lowercase		).manual_seed(__lowercase		)
           __a					       =						np.random.RandomState(__lowercase		).standard_normal((1, 3, 32, 32)		)
           __a					       =						torch.from_numpy(__lowercase		).to(device=__lowercase ,			dtype=torch.floataa		)
           __a					       =						(image > 0).to(device=__lowercase ,			dtype=torch.floataa		)
           __a					       =						{
               """image""": image,
               """mask_image""": mask,
               """generator""": generator,
               """num_inference_steps""": 5,
               """eta""": 0.0,
               """jump_length""": 2,
               """jump_n_sample""": 2,
               """output_type""": """numpy""",
           }
           return inputs





       def   UpperCamelCase_   (				self						:						List[Any]		):

           '''simple docstring'''
           __a					       =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           __a					       =						self.get_dummy_components()
           __a					       =						RePaintPipeline(**__lowercase		)
           __a					       =						sd_pipe.to(__lowercase		)
           sd_pipe.set_progress_bar_config(disable=__lowercase		)

           __a					       =						self.get_dummy_inputs(__lowercase		)
           __a					       =						sd_pipe(**__lowercase		).images
           __a					       =						image[0, -3:, -3:, -1]

           assert image.shape == (1, 32, 32, 3)
           __a					       =						np.array([1.0000, 0.5426, 0.5497, 0.2200, 1.0000, 1.0000, 0.5623, 1.0000, 0.6274]		)

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





       @skip_mps
       def   UpperCamelCase_   (				self						:						int		):

           '''simple docstring'''
           return super().test_save_load_local()





       @unittest.skip("""non-deterministic pipeline"""		)
       def   UpperCamelCase_   (				self						:						Union[str, Any]		):

           '''simple docstring'''
           return super().test_inference_batch_single_identical()





       @skip_mps
       def   UpperCamelCase_   (				self						:						Tuple		):

           '''simple docstring'''
           return super().test_dict_tuple_outputs_equivalent()





       @skip_mps
       def   UpperCamelCase_   (				self						:						str		):

           '''simple docstring'''
           return super().test_save_load_optional_components()







       @skip_mps
       def   UpperCamelCase_   (				self						:						Optional[Any]		):

           '''simple docstring'''
           return super().test_attention_slicing_forward_pass()






@nightly
@require_torch_gpu
class 				SCREAMING_SNAKE_CASE     (							unittest.TestCase   ):





       def   UpperCamelCase_   (				self						:						Any		):

           '''simple docstring'''
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()







       def   UpperCamelCase_   (				self						:						List[Any]		):

           '''simple docstring'''
           __a					       =						load_image(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"""
               """repaint/celeba_hq_256.png"""		)
           __a					       =						load_image(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"""		)
           __a					       =						load_numpy(
               """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"""
               """repaint/celeba_hq_256_result.npy"""		)

           __a					       =						"""google/ddpm-ema-celebahq-256"""
           __a					       =						UNetaDModel.from_pretrained(__lowercase		)
           __a					       =						RePaintScheduler.from_pretrained(__lowercase		)

           __a					       =						RePaintPipeline(unet=__lowercase ,			scheduler=__lowercase		).to(__lowercase		)
           repaint.set_progress_bar_config(disable=__lowercase		)
           repaint.enable_attention_slicing()

           __a					       =						torch.manual_seed(0		)
           __a					       =						repaint(
               __lowercase ,			__lowercase ,			num_inference_steps=250 ,			eta=0.0 ,			jump_length=10 ,			jump_n_sample=10 ,			generator=__lowercase ,			output_type="""np""" ,			)
           __a					       =						output.images[0]

           assert image.shape == (256, 256, 3)
           assert np.abs(expected_image - image		).mean() < 1E-2





import copy

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


lowerCamelCase__     =	logging.get_logger(__name__)

lowerCamelCase__     =	{
    """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 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__   ):





							def __init__(				self						:						List[Any] ,			__lowercase						:						Optional[Any]=256 ,			__lowercase						:						Optional[int]=1024 ,			__lowercase						:						Any=16 ,			__lowercase						:						Optional[Any]=16 ,			__lowercase						:						Optional[Any]=4 ,			__lowercase						:						Optional[int]="gelu" ,			__lowercase						:						List[Any]=1E-6 ,			**__lowercase						:						str ,			):

											'''simple docstring'''
											super().__init__(**__lowercase		)
											__a					       =						hidden_size
											__a					       =						image_size
											__a					       =						patch_size
											__a					       =						image_size // patch_size
											__a					       =						mask_input_channels
											__a					       =						num_point_embeddings
											__a					       =						hidden_act
											__a					       =						layer_norm_eps






class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__   ):





							def __init__(				self						:						int ,			__lowercase						:						Tuple=256 ,			__lowercase						:						List[Any]="relu" ,			__lowercase						:						Any=2048 ,			__lowercase						:						List[Any]=2 ,			__lowercase						:						List[str]=8 ,			__lowercase						:						Any=2 ,			__lowercase						:						List[Any]=3 ,			__lowercase						:						str=3 ,			__lowercase						:						str=256 ,			__lowercase						:						Optional[Any]=1E-6 ,			**__lowercase						:						Tuple ,			):

											'''simple docstring'''
											super().__init__(**__lowercase		)
											__a					       =						hidden_size
											__a					       =						hidden_act
											__a					       =						mlp_dim
											__a					       =						num_hidden_layers
											__a					       =						num_attention_heads
											__a					       =						attention_downsample_rate
											__a					       =						num_multimask_outputs
											__a					       =						iou_head_depth
											__a					       =						iou_head_hidden_dim
											__a					       =						layer_norm_eps






class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__   ):





							def __init__(				self						:						List[Any] ,			__lowercase						:						List[Any]=768 ,			__lowercase						:						List[Any]=256 ,			__lowercase						:						Any=12 ,			__lowercase						:						Tuple=12 ,			__lowercase						:						Union[str, Any]=3 ,			__lowercase						:						Tuple=1024 ,			__lowercase						:						Any=16 ,			__lowercase						:						Dict="gelu" ,			__lowercase						:						List[Any]=1E-06 ,			__lowercase						:						List[str]=0.0 ,			__lowercase						:						Optional[Any]=1E-10 ,			__lowercase						:						Optional[Any]=True ,			__lowercase						:						str=4.0 ,			__lowercase						:						List[Any]=True ,			__lowercase						:						Dict=True ,			__lowercase						:						Optional[Any]=14 ,			__lowercase						:						Optional[int]=[2, 5, 8, 11] ,			__lowercase						:						List[str]=128 ,			__lowercase						:						Optional[int]=None ,			**__lowercase						:						List[Any] ,			):

											'''simple docstring'''
											super().__init__(**__lowercase		)

											__a					       =						hidden_size
											__a					       =						output_channels
											__a					       =						num_hidden_layers
											__a					       =						num_attention_heads
											__a					       =						num_channels
											__a					       =						image_size
											__a					       =						patch_size
											__a					       =						hidden_act
											__a					       =						layer_norm_eps
											__a					       =						attention_dropout
											__a					       =						initializer_range
											__a					       =						qkv_bias
											__a					       =						mlp_ratio
											__a					       =						use_abs_pos
											__a					       =						use_rel_pos
											__a					       =						window_size
											__a					       =						global_attn_indexes
											__a					       =						num_pos_feats
											__a					       =						int(hidden_size * mlp_ratio		) if mlp_dim is None else mlp_dim






class 				SCREAMING_SNAKE_CASE     (							lowerCamelCase__   ):
							__lowerCamelCase		:		Dict      ='sam'
							__lowerCamelCase		:		Any      =True





							def __init__(				self						:						List[str] ,			__lowercase						:						List[str]=None ,			__lowercase						:						Optional[int]=None ,			__lowercase						:						Dict=None ,			__lowercase						:						Optional[int]=0.02 ,			**__lowercase						:						List[str] ,			):

											'''simple docstring'''
											super().__init__(**__lowercase		)
											__a					       =						vision_config if vision_config is not None else {}
											__a					       =						prompt_encoder_config if prompt_encoder_config is not None else {}
											__a					       =						mask_decoder_config if mask_decoder_config is not None else {}

											if isinstance(__lowercase ,			__lowercase		):
															__a					       =						vision_config.to_dict()
											if isinstance(__lowercase ,			__lowercase		):
															__a					       =						prompt_encoder_config.to_dict()
											if isinstance(__lowercase ,			__lowercase		):
															__a					       =						mask_decoder_config.to_dict()

											__a					       =						SamVisionConfig(**__lowercase		)
											__a					       =						SamPromptEncoderConfig(**__lowercase		)
											__a					       =						SamMaskDecoderConfig(**__lowercase		)
											__a					       =						initializer_range







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

											'''simple docstring'''
											__a					       =						copy.deepcopy(self.__dict__		)
											__a					       =						self.vision_config.to_dict()
											__a					       =						self.prompt_encoder_config.to_dict()
											__a					       =						self.mask_decoder_config.to_dict()
											__a					       =						self.__class__.model_type
											return output
