

'''simple docstring'''



import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()







def    a__       (  a__      ,    a__      ,    a__						):




     """simple docstring"""


     return params[F'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]







def    a__       (  a__      ,    a__      ,    a__      ,    a__="attention"						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      __SCREAMING_SNAKE_CASE     =      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]						)
     __SCREAMING_SNAKE_CASE     =      k_tmp.reshape(k_tmp.shape[0]      ,    k_tmp.shape[1] * k_tmp.shape[2]						)
     __SCREAMING_SNAKE_CASE     =      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/out/kernel'][:, i, :, :]						)
     __SCREAMING_SNAKE_CASE     =      o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]      ,    o_tmp.shape[2]						)
     __SCREAMING_SNAKE_CASE     =      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]						)
     __SCREAMING_SNAKE_CASE     =      q_tmp.reshape(q_tmp.shape[0]      ,    q_tmp.shape[1] * q_tmp.shape[2]						)
     __SCREAMING_SNAKE_CASE     =      np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]						)
     __SCREAMING_SNAKE_CASE     =      v_tmp.reshape(v_tmp.shape[0]      ,    v_tmp.shape[1] * v_tmp.shape[2]						)
     return k, o, q, v







def    a__       (  a__      ,    a__      ,    a__      ,    a__=False						):




     """simple docstring"""


     if split_mlp_wi:
          __SCREAMING_SNAKE_CASE     =      params[F'{prefix}/{prefix}/mlp/wi_0/kernel'][:, i, :]
          __SCREAMING_SNAKE_CASE     =      params[F'{prefix}/{prefix}/mlp/wi_1/kernel'][:, i, :]
          __SCREAMING_SNAKE_CASE     =      (wi_a, wi_a)
     else:
          __SCREAMING_SNAKE_CASE     =      params[F'{prefix}/{prefix}/mlp/wi/kernel'][:, i, :]

     __SCREAMING_SNAKE_CASE     =      params[F'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
     return wi, wo







def    a__       (  a__      ,    a__      ,    a__      ,    a__						):




     """simple docstring"""


     return params[F'{prefix}/{prefix}/{layer_name}/scale'][:, i]







def    a__       (  a__      ,    *, a__      ,    a__      ,    a__ = False						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      traverse_util.flatten_dict(variables["""target"""]						)
     __SCREAMING_SNAKE_CASE     =      {"""/""".join(a__						): v for k, v in old.items()}

     # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
     __SCREAMING_SNAKE_CASE     =      """encoder/encoder/mlp/wi_0/kernel""" in old
     print("""Split MLP:"""      ,    a__						)

     __SCREAMING_SNAKE_CASE     =      collections.OrderedDict()

     # Shared embeddings.
     __SCREAMING_SNAKE_CASE     =      old["""token_embedder/embedding"""]

     # Encoder.
     for i in range(a__						):
          # Block i, layer 0 (Self Attention).
          __SCREAMING_SNAKE_CASE     =      tax_layer_norm_lookup(a__      ,    a__      ,    """encoder"""      ,    """pre_attention_layer_norm"""						)
          __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      tax_attention_lookup(a__      ,    a__      ,    """encoder"""      ,    """attention"""						)
          __SCREAMING_SNAKE_CASE     =      layer_norm
          __SCREAMING_SNAKE_CASE     =      k.T
          __SCREAMING_SNAKE_CASE     =      o.T
          __SCREAMING_SNAKE_CASE     =      q.T
          __SCREAMING_SNAKE_CASE     =      v.T

          # Block i, layer 1 (MLP).
          __SCREAMING_SNAKE_CASE     =      tax_layer_norm_lookup(a__      ,    a__      ,    """encoder"""      ,    """pre_mlp_layer_norm"""						)
          __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      tax_mlp_lookup(a__      ,    a__      ,    """encoder"""      ,    a__						)
          __SCREAMING_SNAKE_CASE     =      layer_norm
          if split_mlp_wi:
               __SCREAMING_SNAKE_CASE     =      wi[0].T
               __SCREAMING_SNAKE_CASE     =      wi[1].T
          else:
               __SCREAMING_SNAKE_CASE     =      wi.T
          __SCREAMING_SNAKE_CASE     =      wo.T
          if scalable_attention:
               # convert the rel_embedding of each layer
               __SCREAMING_SNAKE_CASE     =      tax_relpos_bias_lookup(
                   a__      ,    a__      ,    """encoder"""						).T

     __SCREAMING_SNAKE_CASE     =      old["""encoder/encoder_norm/scale"""]

     if not scalable_attention:
          __SCREAMING_SNAKE_CASE     =      tax_relpos_bias_lookup(
              a__      ,    0      ,    """encoder"""						).T
          __SCREAMING_SNAKE_CASE     =      tax_relpos_bias_lookup(
              a__      ,    0      ,    """decoder"""						).T

     if not is_encoder_only:
          # Decoder.
          for i in range(a__						):
               # Block i, layer 0 (Self Attention).
               __SCREAMING_SNAKE_CASE     =      tax_layer_norm_lookup(a__      ,    a__      ,    """decoder"""      ,    """pre_self_attention_layer_norm"""						)
               __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      tax_attention_lookup(a__      ,    a__      ,    """decoder"""      ,    """self_attention"""						)
               __SCREAMING_SNAKE_CASE     =      layer_norm
               __SCREAMING_SNAKE_CASE     =      k.T
               __SCREAMING_SNAKE_CASE     =      o.T
               __SCREAMING_SNAKE_CASE     =      q.T
               __SCREAMING_SNAKE_CASE     =      v.T

               # Block i, layer 1 (Cross Attention).
               __SCREAMING_SNAKE_CASE     =      tax_layer_norm_lookup(a__      ,    a__      ,    """decoder"""      ,    """pre_cross_attention_layer_norm"""						)
               __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      tax_attention_lookup(a__      ,    a__      ,    """decoder"""      ,    """encoder_decoder_attention"""						)
               __SCREAMING_SNAKE_CASE     =      layer_norm
               __SCREAMING_SNAKE_CASE     =      k.T
               __SCREAMING_SNAKE_CASE     =      o.T
               __SCREAMING_SNAKE_CASE     =      q.T
               __SCREAMING_SNAKE_CASE     =      v.T

               # Block i, layer 2 (MLP).
               __SCREAMING_SNAKE_CASE     =      tax_layer_norm_lookup(a__      ,    a__      ,    """decoder"""      ,    """pre_mlp_layer_norm"""						)
               __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      tax_mlp_lookup(a__      ,    a__      ,    """decoder"""      ,    a__						)
               __SCREAMING_SNAKE_CASE     =      layer_norm
               if split_mlp_wi:
                    __SCREAMING_SNAKE_CASE     =      wi[0].T
                    __SCREAMING_SNAKE_CASE     =      wi[1].T
               else:
                    __SCREAMING_SNAKE_CASE     =      wi.T
               __SCREAMING_SNAKE_CASE     =      wo.T

               if scalable_attention:
                    # convert the rel_embedding of each layer
                    __SCREAMING_SNAKE_CASE     =      tax_relpos_bias_lookup(a__      ,    a__      ,    """decoder"""						).T

          __SCREAMING_SNAKE_CASE     =      old["""decoder/decoder_norm/scale"""]

          # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
          if "decoder/logits_dense/kernel" in old:
               __SCREAMING_SNAKE_CASE     =      old["""decoder/logits_dense/kernel"""].T

     return new







def    a__       (  a__      ,    a__						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      collections.OrderedDict([(k, torch.from_numpy(v.copy()						)) for (k, v) in converted_params.items()]						)

     # Add what is missing.
     if "encoder.embed_tokens.weight" not in state_dict:
          __SCREAMING_SNAKE_CASE     =      state_dict["""shared.weight"""]

     if not is_encoder_only:
          if "decoder.embed_tokens.weight" not in state_dict:
               __SCREAMING_SNAKE_CASE     =      state_dict["""shared.weight"""]

          if "lm_head.weight" not in state_dict:  # For old 1.0 models.
               print("""Using shared word embeddings as lm_head."""						)
               __SCREAMING_SNAKE_CASE     =      state_dict["""shared.weight"""]

     return state_dict







def    a__       (  a__      ,    a__      ,    a__      ,    a__      ,    a__						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      checkpoints.load_tax_checkpoint(a__						)
     __SCREAMING_SNAKE_CASE     =      convert_tax_to_pytorch(
         a__      ,    num_layers=config.num_layers      ,    is_encoder_only=a__      ,    scalable_attention=a__						)
     __SCREAMING_SNAKE_CASE     =      make_state_dict(a__      ,    a__						)
     model.load_state_dict(a__      ,    strict=a__						)







def    a__       (  a__      ,    a__      ,    a__      ,    a__ = False      ,    a__ = False      ,    ):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      MTaConfig.from_json_file(a__						)
     print(F'Building PyTorch model from configuration: {config}'						)
     # Non-v1.1 checkpoints could also use T5Model, but this works for all.
     # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
     if is_encoder_only:
          __SCREAMING_SNAKE_CASE     =      UMTaEncoderModel(a__						)
     else:
          __SCREAMING_SNAKE_CASE     =      UMTaForConditionalGeneration(a__						)

     # Load weights from tf checkpoint
     load_tax_weights_in_ta(a__      ,    a__      ,    a__      ,    a__      ,    a__						)

     # Save pytorch-model
     print(F'Save PyTorch model to {pytorch_dump_path}'						)
     model.save_pretrained(a__						)

     # Verify that we can load the checkpoint.
     model.from_pretrained(a__						)
     print("""Done"""						)


if __name__ == "__main__":
  UpperCAmelCase :       Dict			   =     argparse.ArgumentParser(description='Converts a native T5X checkpoint into a PyTorch checkpoint.')
  # Required parameters
  parser.add_argument(
      '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path to the T5X checkpoint.'
  )
  parser.add_argument(
      '--config_file',
      default=None,
      type=str,
      required=True,
      help='The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.',
  )
  parser.add_argument(
      '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
  )
  parser.add_argument(
      '--is_encoder_only', action='store_true', help='Check if the model is encoder-decoder model', default=False
  )
  parser.add_argument(
      '--scalable_attention',
      action='store_true',
      help='Whether the model uses scaled attention (umt5 model)',
      default=False,
  )
  UpperCAmelCase :       Tuple			   =     parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path,
      args.config_file,
      args.pytorch_dump_path,
      args.is_encoder_only,
      args.scalable_attention,
  )









'''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 						lowerCAmelCase__   (      a ,					a ,					a ,					unittest.TestCase					):



 """simple docstring"""



 lowerCAmelCase__  					=      StableDiffusionInpaintPipeline
 lowerCAmelCase__  					=      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
 lowerCAmelCase__  					=      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
 lowerCAmelCase__  					=      frozenset(
     []					)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 lowerCAmelCase__  					=      frozenset([]					)
 def       UpperCAmelCase__		(		self					:       str			)  -> List[Any]:


      """simple docstring"""




      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      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=__SCREAMING_SNAKE_CASE					,							)
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler(skip_prk_steps=__SCREAMING_SNAKE_CASE			)
      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      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			)
      __SCREAMING_SNAKE_CASE     =      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=1_000					,							hidden_act="""gelu"""					,							projection_dim=512					,							)
      __SCREAMING_SNAKE_CASE     =      CLIPTextModel(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

      __SCREAMING_SNAKE_CASE     =      {
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components
 def       UpperCAmelCase__		(		self					:       Any					,							__SCREAMING_SNAKE_CASE					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       List[Any]=0			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      floats_tensor((1, 3, 32, 32)					,							rng=random.Random(__SCREAMING_SNAKE_CASE			)			).to(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      image.cpu().permute(0					,							2					,							3					,							1			)[0]
      __SCREAMING_SNAKE_CASE     =      Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE			)			).convert("""RGB"""			).resize((64, 64)			)
      __SCREAMING_SNAKE_CASE     =      Image.fromarray(np.uinta(image + 4			)			).convert("""RGB"""			).resize((64, 64)			)
      if str(__SCREAMING_SNAKE_CASE			).startswith("""mps"""			):
           __SCREAMING_SNAKE_CASE     =      torch.manual_seed(__SCREAMING_SNAKE_CASE			)
      else:
           __SCREAMING_SNAKE_CASE     =      torch.Generator(device=__SCREAMING_SNAKE_CASE			).manual_seed(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      {
          """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       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
      __SCREAMING_SNAKE_CASE     =      self.get_dummy_components()
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline(**__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      sd_pipe.to(__SCREAMING_SNAKE_CASE			)
      sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)

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

      assert image.shape == (1, 64, 64, 3)
      __SCREAMING_SNAKE_CASE     =      np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]			)

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


 def       UpperCAmelCase__		(		self					:       Tuple			)  -> str:


      """simple docstring"""




      super().test_inference_batch_single_identical(expected_max_diff=3E-3			)







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



 """simple docstring"""



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


      """simple docstring"""




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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
          """/yellow_cat_sitting_on_a_park_bench.npy"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(__SCREAMING_SNAKE_CASE					,							safety_checker=__SCREAMING_SNAKE_CASE			)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing()

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

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							output_type="""np"""					,							)
      __SCREAMING_SNAKE_CASE     =      output.images[0]

      assert image.shape == (512, 512, 3)
      assert np.abs(expected_image - image			).max() < 9E-3
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"""
          """/yellow_cat_sitting_on_a_park_bench_fp16.npy"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(
          __SCREAMING_SNAKE_CASE					,							torch_dtype=torch.floataa					,							safety_checker=__SCREAMING_SNAKE_CASE					,							)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing()

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

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							output_type="""np"""					,							)
      __SCREAMING_SNAKE_CASE     =      output.images[0]

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


 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Any:


      """simple docstring"""




      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/sd2-inpaint/init_image.png"""			)
      __SCREAMING_SNAKE_CASE     =      load_image(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""			)

      __SCREAMING_SNAKE_CASE     =      """stabilityai/stable-diffusion-2-inpainting"""
      __SCREAMING_SNAKE_CASE     =      PNDMScheduler.from_pretrained(__SCREAMING_SNAKE_CASE					,							subfolder="""scheduler"""			)
      __SCREAMING_SNAKE_CASE     =      StableDiffusionInpaintPipeline.from_pretrained(
          __SCREAMING_SNAKE_CASE					,							safety_checker=__SCREAMING_SNAKE_CASE					,							scheduler=__SCREAMING_SNAKE_CASE					,							torch_dtype=torch.floataa					,							)
      pipe.to(__SCREAMING_SNAKE_CASE			)
      pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE			)
      pipe.enable_attention_slicing(1			)
      pipe.enable_sequential_cpu_offload()

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

      __SCREAMING_SNAKE_CASE     =      torch.manual_seed(0			)
      __SCREAMING_SNAKE_CASE     =      pipe(
          prompt=__SCREAMING_SNAKE_CASE					,							image=__SCREAMING_SNAKE_CASE					,							mask_image=__SCREAMING_SNAKE_CASE					,							generator=__SCREAMING_SNAKE_CASE					,							num_inference_steps=2					,							output_type="""np"""					,							)

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






