



'''simple docstring'''




import math
from typing import Callable, List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
from diffusers.schedulers import DDIMScheduler, DDPMScheduler, LMSDiscreteScheduler, PNDMScheduler



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase=[]		)      ->			List[str]:
							UpperCAmelCase     :		Any							     =		size[0] - overlap_pixels * 2
							UpperCAmelCase     :		List[str]							     =		size[1] - overlap_pixels * 2
							for letter in ["l", "r"]:
														if letter in remove_borders:
																					size_x += overlap_pixels
							for letter in ["t", "b"]:
														if letter in remove_borders:
																					size_y += overlap_pixels
							UpperCAmelCase     :		List[Any]							     =		np.ones((size_y, size_x)	,				dtype=np.uinta		) * 2_5_5
							UpperCAmelCase     :		List[str]							     =		np.pad(__lowerCAmelCase	,				mode="""linear_ramp"""	,				pad_width=__lowerCAmelCase	,				end_values=0		)

							if "l" in remove_borders:
														UpperCAmelCase     :		Optional[int]							     =		mask[:, overlap_pixels : mask.shape[1]]
							if "r" in remove_borders:
														UpperCAmelCase     :		Optional[Any]							     =		mask[:, 0 : mask.shape[1] - overlap_pixels]
							if "t" in remove_borders:
														UpperCAmelCase     :		List[str]							     =		mask[overlap_pixels : mask.shape[0], :]
							if "b" in remove_borders:
														UpperCAmelCase     :		Union[str, Any]							     =		mask[0 : mask.shape[0] - overlap_pixels, :]
							return mask



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase		)      ->			int:
							return max(__lowerCAmelCase	,				min(__lowerCAmelCase	,				__lowerCAmelCase		)		)



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase		)      ->			Any:
							return (
							    clamp(rect[0]	,				min[0]	,				max[0]		),
							    clamp(rect[1]	,				min[1]	,				max[1]		),
							    clamp(rect[2]	,				min[0]	,				max[0]		),
							    clamp(rect[3]	,				min[1]	,				max[1]		),
							)



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase		)      ->			Union[str, Any]:
							UpperCAmelCase     :		Union[str, Any]							     =		list(__lowerCAmelCase		)
							rect[0] -= overlap
							rect[1] -= overlap
							rect[2] += overlap
							rect[3] += overlap
							UpperCAmelCase     :		Tuple							     =		clamp_rect(__lowerCAmelCase	,				[0, 0]	,				[image_size[0], image_size[1]]		)
							return rect



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				_lowercase	,				_lowercase		)      ->			List[str]:
							UpperCAmelCase     :		List[str]							     =		Image.new("""RGB"""	,				(tile.size[0] + original_slice, tile.size[1])		)
							result.paste(
							    original_image.resize((tile.size[0], tile.size[1])	,				Image.BICUBIC		).crop(
							        (slice_x, 0, slice_x + original_slice, tile.size[1])		)	,				(0, 0)	,				)
							result.paste(__lowerCAmelCase	,				(original_slice, 0)		)
							return result



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			str:
							UpperCAmelCase     :		List[str]							     =		(original_image_slice * 4, 0, tile.size[0], tile.size[1])
							UpperCAmelCase     :		str							     =		tile.crop(__lowerCAmelCase		)
							return tile



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			List[Any]:
							UpperCAmelCase     :		Tuple							     =		n % d
							return n - divisor

class 			UpperCamelCase_ (							__lowerCamelCase			):





	def __init__(  self		,		A		,		A		,		A		,		A		,		A		,		A		,		A = 350		,		)			->			List[str]:
								super().__init__(
								    vae=__lowercase		,		text_encoder=__lowercase		,		tokenizer=__lowercase		,		unet=__lowercase		,		low_res_scheduler=__lowercase		,		scheduler=__lowercase		,		max_noise_level=__lowercase		,		)





	def  _lowercase(  self		,		A		,		A		,		A		,		A		,		A		,		A		,		A		,		**A    )			->			Optional[int]:
								torch.manual_seed(0    )
								UpperCAmelCase     :		List[str]							     =		(
								    min(image.size[0] - (tile_size + original_image_slice)		,		x * tile_size    ),
								    min(image.size[1] - (tile_size + original_image_slice)		,		y * tile_size    ),
								    min(image.size[0]		,		(x + 1) * tile_size    ),
								    min(image.size[1]		,		(y + 1) * tile_size    ),
								)
								UpperCAmelCase     :		Union[str, Any]							     =		add_overlap_rect(__lowercase		,		__lowercase		,		image.size    )
								UpperCAmelCase     :		Dict							     =		image.crop(__lowercase    )
								UpperCAmelCase     :		List[str]							     =		((crop_rect[0] + ((crop_rect[2] - crop_rect[0]) / 2)) / image.size[0]) * tile.size[0]
								UpperCAmelCase     :		List[Any]							     =		translated_slice_x - (original_image_slice / 2)
								UpperCAmelCase     :		Optional[Any]							     =		max(0		,		__lowercase    )
								UpperCAmelCase     :		List[Any]							     =		squeeze_tile(__lowercase		,		__lowercase		,		__lowercase		,		__lowercase    )
								UpperCAmelCase     :		Dict							     =		to_input.size
								UpperCAmelCase     :		Tuple							     =		to_input.resize((tile_size, tile_size)		,		Image.BICUBIC    )
								UpperCAmelCase     :		Union[str, Any]							     =		super(__lowercase		,		self    ).__call__(image=__lowercase		,		**__lowercase    ).images[0]
								UpperCAmelCase     :		Tuple							     =		upscaled_tile.resize((orig_input_size[0] * 4, orig_input_size[1] * 4)		,		Image.BICUBIC    )
								UpperCAmelCase     :		Union[str, Any]							     =		unsqueeze_tile(__lowercase		,		__lowercase    )
								UpperCAmelCase     :		Optional[int]							     =		upscaled_tile.resize((tile.size[0] * 4, tile.size[1] * 4)		,		Image.BICUBIC    )
								UpperCAmelCase     :		List[str]							     =		[]
								if x == 0:
															remove_borders.append("""l"""    )
								elif crop_rect[2] == image.size[0]:
															remove_borders.append("""r"""    )
								if y == 0:
															remove_borders.append("""t"""    )
								elif crop_rect[3] == image.size[1]:
															remove_borders.append("""b"""    )
								UpperCAmelCase     :		Dict							     =		Image.fromarray(
								    make_transparency_mask(
								        (upscaled_tile.size[0], upscaled_tile.size[1])		,		tile_border * 4		,		remove_borders=__lowercase    )		,		mode="""L"""		,		)
								final_image.paste(
								    __lowercase		,		(crop_rect_with_overlap[0] * 4, crop_rect_with_overlap[1] * 4)		,		__lowercase    )






	@torch.no_grad()
	def __call__(  self		,		A		,		A		,		A = 75		,		A = 9.0		,		A = 50		,		A = None		,		A = 1		,		A = 0.0		,		A = None		,		A = None		,		A = None		,		A = 1		,		A = 128		,		A = 32		,		A = 32		,		)			->			Tuple:
								UpperCAmelCase     :		Optional[Any]							     =		Image.new("""RGB"""		,		(image.size[0] * 4, image.size[1] * 4)    )
								UpperCAmelCase     :		Tuple							     =		math.ceil(image.size[0] / tile_size    )
								UpperCAmelCase     :		Any							     =		math.ceil(image.size[1] / tile_size    )
								UpperCAmelCase     :		Dict							     =		tcx * tcy
								UpperCAmelCase     :		Optional[Any]							     =		0
								for y in range(__lowercase    ):
															for x in range(__lowercase    ):
																						self._process_tile(
																						    __lowercase		,		__lowercase		,		__lowercase		,		__lowercase		,		__lowercase		,		__lowercase		,		__lowercase		,		prompt=__lowercase		,		num_inference_steps=__lowercase		,		guidance_scale=__lowercase		,		noise_level=__lowercase		,		negative_prompt=__lowercase		,		num_images_per_prompt=__lowercase		,		eta=__lowercase		,		generator=__lowercase		,		latents=__lowercase		,		)
																						current_count += 1
																						if callback is not None:
																													callback({"""progress""": current_count / total_tile_count, """image""": final_image}    )
								return final_image




def       __lowerCamelCase  (					)      ->			Optional[Any]:
							UpperCAmelCase     :		str							     =		'''stabilityai/stable-diffusion-x4-upscaler'''
							UpperCAmelCase     :		str							     =		StableDiffusionTiledUpscalePipeline.from_pretrained(__lowerCAmelCase	,				revision="""fp16"""	,				torch_dtype=torch.floataa		)
							UpperCAmelCase     :		Tuple							     =		pipe.to("""cuda"""		)
							UpperCAmelCase     :		Union[str, Any]							     =		Image.open("""../../docs/source/imgs/diffusers_library.jpg"""		)

							def callback(_lowercase		):
														print(F'''progress: {obj['progress']:.4f}'''		)
														obj["image"].save("""diffusers_library_progress.jpg"""		)

							UpperCAmelCase     :		Tuple							     =		pipe(image=__lowerCAmelCase	,				prompt="""Black font, white background, vector"""	,				noise_level=4_0	,				callback=__lowerCAmelCase		)
							final_image.save("""diffusers_library.jpg"""		)


if __name__ == "__main__":
							main()






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


A__								=						['''\nclass''', '''\ndef''', '''\n#''', '''\n@''', '''\nprint''', '''\nif''']


class       a							(			__lowerCamelCase				):



		def __init__(     self     :str   ,__lowercase     :Dict   ,__lowercase     :str   ,__lowercase     :Any=None   ,__lowercase     :Union[str, Any]=1  ):
						snake_case__	:      Optional[int]     		=	tokenizer
						snake_case__	:      List[Any]     		=	dataset
						snake_case__	:      Any     		=	len(__lowercase  ) if n_tasks is None else n_tasks
						snake_case__	:      List[Any]     		=	n_copies






		def __iter__(     self     :Tuple  ):
						snake_case__	:      Union[str, Any]     		=	[]
						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(__lowercase   ,padding=__lowercase   ,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       a							(			__lowerCamelCase				):



		def __init__(     self     :Union[str, Any]   ,__lowercase     :List[Any]   ,__lowercase     :int   ,__lowercase     :List[str]  ):
						snake_case__	:      List[str]     		=	start_length
						snake_case__	:      Any     		=	eof_strings
						snake_case__	:      Optional[Any]     		=	tokenizer






		def __call__(     self     :str   ,__lowercase     :str   ,__lowercase     :Any   ,**__lowercase     :Optional[int]  ):
						snake_case__	:      Tuple     		=	self.tokenizer.batch_decode(input_ids[:, self.start_length :]  )
						snake_case__	:      Union[str, Any]     		=	[]
						for decoded_generation in decoded_generations:
										done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings  )  )
						return all(__lowercase  )



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      List[Any]     		=	re.split('''(%s)''' % '''|'''.join(__lowerCAmelCase      )  ,  __lowerCAmelCase      )
				# last string should be ""
				return "".join(string_list[:-2]      )



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase=20  ,  **__lowerCAmelCase      )   ->					List[Any]:


				"""simple docstring"""

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

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

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

				snake_case__	:      List[Any]     		=	[[] for _ in range(__lowerCAmelCase      )]
				for task, generated_tokens in gen_token_dict.items():
								for s in generated_tokens:
												snake_case__	:      int     		=	tokenizer.decode(__lowerCAmelCase  ,  skip_special_tokens=__lowerCAmelCase  ,  clean_up_tokenization_spaces=__lowerCAmelCase      )
												code_gens[task].append(remove_last_block(__lowerCAmelCase      )      )
				return code_gens



def 			_lowerCAmelCase       (					)   ->					Optional[int]:


				"""simple docstring"""

				snake_case__	:      List[str]     		=	HfArgumentParser(__lowerCAmelCase      )
				snake_case__	:      Optional[Any]     		=	parser.parse_args()

				transformers.logging.set_verbosity_error()
				# enables code execution in code_eval metric
				snake_case__	:      Union[str, Any]     		=	args.HF_ALLOW_CODE_EVAL
				# make sure tokenizer plays nice with multiprocessing
				snake_case__	:      Union[str, Any]     		=	'''false'''

				if args.num_workers is None:
								snake_case__	:      Union[str, Any]     		=	multiprocessing.cpu_count()

				# Use dataset load to feed to accelerate
				snake_case__	:      Dict     		=	Accelerator()
				set_seed(args.seed  ,  device_specific=__lowerCAmelCase      )

				# Load model and tokenizer
				snake_case__	:      str     		=	AutoTokenizer.from_pretrained(args.model_ckpt      )
				snake_case__	:      List[str]     		=	tokenizer.eos_token
				snake_case__	:      Optional[Any]     		=	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  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )]      ),
				}

				# Load evaluation dataset and metric
				snake_case__	:      List[str]     		=	load_dataset('''openai_humaneval'''      )
				snake_case__	:      str     		=	load_metric('''code_eval'''      )

				snake_case__	:      List[Any]     		=	args.num_tasks if args.num_tasks is not None else len(human_eval['''test''']      )
				snake_case__	:      str     		=	args.n_samples // args.batch_size

				snake_case__	:      Tuple     		=	TokenizedDataset(__lowerCAmelCase  ,  human_eval['''test''']  ,  n_copies=__lowerCAmelCase  ,  n_tasks=__lowerCAmelCase      )
				# do not confuse args.batch_size, which is actually the num_return_sequences
				snake_case__	:      Tuple     		=	DataLoader(__lowerCAmelCase  ,  batch_size=1      )

				# Run a quick test to see if code evaluation is enabled
				try:
								snake_case__	:      Optional[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__	:      Optional[Any]     		=	accelerator.prepare(__lowerCAmelCase  ,  __lowerCAmelCase      )

				snake_case__	:      Dict     		=	complete_code(
				    __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  n_tasks=__lowerCAmelCase  ,  batch_size=args.batch_size  ,  **__lowerCAmelCase  ,  )

				if accelerator.is_main_process:
								snake_case__	:      Optional[Any]     		=	[]

								for task in tqdm(range(__lowerCAmelCase      )      ):
												snake_case__	:      str     		=	human_eval['''test'''][task]['''test''']
												snake_case__	:      int     		=	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__	:      Any     		=	code_eval_metric.compute(
								    references=__lowerCAmelCase  ,  predictions=__lowerCAmelCase  ,  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(__lowerCAmelCase  ,  __lowerCAmelCase      )


# 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()

