


"""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      (      SCREAMING_SNAKE_CASE_		,		SCREAMING_SNAKE_CASE_		,		SCREAMING_SNAKE_CASE_		,		unittest.TestCase ):
    a_    :				List[str]             =    StableDiffusionControlNetImgaImgPipeline
    a_    :				Optional[int]             =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    a_    :				List[Any]             =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a_    :				Optional[Any]             =    IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'} )
    a_    :				List[Any]             =    IMAGE_TO_IMAGE_IMAGE_PARAMS
    def      _UpperCamelCase       ( self							:					Optional[int]				):
          torch.manual_seed(0				)
          lowerCamelCase__           =		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__           =		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__           =		DDIMScheduler(
              beta_start=0.0_00_85							,					beta_end=0.0_12							,					beta_schedule='scaled_linear'							,					clip_sample=SCREAMING_SNAKE_CASE__							,					set_alpha_to_one=SCREAMING_SNAKE_CASE__							,					)
          torch.manual_seed(0				)
          lowerCamelCase__           =		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__           =		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=10_00							,					)
          lowerCamelCase__           =		CLIPTextModel(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

          lowerCamelCase__           =		{
              'unet': unet,
              'controlnet': controlnet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'safety_checker': None,
              'feature_extractor': None,
          }
          return components
    def      _UpperCamelCase       ( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Optional[int]							,					SCREAMING_SNAKE_CASE__							:					Optional[int]=0				):
          if str(SCREAMING_SNAKE_CASE__				).startswith('mps'				):
                lowerCamelCase__           =		torch.manual_seed(SCREAMING_SNAKE_CASE__				)
          else:
                lowerCamelCase__           =		torch.Generator(device=SCREAMING_SNAKE_CASE__				).manual_seed(SCREAMING_SNAKE_CASE__				)

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


    def      _UpperCamelCase       ( self							:					Union[str, Any]				):
          self._test_inference_batch_single_identical(expected_max_diff=2e-3				)





class   _a      (      SCREAMING_SNAKE_CASE_		,		SCREAMING_SNAKE_CASE_		,		unittest.TestCase ):
    a_    :				List[Any]             =    StableDiffusionControlNetImgaImgPipeline
    a_    :				List[Any]             =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    a_    :				str             =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a_    :				int             =    frozenset([] )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess
    def      _UpperCamelCase       ( self							:					List[str]				):
          torch.manual_seed(0				)
          lowerCamelCase__           =		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(SCREAMING_SNAKE_CASE__							:					Any				):
                if isinstance(SCREAMING_SNAKE_CASE__							,					torch.nn.Convad				):
                      torch.nn.init.normal(m.weight				)
                      m.bias.data.fill_(1.0				)

          lowerCamelCase__           =		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(SCREAMING_SNAKE_CASE__				)

          torch.manual_seed(0				)
          lowerCamelCase__           =		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(SCREAMING_SNAKE_CASE__				)

          torch.manual_seed(0				)
          lowerCamelCase__           =		DDIMScheduler(
              beta_start=0.0_00_85							,					beta_end=0.0_12							,					beta_schedule='scaled_linear'							,					clip_sample=SCREAMING_SNAKE_CASE__							,					set_alpha_to_one=SCREAMING_SNAKE_CASE__							,					)
          torch.manual_seed(0				)
          lowerCamelCase__           =		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__           =		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=10_00							,					)
          lowerCamelCase__           =		CLIPTextModel(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

          lowerCamelCase__           =		MultiControlNetModel([controlneta, controlneta]				)

          lowerCamelCase__           =		{
              'unet': unet,
              'controlnet': controlnet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'safety_checker': None,
              'feature_extractor': None,
          }
          return components
    def      _UpperCamelCase       ( self							:					Dict							,					SCREAMING_SNAKE_CASE__							:					Optional[int]							,					SCREAMING_SNAKE_CASE__							:					Any=0				):
          if str(SCREAMING_SNAKE_CASE__				).startswith('mps'				):
                lowerCamelCase__           =		torch.manual_seed(SCREAMING_SNAKE_CASE__				)
          else:
                lowerCamelCase__           =		torch.Generator(device=SCREAMING_SNAKE_CASE__				).manual_seed(SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		2

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

          lowerCamelCase__           =		floats_tensor(control_image[0].shape							,					rng=random.Random(SCREAMING_SNAKE_CASE__				)				).to(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		image.cpu().permute(0							,					2							,					3							,					1				)[0]
          lowerCamelCase__           =		Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE__				)				).convert('RGB'				).resize((64, 64)				)
          lowerCamelCase__           =		{
              '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      _UpperCamelCase       ( self							:					Optional[int]				):
          lowerCamelCase__           =		self.get_dummy_components()
          lowerCamelCase__           =		self.pipeline_class(**SCREAMING_SNAKE_CASE__				)
          pipe.to(SCREAMING_SNAKE_CASE__				)

          lowerCamelCase__           =		10.0
          lowerCamelCase__           =		4

          lowerCamelCase__           =		self.get_dummy_inputs(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		steps
          lowerCamelCase__           =		scale
          lowerCamelCase__           =		pipe(**SCREAMING_SNAKE_CASE__				)[0]

          lowerCamelCase__           =		self.get_dummy_inputs(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		steps
          lowerCamelCase__           =		scale
          lowerCamelCase__           =		pipe(**SCREAMING_SNAKE_CASE__							,					control_guidance_start=0.1							,					control_guidance_end=0.2				)[0]

          lowerCamelCase__           =		self.get_dummy_inputs(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		steps
          lowerCamelCase__           =		scale
          lowerCamelCase__           =		pipe(**SCREAMING_SNAKE_CASE__							,					control_guidance_start=[0.1, 0.3]							,					control_guidance_end=[0.2, 0.7]				)[0]

          lowerCamelCase__           =		self.get_dummy_inputs(SCREAMING_SNAKE_CASE__				)
          lowerCamelCase__           =		steps
          lowerCamelCase__           =		scale
          lowerCamelCase__           =		pipe(**SCREAMING_SNAKE_CASE__							,					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      _UpperCamelCase       ( self							:					Optional[int]				):
          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      _UpperCamelCase       ( self							:					str				):
          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3				)
    def      _UpperCamelCase       ( self							:					List[Any]				):
          self._test_inference_batch_single_identical(expected_max_diff=2e-3				)


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





@slow
@require_torch_gpu
class   _a      (      unittest.TestCase ):
    def      _UpperCamelCase       ( self							:					int				):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


    def      _UpperCamelCase       ( self							:					Tuple				):
          lowerCamelCase__           =		ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'				)

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

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

          lowerCamelCase__           =		pipe(
              SCREAMING_SNAKE_CASE__							,					SCREAMING_SNAKE_CASE__							,					control_image=SCREAMING_SNAKE_CASE__							,					generator=SCREAMING_SNAKE_CASE__							,					output_type='np'							,					num_inference_steps=50							,					strength=0.6							,					)

          lowerCamelCase__           =		output.images[0]

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

          lowerCamelCase__           =		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"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_snake_case          =							{
    "configuration_m2m_100": ["M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP", "M2M100Config", "M2M100OnnxConfig"],
    "tokenization_m2m_100": ["M2M100Tokenizer"],
}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _snake_case          =							[
       "M2M_100_PRETRAINED_MODEL_ARCHIVE_LIST",
       "M2M100ForConditionalGeneration",
       "M2M100Model",
       "M2M100PreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_mam_aaa import M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP, MaMaaaConfig, MaMaaaOnnxConfig
   from .tokenization_mam_aaa import MaMaaaTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mam_aaa import (
          M2M_100_PRETRAINED_MODEL_ARCHIVE_LIST,
          MaMaaaForConditionalGeneration,
          MaMaaaModel,
          MaMaaaPreTrainedModel,
      )


else:
   import sys

   _snake_case          =							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



