import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

lowerCAmelCase			:							List[str]     =[
    "cross_validation.py",
    "gradient_accumulation.py",
    "local_sgd.py",
    "multi_process_metrics.py",
    "memory.py",
    "automatic_gradient_accumulation.py",
    "fsdp_with_peak_mem_tracking.py",
    "deepspeed_with_config_support.py",
    "megatron_lm_gpt_pretraining.py",
]




class 	__snake_case	( unittest.TestCase		):





      '''simple docstring'''






      def 	_SCREAMING_SNAKE_CASE    (					self     :				int    ,  _UpperCamelCase     :				str    ,  _UpperCamelCase     :				bool    ,  _UpperCamelCase     :				str = None    ,  _UpperCamelCase     :				list = None)				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Union[str, Any] 						=						None
          _lowerCamelCase     :		str 						=						os.path.abspath(os.path.join("""examples"""    ,  """by_feature"""))
          _lowerCamelCase     :		Any 						=						os.path.abspath("""examples""")
          for item in os.listdir(_UpperCamelCase):
              if item not in EXCLUDE_EXAMPLES:
                  _lowerCamelCase     :		int 						=						os.path.join(_UpperCamelCase    ,  _UpperCamelCase)
                  if os.path.isfile(_UpperCamelCase) and ".py" in item_path:
                      with self.subTest(
                          tested_script=_UpperCamelCase    ,  feature_script=_UpperCamelCase    ,  tested_section="""main()""" if parser_only else """training_function()"""    ,  ):
                          _lowerCamelCase     :		Dict 						=						compare_against_test(
                              os.path.join(_UpperCamelCase    ,  _UpperCamelCase)    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)
                          _lowerCamelCase     :		Optional[Any] 						=						"""\n""".join(_UpperCamelCase)
                          if special_strings is not None:
                              for string in special_strings:
                                  _lowerCamelCase     :		Union[str, Any] 						=						diff.replace(_UpperCamelCase    ,  """""")
                          self.assertEqual(_UpperCamelCase    ,  """""")



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Any:


          """simple docstring"""


          self.one_complete_example("""complete_nlp_example.py"""    ,  _UpperCamelCase)
          self.one_complete_example("""complete_nlp_example.py"""    ,  _UpperCamelCase)




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Union[str, Any])				->str:


          """simple docstring"""


          _lowerCamelCase     :		Union[str, Any] 						=						os.path.abspath(os.path.join("""examples"""    ,  """cv_example.py"""))
          _lowerCamelCase     :		Dict 						=						[
              """ """ * 16 + """{\n\n""",
              """ """ * 20 + """\"accuracy\": eval_metric[\"accuracy\"],\n\n""",
              """ """ * 20 + """\"f1\": eval_metric[\"f1\"],\n\n""",
              """ """ * 20 + """\"train_loss\": total_loss.item() / len(train_dataloader),\n\n""",
              """ """ * 20 + """\"epoch\": epoch,\n\n""",
              """ """ * 16 + """},\n\n""",
              """ """ * 16 + """step=epoch,\n""",
              """ """ * 12,
              """ """ * 8 + """for step, batch in enumerate(active_dataloader):\n""",
          ]
          self.one_complete_example("""complete_cv_example.py"""    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)
          self.one_complete_example("""complete_cv_example.py"""    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)





@mock.patch.dict(os.environ ,  {'TESTING_MOCKED_DATALOADERS': '1'}		)
class 	__snake_case	( __lowerCAmelCase		):





      '''simple docstring'''






      _snake_case    		=       False



      @classmethod
      def 	_SCREAMING_SNAKE_CASE    (					cls     :				Any)				->List[Any]:


          """simple docstring"""


          super().setUpClass()
          _lowerCamelCase     :		str 						=						tempfile.mkdtemp()
          _lowerCamelCase     :		Optional[Any] 						=						os.path.join(cls._tmpdir    ,  """default_config.yml""")

          write_basic_config(save_location=cls.configPath)
          _lowerCamelCase     :		List[str] 						=						["""accelerate""", """launch""", """--config_file""", cls.configPath]



      @classmethod
      def 	_SCREAMING_SNAKE_CASE    (					cls     :				Dict)				->Optional[Any]:


          """simple docstring"""


          super().tearDownClass()
          shutil.rmtree(cls._tmpdir)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Dict)				->Tuple:


          """simple docstring"""


          _lowerCamelCase     :		int 						=						F"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps epoch
        --output_dir {self.tmpdir}
        """.split()
          run_command(self._launch_args + testargs)
          self.assertTrue(os.path.exists(os.path.join(self.tmpdir    ,  """epoch_0""")))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				int)				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		List[str] 						=						F"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        """.split()
          _lowerCamelCase     :		Any 						=						run_command(self._launch_args + testargs)
          self.assertTrue(os.path.exists(os.path.join(self.tmpdir    ,  """step_2""")))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any])				->Tuple:


          """simple docstring"""


          _lowerCamelCase     :		Optional[int] 						=						F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir    ,  'epoch_0')}
        """.split()
          _lowerCamelCase     :		Any 						=						run_command(self._launch_args + testargs    ,  return_stdout=_UpperCamelCase)
          self.assertNotIn("""epoch 0:"""    ,  _UpperCamelCase)
          self.assertIn("""epoch 1:"""    ,  _UpperCamelCase)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->int:


          """simple docstring"""


          _lowerCamelCase     :		Any 						=						F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir    ,  'step_2')}
        """.split()
          _lowerCamelCase     :		Dict 						=						run_command(self._launch_args + testargs    ,  return_stdout=_UpperCamelCase)
          if torch.cuda.is_available():
              _lowerCamelCase     :		List[str] 						=						torch.cuda.device_count()
          else:
              _lowerCamelCase     :		Dict 						=						1
          if num_processes > 1:
              self.assertNotIn("""epoch 0:"""    ,  _UpperCamelCase)
              self.assertIn("""epoch 1:"""    ,  _UpperCamelCase)
          else:
              self.assertIn("""epoch 0:"""    ,  _UpperCamelCase)
              self.assertIn("""epoch 1:"""    ,  _UpperCamelCase)



      @slow
      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str])				->Optional[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Any 						=						"""
        examples/by_feature/cross_validation.py
        --num_folds 2
        """.split()
          with mock.patch.dict(os.environ    ,  {"""TESTING_MOCKED_DATALOADERS""": """0"""}):
              _lowerCamelCase     :		Any 						=						run_command(self._launch_args + testargs    ,  return_stdout=_UpperCamelCase)
              _lowerCamelCase     :		List[Any] 						=						re.findall("""({.+})"""    ,  _UpperCamelCase)
              _lowerCamelCase     :		List[Any] 						=						[r for r in results if """accuracy""" in r][-1]
              _lowerCamelCase     :		Optional[Any] 						=						ast.literal_eval(_UpperCamelCase)
              self.assertGreaterEqual(results["""accuracy"""]    ,  0.7_5)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[int])				->Optional[int]:


          """simple docstring"""


          _lowerCamelCase     :		Union[str, Any] 						=						["""examples/by_feature/multi_process_metrics.py"""]
          run_command(self._launch_args + testargs)



      @require_trackers
      @mock.patch.dict(os.environ    ,  {"""WANDB_MODE""": """offline"""})
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Union[str, Any])				->Any:


          """simple docstring"""


          with tempfile.TemporaryDirectory() as tmpdir:
              _lowerCamelCase     :		int 						=						F"""
            examples/by_feature/tracking.py
            --with_tracking
            --project_dir {tmpdir}
            """.split()
              run_command(self._launch_args + testargs)
              self.assertTrue(os.path.exists(os.path.join(_UpperCamelCase    ,  """tracking""")))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Dict:


          """simple docstring"""


          _lowerCamelCase     :		Tuple 						=						["""examples/by_feature/gradient_accumulation.py"""]
          run_command(self._launch_args + testargs)




      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any])				->List[str]:


          """simple docstring"""


          _lowerCamelCase     :		List[str] 						=						["""examples/by_feature/local_sgd.py"""]
          run_command(self._launch_args + testargs)





import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin




@skip_mps
class 	__snake_case	( __lowerCAmelCase ,  __lowerCAmelCase ,  unittest.TestCase		):





      '''simple docstring'''






      _snake_case    		=       IFPipeline
      _snake_case    		=       TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
      _snake_case    		=       TEXT_TO_IMAGE_BATCH_PARAMS
      _snake_case    		=       PipelineTesterMixin.required_optional_params - {'latents'}



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Dict)				->Optional[Any]:


          """simple docstring"""


          return self._get_dummy_components()



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Tuple    ,  _UpperCamelCase     :				Optional[int]    ,  _UpperCamelCase     :				List[Any]=0)				->Optional[Any]:


          """simple docstring"""


          if str(_UpperCamelCase).startswith("""mps"""):
              _lowerCamelCase     :		int 						=						torch.manual_seed(_UpperCamelCase)
          else:
              _lowerCamelCase     :		List[Any] 						=						torch.Generator(device=_UpperCamelCase).manual_seed(_UpperCamelCase)

          _lowerCamelCase     :		Dict 						=						{
              """prompt""": """A painting of a squirrel eating a burger""",
              """generator""": generator,
              """num_inference_steps""": 2,
              """output_type""": """numpy""",
          }

          return inputs



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Union[str, Any]:


          """simple docstring"""


          self._test_save_load_optional_components()



      @unittest.skipIf(torch_device != """cuda"""    ,  reason="""float16 requires CUDA""")
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Any)				->str:


          """simple docstring"""


          super().test_save_load_floataa(expected_max_diff=1E-1)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				int)				->Any:


          """simple docstring"""


          self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str])				->Union[str, Any]:


          """simple docstring"""


          self._test_save_load_local()



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Dict:


          """simple docstring"""


          self._test_inference_batch_single_identical(
              expected_max_diff=1E-2    ,  )




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()    ,  reason="""XFormers attention is only available with CUDA and `xformers` installed"""    ,  )
      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any])				->int:


          """simple docstring"""


          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)





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





      '''simple docstring'''






      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str])				->List[str]:


          """simple docstring"""


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



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str])				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Optional[int] 						=						IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""    ,  variant="""fp16"""    ,  torch_dtype=torch.floataa)

          _lowerCamelCase     :		Tuple 						=						IFSuperResolutionPipeline.from_pretrained(
              """DeepFloyd/IF-II-L-v1.0"""    ,  variant="""fp16"""    ,  torch_dtype=torch.floataa    ,  text_encoder=_UpperCamelCase    ,  tokenizer=_UpperCamelCase)

          # pre compute text embeddings and remove T5 to save memory

          pipe_a.text_encoder.to("""cuda""")

          _lowerCamelCase						,							_lowerCamelCase     :		str 						=						pipe_a.encode_prompt("""anime turtle"""    ,  device="""cuda""")

          del pipe_a.tokenizer
          del pipe_a.text_encoder
          gc.collect()

          _lowerCamelCase     :		str 						=						None
          _lowerCamelCase     :		str 						=						None

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

          self._test_if(_UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # img2img

          _lowerCamelCase     :		Optional[Any] 						=						IFImgaImgPipeline(**pipe_a.components)
          _lowerCamelCase     :		Optional[Any] 						=						IFImgaImgSuperResolutionPipeline(**pipe_a.components)

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

          self._test_if_imgaimg(_UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # inpainting

          _lowerCamelCase     :		Any 						=						IFInpaintingPipeline(**pipe_a.components)
          _lowerCamelCase     :		Dict 						=						IFInpaintingSuperResolutionPipeline(**pipe_a.components)

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

          self._test_if_inpainting(_UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase    ,  _UpperCamelCase)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any]    ,  _UpperCamelCase     :				List[Any]    ,  _UpperCamelCase     :				Dict    ,  _UpperCamelCase     :				Union[str, Any]    ,  _UpperCamelCase     :				str)				->Tuple:


          """simple docstring"""


          _start_torch_memory_measurement()

          _lowerCamelCase     :		Optional[int] 						=						torch.Generator(device="""cpu""").manual_seed(0)
          _lowerCamelCase     :		Optional[Any] 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  num_inference_steps=2    ,  generator=_UpperCamelCase    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		Optional[int] 						=						output.images[0]

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

          _lowerCamelCase     :		Dict 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 13 * 10**9

          _lowerCamelCase     :		Dict 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)

          # pipeline 2

          _start_torch_memory_measurement()

          _lowerCamelCase     :		Tuple 						=						torch.Generator(device="""cpu""").manual_seed(0)

          _lowerCamelCase     :		int 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(0)).to(_UpperCamelCase)

          _lowerCamelCase     :		str 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  image=_UpperCamelCase    ,  generator=_UpperCamelCase    ,  num_inference_steps=2    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		Any 						=						output.images[0]

          assert image.shape == (256, 256, 3)

          _lowerCamelCase     :		Tuple 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          _lowerCamelCase     :		int 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any]    ,  _UpperCamelCase     :				List[Any]    ,  _UpperCamelCase     :				List[str]    ,  _UpperCamelCase     :				str    ,  _UpperCamelCase     :				List[Any])				->Any:


          """simple docstring"""


          _start_torch_memory_measurement()

          _lowerCamelCase     :		int 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(0)).to(_UpperCamelCase)

          _lowerCamelCase     :		Union[str, Any] 						=						torch.Generator(device="""cpu""").manual_seed(0)

          _lowerCamelCase     :		Dict 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  image=_UpperCamelCase    ,  num_inference_steps=2    ,  generator=_UpperCamelCase    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		Union[str, Any] 						=						output.images[0]

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

          _lowerCamelCase     :		Optional[Any] 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 10 * 10**9

          _lowerCamelCase     :		List[Any] 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)

          # pipeline 2

          _start_torch_memory_measurement()

          _lowerCamelCase     :		Tuple 						=						torch.Generator(device="""cpu""").manual_seed(0)

          _lowerCamelCase     :		List[str] 						=						floats_tensor((1, 3, 256, 256)    ,  rng=random.Random(0)).to(_UpperCamelCase)
          _lowerCamelCase     :		List[Any] 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(0)).to(_UpperCamelCase)

          _lowerCamelCase     :		Optional[Any] 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  image=_UpperCamelCase    ,  original_image=_UpperCamelCase    ,  generator=_UpperCamelCase    ,  num_inference_steps=2    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		List[Any] 						=						output.images[0]

          assert image.shape == (256, 256, 3)

          _lowerCamelCase     :		str 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          _lowerCamelCase     :		int 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Any    ,  _UpperCamelCase     :				Optional[int]    ,  _UpperCamelCase     :				List[str]    ,  _UpperCamelCase     :				Optional[int]    ,  _UpperCamelCase     :				Tuple)				->Optional[int]:


          """simple docstring"""


          _start_torch_memory_measurement()

          _lowerCamelCase     :		Tuple 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(0)).to(_UpperCamelCase)
          _lowerCamelCase     :		Any 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(1)).to(_UpperCamelCase)

          _lowerCamelCase     :		int 						=						torch.Generator(device="""cpu""").manual_seed(0)
          _lowerCamelCase     :		Any 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  image=_UpperCamelCase    ,  mask_image=_UpperCamelCase    ,  num_inference_steps=2    ,  generator=_UpperCamelCase    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		Any 						=						output.images[0]

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

          _lowerCamelCase     :		List[Any] 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 10 * 10**9

          _lowerCamelCase     :		str 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)

          # pipeline 2

          _start_torch_memory_measurement()

          _lowerCamelCase     :		Tuple 						=						torch.Generator(device="""cpu""").manual_seed(0)

          _lowerCamelCase     :		Union[str, Any] 						=						floats_tensor((1, 3, 64, 64)    ,  rng=random.Random(0)).to(_UpperCamelCase)
          _lowerCamelCase     :		Union[str, Any] 						=						floats_tensor((1, 3, 256, 256)    ,  rng=random.Random(0)).to(_UpperCamelCase)
          _lowerCamelCase     :		Optional[int] 						=						floats_tensor((1, 3, 256, 256)    ,  rng=random.Random(1)).to(_UpperCamelCase)

          _lowerCamelCase     :		List[str] 						=						pipe_a(
              prompt_embeds=_UpperCamelCase    ,  negative_prompt_embeds=_UpperCamelCase    ,  image=_UpperCamelCase    ,  mask_image=_UpperCamelCase    ,  original_image=_UpperCamelCase    ,  generator=_UpperCamelCase    ,  num_inference_steps=2    ,  output_type="""np"""    ,  )

          _lowerCamelCase     :		Optional[Any] 						=						output.images[0]

          assert image.shape == (256, 256, 3)

          _lowerCamelCase     :		Optional[Any] 						=						torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          _lowerCamelCase     :		int 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy""")
          assert_mean_pixel_difference(_UpperCamelCase    ,  _UpperCamelCase)







def 						A__  (					):







    '''simple docstring'''
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()




