

'''simple docstring'''
import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


lowerCamelCase :      Any					 =     logging.get_logger(__name__)







class 							A__    (       UpperCamelCase__  ):





  def __init__(		self      :	Tuple							,					_a      :	Union[List[ControlNetModel], Tuple[ControlNetModel]] )	->    Optional[int]:





       '''simple docstring'''






       super().__init__()
       _SCREAMING_SNAKE_CASE							    =nn.ModuleList(__a )





  def 						A  (		self      :	Tuple							,					_a      :	torch.FloatTensor							,					_a      :	Union[torch.Tensor, float, int]							,					_a      :	torch.Tensor							,					_a      :	List[torch.tensor]							,					_a      :	List[float]							,					_a      :	Optional[torch.Tensor] = None							,					_a      :	Optional[torch.Tensor] = None							,					_a      :	Optional[torch.Tensor] = None							,					_a      :	Optional[Dict[str, Any]] = None							,					_a      :	bool = False							,					_a      :	bool = True							,					)	->    Optional[int]:





       '''simple docstring'''






       for i, (image, scale, controlnet) in enumerate(zip(__a							,					__a							,					self.nets ) ):
            _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =controlnet(
                __a							,					__a							,					__a							,					__a							,					__a							,					__a							,					__a							,					__a							,					__a							,					__a							,					__a							,					)

            # merge samples
            if i == 0:
                 _SCREAMING_SNAKE_CASE    ,						_SCREAMING_SNAKE_CASE							    =down_samples, mid_sample
            else:
                 _SCREAMING_SNAKE_CASE							    =[
                     samples_prev + samples_curr
                     for samples_prev, samples_curr in zip(__a							,					__a )
                 ]
                 mid_block_res_sample += mid_sample

       return down_block_res_samples, mid_block_res_sample





  def 						A  (		self      :	List[str]							,					_a      :	Union[str, os.PathLike]							,					_a      :	bool = True							,					_a      :	Callable = None							,					_a      :	bool = False							,					_a      :	Optional[str] = None							,					)	->    Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =0
       _SCREAMING_SNAKE_CASE							    =save_directory
       for controlnet in self.nets:
            controlnet.save_pretrained(
                __a							,					is_main_process=__a							,					save_function=__a							,					safe_serialization=__a							,					variant=__a							,					)

            idx += 1
            _SCREAMING_SNAKE_CASE							    =model_path_to_save + f"_{idx}"





  @classmethod
  def 						A  (		cls      :	Tuple							,					_a      :	Optional[Union[str, os.PathLike]]							,					**_a      :	List[Any] )	->    Union[str, Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							    =0
       _SCREAMING_SNAKE_CASE							    =[]

       # load controlnet and append to list until no controlnet directory exists anymore
       # first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
       # second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
       _SCREAMING_SNAKE_CASE							    =pretrained_model_path
       while os.path.isdir(__a ):
            _SCREAMING_SNAKE_CASE							    =ControlNetModel.from_pretrained(__a							,					**__a )
            controlnets.append(__a )

            idx += 1
            _SCREAMING_SNAKE_CASE							    =pretrained_model_path + f"_{idx}"

       logger.info(f"{len(__a )} controlnets loaded from {pretrained_model_path}." )

       if len(__a ) == 0:
            raise ValueError(
                f"No ControlNets found under {os.path.dirname(__a )}. Expected at least {pretrained_model_path + '_0'}." )

       return cls(__a )

'''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




