



"""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  _UpperCAmelCase (    lowerCAmelCase__					,  lowerCAmelCase__					,  lowerCAmelCase__					,  unittest.TestCase):
 _lowerCAmelCase     :    Any	      =	StableDiffusionControlNetImgaImgPipeline
 _lowerCAmelCase     :    Dict	      =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 _lowerCAmelCase     :    List[str]	      =	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 _lowerCAmelCase     :    str	      =	IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""})
 _lowerCAmelCase     :    int	      =	IMAGE_TO_IMAGE_IMAGE_PARAMS

 def 		_snake_case      (  self	: Optional[int]       ):
     torch.manual_seed(0       )
     snake_case_				:  List[str]       =			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       )
     snake_case_				:  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       )
     snake_case_				:  Tuple       =			DDIMScheduler(
         beta_start=0.0_00_85						,  beta_end=0.0_12						,  beta_schedule='''scaled_linear'''						,  clip_sample=lowercase_						,  set_alpha_to_one=lowercase_						,  )
     torch.manual_seed(0       )
     snake_case_				:  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       )
     snake_case_				:  str       =			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						,  )
     snake_case_				:  int       =			CLIPTextModel(lowercase_       )
     snake_case_				:  Optional[int]       =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

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

 def 		_snake_case      (  self	: List[Any]						,  lowercase_	: Optional[int]						,  lowercase_	: Union[str, Any]=0       ):
     if str(lowercase_       ).startswith('''mps'''       ):
         snake_case_				:  Tuple       =			torch.manual_seed(lowercase_       )
     else:
         snake_case_				:  Tuple       =			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )

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






 def 		_snake_case      (  self	: Tuple       ):
     self._test_inference_batch_single_identical(expected_max_diff=2E-3       )





class  _UpperCAmelCase (    lowerCAmelCase__					,  lowerCAmelCase__					,  unittest.TestCase):
 _lowerCAmelCase     :    Optional[int]	      =	StableDiffusionControlNetImgaImgPipeline
 _lowerCAmelCase     :    str	      =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 _lowerCAmelCase     :    Tuple	      =	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 _lowerCAmelCase     :    List[str]	      =	frozenset([])  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

 def 		_snake_case      (  self	: Union[str, Any]       ):
     torch.manual_seed(0       )
     snake_case_				:  Tuple       =			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(lowercase_	: Tuple       ):
         if isinstance(lowercase_						,  torch.nn.Convad       ):
             torch.nn.init.normal(m.weight       )
             m.bias.data.fill_(1.0       )

     snake_case_				:  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(lowercase_       )

     torch.manual_seed(0       )
     snake_case_				:  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(lowercase_       )

     torch.manual_seed(0       )
     snake_case_				:  Optional[Any]       =			DDIMScheduler(
         beta_start=0.0_00_85						,  beta_end=0.0_12						,  beta_schedule='''scaled_linear'''						,  clip_sample=lowercase_						,  set_alpha_to_one=lowercase_						,  )
     torch.manual_seed(0       )
     snake_case_				:  Union[str, 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       )
     snake_case_				:  List[Any]       =			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						,  )
     snake_case_				:  Union[str, Any]       =			CLIPTextModel(lowercase_       )
     snake_case_				:  List[Any]       =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

     snake_case_				:  Dict       =			MultiControlNetModel([controlneta, controlneta]       )

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

 def 		_snake_case      (  self	: Optional[Any]						,  lowercase_	: List[Any]						,  lowercase_	: Tuple=0       ):
     if str(lowercase_       ).startswith('''mps'''       ):
         snake_case_				:  str       =			torch.manual_seed(lowercase_       )
     else:
         snake_case_				:  Optional[int]       =			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )

     snake_case_				:  List[str]       =			2

     snake_case_				:  List[str]       =			[
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,  generator=lowercase_						,  device=torch.device(lowercase_       )						,  ),
         randn_tensor(
             (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)						,  generator=lowercase_						,  device=torch.device(lowercase_       )						,  ),
     ]

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

     snake_case_				:  List[Any]       =			10.0
     snake_case_				:  Optional[Any]       =			4

     snake_case_				:  str       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  List[str]       =			steps
     snake_case_				:  str       =			scale
     snake_case_				:  Dict       =			pipe(**lowercase_       )[0]

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

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

     snake_case_				:  List[str]       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Dict       =			steps
     snake_case_				:  Optional[int]       =			scale
     snake_case_				:  Optional[Any]       =			pipe(**lowercase_						,  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 		_snake_case      (  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 		_snake_case      (  self	: Optional[int]       ):
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3       )

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






 def 		_snake_case      (  self	: Union[str, Any]       ):
     snake_case_				:  Any       =			self.get_dummy_components()
     snake_case_				:  List[str]       =			self.pipeline_class(**lowercase_       )
     pipe.to(lowercase_       )
     pipe.set_progress_bar_config(disable=lowercase_       )
     with tempfile.TemporaryDirectory() as tmpdir:
         try:
             # save_pretrained is not implemented for Multi-ControlNet
             pipe.save_pretrained(lowercase_       )
         except NotImplementedError:
             pass





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

 def 		_snake_case      (  self	: int       ):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






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

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

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

     snake_case_				:  List[str]       =			pipe(
         lowercase_						,  lowercase_						,  control_image=lowercase_						,  generator=lowercase_						,  output_type='''np'''						,  num_inference_steps=50						,  strength=0.6						,  )

     snake_case_				:  Union[str, Any]       =			output.images[0]

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

     snake_case_				:  int       =			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 sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / '''utils'''))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


lowercase__					:						Any												=			get_tests_dir('''fixtures''')




class  _UpperCAmelCase (    unittest.TestCase):

 def 		_snake_case      (  self	: Tuple       ):
     # A mock response for an HTTP head request to emulate server down
     snake_case_				:  str       =			mock.Mock()
     snake_case_				:  Optional[Any]       =			500
     snake_case_				:  str       =			{}
     snake_case_				:  Optional[int]       =			HTTPError
     snake_case_				:  Tuple       =			{}

     # Download this model to make sure it's in the cache.
     snake_case_				:  List[str]       =			WavaVecaFeatureExtractor.from_pretrained('''hf-internal-testing/tiny-random-wav2vec2'''       )
     # Under the mock environment we get a 500 error when trying to reach the model.
     with mock.patch('''requests.Session.request'''						,  return_value=lowercase_       ) as mock_head:
         snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained('''hf-internal-testing/tiny-random-wav2vec2'''       )
         # This check we did call the fake head request
         mock_head.assert_called()






 def 		_snake_case      (  self	: Union[str, Any]       ):
     # This test is for deprecated behavior and can be removed in v5
     snake_case_				:  str       =			WavaVecaFeatureExtractor.from_pretrained(
         '''https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json'''       )





@is_staging_test
class  _UpperCAmelCase (    unittest.TestCase):

 @classmethod
 def 		_snake_case      (  cls	: int       ):
     snake_case_				:  Dict       =			TOKEN
     HfFolder.save_token(lowercase_       )

 @classmethod
 def 		_snake_case      (  cls	: str       ):
     try:
         delete_repo(token=cls._token						,  repo_id='''test-feature-extractor'''       )
     except HTTPError:
         pass

     try:
         delete_repo(token=cls._token						,  repo_id='''valid_org/test-feature-extractor-org'''       )
     except HTTPError:
         pass

     try:
         delete_repo(token=cls._token						,  repo_id='''test-dynamic-feature-extractor'''       )
     except HTTPError:
         pass

 def 		_snake_case      (  self	: Any       ):
     snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained(lowercase_       )
     feature_extractor.push_to_hub('''test-feature-extractor'''						,  use_auth_token=self._token       )

     snake_case_				:  Optional[Any]       =			WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

     # Reset repo
     delete_repo(token=self._token						,  repo_id='''test-feature-extractor'''       )

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(
             lowercase_						,  repo_id='''test-feature-extractor'''						,  push_to_hub=lowercase_						,  use_auth_token=self._token       )

     snake_case_				:  Optional[int]       =			WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

 def 		_snake_case      (  self	: Dict       ):
     snake_case_				:  Union[str, Any]       =			WavaVecaFeatureExtractor.from_pretrained(lowercase_       )
     feature_extractor.push_to_hub('''valid_org/test-feature-extractor'''						,  use_auth_token=self._token       )

     snake_case_				:  Optional[int]       =			WavaVecaFeatureExtractor.from_pretrained('''valid_org/test-feature-extractor'''       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

     # Reset repo
     delete_repo(token=self._token						,  repo_id='''valid_org/test-feature-extractor'''       )

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(
             lowercase_						,  repo_id='''valid_org/test-feature-extractor-org'''						,  push_to_hub=lowercase_						,  use_auth_token=self._token       )

     snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained('''valid_org/test-feature-extractor-org'''       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )






 def 		_snake_case      (  self	: str       ):
     CustomFeatureExtractor.register_for_auto_class()
     snake_case_				:  Optional[Any]       =			CustomFeatureExtractor.from_pretrained(lowercase_       )

     feature_extractor.push_to_hub('''test-dynamic-feature-extractor'''						,  use_auth_token=self._token       )

     # This has added the proper auto_map field to the config
     self.assertDictEqual(
         feature_extractor.auto_map						,  {'''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor'''}						,  )

     snake_case_				:  Any       =			AutoFeatureExtractor.from_pretrained(
         f"{USER}/test-dynamic-feature-extractor"						,  trust_remote_code=lowercase_       )
     # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
     self.assertEqual(new_feature_extractor.__class__.__name__						,  '''CustomFeatureExtractor'''       )





