






'''simple docstring'''
import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
	UpperCamelCase__ =			argparse.ArgumentParser()

	parser.add_argument(
	    '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
	)
	# !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
	parser.add_argument(
	    '--original_config_file',
	    default=None,
	    type=str,
	    help='The YAML config file corresponding to the original architecture.',
	)
	parser.add_argument(
	    '--num_in_channels',
	    default=None,
	    type=int,
	    help='The number of input channels. If `None` number of input channels will be automatically inferred.',
	)
	parser.add_argument(
	    '--scheduler_type',
	    default='pndm',
	    type=str,
	    help='Type of scheduler to use. Should be one of [\'pndm\', \'lms\', \'ddim\', \'euler\', \'euler-ancestral\', \'dpm\']',
	)
	parser.add_argument(
	    '--pipeline_type',
	    default=None,
	    type=str,
	    help=(
	        'The pipeline type. One of \'FrozenOpenCLIPEmbedder\', \'FrozenCLIPEmbedder\', \'PaintByExample\''
	        '. If `None` pipeline will be automatically inferred.'
	    ),
	)
	parser.add_argument(
	    '--image_size',
	    default=None,
	    type=int,
	    help=(
	        'The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2'
	        ' Base. Use 768 for Stable Diffusion v2.'
	    ),
	)
	parser.add_argument(
	    '--prediction_type',
	    default=None,
	    type=str,
	    help=(
	        'The prediction type that the model was trained on. Use \'epsilon\' for Stable Diffusion v1.X and Stable'
	        ' Diffusion v2 Base. Use \'v_prediction\' for Stable Diffusion v2.'
	    ),
	)
	parser.add_argument(
	    '--extract_ema',
	    action='store_true',
	    help=(
	        'Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights'
	        ' or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield'
	        ' higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning.'
	    ),
	)
	parser.add_argument(
	    '--upcast_attention',
	    action='store_true',
	    help=(
	        'Whether the attention computation should always be upcasted. This is necessary when running stable'
	        ' diffusion 2.1.'
	    ),
	)
	parser.add_argument(
	    '--from_safetensors',
	    action='store_true',
	    help='If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.',
	)
	parser.add_argument(
	    '--to_safetensors',
	    action='store_true',
	    help='Whether to store pipeline in safetensors format or not.',
	)
	parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')
	parser.add_argument('--device', type=str, help='Device to use (e.g. cpu, cuda:0, cuda:1, etc.)')
	parser.add_argument(
	    '--stable_unclip',
	    type=str,
	    default=None,
	    required=False,
	    help='Set if this is a stable unCLIP model. One of \'txt2img\' or \'img2img\'.',
	)
	parser.add_argument(
	    '--stable_unclip_prior',
	    type=str,
	    default=None,
	    required=False,
	    help='Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.',
	)
	parser.add_argument(
	    '--clip_stats_path',
	    type=str,
	    help='Path to the clip stats file. Only required if the stable unclip model\'s config specifies `model.params.noise_aug_config.params.clip_stats_path`.',
	    required=False,
	)
	parser.add_argument(
	    '--controlnet', action='store_true', default=None, help='Set flag if this is a controlnet checkpoint.'
	)
	parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
	parser.add_argument(
	    '--vae_path',
	    type=str,
	    default=None,
	    required=False,
	    help='Set to a path, hub id to an already converted vae to not convert it again.',
	)
	UpperCamelCase__ =			parser.parse_args()

	UpperCamelCase__ =			download_from_original_stable_diffusion_ckpt(
	    checkpoint_path=args.checkpoint_path,
	    original_config_file=args.original_config_file,
	    image_size=args.image_size,
	    prediction_type=args.prediction_type,
	    model_type=args.pipeline_type,
	    extract_ema=args.extract_ema,
	    scheduler_type=args.scheduler_type,
	    num_in_channels=args.num_in_channels,
	    upcast_attention=args.upcast_attention,
	    from_safetensors=args.from_safetensors,
	    device=args.device,
	    stable_unclip=args.stable_unclip,
	    stable_unclip_prior=args.stable_unclip_prior,
	    clip_stats_path=args.clip_stats_path,
	    controlnet=args.controlnet,
	    vae_path=args.vae_path,
	)

	if args.half:
		pipe.to(torch_dtype=torch.floataa)

	if args.controlnet:
		# only save the controlnet model
		pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
	else:
		pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)












'''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__ =			{
    '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 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase=None							):


   """simple docstring"""




   lowercase_	:    Tuple           =     XLNetConfig.from_json_file(_UpperCamelCase							)

   lowercase_	:    Union[str, 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}"""							)
      lowercase_	:    Dict           =     finetuning_task
      lowercase_	:    Any           =     GLUE_TASKS_NUM_LABELS[finetuning_task]
      lowercase_	:    Any           =     XLNetForSequenceClassification(_UpperCamelCase							)
   elif "squad" in finetuning_task:
      lowercase_	:    Optional[int]           =     finetuning_task
      lowercase_	:    Optional[int]           =     XLNetForQuestionAnswering(_UpperCamelCase							)
   else:
      lowercase_	:    Union[str, Any]           =     XLNetLMHeadModel(_UpperCamelCase							)

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

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


if __name__ == "__main__":
 UpperCamelCase__ =			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__ =			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
 )




