






"""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_     :				List[Any]         =				logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase_     :				Optional[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       _A    (__a					,       __a					,       __a=8							)				-> Tuple:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						List[Any]							=	height // scale_factor**2
    if height % scale_factor**2 != 0:
        new_height += 1
    SCREAMING_SNAKE_CASE_   :						Tuple							=	width // scale_factor**2
    if width % scale_factor**2 != 0:
        new_width += 1
    return new_height * scale_factor, new_width * scale_factor





class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def __init__(		self  :						Any		,			lowercase_  :						UNetaDConditionModel		,			lowercase_  :						DDPMScheduler		,			lowercase_  :						VQModel		,			):




          '''simple docstring'''


          super().__init__()

          self.register_modules(
              unet=lowercase_		,			scheduler=lowercase_		,			movq=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						int							=	2 ** (len(self.movq.config.block_out_channels) - 1)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						int		,			lowercase_  :						List[Any]		,			lowercase_  :						Union[str, Any]		,			lowercase_  :						Any		,			lowercase_  :						Dict		,			lowercase_  :						Tuple):




          '''simple docstring'''


          if latents is None:
              SCREAMING_SNAKE_CASE_   :						Dict							=	randn_tensor(lowercase_		,			generator=lowercase_		,			device=lowercase_		,			dtype=lowercase_)
          else:
              if latents.shape != shape:
                  raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}')
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	latents.to(lowercase_)

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	latents * scheduler.init_noise_sigma
          return latents
      def _SCREAMING_SNAKE_CASE				(		self  :						int		,			lowercase_  :						List[str]=0):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_   :						Any							=	[
              self.unet,
              self.movq,
          ]
          for cpu_offloaded_model in models:
              if cpu_offloaded_model is not None:
                  cpu_offload(lowercase_		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						Optional[int]=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.''')

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

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

          SCREAMING_SNAKE_CASE_   :						Tuple							=	None
          for cpu_offloaded_model in [self.unet, self.movq]:
              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Any							=	cpu_offload_with_hook(lowercase_		,			lowercase_		,			prev_module_hook=lowercase_)

          # We'll offload the last model manually.
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	hook
      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          if not hasattr(self.unet		,			'''_hf_hook'''):
              return self.device
          for module in self.unet.modules():
              if (
                  hasattr(lowercase_		,			'''_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(lowercase_)
      def __call__(		self  :						int		,			lowercase_  :						Union[torch.FloatTensor, List[torch.FloatTensor]]		,			lowercase_  :						Union[torch.FloatTensor, List[torch.FloatTensor]]		,			lowercase_  :						int = 512		,			lowercase_  :						int = 512		,			lowercase_  :						int = 100		,			lowercase_  :						float = 4.0		,			lowercase_  :						int = 1		,			lowercase_  :						Optional[Union[torch.Generator, List[torch.Generator]]] = None		,			lowercase_  :						Optional[torch.FloatTensor] = None		,			lowercase_  :						Optional[str] = "pil"		,			lowercase_  :						bool = True		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self._execution_device

          SCREAMING_SNAKE_CASE_   :						int							=	guidance_scale > 1.0

          if isinstance(lowercase_		,			lowercase_):
              SCREAMING_SNAKE_CASE_   :						Dict							=	torch.cat(lowercase_		,			dim=0)
          SCREAMING_SNAKE_CASE_   :						Any							=	image_embeds.shape[0] * num_images_per_prompt
          if isinstance(lowercase_		,			lowercase_):
              SCREAMING_SNAKE_CASE_   :						Dict							=	torch.cat(lowercase_		,			dim=0)

          if do_classifier_free_guidance:
              SCREAMING_SNAKE_CASE_   :						Optional[int]							=	image_embeds.repeat_interleave(lowercase_		,			dim=0)
              SCREAMING_SNAKE_CASE_   :						List[str]							=	negative_image_embeds.repeat_interleave(lowercase_		,			dim=0)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	torch.cat([negative_image_embeds, image_embeds]		,			dim=0).to(dtype=self.unet.dtype		,			device=lowercase_)

          self.scheduler.set_timesteps(lowercase_		,			device=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	self.scheduler.timesteps

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.unet.config.in_channels

          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	downscale_height_and_width(lowercase_		,			lowercase_		,			self.movq_scale_factor)

          # create initial latent
          SCREAMING_SNAKE_CASE_   :						int							=	self.prepare_latents(
              (batch_size, num_channels_latents, height, width)		,			image_embeds.dtype		,			lowercase_		,			lowercase_		,			lowercase_		,			self.scheduler		,			)

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

              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{'''image_embeds''': image_embeds}
              SCREAMING_SNAKE_CASE_   :						Any							=	self.unet(
                  sample=lowercase_		,			timestep=lowercase_		,			encoder_hidden_states=lowercase_		,			added_cond_kwargs=lowercase_		,			return_dict=lowercase_		,			)[0]

              if do_classifier_free_guidance:
                  SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Any							=	noise_pred.split(latents.shape[1]		,			dim=1)
                  SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						str							=	noise_pred.chunk(2)
                  SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						List[str]							=	variance_pred.chunk(2)
                  SCREAMING_SNAKE_CASE_   :						Tuple							=	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                  SCREAMING_SNAKE_CASE_   :						str							=	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"]
              ):
                  SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	noise_pred.split(latents.shape[1]		,			dim=1)

              # compute the previous noisy sample x_t -> x_t-1
              SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.scheduler.step(
                  lowercase_		,			lowercase_		,			lowercase_		,			generator=lowercase_		,			)[0]
          # post-processing
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.movq.decode(lowercase_		,			force_not_quantize=lowercase_)['''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"]:
              SCREAMING_SNAKE_CASE_   :						Dict							=	image * 0.5 + 0.5
              SCREAMING_SNAKE_CASE_   :						Tuple							=	image.clamp(0		,			1)
              SCREAMING_SNAKE_CASE_   :						Dict							=	image.cpu().permute(0		,			2		,			3		,			1).float().numpy()

          if output_type == "pil":
              SCREAMING_SNAKE_CASE_   :						str							=	self.numpy_to_pil(lowercase_)

          if not return_dict:
              return (image,)

          return ImagePipelineOutput(images=lowercase_)










"""simple docstring"""






import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


UpperCAmelCase_     :				int         =				{
    """cola""": 2,
    """mnli""": 3,
    """mrpc""": 2,
    """sst-2""": 2,
    """sts-b""": 1,
    """qqp""": 2,
    """qnli""": 2,
    """rte""": 2,
    """wnli""": 2,
}


logging.set_verbosity_info()



def       _A    (__a					,       __a					,       __a					,       __a=None							)				-> int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						str							=	XLNetConfig.from_json_file(__a							)

    SCREAMING_SNAKE_CASE_   :						List[Any]							=	finetuning_task.lower() if finetuning_task is not None else ''''''
    if finetuning_task in GLUE_TASKS_NUM_LABELS:
        print(f'Building PyTorch XLNetForSequenceClassification model from configuration: {config}'							)
        SCREAMING_SNAKE_CASE_   :						List[Any]							=	finetuning_task
        SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	GLUE_TASKS_NUM_LABELS[finetuning_task]
        SCREAMING_SNAKE_CASE_   :						Dict							=	XLNetForSequenceClassification(__a							)
    elif "squad" in finetuning_task:
        SCREAMING_SNAKE_CASE_   :						str							=	finetuning_task
        SCREAMING_SNAKE_CASE_   :						Dict							=	XLNetForQuestionAnswering(__a							)
    else:
        SCREAMING_SNAKE_CASE_   :						List[str]							=	XLNetLMHeadModel(__a							)

    # Load weights from tf checkpoint
    load_tf_weights_in_xlnet(__a					,       __a					,       __a							)

    # Save pytorch-model
    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	os.path.join(__a					,       __a							)
    SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	os.path.join(__a					,       __a							)
    print(f'Save PyTorch model to {os.path.abspath(__a							)}'							)
    torch.save(model.state_dict()					,       __a							)
    print(f'Save configuration file to {os.path.abspath(__a							)}'							)
    with open(__a					,       '''w'''					,       encoding='''utf-8'''							) as f:
        f.write(config.to_json_string()							)


if __name__ == "__main__":
  UpperCAmelCase_     :				str         =				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
  )
  parser.add_argument(
      """--xlnet_config_file""",
      default=None,
      type=str,
      required=True,
      help=(
          """The config json file corresponding to the pre-trained XLNet model. \n"""
          """This specifies the model architecture."""
      ),
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""",
      default=None,
      type=str,
      required=True,
      help="""Path to the folder to store the PyTorch model or dataset/vocab.""",
  )
  parser.add_argument(
      """--finetuning_task""",
      default=None,
      type=str,
      help="""Name of a task on which the XLNet TensorFlow model was fine-tuned""",
  )
  UpperCAmelCase_     :				int         =				parser.parse_args()
  print(args)

  convert_xlnet_checkpoint_to_pytorch(
      args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
  )


