





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 SCREAMING_SNAKE_CASE_  (   snake_case_    ,			unittest.TestCase   ):
						# TODO: is there an appropriate internal test set?
						__magic_name__:     Tuple	 =    "ssube/stable-diffusion-x4-upscaler-onnx"






						def UpperCAmelCase_			( self : str						,				_A : Dict=0		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     List[str]     =		floats_tensor((1, 3, 128, 128)						,				rng=random.Random(_A		)		)
							snake_case_						:     Optional[Any]     =		torch.manual_seed(_A		)
							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 UpperCAmelCase_			( self : List[Any]		)							->      List[Any]:







							"""simple docstring"""


							snake_case_						:     Dict     =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,				provider='CPUExecutionProvider'		)
							pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     Union[str, Any]     =		self.get_dummy_inputs()
							snake_case_						:     Union[str, Any]     =		pipe(**_A		).images
							snake_case_						:     int     =		image[0, -3:, -3:, -1].flatten()

							# started as 128, should now be 512
							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     Tuple     =		np.array(
							    [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]		)
							assert np.abs(image_slice - expected_slice		).max() < 1E-1






						def UpperCAmelCase_			( self : Union[str, Any]		)							->      Dict:







							"""simple docstring"""


							snake_case_						:     Dict     =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,				provider='CPUExecutionProvider'		)
							snake_case_						:     Dict     =		PNDMScheduler.from_config(pipe.scheduler.config						,				skip_prk_steps=_A		)
							pipe.set_progress_bar_config(disable=_A		)

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

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     Dict     =		np.array(
							    [0.6_8_9_8_8_9_2, 0.5_9_2_4_0_5_5_6, 0.5_2_4_9_9_5_2_7, 0.5_8_8_6_6_2_1_5, 0.5_2_2_5_8_2_3_5, 0.5_2_5_7_2_7_1_5, 0.6_2_4_1_4_4_7_3, 0.6_1_7_4_3_8_7, 0.6_2_1_4_9_6_4]		)
							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1






						def UpperCAmelCase_			( self : Dict		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     Any     =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,				provider='CPUExecutionProvider'		)
							snake_case_						:     int     =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=_A		)

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

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     Dict     =		np.array(
							    [0.7_6_5_9_2_7_8, 0.7_6_4_3_7_6_6_4, 0.7_5_5_7_9_1_0_7, 0.7_6_9_1_1_1_6, 0.7_7_6_6_6_9_8_6, 0.7_7_2_7_6_7_2, 0.7_7_5_8_6_6_4, 0.7_8_1_2_2_2_6, 0.7_6_9_4_2_5_1_5]		)

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






						def UpperCAmelCase_			( self : Optional[Any]		)							->      str:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,				provider='CPUExecutionProvider'		)
							snake_case_						:     Union[str, Any]     =		EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=_A		)

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

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     List[str]     =		np.array(
							    [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]		)
							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1






						def UpperCAmelCase_			( self : List[Any]		)							->      Any:







							"""simple docstring"""


							snake_case_						:     str     =		OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,				provider='CPUExecutionProvider'		)
							snake_case_						:     Any     =		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
							pipe.set_progress_bar_config(disable=_A		)

							snake_case_						:     List[Any]     =		self.get_dummy_inputs()
							snake_case_						:     List[Any]     =		pipe(**_A		).images
							snake_case_						:     Tuple     =		image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case_						:     List[str]     =		np.array(
							    [0.7_7_4_2_4_4_9_6, 0.7_7_3_6_0_1, 0.7_6_4_5_2_8_8, 0.7_7_6_9_5_9_8, 0.7_7_7_2_7_3_9, 0.7_7_3_8_6_8_8, 0.7_8_1_8_7_2_3_3, 0.7_7_8_7_9_5_8_4, 0.7_6_7_0_4_3]		)

							assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-1
@nightly
@require_onnxruntime
@require_torch_gpu
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):






						@property
						def UpperCAmelCase_			( self : List[str]		)							->      Dict:







							"""simple docstring"""


							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)






						@property
						def UpperCAmelCase_			( self : str		)							->      Any:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		ort.SessionOptions()
							snake_case_						:     Union[str, Any]     =		False
							return options






						def UpperCAmelCase_			( self : int		)							->      Tuple:







							"""simple docstring"""


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

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

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

							assert images.shape == (1, 512, 512, 3)
							snake_case_						:     List[str]     =		np.array([0.4_8_8_3, 0.4_9_4_7, 0.4_9_8_0, 0.4_9_7_5, 0.4_9_8_2, 0.4_9_8_0, 0.5_0_0_0, 0.5_0_0_6, 0.4_9_7_2]		)
							# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






						def UpperCAmelCase_			( self : Dict		)							->      Any:







							"""simple docstring"""


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

							snake_case_						:     Dict     =		'A fantasy landscape, trending on artstation'

							snake_case_						:     Dict     =		torch.manual_seed(0		)
							snake_case_						:     Tuple     =		pipe(
							    prompt=_A						,				image=_A						,				guidance_scale=7.5						,				num_inference_steps=20						,				generator=_A						,				output_type='np'						,				)
							snake_case_						:     List[Any]     =		output.images
							snake_case_						:     int     =		images[0, 255:258, 383:386, -1]

							assert images.shape == (1, 512, 512, 3)
							snake_case_						:     Optional[int]     =		np.array(
							    [0.5_0_1_7_3_7_5_3, 0.5_0_2_2_3_3_5_6, 0.5_0_2_0_3_9, 0.5_0_2_3_3_0_3_6, 0.5_0_2_3_7_2_5, 0.5_0_2_2_6_0_1, 0.5_0_1_8_7_5_8, 0.5_0_2_3_4_0_8_5, 0.5_0_2_4_1_5_6_6]		)
							# 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
import multiprocessing
import os
import re
from collections import defaultdict

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from arguments import HumanEvalArguments
from datasets import load_dataset, load_metric
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, StoppingCriteria, StoppingCriteriaList


_SCREAMING_SNAKE_CASE 		=						["""\nclass""", """\ndef""", """\n#""", """\n@""", """\nprint""", """\nif"""]
class SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : Tuple						,				_A : Any						,				_A : str						,				_A : int=None						,				_A : str=1		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Union[str, Any]     =		tokenizer
							snake_case_						:     Optional[int]     =		dataset
							snake_case_						:     List[str]     =		len(_A		) if n_tasks is None else n_tasks
							snake_case_						:     List[Any]     =		n_copies






						def __iter__( self : Any		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     List[str]     =		[]
							for task in range(self.n_tasks		):
								# without strip, the model generate commented codes ...
								prompts.append(self.tokenizer.eos_token + self.dataset[task]['prompt'].strip()		)
							snake_case_						:     Optional[int]     =		self.tokenizer(_A						,				padding=_A						,				return_tensors='pt'		)
							for task in range(self.n_tasks		):
								for _ in range(self.n_copies		):
									yield {
									    "ids": outputs.input_ids[task],
									    "task_id": task,
									    "input_len": outputs.attention_mask[task].sum(),
									}
class SCREAMING_SNAKE_CASE_  (   snake_case_   ):






						def __init__( self : List[Any]						,				_A : Optional[int]						,				_A : str						,				_A : Dict		)							->      Any:







							"""simple docstring"""


							snake_case_						:     List[str]     =		start_length
							snake_case_						:     int     =		eof_strings
							snake_case_						:     Dict     =		tokenizer






						def __call__( self : Any						,				_A : Union[str, Any]						,				_A : Dict						,				**_A : List[Any]		)							->      List[str]:







							"""simple docstring"""


							snake_case_						:     Optional[int]     =		self.tokenizer.batch_decode(input_ids[:, self.start_length :]		)
							snake_case_						:     Optional[int]     =		[]
							for decoded_generation in decoded_generations:
								done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings		)		)
							return all(_A		)


def 							SCREAMING_SNAKE_CASE__  (      __a				):
	snake_case_						:     List[str]     =		re.split('(%s)' % '|'.join(__a				)				,	__a				)
	# last string should be ""
	return "".join(string_list[:-2]				)


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				,	__a				,	__a				,	__a=20				,	**__a				):
	snake_case_						:     Tuple     =		defaultdict(__a				)  # dict of list of generated tokens
	for step, batch in tqdm(enumerate(__a				)				):
		with torch.no_grad():
			snake_case_						:     Optional[Any]     =		batch['ids'].shape[-1]
			snake_case_						:     List[str]     =		accelerator.unwrap_model(__a				).generate(
			    input_ids=batch['ids'][:, : batch['input_len']]				,	num_return_sequences=__a				,	**__a				)
			# each task is generated batch_size times
			snake_case_						:     List[str]     =		batch['task_id'].repeat(__a				)
			snake_case_						:     Union[str, Any]     =		accelerator.pad_across_processes(
			    __a				,	dim=1				,	pad_index=tokenizer.pad_token_id				)

			snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.gather((generated_tokens, generated_tasks)				)
			snake_case_						:     Optional[Any]     =		generated_tokens.cpu().numpy()
			snake_case_						:     Dict     =		generated_tasks.cpu().numpy()

			for task, generated_tokens in zip(__a				,	__a				):
				gen_token_dict[task].append(__a				)

	snake_case_						:     Tuple     =		[[] for _ in range(__a				)]
	for task, generated_tokens in gen_token_dict.items():
		for s in generated_tokens:
			snake_case_						:     int     =		tokenizer.decode(__a				,	skip_special_tokens=__a				,	clean_up_tokenization_spaces=__a				)
			code_gens[task].append(remove_last_block(__a				)				)
	return code_gens


def 							SCREAMING_SNAKE_CASE__  (      ):
	# Setup configuration
	snake_case_						:     Optional[int]     =		HfArgumentParser(__a				)
	snake_case_						:     Tuple     =		parser.parse_args()

	transformers.logging.set_verbosity_error()
	# enables code execution in code_eval metric
	snake_case_						:     Optional[int]     =		args.HF_ALLOW_CODE_EVAL
	# make sure tokenizer plays nice with multiprocessing
	snake_case_						:     int     =		'false'

	if args.num_workers is None:
		snake_case_						:     Any     =		multiprocessing.cpu_count()

	# Use dataset load to feed to accelerate
	snake_case_						:     List[Any]     =		Accelerator()
	set_seed(args.seed				,	device_specific=__a				)

	# Load model and tokenizer
	snake_case_						:     Any     =		AutoTokenizer.from_pretrained(args.model_ckpt				)
	snake_case_						:     int     =		tokenizer.eos_token
	snake_case_						:     List[str]     =		AutoModelForCausalLM.from_pretrained(args.model_ckpt				)

	# Generation settings
	snake_case_						:     List[Any]     =		{
	    'do_sample': args.do_sample,
	    'temperature': args.temperature,
	    'max_new_tokens': args.max_new_tokens,
	    'top_p': args.top_p,
	    'top_k': args.top_k,
	    'stopping_criteria': StoppingCriteriaList([EndOfFunctionCriteria(0				,	__a				,	__a				)]				),
	}

	# Load evaluation dataset and metric
	snake_case_						:     Dict     =		load_dataset('openai_humaneval'				)
	snake_case_						:     Optional[Any]     =		load_metric('code_eval'				)

	snake_case_						:     Union[str, Any]     =		args.num_tasks if args.num_tasks is not None else len(human_eval['test']				)
	snake_case_						:     Optional[int]     =		args.n_samples // args.batch_size

	snake_case_						:     Dict     =		TokenizedDataset(__a				,	human_eval['test']				,	n_copies=__a				,	n_tasks=__a				)
	# do not confuse args.batch_size, which is actually the num_return_sequences
	snake_case_						:     Optional[Any]     =		DataLoader(__a				,	batch_size=1				)

	# Run a quick test to see if code evaluation is enabled
	try:
		snake_case_						:     Union[str, Any]     =		code_eval_metric.compute(references=['']				,	predictions=[['']]				)
	except ValueError as exception:
		print(
		    'Code evaluation not enabled. Read the warning below carefully and then use `--HF_ALLOW_CODE_EVAL="1"`'
		    ' flag to enable code evaluation.'				)
		raise exception

	snake_case_  ,snake_case_						:     Union[str, Any]     =		accelerator.prepare(__a				,	__a				)

	snake_case_						:     str     =		complete_code(
	    __a				,	__a				,	__a				,	__a				,	n_tasks=__a				,	batch_size=args.batch_size				,	**__a				,	)

	if accelerator.is_main_process:
		snake_case_						:     Tuple     =		[]

		for task in tqdm(range(__a				)				):
			snake_case_						:     Union[str, Any]     =		human_eval['test'][task]['test']
			snake_case_						:     Union[str, Any]     =		f"""check({human_eval['test'][task]['entry_point']})"""
			references.append('\n' + test_func + '\n' + entry_point				)

		# Evaluate completions with "code_eval" metric
		snake_case_  ,snake_case_						:     int     =		code_eval_metric.compute(
		    references=__a				,	predictions=__a				,	num_workers=args.num_workers				)
		print(f"""Results: {pass_at_k}"""				)

		# Save results to json file
		with open(args.output_file				,	'w'				) as fp:
			json.dump(__a				,	__a				)


# For some reason the folliwng seems to be necessary sometimes for code_eval to work nice with multiprocessing
# https://stackoverflow.com/questions/60804599/python-multiprocessing-keeps-spawning-the-whole-script
if __name__ == "__main__":
						main()

