


'''simple docstring'''






import math
from collections import defaultdict
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput






def 		_A (lowerCAmelCase__     :Optional[int] ,  lowerCAmelCase__     :str=0.9_9_9 ,  lowerCAmelCase__     :str="cosine" ,  )      ->     List[str]:




     '''simple docstring'''



     if alpha_transform_type == "cosine":

          def alpha_bar_fn(lowerCAmelCase__     :Union[str, Any]						):
               return math.cos((t + 0.0_0_8) / 1.0_0_8 * math.pi / 2						) ** 2

     elif alpha_transform_type == "exp":

          def alpha_bar_fn(lowerCAmelCase__     :List[str]						):
               return math.exp(t * -1_2.0						)

     else:
          raise ValueError(f'Unsupported alpha_tranform_type: {alpha_transform_type}'						)

     _a						=							[]
     for i in range(snake_case_						):
          _a						=							i / num_diffusion_timesteps
          _a						=							(i + 1) / num_diffusion_timesteps
          betas.append(min(1 - alpha_bar_fn(snake_case_						) / alpha_bar_fn(snake_case_						) ,  snake_case_						)						)
     return torch.tensor(snake_case_ ,  dtype=torch.floataa						)






class        a   (   SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__				):
 _lowerCAmelCase           =		[e.name for e in KarrasDiffusionSchedulers]
 _lowerCAmelCase           =		2
 @register_to_config
 def __init__(       self					, __magic_name__ = 10_00					, __magic_name__ = 0.0_0_0_8_5					, __magic_name__ = 0.0_1_2					, __magic_name__ = "linear"					, __magic_name__ = None					, __magic_name__ = "epsilon"					, __magic_name__ = False					, __magic_name__ = False					, __magic_name__ = 1.0					, __magic_name__ = "linspace"					, __magic_name__ = 0					, )    ->							List[Any]:




      if trained_betas is not None:
           _a						=							torch.tensor(__lowerCamelCase					, dtype=torch.floataa  )
      elif beta_schedule == "linear":
           _a						=							torch.linspace(__lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase					, dtype=torch.floataa  )
      elif beta_schedule == "scaled_linear":
           # this schedule is very specific to the latent diffusion model.
           _a						=							(
               torch.linspace(beta_start**0.5					, beta_end**0.5					, __lowerCamelCase					, dtype=torch.floataa  ) ** 2
           )
      elif beta_schedule == "squaredcos_cap_v2":
           # Glide cosine schedule
           _a						=							betas_for_alpha_bar(__lowerCamelCase					, alpha_transform_type='cosine'  )
      elif beta_schedule == "exp":
           _a						=							betas_for_alpha_bar(__lowerCamelCase					, alpha_transform_type='exp'  )
      else:
           raise NotImplementedError(f'{beta_schedule} does is not implemented for {self.__class__}'  )

      _a						=							1.0 - self.betas
      _a						=							torch.cumprod(self.alphas					, dim=0  )

      #  set all values
      self.set_timesteps(__lowerCamelCase					, __lowerCamelCase					, __lowerCamelCase  )
      _a						=							use_karras_sigmas
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__=None  )    ->							Tuple:




      if schedule_timesteps is None:
           _a						=							self.timesteps

      _a						=							(schedule_timesteps == timestep).nonzero()

      # The sigma index that is taken for the **very** first `step`
      # is always the second index (or the last index if there is only 1)
      # This way we can ensure we don't accidentally skip a sigma in
      # case we start in the middle of the denoising schedule (e.g. for image-to-image)
      if len(self._index_counter  ) == 0:
           _a						=							1 if len(__lowerCamelCase  ) > 1 else 0
      else:
           _a						=							timestep.cpu().item() if torch.is_tensor(__lowerCamelCase  ) else timestep
           _a						=							self._index_counter[timestep_int]

      return indices[pos].item()
 @property
 def 							__UpperCAmelCase						(       self  )    ->							Union[str, Any]:




      if self.config.timestep_spacing in ["linspace", "trailing"]:
           return self.sigmas.max()

      return (self.sigmas.max() ** 2 + 1) ** 0.5
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__					, )    ->							List[Any]:
      _a						=							self.index_for_timestep(__lowerCamelCase  )

      _a						=							self.sigmas[step_index]
      _a						=							sample / ((sigma**2 + 1) ** 0.5)
      return sample
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__ = None					, __magic_name__ = None					, )    ->							str:
      _a						=							num_inference_steps

      _a						=							num_train_timesteps or self.config.num_train_timesteps

      # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
      if self.config.timestep_spacing == "linspace":
           _a						=							np.linspace(0					, num_train_timesteps - 1					, __lowerCamelCase					, dtype=__lowerCamelCase  )[::-1].copy()
      elif self.config.timestep_spacing == "leading":
           _a						=							num_train_timesteps // self.num_inference_steps
           # creates integer timesteps by multiplying by ratio
           # casting to int to avoid issues when num_inference_step is power of 3
           _a						=							(np.arange(0					, __lowerCamelCase  ) * step_ratio).round()[::-1].copy().astype(__lowerCamelCase  )
           timesteps += self.config.steps_offset
      elif self.config.timestep_spacing == "trailing":
           _a						=							num_train_timesteps / self.num_inference_steps
           # creates integer timesteps by multiplying by ratio
           # casting to int to avoid issues when num_inference_step is power of 3
           _a						=							(np.arange(__lowerCamelCase					, 0					, -step_ratio  )).round().copy().astype(__lowerCamelCase  )
           timesteps -= 1
      else:
           raise ValueError(
               f'{self.config.timestep_spacing} is not supported. Please make sure to choose one of \'linspace\', \'leading\' or \'trailing\'.'  )

      _a						=							np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5  )
      _a						=							np.log(__lowerCamelCase  )
      _a						=							np.interp(__lowerCamelCase					, np.arange(0					, len(__lowerCamelCase  )  )					, __lowerCamelCase  )

      if self.config.use_karras_sigmas:
           _a						=							self._convert_to_karras(in_sigmas=__lowerCamelCase					, num_inference_steps=self.num_inference_steps  )
           _a						=							np.array([self._sigma_to_t(__lowerCamelCase					, __lowerCamelCase  ) for sigma in sigmas]  )

      _a						=							np.concatenate([sigmas, [0.0]]  ).astype(np.floataa  )
      _a						=							torch.from_numpy(__lowerCamelCase  ).to(device=__lowerCamelCase  )
      _a						=							torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2  ), sigmas[-1:]]  )

      _a						=							torch.from_numpy(__lowerCamelCase  )
      _a						=							torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2  )]  )

      if str(__lowerCamelCase  ).startswith('mps'  ):
           # mps does not support float64
           _a						=							timesteps.to(__lowerCamelCase					, dtype=torch.floataa  )
      else:
           _a						=							timesteps.to(device=__lowerCamelCase  )

      # empty dt and derivative
      _a						=							None
      _a						=							None

      # for exp beta schedules, such as the one for `pipeline_shap_e.py`
      # we need an index counter
      _a						=							defaultdict(__lowerCamelCase  )
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__  )    ->							Tuple:
      _a						=							np.log(__lowerCamelCase  )

      # get distribution
      _a						=							log_sigma - log_sigmas[:, np.newaxis]

      # get sigmas range
      _a						=							np.cumsum((dists >= 0)					, axis=0  ).argmax(axis=0  ).clip(max=log_sigmas.shape[0] - 2  )
      _a						=							low_idx + 1

      _a						=							log_sigmas[low_idx]
      _a						=							log_sigmas[high_idx]

      # interpolate sigmas
      _a						=							(low - log_sigma) / (low - high)
      _a						=							np.clip(__lowerCamelCase					, 0					, 1  )

      # transform interpolation to time range
      _a						=							(1 - w) * low_idx + w * high_idx
      _a						=							t.reshape(sigma.shape  )
      return t
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__  )    ->							Any:
      _a						=							in_sigmas[-1].item()
      _a						=							in_sigmas[0].item()

      _a						=							7.0  # 7.0 is the value used in the paper
      _a						=							np.linspace(0					, 1					, __lowerCamelCase  )
      _a						=							sigma_min ** (1 / rho)
      _a						=							sigma_max ** (1 / rho)
      _a						=							(max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
      return sigmas
 @property
 def 							__UpperCAmelCase						(       self  )    ->							int:
      return self.dt is None
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__					, __magic_name__					, __magic_name__ = True					, )    ->							List[Any]:
      _a						=							self.index_for_timestep(__lowerCamelCase  )

      # advance index counter by 1
      _a						=							timestep.cpu().item() if torch.is_tensor(__lowerCamelCase  ) else timestep
      self._index_counter[timestep_int] += 1

      if self.state_in_first_order:
           _a						=							self.sigmas[step_index]
           _a						=							self.sigmas[step_index + 1]
      else:
           # 2nd order / Heun's method
           _a						=							self.sigmas[step_index - 1]
           _a						=							self.sigmas[step_index]

      # currently only gamma=0 is supported. This usually works best anyways.
      # We can support gamma in the future but then need to scale the timestep before
      # passing it to the model which requires a change in API
      _a						=							0
      _a						=							sigma * (gamma + 1)  # Note: sigma_hat == sigma for now

      # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
      if self.config.prediction_type == "epsilon":
           _a						=							sigma_hat if self.state_in_first_order else sigma_next
           _a						=							sample - sigma_input * model_output
      elif self.config.prediction_type == "v_prediction":
           _a						=							sigma_hat if self.state_in_first_order else sigma_next
           _a						=							model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + (
               sample / (sigma_input**2 + 1)
           )
      elif self.config.prediction_type == "sample":
           _a						=							model_output
      else:
           raise ValueError(
               f'prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'  )

      if self.config.clip_sample:
           _a						=							pred_original_sample.clamp(
               -self.config.clip_sample_range					, self.config.clip_sample_range  )

      if self.state_in_first_order:
           # 2. Convert to an ODE derivative for 1st order
           _a						=							(sample - pred_original_sample) / sigma_hat
           # 3. delta timestep
           _a						=							sigma_next - sigma_hat

           # store for 2nd order step
           _a						=							derivative
           _a						=							dt
           _a						=							sample
      else:
           # 2. 2nd order / Heun's method
           _a						=							(sample - pred_original_sample) / sigma_next
           _a						=							(self.prev_derivative + derivative) / 2

           # 3. take prev timestep & sample
           _a						=							self.dt
           _a						=							self.sample

           # free dt and derivative
           # Note, this puts the scheduler in "first order mode"
           _a						=							None
           _a						=							None
           _a						=							None

      _a						=							sample + derivative * dt

      if not return_dict:
           return (prev_sample,)

      return SchedulerOutput(prev_sample=__lowerCamelCase  )
 def 							__UpperCAmelCase						(       self					, __magic_name__					, __magic_name__					, __magic_name__					, )    ->							Optional[int]:
      _a						=							self.sigmas.to(device=original_samples.device					, dtype=original_samples.dtype  )
      if original_samples.device.type == "mps" and torch.is_floating_point(__lowerCamelCase  ):
           # mps does not support float64
           _a						=							self.timesteps.to(original_samples.device					, dtype=torch.floataa  )
           _a						=							timesteps.to(original_samples.device					, dtype=torch.floataa  )
      else:
           _a						=							self.timesteps.to(original_samples.device  )
           _a						=							timesteps.to(original_samples.device  )

      _a						=							[self.index_for_timestep(__lowerCamelCase					, __lowerCamelCase  ) for t in timesteps]

      _a						=							sigmas[step_indices].flatten()
      while len(sigma.shape  ) < len(original_samples.shape  ):
           _a						=							sigma.unsqueeze(-1  )

      _a						=							original_samples + noise * sigma
      return noisy_samples




 def __len__(       self  )    ->							Optional[int]:
      return self.config.num_train_timesteps








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





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     1
     __A    :						Any						 =     3
     __A    :						List[str]						 =     (32, 32)

     __A    :						List[str]						 =     floats_tensor((batch_size, num_channels) + sizes					,						rng=random.Random(0  )  ).to(__lowerCamelCase  )
     return image

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						List[Any]						 =     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=__lowerCamelCase					,						only_cross_attention=(True, True, False)					,						num_class_embeds=100					,						)
     return model

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						Any						 =     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 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     torch.manual_seed(0  )
     __A    :						Tuple						 =     CLIPTextConfig(
         bos_token_id=0					,						eos_token_id=2					,						hidden_size=32					,						intermediate_size=37					,						layer_norm_eps=1e-0_5					,						num_attention_heads=4					,						num_hidden_layers=5					,						pad_token_id=1					,						vocab_size=1000					,						hidden_act='''gelu'''					,						projection_dim=512					,						)
     return CLIPTextModel(__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

     __A    :						Tuple						 =     self.dummy_image.cpu().permute(0					,						2					,						3					,						1  )[0]
     __A    :						Any						 =     Image.fromarray(np.uinta(__lowerCamelCase  )  ).convert('''RGB'''  ).resize((64, 64)  )

     # make sure here that pndm scheduler skips prk
     __A    :						Dict						 =     StableDiffusionUpscalePipeline(
         unet=__lowerCamelCase					,						low_res_scheduler=__lowerCamelCase					,						scheduler=__lowerCamelCase					,						vae=__lowerCamelCase					,						text_encoder=__lowerCamelCase					,						tokenizer=__lowerCamelCase					,						max_noise_level=350					,						)
     __A    :						str						 =     sd_pipe.to(__lowerCamelCase  )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

     __A    :						List[str]						 =     '''A painting of a squirrel eating a burger'''
     __A    :						Any						 =     torch.Generator(device=__lowerCamelCase  ).manual_seed(0  )
     __A    :						List[str]						 =     sd_pipe(
         [prompt]					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						guidance_scale=6.0					,						noise_level=20					,						num_inference_steps=2					,						output_type='''np'''					,						)

     __A    :						Union[str, Any]						 =     output.images

     __A    :						List[str]						 =     torch.Generator(device=__lowerCamelCase  ).manual_seed(0  )
     __A    :						str						 =     sd_pipe(
         [prompt]					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						guidance_scale=6.0					,						noise_level=20					,						num_inference_steps=2					,						output_type='''np'''					,						return_dict=__lowerCamelCase					,						)[0]

     __A    :						Tuple						 =     image[0, -3:, -3:, -1]
     __A    :						int						 =     image_from_tuple[0, -3:, -3:, -1]

     __A    :						Dict						 =     low_res_image.size[0] * 4
     assert image.shape == (1, expected_height_width, expected_height_width, 3)
     __A    :						str						 =     np.array([0.3_1_1_3, 0.3_9_1_0, 0.4_2_7_2, 0.4_8_5_9, 0.5_0_6_1, 0.4_6_5_2, 0.5_3_6_2, 0.5_7_1_5, 0.5_6_6_1]  )

     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 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

     __A    :						List[Any]						 =     self.dummy_image.cpu().permute(0					,						2					,						3					,						1  )[0]
     __A    :						int						 =     Image.fromarray(np.uinta(__lowerCamelCase  )  ).convert('''RGB'''  ).resize((64, 64)  )

     # make sure here that pndm scheduler skips prk
     __A    :						Any						 =     StableDiffusionUpscalePipeline(
         unet=__lowerCamelCase					,						low_res_scheduler=__lowerCamelCase					,						scheduler=__lowerCamelCase					,						vae=__lowerCamelCase					,						text_encoder=__lowerCamelCase					,						tokenizer=__lowerCamelCase					,						max_noise_level=350					,						)
     __A    :						Any						 =     sd_pipe.to(__lowerCamelCase  )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

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

     __A    :						Optional[Any]						 =     torch.Generator(device=__lowerCamelCase  ).manual_seed(0  )
     __A    :						Any						 =     sd_pipe(
         [prompt]					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						num_images_per_prompt=2					,						guidance_scale=6.0					,						noise_level=20					,						num_inference_steps=2					,						output_type='''np'''					,						)
     __A    :						Union[str, Any]						 =     output.images
     assert image.shape[0] == 2







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



     '''simple docstring'''




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

     __A    :						Union[str, Any]						 =     self.dummy_image.cpu().permute(0					,						2					,						3					,						1  )[0]
     __A    :						int						 =     Image.fromarray(np.uinta(__lowerCamelCase  )  ).convert('''RGB'''  ).resize((64, 64)  )

     # put models in fp16, except vae as it overflows in fp16
     __A    :						Union[str, Any]						 =     unet.half()
     __A    :						Optional[int]						 =     text_encoder.half()

     # make sure here that pndm scheduler skips prk
     __A    :						Optional[int]						 =     StableDiffusionUpscalePipeline(
         unet=__lowerCamelCase					,						low_res_scheduler=__lowerCamelCase					,						scheduler=__lowerCamelCase					,						vae=__lowerCamelCase					,						text_encoder=__lowerCamelCase					,						tokenizer=__lowerCamelCase					,						max_noise_level=350					,						)
     __A    :						Union[str, Any]						 =     sd_pipe.to(__lowerCamelCase  )
     sd_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     '''A painting of a squirrel eating a burger'''
     __A    :						Optional[Any]						 =     torch.manual_seed(0  )
     __A    :						Tuple						 =     sd_pipe(
         [prompt]					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						num_inference_steps=2					,						output_type='''np'''					,						).images

     __A    :						str						 =     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     ):





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

     __A    :						str						 =     '''stabilityai/stable-diffusion-x4-upscaler'''
     __A    :						Optional[Any]						 =     StableDiffusionUpscalePipeline.from_pretrained(__lowerCamelCase  )
     pipe.to(__lowerCamelCase  )
     pipe.set_progress_bar_config(disable=__lowerCamelCase  )
     pipe.enable_attention_slicing()

     __A    :						Union[str, Any]						 =     '''a cat sitting on a park bench'''

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

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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Dict						 =     load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/sd2-upscale/low_res_cat.png'''  )
     __A    :						List[Any]						 =     load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
         '''/upsampled_cat_fp16.npy'''  )

     __A    :						Optional[int]						 =     '''stabilityai/stable-diffusion-x4-upscaler'''
     __A    :						Optional[int]						 =     StableDiffusionUpscalePipeline.from_pretrained(
         __lowerCamelCase					,						torch_dtype=torch.floataa					,						)
     pipe.to(__lowerCamelCase  )
     pipe.set_progress_bar_config(disable=__lowerCamelCase  )
     pipe.enable_attention_slicing()

     __A    :						Dict						 =     '''a cat sitting on a park bench'''

     __A    :						Any						 =     torch.manual_seed(0  )
     __A    :						Optional[int]						 =     pipe(
         prompt=__lowerCamelCase					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						output_type='''np'''					,						)
     __A    :						Any						 =     output.images[0]

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







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




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

     __A    :						Union[str, Any]						 =     load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/sd2-upscale/low_res_cat.png'''  )

     __A    :						List[str]						 =     '''stabilityai/stable-diffusion-x4-upscaler'''
     __A    :						Dict						 =     StableDiffusionUpscalePipeline.from_pretrained(
         __lowerCamelCase					,						torch_dtype=torch.floataa					,						)
     pipe.to(__lowerCamelCase  )
     pipe.set_progress_bar_config(disable=__lowerCamelCase  )
     pipe.enable_attention_slicing(1  )
     pipe.enable_sequential_cpu_offload()

     __A    :						Tuple						 =     '''a cat sitting on a park bench'''

     __A    :						Tuple						 =     torch.manual_seed(0  )
     __A    :						List[str]						 =     pipe(
         prompt=__lowerCamelCase					,						image=__lowerCamelCase					,						generator=__lowerCamelCase					,						num_inference_steps=5					,						output_type='''np'''					,						)

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




