

'''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__  (							lowerCamelCase_  ,				lowerCamelCase_  ,				lowerCamelCase_  ,				unittest.TestCase  ):
   _SCREAMING_SNAKE_CASE   :					Any     	=							StableDiffusionControlNetImgaImgPipeline
   _SCREAMING_SNAKE_CASE   :					Optional[int]     	=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
   _SCREAMING_SNAKE_CASE   :					Any     	=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   _SCREAMING_SNAKE_CASE   :					int     	=							IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'}  )
   _SCREAMING_SNAKE_CASE   :					Any     	=							IMAGE_TO_IMAGE_IMAGE_PARAMS


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





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

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


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase=0  ):





      """simple docstring"""





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

      _lowercase      :				List[Any]      					=       2
      _lowercase      :				int      					=       randn_tensor(
          (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)      ,    generator=_UpperCamelCase      ,    device=torch.device(_UpperCamelCase  )      ,    )
      _lowercase      :				int      					=       floats_tensor(control_image.shape      ,    rng=random.Random(_UpperCamelCase  )  ).to(_UpperCamelCase  )
      _lowercase      :				List[Any]      					=       image.cpu().permute(0      ,    2      ,    3      ,    1  )[0]
      _lowercase      :				Optional[int]      					=       Image.fromarray(np.uinta(_UpperCamelCase  )  ).convert("RGB"  ).resize((64, 64)  )
      _lowercase      :				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 						_lowerCamelCase       (			self  ):





      """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 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )






   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      self._test_inference_batch_single_identical(expected_max_diff=2E-3  )




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


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      torch.manual_seed(0  )
      _lowercase      :				List[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  ):
         if isinstance(_UpperCamelCase      ,    torch.nn.Convad  ):
            torch.nn.init.normal(m.weight  )
            m.bias.data.fill_(1.0  )

      _lowercase      :				Tuple      					=       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(_UpperCamelCase  )

      torch.manual_seed(0  )
      _lowercase      :				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)      ,    )
      controlneta.controlnet_down_blocks.apply(_UpperCamelCase  )

      torch.manual_seed(0  )
      _lowercase      :				Union[str, Any]      					=       DDIMScheduler(
          beta_start=0.0_0_0_8_5      ,    beta_end=0.0_1_2      ,    beta_schedule="scaled_linear"      ,    clip_sample=_UpperCamelCase      ,    set_alpha_to_one=_UpperCamelCase      ,    )
      torch.manual_seed(0  )
      _lowercase      :				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  )
      _lowercase      :				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=1000      ,    )
      _lowercase      :				Union[str, Any]      					=       CLIPTextModel(_UpperCamelCase  )
      _lowercase      :				Any      					=       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

      _lowercase      :				Optional[Any]      					=       MultiControlNetModel([controlneta, controlneta]  )

      _lowercase      :				List[Any]      					=       {
          "unet": unet,
          "controlnet": controlnet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      return components


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase=0  ):





      """simple docstring"""





      if str(_UpperCamelCase  ).startswith("mps"  ):
         _lowercase      :				Optional[Any]      					=       torch.manual_seed(_UpperCamelCase  )
      else:
         _lowercase      :				Optional[Any]      					=       torch.Generator(device=_UpperCamelCase  ).manual_seed(_UpperCamelCase  )

      _lowercase      :				Any      					=       2

      _lowercase      :				Any      					=       [
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)      ,    generator=_UpperCamelCase      ,    device=torch.device(_UpperCamelCase  )      ,    ),
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)      ,    generator=_UpperCamelCase      ,    device=torch.device(_UpperCamelCase  )      ,    ),
      ]

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





      """simple docstring"""





      _lowercase      :				List[Any]      					=       self.get_dummy_components()
      _lowercase      :				Optional[int]      					=       self.pipeline_class(**_UpperCamelCase  )
      pipe.to(_UpperCamelCase  )

      _lowercase      :				int      					=       1_0.0
      _lowercase      :				Optional[int]      					=       4

      _lowercase      :				Dict      					=       self.get_dummy_inputs(_UpperCamelCase  )
      _lowercase      :				List[Any]      					=       steps
      _lowercase      :				List[Any]      					=       scale
      _lowercase      :				str      					=       pipe(**_UpperCamelCase  )[0]

      _lowercase      :				Any      					=       self.get_dummy_inputs(_UpperCamelCase  )
      _lowercase      :				Tuple      					=       steps
      _lowercase      :				str      					=       scale
      _lowercase      :				Optional[int]      					=       pipe(**_UpperCamelCase      ,    control_guidance_start=0.1      ,    control_guidance_end=0.2  )[0]

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

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





      """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 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      self._test_inference_batch_single_identical(expected_max_diff=2E-3  )






   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





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




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


   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





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






   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      _lowercase      :				Dict      					=       ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"  )

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

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

      _lowercase      :				List[Any]      					=       pipe(
          _UpperCamelCase      ,    _UpperCamelCase      ,    control_image=_UpperCamelCase      ,    generator=_UpperCamelCase      ,    output_type="np"      ,    num_inference_steps=50      ,    strength=0.6      ,    )

      _lowercase      :				Optional[int]      					=       output.images[0]

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

      _lowercase      :				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 warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer


class        a__  (							lowerCamelCase_  ):
   _SCREAMING_SNAKE_CASE   :					int     	=							'Wav2Vec2FeatureExtractor'
   _SCREAMING_SNAKE_CASE   :					List[str]     	=							'AutoTokenizer'


   def __init__(			self      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      super().__init__(_UpperCamelCase      ,    _UpperCamelCase  )
      _lowercase      :				List[Any]      					=       self.feature_extractor
      _lowercase      :				Optional[Any]      					=       False


   @classmethod
   def 						_lowerCamelCase       (			cls      ,    _UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      try:
         return super().from_pretrained(_UpperCamelCase      ,    **_UpperCamelCase  )
      except OSError:
         warnings.warn(
             f'''Loading a tokenizer inside {cls.__name__} from a config that does not'''
             " include a `tokenizer_class` attribute is deprecated and will be "
             "removed in v5. Please add `'tokenizer_class': 'Wav2Vec2CTCTokenizer'`"
             " attribute to either your `config.json` or `tokenizer_config.json` "
             "file to suppress this warning: "      ,    _UpperCamelCase      ,    )

         _lowercase      :				Dict      					=       WavaVecaFeatureExtractor.from_pretrained(_UpperCamelCase      ,    **_UpperCamelCase  )
         _lowercase      :				str      					=       WavaVecaCTCTokenizer.from_pretrained(_UpperCamelCase      ,    **_UpperCamelCase  )

         return cls(feature_extractor=_UpperCamelCase      ,    tokenizer=_UpperCamelCase  )


   def __call__(			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      if self._in_target_context_manager:
         return self.current_processor(*_UpperCamelCase      ,    **_UpperCamelCase  )

      if "raw_speech" in kwargs:
         warnings.warn("Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead."  )
         _lowercase      :				int      					=       kwargs.pop("raw_speech"  )
      else:
         _lowercase      :				List[Any]      					=       kwargs.pop("audio"      ,    _UpperCamelCase  )
      _lowercase      :				List[Any]      					=       kwargs.pop("sampling_rate"      ,    _UpperCamelCase  )
      _lowercase      :				Union[str, Any]      					=       kwargs.pop("text"      ,    _UpperCamelCase  )
      if len(_UpperCamelCase  ) > 0:
         _lowercase      :				int      					=       args[0]
         _lowercase      :				Any      					=       args[1:]

      if audio is None and text is None:
         raise ValueError("You need to specify either an `audio` or `text` input to process."  )

      if audio is not None:
         _lowercase      :				Dict      					=       self.feature_extractor(_UpperCamelCase      ,    *_UpperCamelCase      ,    sampling_rate=_UpperCamelCase      ,    **_UpperCamelCase  )
      if text is not None:
         _lowercase      :				Union[str, Any]      					=       self.tokenizer(_UpperCamelCase      ,    **_UpperCamelCase  )

      if text is None:
         return inputs
      elif audio is None:
         return encodings
      else:
         _lowercase      :				int      					=       encodings["input_ids"]
         return inputs


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      if self._in_target_context_manager:
         return self.current_processor.pad(*_UpperCamelCase      ,    **_UpperCamelCase  )

      _lowercase      :				List[Any]      					=       kwargs.pop("input_features"      ,    _UpperCamelCase  )
      _lowercase      :				Any      					=       kwargs.pop("labels"      ,    _UpperCamelCase  )
      if len(_UpperCamelCase  ) > 0:
         _lowercase      :				Any      					=       args[0]
         _lowercase      :				Any      					=       args[1:]

      if input_features is not None:
         _lowercase      :				Any      					=       self.feature_extractor.pad(_UpperCamelCase      ,    *_UpperCamelCase      ,    **_UpperCamelCase  )
      if labels is not None:
         _lowercase      :				int      					=       self.tokenizer.pad(_UpperCamelCase      ,    **_UpperCamelCase  )

      if labels is None:
         return input_features
      elif input_features is None:
         return labels
      else:
         _lowercase      :				Optional[Any]      					=       labels["input_ids"]
         return input_features


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.tokenizer.batch_decode(*_UpperCamelCase      ,    **_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.tokenizer.decode(*_UpperCamelCase      ,    **_UpperCamelCase  )






   @contextmanager
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      warnings.warn(
          "`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your "
          "labels by using the argument `text` of the regular `__call__` method (either in the same call as "
          "your audio inputs, or in a separate call."  )
      _lowercase      :				Optional[Any]      					=       True
      _lowercase      :				Dict      					=       self.tokenizer
      yield
      _lowercase      :				List[str]      					=       self.feature_extractor
      _lowercase      :				List[str]      					=       False






