

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 					__UpperCAmelCase     (       SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  unittest.TestCase						):
   UpperCamelCase        =     StableDiffusionControlNetImgaImgPipeline
   UpperCamelCase        =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
   UpperCamelCase        =     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   UpperCamelCase        =     IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""}						)
   UpperCamelCase        =     IMAGE_TO_IMAGE_IMAGE_PARAMS



   def        __magic_name__    (							self							:     Union[str, Any]						):
     torch.manual_seed(0						)
     UpperCAmelCase       :							Tuple  = UNetaDConditionModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		sample_size=3_2,		in_channels=4,		out_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),		cross_attention_dim=3_2,		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							Any  = ControlNetModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		in_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		cross_attention_dim=3_2,		conditioning_embedding_out_channels=(1_6, 3_2),		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							Any  = DDIMScheduler(
         beta_start=0.0_0_0_8_5,		beta_end=0.0_1_2,		beta_schedule='''scaled_linear''',		clip_sample=_a,		set_alpha_to_one=_a,		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							Dict  = AutoencoderKL(
         block_out_channels=[3_2, 6_4],		in_channels=3,		out_channels=3,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],		latent_channels=4,		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							Dict  = CLIPTextConfig(
         bos_token_id=0,		eos_token_id=2,		hidden_size=3_2,		intermediate_size=3_7,		layer_norm_eps=1E-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=1_0_0_0,		)
     UpperCAmelCase       :							Union[str, Any]  = CLIPTextModel(_a						)
     UpperCAmelCase       :							Any  = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

     UpperCAmelCase       :							int  = {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components



   def        __magic_name__    (							self							:     Union[str, Any],		__A							:     Any,		__A							:     List[str]=0						):
     if str(_a						).startswith('''mps'''						):
       UpperCAmelCase       :							Any  = torch.manual_seed(_a						)
     else:
       UpperCAmelCase       :							Any  = torch.Generator(device=_a						).manual_seed(_a						)

     UpperCAmelCase       :							Union[str, Any]  = 2
     UpperCAmelCase       :							Any  = randn_tensor(
         (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),		generator=_a,		device=torch.device(_a						),		)
     UpperCAmelCase       :							Optional[Any]  = floats_tensor(control_image.shape,		rng=random.Random(_a						)						).to(_a						)
     UpperCAmelCase       :							List[Any]  = image.cpu().permute(0,		2,		3,		1						)[0]
     UpperCAmelCase       :							int  = Image.fromarray(np.uinta(_a						)						).convert('''RGB'''						).resize((6_4, 6_4)						)
     UpperCAmelCase       :							int  = {
         '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        __magic_name__    (							self							:     List[str]						):
     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        __magic_name__    (							self							:     str						):
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3						)





   def        __magic_name__    (							self							:     Any						):
     self._test_inference_batch_single_identical(expected_max_diff=2E-3						)






class 					__UpperCAmelCase     (       SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_			,  unittest.TestCase						):
   UpperCamelCase        =     StableDiffusionControlNetImgaImgPipeline
   UpperCamelCase        =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
   UpperCamelCase        =     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   UpperCamelCase        =     frozenset([]						)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess



   def        __magic_name__    (							self							:     List[Any]						):
     torch.manual_seed(0						)
     UpperCAmelCase       :							Any  = UNetaDConditionModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		sample_size=3_2,		in_channels=4,		out_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),		cross_attention_dim=3_2,		)
     torch.manual_seed(0						)

     def init_weights(__A							:     str						):
       if isinstance(_a,		torch.nn.Convad						):
         torch.nn.init.normal(m.weight						)
         m.bias.data.fill_(1.0						)

     UpperCAmelCase       :							Any  = ControlNetModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		in_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		cross_attention_dim=3_2,		conditioning_embedding_out_channels=(1_6, 3_2),		)
     controlneta.controlnet_down_blocks.apply(_a						)

     torch.manual_seed(0						)
     UpperCAmelCase       :							Any  = ControlNetModel(
         block_out_channels=(3_2, 6_4),		layers_per_block=2,		in_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		cross_attention_dim=3_2,		conditioning_embedding_out_channels=(1_6, 3_2),		)
     controlneta.controlnet_down_blocks.apply(_a						)

     torch.manual_seed(0						)
     UpperCAmelCase       :							List[Any]  = DDIMScheduler(
         beta_start=0.0_0_0_8_5,		beta_end=0.0_1_2,		beta_schedule='''scaled_linear''',		clip_sample=_a,		set_alpha_to_one=_a,		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							str  = AutoencoderKL(
         block_out_channels=[3_2, 6_4],		in_channels=3,		out_channels=3,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],		latent_channels=4,		)
     torch.manual_seed(0						)
     UpperCAmelCase       :							Optional[int]  = CLIPTextConfig(
         bos_token_id=0,		eos_token_id=2,		hidden_size=3_2,		intermediate_size=3_7,		layer_norm_eps=1E-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=1_0_0_0,		)
     UpperCAmelCase       :							Dict  = CLIPTextModel(_a						)
     UpperCAmelCase       :							Optional[Any]  = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

     UpperCAmelCase       :							List[str]  = MultiControlNetModel([controlneta, controlneta]						)

     UpperCAmelCase       :							str  = {
         'unet': unet,
         'controlnet': controlnet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components



   def        __magic_name__    (							self							:     Dict,		__A							:     int,		__A							:     Union[str, Any]=0						):
     if str(_a						).startswith('''mps'''						):
       UpperCAmelCase       :							Optional[Any]  = torch.manual_seed(_a						)
     else:
       UpperCAmelCase       :							Union[str, Any]  = torch.Generator(device=_a						).manual_seed(_a						)

     UpperCAmelCase       :							Any  = 2

     UpperCAmelCase       :							Any  = [
         randn_tensor(
             (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),		generator=_a,		device=torch.device(_a						),		),
         randn_tensor(
             (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),		generator=_a,		device=torch.device(_a						),		),
     ]

     UpperCAmelCase       :							List[Any]  = floats_tensor(control_image[0].shape,		rng=random.Random(_a						)						).to(_a						)
     UpperCAmelCase       :							Union[str, Any]  = image.cpu().permute(0,		2,		3,		1						)[0]
     UpperCAmelCase       :							Dict  = Image.fromarray(np.uinta(_a						)						).convert('''RGB'''						).resize((6_4, 6_4)						)
     UpperCAmelCase       :							Dict  = {
         '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        __magic_name__    (							self							:     Tuple						):
     UpperCAmelCase       :							Any  = self.get_dummy_components()
     UpperCAmelCase       :							Any  = self.pipeline_class(**_a						)
     pipe.to(_a						)

     UpperCAmelCase       :							Optional[int]  = 1_0.0
     UpperCAmelCase       :							Tuple  = 4

     UpperCAmelCase       :							str  = self.get_dummy_inputs(_a						)
     UpperCAmelCase       :							int  = steps
     UpperCAmelCase       :							List[Any]  = scale
     UpperCAmelCase       :							Optional[Any]  = pipe(**_a						)[0]

     UpperCAmelCase       :							Dict  = self.get_dummy_inputs(_a						)
     UpperCAmelCase       :							List[str]  = steps
     UpperCAmelCase       :							List[Any]  = scale
     UpperCAmelCase       :							Dict  = pipe(**_a,		control_guidance_start=0.1,		control_guidance_end=0.2						)[0]

     UpperCAmelCase       :							Any  = self.get_dummy_inputs(_a						)
     UpperCAmelCase       :							Union[str, Any]  = steps
     UpperCAmelCase       :							Any  = scale
     UpperCAmelCase       :							Optional[int]  = pipe(**_a,		control_guidance_start=[0.1, 0.3],		control_guidance_end=[0.2, 0.7]						)[0]

     UpperCAmelCase       :							Any  = self.get_dummy_inputs(_a						)
     UpperCAmelCase       :							Tuple  = steps
     UpperCAmelCase       :							List[Any]  = scale
     UpperCAmelCase       :							List[str]  = 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        __magic_name__    (							self							:     Optional[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        __magic_name__    (							self							:     Dict						):
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3						)



   def        __magic_name__    (							self							:     str						):
     self._test_inference_batch_single_identical(expected_max_diff=2E-3						)





   def        __magic_name__    (							self							:     List[Any]						):
     UpperCAmelCase       :							int  = self.get_dummy_components()
     UpperCAmelCase       :							List[Any]  = 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 					__UpperCAmelCase     (       unittest.TestCase						):



   def        __magic_name__    (							self							:     List[str]						):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





   def        __magic_name__    (							self							:     Optional[Any]						):
     UpperCAmelCase       :							Optional[Any]  = ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''						)

     UpperCAmelCase       :							List[str]  = 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       :							Optional[int]  = torch.Generator(device='''cpu'''						).manual_seed(0						)
     UpperCAmelCase       :							List[Any]  = 'evil space-punk bird'
     UpperCAmelCase       :							List[str]  = load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''						).resize((5_1_2, 5_1_2)						)
     UpperCAmelCase       :							Any  = load_image(
         '''https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'''						).resize((5_1_2, 5_1_2)						)

     UpperCAmelCase       :							Dict  = pipe(
         _a,		_a,		control_image=_a,		generator=_a,		output_type='''np''',		num_inference_steps=5_0,		strength=0.6,		)

     UpperCAmelCase       :							Optional[Any]  = output.images[0]

     assert image.shape == (5_1_2, 5_1_2, 3)

     UpperCAmelCase       :							List[str]  = 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 inspect
from typing import List, Optional, Tuple, Union

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput

def        a_     (							_lowerCAmelCase       )      ->      str:
  __lowerCamelCase       ,__lowerCamelCase       :       List[Any]           = image.size
  __lowerCamelCase       ,__lowerCamelCase       :       List[str]           = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
  __lowerCamelCase       :       Optional[Any]           = image.resize((w, h)       ,resample=PIL_INTERPOLATION['lanczos']       )
  __lowerCamelCase       :       List[Any]           = np.array(_lowerCAmelCase       ).astype(np.floataa       ) / 255.0
  __lowerCamelCase       :       Optional[Any]           = image[None].transpose(0       ,3       ,1       ,2       )
  __lowerCamelCase       :       int           = torch.from_numpy(_lowerCAmelCase       )
  return 2.0 * image - 1.0






class 					lowerCamelCase_    (				SCREAMING_SNAKE_CASE_     ):
  """simple docstring"""




  def __init__( self      :	str			,	_a      :	VQModel			,	_a      :	UNetaDModel			,	_a      :	Union[
          DDIMScheduler,
          PNDMScheduler,
          LMSDiscreteScheduler,
          EulerDiscreteScheduler,
          EulerAncestralDiscreteScheduler,
          DPMSolverMultistepScheduler,
      ]			,	)					->					Optional[Any]:
    super().__init__()
    self.register_modules(vqvae=_a			,	unet=_a			,	scheduler=_a   )




  @torch.no_grad()
  def __call__( self      :	List[Any]			,	_a      :	Union[torch.Tensor, PIL.Image.Image] = None			,	_a      :	Optional[int] = 1			,	_a      :	Optional[int] = 100			,	_a      :	Optional[float] = 0.0			,	_a      :	Optional[Union[torch.Generator, List[torch.Generator]]] = None			,	_a      :	Optional[str] = "pil"			,	_a      :	bool = True			,	)					->					Union[Tuple, ImagePipelineOutput]:
    if isinstance(_a			,	PIL.Image.Image   ):
      __lowerCamelCase       :       Any           = 1
    elif isinstance(_a			,	torch.Tensor   ):
      __lowerCamelCase       :       Any           = image.shape[0]
    else:
      raise ValueError(f'`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(_a   )}'   )

    if isinstance(_a			,	PIL.Image.Image   ):
      __lowerCamelCase       :       List[str]           = preprocess(_a   )

    __lowerCamelCase       ,__lowerCamelCase       :       List[str]           = image.shape[-2:]

    # in_channels should be 6: 3 for latents, 3 for low resolution image
    __lowerCamelCase       :       Union[str, Any]           = (batch_size, self.unet.config.in_channels // 2, height, width)
    __lowerCamelCase       :       Tuple           = next(self.unet.parameters()   ).dtype

    __lowerCamelCase       :       Optional[int]           = randn_tensor(_a			,	generator=_a			,	device=self.device			,	dtype=_a   )

    __lowerCamelCase       :       Optional[int]           = image.to(device=self.device			,	dtype=_a   )

    # set timesteps and move to the correct device
    self.scheduler.set_timesteps(_a			,	device=self.device   )
    __lowerCamelCase       :       Union[str, Any]           = self.scheduler.timesteps

    # scale the initial noise by the standard deviation required by the scheduler
    __lowerCamelCase       :       List[Any]           = latents * self.scheduler.init_noise_sigma

    # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
    # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
    # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
    # and should be between [0, 1]
    __lowerCamelCase       :       int           = 'eta' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
    __lowerCamelCase       :       List[str]           = {}
    if accepts_eta:
      __lowerCamelCase       :       str           = eta

    for t in self.progress_bar(_a   ):
      # concat latents and low resolution image in the channel dimension.
      __lowerCamelCase       :       str           = torch.cat([latents, image]			,	dim=1   )
      __lowerCamelCase       :       Union[str, Any]           = self.scheduler.scale_model_input(_a			,	_a   )
      # predict the noise residual
      __lowerCamelCase       :       Optional[int]           = self.unet(_a			,	_a   ).sample
      # compute the previous noisy sample x_t -> x_t-1
      __lowerCamelCase       :       List[Any]           = self.scheduler.step(_a			,	_a			,	_a			,	**_a   ).prev_sample

    # decode the image latents with the VQVAE
    __lowerCamelCase       :       Union[str, Any]           = self.vqvae.decode(_a   ).sample
    __lowerCamelCase       :       Union[str, Any]           = torch.clamp(_a			,	-1.0			,	1.0   )
    __lowerCamelCase       :       List[str]           = image / 2 + 0.5
    __lowerCamelCase       :       Optional[Any]           = image.cpu().permute(0			,	2			,	3			,	1   ).numpy()

    if output_type == "pil":
      __lowerCamelCase       :       Tuple           = self.numpy_to_pil(_a   )

    if not return_dict:
      return (image,)

    return ImagePipelineOutput(images=_a   )






