


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        _lowercase      (			A__						,			unittest.TestCase       ):






				'''simple docstring'''



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


				def 		__magic_name__(       self			:List[str]			,			lowerCAmelCase__			:List[Any]=0	)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		floats_tensor((1, 3, 128, 128)			,			rng=random.Random(lowerCAmelCase__	)	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		torch.manual_seed(lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		{
										    '''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 		__magic_name__(       self			:Any	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     int					=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,			provider='''CPUExecutionProvider'''	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		self.get_dummy_inputs()
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		pipe(**lowerCAmelCase__	).images
										__SCREAMING_SNAKE_CASE						:     str					=		image[0, -3:, -3:, -1].flatten()

										# started as 128, should now be 512
										assert image.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     int					=		np.array(
										    [0.697_4782, 0.6890_2093, 0.7013_5885, 0.758_3618, 0.780_4545, 0.785_4912, 0.7866_7426, 0.7874_3863, 0.7807_0223]	)
										assert np.abs(image_slice - expected_slice	).max() < 1E-1


				def 		__magic_name__(       self			:Any	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,			provider='''CPUExecutionProvider'''	)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		PNDMScheduler.from_config(pipe.scheduler.config			,			skip_prk_steps=lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     int					=		self.get_dummy_inputs()
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		pipe(**lowerCAmelCase__	).images
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		np.array(
										    [0.689_8892, 0.5924_0556, 0.5249_9527, 0.5886_6215, 0.5225_8235, 0.5257_2715, 0.6241_4473, 0.617_4387, 0.621_4964]	)
										assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-1


				def 		__magic_name__(       self			:int	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     Any					=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,			provider='''CPUExecutionProvider'''	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		self.get_dummy_inputs()
										__SCREAMING_SNAKE_CASE						:     int					=		pipe(**lowerCAmelCase__	).images
										__SCREAMING_SNAKE_CASE						:     Tuple					=		image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		np.array(
										    [0.765_9278, 0.7643_7664, 0.7557_9107, 0.769_1116, 0.7766_6986, 0.772_7672, 0.775_8664, 0.781_2226, 0.7694_2515]	)

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


				def 		__magic_name__(       self			:Dict	)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     Dict					=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,			provider='''CPUExecutionProvider'''	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Dict					=		self.get_dummy_inputs()
										__SCREAMING_SNAKE_CASE						:     List[str]					=		pipe(**lowerCAmelCase__	).images
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		np.array(
										    [0.697_4782, 0.6890_2093, 0.7013_5885, 0.758_3618, 0.780_4545, 0.785_4912, 0.7866_7426, 0.7874_3863, 0.7807_0223]	)
										assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-1



				def 		__magic_name__(       self			:Dict	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint			,			provider='''CPUExecutionProvider'''	)
										__SCREAMING_SNAKE_CASE						:     Any					=		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		self.get_dummy_inputs()
										__SCREAMING_SNAKE_CASE						:     Any					=		pipe(**lowerCAmelCase__	).images
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		np.array(
										    [0.7742_4496, 0.77_3601, 0.764_5288, 0.776_9598, 0.777_2739, 0.773_8688, 0.7818_7233, 0.7787_9584, 0.76_7043]	)

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




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






				'''simple docstring'''



				@property
				def 		__magic_name__(       self			:Any	)		->			Tuple:
										return (
										    "CUDAExecutionProvider",
										    {
										        "gpu_mem_limit": "15000000000",  # 15GB
										        "arena_extend_strategy": "kSameAsRequested",
										    },
										)


				@property
				def 		__magic_name__(       self			:Tuple	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		ort.SessionOptions()
										__SCREAMING_SNAKE_CASE						:     Dict					=		False
										return options


				def 		__magic_name__(       self			:Any	)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''	)
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		init_image.resize((128, 128)	)
										# using the PNDM scheduler by default
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		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__	)

										__SCREAMING_SNAKE_CASE						:     str					=		'''A fantasy landscape, trending on artstation'''

										__SCREAMING_SNAKE_CASE						:     List[str]					=		torch.manual_seed(0	)
										__SCREAMING_SNAKE_CASE						:     Any					=		pipe(
										    prompt=lowerCAmelCase__			,			image=lowerCAmelCase__			,			guidance_scale=7.5			,			num_inference_steps=10			,			generator=lowerCAmelCase__			,			output_type='''np'''			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		output.images
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     Any					=		np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]	)
										# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



				def 		__magic_name__(       self			:int	)		->			str:
										__SCREAMING_SNAKE_CASE						:     int					=		load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''	)
										__SCREAMING_SNAKE_CASE						:     Dict					=		init_image.resize((128, 128)	)
										__SCREAMING_SNAKE_CASE						:     int					=		LMSDiscreteScheduler.from_pretrained(
										    '''ssube/stable-diffusion-x4-upscaler-onnx'''			,			subfolder='''scheduler'''	)
										__SCREAMING_SNAKE_CASE						:     int					=		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__	)

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		'''A fantasy landscape, trending on artstation'''

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		torch.manual_seed(0	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		pipe(
										    prompt=lowerCAmelCase__			,			image=lowerCAmelCase__			,			guidance_scale=7.5			,			num_inference_steps=20			,			generator=lowerCAmelCase__			,			output_type='''np'''			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		output.images
										__SCREAMING_SNAKE_CASE						:     str					=		images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 512, 3)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		np.array(
										    [0.5017_3753, 0.5022_3356, 0.50_2039, 0.5023_3036, 0.502_3725, 0.502_2601, 0.501_8758, 0.5023_4085, 0.5024_1566]	)
										# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
						from transformers.pipelines.conversational import Conversation


__lowerCAmelCase	: Optional[int]  				=logging.get_logger(__name__)

__lowerCAmelCase	: Optional[Any]  				={'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

__lowerCAmelCase	: List[str]  				={
    'tokenizer_file': {
        'EleutherAI/gpt-neox-20b': 'https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/tokenizer.json',
    },
}

__lowerCAmelCase	: Optional[int]  				={
    'gpt-neox-20b': 2_0_4_8,
}
class        _lowercase      (			A__       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			str						=		VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE__				:			Dict						=		PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE__				:			Any						=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		['''input_ids''', '''attention_mask''']


				def __init__(       self			:int			,			lowerCAmelCase__			:Any=None			,			lowerCAmelCase__			:Optional[Any]=None			,			lowerCAmelCase__			:List[Any]=None			,			lowerCAmelCase__			:str="<|endoftext|>"			,			lowerCAmelCase__			:str="<|endoftext|>"			,			lowerCAmelCase__			:Dict="<|endoftext|>"			,			lowerCAmelCase__			:Union[str, Any]=False			,			**lowerCAmelCase__			:List[str]			,			)		->			Any:
										super().__init__(
										    lowerCAmelCase__			,			lowerCAmelCase__			,			tokenizer_file=lowerCAmelCase__			,			unk_token=lowerCAmelCase__			,			bos_token=lowerCAmelCase__			,			eos_token=lowerCAmelCase__			,			add_prefix_space=lowerCAmelCase__			,			**lowerCAmelCase__			,			)

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
										if pre_tok_state.get('''add_prefix_space'''			,			lowerCAmelCase__	) != add_prefix_space:
																__SCREAMING_SNAKE_CASE						:     List[str]					=		getattr(lowerCAmelCase__			,			pre_tok_state.pop('''type'''	)	)
																__SCREAMING_SNAKE_CASE						:     str					=		add_prefix_space
																__SCREAMING_SNAKE_CASE						:     Any					=		pre_tok_class(**lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		add_prefix_space


				def 		__magic_name__(       self			:Union[str, Any]			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Optional[str] = None	)		->			Tuple[str]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self._tokenizer.model.save(lowerCAmelCase__			,			name=lowerCAmelCase__	)
										return tuple(lowerCAmelCase__	)



				def 		__magic_name__(       self			:Optional[Any]			,			lowerCAmelCase__			:"Conversation"	)		->			List[int]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		[]
										for is_user, text in conversation.iter_texts():
																input_ids.extend(self.encode(lowerCAmelCase__			,			add_special_tokens=lowerCAmelCase__	) + [self.eos_token_id]	)

										if len(lowerCAmelCase__	) > self.model_max_length:
																__SCREAMING_SNAKE_CASE						:     List[str]					=		input_ids[-self.model_max_length :]
										return input_ids
