
"""simple docstring"""


import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
		import onnxruntime as ort







class 				lowerCAmelCase_				(      lowercase__		,					unittest.TestCase    ):



						'''simple docstring'''


						_lowerCamelCase:    Optional[int]      =    '''ssube/stable-diffusion-x4-upscaler-onnx'''


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[str]  ,A_     :      Dict=0      )     ->   Tuple:
							A						   =   floats_tensor((1, 3, 128, 128)  ,rng=random.Random(__lowerCamelCase      )      )
							A						   =   torch.manual_seed(__lowerCamelCase      )
							A						   =   {
							    "prompt": "A painting of a squirrel eating a burger",
							    "image": image,
							    "generator": generator,
							    "num_inference_steps": 3,
							    "guidance_scale": 7.5,
							    "output_type": "numpy",
							}
							return inputs


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   Union[str, Any]:
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider='CPUExecutionProvider'      )
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   self.get_dummy_inputs()
							A						   =   pipe(**__lowerCamelCase      ).images
							A						   =   image[0, -3:, -3:, -1].flatten()

							# started as 128, should now be 512
							assert image.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]      )
							assert np.abs(image_slice - expected_slice      ).max() < 1e-1


						def 						_SCREAMING_SNAKE_CASE      ( self     :      Tuple      )     ->   Optional[int]:
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider='CPUExecutionProvider'      )
							A						   =   PNDMScheduler.from_config(pipe.scheduler.config  ,skip_prk_steps=__lowerCamelCase      )
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   self.get_dummy_inputs()
							A						   =   pipe(**__lowerCamelCase      ).images
							A						   =   image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.6_89_88_92, 0.59_24_05_56, 0.52_49_95_27, 0.58_86_62_15, 0.52_25_82_35, 0.52_57_27_15, 0.62_41_44_73, 0.6_17_43_87, 0.6_21_49_64]      )
							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1


						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Any:
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider='CPUExecutionProvider'      )
							A						   =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   self.get_dummy_inputs()
							A						   =   pipe(**__lowerCamelCase      ).images
							A						   =   image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.7_65_92_78, 0.76_43_76_64, 0.75_57_91_07, 0.7_69_11_16, 0.77_66_69_86, 0.7_72_76_72, 0.7_75_86_64, 0.7_81_22_26, 0.76_94_25_15]      )

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


						def 						_SCREAMING_SNAKE_CASE      ( self     :      List[Any]      )     ->   Any:
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider='CPUExecutionProvider'      )
							A						   =   EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   self.get_dummy_inputs()
							A						   =   pipe(**__lowerCamelCase      ).images
							A						   =   image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.6_97_47_82, 0.68_90_20_93, 0.70_13_58_85, 0.7_58_36_18, 0.7_80_45_45, 0.7_85_49_12, 0.78_66_74_26, 0.78_74_38_63, 0.78_07_02_23]      )
							assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1




						def 						_SCREAMING_SNAKE_CASE      ( self     :      int      )     ->   Dict:
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint  ,provider='CPUExecutionProvider'      )
							A						   =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   self.get_dummy_inputs()
							A						   =   pipe(**__lowerCamelCase      ).images
							A						   =   image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.77_42_44_96, 0.77_36_01, 0.7_64_52_88, 0.7_76_95_98, 0.7_77_27_39, 0.7_73_86_88, 0.78_18_72_33, 0.77_87_95_84, 0.76_70_43]      )

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







@nightly
@require_onnxruntime
@require_torch_gpu
class 				lowerCAmelCase_				(      unittest.TestCase    ):



						'''simple docstring'''


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[Any]      )     ->   List[Any]:
							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)


						@property
						def 						_SCREAMING_SNAKE_CASE      ( self     :      Optional[int]      )     ->   int:
							A						   =   ort.SessionOptions()
							A						   =   False
							return options


						def 						_SCREAMING_SNAKE_CASE      ( self     :      str      )     ->   List[Any]:
							A						   =   load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'      )
							A						   =   init_image.resize((128, 128)      )
							# using the PNDM scheduler by default
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'  ,provider=self.gpu_provider  ,sess_options=self.gpu_options  ,)
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   "A fantasy landscape, trending on artstation"

							A						   =   torch.manual_seed(0      )
							A						   =   pipe(
							    prompt=__lowerCamelCase  ,image=__lowerCamelCase  ,guidance_scale=7.5  ,num_inference_steps=10  ,generator=__lowerCamelCase  ,output_type='np'  ,)
							A						   =   output.images
							A						   =   images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							A						   =   np.array([0.48_83, 0.49_47, 0.49_80, 0.49_75, 0.49_82, 0.49_80, 0.50_00, 0.50_06, 0.49_72]      )
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




						def 						_SCREAMING_SNAKE_CASE      ( self     :      Dict      )     ->   Optional[int]:
							A						   =   load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'      )
							A						   =   init_image.resize((128, 128)      )
							A						   =   LMSDiscreteScheduler.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'  ,subfolder='scheduler'      )
							A						   =   OnnxStableDiffusionUpscalePipeline.from_pretrained(
							    'ssube/stable-diffusion-x4-upscaler-onnx'  ,scheduler=__lowerCamelCase  ,provider=self.gpu_provider  ,sess_options=self.gpu_options  ,)
							pipe.set_progress_bar_config(disable=__lowerCamelCase      )

							A						   =   "A fantasy landscape, trending on artstation"

							A						   =   torch.manual_seed(0      )
							A						   =   pipe(
							    prompt=__lowerCamelCase  ,image=__lowerCamelCase  ,guidance_scale=7.5  ,num_inference_steps=20  ,generator=__lowerCamelCase  ,output_type='np'  ,)
							A						   =   output.images
							A						   =   images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							A						   =   np.array(
							    [0.50_17_37_53, 0.50_22_33_56, 0.50_20_39, 0.50_23_30_36, 0.5_02_37_25, 0.5_02_26_01, 0.5_01_87_58, 0.50_23_40_85, 0.50_24_15_66]      )
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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

"""simple docstring"""


import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase							   =				data_utils.TransfoXLTokenizer
_lowercase							   =				data_utils.TransfoXLCorpus
_lowercase							   =				data_utils
_lowercase							   =				data_utils






def       _snake_case				(		snake_case__   :     Union[str, Any]	,      snake_case__   :     Dict	,      snake_case__   :     Tuple	,      snake_case__   :     int     ):
	if transfo_xl_dataset_file:
		# Convert a pre-processed corpus (see original TensorFlow repo)
		with open(snake_case__	,      'rb'     ) as fp:
			A						   =   pickle.load(snake_case__	,      encoding='latin1'     )
		# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
		A						   =   pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
		print(F'Save vocabulary to {pytorch_vocab_dump_path}'     )
		A						   =   corpus.vocab.__dict__
		torch.save(snake_case__	,      snake_case__     )

		A						   =   corpus.__dict__
		corpus_dict_no_vocab.pop('vocab'	,      snake_case__     )
		A						   =   pytorch_dump_folder_path + '/' + CORPUS_NAME
		print(F'Save dataset to {pytorch_dataset_dump_path}'     )
		torch.save(snake_case__	,      snake_case__     )

	if tf_checkpoint_path:
		# Convert a pre-trained TensorFlow model
		A						   =   os.path.abspath(snake_case__     )
		A						   =   os.path.abspath(snake_case__     )

		print(F'Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'     )
		# Initialise PyTorch model
		if transfo_xl_config_file == "":
			A						   =   TransfoXLConfig()
		else:
			A						   =   TransfoXLConfig.from_json_file(snake_case__     )
		print(F'Building PyTorch model from configuration: {config}'     )
		A						   =   TransfoXLLMHeadModel(snake_case__     )

		A						   =   load_tf_weights_in_transfo_xl(snake_case__	,      snake_case__	,      snake_case__     )
		# Save pytorch-model
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		A						   =   os.path.join(snake_case__	,      snake_case__     )
		print(F'Save PyTorch model to {os.path.abspath(snake_case__     )}'     )
		torch.save(model.state_dict()	,      snake_case__     )
		print(F'Save configuration file to {os.path.abspath(snake_case__     )}'     )
		with open(snake_case__	,      'w'	,      encoding='utf-8'     ) as f:
			f.write(config.to_json_string()     )


if __name__ == "__main__":
		_lowercase							   =				argparse.ArgumentParser()
		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(
		    '''--tf_checkpoint_path''',
		    default='''''',
		    type=str,
		    help='''An optional path to a TensorFlow checkpoint path to be converted.''',
		)
		parser.add_argument(
		    '''--transfo_xl_config_file''',
		    default='''''',
		    type=str,
		    help=(
		        '''An optional config json file corresponding to the pre-trained BERT model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--transfo_xl_dataset_file''',
		    default='''''',
		    type=str,
		    help='''An optional dataset file to be converted in a vocabulary.''',
		)
		_lowercase							   =				parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)