
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      _lowerCAmelCase (     __a ,				unittest.TestCase	):
				_lowercase      		=RePaintPipeline
				_lowercase      		=IMAGE_INPAINTING_PARAMS - {'''width''', '''height''', '''guidance_scale'''}
				_lowercase      		=PipelineTesterMixin.required_optional_params - {
				    '''latents''',
				    '''num_images_per_prompt''',
				    '''callback''',
				    '''callback_steps''',
				}
				_lowercase      		=IMAGE_INPAINTING_BATCH_PARAMS
				_lowercase      		=False




				def 					__a						( self					)    ->				Dict:
							torch.manual_seed(0					)
							torch.manual_seed(0					)
							lowerCAmelCase_      =     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")    ,    )
							lowerCAmelCase_      =     RePaintScheduler()
							lowerCAmelCase_      =     {"unet": unet, "scheduler": scheduler}
							return components




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase=0					)    ->				Optional[int]:
							if str(_UpperCamelCase					).startswith("mps"					):
										lowerCAmelCase_      =     torch.manual_seed(_UpperCamelCase					)
							else:
										lowerCAmelCase_      =     torch.Generator(device=_UpperCamelCase					).manual_seed(_UpperCamelCase					)
							lowerCAmelCase_      =     np.random.RandomState(_UpperCamelCase					).standard_normal((1, 3, 32, 32)					)
							lowerCAmelCase_      =     torch.from_numpy(_UpperCamelCase					).to(device=_UpperCamelCase    ,    dtype=torch.floataa					)
							lowerCAmelCase_      =     (image > 0).to(device=_UpperCamelCase    ,    dtype=torch.floataa					)
							lowerCAmelCase_      =     {
							    "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 					__a						( self					)    ->				Tuple:
							lowerCAmelCase_      =     "cpu"  # ensure determinism for the device-dependent torch.Generator
							lowerCAmelCase_      =     self.get_dummy_components()
							lowerCAmelCase_      =     RePaintPipeline(**_UpperCamelCase					)
							lowerCAmelCase_      =     sd_pipe.to(_UpperCamelCase					)
							sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

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

							assert image.shape == (1, 32, 32, 3)
							lowerCAmelCase_      =     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 					__a						( self					)    ->				int:
							return super().test_save_load_local()




				@unittest.skip("non-deterministic pipeline"					)
				def 					__a						( self					)    ->				Any:
							return super().test_inference_batch_single_identical()




				@skip_mps
				def 					__a						( self					)    ->				List[str]:
							return super().test_dict_tuple_outputs_equivalent()




				@skip_mps
				def 					__a						( self					)    ->				Any:
							return super().test_save_load_optional_components()




				@skip_mps
				def 					__a						( self					)    ->				Optional[Any]:
							return super().test_attention_slicing_forward_pass()





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




				def 					__a						( self					)    ->				List[str]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




				def 					__a						( self					)    ->				Tuple:
							lowerCAmelCase_      =     load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
							    "repaint/celeba_hq_256.png"					)
							lowerCAmelCase_      =     load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"					)
							lowerCAmelCase_      =     load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
							    "repaint/celeba_hq_256_result.npy"					)

							lowerCAmelCase_      =     "google/ddpm-ema-celebahq-256"
							lowerCAmelCase_      =     UNetaDModel.from_pretrained(_UpperCamelCase					)
							lowerCAmelCase_      =     RePaintScheduler.from_pretrained(_UpperCamelCase					)

							lowerCAmelCase_      =     RePaintPipeline(unet=_UpperCamelCase    ,    scheduler=_UpperCamelCase					).to(_UpperCamelCase					)
							repaint.set_progress_bar_config(disable=_UpperCamelCase					)
							repaint.enable_attention_slicing()

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

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





import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest





class      _lowerCAmelCase (     __a	):
				_lowercase      		=(VQDiffusionScheduler,)




				def 					__a						( self    ,    **_UpperCamelCase					)    ->				List[str]:
							lowerCAmelCase_      =     {
							    "num_vec_classes": 4_097,
							    "num_train_timesteps": 100,
							}

							config.update(**_UpperCamelCase					)
							return config




				def 					__a						( self    ,    _UpperCamelCase					)    ->				str:
							lowerCAmelCase_      =     4
							lowerCAmelCase_      =     8
							lowerCAmelCase_      =     8

							lowerCAmelCase_      =     torch.randint(0    ,    _UpperCamelCase    ,    (batch_size, height * width)					)

							return sample




				@property
				def 					__a						( self					)    ->				List[str]:
							assert False




				def 					__a						( self    ,    _UpperCamelCase					)    ->				Tuple:
							def model(_UpperCamelCase    ,    _UpperCamelCase    ,    *_UpperCamelCase					):
										lowerCAmelCase_     ,  lowerCAmelCase_      =     sample.shape
										lowerCAmelCase_      =     torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)					)
										lowerCAmelCase_      =     F.log_softmax(logits.double()    ,    dim=1					).float()
										return return_value

							return model




				def 					__a						( self					)    ->				str:
							for timesteps in [2, 5, 100, 1_000]:
										self.check_over_configs(num_train_timesteps=_UpperCamelCase					)




				def 					__a						( self					)    ->				Dict:
							for num_vec_classes in [5, 100, 1_000, 4_000]:
										self.check_over_configs(num_vec_classes=_UpperCamelCase					)




				def 					__a						( self					)    ->				Tuple:
							for t in [0, 50, 99]:
										self.check_over_forward(time_step=_UpperCamelCase					)




				def 					__a						( self					)    ->				List[str]:
							pass



