'''simple docstring'''




import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)



def     __snake_case			(							lowerCAmelCase							:  Any   ,	lowerCAmelCase							:  int      ):
 __UpperCAmelCase								=       XCLIPTextConfig()

 # derive patch size from model name
 __UpperCAmelCase								=       model_name.find('patch'      )
 __UpperCAmelCase								=       int(model_name[start_idx + len('patch'      ) : start_idx + len('patch'      ) + 2]      )
 __UpperCAmelCase								=       XCLIPVisionConfig(patch_size=_lowerCAmelCase   ,	num_frames=_lowerCAmelCase      )

 if "large" in model_name:
  __UpperCAmelCase								=       768
  __UpperCAmelCase								=       3072
  __UpperCAmelCase								=       12

  __UpperCAmelCase								=       1024
  __UpperCAmelCase								=       4096
  __UpperCAmelCase								=       16
  __UpperCAmelCase								=       24
  __UpperCAmelCase								=       768
  __UpperCAmelCase								=       3072

 if model_name == "xclip-large-patch14-16-frames":
  __UpperCAmelCase								=       336

 __UpperCAmelCase								=       XCLIPConfig.from_text_vision_configs(_lowerCAmelCase   ,	_lowerCAmelCase      )

 if "large" in model_name:
  __UpperCAmelCase								=       768

 return config



def     __snake_case			(							lowerCAmelCase							:  Optional[int]      ):
 # text encoder
 if name == "token_embedding.weight":
  __UpperCAmelCase								=       name.replace('token_embedding.weight'   ,	'text_model.embeddings.token_embedding.weight'      )
 if name == "positional_embedding":
  __UpperCAmelCase								=       name.replace('positional_embedding'   ,	'text_model.embeddings.position_embedding.weight'      )
 if "ln_1" in name:
  __UpperCAmelCase								=       name.replace('ln_1'   ,	'layer_norm1'      )
 if "ln_2" in name:
  __UpperCAmelCase								=       name.replace('ln_2'   ,	'layer_norm2'      )
 if "c_fc" in name:
  __UpperCAmelCase								=       name.replace('c_fc'   ,	'fc1'      )
 if "c_proj" in name:
  __UpperCAmelCase								=       name.replace('c_proj'   ,	'fc2'      )
 if name.startswith('transformer.resblocks'      ):
  __UpperCAmelCase								=       name.replace('transformer.resblocks'   ,	'text_model.encoder.layers'      )
 if "attn.out_proj" in name and "message" not in name:
  __UpperCAmelCase								=       name.replace('attn.out_proj'   ,	'self_attn.out_proj'      )
 if "ln_final" in name:
  __UpperCAmelCase								=       name.replace('ln_final'   ,	'text_model.final_layer_norm'      )
 # visual encoder
 if name == "visual.class_embedding":
  __UpperCAmelCase								=       name.replace('visual.class_embedding'   ,	'vision_model.embeddings.class_embedding'      )
 if name == "visual.positional_embedding":
  __UpperCAmelCase								=       name.replace('visual.positional_embedding'   ,	'vision_model.embeddings.position_embedding.weight'      )
 if name.startswith('visual.transformer.resblocks'      ):
  __UpperCAmelCase								=       name.replace('visual.transformer.resblocks'   ,	'vision_model.encoder.layers'      )
 if "visual.conv1" in name:
  __UpperCAmelCase								=       name.replace('visual.conv1'   ,	'vision_model.embeddings.patch_embedding'      )
 if "visual.ln_pre" in name:
  __UpperCAmelCase								=       name.replace('visual.ln_pre'   ,	'vision_model.pre_layernorm'      )
 if "visual.ln_post" in name:
  __UpperCAmelCase								=       name.replace('visual.ln_post'   ,	'vision_model.post_layernorm'      )
 if "visual.proj" in name:
  __UpperCAmelCase								=       name.replace('visual.proj'   ,	'visual_projection.weight'      )
 if "text_projection" in name:
  __UpperCAmelCase								=       name.replace('text_projection'   ,	'text_projection.weight'      )
 # things on top
 if "prompts_visual_proj" in name:
  __UpperCAmelCase								=       name.replace('prompts_visual_proj'   ,	'prompts_visual_projection'      )
 if "prompts_visual_ln" in name:
  __UpperCAmelCase								=       name.replace('prompts_visual_ln'   ,	'prompts_visual_layernorm'      )
 # mit
 if name == "mit.positional_embedding":
  __UpperCAmelCase								=       name.replace('positional'   ,	'position'      )
 if name.startswith('mit.resblocks'      ):
  __UpperCAmelCase								=       name.replace('mit.resblocks'   ,	'mit.encoder.layers'      )
 # prompts generator
 if name.startswith('prompts_generator.norm'      ):
  __UpperCAmelCase								=       name.replace('prompts_generator.norm'   ,	'prompts_generator.layernorm'      )

 return name



def     __snake_case			(							lowerCAmelCase							:  Union[str, Any]   ,	lowerCAmelCase							:  List[Any]      ):
 for key in orig_state_dict.copy().keys():
  __UpperCAmelCase								=       orig_state_dict.pop(_lowerCAmelCase      )

  if "attn.in_proj" in key:
   __UpperCAmelCase								=       key.split('.'      )
   if key.startswith('visual'      ):
    __UpperCAmelCase								=       key_split[3]
    __UpperCAmelCase								=       config.vision_config.hidden_size
    if "message_attn" in key:
     if "weight" in key:
      __UpperCAmelCase								=       val[
          :dim, :
      ]
      __UpperCAmelCase								=       val[
          dim : dim * 2, :
      ]
      __UpperCAmelCase								=       val[
          -dim:, :
      ]
     else:
      __UpperCAmelCase								=       val[
          :dim
      ]
      __UpperCAmelCase								=       val[
          dim : dim * 2
      ]
      __UpperCAmelCase								=       val[
          -dim:
      ]
    else:
     if "weight" in key:
      __UpperCAmelCase								=       val[
          :dim, :
      ]
      __UpperCAmelCase								=       val[
          dim : dim * 2, :
      ]
      __UpperCAmelCase								=       val[
          -dim:, :
      ]
     else:
      __UpperCAmelCase								=       val[:dim]
      __UpperCAmelCase								=       val[
          dim : dim * 2
      ]
      __UpperCAmelCase								=       val[-dim:]
   elif key.startswith('mit'      ):
    __UpperCAmelCase								=       key_split[2]
    __UpperCAmelCase								=       config.vision_config.mit_hidden_size
    if "weight" in key:
     __UpperCAmelCase								=       val[:dim, :]
     __UpperCAmelCase								=       val[dim : dim * 2, :]
     __UpperCAmelCase								=       val[-dim:, :]
    else:
     __UpperCAmelCase								=       val[:dim]
     __UpperCAmelCase								=       val[dim : dim * 2]
     __UpperCAmelCase								=       val[-dim:]
   else:
    __UpperCAmelCase								=       key_split[2]
    __UpperCAmelCase								=       config.text_config.hidden_size
    if "weight" in key:
     __UpperCAmelCase								=       val[:dim, :]
     __UpperCAmelCase								=       val[
         dim : dim * 2, :
     ]
     __UpperCAmelCase								=       val[-dim:, :]
    else:
     __UpperCAmelCase								=       val[:dim]
     __UpperCAmelCase								=       val[
         dim : dim * 2
     ]
     __UpperCAmelCase								=       val[-dim:]
  else:
   __UpperCAmelCase								=       rename_key(_lowerCAmelCase      )
   if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
    __UpperCAmelCase								=       val.T
   __UpperCAmelCase								=       val

 return orig_state_dict



def     __snake_case			(							lowerCAmelCase							:  List[str]      ):
 if num_frames == 8:
  __UpperCAmelCase								=       'eating_spaghetti_8_frames.npy'
 elif num_frames == 16:
  __UpperCAmelCase								=       'eating_spaghetti.npy'
 elif num_frames == 32:
  __UpperCAmelCase								=       'eating_spaghetti_32_frames.npy'
 __UpperCAmelCase								=       hf_hub_download(
     repo_id='hf-internal-testing/spaghetti-video'   ,	filename=_lowerCAmelCase   ,	repo_type='dataset'   ,	)
 __UpperCAmelCase								=       np.load(_lowerCAmelCase      )
 return list(_lowerCAmelCase      )



def     __snake_case			(							lowerCAmelCase							:  Any   ,	lowerCAmelCase							:  Dict=None   ,	lowerCAmelCase							:  Optional[int]=False      ):
 __UpperCAmelCase								=       {
     # fully supervised kinetics-400 checkpoints
     'xclip-base-patch32': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth',
     'xclip-base-patch32-16-frames': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'
     ),
     'xclip-base-patch16': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth',
     'xclip-base-patch16-16-frames': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'
     ),
     'xclip-large-patch14': 'https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb',
     'xclip-large-patch14-16-frames': 'https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f',
     # fully supervised kinetics-600 checkpoints
     'xclip-base-patch16-kinetics-600': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'
     ),
     'xclip-base-patch16-kinetics-600-16-frames': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'
     ),
     'xclip-large-patch14-kinetics-600': 'https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be',
     # few shot
     'xclip-base-patch16-hmdb-2-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'
     ),
     'xclip-base-patch16-hmdb-4-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'
     ),
     'xclip-base-patch16-hmdb-8-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'
     ),
     'xclip-base-patch16-hmdb-16-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'
     ),
     'xclip-base-patch16-ucf-2-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'
     ),
     'xclip-base-patch16-ucf-4-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'
     ),
     'xclip-base-patch16-ucf-8-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'
     ),
     'xclip-base-patch16-ucf-16-shot': (
         'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'
     ),
     # zero shot
     'xclip-base-patch16-zero-shot': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth',
 }

 __UpperCAmelCase								=       model_to_url[model_name]
 __UpperCAmelCase								=       8
 if "16-frames" in model_name:
  __UpperCAmelCase								=       16
 elif "shot" in model_name:
  __UpperCAmelCase								=       32

 __UpperCAmelCase								=       get_xclip_config(_lowerCAmelCase   ,	_lowerCAmelCase      )
 __UpperCAmelCase								=       XCLIPModel(_lowerCAmelCase      )
 model.eval()

 if "drive" in checkpoint_url:
  __UpperCAmelCase								=       'pytorch_model.bin'
  gdown.cached_download(_lowerCAmelCase   ,	_lowerCAmelCase   ,	quiet=_lowerCAmelCase      )
  __UpperCAmelCase								=       torch.load(_lowerCAmelCase   ,	map_location='cpu'      )['model']
 else:
  __UpperCAmelCase								=       torch.hub.load_state_dict_from_url(_lowerCAmelCase      )['model']

 __UpperCAmelCase								=       convert_state_dict(_lowerCAmelCase   ,	_lowerCAmelCase      )

 __UpperCAmelCase								=       XCLIPModel(_lowerCAmelCase      )
 __UpperCAmelCase								=       model.load_state_dict(_lowerCAmelCase   ,	strict=_lowerCAmelCase      )
 assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
 model.eval()

 __UpperCAmelCase								=       336 if model_name == 'xclip-large-patch14-16-frames' else 224
 __UpperCAmelCase								=       VideoMAEImageProcessor(size=_lowerCAmelCase      )
 __UpperCAmelCase								=       CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'      )
 __UpperCAmelCase								=       CLIPTokenizerFast.from_pretrained('openai/clip-vit-base-patch32'      )
 __UpperCAmelCase								=       XCLIPProcessor(image_processor=_lowerCAmelCase   ,	tokenizer=_lowerCAmelCase      )

 __UpperCAmelCase								=       prepare_video(_lowerCAmelCase      )
 __UpperCAmelCase								=       processor(
     text=['playing sports', 'eating spaghetti', 'go shopping']   ,	videos=_lowerCAmelCase   ,	return_tensors='pt'   ,	padding=_lowerCAmelCase      )

 print('Shape of pixel values:'   ,	inputs.pixel_values.shape      )

 with torch.no_grad():
  __UpperCAmelCase								=       model(**_lowerCAmelCase      )

 # Verify outputs
 __UpperCAmelCase								=       outputs.logits_per_video
 __UpperCAmelCase								=       logits_per_video.softmax(dim=1      )
 print('Probs:'   ,	_lowerCAmelCase      )
 # kinetics-400
 if model_name == "xclip-base-patch32":
  __UpperCAmelCase								=       torch.tensor([[0.00_19, 0.99_51, 0.00_30]]      )
 elif model_name == "xclip-base-patch32-16-frames":
  __UpperCAmelCase								=       torch.tensor([[7.0_999E-04, 9.9_883E-01, 4.5_580E-04]]      )
 elif model_name == "xclip-base-patch16":
  __UpperCAmelCase								=       torch.tensor([[0.00_83, 0.96_81, 0.02_36]]      )
 elif model_name == "xclip-base-patch16-16-frames":
  __UpperCAmelCase								=       torch.tensor([[7.6_937E-04, 9.9_728E-01, 1.9_473E-03]]      )
 elif model_name == "xclip-large-patch14":
  __UpperCAmelCase								=       torch.tensor([[0.00_62, 0.98_64, 0.00_75]]      )
 elif model_name == "xclip-large-patch14-16-frames":
  __UpperCAmelCase								=       torch.tensor([[3.3_877E-04, 9.9_937E-01, 2.8_888E-04]]      )
 # kinetics-600
 elif model_name == "xclip-base-patch16-kinetics-600":
  __UpperCAmelCase								=       torch.tensor([[0.05_55, 0.89_14, 0.05_31]]      )
 elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
  __UpperCAmelCase								=       torch.tensor([[3.8_554E-04, 9.9_929E-01, 3.2_754E-04]]      )
 elif model_name == "xclip-large-patch14-kinetics-600":
  __UpperCAmelCase								=       torch.tensor([[0.00_36, 0.99_20, 0.00_45]]      )
 # few shot
 elif model_name == "xclip-base-patch16-hmdb-2-shot":
  __UpperCAmelCase								=       torch.tensor([[7.1_890E-06, 9.9_994E-01, 5.6_559E-05]]      )
 elif model_name == "xclip-base-patch16-hmdb-4-shot":
  __UpperCAmelCase								=       torch.tensor([[1.0_320E-05, 9.9_993E-01, 6.2_435E-05]]      )
 elif model_name == "xclip-base-patch16-hmdb-8-shot":
  __UpperCAmelCase								=       torch.tensor([[4.1_377E-06, 9.9_990E-01, 9.8_386E-05]]      )
 elif model_name == "xclip-base-patch16-hmdb-16-shot":
  __UpperCAmelCase								=       torch.tensor([[4.1_347E-05, 9.9_962E-01, 3.3_411E-04]]      )
 elif model_name == "xclip-base-patch16-ucf-2-shot":
  __UpperCAmelCase								=       torch.tensor([[8.5_857E-05, 9.9_928E-01, 6.3_291E-04]]      )
 elif model_name == "xclip-base-patch16-ucf-4-shot":
  __UpperCAmelCase								=       torch.tensor([[8.5_857E-05, 9.9_928E-01, 6.3_291E-04]]      )
 elif model_name == "xclip-base-patch16-ucf-8-shot":
  __UpperCAmelCase								=       torch.tensor([[0.00_27, 0.99_04, 0.00_70]]      )
 elif model_name == "xclip-base-patch16-ucf-16-shot":
  __UpperCAmelCase								=       torch.tensor([[9.8_219E-04, 9.9_593E-01, 3.0_863E-03]]      )
 # zero shot
 elif model_name == "xclip-base-patch16-zero-shot":
  __UpperCAmelCase								=       torch.tensor([[3.5_082E-04, 9.9_785E-01, 1.7_966E-03]]      )
 else:
  raise ValueError(F"""Model name {model_name} not supported"""      )
 assert torch.allclose(_lowerCAmelCase   ,	_lowerCAmelCase   ,	atol=1E-3      )
 print('Looks ok!'      )

 if pytorch_dump_folder_path is not None:
  print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""      )
  model.save_pretrained(_lowerCAmelCase      )

 if push_to_hub:
  print('Pushing model, processor and slow tokenizer files to the hub...'      )
  model.push_to_hub(_lowerCAmelCase   ,	organization='nielsr'      )
  processor.push_to_hub(_lowerCAmelCase   ,	organization='nielsr'      )
  slow_tokenizer.push_to_hub(_lowerCAmelCase   ,	organization='nielsr'      )


if __name__ == "__main__":
     _UpperCamelCase					:		Union[str, Any]							      =		argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default='xclip-base-patch32',
         type=str,
         help='Name of the model.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )
     parser.add_argument(
         '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
     )

     _UpperCamelCase					:		str							      =		parser.parse_args()
     convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








'''simple docstring'''





from typing import Callable, List, Optional, Union

import PIL
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPTextModel,
    CLIPTokenizer,
)

from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import deprecate, is_accelerate_available, logging


_UpperCamelCase											=		logging.get_logger(__name__)  # pylint: disable=invalid-name






class    lowerCamelCase_       (							SCREAMING_SNAKE_CASE_	):
  """simple docstring"""




  def __init__(					self :      Union[str, Any]     ,      _a :      CLIPSegForImageSegmentation     ,      _a :      CLIPSegProcessor     ,      _a :      AutoencoderKL     ,      _a :      CLIPTextModel     ,      _a :      CLIPTokenizer     ,      _a :      UNetaDConditionModel     ,      _a :      Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]     ,      _a :      StableDiffusionSafetyChecker     ,      _a :      CLIPImageProcessor     ,      )   ->   int:
    super().__init__()

    if hasattr(scheduler.config     ,      'steps_offset'  ) and scheduler.config.steps_offset != 1:
      __lowerCamelCase			:			int       			=					(
          f'The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`'
          f' should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure '
          'to update the config accordingly as leaving `steps_offset` might led to incorrect results'
          ' in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,'
          ' it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`'
          ' file'
      )
      deprecate('steps_offset!=1'     ,      '1.0.0'     ,      _a     ,      standard_warn=_a  )
      __lowerCamelCase			:			int       			=					dict(scheduler.config  )
      __lowerCamelCase			:			Any       			=					1
      __lowerCamelCase			:			List[Any]       			=					FrozenDict(_a  )

    if hasattr(scheduler.config     ,      'skip_prk_steps'  ) and scheduler.config.skip_prk_steps is False:
      __lowerCamelCase			:			str       			=					(
          f'The configuration file of this scheduler: {scheduler} has not set the configuration'
          ' `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make'
          ' sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to'
          ' incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face'
          ' Hub, it would be very nice if you could open a Pull request for the'
          ' `scheduler/scheduler_config.json` file'
      )
      deprecate('skip_prk_steps not set'     ,      '1.0.0'     ,      _a     ,      standard_warn=_a  )
      __lowerCamelCase			:			Dict       			=					dict(scheduler.config  )
      __lowerCamelCase			:			Tuple       			=					True
      __lowerCamelCase			:			Union[str, Any]       			=					FrozenDict(_a  )

    if safety_checker is None:
      logger.warning(
          f'You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure'
          ' that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered'
          ' results in services or applications open to the public. Both the diffusers team and Hugging Face'
          ' strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling'
          ' it only for use-cases that involve analyzing network behavior or auditing its results. For more'
          ' information, please have a look at https://github.com/huggingface/diffusers/pull/254 .'  )

    self.register_modules(
        segmentation_model=_a     ,      segmentation_processor=_a     ,      vae=_a     ,      text_encoder=_a     ,      tokenizer=_a     ,      unet=_a     ,      scheduler=_a     ,      safety_checker=_a     ,      feature_extractor=_a     ,      )




  def 		_lowercase	(					self :      Optional[int]     ,      _a :      Optional[Union[str, int]] = "auto"  )   ->   List[str]:
    if slice_size == "auto":
      # half the attention head size is usually a good trade-off between
      # speed and memory
      __lowerCamelCase			:			List[Any]       			=					self.unet.config.attention_head_dim // 2
    self.unet.set_attention_slice(_a  )




  def 		_lowercase	(					self :      Optional[Any]  )   ->   int:
    self.enable_attention_slicing(_a  )




  def 		_lowercase	(					self :      str  )   ->   Any:
    if is_accelerate_available():
      from accelerate import cpu_offload
    else:
      raise ImportError('Please install accelerate via `pip install accelerate`'  )

    __lowerCamelCase			:			Dict       			=					torch.device('cuda'  )

    for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
      if cpu_offloaded_model is not None:
        cpu_offload(_a     ,      _a  )




  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 		_lowercase	(					self :      Dict  )   ->   int:
    if self.device != torch.device('meta'  ) or not hasattr(self.unet     ,      '_hf_hook'  ):
      return self.device
    for module in self.unet.modules():
      if (
          hasattr(_a     ,      '_hf_hook'  )
          and hasattr(module._hf_hook     ,      'execution_device'  )
          and module._hf_hook.execution_device is not None
      ):
        return torch.device(module._hf_hook.execution_device  )
    return self.device




  @torch.no_grad()
  def __call__(					self :      List[str]     ,      _a :      Union[str, List[str]]     ,      _a :      Union[torch.FloatTensor, PIL.Image.Image]     ,      _a :      str     ,      _a :      int = 512     ,      _a :      int = 512     ,      _a :      int = 50     ,      _a :      float = 7.5     ,      _a :      Optional[Union[str, List[str]]] = None     ,      _a :      Optional[int] = 1     ,      _a :      float = 0.0     ,      _a :      Optional[torch.Generator] = None     ,      _a :      Optional[torch.FloatTensor] = None     ,      _a :      Optional[str] = "pil"     ,      _a :      bool = True     ,      _a :      Optional[Callable[[int, int, torch.FloatTensor], None]] = None     ,      _a :      int = 1     ,      **_a :      Optional[Any]     ,      )   ->   int:
    __lowerCamelCase			:			Union[str, Any]       			=					self.segmentation_processor(
        text=[text]     ,      images=[image]     ,      padding='max_length'     ,      return_tensors='pt'  ).to(self.device  )
    __lowerCamelCase			:			Union[str, Any]       			=					self.segmentation_model(**_a  )
    __lowerCamelCase			:			List[str]       			=					torch.sigmoid(outputs.logits  ).cpu().detach().unsqueeze(-1  ).numpy()
    __lowerCamelCase			:			Optional[Any]       			=					self.numpy_to_pil(_a  )[0].resize(image.size  )

    # Run inpainting pipeline with the generated mask
    __lowerCamelCase			:			Tuple       			=					StableDiffusionInpaintPipeline(
        vae=self.vae     ,      text_encoder=self.text_encoder     ,      tokenizer=self.tokenizer     ,      unet=self.unet     ,      scheduler=self.scheduler     ,      safety_checker=self.safety_checker     ,      feature_extractor=self.feature_extractor     ,      )
    return inpainting_pipeline(
        prompt=_a     ,      image=_a     ,      mask_image=_a     ,      height=_a     ,      width=_a     ,      num_inference_steps=_a     ,      guidance_scale=_a     ,      negative_prompt=_a     ,      num_images_per_prompt=_a     ,      eta=_a     ,      generator=_a     ,      latents=_a     ,      output_type=_a     ,      return_dict=_a     ,      callback=_a     ,      callback_steps=_a     ,      )






