






import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class     UpperCamelCase_ (__A       ,		unittest.TestCase  ):
	__magic_name__			   =       ConsistencyModelPipeline
	__magic_name__			   =       UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	__magic_name__			   =       UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

	# Override required_optional_params to remove num_images_per_prompt
	__magic_name__			   =       frozenset(
	    [
	        '''num_inference_steps''',
	        '''generator''',
	        '''latents''',
	        '''output_type''',
	        '''return_dict''',
	        '''callback''',
	        '''callback_steps''',
	    ]  )





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Any		)		->    Tuple:
				UpperCAmelCase_  :				int      =      UNetaDModel.from_pretrained(
				    "diffusers/consistency-models-test"			,			subfolder="test_unet"			,			)
				return unet





	@property
	def   _SCREAMING_SNAKE_CASE   (   self      :				Union[str, Any]		)		->    int:
				UpperCAmelCase_  :				Optional[Any]      =      UNetaDModel.from_pretrained(
				    "diffusers/consistency-models-test"			,			subfolder="test_unet_class_cond"			,			)
				return unet





	def   _SCREAMING_SNAKE_CASE   (   self      :				Dict			,			lowerCAmelCase_      :				int=False		)		->    List[Any]:
				if class_cond:
							UpperCAmelCase_  :				str      =      self.dummy_cond_unet
				else:
							UpperCAmelCase_  :				Optional[int]      =      self.dummy_uncond_unet

				# Default to CM multistep sampler
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=80.0			,			)

				UpperCAmelCase_  :				Union[str, Any]      =      {
				    "unet": unet,
				    "scheduler": scheduler,
				}

				return components





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]			,			lowerCAmelCase_      :				str			,			lowerCAmelCase_      :				str=0		)		->    Optional[Any]:
				if str(lowerCAmelCase_		).startswith("mps"		):
							UpperCAmelCase_  :				Tuple      =      torch.manual_seed(lowerCAmelCase_		)
				else:
							UpperCAmelCase_  :				Tuple      =      torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[Any]      =      {
				    "batch_size": 1,
				    "num_inference_steps": None,
				    "timesteps": [22, 0],
				    "generator": generator,
				    "output_type": "np",
				}

				return inputs





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Dict:
				UpperCAmelCase_  :				Union[str, Any]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				Tuple      =      self.get_dummy_components()
				UpperCAmelCase_  :				Tuple      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[int]      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				int      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				List[str]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				List[Any]      =      np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[str]		)		->    Any:
				UpperCAmelCase_  :				List[str]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				List[str]      =      self.get_dummy_components(class_cond=lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				List[str]      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      0
				UpperCAmelCase_  :				str      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				List[Any]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				Optional[Any]      =      np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Optional[Any]:
				UpperCAmelCase_  :				List[str]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				List[Any]      =      self.get_dummy_components()
				UpperCAmelCase_  :				List[str]      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Optional[Any]      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[Any]      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				Tuple      =      1
				UpperCAmelCase_  :				Dict      =      None
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				Any      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				List[Any]      =      np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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





	def   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]		)		->    Union[str, Any]:
				UpperCAmelCase_  :				Optional[Any]      =      "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_  :				Any      =      self.get_dummy_components(class_cond=lowerCAmelCase_		)
				UpperCAmelCase_  :				List[Any]      =      ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_  :				Dict      =      pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				str      =      self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_  :				List[Any]      =      1
				UpperCAmelCase_  :				List[str]      =      None
				UpperCAmelCase_  :				Union[str, Any]      =      0
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_  :				Optional[int]      =      image[0, -3:, -3:, -1]
				UpperCAmelCase_  :				Any      =      np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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







@slow
@require_torch_gpu
class     UpperCamelCase_ (unittest.TestCase  ):





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    Union[str, Any]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def   _SCREAMING_SNAKE_CASE   (   self      :				int			,			lowerCAmelCase_      :				Dict=0			,			lowerCAmelCase_      :				Union[str, Any]=False			,			lowerCAmelCase_      :				Any="cpu"			,			lowerCAmelCase_      :				int=torch.floataa			,			lowerCAmelCase_      :				Optional[int]=(1, 3, 64, 64)		)		->    str:
				UpperCAmelCase_  :				Optional[Any]      =      torch.manual_seed(lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      {
				    "num_inference_steps": None,
				    "timesteps": [22, 0],
				    "class_labels": 0,
				    "generator": generator,
				    "output_type": "np",
				}

				if get_fixed_latents:
							UpperCAmelCase_  :				Dict      =      self.get_fixed_latents(seed=lowerCAmelCase_			,			device=lowerCAmelCase_			,			dtype=lowerCAmelCase_			,			shape=lowerCAmelCase_		)
							UpperCAmelCase_  :				List[Any]      =      latents

				return inputs





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]			,			lowerCAmelCase_      :				Optional[Any]=0			,			lowerCAmelCase_      :				List[Any]="cpu"			,			lowerCAmelCase_      :				Optional[Any]=torch.floataa			,			lowerCAmelCase_      :				Optional[Any]=(1, 3, 64, 64)		)		->    Tuple:
				if type(lowerCAmelCase_		) == str:
							UpperCAmelCase_  :				List[str]      =      torch.device(lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)
				UpperCAmelCase_  :				Dict      =      randn_tensor(lowerCAmelCase_			,			generator=lowerCAmelCase_			,			device=lowerCAmelCase_			,			dtype=lowerCAmelCase_		)
				return latents





	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[int]		)		->    List[Any]:
				UpperCAmelCase_  :				Optional[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Dict      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=80.0			,			)
				UpperCAmelCase_  :				Optional[int]      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Optional[int]      =      self.get_inputs()
				UpperCAmelCase_  :				List[Any]      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_  :				List[str]      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				str      =      np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4]		)

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





	def   _SCREAMING_SNAKE_CASE   (   self      :				Tuple		)		->    int:
				UpperCAmelCase_  :				List[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=80.0			,			)
				UpperCAmelCase_  :				Any      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Dict      =      self.get_inputs()
				UpperCAmelCase_  :				List[Any]      =      1
				UpperCAmelCase_  :				Union[str, Any]      =      None
				UpperCAmelCase_  :				Dict      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_  :				int      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				str      =      np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7]		)

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





	@require_torch_a
	def   _SCREAMING_SNAKE_CASE   (   self      :				List[Any]		)		->    Union[str, Any]:
				UpperCAmelCase_  :				Dict      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				int      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=80.0			,			)
				UpperCAmelCase_  :				Any      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_			,			torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Tuple      =      self.get_inputs(get_fixed_latents=lowerCAmelCase_			,			device=lowerCAmelCase_		)
				# Ensure usage of flash attention in torch 2.0
				with sdp_kernel(enable_flash=lowerCAmelCase_			,			enable_math=lowerCAmelCase_			,			enable_mem_efficient=lowerCAmelCase_		):
							UpperCAmelCase_  :				Optional[int]      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_  :				Tuple      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				Optional[int]      =      np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3]		)

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





	@require_torch_a
	def   _SCREAMING_SNAKE_CASE   (   self      :				Optional[Any]		)		->    Optional[Any]:
				UpperCAmelCase_  :				List[Any]      =      UNetaDModel.from_pretrained("diffusers/consistency_models"			,			subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_  :				Optional[int]      =      CMStochasticIterativeScheduler(
				    num_train_timesteps=40			,			sigma_min=0.0_0_2			,			sigma_max=80.0			,			)
				UpperCAmelCase_  :				List[Any]      =      ConsistencyModelPipeline(unet=lowerCAmelCase_			,			scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_			,			torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_  :				Any      =      self.get_inputs(get_fixed_latents=lowerCAmelCase_			,			device=lowerCAmelCase_		)
				UpperCAmelCase_  :				Union[str, Any]      =      1
				UpperCAmelCase_  :				str      =      None
				# Ensure usage of flash attention in torch 2.0
				with sdp_kernel(enable_flash=lowerCAmelCase_			,			enable_math=lowerCAmelCase_			,			enable_mem_efficient=lowerCAmelCase_		):
							UpperCAmelCase_  :				List[str]      =      pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_  :				Any      =      image[0, -3:, -3:, -1]

				UpperCAmelCase_  :				Tuple      =      np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5]		)

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











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



