



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      __lowerCAmelCase     (		UpperCAmelCase_							,     unittest.TestCase):
					"""simple docstring"""






					A__				:		Optional[Any]			       =    "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"



					def     _a							(     self					:       Optional[Any]				,					_snake_case					:       str=0					):






								"""simple docstring"""


								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     _a							(     self					:       Any					):






								"""simple docstring"""


								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     _a							(     self					:       Union[str, Any]					):






								"""simple docstring"""


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






								"""simple docstring"""


								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     _a							(     self					:       Optional[int]					):






								"""simple docstring"""


								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     _a							(     self					:       Dict					):






								"""simple docstring"""


								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     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								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     _a							(     self					:       Any					):






								"""simple docstring"""


								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     _a							(     self					:       List[Any]					):






								"""simple docstring"""


								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      __lowerCAmelCase     (		unittest.TestCase):
					"""simple docstring"""






					@property
					def     _a							(     self					:       Optional[int]					):






								"""simple docstring"""


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



					@property
					def     _a							(     self					:       Union[str, Any]					):






								"""simple docstring"""


								A__       =    ort.SessionOptions()
								A__       =    False
								return options



					def     _a							(     self					:       List[Any]					):






								"""simple docstring"""


								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     _a							(     self					:       Any					):






								"""simple docstring"""


								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     _a							(     self					:       Any					):






								"""simple docstring"""


								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     _a							(     self					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    0

								def test_callback_fn(_snake_case					:       Dict				,					_snake_case					:       List[Any]				,					_snake_case					:       Union[str, Any]					) -> 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     _a							(     self					:       Tuple					):






								"""simple docstring"""


								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 unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
							import optax
							from flax.training.common_utils import onehot

							from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
							from transformers.models.ta.modeling_flax_ta import shift_tokens_right


@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class      __lowerCAmelCase     (		unittest.TestCase      ):
					"""simple docstring"""






					@slow
					def     _a							(     self					:       int					):






								"""simple docstring"""


								A__       =    FlaxMTaForConditionalGeneration.from_pretrained('google/mt5-small'					)
								A__       =    AutoTokenizer.from_pretrained('google/mt5-small'					)

								A__       =    tokenizer('Hello there'				,					return_tensors='np'					).input_ids
								A__       =    tokenizer('Hi I am'				,					return_tensors='np'					).input_ids

								A__       =    shift_tokens_right(_snake_case				,					model.config.pad_token_id				,					model.config.decoder_start_token_id					)

								A__       =    model(_snake_case				,					decoder_input_ids=_snake_case					).logits
								A__       =    optax.softmax_cross_entropy(_snake_case				,					onehot(_snake_case				,					logits.shape[-1]					)					).mean()

								A__       =    -(labels.shape[-1] * loss.item())

								A__       =    -84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE					) < 1E-4					)




