

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 SCREAMING_SNAKE_CASE	(		__snake_case				,    unittest.TestCase				):






			"""simple docstring"""




			__A    	=	"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase=0				):

				"""simple docstring"""


				snake_case_   		=			np.random.RandomState(__UpperCamelCase				)
				snake_case_   		=			{
				    '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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				snake_case_   		=			PNDMScheduler.from_config(pipe.scheduler.config			,					skip_prk_steps=__UpperCamelCase				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				snake_case_   		=			LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				snake_case_   		=			EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				snake_case_   		=			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				snake_case_   		=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 1_28, 1_28, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_dummy_inputs()
				snake_case_   		=			3 * [inputs['prompt']]

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

				snake_case_   		=			self.get_dummy_inputs()
				snake_case_   		=			3 * [inputs.pop('prompt'				)]

				snake_case_   		=			pipe.tokenizer(
				    __UpperCamelCase			,					padding='max_length'			,					max_length=pipe.tokenizer.model_max_length			,					truncation=__UpperCamelCase			,					return_tensors='np'			,					)
				snake_case_   		=			text_inputs['input_ids']

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

				snake_case_   		=			prompt_embeds

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

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







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint			,					provider='CPUExecutionProvider'				)
				pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

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

				snake_case_   		=			self.get_dummy_inputs()
				snake_case_   		=			3 * [inputs.pop('prompt'				)]

				snake_case_   		=			[]
				for p in [prompt, negative_prompt]:
					snake_case_   		=			pipe.tokenizer(
					    __UpperCamelCase			,					padding='max_length'			,					max_length=pipe.tokenizer.model_max_length			,					truncation=__UpperCamelCase			,					return_tensors='np'			,					)
					snake_case_   		=			text_inputs['input_ids']

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

				snake_case_		,							snake_case_   		=			embeds

				# forward
				snake_case_   		=			pipe(**__UpperCamelCase				)
				snake_case_   		=			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 SCREAMING_SNAKE_CASE	(		unittest.TestCase				):






			"""simple docstring"""




			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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


			@property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(
				    'CompVis/stable-diffusion-v1-4'			,					revision='onnx'			,					safety_checker=__UpperCamelCase			,					feature_extractor=__UpperCamelCase			,					provider=self.gpu_provider			,					sess_options=self.gpu_options			,					)
				sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

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

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			DDIMScheduler.from_pretrained(
				    'runwayml/stable-diffusion-v1-5'			,					subfolder='scheduler'			,					revision='onnx'				)
				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(
				    'runwayml/stable-diffusion-v1-5'			,					revision='onnx'			,					scheduler=__UpperCamelCase			,					safety_checker=__UpperCamelCase			,					feature_extractor=__UpperCamelCase			,					provider=self.gpu_provider			,					sess_options=self.gpu_options			,					)
				sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			LMSDiscreteScheduler.from_pretrained(
				    'runwayml/stable-diffusion-v1-5'			,					subfolder='scheduler'			,					revision='onnx'				)
				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(
				    'runwayml/stable-diffusion-v1-5'			,					revision='onnx'			,					scheduler=__UpperCamelCase			,					safety_checker=__UpperCamelCase			,					feature_extractor=__UpperCamelCase			,					provider=self.gpu_provider			,					sess_options=self.gpu_options			,					)
				sd_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

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

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_   		=			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       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			0

				def test_callback_fn(__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase				) -> None:
					snake_case_   		=			True
					nonlocal number_of_steps
					number_of_steps += 1
					if step == 0:
						assert latents.shape == (1, 4, 64, 64)
						snake_case_   		=			latents[0, -3:, -3:, -1]
						snake_case_   		=			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)
						snake_case_   		=			latents[0, -3:, -3:, -1]
						snake_case_   		=			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

				snake_case_   		=			False

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

				snake_case_   		=			'Andromeda galaxy in a bottle'

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







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(
				    'runwayml/stable-diffusion-v1-5'			,					revision='onnx'			,					safety_checker=__UpperCamelCase			,					feature_extractor=__UpperCamelCase			,					provider=self.gpu_provider			,					sess_options=self.gpu_options			,					)
				assert isinstance(__UpperCamelCase			,					__UpperCamelCase				)
				assert pipe.safety_checker is None

				snake_case_   		=			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(__UpperCamelCase				)
					snake_case_   		=			OnnxStableDiffusionPipeline.from_pretrained(__UpperCamelCase				)

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



import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
A									=  logging.getLogger()

A									=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)


class SCREAMING_SNAKE_CASE	(		__snake_case				):






			"""simple docstring"""




			def       __lowerCAmelCase			(     self			,					__UpperCamelCase				):

				"""simple docstring"""


				os.makedirs(__UpperCamelCase			,					exist_ok=__UpperCamelCase				)
				snake_case_   		=			{'source': 'What is love ?', 'target': 'life'}
				snake_case_   		=			{'train': 12, 'val': 2, 'test': 2}
				for split in ["train", "test", "val"]:
					for field in ["source", "target"]:
						snake_case_   		=			'\n'.join([contents[field]] * n_lines[split]				)
						with open(os.path.join(__UpperCamelCase			,					f"""{split}.{field}"""				)			,					'w'				) as f:
							f.write(__UpperCamelCase				)


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase = "pytorch"				):

				"""simple docstring"""


				snake_case_   		=			self.get_auto_remove_tmp_dir()
				snake_case_   		=			os.path.join(__UpperCamelCase			,					'output'				)
				snake_case_   		=			os.path.join(__UpperCamelCase			,					'data'				)
				self._create_dummy_data(data_dir=__UpperCamelCase				)

				snake_case_   		=			f"""
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            """.split()

				if gpus > 0:
					testargs.append(f"""--gpus={gpus}"""				)
					if is_apex_available():
						testargs.append('--fp16'				)
				else:
					testargs.append('--gpus=0'				)
					testargs.append('--distributed_backend=ddp_cpu'				)
					testargs.append('--num_processes=2'				)

				snake_case_   		=			[sys.executable, str(Path(finetune_rag.__file__				).resolve()				)] + testargs
				execute_subprocess_async(__UpperCamelCase			,					env=self.get_env()				)

				snake_case_   		=			os.path.join(__UpperCamelCase			,					'metrics.json'				)
				with open(__UpperCamelCase				) as f:
					snake_case_   		=			json.load(__UpperCamelCase				)
				return result


			@require_torch_gpu
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)


			@require_torch_multi_gpu
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=2				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)


			@require_torch_gpu
			@require_ray
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1			,					distributed_retriever='ray'				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)







			@require_torch_multi_gpu
			@require_ray
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self._run_finetune(gpus=1			,					distributed_retriever='ray'				)
				self.assertGreaterEqual(result['test'][0]['test_avg_em']			,					0.2				)
