





# 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      _SCREAMING_SNAKE_CASE		(			__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase					,							unittest.TestCase					):







 '''simple docstring'''
 lowercase_	=    StableDiffusionControlNetImgaImgPipeline
 lowercase_	=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
 lowercase_	=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowercase_	=    IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}					)
 lowercase_	=    IMAGE_TO_IMAGE_IMAGE_PARAMS

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->Dict:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:   Dict    						=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)
   lowerCamelCase__:   Optional[int]    						=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)
   lowerCamelCase__:   Optional[Any]    						=DDIMScheduler(
       beta_start=0.0_0085 ,				beta_end=0.012 ,				beta_schedule="scaled_linear" ,				clip_sample=lowerCamelCase__ ,				set_alpha_to_one=lowerCamelCase__ ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   List[str]    						=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)
   lowerCamelCase__:   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=1_000 ,				)
   lowerCamelCase__:   Any    						=CLIPTextModel(lowerCamelCase__)
   lowerCamelCase__:   int    						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:   List[str]    						={
       '''unet''': unet,
       '''controlnet''': controlnet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_      (self       :     Any ,				UpperCAmelCase_       :     Union[str, Any] ,				UpperCAmelCase_       :     List[Any]=0) ->List[Any]:



   '''simple docstring'''



   if str(lowerCamelCase__).startswith("mps"):
     lowerCamelCase__:   Tuple    						=torch.manual_seed(lowerCamelCase__)
   else:
     lowerCamelCase__:   Dict    						=torch.Generator(device=lowerCamelCase__).manual_seed(lowerCamelCase__)

   lowerCamelCase__:   Tuple    						=2
   lowerCamelCase__:   Dict    						=randn_tensor(
       (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor) ,				generator=lowerCamelCase__ ,				device=torch.device(lowerCamelCase__) ,				)
   lowerCamelCase__:   List[Any]    						=floats_tensor(control_image.shape ,				rng=random.Random(lowerCamelCase__)).to(lowerCamelCase__)
   lowerCamelCase__:   Tuple    						=image.cpu().permute(0 ,				2 ,				3 ,				1)[0]
   lowerCamelCase__:   List[Any]    						=Image.fromarray(np.uinta(lowerCamelCase__)).convert("RGB").resize((64, 64))
   lowerCamelCase__:   List[Any]    						={
       '''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 							SCREAMING_SNAKE_CASE_      (self       :     Union[str, Any]) ->List[Any]:



   '''simple docstring'''



   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 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->Union[str, Any]:



   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3)







 def 							SCREAMING_SNAKE_CASE_      (self       :     List[Any]) ->List[Any]:



   '''simple docstring'''



   self._test_inference_batch_single_identical(expected_max_diff=2E-3)



class      _SCREAMING_SNAKE_CASE		(			__lowerCAmelCase					,							__lowerCAmelCase					,							unittest.TestCase					):







 '''simple docstring'''
 lowercase_	=    StableDiffusionControlNetImgaImgPipeline
 lowercase_	=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
 lowercase_	=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowercase_	=    frozenset([]					)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

 def 							SCREAMING_SNAKE_CASE_      (self       :     Tuple) ->Dict:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:   Any    						=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(UpperCAmelCase_       :     int):
     if isinstance(lowerCamelCase__ ,				torch.nn.Convad):
       torch.nn.init.normal(m.weight)
       m.bias.data.fill_(1.0)

   lowerCamelCase__:   Union[str, Any]    						=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(lowerCamelCase__)

   torch.manual_seed(0)
   lowerCamelCase__:   Union[str, Any]    						=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(lowerCamelCase__)

   torch.manual_seed(0)
   lowerCamelCase__:   int    						=DDIMScheduler(
       beta_start=0.0_0085 ,				beta_end=0.012 ,				beta_schedule="scaled_linear" ,				clip_sample=lowerCamelCase__ ,				set_alpha_to_one=lowerCamelCase__ ,				)
   torch.manual_seed(0)
   lowerCamelCase__:   List[Any]    						=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)
   lowerCamelCase__:   Union[str, Any]    						=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=1_000 ,				)
   lowerCamelCase__:   Union[str, Any]    						=CLIPTextModel(lowerCamelCase__)
   lowerCamelCase__:   Optional[int]    						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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

   lowerCamelCase__:   List[str]    						={
       '''unet''': unet,
       '''controlnet''': controlnet,
       '''scheduler''': scheduler,
       '''vae''': vae,
       '''text_encoder''': text_encoder,
       '''tokenizer''': tokenizer,
       '''safety_checker''': None,
       '''feature_extractor''': None,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict ,				UpperCAmelCase_       :     Optional[int] ,				UpperCAmelCase_       :     Tuple=0) ->Dict:



   '''simple docstring'''



   if str(lowerCamelCase__).startswith("mps"):
     lowerCamelCase__:   Tuple    						=torch.manual_seed(lowerCamelCase__)
   else:
     lowerCamelCase__:   Union[str, Any]    						=torch.Generator(device=lowerCamelCase__).manual_seed(lowerCamelCase__)

   lowerCamelCase__:   str    						=2

   lowerCamelCase__:   Tuple    						=[
       randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor) ,				generator=lowerCamelCase__ ,				device=torch.device(lowerCamelCase__) ,				),
       randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor) ,				generator=lowerCamelCase__ ,				device=torch.device(lowerCamelCase__) ,				),
   ]

   lowerCamelCase__:   List[str]    						=floats_tensor(control_image[0].shape ,				rng=random.Random(lowerCamelCase__)).to(lowerCamelCase__)
   lowerCamelCase__:   Tuple    						=image.cpu().permute(0 ,				2 ,				3 ,				1)[0]
   lowerCamelCase__:   Optional[int]    						=Image.fromarray(np.uinta(lowerCamelCase__)).convert("RGB").resize((64, 64))
   lowerCamelCase__:   List[Any]    						={
       '''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 							SCREAMING_SNAKE_CASE_      (self       :     int) ->int:



   '''simple docstring'''



   lowerCamelCase__:   List[Any]    						=self.get_dummy_components()
   lowerCamelCase__:   List[Any]    						=self.pipeline_class(**lowerCamelCase__)
   pipe.to(lowerCamelCase__)

   lowerCamelCase__:   Union[str, Any]    						=10.0
   lowerCamelCase__:   Optional[int]    						=4

   lowerCamelCase__:   Any    						=self.get_dummy_inputs(lowerCamelCase__)
   lowerCamelCase__:   Dict    						=steps
   lowerCamelCase__:   Optional[Any]    						=scale
   lowerCamelCase__:   Union[str, Any]    						=pipe(**lowerCamelCase__)[0]

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

   lowerCamelCase__:   Dict    						=self.get_dummy_inputs(lowerCamelCase__)
   lowerCamelCase__:   Dict    						=steps
   lowerCamelCase__:   Optional[Any]    						=scale
   lowerCamelCase__:   List[Any]    						=pipe(**lowerCamelCase__ ,				control_guidance_start=[0.1, 0.3] ,				control_guidance_end=[0.2, 0.7])[0]

   lowerCamelCase__:   Tuple    						=self.get_dummy_inputs(lowerCamelCase__)
   lowerCamelCase__:   Optional[int]    						=steps
   lowerCamelCase__:   Optional[Any]    						=scale
   lowerCamelCase__:   Optional[int]    						=pipe(**lowerCamelCase__ ,				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 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->Tuple:



   '''simple docstring'''



   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 							SCREAMING_SNAKE_CASE_      (self       :     Union[str, Any]) ->List[Any]:



   '''simple docstring'''



   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3)

 def 							SCREAMING_SNAKE_CASE_      (self       :     Any) ->List[Any]:



   '''simple docstring'''



   self._test_inference_batch_single_identical(expected_max_diff=2E-3)







 def 							SCREAMING_SNAKE_CASE_      (self       :     Tuple) ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:   Any    						=self.get_dummy_components()
   lowerCamelCase__:   int    						=self.pipeline_class(**lowerCamelCase__)
   pipe.to(lowerCamelCase__)
   pipe.set_progress_bar_config(disable=lowerCamelCase__)
   with tempfile.TemporaryDirectory() as tmpdir:
     try:
       # save_pretrained is not implemented for Multi-ControlNet
       pipe.save_pretrained(lowerCamelCase__)
     except NotImplementedError:
       pass



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







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_      (self       :     List[str]) ->List[Any]:



   '''simple docstring'''



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







 def 							SCREAMING_SNAKE_CASE_      (self       :     Dict) ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:   str    						=ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")

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

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

   lowerCamelCase__:   Tuple    						=pipe(
       lowerCamelCase__ ,				lowerCamelCase__ ,				control_image=lowerCamelCase__ ,				generator=lowerCamelCase__ ,				output_type="np" ,				num_inference_steps=50 ,				strength=0.6 ,				)

   lowerCamelCase__:   Dict    						=output.images[0]

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

   lowerCamelCase__:   Optional[Any]    						=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








from bisect import bisect
from itertools import accumulate





def   lowerCAmelCase_				( __a			,					__a			,					__a			,					__a	)  ->	List[Any]:







  """simple docstring"""





  lowerCamelCase__:   str    						=sorted(zip(__a			,					__a	)			,					key=lambda __a	: x[0] / x[1]			,					reverse=__a	)
  lowerCamelCase__						,      lowerCamelCase__:   str    						=[i[0] for i in r], [i[1] for i in r]
  lowerCamelCase__:   Tuple    						=list(accumulate(__a	)	)
  lowerCamelCase__:   Any    						=bisect(__a			,					__a	)
  return (
      0
      if k == 0
      else sum(vl[:k]	) + (w - acc[k - 1]) * (vl[k]) / (wt[k])
      if k != n
      else sum(vl[:k]	)
  )


if __name__ == "__main__":
     import doctest

     doctest.testmod()

