


import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch
			from torch import nn

			from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
			from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			from PIL import Image

class 							a						:




							"""simple docstring"""


							def __init__(				self			:    Optional[Any]       ,      lowerCamelCase			:    Tuple       ,      lowerCamelCase			:    str=3       ,      lowerCamelCase			:    Tuple=32       ,      lowerCamelCase			:    Tuple=3       ,      lowerCamelCase			:    Any=10       ,      lowerCamelCase			:    int=[8, 16, 32, 64]       ,      lowerCamelCase			:    Optional[int]=[1, 1, 2, 1]       ,      lowerCamelCase			:    int=True       ,      lowerCamelCase			:    Optional[Any]=True       ,      lowerCamelCase			:    Any="relu"       ,      lowerCamelCase			:    Tuple=3       ,      lowerCamelCase			:    str=None       ,      lowerCamelCase			:    Optional[Any]=["stage2", "stage3", "stage4"]       ,      lowerCamelCase			:    Union[str, Any]=[2, 3, 4]       ,      lowerCamelCase			:    str=1       ,      )      ->       Union[str, Any]:
														__snake_case						:     Optional[int]									=  parent
														__snake_case						:     Optional[int]									=  batch_size
														__snake_case						:     Any									=  image_size
														__snake_case						:     Optional[int]									=  num_channels
														__snake_case						:     int									=  embeddings_size
														__snake_case						:     Tuple									=  hidden_sizes
														__snake_case						:     Optional[int]									=  depths
														__snake_case						:     Union[str, Any]									=  is_training
														__snake_case						:     Union[str, Any]									=  use_labels
														__snake_case						:     Optional[int]									=  hidden_act
														__snake_case						:     List[str]									=  num_labels
														__snake_case						:     Union[str, Any]									=  scope
														__snake_case						:     List[Any]									=  len(lowerCamelCase							)
														__snake_case						:     int									=  out_features
														__snake_case						:     Tuple									=  out_indices
														__snake_case						:     int									=  num_groups

							def        __snake_case   (				self			:    Optional[Any]							)      ->       Optional[int]:
														__snake_case						:     Union[str, Any]									=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

														__snake_case						:     Tuple									=  None
														if self.use_labels:
																					__snake_case						:     List[Any]									=  ids_tensor([self.batch_size]       ,      self.num_labels							)

														__snake_case						:     str									=  self.get_config()

														return config, pixel_values, labels

							def        __snake_case   (				self			:    int							)      ->       List[str]:
														return BitConfig(
														    num_channels=self.num_channels       ,      embeddings_size=self.embeddings_size       ,      hidden_sizes=self.hidden_sizes       ,      depths=self.depths       ,      hidden_act=self.hidden_act       ,      num_labels=self.num_labels       ,      out_features=self.out_features       ,      out_indices=self.out_indices       ,      num_groups=self.num_groups       ,      )

							def        __snake_case   (				self			:    Optional[Any]       ,      lowerCamelCase			:    Dict       ,      lowerCamelCase			:    str       ,      lowerCamelCase			:    List[Any]							)      ->       Any:
														__snake_case						:     int									=  BitModel(config=lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.eval()
														__snake_case						:     Tuple									=  model(lowerCamelCase							)
														self.parent.assertEqual(
														    result.last_hidden_state.shape       ,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,      )

							def        __snake_case   (				self			:    List[str]       ,      lowerCamelCase			:    Optional[int]       ,      lowerCamelCase			:    Any       ,      lowerCamelCase			:    Optional[int]							)      ->       int:
														__snake_case						:     List[str]									=  self.num_labels
														__snake_case						:     Union[str, Any]									=  BitForImageClassification(lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.eval()
														__snake_case						:     Any									=  model(lowerCamelCase       ,      labels=lowerCamelCase							)
														self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_labels)							)

							def        __snake_case   (				self			:    Dict       ,      lowerCamelCase			:    str       ,      lowerCamelCase			:    Optional[Any]       ,      lowerCamelCase			:    int							)      ->       str:
														__snake_case						:     List[Any]									=  BitBackbone(config=lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.eval()
														__snake_case						:     Union[str, Any]									=  model(lowerCamelCase							)

														# verify feature maps
														self.parent.assertEqual(len(result.feature_maps							)       ,      len(config.out_features							)							)
														self.parent.assertListEqual(list(result.feature_maps[0].shape							)       ,      [self.batch_size, self.hidden_sizes[1], 4, 4]							)

														# verify channels
														self.parent.assertEqual(len(model.channels							)       ,      len(config.out_features							)							)
														self.parent.assertListEqual(model.channels       ,      config.hidden_sizes[1:]							)

														# verify backbone works with out_features=None
														__snake_case						:     Dict									=  None
														__snake_case						:     str									=  BitBackbone(config=lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.eval()
														__snake_case						:     int									=  model(lowerCamelCase							)

														# verify feature maps
														self.parent.assertEqual(len(result.feature_maps							)       ,      1							)
														self.parent.assertListEqual(list(result.feature_maps[0].shape							)       ,      [self.batch_size, self.hidden_sizes[-1], 1, 1]							)

														# verify channels
														self.parent.assertEqual(len(model.channels							)       ,      1							)
														self.parent.assertListEqual(model.channels       ,      [config.hidden_sizes[-1]]							)






							def        __snake_case   (				self			:    Union[str, Any]							)      ->       List[Any]:
														__snake_case						:     Tuple									=  self.prepare_config_and_inputs()
														__snake_case					,     __snake_case					,     __snake_case						:     Optional[Any]									=  config_and_inputs
														__snake_case						:     int									=  {"pixel_values": pixel_values}
														return config, inputs_dict

@require_torch
class 							a						(_lowerCAmelCase	,							_lowerCAmelCase	,							unittest.TestCase ):




							"""simple docstring"""


							__UpperCAmelCase       :				str	    =    (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
							__UpperCAmelCase       :				List[str]	    =    (
							    {"feature-extraction": BitModel, "image-classification": BitForImageClassification}
							    if is_torch_available()
							    else {}
							)

							__UpperCAmelCase       :				int	    =    False
							__UpperCAmelCase       :				Tuple	    =    False
							__UpperCAmelCase       :				Any	    =    False
							__UpperCAmelCase       :				List[str]	    =    False
							__UpperCAmelCase       :				Optional[int]	    =    False

							def        __snake_case   (				self			:    int							)      ->       Optional[Any]:
														__snake_case						:     Any									=  BitModelTester(self							)
														__snake_case						:     Optional[Any]									=  ConfigTester(self       ,      config_class=lowerCamelCase       ,      has_text_modality=lowerCamelCase							)

							def        __snake_case   (				self			:    str							)      ->       int:
														self.create_and_test_config_common_properties()
														self.config_tester.create_and_test_config_to_json_string()
														self.config_tester.create_and_test_config_to_json_file()
														self.config_tester.create_and_test_config_from_and_save_pretrained()
														self.config_tester.create_and_test_config_with_num_labels()
														self.config_tester.check_config_can_be_init_without_params()
														self.config_tester.check_config_arguments_init()

							def        __snake_case   (				self			:    Optional[int]							)      ->       int:
														return

							@unittest.skip(reason="Bit does not output attentions"							)
							def        __snake_case   (				self			:    int							)      ->       Union[str, Any]:
														pass

							@unittest.skip(reason="Bit does not use inputs_embeds"							)
							def        __snake_case   (				self			:    int							)      ->       Dict:
														pass

							@unittest.skip(reason="Bit does not support input and output embeddings"							)
							def        __snake_case   (				self			:    str							)      ->       Optional[Any]:
														pass

							def        __snake_case   (				self			:    List[Any]							)      ->       List[Any]:
														__snake_case					,     __snake_case						:     List[Any]									=  self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					__snake_case						:     Tuple									=  model_class(lowerCamelCase							)
																					__snake_case						:     Optional[Any]									=  inspect.signature(model.forward							)
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					__snake_case						:     Dict									=  [*signature.parameters.keys()]

																					__snake_case						:     int									=  ["pixel_values"]
																					self.assertListEqual(arg_names[:1]       ,      lowerCamelCase							)

							def        __snake_case   (				self			:    str							)      ->       List[str]:
														__snake_case						:     Any									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCamelCase							)

							def        __snake_case   (				self			:    str							)      ->       Any:
														__snake_case						:     List[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_backbone(*lowerCamelCase							)

							def        __snake_case   (				self			:    List[Any]							)      ->       Tuple:
														__snake_case					,     __snake_case						:     Optional[int]									=  self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					__snake_case						:     Union[str, Any]									=  model_class(config=lowerCamelCase							)
																					for name, module in model.named_modules():
																												if isinstance(lowerCamelCase       ,      (nn.BatchNormad, nn.GroupNorm)							):
																																			self.assertTrue(
																																			    torch.all(module.weight == 1							)       ,      msg=F'Parameter {name} of model {model_class} seems not properly initialized'       ,      )
																																			self.assertTrue(
																																			    torch.all(module.bias == 0							)       ,      msg=F'Parameter {name} of model {model_class} seems not properly initialized'       ,      )

							def        __snake_case   (				self			:    Dict							)      ->       List[Any]:
														def check_hidden_states_output(lowerCamelCase			:    Dict       ,      lowerCamelCase			:    Dict       ,      lowerCamelCase			:    Any							):
																					__snake_case						:     Union[str, Any]									=  model_class(lowerCamelCase							)
																					model.to(lowerCamelCase							)
																					model.eval()

																					with torch.no_grad():
																												__snake_case						:     List[str]									=  model(**self._prepare_for_class(lowerCamelCase       ,      lowerCamelCase							)							)

																					__snake_case						:     Union[str, Any]									=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																					__snake_case						:     List[Any]									=  self.model_tester.num_stages
																					self.assertEqual(len(lowerCamelCase							)       ,      expected_num_stages + 1							)

																					# Bit's feature maps are of shape (batch_size, num_channels, height, width)
																					self.assertListEqual(
																					    list(hidden_states[0].shape[-2:]							)       ,      [self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,      )

														__snake_case					,     __snake_case						:     Tuple									=  self.model_tester.prepare_config_and_inputs_for_common()
														__snake_case						:     Tuple									=  ["preactivation", "bottleneck"]
														for model_class in self.all_model_classes:
																					for layer_type in layers_type:
																												__snake_case						:     Optional[Any]									=  layer_type
																												__snake_case						:     str									=  True
																												check_hidden_states_output(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase							)

																												# check that output_hidden_states also work using config
																												del inputs_dict["output_hidden_states"]
																												__snake_case						:     str									=  True

																												check_hidden_states_output(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase							)

							@unittest.skip(reason="Bit does not use feedforward chunking"							)
							def        __snake_case   (				self			:    str							)      ->       Union[str, Any]:
														pass

							def        __snake_case   (				self			:    Dict							)      ->       str:
														__snake_case						:     int									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_image_classification(*lowerCamelCase							)






							@slow
							def        __snake_case   (				self			:    List[Any]							)      ->       List[Any]:
														for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					__snake_case						:     Tuple									=  BitModel.from_pretrained(lowerCamelCase							)
																					self.assertIsNotNone(lowerCamelCase							)







def 							lowerCAmelCase_	(   ):
							__snake_case						:     Dict									=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
							return image

@require_torch
@require_vision
class 							a						(unittest.TestCase ):




							"""simple docstring"""


							@cached_property
							def        __snake_case   (				self			:    Any							)      ->       Tuple:
														return (
														    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]							) if is_vision_available() else None
														)






							@slow
							def        __snake_case   (				self			:    List[Any]							)      ->       List[str]:
														__snake_case						:     str									=  BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]							).to(lowerCamelCase							)

														__snake_case						:     str									=  self.default_image_processor
														__snake_case						:     Dict									=  prepare_img()
														__snake_case						:     Tuple									=  image_processor(images=lowerCamelCase       ,      return_tensors="pt"							).to(lowerCamelCase							)

														# forward pass
														with torch.no_grad():
																					__snake_case						:     Optional[int]									=  model(**lowerCamelCase							)

														# verify the logits
														__snake_case						:     Tuple									=  torch.Size((1, 1000)							)
														self.assertEqual(outputs.logits.shape       ,      lowerCamelCase							)

														__snake_case						:     Dict									=  torch.tensor([[-0.65_26, -0.52_63, -1.43_98]]							).to(lowerCamelCase							)

														self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,      lowerCamelCase       ,      atol=1E-4							)							)

@require_torch
class 							a						(_lowerCAmelCase	,							unittest.TestCase ):




							"""simple docstring"""


							__UpperCAmelCase       :				str	    =    (BitBackbone,) if is_torch_available() else ()
							__UpperCAmelCase       :				str	    =    BitConfig

							__UpperCAmelCase       :				List[Any]	    =    False

							def        __snake_case   (				self			:    Union[str, Any]							)      ->       Optional[int]:
														__snake_case						:     Optional[int]									=  BitModelTester(self							)






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  lowerCamelCase__	(    unittest.TestCase):
     """simple docstring"""



     def     _a   (self     ):






           '''simple docstring'''



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







     @property
     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		1
           lowerCamelCase             =		3
           lowerCamelCase             =		(32, 32)

           lowerCamelCase             =		floats_tensor((batch_size, num_channels) + sizes					,      rng=random.Random(0     )     ).to(__a     )
           return image







     @property
     def     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		UNetaDConditionModel(
               block_out_channels=(32, 32, 64)					,      layers_per_block=2					,      sample_size=32					,      in_channels=7					,      out_channels=4					,      down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")					,      up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")					,      cross_attention_dim=32					,      attention_head_dim=8					,      use_linear_projection=__a					,      only_cross_attention=(True, True, False)					,      num_class_embeds=1_00					,      )
           return model







     @property
     def     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		AutoencoderKL(
               block_out_channels=[32, 32, 64]					,      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     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		CLIPTextConfig(
               bos_token_id=0					,      eos_token_id=2					,      hidden_size=32					,      intermediate_size=37					,      layer_norm_eps=1E-05					,      num_attention_heads=4					,      num_hidden_layers=5					,      pad_token_id=1					,      vocab_size=10_00					,      hidden_act="gelu"					,      projection_dim=5_12					,      )
           return CLIPTextModel(__a     )







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

           lowerCamelCase             =		"A painting of a squirrel eating a burger"
           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      )

           lowerCamelCase             =		output.images

           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      return_dict=__a					,      )[0]

           lowerCamelCase             =		image[0, -3:, -3:, -1]
           lowerCamelCase             =		image_from_tuple[0, -3:, -3:, -1]

           lowerCamelCase             =		low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)
           lowerCamelCase             =		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     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

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

           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      num_images_per_prompt=2					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      )
           lowerCamelCase             =		output.images
           assert image.shape[0] == 2







     @unittest.skipIf(torch_device != "cuda"					,      "This test requires a GPU"     )
     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # put models in fp16, except vae as it overflows in fp16
           lowerCamelCase             =		unet.half()
           lowerCamelCase             =		text_encoder.half()

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

           lowerCamelCase             =		"A painting of a squirrel eating a burger"
           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      num_inference_steps=2					,      output_type="np"					,      ).images

           lowerCamelCase             =		low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)





@slow
@require_torch_gpu
class  lowerCamelCase__	(    unittest.TestCase):
     """simple docstring"""



     def     _a   (self     ):






           '''simple docstring'''



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







     def     _a   (self     ):






           '''simple docstring'''



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

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(__a     )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      output_type="np"					,      )
           lowerCamelCase             =		output.images[0]

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







     def     _a   (self     ):






           '''simple docstring'''



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

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(
               __a					,      torch_dtype=torch.floataa					,      )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      output_type="np"					,      )
           lowerCamelCase             =		output.images[0]

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







     def     _a   (self     ):






           '''simple docstring'''



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

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

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(
               __a					,      torch_dtype=torch.floataa					,      )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing(1     )
           pipe.enable_sequential_cpu_offload()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      num_inference_steps=5					,      output_type="np"					,      )

           lowerCamelCase             =		torch.cuda.max_memory_allocated()
           # make sure that less than 2.9 GB is allocated
           assert mem_bytes < 2.9 * 10**9