



import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
						import onnxruntime as ort



class   UpperCAmelCase_    (       _lowerCAmelCase   ,					unittest.TestCase):
							lowerCamelCase__   			=		'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'







							def    snake_case__      (   self,			__a=0):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    floats_tensor((1, 3, 128, 128),			rng=random.Random(__a))
														_lowerCAmelCase     :							Optional[Any]    =    np.random.RandomState(__a)
														_lowerCAmelCase     :							Dict    =    {
														    "prompt": "A painting of a squirrel eating a burger",
														    "image": image,
														    "generator": generator,
														    "num_inference_steps": 3,
														    "strength": 0.75,
														    "guidance_scale": 7.5,
														    "output_type": "numpy",
														}
														return inputs







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[str]    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Optional[int]    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							int    =    pipe(**__a).images
														_lowerCAmelCase     :							int    =    image[0, -3:, -3:, -1].flatten()

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							str    =    np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087])
														assert np.abs(image_slice - expected_slice).max() < 1E-1







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Any    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														_lowerCAmelCase     :							str    =    PNDMScheduler.from_config(pipe.scheduler.config,			skip_prk_steps=__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							str    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							Optional[int]    =    pipe(**__a).images
														_lowerCAmelCase     :							Optional[Any]    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							Optional[Any]    =    np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154])

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														_lowerCAmelCase     :							int    =    LMSDiscreteScheduler.from_config(pipe.scheduler.config)
														pipe.set_progress_bar_config(disable=__a)

														# warmup pass to apply optimizations
														_lowerCAmelCase     :							Optional[Any]    =    pipe(**self.get_dummy_inputs())

														_lowerCAmelCase     :							List[str]    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							List[str]    =    pipe(**__a).images
														_lowerCAmelCase     :							Any    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							Dict    =    np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203])

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														_lowerCAmelCase     :							Any    =    EulerDiscreteScheduler.from_config(pipe.scheduler.config)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							str    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							Tuple    =    pipe(**__a).images
														_lowerCAmelCase     :							Optional[int]    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							Union[str, Any]    =    np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304])

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														_lowerCAmelCase     :							Union[str, Any]    =    EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Union[str, Any]    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							int    =    pipe(**__a).images
														_lowerCAmelCase     :							Tuple    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							List[str]    =    np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304])

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Any    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,			provider="CPUExecutionProvider")
														_lowerCAmelCase     :							Optional[int]    =    DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Optional[int]    =    self.get_dummy_inputs()
														_lowerCAmelCase     :							List[str]    =    pipe(**__a).images
														_lowerCAmelCase     :							Tuple    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 128, 128, 3)
														_lowerCAmelCase     :							Union[str, Any]    =    np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552])

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



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







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														return (
														    "CUDAExecutionProvider",
														    {
														        "gpu_mem_limit": "15000000000",  # 15GB
														        "arena_extend_strategy": "kSameAsRequested",
														    },
														)







							@property
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    ort.SessionOptions()
														_lowerCAmelCase     :							Union[str, Any]    =    False
														return options







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg")
														_lowerCAmelCase     :							List[Any]    =    init_image.resize((768, 512))
														# using the PNDM scheduler by default
														_lowerCAmelCase     :							Optional[Any]    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(
														    "CompVis/stable-diffusion-v1-4",			revision="onnx",			safety_checker=__a,			feature_extractor=__a,			provider=self.gpu_provider,			sess_options=self.gpu_options,			)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Optional[int]    =    "A fantasy landscape, trending on artstation"

														_lowerCAmelCase     :							Optional[Any]    =    np.random.RandomState(0)
														_lowerCAmelCase     :							int    =    pipe(
														    prompt=__a,			image=__a,			strength=0.75,			guidance_scale=7.5,			num_inference_steps=10,			generator=__a,			output_type="np",			)
														_lowerCAmelCase     :							int    =    output.images
														_lowerCAmelCase     :							Optional[int]    =    images[0, 255:258, 383:386, -1]

														assert images.shape == (1, 512, 768, 3)
														_lowerCAmelCase     :							Dict    =    np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019])
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    load_image(
														    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
														    "/img2img/sketch-mountains-input.jpg")
														_lowerCAmelCase     :							Any    =    init_image.resize((768, 512))
														_lowerCAmelCase     :							Optional[Any]    =    LMSDiscreteScheduler.from_pretrained(
														    "runwayml/stable-diffusion-v1-5",			subfolder="scheduler",			revision="onnx")
														_lowerCAmelCase     :							Dict    =    OnnxStableDiffusionImgaImgPipeline.from_pretrained(
														    "runwayml/stable-diffusion-v1-5",			revision="onnx",			scheduler=__a,			safety_checker=__a,			feature_extractor=__a,			provider=self.gpu_provider,			sess_options=self.gpu_options,			)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Tuple    =    "A fantasy landscape, trending on artstation"

														_lowerCAmelCase     :							Optional[Any]    =    np.random.RandomState(0)
														_lowerCAmelCase     :							Optional[Any]    =    pipe(
														    prompt=__a,			image=__a,			strength=0.75,			guidance_scale=7.5,			num_inference_steps=20,			generator=__a,			output_type="np",			)
														_lowerCAmelCase     :							Tuple    =    output.images
														_lowerCAmelCase     :							Optional[Any]    =    images[0, 255:258, 383:386, -1]

														assert images.shape == (1, 512, 768, 3)
														_lowerCAmelCase     :							Optional[Any]    =    np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431])
														# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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







'''simple docstring'''


import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
 from transformers.models.bert.modeling_flax_bert import (
     FlaxBertForMaskedLM,
     FlaxBertForMultipleChoice,
     FlaxBertForNextSentencePrediction,
     FlaxBertForPreTraining,
     FlaxBertForQuestionAnswering,
     FlaxBertForSequenceClassification,
     FlaxBertForTokenClassification,
     FlaxBertModel,
 )
class       a_   (				unittest.TestCase							):





     def __init__( self		:					List[str]	,      lowercase		:					str	,      lowercase		:					Union[str, Any]=13	,      lowercase		:					int=7	,      lowercase		:					List[str]=True	,      lowercase		:					int=True	,      lowercase		:					str=True	,      lowercase		:					Any=True	,      lowercase		:					List[str]=99	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Optional[Any]=5	,      lowercase		:					Dict=4	,      lowercase		:					Dict=37	,      lowercase		:					Dict="gelu"	,      lowercase		:					Optional[int]=0.1	,      lowercase		:					str=0.1	,      lowercase		:					List[Any]=512	,      lowercase		:					str=16	,      lowercase		:					Dict=2	,      lowercase		:					Any=0.02	,      lowercase		:					Any=4	,      ):





        """simple docstring"""




        lowercase_			:List[str]    			= parent
        lowercase_			:Any    			= batch_size
        lowercase_			:Dict    			= seq_length
        lowercase_			:Union[str, Any]    			= is_training
        lowercase_			:Optional[int]    			= use_attention_mask
        lowercase_			:Any    			= use_token_type_ids
        lowercase_			:Union[str, Any]    			= use_labels
        lowercase_			:Dict    			= vocab_size
        lowercase_			:Tuple    			= hidden_size
        lowercase_			:Tuple    			= num_hidden_layers
        lowercase_			:Optional[int]    			= num_attention_heads
        lowercase_			:Optional[Any]    			= intermediate_size
        lowercase_			:str    			= hidden_act
        lowercase_			:Tuple    			= hidden_dropout_prob
        lowercase_			:Optional[Any]    			= attention_probs_dropout_prob
        lowercase_			:Tuple    			= max_position_embeddings
        lowercase_			:Any    			= type_vocab_size
        lowercase_			:int    			= type_sequence_label_size
        lowercase_			:Tuple    			= initializer_range
        lowercase_			:Optional[Any]    			= num_choices





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:int    			= ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

        lowercase_			:Union[str, Any]    			= None
        if self.use_attention_mask:
           lowercase_			:Dict    			= random_attention_mask([self.batch_size, self.seq_length]						)

        lowercase_			:List[str]    			= None
        if self.use_token_type_ids:
           lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      self.type_vocab_size						)

        lowercase_			:Optional[Any]    			= BertConfig(
            vocab_size=self.vocab_size	,      hidden_size=self.hidden_size	,      num_hidden_layers=self.num_hidden_layers	,      num_attention_heads=self.num_attention_heads	,      intermediate_size=self.intermediate_size	,      hidden_act=self.hidden_act	,      hidden_dropout_prob=self.hidden_dropout_prob	,      attention_probs_dropout_prob=self.attention_probs_dropout_prob	,      max_position_embeddings=self.max_position_embeddings	,      type_vocab_size=self.type_vocab_size	,      is_decoder=lowercase	,      initializer_range=self.initializer_range	,      )

        return config, input_ids, token_type_ids, attention_mask





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





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Tuple    			= config_and_inputs
        lowercase_			:Tuple    			= {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
        return config, inputs_dict







     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Any    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Union[str, Any]    			= config_and_inputs

        lowercase_			:Dict    			= True
        lowercase_			:Optional[int]    			= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
        lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2						)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )







@require_flax
class       a_   (				_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			True

     __A      =			(
         (
             FlaxBertModel,
             FlaxBertForPreTraining,
             FlaxBertForMaskedLM,
             FlaxBertForMultipleChoice,
             FlaxBertForQuestionAnswering,
             FlaxBertForNextSentencePrediction,
             FlaxBertForSequenceClassification,
             FlaxBertForTokenClassification,
             FlaxBertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= FlaxBertModelTester(self						)







     @slow
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBertModel.from_pretrained("bert-base-cased"						)
        lowercase_			:str    			= model(np.ones((1, 1)						)						)
        self.assertIsNotNone(lowercase						)


