





"""simple docstring"""






import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


__UpperCAmelCase						=			16
__UpperCAmelCase						=			32






def      lowercase__ (					lowerCAmelCase__   :			Accelerator							,			lowerCAmelCase__   :			DatasetDict							,			lowerCAmelCase__   :			List[int]							,			lowerCAmelCase__   :			List[int]							,			lowerCAmelCase__   :			int = 1_6			)    ->			Tuple:



			'''simple docstring'''
			a__  :	List[Any]		 =			AutoTokenizer.from_pretrained("bert-base-cased"			)
			a__  :	str		 =			DatasetDict(
			    {
			        "train": dataset["train"].select(lowerCAmelCase__			),
			        "validation": dataset["train"].select(lowerCAmelCase__			),
			        "test": dataset["validation"],
			    }			)

			def tokenize_function(lowerCAmelCase__   :			Any			):
						# max_length=None => use the model max length (it's actually the default)
						a__  :	List[Any]		 =			tokenizer(examples["sentence1"]							,			examples["sentence2"]							,			truncation=lowerCAmelCase__							,			max_length=lowerCAmelCase__			)
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						a__  :	List[Any]		 =			datasets.map(
						    lowerCAmelCase__							,			batched=lowerCAmelCase__							,			remove_columns=["idx", "sentence1", "sentence2"]							,			)

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			a__  :	Any		 =			tokenized_datasets.rename_column("label"							,			"labels"			)

			def collate_fn(lowerCAmelCase__   :			Optional[int]			):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						a__  :	Union[str, Any]		 =			1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									a__  :	List[Any]		 =			1_6
						elif accelerator.mixed_precision != "no":
									a__  :	Tuple		 =			8
						else:
									a__  :	List[str]		 =			None

						return tokenizer.pad(
						    lowerCAmelCase__							,			padding="longest"							,			max_length=lowerCAmelCase__							,			pad_to_multiple_of=lowerCAmelCase__							,			return_tensors="pt"							,			)

			# Instantiate dataloaders.
			a__  :	Dict		 =			DataLoader(
			    tokenized_datasets["train"]							,			shuffle=lowerCAmelCase__							,			collate_fn=lowerCAmelCase__							,			batch_size=lowerCAmelCase__			)
			a__  :	int		 =			DataLoader(
			    tokenized_datasets["validation"]							,			shuffle=lowerCAmelCase__							,			collate_fn=lowerCAmelCase__							,			batch_size=lowerCAmelCase__			)

			a__  :	str		 =			DataLoader(
			    tokenized_datasets["test"]							,			shuffle=lowerCAmelCase__							,			collate_fn=lowerCAmelCase__							,			batch_size=lowerCAmelCase__			)

			return train_dataloader, eval_dataloader, test_dataloader






def      lowercase__ (					lowerCAmelCase__   :			Dict							,			lowerCAmelCase__   :			Optional[int]			)    ->			Any:



			'''simple docstring'''
			# New Code #
			a__  :	Optional[Any]		 =			[]
			# Download the dataset
			a__  :	Any		 =			load_dataset("glue"							,			"mrpc"			)
			# Create our splits
			a__  :	int		 =			StratifiedKFold(n_splits=int(args.num_folds			)			)
			# Initialize accelerator
			a__  :	Tuple		 =			Accelerator(cpu=args.cpu							,			mixed_precision=args.mixed_precision			)
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			a__  :	Dict		 =			config["lr"]
			a__  :	str		 =			int(config["num_epochs"]			)
			a__  :	Dict		 =			int(config["seed"]			)
			a__  :	str		 =			int(config["batch_size"]			)

			a__  :	Union[str, Any]		 =			evaluate.load("glue"							,			"mrpc"			)

			# If the batch size is too big we use gradient accumulation
			a__  :	Optional[int]		 =			1
			if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
						a__  :	Optional[Any]		 =			batch_size // MAX_GPU_BATCH_SIZE
						a__  :	Optional[Any]		 =			MAX_GPU_BATCH_SIZE

			set_seed(lowerCAmelCase__			)

			# New Code #
			# Create our folds:
			a__  :	Tuple		 =			kfold.split(np.zeros(datasets["train"].num_rows			)							,			datasets["train"]["label"]			)
			a__  :	Union[str, Any]		 =			[]
			# Iterate over them
			for i, (train_idxs, valid_idxs) in enumerate(lowerCAmelCase__			):
						a__ ,	a__ ,	a__  :	str		 =			get_fold_dataloaders(
						    lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			)
						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						a__  :	Tuple		 =			AutoModelForSequenceClassification.from_pretrained("bert-base-cased"							,			return_dict=lowerCAmelCase__			)

						# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
						# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
						# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
						a__  :	Dict		 =			model.to(accelerator.device			)

						# Instantiate optimizer
						a__  :	int		 =			AdamW(params=model.parameters()							,			lr=lowerCAmelCase__			)

						# Instantiate scheduler
						a__  :	Any		 =			get_linear_schedule_with_warmup(
						    optimizer=lowerCAmelCase__							,			num_warmup_steps=1_0_0							,			num_training_steps=(len(lowerCAmelCase__			) * num_epochs) // gradient_accumulation_steps							,			)

						# Prepare everything
						# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
						# prepare method.
						a__ ,	a__ ,	a__ ,	a__ ,	a__  :	Any		 =			accelerator.prepare(
						    lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__			)

						# Now we train the model
						for epoch in range(lowerCAmelCase__			):
									model.train()
									for step, batch in enumerate(lowerCAmelCase__			):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device			)
												a__  :	List[Any]		 =			model(**lowerCAmelCase__			)
												a__  :	Tuple		 =			outputs.loss
												a__  :	Optional[Any]		 =			loss / gradient_accumulation_steps
												accelerator.backward(lowerCAmelCase__			)
												if step % gradient_accumulation_steps == 0:
															optimizer.step()
															lr_scheduler.step()
															optimizer.zero_grad()

									model.eval()
									for step, batch in enumerate(lowerCAmelCase__			):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device			)
												with torch.no_grad():
															a__  :	Union[str, Any]		 =			model(**lowerCAmelCase__			)
												a__  :	List[str]		 =			outputs.logits.argmax(dim=-1			)
												a__ ,	a__  :	Tuple		 =			accelerator.gather_for_metrics((predictions, batch["labels"])			)
												metric.add_batch(
												    predictions=lowerCAmelCase__							,			references=lowerCAmelCase__							,			)

									a__  :	int		 =			metric.compute()
									# Use accelerator.print to print only on the main process.
									accelerator.print(F"epoch {epoch}:"							,			lowerCAmelCase__			)

						# New Code #
						# We also run predictions on the test set at the very end
						a__  :	Optional[Any]		 =			[]
						for step, batch in enumerate(lowerCAmelCase__			):
									# We could avoid this line since we set the accelerator with `device_placement=True`.
									batch.to(accelerator.device			)
									with torch.no_grad():
												a__  :	Union[str, Any]		 =			model(**lowerCAmelCase__			)
									a__  :	List[str]		 =			outputs.logits
									a__ ,	a__  :	List[Any]		 =			accelerator.gather_for_metrics((predictions, batch["labels"])			)
									fold_predictions.append(predictions.cpu()			)
									if i == 0:
												# We need all of the test predictions
												test_references.append(references.cpu()			)
        # Use accelerator.print to print only on the main process.
						test_predictions.append(torch.cat(lowerCAmelCase__							,			dim=0			)			)
						# We now need to release all our memory and get rid of the current model, optimizer, etc
						accelerator.free_memory()
			# New Code #
			# Finally we check the accuracy of our folded results:
			a__  :	int		 =			torch.cat(lowerCAmelCase__							,			dim=0			)
			a__  :	Any		 =			torch.stack(lowerCAmelCase__							,			dim=0			).sum(dim=0			).div(int(args.num_folds			)			).argmax(dim=-1			)
			a__  :	Any		 =			metric.compute(predictions=lowerCAmelCase__							,			references=lowerCAmelCase__			)
			accelerator.print("Average test metrics from all folds:"							,			lowerCAmelCase__			)






def      lowercase__ (					)    ->			List[str]:



			'''simple docstring'''
			a__  :	Optional[Any]		 =			argparse.ArgumentParser(description="Simple example of training script."			)
			parser.add_argument(
			    "--mixed_precision"							,			type=lowerCAmelCase__							,			default=lowerCAmelCase__							,			choices=["no", "fp16", "bf16", "fp8"]							,			help="Whether to use mixed precision. Choose"
			    "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
			    "and an Nvidia Ampere GPU."							,			)
			parser.add_argument("--cpu"							,			action="store_true"							,			help="If passed, will train on the CPU."			)
			# New Code #
			parser.add_argument("--num_folds"							,			type=lowerCAmelCase__							,			default=3							,			help="The number of splits to perform across the dataset"			)
			a__  :	Optional[Any]		 =			parser.parse_args()
			a__  :	int		 =			{"lr": 2E-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
			training_function(lowerCAmelCase__							,			lowerCAmelCase__			)


if __name__ == "__main__":
		main()






"""simple docstring"""






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 					__UpperCAmelCase							(						_UpperCamelCase						,  unittest.TestCase  ):
				__lowerCamelCase				:	List[Any]				     =				"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"
				def       UpperCAmelCase      (		self    :			str	, a_    :			Optional[Any]=0    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	Union[str, Any]		 =			np.random.RandomState(a_    )
							a__  :	List[Any]		 =			{
							    "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    :			Dict    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Tuple		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							pipe.set_progress_bar_config(disable=a_    )

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	List[str]		 =			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    :			Tuple    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	Dict		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							a__  :	List[str]		 =			PNDMScheduler.from_config(pipe.scheduler.config	, skip_prk_steps=a_    )
							pipe.set_progress_bar_config(disable=a_    )

							a__  :	Optional[int]		 =			self.get_dummy_inputs()
							a__  :	List[Any]		 =			pipe(**a_    ).images
							a__  :	Union[str, Any]		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	List[str]		 =			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    :			Tuple    )	->	Tuple:



							'''simple docstring'''



							a__  :	Dict		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							a__  :	Tuple		 =			LMSDiscreteScheduler.from_config(pipe.scheduler.config    )
							pipe.set_progress_bar_config(disable=a_    )

							a__  :	List[Any]		 =			self.get_dummy_inputs()
							a__  :	Optional[Any]		 =			pipe(**a_    ).images
							a__  :	int		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	Dict		 =			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    :			Optional[Any]    )	->	Dict:



							'''simple docstring'''



							a__  :	Union[str, Any]		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							a__  :	Dict		 =			EulerDiscreteScheduler.from_config(pipe.scheduler.config    )
							pipe.set_progress_bar_config(disable=a_    )

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	str		 =			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    :			Union[str, Any]    )	->	str:



							'''simple docstring'''



							a__  :	str		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							a__  :	Dict		 =			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
							pipe.set_progress_bar_config(disable=a_    )

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	Any		 =			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    :			Tuple    )	->	Any:



							'''simple docstring'''



							a__  :	str		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							a__  :	str		 =			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
							pipe.set_progress_bar_config(disable=a_    )

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a__  :	Any		 =			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    )	->	Tuple:



							'''simple docstring'''



							a__  :	Dict		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							pipe.set_progress_bar_config(disable=a_    )

							a__  :	Any		 =			self.get_dummy_inputs()
							a__  :	Any		 =			3 * [inputs["prompt"]]

							# forward
							a__  :	Union[str, Any]		 =			pipe(**a_    )
							a__  :	int		 =			output.images[0, -3:, -3:, -1]

							a__  :	Union[str, Any]		 =			self.get_dummy_inputs()
							a__  :	List[Any]		 =			3 * [inputs.pop("prompt"    )]

							a__  :	Optional[Any]		 =			pipe.tokenizer(
							    a_	, padding="max_length"	, max_length=pipe.tokenizer.model_max_length	, truncation=a_	, return_tensors="np"	, )
							a__  :	List[str]		 =			text_inputs["input_ids"]

							a__  :	int		 =			pipe.text_encoder(input_ids=text_inputs.astype(np.intaa    )    )[0]

							a__  :	List[Any]		 =			prompt_embeds

							# forward
							a__  :	List[Any]		 =			pipe(**a_    )
							a__  :	List[str]		 =			output.images[0, -3:, -3:, -1]

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






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



							'''simple docstring'''



							a__  :	Tuple		 =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	, provider="CPUExecutionProvider"    )
							pipe.set_progress_bar_config(disable=a_    )

							a__  :	Tuple		 =			self.get_dummy_inputs()
							a__  :	Dict		 =			3 * ["this is a negative prompt"]
							a__  :	Optional[Any]		 =			negative_prompt
							a__  :	Any		 =			3 * [inputs["prompt"]]

							# forward
							a__  :	str		 =			pipe(**a_    )
							a__  :	List[str]		 =			output.images[0, -3:, -3:, -1]

							a__  :	Union[str, Any]		 =			self.get_dummy_inputs()
							a__  :	Union[str, Any]		 =			3 * [inputs.pop("prompt"    )]

							a__  :	List[Any]		 =			[]
							for p in [prompt, negative_prompt]:
										a__  :	int		 =			pipe.tokenizer(
										    a_	, padding="max_length"	, max_length=pipe.tokenizer.model_max_length	, truncation=a_	, return_tensors="np"	, )
										a__  :	Any		 =			text_inputs["input_ids"]

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

							a__ ,	a__  :	Union[str, Any]		 =			embeds

							# forward
							a__  :	Dict		 =			pipe(**a_    )
							a__  :	Any		 =			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 					__UpperCAmelCase							(						unittest.TestCase  ):
				@property
				def       UpperCAmelCase      (		self    :			Optional[Any]    )	->	Optional[int]:



							'''simple docstring'''



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



							'''simple docstring'''



							a__  :	List[str]		 =			ort.SessionOptions()
							a__  :	List[str]		 =			False
							return options
				def       UpperCAmelCase      (		self    :			Optional[int]    )	->	List[Any]:



							'''simple docstring'''



							a__  :	Dict		 =			OnnxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"	, revision="onnx"	, safety_checker=a_	, feature_extractor=a_	, provider=self.gpu_provider	, sess_options=self.gpu_options	, )
							sd_pipe.set_progress_bar_config(disable=a_    )

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

							a__  :	Optional[int]		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a__  :	str		 =			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    :			Union[str, Any]    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	Any		 =			DDIMScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"	, subfolder="scheduler"	, revision="onnx"    )
							a__  :	str		 =			OnnxStableDiffusionPipeline.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	, )
							sd_pipe.set_progress_bar_config(disable=a_    )

							a__  :	str		 =			"open neural network exchange"
							a__  :	Tuple		 =			np.random.RandomState(0    )
							a__  :	Dict		 =			sd_pipe([prompt]	, guidance_scale=7.5	, num_inference_steps=10	, generator=a_	, output_type="np"    )
							a__  :	Dict		 =			output.images
							a__  :	Union[str, Any]		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a__  :	Dict		 =			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[int]    )	->	Dict:



							'''simple docstring'''



							a__  :	List[Any]		 =			LMSDiscreteScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"	, subfolder="scheduler"	, revision="onnx"    )
							a__  :	List[Any]		 =			OnnxStableDiffusionPipeline.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	, )
							sd_pipe.set_progress_bar_config(disable=a_    )

							a__  :	Any		 =			"open neural network exchange"
							a__  :	Optional[Any]		 =			np.random.RandomState(0    )
							a__  :	Optional[int]		 =			sd_pipe([prompt]	, guidance_scale=7.5	, num_inference_steps=10	, generator=a_	, output_type="np"    )
							a__  :	int		 =			output.images
							a__  :	Tuple		 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a__  :	Dict		 =			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    :			int    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	List[str]		 =			0

							def test_callback_fn(a_    :			int	, a_    :			int	, a_    :			np.ndarray    ) -> None:
										a__  :	Optional[int]		 =			True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 0:
													assert latents.shape == (1, 4, 64, 64)
													a__  :	Any		 =			latents[0, -3:, -3:, -1]
													a__  :	Union[str, Any]		 =			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__  :	Union[str, Any]		 =			latents[0, -3:, -3:, -1]
													a__  :	Optional[int]		 =			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__  :	Tuple		 =			False

							a__  :	Optional[int]		 =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"	, revision="onnx"	, safety_checker=a_	, feature_extractor=a_	, provider=self.gpu_provider	, sess_options=self.gpu_options	, )
							pipe.set_progress_bar_config(disable=a_    )

							a__  :	List[Any]		 =			"Andromeda galaxy in a bottle"

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






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



							'''simple docstring'''



							a__  :	Optional[Any]		 =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"	, revision="onnx"	, safety_checker=a_	, feature_extractor=a_	, provider=self.gpu_provider	, sess_options=self.gpu_options	, )
							assert isinstance(a_	, a_    )
							assert pipe.safety_checker is None

							a__  :	Tuple		 =			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(a_    )
										a__  :	Optional[int]		 =			OnnxStableDiffusionPipeline.from_pretrained(a_    )

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