


'''simple docstring'''

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       a__				( __A							,						unittest.TestCase      ):







				"""simple docstring"""



				__UpperCamelCase    :		Dict	=	'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'




				def 					_snake_case  (self       ,			__lowercase=0		):
										__lowerCAmelCase							       =				floats_tensor((1, 3, 1_28, 1_28)       ,			rng=random.Random(__lowercase		)		)
										__lowerCAmelCase							       =				np.random.RandomState(__lowercase		)
										__lowerCAmelCase							       =				{
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''image''': image,
										    '''generator''': generator,
										    '''num_inference_steps''': 3,
										    '''strength''': 0.7_5,
										    '''guidance_scale''': 7.5,
										    '''output_type''': '''numpy''',
										}
										return inputs




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										pipe.set_progress_bar_config(disable=__lowercase		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]		)
										assert np.abs(image_slice - expected_slice		).max() < 1e-1




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										__lowerCAmelCase							       =				PNDMScheduler.from_config(pipe.scheduler.config       ,			skip_prk_steps=__lowercase		)
										pipe.set_progress_bar_config(disable=__lowercase		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										__lowerCAmelCase							       =				LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
										pipe.set_progress_bar_config(disable=__lowercase		)

										# warmup pass to apply optimizations
										__lowerCAmelCase							       =				pipe(**self.get_dummy_inputs()		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										__lowerCAmelCase							       =				EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
										pipe.set_progress_bar_config(disable=__lowercase		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										__lowerCAmelCase							       =				EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
										pipe.set_progress_bar_config(disable=__lowercase		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]		)

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




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,			provider='''CPUExecutionProvider'''		)
										__lowerCAmelCase							       =				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
										pipe.set_progress_bar_config(disable=__lowercase		)

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =				np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]		)

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


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







				"""simple docstring"""



				@property
				def 					_snake_case  (self		):
										return (
										    "CUDAExecutionProvider",
										    {
										        "gpu_mem_limit": "15000000000",  # 15GB
										        "arena_extend_strategy": "kSameAsRequested",
										    },
										)




				@property
				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				ort.SessionOptions()
										__lowerCAmelCase							       =				False
										return options




				def 					_snake_case  (self		):
										__lowerCAmelCase							       =				load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''		)
										__lowerCAmelCase							       =				init_image.resize((7_68, 5_12)		)
										# using the PNDM scheduler by default
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(
										    '''CompVis/stable-diffusion-v1-4'''       ,			revision='''onnx'''       ,			safety_checker=__lowercase       ,			feature_extractor=__lowercase       ,			provider=self.gpu_provider       ,			sess_options=self.gpu_options       ,			)
										pipe.set_progress_bar_config(disable=__lowercase		)

										__lowerCAmelCase							       =				'''A fantasy landscape, trending on artstation'''

										__lowerCAmelCase							       =				np.random.RandomState(0		)
										__lowerCAmelCase							       =				pipe(
										    prompt=__lowercase       ,			image=__lowercase       ,			strength=0.7_5       ,			guidance_scale=7.5       ,			num_inference_steps=10       ,			generator=__lowercase       ,			output_type='''np'''       ,			)
										__lowerCAmelCase							       =				output.images
										__lowerCAmelCase							       =				images[0, 2_55:2_58, 3_83:3_86, -1]

										assert images.shape == (1, 5_12, 7_68, 3)
										__lowerCAmelCase							       =				np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]		)
										# 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		):
										__lowerCAmelCase							       =				load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/img2img/sketch-mountains-input.jpg'''		)
										__lowerCAmelCase							       =				init_image.resize((7_68, 5_12)		)
										__lowerCAmelCase							       =				LMSDiscreteScheduler.from_pretrained(
										    '''runwayml/stable-diffusion-v1-5'''       ,			subfolder='''scheduler'''       ,			revision='''onnx'''		)
										__lowerCAmelCase							       =				OnnxStableDiffusionImgaImgPipeline.from_pretrained(
										    '''runwayml/stable-diffusion-v1-5'''       ,			revision='''onnx'''       ,			scheduler=__lowercase       ,			safety_checker=__lowercase       ,			feature_extractor=__lowercase       ,			provider=self.gpu_provider       ,			sess_options=self.gpu_options       ,			)
										pipe.set_progress_bar_config(disable=__lowercase		)

										__lowerCAmelCase							       =				'''A fantasy landscape, trending on artstation'''

										__lowerCAmelCase							       =				np.random.RandomState(0		)
										__lowerCAmelCase							       =				pipe(
										    prompt=__lowercase       ,			image=__lowercase       ,			strength=0.7_5       ,			guidance_scale=7.5       ,			num_inference_steps=20       ,			generator=__lowercase       ,			output_type='''np'''       ,			)
										__lowerCAmelCase							       =				output.images
										__lowerCAmelCase							       =				images[0, 2_55:2_58, 3_83:3_86, -1]

										assert images.shape == (1, 5_12, 7_68, 3)
										__lowerCAmelCase							       =				np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]		)
										# 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'''

from typing import List, Optional, Tuple, Union

import PIL
import torch
from torchvision import transforms

from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import DDIMScheduler
from diffusers.utils import randn_tensor


_UpperCAmelCase				: Tuple							    =				transforms.Compose(
    [
        transforms.Resize((2_5_6, 2_5_6)),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)

def 							__magic_name__(						lowerCamelCase):
						if isinstance(lowerCamelCase,				torch.Tensor):
												return image
						elif isinstance(lowerCamelCase,				PIL.Image.Image):
												__lowerCAmelCase							       =				[image]

						__lowerCAmelCase							       =				[trans(img.convert('''RGB''')) for img in image]
						__lowerCAmelCase							       =				torch.stack(lowerCamelCase)
						return image




class       a__				( __A      ):







				"""simple docstring"""



				def __init__(self       ,			__lowercase       ,			__lowercase		):
										super().__init__()

										# make sure scheduler can always be converted to DDIM
										__lowerCAmelCase							       =				DDIMScheduler.from_config(scheduler.config		)

										self.register_modules(unet=__lowercase       ,			scheduler=__lowercase		)




				def 					_snake_case  (self       ,			__lowercase		):
										if strength < 0 or strength > 1:
																raise ValueError(F"""The value of strength should in [0.0, 1.0] but is {strength}"""		)




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
										# get the original timestep using init_timestep
										__lowerCAmelCase							       =				min(int(num_inference_steps * strength		)       ,			__lowercase		)

										__lowerCAmelCase							       =				max(num_inference_steps - init_timestep       ,			0		)
										__lowerCAmelCase							       =				self.scheduler.timesteps[t_start:]

										return timesteps, num_inference_steps - t_start




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase=None		):
										if not isinstance(__lowercase       ,			(torch.Tensor, PIL.Image.Image, list)		):
																raise ValueError(
																    F"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(__lowercase		)}"""		)

										__lowerCAmelCase							       =				image.to(device=__lowercase       ,			dtype=__lowercase		)

										if isinstance(__lowercase       ,			__lowercase		) and len(__lowercase		) != batch_size:
																raise ValueError(
																    F"""You have passed a list of generators of length {len(__lowercase		)}, but requested an effective batch"""
																    F""" size of {batch_size}. Make sure the batch size matches the length of the generators."""		)

										__lowerCAmelCase							       =				init_latents.shape
										__lowerCAmelCase							       =				randn_tensor(__lowercase       ,			generator=__lowercase       ,			device=__lowercase       ,			dtype=__lowercase		)

										# get latents
										print('''add noise to latents at timestep'''       ,			__lowercase		)
										__lowerCAmelCase							       =				self.scheduler.add_noise(__lowercase       ,			__lowercase       ,			__lowercase		)
										__lowerCAmelCase							       =				init_latents

										return latents




				@torch.no_grad()
				def __call__(self       ,			__lowercase = None       ,			__lowercase = 0.8       ,			__lowercase = 1       ,			__lowercase = None       ,			__lowercase = 0.0       ,			__lowercase = 50       ,			__lowercase = None       ,			__lowercase = "pil"       ,			__lowercase = True       ,			):
										self.check_inputs(__lowercase		)

										# 2. Preprocess image
										__lowerCAmelCase							       =				preprocess(__lowercase		)

										# 3. set timesteps
										self.scheduler.set_timesteps(__lowercase       ,			device=self.device		)
										__lowerCAmelCase       ,					__lowerCAmelCase							       =				self.get_timesteps(__lowercase       ,			__lowercase       ,			self.device		)
										__lowerCAmelCase							       =				timesteps[:1].repeat(__lowercase		)

										# 4. Prepare latent variables
										__lowerCAmelCase							       =				self.prepare_latents(__lowercase       ,			__lowercase       ,			__lowercase       ,			self.unet.dtype       ,			self.device       ,			__lowercase		)
										__lowerCAmelCase							       =				latents

										# 5. Denoising loop
										for t in self.progress_bar(__lowercase		):
																# 1. predict noise model_output
																__lowerCAmelCase							       =				self.unet(__lowercase       ,			__lowercase		).sample

																# 2. predict previous mean of image x_t-1 and add variance depending on eta
																# eta corresponds to η in paper and should be between [0, 1]
																# do x_t -> x_t-1
																__lowerCAmelCase							       =				self.scheduler.step(
																    __lowercase       ,			__lowercase       ,			__lowercase       ,			eta=__lowercase       ,			use_clipped_model_output=__lowercase       ,			generator=__lowercase       ,			).prev_sample

										__lowerCAmelCase							       =				(image / 2 + 0.5).clamp(0       ,			1		)
										__lowerCAmelCase							       =				image.cpu().permute(0       ,			2       ,			3       ,			1		).numpy()
										if output_type == "pil":
																__lowerCAmelCase							       =				self.numpy_to_pil(__lowercase		)

										if not return_dict:
																return (image, latent_timestep.item())

										return ImagePipelineOutput(images=__lowercase		)

