






import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM







@require_torch
class     UpperCamelCase_ :





	def __init__(   self      :				Optional[int]			,			lowerCAmelCase_      :				Tuple			,			lowerCAmelCase_      :				int=99			,			lowerCAmelCase_      :				int=13			,			lowerCAmelCase_      :				Any=16			,			lowerCAmelCase_      :				Dict=7			,			lowerCAmelCase_      :				Any=True			,			lowerCAmelCase_      :				Tuple=True			,			lowerCAmelCase_      :				Union[str, Any]=True			,			lowerCAmelCase_      :				str=False			,			lowerCAmelCase_      :				Union[str, Any]=True			,			lowerCAmelCase_      :				List[str]=2			,			lowerCAmelCase_      :				Dict=32			,			lowerCAmelCase_      :				int=4			,			lowerCAmelCase_      :				List[str]=4			,			lowerCAmelCase_      :				Union[str, Any]=30			,			lowerCAmelCase_      :				str=0			,			lowerCAmelCase_      :				Union[str, Any]=1			,			lowerCAmelCase_      :				Any=2			,			lowerCAmelCase_      :				Any=None			,			)		->    Union[str, Any]:
				UpperCAmelCase_  :				Tuple      =      parent
				UpperCAmelCase_  :				Tuple      =      batch_size
				UpperCAmelCase_  :				List[str]      =      decoder_seq_length
				# For common tests
				UpperCAmelCase_  :				List[str]      =      self.decoder_seq_length
				UpperCAmelCase_  :				List[Any]      =      is_training
				UpperCAmelCase_  :				List[Any]      =      use_attention_mask
				UpperCAmelCase_  :				Dict      =      use_labels

				UpperCAmelCase_  :				Union[str, Any]      =      vocab_size
				UpperCAmelCase_  :				List[str]      =      d_model
				UpperCAmelCase_  :				str      =      d_model
				UpperCAmelCase_  :				List[str]      =      decoder_layers
				UpperCAmelCase_  :				Any      =      decoder_layers
				UpperCAmelCase_  :				int      =      decoder_ffn_dim
				UpperCAmelCase_  :				Tuple      =      decoder_attention_heads
				UpperCAmelCase_  :				Union[str, Any]      =      decoder_attention_heads
				UpperCAmelCase_  :				Optional[int]      =      eos_token_id
				UpperCAmelCase_  :				Dict      =      bos_token_id
				UpperCAmelCase_  :				int      =      pad_token_id
				UpperCAmelCase_  :				List[Any]      =      decoder_start_token_id
				UpperCAmelCase_  :				Union[str, Any]      =      use_cache
				UpperCAmelCase_  :				List[str]      =      max_position_embeddings

				UpperCAmelCase_  :				Union[str, Any]      =      None
				UpperCAmelCase_  :				Tuple      =      decoder_seq_length
				UpperCAmelCase_  :				int      =      2
				UpperCAmelCase_  :				Optional[Any]      =      1





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]		)		->    List[str]:
				UpperCAmelCase_  :				int      =      ids_tensor([self.batch_size, self.decoder_seq_length]			,			self.vocab_size		)

				UpperCAmelCase_  :				Dict      =      None
				if self.use_attention_mask:
							UpperCAmelCase_  :				Optional[Any]      =      ids_tensor([self.batch_size, self.decoder_seq_length]			,			vocab_size=2		)

				UpperCAmelCase_  :				Tuple      =      None
				if self.use_labels:
							UpperCAmelCase_  :				List[Any]      =      ids_tensor([self.batch_size, self.decoder_seq_length]			,			self.vocab_size		)

				UpperCAmelCase_  :				int      =      TrOCRConfig(
				    vocab_size=self.vocab_size			,			d_model=self.d_model			,			decoder_layers=self.decoder_layers			,			decoder_ffn_dim=self.decoder_ffn_dim			,			decoder_attention_heads=self.decoder_attention_heads			,			eos_token_id=self.eos_token_id			,			bos_token_id=self.bos_token_id			,			use_cache=self.use_cache			,			pad_token_id=self.pad_token_id			,			decoder_start_token_id=self.decoder_start_token_id			,			max_position_embeddings=self.max_position_embeddings			,			)

				return (config, input_ids, attention_mask, lm_labels)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any			,			lowerCAmelCase_      :				Dict			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				Tuple			,			)		->    List[Any]:
				UpperCAmelCase_  :				Optional[int]      =      True
				UpperCAmelCase_  :				Optional[int]      =      TrOCRDecoder(config=lowercase__		).to(lowercase__		).eval()
				UpperCAmelCase_  :				List[str]      =      input_ids[:2]

				input_ids[input_ids == 0] += 1
				# first forward pass
				UpperCAmelCase_  :				Tuple      =      model(lowercase__			,			use_cache=lowercase__		)
				UpperCAmelCase_  :				Optional[Any]      =      model(lowercase__		)
				UpperCAmelCase_  :				Optional[Any]      =      model(lowercase__			,			use_cache=lowercase__		)

				self.parent.assertTrue(len(lowercase__		) == len(lowercase__		)		)
				self.parent.assertTrue(len(lowercase__		) == len(lowercase__		) + 1		)

				UpperCAmelCase_  :				List[str]      =      outputs["""past_key_values"""]

				# create hypothetical next token and extent to next_input_ids
				UpperCAmelCase_  :				int      =      ids_tensor((2, 1)			,			config.vocab_size - 1		) + 1

				# append to next input_ids and
				UpperCAmelCase_  :				Dict      =      torch.cat([input_ids, next_tokens]			,			dim=-1		)

				UpperCAmelCase_  :				Optional[Any]      =      model(lowercase__		)["""last_hidden_state"""]
				UpperCAmelCase_  :				List[Any]      =      model(lowercase__			,			past_key_values=lowercase__		)["""last_hidden_state"""]

				# select random slice
				UpperCAmelCase_  :				str      =      ids_tensor((1,)			,			output_from_past.shape[-1]		).item()
				UpperCAmelCase_  :				List[str]      =      output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
				UpperCAmelCase_  :				str      =      output_from_past[:, 0, random_slice_idx].detach()

				# test that outputs are equal for slice
				assert torch.allclose(lowercase__			,			lowercase__			,			atol=1e-3		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict		)		->    str:
				UpperCAmelCase_  :				Optional[Any]      =      self.prepare_config_and_inputs()
				UpperCAmelCase_  :				List[str]      =      config_and_inputs

				UpperCAmelCase_  :				Union[str, Any]      =      {"""input_ids""": input_ids, """attention_mask""": attention_mask}
				return config, inputs_dict







@require_torch
class     UpperCamelCase_ (_UpperCAmelCase       ,		_UpperCAmelCase       ,		_UpperCAmelCase       ,		unittest.TestCase  ):
	__magic_name__			   =       (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
	__magic_name__			   =       (TrOCRForCausalLM,) if is_torch_available() else ()
	__magic_name__			   =       {"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
	__magic_name__			   =       True
	__magic_name__			   =       False





	def   _SCREAMING_SNAKE_CASE   (   self      :				int		)		->    Union[str, Any]:
				UpperCAmelCase_  :				Dict      =      TrOCRStandaloneDecoderModelTester(self			,			is_training=lowercase__		)
				UpperCAmelCase_  :				Tuple      =      ConfigTester(self			,			config_class=lowercase__		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    Union[str, Any]:
				pass





	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    str:
				pass





	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    Any:
				pass





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Optional[Any]:
				self.config_tester.run_common_tests()





	def   _SCREAMING_SNAKE_CASE   (   self      :				str		)		->    List[str]:
				UpperCAmelCase_  :				Optional[Any]      =      self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_decoder_model_past(*lowercase__		)





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict		)		->    Any:
				return





	@unittest.skip("The model doesn't support left padding"		)  # and it's not used enough to be worth fixing :)
	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    int:
				pass











"""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        snake_case   (     A__			,A__			,A__=[]		):
			UpperCAmelCase_  :				int      =      size[0] - overlap_pixels * 2
			UpperCAmelCase_  :				Dict      =      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_  :				Union[str, Any]      =      np.ones((size_y, size_x)			,dtype=np.uinta		) * 2_55
			UpperCAmelCase_  :				Optional[int]      =      np.pad(A__			,mode="linear_ramp"			,pad_width=A__			,end_values=0		)

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

def        snake_case   (     A__			,A__			,A__		):
			return max(A__			,min(A__			,A__		)		)

def        snake_case   (     A__			,A__			,A__		):
			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        snake_case   (     A__			,A__			,A__		):
			UpperCAmelCase_  :				List[str]      =      list(A__		)
			rect[0] -= overlap
			rect[1] -= overlap
			rect[2] += overlap
			rect[3] += overlap
			UpperCAmelCase_  :				Optional[Any]      =      clamp_rect(A__			,[0, 0]			,[image_size[0], image_size[1]]		)
			return rect

def        snake_case   (     A__			,A__			,A__			,A__		):
			UpperCAmelCase_  :				Dict      =      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(A__			,(original_slice, 0)		)
			return result

def        snake_case   (     A__			,A__		):
			UpperCAmelCase_  :				Optional[int]      =      (original_image_slice * 4, 0, tile.size[0], tile.size[1])
			UpperCAmelCase_  :				Optional[int]      =      tile.crop(A__		)
			return tile

def        snake_case   (     A__			,A__		):
			UpperCAmelCase_  :				List[str]      =      n % d
			return n - divisor







class     UpperCamelCase_ (__A  ):





	def __init__(   self      :				List[str]			,			lowerCAmelCase_      :				AutoencoderKL			,			lowerCAmelCase_      :				CLIPTextModel			,			lowerCAmelCase_      :				CLIPTokenizer			,			lowerCAmelCase_      :				UNetaDConditionModel			,			lowerCAmelCase_      :				DDPMScheduler			,			lowerCAmelCase_      :				Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]			,			lowerCAmelCase_      :				int = 350			,			)		->    Optional[Any]:
				super().__init__(
				    vae=lowerCAmelCase_			,			text_encoder=lowerCAmelCase_			,			tokenizer=lowerCAmelCase_			,			unet=lowerCAmelCase_			,			low_res_scheduler=lowerCAmelCase_			,			scheduler=lowerCAmelCase_			,			max_noise_level=lowerCAmelCase_			,			)





	def   _SCREAMING_SNAKE_CASE   (   self      :				str			,			lowerCAmelCase_      :				Union[str, Any]			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				Tuple			,			lowerCAmelCase_      :				Any			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				int			,			lowerCAmelCase_      :				Any			,			**lowerCAmelCase_      :				List[str]		)		->    Dict:
				torch.manual_seed(0		)
				UpperCAmelCase_  :				Optional[Any]      =      (
				    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_  :				Dict      =      add_overlap_rect(lowerCAmelCase_			,			lowerCAmelCase_			,			image.size		)
				UpperCAmelCase_  :				Tuple      =      image.crop(lowerCAmelCase_		)
				UpperCAmelCase_  :				Any      =      ((crop_rect[0] + ((crop_rect[2] - crop_rect[0]) / 2)) / image.size[0]) * tile.size[0]
				UpperCAmelCase_  :				int      =      translated_slice_x - (original_image_slice / 2)
				UpperCAmelCase_  :				Any      =      max(0			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      squeeze_tile(lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				Any      =      to_input.size
				UpperCAmelCase_  :				Optional[int]      =      to_input.resize((tile_size, tile_size)			,			Image.BICUBIC		)
				UpperCAmelCase_  :				Optional[Any]      =      super(lowerCAmelCase_			,			self		).__call__(image=lowerCAmelCase_			,			**lowerCAmelCase_		).images[0]
				UpperCAmelCase_  :				Union[str, Any]      =      upscaled_tile.resize((orig_input_size[0] * 4, orig_input_size[1] * 4)			,			Image.BICUBIC		)
				UpperCAmelCase_  :				List[Any]      =      unsqueeze_tile(lowerCAmelCase_			,			lowerCAmelCase_		)
				UpperCAmelCase_  :				List[str]      =      upscaled_tile.resize((tile.size[0] * 4, tile.size[1] * 4)			,			Image.BICUBIC		)
				UpperCAmelCase_  :				List[Any]      =      []
				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_  :				str      =      Image.fromarray(
				    make_transparency_mask(
				        (upscaled_tile.size[0], upscaled_tile.size[1])			,			tile_border * 4			,			remove_borders=lowerCAmelCase_		)			,			mode="L"			,			)
				final_image.paste(
				    lowerCAmelCase_			,			(crop_rect_with_overlap[0] * 4, crop_rect_with_overlap[1] * 4)			,			lowerCAmelCase_		)





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





def        snake_case   (     ):
			# Run a demo
			UpperCAmelCase_  :				List[str]      =      "stabilityai/stable-diffusion-x4-upscaler"
			UpperCAmelCase_  :				Optional[Any]      =      StableDiffusionTiledUpscalePipeline.from_pretrained(A__			,revision="fp16"			,torch_dtype=torch.floataa		)
			UpperCAmelCase_  :				Optional[int]      =      pipe.to("cuda"		)
			UpperCAmelCase_  :				Tuple      =      Image.open("../../docs/source/imgs/diffusers_library.jpg"		)

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

			UpperCAmelCase_  :				Union[str, Any]      =      pipe(image=A__			,prompt="Black font, white background, vector"			,noise_level=40			,callback=A__		)
			final_image.save("diffusers_library.jpg"		)


if __name__ == "__main__":
			main()



