import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
    import onnxruntime as ort



class 				a							(			__lowerCAmelCase       ,			unittest.TestCase       ):







 """simple docstring"""






 lowerCamelCase	:Tuple    			=		'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''



 def      UpperCAmelCase							(					self		,				lowerCAmelCase_=0							)  ->					Any:
        _A        =  np.random.RandomState(lowerCAmelCase_							)
        _A        =  {
            """prompt""": """A painting of a squirrel eating a burger""",
            """generator""": generator,
            """num_inference_steps""": 2,
            """guidance_scale""": 7.5,
            """output_type""": """numpy""",
        }
        return inputs



 def      UpperCAmelCase							(					self							)  ->					List[str]:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					Any:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        _A        =  PNDMScheduler.from_config(pipe.scheduler.config		,				skip_prk_steps=lowerCAmelCase_							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					List[Any]:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        _A        =  LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					int:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        _A        =  EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					List[str]:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        _A        =  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					Optional[Any]:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        _A        =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  pipe(**lowerCAmelCase_							).images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 1_28, 1_28, 3)
        _A        =  np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2



 def      UpperCAmelCase							(					self							)  ->					str:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  3 * [inputs["""prompt"""]]

        # forward
        _A        =  pipe(**lowerCAmelCase_							)
        _A        =  output.images[0, -3:, -3:, -1]

        _A        =  self.get_dummy_inputs()
        _A        =  3 * [inputs.pop("""prompt"""							)]

        _A        =  pipe.tokenizer(
            lowerCAmelCase_		,				padding="""max_length"""		,				max_length=pipe.tokenizer.model_max_length		,				truncation=lowerCAmelCase_		,				return_tensors="""np"""		,				)
        _A        =  text_inputs["""input_ids"""]

        _A        =  pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]

        _A        =  prompt_embeds

        # forward
        _A        =  pipe(**lowerCAmelCase_							)
        _A        =  output.images[0, -3:, -3:, -1]

        assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()							).max() < 1E-4



 def      UpperCAmelCase							(					self							)  ->					Optional[int]:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint		,				provider="""CPUExecutionProvider"""							)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  self.get_dummy_inputs()
        _A        =  3 * ["""this is a negative prompt"""]
        _A        =  negative_prompt
        _A        =  3 * [inputs["""prompt"""]]

        # forward
        _A        =  pipe(**lowerCAmelCase_							)
        _A        =  output.images[0, -3:, -3:, -1]

        _A        =  self.get_dummy_inputs()
        _A        =  3 * [inputs.pop("""prompt"""							)]

        _A        =  []
        for p in [prompt, negative_prompt]:
               _A        =  pipe.tokenizer(
                   lowerCAmelCase_		,				padding="""max_length"""		,				max_length=pipe.tokenizer.model_max_length		,				truncation=lowerCAmelCase_		,				return_tensors="""np"""		,				)
               _A        =  text_inputs["""input_ids"""]

               embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]							)

        _A   ,  _A        =  embeds

        # forward
        _A        =  pipe(**lowerCAmelCase_							)
        _A        =  output.images[0, -3:, -3:, -1]

        assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()							).max() < 1E-4



@nightly
@require_onnxruntime
@require_torch_gpu
class 				a							(			unittest.TestCase       ):







 """simple docstring"""






 @property
 def      UpperCAmelCase							(					self							)  ->					List[str]:
        return (
            "CUDAExecutionProvider",
            {
                "gpu_mem_limit": "15000000000",  # 15GB
                "arena_extend_strategy": "kSameAsRequested",
            },
        )



 @property
 def      UpperCAmelCase							(					self							)  ->					List[Any]:
        _A        =  ort.SessionOptions()
        _A        =  False
        return options



 def      UpperCAmelCase							(					self							)  ->					Union[str, Any]:
        # using the PNDM scheduler by default
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(
            """CompVis/stable-diffusion-v1-4"""		,				revision="""onnx"""		,				safety_checker=lowerCAmelCase_		,				feature_extractor=lowerCAmelCase_		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
        sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  """A painting of a squirrel eating a burger"""
        np.random.seed(0							)
        _A        =  sd_pipe([prompt]		,				guidance_scale=6.0		,				num_inference_steps=10		,				output_type="""np"""							)
        _A        =  output.images

        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        _A        =  np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3



 def      UpperCAmelCase							(					self							)  ->					int:
        _A        =  DDIMScheduler.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				subfolder="""scheduler"""		,				revision="""onnx"""							)
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				revision="""onnx"""		,				scheduler=lowerCAmelCase_		,				safety_checker=lowerCAmelCase_		,				feature_extractor=lowerCAmelCase_		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
        sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  """open neural network exchange"""
        _A        =  np.random.RandomState(0							)
        _A        =  sd_pipe([prompt]		,				guidance_scale=7.5		,				num_inference_steps=10		,				generator=lowerCAmelCase_		,				output_type="""np"""							)
        _A        =  output.images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        _A        =  np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3



 def      UpperCAmelCase							(					self							)  ->					Optional[Any]:
        _A        =  LMSDiscreteScheduler.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				subfolder="""scheduler"""		,				revision="""onnx"""							)
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				revision="""onnx"""		,				scheduler=lowerCAmelCase_		,				safety_checker=lowerCAmelCase_		,				feature_extractor=lowerCAmelCase_		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
        sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  """open neural network exchange"""
        _A        =  np.random.RandomState(0							)
        _A        =  sd_pipe([prompt]		,				guidance_scale=7.5		,				num_inference_steps=10		,				generator=lowerCAmelCase_		,				output_type="""np"""							)
        _A        =  output.images
        _A        =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        _A        =  np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161]							)

        assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3



 def      UpperCAmelCase							(					self							)  ->					Union[str, Any]:
        _A        =  0

        def test_callback_fn(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							) -> None:
               _A        =  True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 0:
                      assert latents.shape == (1, 4, 64, 64)
                      _A        =  latents[0, -3:, -3:, -1]
                      _A        =  np.array(
                          [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167]							)

                      assert np.abs(latents_slice.flatten() - expected_slice							).max() < 1E-3
               elif step == 5:
                      assert latents.shape == (1, 4, 64, 64)
                      _A        =  latents[0, -3:, -3:, -1]
                      _A        =  np.array(
                          [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875]							)

                      assert np.abs(latents_slice.flatten() - expected_slice							).max() < 1E-3

        _A        =  False

        _A        =  OnnxStableDiffusionPipeline.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				revision="""onnx"""		,				safety_checker=lowerCAmelCase_		,				feature_extractor=lowerCAmelCase_		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
        pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

        _A        =  """Andromeda galaxy in a bottle"""

        _A        =  np.random.RandomState(0							)
        pipe(
            prompt=lowerCAmelCase_		,				num_inference_steps=5		,				guidance_scale=7.5		,				generator=lowerCAmelCase_		,				callback=lowerCAmelCase_		,				callback_steps=1		,				)
        assert test_callback_fn.has_been_called
        assert number_of_steps == 6



 def      UpperCAmelCase							(					self							)  ->					Any:
        _A        =  OnnxStableDiffusionPipeline.from_pretrained(
            """runwayml/stable-diffusion-v1-5"""		,				revision="""onnx"""		,				safety_checker=lowerCAmelCase_		,				feature_extractor=lowerCAmelCase_		,				provider=self.gpu_provider		,				sess_options=self.gpu_options		,				)
        assert isinstance(lowerCAmelCase_		,				lowerCAmelCase_							)
        assert pipe.safety_checker is None

        _A        =  pipe("""example prompt"""		,				num_inference_steps=2							).images[0]
        assert image is not None

        # check that there's no error when saving a pipeline with one of the models being None
        with tempfile.TemporaryDirectory() as tmpdirname:
               pipe.save_pretrained(lowerCAmelCase_							)
               _A        =  OnnxStableDiffusionPipeline.from_pretrained(lowerCAmelCase_							)

        # sanity check that the pipeline still works
        assert pipe.safety_checker is None
        _A        =  pipe("""example prompt"""		,				num_inference_steps=2							).images[0]
        assert image is not None





import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE										=						logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE										=						{
    'BridgeTower/bridgetower-base': 'https://huggingface.co/BridgeTower/bridgetower-base/blob/main/config.json',
    'BridgeTower/bridgetower-base-itm-mlm': (
        'https://huggingface.co/BridgeTower/bridgetower-base-itm-mlm/blob/main/config.json'
    ),
}



class 				a							(			__lowerCAmelCase       ):







 """simple docstring"""






 lowerCamelCase	:Optional[Any]    			=		'''bridgetower_vision_model'''



 def __init__(					self		,				lowerCAmelCase_=7_68		,				lowerCAmelCase_=12		,				lowerCAmelCase_=3		,				lowerCAmelCase_=16		,				lowerCAmelCase_=2_88		,				lowerCAmelCase_=1		,				lowerCAmelCase_=1E-05		,				lowerCAmelCase_=False		,				lowerCAmelCase_=True		,				lowerCAmelCase_=False		,				**lowerCAmelCase_		,				)  ->					Any:
        super().__init__(**lowerCAmelCase_							)
        _A        =  hidden_size
        _A        =  num_hidden_layers
        _A        =  num_channels
        _A        =  patch_size
        _A        =  image_size
        _A        =  initializer_factor
        _A        =  layer_norm_eps
        _A        =  stop_gradient
        _A        =  share_layernorm
        _A        =  remove_last_layer



 @classmethod
 def      UpperCAmelCase							(					cls		,				lowerCAmelCase_		,				**lowerCAmelCase_							)  ->					"PretrainedConfig":
        _A   ,  _A        =  cls.get_config_dict(lowerCAmelCase_		,				**lowerCAmelCase_							)

        if config_dict.get("""model_type"""							) == "bridgetower":
               _A        =  config_dict["""text_config"""]

        if "model_type" in config_dict and hasattr(cls		,				"""model_type"""							) and config_dict["model_type"] != cls.model_type:
               logger.warning(
                   F'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
                   F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

        return cls.from_dict(lowerCAmelCase_		,				**lowerCAmelCase_							)



class 				a							(			__lowerCAmelCase       ):







 """simple docstring"""






 lowerCamelCase	:Union[str, Any]    			=		'''bridgetower_text_model'''



 def __init__(					self		,				lowerCAmelCase_=5_02_65		,				lowerCAmelCase_=7_68		,				lowerCAmelCase_=12		,				lowerCAmelCase_=12		,				lowerCAmelCase_=1		,				lowerCAmelCase_=30_72		,				lowerCAmelCase_="gelu"		,				lowerCAmelCase_=0.1		,				lowerCAmelCase_=0.1		,				lowerCAmelCase_=5_14		,				lowerCAmelCase_=1		,				lowerCAmelCase_=1E-05		,				lowerCAmelCase_=1		,				lowerCAmelCase_=0		,				lowerCAmelCase_=2		,				lowerCAmelCase_="absolute"		,				lowerCAmelCase_=True		,				**lowerCAmelCase_		,				)  ->					Optional[Any]:
        super().__init__(**lowerCAmelCase_							)

        _A        =  vocab_size
        _A        =  hidden_size
        _A        =  num_hidden_layers
        _A        =  num_attention_heads
        _A        =  hidden_act
        _A        =  initializer_factor
        _A        =  intermediate_size
        _A        =  hidden_dropout_prob
        _A        =  attention_probs_dropout_prob
        _A        =  max_position_embeddings
        _A        =  type_vocab_size
        _A        =  layer_norm_eps
        _A        =  position_embedding_type
        _A        =  use_cache
        _A        =  pad_token_id
        _A        =  bos_token_id
        _A        =  eos_token_id



 @classmethod
 def      UpperCAmelCase							(					cls		,				lowerCAmelCase_		,				**lowerCAmelCase_							)  ->					"PretrainedConfig":
        _A   ,  _A        =  cls.get_config_dict(lowerCAmelCase_		,				**lowerCAmelCase_							)

        if config_dict.get("""model_type"""							) == "bridgetower":
               _A        =  config_dict["""text_config"""]

        if "model_type" in config_dict and hasattr(cls		,				"""model_type"""							) and config_dict["model_type"] != cls.model_type:
               logger.warning(
                   F'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
                   F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

        return cls.from_dict(lowerCAmelCase_		,				**lowerCAmelCase_							)



class 				a							(			__lowerCAmelCase       ):







 """simple docstring"""






 lowerCamelCase	:Optional[int]    			=		'''bridgetower'''



 def __init__(					self		,				lowerCAmelCase_=True		,				lowerCAmelCase_="gelu"		,				lowerCAmelCase_=7_68		,				lowerCAmelCase_=1		,				lowerCAmelCase_=1E-05		,				lowerCAmelCase_=False		,				lowerCAmelCase_="add"		,				lowerCAmelCase_=12		,				lowerCAmelCase_=6		,				lowerCAmelCase_=False		,				lowerCAmelCase_=False		,				lowerCAmelCase_=None		,				lowerCAmelCase_=None		,				**lowerCAmelCase_		,				)  ->					int:
        # TODO: remove this once the Hub files are updated.
        _A        =  kwargs.pop("""text_config_dict"""		,				lowerCAmelCase_							)
        _A        =  kwargs.pop("""vision_config_dict"""		,				lowerCAmelCase_							)

        super().__init__(**lowerCAmelCase_							)
        _A        =  share_cross_modal_transformer_layers
        _A        =  hidden_act
        _A        =  hidden_size
        _A        =  initializer_factor
        _A        =  layer_norm_eps
        _A        =  share_link_tower_layers
        _A        =  link_tower_type
        _A        =  num_attention_heads
        _A        =  num_hidden_layers
        _A        =  tie_word_embeddings
        _A        =  init_layernorm_from_vision_encoder

        if text_config is None:
               _A        =  {}
               logger.info("""`text_config` is `None`. Initializing the `BridgeTowerTextConfig` with default values."""							)

        if vision_config is None:
               _A        =  {}
               logger.info("""`vision_config` is `None`. Initializing the `BridgeTowerVisionConfig` with default values."""							)

        _A        =  BridgeTowerTextConfig(**lowerCAmelCase_							)
        _A        =  BridgeTowerVisionConfig(**lowerCAmelCase_							)



 @classmethod
 def      UpperCAmelCase							(					cls		,				lowerCAmelCase_		,				lowerCAmelCase_		,				**lowerCAmelCase_							)  ->					Any:
        return cls(text_config=text_config.to_dict()		,				vision_config=vision_config.to_dict()		,				**lowerCAmelCase_							)



 def      UpperCAmelCase							(					self							)  ->					Union[str, Any]:
        _A        =  copy.deepcopy(self.__dict__							)
        _A        =  self.text_config.to_dict()
        _A        =  self.vision_config.to_dict()
        _A        =  self.__class__.model_type
        return output




