





import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment






def lowerCamelCase__    (							)      ->				Tuple:
   __snake_case							       =   torch.nn.Linear(2       ,      4       )
   __snake_case							       =   torch.optim.AdamW(model.parameters()       ,      lr=1.0       )
   __snake_case							       =   torch.optim.lr_scheduler.OneCycleLR(snake_case_       ,      max_lr=0.01       ,      steps_per_epoch=2       ,      epochs=1       )
   __snake_case							       =   DataLoader(TensorDataset(torch.tensor([1, 2, 3]       )       )       )
   __snake_case							       =   DataLoader(TensorDataset(torch.tensor([4, 5, 6]       )       )       )

   return model, optimizer, scheduler, train_dl, valid_dl






def lowerCamelCase__    (							snake_case_	:  Any       )      ->				Any:
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()






def lowerCamelCase__    (							snake_case_	:  Union[str, Any]       )      ->				Union[str, Any]:
   __snake_case							       =   torch.nn.Linear(*tuple(model.weight.T.shape       )       ).state_dict()
   model.load_state_dict(snake_case_       )







class 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase  ):


      @require_cuda
      def    a (self       :      Tuple      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(a__      ):
            __snake_case							       =   Accelerator(cpu=a__      )


      def    a (self       :      Tuple      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case							       =   GradientState()
         assert state.num_steps == 1
         __snake_case							       =   4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         __snake_case							       =   False
         assert state.sync_gradients is False
         GradientState._reset_state()


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()

         (
             (
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	(
             __snake_case
         )       ,	
         )							       =   accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         self.assertTrue(prepared_model in accelerator._models      )
         self.assertTrue(prepared_optimizer in accelerator._optimizers      )
         self.assertTrue(prepared_scheduler in accelerator._schedulers      )
         self.assertTrue(prepared_train_dl in accelerator._dataloaders      )
         self.assertTrue(prepared_valid_dl in accelerator._dataloaders      )


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )
         accelerator.free_memory()

         self.assertTrue(len(accelerator._models      ) == 0      )
         self.assertTrue(len(accelerator._optimizers      ) == 0      )
         self.assertTrue(len(accelerator._schedulers      ) == 0      )
         self.assertTrue(len(accelerator._dataloaders      ) == 0      )


      def    a (self       :      str      ):





         """simple docstring"""



         PartialState._reset_state()

         # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
         def noop(*a__       :      List[Any]				,			**a__       :      List[Any]      ):
            pass

         with patch('''torch.cuda.set_device'''				,			a__      ), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64'''      ):
            __snake_case							       =   Accelerator()
            self.assertEqual(str(accelerator.state.device      )				,			'''cuda:64'''      )


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





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         __snake_case							       =   get_signature(a__      )

         with tempfile.TemporaryDirectory() as tmpdirname:
            accelerator.save_state(a__      )

            # make sure random weights don't match
            load_random_weights(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) > 1E-3      )

            # make sure loaded weights match
            accelerator.load_state(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) < 1E-3      )


      def    a (self       :      int      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         accelerator.prepare(a__				,			a__				,			a__				,			a__				,			a__      )

         __snake_case							       =   get_signature(a__      )

         # saving hook
         def save_config(a__       :      Any				,			a__       :      Any				,			a__       :      Optional[int]      ):
            __snake_case							       =   {'''class_name''': models[0].__class__.__name__}

            with open(os.path.join(a__				,			'''data.json'''      )				,			'''w'''      ) as f:
               json.dump(a__				,			a__      )

        # loading hook
         def load_config(a__       :      str				,			a__       :      Optional[Any]      ):
            with open(os.path.join(a__				,			'''data.json'''      )				,			'''r'''      ) as f:
               __snake_case							       =   json.load(a__      )

            __snake_case							       =   config['''class_name''']

         __snake_case							       =   accelerator.register_save_state_pre_hook(a__      )
         __snake_case							       =   accelerator.register_load_state_pre_hook(a__      )

         with tempfile.TemporaryDirectory() as tmpdirname:
            accelerator.save_state(a__      )

            # make sure random weights don't match with hooks
            load_random_weights(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) > 1E-3      )

            # random class name to verify correct one is loaded
            __snake_case							       =   '''random'''

            # make sure loaded weights match with hooks
            accelerator.load_state(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) < 1E-3      )

            # mode.class_name is loaded from config
            self.assertTrue(model.class_name == model.__class__.__name__      )

         # remove hooks
         save_hook.remove()
         load_hook.remove()

         with tempfile.TemporaryDirectory() as tmpdirname:
            accelerator.save_state(a__      )

            # make sure random weights don't match with hooks removed
            load_random_weights(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) > 1E-3      )

            # random class name to verify correct one is loaded
            __snake_case							       =   '''random'''

            # make sure loaded weights match with hooks removed
            accelerator.load_state(a__      )
            self.assertTrue(abs(model_signature - get_signature(a__      )      ) < 1E-3      )

            # mode.class_name is NOT loaded from config
            self.assertTrue(model.class_name != model.__class__.__name__      )


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         __snake_case							       =   None

         # This should work
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   accelerator.prepare(
             a__				,			a__				,			a__				,			a__				,			a__				,			a__      )
         self.assertTrue(dummy_obj is None      )


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





         """simple docstring"""



         __snake_case							       =   Accelerator()
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   create_components()
         __snake_case							       =   [1, 2, 3]

         # This should work
         __snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case       ,	__snake_case							       =   accelerator.prepare(
             a__				,			a__				,			a__				,			a__				,			a__				,			a__      )
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Dummy object should have `_is_accelerate_prepared` set to `True`'''				,			)
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Model is missing `_is_accelerator_prepared` or is set to `False`'''				,			)
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Optimizer is missing `_is_accelerator_prepared` or is set to `False`'''				,			)
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Scheduler is missing `_is_accelerator_prepared` or is set to `False`'''				,			)
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''				,			)
         self.assertEqual(
             getattr(a__				,			'''_is_accelerate_prepared'''				,			a__      )				,			a__				,			'''Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''				,			)


      @slow
      @require_bnb
      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

         __snake_case							       =   AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''				,			load_in_abit=a__				,			device_map={'''''': 0}				,			)
         __snake_case							       =   Accelerator()

         # This should work
         __snake_case							       =   accelerator.prepare(a__      )


      @slow
      @require_bnb
      def    a (self       :      Any      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

         __snake_case							       =   Accelerator()

         with init_empty_weights():
            __snake_case							       =   AutoModelForCausalLM.from_pretrained(
                '''EleutherAI/gpt-neo-125m'''				,			)
            model.tie_weights()
            __snake_case							       =   infer_auto_device_map(a__      )
            __snake_case							       =   '''cpu'''

         __snake_case							       =   AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''				,			device_map=a__				,			load_in_abit=a__				,			llm_inta_enable_fpaa_cpu_offload=a__      )

         # This should not work and get value error
         with self.assertRaises(a__      ):
            __snake_case							       =   accelerator.prepare(a__      )


      @slow
      @require_bnb
      @require_multi_gpu
      def    a (self       :      int      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

         __snake_case							       =   {'''distributed_type''': DistributedType.MULTI_GPU}

         with init_empty_weights():
            __snake_case							       =   AutoModelForCausalLM.from_pretrained(
                '''EleutherAI/gpt-neo-125m'''				,			)
            model.tie_weights()
            __snake_case							       =   infer_auto_device_map(a__      )
            __snake_case							       =   1

         __snake_case							       =   AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''				,			load_in_abit=a__				,			device_map=a__				,			)
         __snake_case							       =   Accelerator()

         # This should not work and get value error
         with self.assertRaises(a__      ):
            __snake_case							       =   accelerator.prepare(a__      )

         PartialState._reset_state()


      @slow
      @require_bnb
      @require_multi_gpu
      def    a (self       :      int      ):





         """simple docstring"""



         from transformers import AutoModelForCausalLM

         with init_empty_weights():
            __snake_case							       =   AutoModelForCausalLM.from_pretrained(
                '''EleutherAI/gpt-neo-125m'''				,			)
            __snake_case							       =   infer_auto_device_map(a__      )
            __snake_case							       =   1

         __snake_case							       =   AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''				,			load_in_abit=a__				,			device_map=a__				,			)
         __snake_case							       =   Accelerator()

         # This should work
         __snake_case							       =   accelerator.prepare(a__      )


      @require_cuda
      def    a (self       :      str      ):





         """simple docstring"""



         __snake_case							       =   torch.nn.Linear(10				,			10      )
         __snake_case							       =   torch.optim.SGD(model.parameters()				,			lr=0.0_1      )
         __snake_case							       =   Accelerator(cpu=a__      )
         __snake_case							       =   accelerator.prepare(a__      )





"""simple docstring"""







import gc
import random
import unittest

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

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 						_A  (  _a    ,_a    ,_a    ,unittest.TestCase  ):







    """simple docstring"""




    UpperCAmelCase       :							str	=       StableDiffusionInpaintPipeline
    UpperCAmelCase       :							Optional[Any]	=       TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
    UpperCAmelCase       :							Any	=       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
    UpperCAmelCase       :							Union[str, Any]	=       frozenset(
        []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    UpperCAmelCase       :							int	=       frozenset([]  )






    def 	__snake_case							(					self				:    Dict):
      torch.manual_seed(0)
      a						:		List[Any]      							=			UNetaDConditionModel(
          block_out_channels=(32, 64)				, layers_per_block=2				, sample_size=32				, in_channels=9				, 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				, )
      a						:		Tuple      							=			PNDMScheduler(skip_prk_steps=__UpperCAmelCase)
      torch.manual_seed(0)
      a						:		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=128				, )
      torch.manual_seed(0)
      a						:		List[str]      							=			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=1000				, hidden_act="gelu"				, projection_dim=512				, )
      a						:		Any      							=			CLIPTextModel(__UpperCAmelCase)
      a						:		List[str]      							=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

      a						:		Optional[Any]      							=			{
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      return components






    def 	__snake_case							(					self				:    Union[str, Any]				, __UpperCAmelCase				:    Any				, __UpperCAmelCase				:    Union[str, Any]=0):
      # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched
      a						:		Dict      							=			floats_tensor((1, 3, 32, 32)				, rng=random.Random(__UpperCAmelCase)).to(__UpperCAmelCase)
      a						:		List[str]      							=			image.cpu().permute(0				, 2				, 3				, 1)[0]
      a						:		Union[str, Any]      							=			Image.fromarray(np.uinta(__UpperCAmelCase)).convert("RGB").resize((64, 64))
      a						:		Dict      							=			Image.fromarray(np.uinta(image + 4)).convert("RGB").resize((64, 64))
      if str(__UpperCAmelCase).startswith("mps"):
        a						:		Tuple      							=			torch.manual_seed(__UpperCAmelCase)
      else:
        a						:		Tuple      							=			torch.Generator(device=__UpperCAmelCase).manual_seed(__UpperCAmelCase)
      a						:		Optional[Any]      							=			{
          "prompt": "A painting of a squirrel eating a burger",
          "image": init_image,
          "mask_image": mask_image,
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
          "output_type": "numpy",
      }
      return inputs






    def 	__snake_case							(					self				:    List[str]):
      a						:		Union[str, Any]      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
      a						:		Tuple      							=			self.get_dummy_components()
      a						:		Optional[int]      							=			StableDiffusionInpaintPipeline(**__UpperCAmelCase)
      a						:		int      							=			sd_pipe.to(__UpperCAmelCase)
      sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase)

      a						:		Any      							=			self.get_dummy_inputs(__UpperCAmelCase)
      a						:		Optional[int]      							=			sd_pipe(**__UpperCAmelCase).images
      a						:		Optional[int]      							=			image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      a						:		int      							=			np.array([0.4_727, 0.5_735, 0.3_941, 0.5_446, 0.5_926, 0.4_394, 0.5_062, 0.4_654, 0.4_476])

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






    def 	__snake_case							(					self				:    str):
      super().test_inference_batch_single_identical(expected_max_diff=3e-3)






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







    """simple docstring"""




    def 	__snake_case							(					self				:    Union[str, Any]):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()






    def 	__snake_case							(					self				:    Dict):
      a						:		str      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-inpaint/init_image.png")
      a						:		str      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")
      a						:		List[Any]      							=			load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
          "/yellow_cat_sitting_on_a_park_bench.npy")

      a						:		Tuple      							=			"stabilityai/stable-diffusion-2-inpainting"
      a						:		Optional[Any]      							=			StableDiffusionInpaintPipeline.from_pretrained(__UpperCAmelCase				, safety_checker=__UpperCAmelCase)
      pipe.to(__UpperCAmelCase)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase)
      pipe.enable_attention_slicing()

      a						:		Any      							=			"Face of a yellow cat, high resolution, sitting on a park bench"

      a						:		str      							=			torch.manual_seed(0)
      a						:		Union[str, Any]      							=			pipe(
          prompt=__UpperCAmelCase				, image=__UpperCAmelCase				, mask_image=__UpperCAmelCase				, generator=__UpperCAmelCase				, output_type="np"				, )
      a						:		List[str]      							=			output.images[0]

      assert image.shape == (512, 512, 3)
      assert np.abs(expected_image - image).max() < 9e-3






    def 	__snake_case							(					self				:    Any):
      a						:		str      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-inpaint/init_image.png")
      a						:		Optional[Any]      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")
      a						:		str      							=			load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
          "/yellow_cat_sitting_on_a_park_bench_fp16.npy")

      a						:		Optional[Any]      							=			"stabilityai/stable-diffusion-2-inpainting"
      a						:		Any      							=			StableDiffusionInpaintPipeline.from_pretrained(
          __UpperCAmelCase				, torch_dtype=torch.floataa				, safety_checker=__UpperCAmelCase				, )
      pipe.to(__UpperCAmelCase)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase)
      pipe.enable_attention_slicing()

      a						:		Optional[int]      							=			"Face of a yellow cat, high resolution, sitting on a park bench"

      a						:		Dict      							=			torch.manual_seed(0)
      a						:		List[Any]      							=			pipe(
          prompt=__UpperCAmelCase				, image=__UpperCAmelCase				, mask_image=__UpperCAmelCase				, generator=__UpperCAmelCase				, output_type="np"				, )
      a						:		Optional[Any]      							=			output.images[0]

      assert image.shape == (512, 512, 3)
      assert np.abs(expected_image - image).max() < 5e-1






    def 	__snake_case							(					self				:    int):
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      a						:		Tuple      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/sd2-inpaint/init_image.png")
      a						:		Optional[int]      							=			load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png")

      a						:		Optional[Any]      							=			"stabilityai/stable-diffusion-2-inpainting"
      a						:		Optional[int]      							=			PNDMScheduler.from_pretrained(__UpperCAmelCase				, subfolder="scheduler")
      a						:		int      							=			StableDiffusionInpaintPipeline.from_pretrained(
          __UpperCAmelCase				, safety_checker=__UpperCAmelCase				, scheduler=__UpperCAmelCase				, torch_dtype=torch.floataa				, )
      pipe.to(__UpperCAmelCase)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase)
      pipe.enable_attention_slicing(1)
      pipe.enable_sequential_cpu_offload()

      a						:		Optional[int]      							=			"Face of a yellow cat, high resolution, sitting on a park bench"

      a						:		Optional[int]      							=			torch.manual_seed(0)
      a						:		str      							=			pipe(
          prompt=__UpperCAmelCase				, image=__UpperCAmelCase				, mask_image=__UpperCAmelCase				, generator=__UpperCAmelCase				, num_inference_steps=2				, output_type="np"				, )

      a						:		int      							=			torch.cuda.max_memory_allocated()
      # make sure that less than 2.65 GB is allocated
      assert mem_bytes < 2.65 * 10**9






