
"""simple docstring"""





import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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 __lowerCAmelCase					(       _lowercase  ,					unittest.TestCase   ):




			"""simple docstring"""







			snake_case						   =					"ssube/stable-diffusion-x4-upscaler-onnx"




			def   lowerCamelCase__  (					self    :       str							, _snake_case    :       Any=0						)   ->			Tuple:







								"""simple docstring"""
								A_								=    floats_tensor((1, 3, 128, 128)							, rng=random.Random(_snake_case						)						)
								A_								=    torch.manual_seed(_snake_case						)
								A_								=    {
								    "prompt": "A painting of a squirrel eating a burger",
								    "image": image,
								    "generator": generator,
								    "num_inference_steps": 3,
								    "guidance_scale": 7.5,
								    "output_type": "numpy",
								}
								return inputs




			def   lowerCamelCase__  (					self    :       int						)   ->			str:







								"""simple docstring"""
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider="CPUExecutionProvider"						)
								pipe.set_progress_bar_config(disable=_snake_case						)

								A_								=    self.get_dummy_inputs()
								A_								=    pipe(**_snake_case						).images
								A_								=    image[0, -3:, -3:, -1].flatten()

								# started as 128, should now be 512
								assert image.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]						)
								assert np.abs(image_slice - expected_slice						).max() < 1e-1




			def   lowerCamelCase__  (					self    :       Optional[int]						)   ->			Tuple:







								"""simple docstring"""
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider="CPUExecutionProvider"						)
								A_								=    PNDMScheduler.from_config(pipe.scheduler.config							, skip_prk_steps=_snake_case						)
								pipe.set_progress_bar_config(disable=_snake_case						)

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

								assert image.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.6_8_9_8_8_9_2, 0.5_9_2_4_0_5_5_6, 0.5_2_4_9_9_5_2_7, 0.5_8_8_6_6_2_1_5, 0.5_2_2_5_8_2_3_5, 0.5_2_5_7_2_7_1_5, 0.6_2_4_1_4_4_7_3, 0.6_1_7_4_3_8_7, 0.6_2_1_4_9_6_4]						)
								assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1




			def   lowerCamelCase__  (					self    :       int						)   ->			str:







								"""simple docstring"""
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider="CPUExecutionProvider"						)
								A_								=    DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
								pipe.set_progress_bar_config(disable=_snake_case						)

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

								assert image.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.7_6_5_9_2_7_8, 0.7_6_4_3_7_6_6_4, 0.7_5_5_7_9_1_0_7, 0.7_6_9_1_1_1_6, 0.7_7_6_6_6_9_8_6, 0.7_7_2_7_6_7_2, 0.7_7_5_8_6_6_4, 0.7_8_1_2_2_2_6, 0.7_6_9_4_2_5_1_5]						)

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




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Tuple:







								"""simple docstring"""
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider="CPUExecutionProvider"						)
								A_								=    EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
								pipe.set_progress_bar_config(disable=_snake_case						)

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

								assert image.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]						)
								assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1






			def   lowerCamelCase__  (					self    :       Tuple						)   ->			int:







								"""simple docstring"""
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							, provider="CPUExecutionProvider"						)
								A_								=    EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
								pipe.set_progress_bar_config(disable=_snake_case						)

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

								assert image.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.7_7_4_2_4_4_9_6, 0.7_7_3_6_0_1, 0.7_6_4_5_2_8_8, 0.7_7_6_9_5_9_8, 0.7_7_7_2_7_3_9, 0.7_7_3_8_6_8_8, 0.7_8_1_8_7_2_3_3, 0.7_7_8_7_9_5_8_4, 0.7_6_7_0_4_3]						)

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





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




			"""simple docstring"""







			@property
			def   lowerCamelCase__  (					self    :       str						)   ->			str:







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




			@property
			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Any:







								"""simple docstring"""
								A_								=    ort.SessionOptions()
								A_								=    False
								return options




			def   lowerCamelCase__  (					self    :       int						)   ->			Dict:







								"""simple docstring"""
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/img2img/sketch-mountains-input.jpg"						)
								A_								=    init_image.resize((128, 128)						)
								# using the PNDM scheduler by default
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(
								    "ssube/stable-diffusion-x4-upscaler-onnx"							, provider=self.gpu_provider							, sess_options=self.gpu_options							, )
								pipe.set_progress_bar_config(disable=_snake_case						)

								A_								=    "A fantasy landscape, trending on artstation"

								A_								=    torch.manual_seed(0						)
								A_								=    pipe(
								    prompt=_snake_case							, image=_snake_case							, guidance_scale=7.5							, num_inference_steps=10							, generator=_snake_case							, output_type="np"							, )
								A_								=    output.images
								A_								=    images[0, 255:258, 383:386, -1]

								assert images.shape == (1, 512, 512, 3)
								A_								=    np.array([0.4_8_8_3, 0.4_9_4_7, 0.4_9_8_0, 0.4_9_7_5, 0.4_9_8_2, 0.4_9_8_0, 0.5_0_0_0, 0.5_0_0_6, 0.4_9_7_2]						)
								# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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






			def   lowerCamelCase__  (					self    :       Optional[int]						)   ->			Dict:







								"""simple docstring"""
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/img2img/sketch-mountains-input.jpg"						)
								A_								=    init_image.resize((128, 128)						)
								A_								=    LMSDiscreteScheduler.from_pretrained(
								    "ssube/stable-diffusion-x4-upscaler-onnx"							, subfolder="scheduler"						)
								A_								=    OnnxStableDiffusionUpscalePipeline.from_pretrained(
								    "ssube/stable-diffusion-x4-upscaler-onnx"							, scheduler=_snake_case							, provider=self.gpu_provider							, sess_options=self.gpu_options							, )
								pipe.set_progress_bar_config(disable=_snake_case						)

								A_								=    "A fantasy landscape, trending on artstation"

								A_								=    torch.manual_seed(0						)
								A_								=    pipe(
								    prompt=_snake_case							, image=_snake_case							, guidance_scale=7.5							, num_inference_steps=20							, generator=_snake_case							, output_type="np"							, )
								A_								=    output.images
								A_								=    images[0, 255:258, 383:386, -1]

								assert images.shape == (1, 512, 512, 3)
								A_								=    np.array(
								    [0.5_0_1_7_3_7_5_3, 0.5_0_2_2_3_3_5_6, 0.5_0_2_0_3_9, 0.5_0_2_3_3_0_3_6, 0.5_0_2_3_7_2_5, 0.5_0_2_2_6_0_1, 0.5_0_1_8_7_5_8, 0.5_0_2_3_4_0_8_5, 0.5_0_2_4_1_5_6_6]						)
								# 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 random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
		import os

		import jax.numpy as jnp
		from jax import jit

		from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
		from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

		UpperCamelCase_       :    List[Any]						      =			'''0.12'''  # assumed parallelism: 8


if is_torch_available():
		import torch
def 				A_  (__a					, __a					, __a=None					):

					'''simple docstring'''





					if rng is None:
										A_								=    random.Random()

					A_								=    1
					for dim in shape:
										total_dims *= dim

					A_								=    []
					for _ in range(__a					):
										values.append(rng.randint(0					, vocab_size - 1					)					)

					A_								=    np.array(__a					, dtype=jnp.intaa					).reshape(__a					)

					return output
def 				A_  (__a					, __a=None					):

					'''simple docstring'''





					A_								=    ids_tensor(__a					, vocab_size=2					, rng=__a					)
					# make sure that at least one token is attended to for each batch
					A_								=    1
					return attn_mask



@require_flax
class __lowerCAmelCase					:




			"""simple docstring"""







			snake_case						   =					None
			snake_case						   =					()




			def   lowerCamelCase__  (					self    :       Any						)   ->			Tuple:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								# cut to half length & take max batch_size 3
								A_								=    2
								A_								=    inputs["input_ids"].shape[-1] // 2
								A_								=    inputs["input_ids"][:max_batch_size, :sequence_length]

								A_								=    jnp.ones_like(_snake_case						)
								A_								=    attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								A_								=    input_ids.shape[-1] + 5
								if config.eos_token_id is not None and config.pad_token_id is None:
													# hack to allow generate for models such as GPT2 as is done in `generate()`
													A_								=    config.eos_token_id
								return config, input_ids, attention_mask, max_length




			@is_pt_flax_cross_test
			def   lowerCamelCase__  (					self    :       List[Any]						)   ->			List[Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    0

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
													A_								=    getattr(_snake_case							, _snake_case						)
													A_								=    pt_model_class(_snake_case						).eval()
													A_								=    load_flax_weights_in_pytorch_model(_snake_case							, flax_model.params						)

													A_								=    flax_model.generate(_snake_case						).sequences
													A_								=    pt_model.generate(torch.tensor(_snake_case							, dtype=torch.long						)						)

													if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																		A_								=    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

													self.assertListEqual(pt_generation_outputs.numpy().tolist()							, flax_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Any						)   ->			Optional[Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			List[str]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       str						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[0]							, input_ids.shape[0] * config.num_return_sequences						)




			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length
								A_								=    0.8
								A_								=    10
								A_								=    0.3
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    2
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)






			def   lowerCamelCase__  (					self    :       Dict						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    2
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)





@require_flax
class __lowerCAmelCase					(       unittest.TestCase   ):




			"""simple docstring"""







			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			Any:







								"""simple docstring"""
								A_								=    AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"						)
								A_								=    FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"						)

								A_								=    "Hello world"
								A_								=    tokenizer(_snake_case							, return_tensors="np"						).input_ids

								# typos are quickly detected (the correct argument is `do_sample`)
								with self.assertRaisesRegex(_snake_case							, "do_samples"						):
													model.generate(_snake_case							, do_samples=_snake_case						)

								# arbitrary arguments that will not be used anywhere are also not accepted
								with self.assertRaisesRegex(_snake_case							, "foo"						):
													A_								=    {"foo": "bar"}
													model.generate(_snake_case							, **_snake_case						)




