'''simple docstring'''





# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()




class        __A      ( UpperCamelCase__		,	UpperCamelCase__		,	UpperCamelCase__		,	unittest.TestCase	):
    a__       :						str          =  StableDiffusionControlNetImgaImgPipeline
    a__       :						Optional[Any]          =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
    a__       :						Any          =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a__       :						Dict          =  IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""}	)
    a__       :						Dict          =  IMAGE_TO_IMAGE_IMAGE_PARAMS



    def 					_lowercase						(self    :    List[str] ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   UNetaDConditionModel(
             block_out_channels=(32, 64)   ,			layers_per_block=2   ,			sample_size=32   ,			in_channels=4   ,			out_channels=4   ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,			cross_attention_dim=32   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   ControlNetModel(
             block_out_channels=(32, 64)   ,			layers_per_block=2   ,			in_channels=4   ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,			cross_attention_dim=32   ,			conditioning_embedding_out_channels=(16, 32)   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   DDIMScheduler(
             beta_start=0.0_00_85   ,			beta_end=0.0_12   ,			beta_schedule="scaled_linear"   ,			clip_sample=__a   ,			set_alpha_to_one=__a   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   AutoencoderKL(
             block_out_channels=[32, 64]   ,			in_channels=3   ,			out_channels=3   ,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,			latent_channels=4   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   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=1000   ,			)
         UpperCAmelCase_           =   CLIPTextModel(__a )
         UpperCAmelCase_           =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

         UpperCAmelCase_           =   {
             "unet": unet,
             "controlnet": controlnet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components



    def 					_lowercase						(self    :    List[str]   ,			__a    :    Optional[Any]   ,			__a    :    List[Any]=0 ):
         if str(__a ).startswith("mps" ):
              UpperCAmelCase_           =   torch.manual_seed(__a )
         else:
              UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(__a )

         UpperCAmelCase_           =   2
         UpperCAmelCase_           =   randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,			generator=__a   ,			device=torch.device(__a )   ,			)
         UpperCAmelCase_           =   floats_tensor(control_image.shape   ,			rng=random.Random(__a ) ).to(__a )
         UpperCAmelCase_           =   image.cpu().permute(0   ,			2   ,			3   ,			1 )[0]
         UpperCAmelCase_           =   Image.fromarray(np.uinta(__a ) ).convert("RGB" ).resize((64, 64) )
         UpperCAmelCase_           =   {
             "prompt": "A painting of a squirrel eating a burger",
             "generator": generator,
             "num_inference_steps": 2,
             "guidance_scale": 6.0,
             "output_type": "numpy",
             "image": image,
             "control_image": control_image,
         }

         return inputs



    def 					_lowercase						(self    :    Optional[int] ):
         return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3 )



    @unittest.skipIf(
        torch_device != "cuda" or not is_xformers_available()   ,			reason="XFormers attention is only available with CUDA and `xformers` installed"   ,			)
    def 					_lowercase						(self    :    List[Any] ):
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3 )






    def 					_lowercase						(self    :    Dict ):
         self._test_inference_batch_single_identical(expected_max_diff=2E-3 )




class        __A      ( UpperCamelCase__		,	UpperCamelCase__		,	unittest.TestCase	):
    a__       :						List[str]          =  StableDiffusionControlNetImgaImgPipeline
    a__       :						Any          =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
    a__       :						Dict          =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a__       :						Optional[int]          =  frozenset([]	)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess



    def 					_lowercase						(self    :    Dict ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   UNetaDConditionModel(
             block_out_channels=(32, 64)   ,			layers_per_block=2   ,			sample_size=32   ,			in_channels=4   ,			out_channels=4   ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,			cross_attention_dim=32   ,			)
         torch.manual_seed(0 )

         def init_weights(__a    :    Union[str, Any] ):
              if isinstance(__a   ,			torch.nn.Convad ):
                   torch.nn.init.normal(m.weight )
                   m.bias.data.fill_(1.0 )

         UpperCAmelCase_           =   ControlNetModel(
             block_out_channels=(32, 64)   ,			layers_per_block=2   ,			in_channels=4   ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,			cross_attention_dim=32   ,			conditioning_embedding_out_channels=(16, 32)   ,			)
         controlneta.controlnet_down_blocks.apply(__a )

         torch.manual_seed(0 )
         UpperCAmelCase_           =   ControlNetModel(
             block_out_channels=(32, 64)   ,			layers_per_block=2   ,			in_channels=4   ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,			cross_attention_dim=32   ,			conditioning_embedding_out_channels=(16, 32)   ,			)
         controlneta.controlnet_down_blocks.apply(__a )

         torch.manual_seed(0 )
         UpperCAmelCase_           =   DDIMScheduler(
             beta_start=0.0_00_85   ,			beta_end=0.0_12   ,			beta_schedule="scaled_linear"   ,			clip_sample=__a   ,			set_alpha_to_one=__a   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   AutoencoderKL(
             block_out_channels=[32, 64]   ,			in_channels=3   ,			out_channels=3   ,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,			latent_channels=4   ,			)
         torch.manual_seed(0 )
         UpperCAmelCase_           =   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=1000   ,			)
         UpperCAmelCase_           =   CLIPTextModel(__a )
         UpperCAmelCase_           =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

         UpperCAmelCase_           =   MultiControlNetModel([controlneta, controlneta] )

         UpperCAmelCase_           =   {
             "unet": unet,
             "controlnet": controlnet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components



    def 					_lowercase						(self    :    List[str]   ,			__a    :    int   ,			__a    :    str=0 ):
         if str(__a ).startswith("mps" ):
              UpperCAmelCase_           =   torch.manual_seed(__a )
         else:
              UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(__a )

         UpperCAmelCase_           =   2

         UpperCAmelCase_           =   [
             randn_tensor(
                 (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,			generator=__a   ,			device=torch.device(__a )   ,			),
             randn_tensor(
                 (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,			generator=__a   ,			device=torch.device(__a )   ,			),
         ]

         UpperCAmelCase_           =   floats_tensor(control_image[0].shape   ,			rng=random.Random(__a ) ).to(__a )
         UpperCAmelCase_           =   image.cpu().permute(0   ,			2   ,			3   ,			1 )[0]
         UpperCAmelCase_           =   Image.fromarray(np.uinta(__a ) ).convert("RGB" ).resize((64, 64) )
         UpperCAmelCase_           =   {
             "prompt": "A painting of a squirrel eating a burger",
             "generator": generator,
             "num_inference_steps": 2,
             "guidance_scale": 6.0,
             "output_type": "numpy",
             "image": image,
             "control_image": control_image,
         }

         return inputs



    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_           =   self.get_dummy_components()
         UpperCAmelCase_           =   self.pipeline_class(**__a )
         pipe.to(__a )

         UpperCAmelCase_           =   10.0
         UpperCAmelCase_           =   4

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   steps
         UpperCAmelCase_           =   scale
         UpperCAmelCase_           =   pipe(**__a )[0]

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   steps
         UpperCAmelCase_           =   scale
         UpperCAmelCase_           =   pipe(**__a   ,			control_guidance_start=0.1   ,			control_guidance_end=0.2 )[0]

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   steps
         UpperCAmelCase_           =   scale
         UpperCAmelCase_           =   pipe(**__a   ,			control_guidance_start=[0.1, 0.3]   ,			control_guidance_end=[0.2, 0.7] )[0]

         UpperCAmelCase_           =   self.get_dummy_inputs(__a )
         UpperCAmelCase_           =   steps
         UpperCAmelCase_           =   scale
         UpperCAmelCase_           =   pipe(**__a   ,			control_guidance_start=0.4   ,			control_guidance_end=[0.5, 0.8] )[0]

         # make sure that all outputs are different
         assert np.sum(np.abs(output_a - output_a ) ) > 1E-3
         assert np.sum(np.abs(output_a - output_a ) ) > 1E-3
         assert np.sum(np.abs(output_a - output_a ) ) > 1E-3



    def 					_lowercase						(self    :    Any ):
         return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3 )



    @unittest.skipIf(
        torch_device != "cuda" or not is_xformers_available()   ,			reason="XFormers attention is only available with CUDA and `xformers` installed"   ,			)
    def 					_lowercase						(self    :    List[Any] ):
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3 )



    def 					_lowercase						(self    :    List[str] ):
         self._test_inference_batch_single_identical(expected_max_diff=2E-3 )






    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.get_dummy_components()
         UpperCAmelCase_           =   self.pipeline_class(**__a )
         pipe.to(__a )
         pipe.set_progress_bar_config(disable=__a )
         with tempfile.TemporaryDirectory() as tmpdir:
              try:
                   # save_pretrained is not implemented for Multi-ControlNet
                   pipe.save_pretrained(__a )
              except NotImplementedError:
                   pass




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



    def 					_lowercase						(self    :    Optional[int] ):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_           =   ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny" )

         UpperCAmelCase_           =   StableDiffusionControlNetImgaImgPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"   ,			safety_checker=__a   ,			controlnet=__a )
         pipe.enable_model_cpu_offload()
         pipe.set_progress_bar_config(disable=__a )

         UpperCAmelCase_           =   torch.Generator(device="cpu" ).manual_seed(0 )
         UpperCAmelCase_           =   "evil space-punk bird"
         UpperCAmelCase_           =   load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png" ).resize((512, 512) )
         UpperCAmelCase_           =   load_image(
             "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png" ).resize((512, 512) )

         UpperCAmelCase_           =   pipe(
             __a   ,			__a   ,			control_image=__a   ,			generator=__a   ,			output_type="np"   ,			num_inference_steps=50   ,			strength=0.6   ,			)

         UpperCAmelCase_           =   output.images[0]

         assert image.shape == (512, 512, 3)

         UpperCAmelCase_           =   load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy" )

         assert np.abs(expected_image - image ).max() < 9E-2





'''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        __A      ( unittest.TestCase	):



    def 					_lowercase						(self    :    Tuple ):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



    @property
    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   3
         UpperCAmelCase_           =   (32, 32)

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



    @property
    def 					_lowercase						(self    :    int ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   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=100   ,			)
         return model



    @property
    def 					_lowercase						(self    :    Any ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   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 					_lowercase						(self    :    Optional[Any] ):
         torch.manual_seed(0 )
         UpperCAmelCase_           =   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=1000   ,			hidden_act="gelu"   ,			projection_dim=512   ,			)
         return CLIPTextModel(__a )



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

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

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

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

         UpperCAmelCase_           =   output.images

         UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(0 )
         UpperCAmelCase_           =   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]

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

         UpperCAmelCase_           =   low_res_image.size[0] * 4
         assert image.shape == (1, expected_height_width, expected_height_width, 3)
         UpperCAmelCase_           =   np.array([0.31_13, 0.39_10, 0.42_72, 0.48_59, 0.50_61, 0.46_52, 0.53_62, 0.57_15, 0.56_61] )

         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    :    Optional[int] ):
         UpperCAmelCase_           =   "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase_           =   self.dummy_cond_unet_upscale
         UpperCAmelCase_           =   DDPMScheduler()
         UpperCAmelCase_           =   DDIMScheduler(prediction_type="v_prediction" )
         UpperCAmelCase_           =   self.dummy_vae
         UpperCAmelCase_           =   self.dummy_text_encoder
         UpperCAmelCase_           =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

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

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

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

         UpperCAmelCase_           =   torch.Generator(device=__a ).manual_seed(0 )
         UpperCAmelCase_           =   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"   ,			)
         UpperCAmelCase_           =   output.images
         assert image.shape[0] == 2






    @unittest.skipIf(torch_device != "cuda"   ,			"This test requires a GPU" )
    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.dummy_cond_unet_upscale
         UpperCAmelCase_           =   DDPMScheduler()
         UpperCAmelCase_           =   DDIMScheduler(prediction_type="v_prediction" )
         UpperCAmelCase_           =   self.dummy_vae
         UpperCAmelCase_           =   self.dummy_text_encoder
         UpperCAmelCase_           =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

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

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

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

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

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




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



    def 					_lowercase						(self    :    List[str] ):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



    def 					_lowercase						(self    :    List[Any] ):
         UpperCAmelCase_           =   load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/sd2-upscale/low_res_cat.png" )
         UpperCAmelCase_           =   load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
             "/upsampled_cat.npy" )

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

         UpperCAmelCase_           =   "a cat sitting on a park bench"

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

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



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

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

         UpperCAmelCase_           =   "a cat sitting on a park bench"

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

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






    def 					_lowercase						(self    :    List[Any] ):
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

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

         UpperCAmelCase_           =   "stabilityai/stable-diffusion-x4-upscaler"
         UpperCAmelCase_           =   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()

         UpperCAmelCase_           =   "a cat sitting on a park bench"

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

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




