




'''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_ (			lowercase_      ,	lowercase_      ,	lowercase_      ,	unittest.TestCase ):







  '''simple docstring'''







  lowerCAmelCase_					:       List[Any]					  =    StableDiffusionControlNetImgaImgPipeline
  lowerCAmelCase_					:       Union[str, Any]					  =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
  lowerCAmelCase_					:       Optional[int]					  =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowerCAmelCase_					:       str					  =    IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"""control_image"""} )
  lowerCAmelCase_					:       Tuple					  =    IMAGE_TO_IMAGE_IMAGE_PARAMS


  def 				SCREAMING_SNAKE_CASE__      (    self     :						Dict			):






   """simple docstring"""






   torch.manual_seed(0			)
   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     DDIMScheduler(
       beta_start=0.0_0085    ,		beta_end=0.012    ,		beta_schedule="""scaled_linear"""    ,		clip_sample=UpperCamelCase_    ,		set_alpha_to_one=UpperCamelCase_    ,		)
   torch.manual_seed(0			)
   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     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    ,		)
   UpperCAmelCase__      				=     CLIPTextModel(UpperCamelCase_			)
   UpperCAmelCase__      				=     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

   UpperCAmelCase__      				=     {
       """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     :						Tuple    ,		_UpperCAmelCase     :						Tuple    ,		_UpperCAmelCase     :						List[str]=0			):






   """simple docstring"""






   if str(UpperCamelCase_			).startswith("""mps"""			):
    UpperCAmelCase__      				=     torch.manual_seed(UpperCamelCase_			)
   else:
    UpperCAmelCase__      				=     torch.Generator(device=UpperCamelCase_			).manual_seed(UpperCamelCase_			)

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






   """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     :						Optional[Any]			):






   """simple docstring"""






   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3			)



  def 				SCREAMING_SNAKE_CASE__      (    self     :						Any			):






   """simple docstring"""






   self._test_inference_batch_single_identical(expected_max_diff=2E-3			)

class     lowerCAmelCase_ (			lowercase_      ,	lowercase_      ,	unittest.TestCase ):







  '''simple docstring'''







  lowerCAmelCase_					:       Optional[Any]					  =    StableDiffusionControlNetImgaImgPipeline
  lowerCAmelCase_					:       Optional[int]					  =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
  lowerCAmelCase_					:       str					  =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowerCAmelCase_					:       List[str]					  =    frozenset([] )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess


  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[Any]			):






   """simple docstring"""






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

   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     DDIMScheduler(
       beta_start=0.0_0085    ,		beta_end=0.012    ,		beta_schedule="""scaled_linear"""    ,		clip_sample=UpperCamelCase_    ,		set_alpha_to_one=UpperCamelCase_    ,		)
   torch.manual_seed(0			)
   UpperCAmelCase__      				=     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			)
   UpperCAmelCase__      				=     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    ,		)
   UpperCAmelCase__      				=     CLIPTextModel(UpperCamelCase_			)
   UpperCAmelCase__      				=     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

   UpperCAmelCase__      				=     MultiControlNetModel([controlneta, controlneta]			)

   UpperCAmelCase__      				=     {
       """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     :						Optional[int]    ,		_UpperCAmelCase     :						Dict    ,		_UpperCAmelCase     :						Optional[Any]=0			):






   """simple docstring"""






   if str(UpperCamelCase_			).startswith("""mps"""			):
    UpperCAmelCase__      				=     torch.manual_seed(UpperCamelCase_			)
   else:
    UpperCAmelCase__      				=     torch.Generator(device=UpperCamelCase_			).manual_seed(UpperCamelCase_			)

   UpperCAmelCase__      				=     2

   UpperCAmelCase__      				=     [
       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_			)    ,		),
   ]

   UpperCAmelCase__      				=     floats_tensor(control_image[0].shape    ,		rng=random.Random(UpperCamelCase_			)			).to(UpperCamelCase_			)
   UpperCAmelCase__      				=     image.cpu().permute(0    ,		2    ,		3    ,		1			)[0]
   UpperCAmelCase__      				=     Image.fromarray(np.uinta(UpperCamelCase_			)			).convert("""RGB"""			).resize((64, 64)			)
   UpperCAmelCase__      				=     {
       """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			):






   """simple docstring"""






   UpperCAmelCase__      				=     self.get_dummy_components()
   UpperCAmelCase__      				=     self.pipeline_class(**UpperCamelCase_			)
   pipe.to(UpperCamelCase_			)

   UpperCAmelCase__      				=     10.0
   UpperCAmelCase__      				=     4

   UpperCAmelCase__      				=     self.get_dummy_inputs(UpperCamelCase_			)
   UpperCAmelCase__      				=     steps
   UpperCAmelCase__      				=     scale
   UpperCAmelCase__      				=     pipe(**UpperCamelCase_			)[0]

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

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

   UpperCAmelCase__      				=     self.get_dummy_inputs(UpperCamelCase_			)
   UpperCAmelCase__      				=     steps
   UpperCAmelCase__      				=     scale
   UpperCAmelCase__      				=     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 				SCREAMING_SNAKE_CASE__      (    self     :						Optional[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]			):






   """simple docstring"""






   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3			)


  def 				SCREAMING_SNAKE_CASE__      (    self     :						Optional[int]			):






   """simple docstring"""






   self._test_inference_batch_single_identical(expected_max_diff=2E-3			)



  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[Any]			):






   """simple docstring"""






   UpperCAmelCase__      				=     self.get_dummy_components()
   UpperCAmelCase__      				=     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     lowerCAmelCase_ (			unittest.TestCase ):







  '''simple docstring'''







  def 				SCREAMING_SNAKE_CASE__      (    self     :						Optional[int]			):






   """simple docstring"""






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



  def 				SCREAMING_SNAKE_CASE__      (    self     :						str			):






   """simple docstring"""






   UpperCAmelCase__      				=     ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""			)

   UpperCAmelCase__      				=     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_			)

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

   UpperCAmelCase__      				=     pipe(
       UpperCamelCase_    ,		UpperCamelCase_    ,		control_image=UpperCamelCase_    ,		generator=UpperCamelCase_    ,		output_type="""np"""    ,		num_inference_steps=50    ,		strength=0.6    ,		)

   UpperCAmelCase__      				=     output.images[0]

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

   UpperCAmelCase__      				=     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 ...configuration_utils import PretrainedConfig
from ...utils import logging


__a   :							Optional[Any]						= logging.get_logger(__name__)

__a   :							Dict						= {
    "google/realm-cc-news-pretrained-embedder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-encoder": (
        "https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-scorer": (
        "https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json"
    ),
    "google/realm-cc-news-pretrained-openqa": (
        "https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json"
    ),
    "google/realm-orqa-nq-openqa": "https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json",
    "google/realm-orqa-nq-reader": "https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json",
    "google/realm-orqa-wq-openqa": "https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json",
    "google/realm-orqa-wq-reader": "https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json",
    # See all REALM models at https://huggingface.co/models?filter=realm
}


class    __lowercase   (						lowercase_	):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE     =					"realm"




     def __init__(					self							:      str			,						UpperCamelCase_							:      List[Any]=30_522			,						UpperCamelCase_							:      Dict=768			,						UpperCamelCase_							:      Union[str, Any]=128			,						UpperCamelCase_							:      Tuple=12			,						UpperCamelCase_							:      Optional[Any]=12			,						UpperCamelCase_							:      Optional[int]=8			,						UpperCamelCase_							:      str=3_072			,						UpperCamelCase_							:      List[str]="gelu_new"			,						UpperCamelCase_							:      Optional[int]=0.1			,						UpperCamelCase_							:      Optional[int]=0.1			,						UpperCamelCase_							:      str=512			,						UpperCamelCase_							:      int=2			,						UpperCamelCase_							:      str=0.02			,						UpperCamelCase_							:      List[Any]=1e-12			,						UpperCamelCase_							:      str=256			,						UpperCamelCase_							:      List[str]=10			,						UpperCamelCase_							:      int=1e-3			,						UpperCamelCase_							:      List[str]=5			,						UpperCamelCase_							:      Tuple=320			,						UpperCamelCase_							:      List[str]=13_353_718			,						UpperCamelCase_							:      Tuple=5_000			,						UpperCamelCase_							:      Optional[int]=1			,						UpperCamelCase_							:      List[str]=0			,						UpperCamelCase_							:      Union[str, Any]=2			,						**UpperCamelCase_							:      List[Any]			,						):


         """simple docstring"""



         super().__init__(pad_token_id=UpperCamelCase_			,						bos_token_id=UpperCamelCase_			,						eos_token_id=UpperCamelCase_			,						**UpperCamelCase_					)

         # Common config
         __A       =						vocab_size
         __A       =						max_position_embeddings
         __A       =						hidden_size
         __A       =						retriever_proj_size
         __A       =						num_hidden_layers
         __A       =						num_attention_heads
         __A       =						num_candidates
         __A       =						intermediate_size
         __A       =						hidden_act
         __A       =						hidden_dropout_prob
         __A       =						attention_probs_dropout_prob
         __A       =						initializer_range
         __A       =						type_vocab_size
         __A       =						layer_norm_eps

         # Reader config
         __A       =						span_hidden_size
         __A       =						max_span_width
         __A       =						reader_layer_norm_eps
         __A       =						reader_beam_size
         __A       =						reader_seq_len

         # Retrieval config
         __A       =						num_block_records
         __A       =						searcher_beam_size




