

"""simple docstring"""


import logging
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

from filelock import FileLock

from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available


__A						   =			logging.getLogger(__name__)





@dataclass
class     snake_case	:
				SCREAMING_SNAKE_CASE_   :      str
				SCREAMING_SNAKE_CASE_   :      List[str]
				SCREAMING_SNAKE_CASE_   :      Optional[List[str]]





@dataclass
class     snake_case	:
				SCREAMING_SNAKE_CASE_   :      List[int]
				SCREAMING_SNAKE_CASE_   :      List[int]
				SCREAMING_SNAKE_CASE_   :      Optional[List[int]]         =  None
				SCREAMING_SNAKE_CASE_   :      Optional[List[int]]         =  None





class     snake_case	(				__snake_case  ):
				SCREAMING_SNAKE_CASE_   :      Optional[Any]         =  """train"""
				SCREAMING_SNAKE_CASE_   :      Tuple         =  """dev"""
				SCREAMING_SNAKE_CASE_   :      Optional[int]         =  """test"""





class     snake_case	:

				@staticmethod
				def 							lowercase_     (	UpperCamelCase__							:				int     ,  UpperCamelCase__							:				Union[Split, str])->   List[InputExample]:



						'''simple docstring'''


						raise NotImplementedError

				@staticmethod
				def 							lowercase_     (	UpperCamelCase__							:				str)->   List[str]:



						'''simple docstring'''


						raise NotImplementedError




				@staticmethod
				def 							lowercase_     (	UpperCamelCase__							:				List[InputExample]     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				PreTrainedTokenizer     ,  UpperCamelCase__							:				Optional[Any]=False     ,  UpperCamelCase__							:				int="[CLS]"     ,  UpperCamelCase__							:				str=1     ,  UpperCamelCase__							:				Optional[Any]="[SEP]"     ,  UpperCamelCase__							:				Optional[int]=False     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				Union[str, Any]=0     ,  UpperCamelCase__							:				List[Any]=0     ,  UpperCamelCase__							:				Dict=-1_0_0     ,  UpperCamelCase__							:				Optional[int]=0     ,  UpperCamelCase__							:				Union[str, Any]=True     ,  )->   List[InputFeatures]:



						'''simple docstring'''


						__lowerCAmelCase:       Tuple          =			{label: i for i, label in enumerate(UpperCamelCase__)}

						__lowerCAmelCase:       Dict          =			[]
						for ex_index, example in enumerate(UpperCamelCase__):
								if ex_index % 1_0_0_0_0 == 0:
										logger.info("Writing example %d of %d"     ,  UpperCamelCase__     ,  len(UpperCamelCase__))

								__lowerCAmelCase:       List[Any]          =			[]
								__lowerCAmelCase:       str          =			[]
								for word, label in zip(example.words     ,  example.labels):
										__lowerCAmelCase:       Optional[int]          =			tokenizer.tokenize(UpperCamelCase__)

										# bert-base-multilingual-cased sometimes output "nothing ([]) when calling tokenize with just a space.
										if len(UpperCamelCase__) > 0:
												tokens.extend(UpperCamelCase__)
												# Use the real label id for the first token of the word, and padding ids for the remaining tokens
												label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(UpperCamelCase__) - 1))

            # Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
								__lowerCAmelCase:       str          =			tokenizer.num_special_tokens_to_add()
								if len(UpperCamelCase__) > max_seq_length - special_tokens_count:
										__lowerCAmelCase:       str          =			tokens[: (max_seq_length - special_tokens_count)]
										__lowerCAmelCase:       Optional[Any]          =			label_ids[: (max_seq_length - special_tokens_count)]

								# The convention in BERT is:
								# (a) For sequence pairs:
								#  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
								#  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1
								# (b) For single sequences:
								#  tokens:   [CLS] the dog is hairy . [SEP]
								#  type_ids:   0   0   0   0  0     0   0
								#
								# Where "type_ids" are used to indicate whether this is the first
								# sequence or the second sequence. The embedding vectors for `type=0` and
								# `type=1` were learned during pre-training and are added to the wordpiece
								# embedding vector (and position vector). This is not *strictly* necessary
								# since the [SEP] token unambiguously separates the sequences, but it makes
								# it easier for the model to learn the concept of sequences.
								#
								# For classification tasks, the first vector (corresponding to [CLS]) is
								# used as the "sentence vector". Note that this only makes sense because
								# the entire model is fine-tuned.
								tokens += [sep_token]
								label_ids += [pad_token_label_id]
								if sep_token_extra:
										# roberta uses an extra separator b/w pairs of sentences
										tokens += [sep_token]
										label_ids += [pad_token_label_id]
								__lowerCAmelCase:       Union[str, Any]          =			[sequence_a_segment_id] * len(UpperCamelCase__)

								if cls_token_at_end:
										tokens += [cls_token]
										label_ids += [pad_token_label_id]
										segment_ids += [cls_token_segment_id]
								else:
										__lowerCAmelCase:       List[Any]          =			[cls_token] + tokens
										__lowerCAmelCase:       Optional[Any]          =			[pad_token_label_id] + label_ids
										__lowerCAmelCase:       List[Any]          =			[cls_token_segment_id] + segment_ids

								__lowerCAmelCase:       Union[str, Any]          =			tokenizer.convert_tokens_to_ids(UpperCamelCase__)

								# The mask has 1 for real tokens and 0 for padding tokens. Only real
								# tokens are attended to.
								__lowerCAmelCase:       Tuple          =			[1 if mask_padding_with_zero else 0] * len(UpperCamelCase__)

								# Zero-pad up to the sequence length.
								__lowerCAmelCase:       Dict          =			max_seq_length - len(UpperCamelCase__)
								if pad_on_left:
										__lowerCAmelCase:       str          =			([pad_token] * padding_length) + input_ids
										__lowerCAmelCase:       str          =			([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
										__lowerCAmelCase:       List[str]          =			([pad_token_segment_id] * padding_length) + segment_ids
										__lowerCAmelCase:       Tuple          =			([pad_token_label_id] * padding_length) + label_ids
								else:
										input_ids += [pad_token] * padding_length
										input_mask += [0 if mask_padding_with_zero else 1] * padding_length
										segment_ids += [pad_token_segment_id] * padding_length
										label_ids += [pad_token_label_id] * padding_length

								assert len(UpperCamelCase__) == max_seq_length
								assert len(UpperCamelCase__) == max_seq_length
								assert len(UpperCamelCase__) == max_seq_length
								assert len(UpperCamelCase__) == max_seq_length

								if ex_index < 5:
										logger.info("*** Example ***")
										logger.info("guid: %s"     ,  example.guid)
										logger.info("tokens: %s"     ,  " ".join([str(UpperCamelCase__) for x in tokens]))
										logger.info("input_ids: %s"     ,  " ".join([str(UpperCamelCase__) for x in input_ids]))
										logger.info("input_mask: %s"     ,  " ".join([str(UpperCamelCase__) for x in input_mask]))
										logger.info("segment_ids: %s"     ,  " ".join([str(UpperCamelCase__) for x in segment_ids]))
										logger.info("label_ids: %s"     ,  " ".join([str(UpperCamelCase__) for x in label_ids]))

								if "token_type_ids" not in tokenizer.model_input_names:
										__lowerCAmelCase:       Optional[Any]          =			None

								features.append(
								    InputFeatures(
								        input_ids=UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  token_type_ids=UpperCamelCase__     ,  label_ids=UpperCamelCase__))
						return features


if is_torch_available():
			import torch
			from torch import nn
			from torch.utils.data import Dataset





			class     snake_case	(				__snake_case  ):
							SCREAMING_SNAKE_CASE_   :      List[InputFeatures]
							SCREAMING_SNAKE_CASE_   :      int         =  nn.CrossEntropyLoss().ignore_index

							def __init__(	self							:				Tuple     ,  UpperCamelCase__							:				TokenClassificationTask     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				PreTrainedTokenizer     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				Optional[int] = None     ,  UpperCamelCase__							:				int=False     ,  UpperCamelCase__							:				Split = Split.train     ,  )->   Optional[Any]:



									'''simple docstring'''


									__lowerCAmelCase:       Dict          =			os.path.join(
									    UpperCamelCase__     ,  "cached_{}_{}_{}".format(mode.value     ,  tokenizer.__class__.__name__     ,  str(UpperCamelCase__))     ,  )

									# Make sure only the first process in distributed training processes the dataset,
									# and the others will use the cache.
									__lowerCAmelCase:       List[Any]          =			cached_features_file + ".lock"
									with FileLock(UpperCamelCase__):
											if os.path.exists(UpperCamelCase__) and not overwrite_cache:
													logger.info(f"Loading features from cached file {cached_features_file}")
													__lowerCAmelCase:       List[Any]          =			torch.load(UpperCamelCase__)
											else:
													logger.info(f"Creating features from dataset file at {data_dir}")
													__lowerCAmelCase:       Optional[Any]          =			token_classification_task.read_examples_from_file(UpperCamelCase__     ,  UpperCamelCase__)
													# TODO clean up all this to leverage built-in features of tokenizers
													__lowerCAmelCase:       List[Any]          =			token_classification_task.convert_examples_to_features(
													    UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  cls_token_at_end=bool(model_type in ["xlnet"])     ,  cls_token=tokenizer.cls_token     ,  cls_token_segment_id=2 if model_type in ["xlnet"] else 0     ,  sep_token=tokenizer.sep_token     ,  sep_token_extra=UpperCamelCase__     ,  pad_on_left=bool(tokenizer.padding_side == "left")     ,  pad_token=tokenizer.pad_token_id     ,  pad_token_segment_id=tokenizer.pad_token_type_id     ,  pad_token_label_id=self.pad_token_label_id     ,  )
													logger.info(f"Saving features into cached file {cached_features_file}")
													torch.save(self.features     ,  UpperCamelCase__)

							def __len__(	self							:				str)->   Optional[int]:



									'''simple docstring'''


									return len(self.features)




							def __getitem__(	self							:				Any     ,  UpperCamelCase__							:				Optional[int])->   InputFeatures:



									'''simple docstring'''


									return self.features[i]


if is_tf_available():
			import tensorflow as tf





			class     snake_case	:
							SCREAMING_SNAKE_CASE_   :      List[InputFeatures]
							SCREAMING_SNAKE_CASE_   :      int         =  -1_00

							def __init__(	self							:				Dict     ,  UpperCamelCase__							:				TokenClassificationTask     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				PreTrainedTokenizer     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				Optional[int] = None     ,  UpperCamelCase__							:				List[Any]=False     ,  UpperCamelCase__							:				Split = Split.train     ,  )->   Optional[int]:



									'''simple docstring'''


									__lowerCAmelCase:       Optional[Any]          =			token_classification_task.read_examples_from_file(UpperCamelCase__     ,  UpperCamelCase__)
									# TODO clean up all this to leverage built-in features of tokenizers
									__lowerCAmelCase:       Dict          =			token_classification_task.convert_examples_to_features(
									    UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  cls_token_at_end=bool(model_type in ["xlnet"])     ,  cls_token=tokenizer.cls_token     ,  cls_token_segment_id=2 if model_type in ["xlnet"] else 0     ,  sep_token=tokenizer.sep_token     ,  sep_token_extra=UpperCamelCase__     ,  pad_on_left=bool(tokenizer.padding_side == "left")     ,  pad_token=tokenizer.pad_token_id     ,  pad_token_segment_id=tokenizer.pad_token_type_id     ,  pad_token_label_id=self.pad_token_label_id     ,  )

									def gen():
											for ex in self.features:
													if ex.token_type_ids is None:
															yield (
															    {"input_ids": ex.input_ids, "attention_mask": ex.attention_mask},
															    ex.label_ids,
															)
													else:
															yield (
															    {
															        "input_ids": ex.input_ids,
															        "attention_mask": ex.attention_mask,
															        "token_type_ids": ex.token_type_ids,
															    },
															    ex.label_ids,
															)

									if "token_type_ids" not in tokenizer.model_input_names:
											__lowerCAmelCase:       Optional[Any]          =			tf.data.Dataset.from_generator(
											    UpperCamelCase__     ,  ({"input_ids": tf.intaa, "attention_mask": tf.intaa}, tf.intaa)     ,  (
											        {"input_ids": tf.TensorShape([None]), "attention_mask": tf.TensorShape([None])},
											        tf.TensorShape([None]),
											    )     ,  )
									else:
											__lowerCAmelCase:       Optional[Any]          =			tf.data.Dataset.from_generator(
											    UpperCamelCase__     ,  ({"input_ids": tf.intaa, "attention_mask": tf.intaa, "token_type_ids": tf.intaa}, tf.intaa)     ,  (
											        {
											            "input_ids": tf.TensorShape([None]),
											            "attention_mask": tf.TensorShape([None]),
											            "token_type_ids": tf.TensorShape([None]),
											        },
											        tf.TensorShape([None]),
											    )     ,  )

							def 							lowercase_     (	self							:				List[Any])->   Union[str, Any]:



									'''simple docstring'''


									__lowerCAmelCase:       Dict          =			self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features)))

									return self.dataset

							def __len__(	self							:				Tuple)->   Optional[Any]:



									'''simple docstring'''


									return len(self.features)




							def __getitem__(	self							:				str     ,  UpperCamelCase__							:				Dict)->   InputFeatures:



									'''simple docstring'''


									return self.features[i]





"""simple docstring"""


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class     snake_case	(				unittest.TestCase  ):

				def 							lowercase_     (	self							:				Optional[int])->   int:



						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

				@property
				def 							lowercase_     (	self							:				int)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			1
						__lowerCAmelCase:       Union[str, Any]          =			3
						__lowerCAmelCase:       Union[str, Any]          =			(3_2, 3_2)

						__lowerCAmelCase:       Union[str, Any]          =			floats_tensor((batch_size, num_channels) + sizes     ,  rng=random.Random(0)).to(UpperCamelCase__)
						return image

				@property
				def 							lowercase_     (	self							:				Tuple)->   str:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			UNetaDConditionModel(
						    block_out_channels=(3_2, 3_2, 6_4)     ,  layers_per_block=2     ,  sample_size=3_2     ,  in_channels=7     ,  out_channels=4     ,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")     ,  up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")     ,  cross_attention_dim=3_2     ,  attention_head_dim=8     ,  use_linear_projection=UpperCamelCase__     ,  only_cross_attention=(True, True, False)     ,  num_class_embeds=1_0_0     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Tuple          =			AutoencoderKL(
						    block_out_channels=[3_2, 3_2, 6_4]     ,  in_channels=3     ,  out_channels=3     ,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]     ,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]     ,  latent_channels=4     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       List[str]          =			CLIPTextConfig(
						    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=3_2     ,  intermediate_size=3_7     ,  layer_norm_eps=1e-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=1_0_0_0     ,  hidden_act="gelu"     ,  projection_dim=5_1_2     ,  )
						return CLIPTextModel(UpperCamelCase__)

				def 							lowercase_     (	self							:				List[str])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       Tuple          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       int          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       List[str]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Tuple          =			self.dummy_vae
						__lowerCAmelCase:       Optional[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Any          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       Tuple          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       List[Any]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Tuple          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Any          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[int]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )

						__lowerCAmelCase:       List[str]          =			output.images

						__lowerCAmelCase:       Union[str, Any]          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       List[str]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  return_dict=UpperCamelCase__     ,  )[0]

						__lowerCAmelCase:       int          =			image[0, -3:, -3:, -1]
						__lowerCAmelCase:       Dict          =			image_from_tuple[0, -3:, -3:, -1]

						__lowerCAmelCase:       Dict          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)
						__lowerCAmelCase:       List[Any]          =			np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661])

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

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



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       Dict          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       List[str]          =			DDPMScheduler()
						__lowerCAmelCase:       Union[str, Any]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Optional[int]          =			self.dummy_vae
						__lowerCAmelCase:       List[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Optional[int]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       str          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Optional[int]          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       List[str]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       List[Any]          =			sd_pipe(
						    2 * [prompt]     ,  image=2 * [low_res_image]     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2

						__lowerCAmelCase:       Dict          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_images_per_prompt=2     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2




				@unittest.skipIf(torch_device != "cuda"     ,  "This test requires a GPU")
				def 							lowercase_     (	self							:				Tuple)->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       int          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Dict          =			self.dummy_vae
						__lowerCAmelCase:       int          =			self.dummy_text_encoder
						__lowerCAmelCase:       List[Any]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       Optional[int]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# put models in fp16, except vae as it overflows in fp16
						__lowerCAmelCase:       List[Any]          =			unet.half()
						__lowerCAmelCase:       List[str]          =			text_encoder.half()

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       List[Any]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       str          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Optional[Any]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=2     ,  output_type="np"     ,  ).images

						__lowerCAmelCase:       Optional[Any]          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)





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

				def 							lowercase_     (	self							:				Tuple)->   Tuple:



						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

				def 							lowercase_     (	self							:				List[Any])->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Optional[Any]          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat.npy")

						__lowerCAmelCase:       str          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline.from_pretrained(UpperCamelCase__)
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       Tuple          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       int          =			torch.manual_seed(0)
						__lowerCAmelCase:       List[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Dict          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 1e-3

				def 							lowercase_     (	self							:				Optional[int])->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Tuple          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat_fp16.npy")

						__lowerCAmelCase:       Optional[Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Tuple          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       str          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       List[str]          =			torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Union[str, Any]          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 5e-1




				def 							lowercase_     (	self							:				Optional[int])->   Dict:



						'''simple docstring'''


						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

						__lowerCAmelCase:       Tuple          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")

						__lowerCAmelCase:       Union[str, Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Any          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing(1)
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       int          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       Dict          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=5     ,  output_type="np"     ,  )

						__lowerCAmelCase:       Optional[int]          =			torch.cuda.max_memory_allocated()
						# make sure that less than 2.9 GB is allocated
						assert mem_bytes < 2.9 * 1_0**9


