






"""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 			lowerCAmelCase__			(    __magic_name__							,						__magic_name__							,						__magic_name__							,						unittest.TestCase   ):
   SCREAMING_SNAKE_CASE_    =StableDiffusionControlNetImgaImgPipeline
   SCREAMING_SNAKE_CASE_    =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
   SCREAMING_SNAKE_CASE_    =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   SCREAMING_SNAKE_CASE_    =IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}   )
   SCREAMING_SNAKE_CASE_    =IMAGE_TO_IMAGE_IMAGE_PARAMS
   def       __a							(		self      :	Optional[Any]							):


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

        UpperCAmelCase__						:    Optional[Any]							=							{
            "unet": unet,
            "controlnet": controlnet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components
   def       __a							(		self      :	Optional[int]							,			snake_case__      :	Union[str, Any]							,			snake_case__      :	Optional[int]=0							):


        '''simple docstring'''
        if str(snake_case__							).startswith("mps"							):
             UpperCAmelCase__						:    Optional[Any]							=							torch.manual_seed(snake_case__							)
        else:
             UpperCAmelCase__						:    str							=							torch.Generator(device=snake_case__							).manual_seed(snake_case__							)

        UpperCAmelCase__						:    Dict							=							2
        UpperCAmelCase__						:    Any							=							randn_tensor(
            (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)							,			generator=snake_case__							,			device=torch.device(snake_case__							)							,			)
        UpperCAmelCase__						:    int							=							floats_tensor(control_image.shape							,			rng=random.Random(snake_case__							)							).to(snake_case__							)
        UpperCAmelCase__						:    Union[str, Any]							=							image.cpu().permute(0							,			2							,			3							,			1							)[0]
        UpperCAmelCase__						:    List[str]							=							Image.fromarray(np.uinta(snake_case__							)							).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       __a							(		self      :	Optional[int]							):


        '''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       __a							(		self      :	List[str]							):


        '''simple docstring'''
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3							)




   def       __a							(		self      :	Union[str, Any]							):


        '''simple docstring'''
        self._test_inference_batch_single_identical(expected_max_diff=2e-3							)




class 			lowerCAmelCase__			(    __magic_name__							,						__magic_name__							,						unittest.TestCase   ):
   SCREAMING_SNAKE_CASE_    =StableDiffusionControlNetImgaImgPipeline
   SCREAMING_SNAKE_CASE_    =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
   SCREAMING_SNAKE_CASE_    =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   SCREAMING_SNAKE_CASE_    =frozenset([]   )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess
   def       __a							(		self      :	Tuple							):


        '''simple docstring'''
        torch.manual_seed(0							)
        UpperCAmelCase__						:    List[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(snake_case__      :	Optional[int]							):
             if isinstance(snake_case__							,			torch.nn.Convad							):
                  torch.nn.init.normal(m.weight							)
                  m.bias.data.fill_(1.0							)

        UpperCAmelCase__						:    str							=							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(snake_case__							)

        torch.manual_seed(0							)
        UpperCAmelCase__						:    Union[str, 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(snake_case__							)

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

        UpperCAmelCase__						:    Dict							=							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       __a							(		self      :	Optional[Any]							,			snake_case__      :	Dict							,			snake_case__      :	Union[str, Any]=0							):


        '''simple docstring'''
        if str(snake_case__							).startswith("mps"							):
             UpperCAmelCase__						:    Optional[int]							=							torch.manual_seed(snake_case__							)
        else:
             UpperCAmelCase__						:    Dict							=							torch.Generator(device=snake_case__							).manual_seed(snake_case__							)

        UpperCAmelCase__						:    List[str]							=							2

        UpperCAmelCase__						:    List[str]							=							[
            randn_tensor(
                (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)							,			generator=snake_case__							,			device=torch.device(snake_case__							)							,			),
            randn_tensor(
                (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor)							,			generator=snake_case__							,			device=torch.device(snake_case__							)							,			),
        ]

        UpperCAmelCase__						:    int							=							floats_tensor(control_image[0].shape							,			rng=random.Random(snake_case__							)							).to(snake_case__							)
        UpperCAmelCase__						:    Any							=							image.cpu().permute(0							,			2							,			3							,			1							)[0]
        UpperCAmelCase__						:    Optional[Any]							=							Image.fromarray(np.uinta(snake_case__							)							).convert("RGB"							).resize((6_4, 6_4)							)
        UpperCAmelCase__						:    List[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       __a							(		self      :	Optional[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    int							=							self.get_dummy_components()
        UpperCAmelCase__						:    Optional[int]							=							self.pipeline_class(**snake_case__							)
        pipe.to(snake_case__							)

        UpperCAmelCase__						:    Dict							=							10.0
        UpperCAmelCase__						:    Tuple							=							4

        UpperCAmelCase__						:    Dict							=							self.get_dummy_inputs(snake_case__							)
        UpperCAmelCase__						:    Any							=							steps
        UpperCAmelCase__						:    str							=							scale
        UpperCAmelCase__						:    Optional[Any]							=							pipe(**snake_case__							)[0]

        UpperCAmelCase__						:    Union[str, Any]							=							self.get_dummy_inputs(snake_case__							)
        UpperCAmelCase__						:    Optional[Any]							=							steps
        UpperCAmelCase__						:    Tuple							=							scale
        UpperCAmelCase__						:    List[Any]							=							pipe(**snake_case__							,			control_guidance_start=0.1							,			control_guidance_end=0.2							)[0]

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

        UpperCAmelCase__						:    Union[str, Any]							=							self.get_dummy_inputs(snake_case__							)
        UpperCAmelCase__						:    Any							=							steps
        UpperCAmelCase__						:    int							=							scale
        UpperCAmelCase__						:    Dict							=							pipe(**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       __a							(		self      :	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       __a							(		self      :	List[str]							):


        '''simple docstring'''
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3							)
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        self._test_inference_batch_single_identical(expected_max_diff=2e-3							)




   def       __a							(		self      :	Optional[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[Any]							=							self.get_dummy_components()
        UpperCAmelCase__						:    Any							=							self.pipeline_class(**snake_case__							)
        pipe.to(snake_case__							)
        pipe.set_progress_bar_config(disable=snake_case__							)
        with tempfile.TemporaryDirectory() as tmpdir:
             try:
                  # save_pretrained is not implemented for Multi-ControlNet
                  pipe.save_pretrained(snake_case__							)
             except NotImplementedError:
                  pass




@slow
@require_torch_gpu
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"							)

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

        UpperCAmelCase__						:    str							=							torch.Generator(device="cpu"							).manual_seed(0							)
        UpperCAmelCase__						:    Optional[Any]							=							"evil space-punk bird"
        UpperCAmelCase__						:    int							=							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__						:    Optional[Any]							=							load_image(
            "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"							).resize((5_1_2, 5_1_2)							)

        UpperCAmelCase__						:    List[Any]							=							pipe(
            snake_case__							,			snake_case__							,			control_image=snake_case__							,			generator=snake_case__							,			output_type="np"							,			num_inference_steps=5_0							,			strength=0.6							,			)

        UpperCAmelCase__						:    Tuple							=							output.images[0]

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

        UpperCAmelCase__						:    Optional[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"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase	:	Tuple				=      {
    """configuration_lxmert""": ["""LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LxmertConfig"""],
    """tokenization_lxmert""": ["""LxmertTokenizer"""],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	int				=      ["""LxmertTokenizerFast"""]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	Tuple				=      [
     """LxmertEncoder""",
     """LxmertForPreTraining""",
     """LxmertForQuestionAnswering""",
     """LxmertModel""",
     """LxmertPreTrainedModel""",
     """LxmertVisualFeatureEncoder""",
     """LxmertXLayer""",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	Optional[int]				=      [
     """TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """TFLxmertForPreTraining""",
     """TFLxmertMainLayer""",
     """TFLxmertModel""",
     """TFLxmertPreTrainedModel""",
     """TFLxmertVisualFeatureEncoder""",
 ]


if TYPE_CHECKING:
 from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
 from .tokenization_lxmert import LxmertTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_lxmert_fast import LxmertTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_lxmert import (
      LxmertEncoder,
      LxmertForPreTraining,
      LxmertForQuestionAnswering,
      LxmertModel,
      LxmertPreTrainedModel,
      LxmertVisualFeatureEncoder,
      LxmertXLayer,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_lxmert import (
      TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFLxmertForPreTraining,
      TFLxmertMainLayer,
      TFLxmertModel,
      TFLxmertPreTrainedModel,
      TFLxmertVisualFeatureEncoder,
  )

else:
 import sys

 _lowerCAmelCase	:	Any				=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


