





"""simple docstring"""
from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


__UpperCamelCase     :						Union[str, Any]						       =  logging.get_logger(__name__)  # pylint: disable=invalid-name

__UpperCamelCase     :						List[Any]						       =  '''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")
        >>> pipe_prior.to(\"cuda\")
        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")
        >>> pipe.to(\"cuda\")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save(\"cat.png\")
        ```
'''

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Union[str, Any]  ,     _UpperCAmelCase					:    Dict  ,     _UpperCAmelCase					:    int=8			):
 lowerCAmelCase    				= height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 lowerCAmelCase    				= width // scale_factor**2
 if width % scale_factor**2 != 0:
  new_width += 1
 return new_height * scale_factor, new_width * scale_factor







class       a							(   A_  ):



 def __init__(     self				,       _snake_case				,       _snake_case				,       _snake_case				,       ):







  """simple docstring"""






  super().__init__()

  self.register_modules(
      unet=snake_case__				,       scheduler=snake_case__				,       movq=snake_case__				,       )
  lowerCAmelCase    				= 2 ** (len(self.movq.config.block_out_channels     ) - 1)



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  if latents is None:
   lowerCAmelCase    				= randn_tensor(snake_case__				,       generator=snake_case__				,       device=snake_case__				,       dtype=snake_case__     )
  else:
   if latents.shape != shape:
    raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'     )
   lowerCAmelCase    				= latents.to(snake_case__     )

  lowerCAmelCase    				= latents * scheduler.init_noise_sigma
  return latents



 def    UpperCamelCase__  (     self				,       _snake_case=0     ):







  """simple docstring"""






  if is_accelerate_available():
   from accelerate import cpu_offload
  else:
   raise ImportError('Please install accelerate via `pip install accelerate`'     )

  lowerCAmelCase    				= torch.device(F'cuda:{gpu_id}'     )

  lowerCAmelCase    				= [
      self.unet,
      self.movq,
  ]
  for cpu_offloaded_model in models:
   if cpu_offloaded_model is not None:
    cpu_offload(snake_case__				,       snake_case__     )



 def    UpperCamelCase__  (     self				,       _snake_case=0     ):







  """simple docstring"""






  if is_accelerate_available() and is_accelerate_version('>='				,       '0.17.0.dev0'     ):
   from accelerate import cpu_offload_with_hook
  else:
   raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'     )

  lowerCAmelCase    				= torch.device(F'cuda:{gpu_id}'     )

  if self.device.type != "cpu":
   self.to('cpu'				,       silence_dtype_warnings=snake_case__     )
   torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

  lowerCAmelCase    				= None
  for cpu_offloaded_model in [self.unet, self.movq]:
   lowerCAmelCase    				= cpu_offload_with_hook(snake_case__				,       snake_case__				,       prev_module_hook=snake_case__     )

  # We'll offload the last model manually.
  lowerCAmelCase    				= hook



 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  if not hasattr(self.unet				,       '_hf_hook'     ):
   return self.device
  for module in self.unet.modules():
   if (
       hasattr(snake_case__				,       '_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()
 @replace_example_docstring(snake_case__     )
 def __call__(     self				,       _snake_case				,       _snake_case				,       _snake_case = 5_12				,       _snake_case = 5_12				,       _snake_case = 1_00				,       _snake_case = 4.0				,       _snake_case = 1				,       _snake_case = None				,       _snake_case = None				,       _snake_case = "pil"				,       _snake_case = True				,       ):







  """simple docstring"""






  lowerCAmelCase    				= self._execution_device

  lowerCAmelCase    				= guidance_scale > 1.0

  if isinstance(snake_case__				,       snake_case__     ):
   lowerCAmelCase    				= torch.cat(snake_case__				,       dim=0     )
  lowerCAmelCase    				= image_embeds.shape[0] * num_images_per_prompt
  if isinstance(snake_case__				,       snake_case__     ):
   lowerCAmelCase    				= torch.cat(snake_case__				,       dim=0     )

  if do_classifier_free_guidance:
   lowerCAmelCase    				= image_embeds.repeat_interleave(snake_case__				,       dim=0     )
   lowerCAmelCase    				= negative_image_embeds.repeat_interleave(snake_case__				,       dim=0     )

  lowerCAmelCase    				= torch.cat([negative_image_embeds, image_embeds]				,       dim=0     ).to(dtype=self.unet.dtype				,       device=snake_case__     )

  self.scheduler.set_timesteps(snake_case__				,       device=snake_case__     )
  lowerCAmelCase    				= self.scheduler.timesteps

  lowerCAmelCase    				= self.unet.config.in_channels

  lowerCAmelCase    				= downscale_height_and_width(snake_case__				,       snake_case__				,       self.movq_scale_factor     )

  # create initial latent
  lowerCAmelCase    				= self.prepare_latents(
      (batch_size, num_channels_latents, height, width)				,       image_embeds.dtype				,       snake_case__				,       snake_case__				,       snake_case__				,       self.scheduler				,       )

  for i, t in enumerate(self.progress_bar(snake_case__     )     ):
   # expand the latents if we are doing classifier free guidance
   lowerCAmelCase    				= torch.cat([latents] * 2     ) if do_classifier_free_guidance else latents

   lowerCAmelCase    				= {"image_embeds": image_embeds}
   lowerCAmelCase    				= self.unet(
       sample=snake_case__				,       timestep=snake_case__				,       encoder_hidden_states=snake_case__				,       added_cond_kwargs=snake_case__				,       return_dict=snake_case__				,       )[0]

   if do_classifier_free_guidance:
    lowerCAmelCase    				= noise_pred.split(latents.shape[1]				,       dim=1     )
    lowerCAmelCase    				= noise_pred.chunk(2     )
    lowerCAmelCase    				= variance_pred.chunk(2     )
    lowerCAmelCase    				= noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    lowerCAmelCase    				= torch.cat([noise_pred, variance_pred_text]				,       dim=1     )

   if not (
       hasattr(self.scheduler.config				,       'variance_type'     )
       and self.scheduler.config.variance_type in ["learned", "learned_range"]
   ):
    lowerCAmelCase    				= noise_pred.split(latents.shape[1]				,       dim=1     )

   # compute the previous noisy sample x_t -> x_t-1
   lowerCAmelCase    				= self.scheduler.step(
       snake_case__				,       snake_case__				,       snake_case__				,       generator=snake_case__				,       )[0]
  # post-processing
  lowerCAmelCase    				= self.movq.decode(snake_case__				,       force_not_quantize=snake_case__     )["sample"]

  if output_type not in ["pt", "np", "pil"]:
   raise ValueError(F'Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'     )

  if output_type in ["np", "pil"]:
   lowerCAmelCase    				= image * 0.5 + 0.5
   lowerCAmelCase    				= image.clamp(0				,       1     )
   lowerCAmelCase    				= image.cpu().permute(0				,       2				,       3				,       1     ).float().numpy()

  if output_type == "pil":
   lowerCAmelCase    				= self.numpy_to_pil(snake_case__     )

  if not return_dict:
   return (image,)

  return ImagePipelineOutput(images=snake_case__     )







import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCamelCase				=			logging.get_logger(__name__)

__lowerCamelCase				=			[
    ("""bert.bert""", """visual_bert"""),
    ("""bert.cls""", """cls"""),
    ("""bert.classifier""", """cls"""),
    ("""token_type_embeddings_visual""", """visual_token_type_embeddings"""),
    ("""position_embeddings_visual""", """visual_position_embeddings"""),
    ("""projection""", """visual_projection"""),
]

__lowerCamelCase				=			[
    """nlvr2_coco_pre_trained.th""",
    """nlvr2_fine_tuned.th""",
    """nlvr2_pre_trained.th""",
    """vcr_coco_pre_train.th""",
    """vcr_fine_tune.th""",
    """vcr_pre_train.th""",
    """vqa_coco_pre_trained.th""",
    """vqa_fine_tuned.th""",
    """vqa_pre_trained.th""",
]


def      UpperCamelCase   (       __lowerCamelCase	:					int       ):
       snake_case			:  List[Any]     							= torch.load(__lowerCamelCase							,   map_location="cpu"       )
       return sd


def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[int]							,   __lowerCamelCase	:					Optional[Any]=rename_keys_prefix       ):
       snake_case			:  Optional[Any]     							= OrderedDict()
       snake_case			:  Optional[int]     							= torch.arange(config.max_position_embeddings       ).expand((1, -1)       )
       # detector_d = OrderedDict()
       for key in d:
              if "detector" in key:
                     # detector_d[key.replace('detector.','')] = d[key]
                     continue
              snake_case			:  Any     							= key
              for name_pair in rename_keys_prefix:
                     snake_case			:  List[Any]     							= new_key.replace(name_pair[0]							,   name_pair[1]       )
              snake_case			:  List[str]     							= d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     snake_case			:  int     							= new_d["cls.predictions.bias"]
       return new_d


@torch.no_grad()
def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[Any]       ):
       assert (
           checkpoint_path.split("/"       )[-1] in ACCEPTABLE_CHECKPOINTS
       ), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

       # Get Config
       if "pre" in checkpoint_path:
              snake_case			:  Optional[Any]     							= "pretraining"
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[int]     							= {"visual_embedding_dim": 512}
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  int     							= {"visual_embedding_dim": 2048}
              elif "vqa" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
              elif "nlvr" in checkpoint_path:
                     snake_case			:  Tuple     							= {"visual_embedding_dim": 1024}
              else:
                     raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""       )
       else:
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 512}
                     snake_case			:  List[str]     							= "multichoice"
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
                     snake_case			:  Tuple     							= "vqa_advanced"
              elif "vqa" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 2048, "num_labels": 3129}
                     snake_case			:  Tuple     							= "vqa"
              elif "nlvr" in checkpoint_path:
                     snake_case			:  str     							= {
                         "visual_embedding_dim": 1024,
                         "num_labels": 2,
                     }
                     snake_case			:  str     							= "nlvr"

       snake_case			:  int     							= VisualBertConfig(**__lowerCamelCase       )

       # Load State Dict
       snake_case			:  Any     							= load_state_dict(__lowerCamelCase       )

       snake_case			:  List[Any]     							= get_new_dict(__lowerCamelCase							,   __lowerCamelCase       )

       if model_type == "pretraining":
              snake_case			:  str     							= VisualBertForPreTraining(__lowerCamelCase       )
       elif model_type == "vqa":
              snake_case			:  Any     							= VisualBertForQuestionAnswering(__lowerCamelCase       )
       elif model_type == "nlvr":
              snake_case			:  Optional[Any]     							= VisualBertForVisualReasoning(__lowerCamelCase       )
       elif model_type == "multichoice":
              snake_case			:  Optional[int]     							= VisualBertForMultipleChoice(__lowerCamelCase       )

       model.load_state_dict(__lowerCamelCase       )
       # Save Checkpoints
       Path(__lowerCamelCase       ).mkdir(exist_ok=__lowerCamelCase       )
       model.save_pretrained(__lowerCamelCase       )


if __name__ == "__main__":
   __lowerCamelCase				=			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument("""orig_checkpoint_path""", type=str, help="""A path to .th on local filesystem.""")
   parser.add_argument("""pytorch_dump_folder_path""", type=str, help="""Path to the output PyTorch model.""")
   __lowerCamelCase				=			parser.parse_args()
   convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)





