






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        lowercase__      (							_a,       unittest.TestCase       ):
			_UpperCAmelCase		:Any             =			"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"


			def      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Tuple=0 ):
										lowerCamelCase_	:							Any									=floats_tensor((1, 3, 128, 128)       ,	rng=random.Random(snake_case_ ) )
										lowerCamelCase_	:							List[str]									=np.random.RandomState(snake_case_ )
										lowerCamelCase_	:							Optional[int]									={
										    """prompt""": """A painting of a squirrel eating a burger""",
										    """image""": image,
										    """generator""": generator,
										    """num_inference_steps""": 3,
										    """strength""": 0.75,
										    """guidance_scale""": 7.5,
										    """output_type""": """numpy""",
										}
										return inputs


			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										lowerCamelCase_	:							str									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										pipe.set_progress_bar_config(disable=snake_case_ )

										lowerCamelCase_	:							Tuple									=self.get_dummy_inputs()
										lowerCamelCase_	:							Union[str, Any]									=pipe(**snake_case_ ).images
										lowerCamelCase_	:							List[Any]									=image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							int									=np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087] )
										assert np.abs(image_slice - expected_slice ).max() < 1E-1


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							str									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							int									=PNDMScheduler.from_config(pipe.scheduler.config       ,	skip_prk_steps=snake_case_ )
										pipe.set_progress_bar_config(disable=snake_case_ )

										lowerCamelCase_	:							Dict									=self.get_dummy_inputs()
										lowerCamelCase_	:							int									=pipe(**snake_case_ ).images
										lowerCamelCase_	:							Tuple									=image[0, -3:, -3:, -1]

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							Tuple									=np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154] )

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


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							int									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Dict									=LMSDiscreteScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=snake_case_ )

										# warmup pass to apply optimizations
										lowerCamelCase_	:							List[Any]									=pipe(**self.get_dummy_inputs() )

										lowerCamelCase_	:							List[str]									=self.get_dummy_inputs()
										lowerCamelCase_	:							Optional[int]									=pipe(**snake_case_ ).images
										lowerCamelCase_	:							Tuple									=image[0, -3:, -3:, -1]

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							Any									=np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203] )

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


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							List[str]									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Dict									=EulerDiscreteScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=snake_case_ )

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

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							Optional[Any]									=np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304] )

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


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							List[str]									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							List[Any]									=EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=snake_case_ )

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

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							int									=np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304] )

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


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Optional[Any]									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Union[str, Any]									=DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=snake_case_ )

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

										assert image.shape == (1, 128, 128, 3)
										lowerCamelCase_	:							List[str]									=np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552] )

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







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


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


			@property
			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Tuple									=ort.SessionOptions()
										lowerCamelCase_	:							Optional[Any]									=False
										return options


			def      UpperCAmelCase__						(				self					:  List[str] ):
										lowerCamelCase_	:							Any									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										lowerCamelCase_	:							str									=init_image.resize((768, 512) )
										# using the PNDM scheduler by default
										lowerCamelCase_	:							Optional[Any]									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"       ,	revision="onnx"       ,	safety_checker=snake_case_       ,	feature_extractor=snake_case_       ,	provider=self.gpu_provider       ,	sess_options=self.gpu_options       ,	)
										pipe.set_progress_bar_config(disable=snake_case_ )

										lowerCamelCase_	:							Dict									="""A fantasy landscape, trending on artstation"""

										lowerCamelCase_	:							str									=np.random.RandomState(0 )
										lowerCamelCase_	:							Union[str, Any]									=pipe(
										    prompt=snake_case_       ,	image=snake_case_       ,	strength=0.75       ,	guidance_scale=7.5       ,	num_inference_steps=10       ,	generator=snake_case_       ,	output_type="np"       ,	)
										lowerCamelCase_	:							str									=output.images
										lowerCamelCase_	:							Optional[Any]									=images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 768, 3)
										lowerCamelCase_	:							Optional[Any]									=np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019] )
										# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


			def      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							List[Any]									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										lowerCamelCase_	:							List[Any]									=init_image.resize((768, 512) )
										lowerCamelCase_	:							Tuple									=LMSDiscreteScheduler.from_pretrained(
										    "runwayml/stable-diffusion-v1-5"       ,	subfolder="scheduler"       ,	revision="onnx" )
										lowerCamelCase_	:							str									=OnnxStableDiffusionImgaImgPipeline.from_pretrained(
										    "runwayml/stable-diffusion-v1-5"       ,	revision="onnx"       ,	scheduler=snake_case_       ,	safety_checker=snake_case_       ,	feature_extractor=snake_case_       ,	provider=self.gpu_provider       ,	sess_options=self.gpu_options       ,	)
										pipe.set_progress_bar_config(disable=snake_case_ )

										lowerCamelCase_	:							Union[str, Any]									="""A fantasy landscape, trending on artstation"""

										lowerCamelCase_	:							Optional[int]									=np.random.RandomState(0 )
										lowerCamelCase_	:							Optional[int]									=pipe(
										    prompt=snake_case_       ,	image=snake_case_       ,	strength=0.75       ,	guidance_scale=7.5       ,	num_inference_steps=20       ,	generator=snake_case_       ,	output_type="np"       ,	)
										lowerCamelCase_	:							Any									=output.images
										lowerCamelCase_	:							Tuple									=images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 768, 3)
										lowerCamelCase_	:							Tuple									=np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431] )
										# 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"""



# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)







@flax.struct.dataclass
class        lowercase__      :
			_UpperCAmelCase		:CommonSchedulerState

			# setable values
			_UpperCAmelCase		:jnp.ndarray
			_UpperCAmelCase		:jnp.ndarray
			_UpperCAmelCase		:Optional[int]             =			None


			@classmethod
			def      UpperCAmelCase__						(				cls					:  int       ,	snake_case__					:  CommonSchedulerState       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  jnp.ndarray ):
										return cls(common=snake_case__       ,	init_noise_sigma=snake_case__       ,	timesteps=snake_case__ )







@dataclass
class        lowercase__      (							snake_case__       ):
			_UpperCAmelCase		:DDPMSchedulerState







class        lowercase__      (							snake_case__,       snake_case__       ):
			_UpperCAmelCase		:Any             =			[e.name for e in FlaxKarrasDiffusionSchedulers]

			_UpperCAmelCase		:jnp.dtype


			@property
			def      UpperCAmelCase__						(				self					:  Optional[Any] ):
										return True


			@register_to_config
			def __init__(				self					:  Optional[int]       ,	snake_case__					:  int = 1000       ,	snake_case__					:  float = 0.0_001       ,	snake_case__					:  float = 0.02       ,	snake_case__					:  str = "linear"       ,	snake_case__					:  Optional[jnp.ndarray] = None       ,	snake_case__					:  str = "fixed_small"       ,	snake_case__					:  bool = True       ,	snake_case__					:  str = "epsilon"       ,	snake_case__					:  jnp.dtype = jnp.floataa       ,	):
										lowerCamelCase_	:							str									=dtype


			def      UpperCAmelCase__						(				self					:  List[str]       ,	snake_case__					:  Optional[CommonSchedulerState] = None ):
										if common is None:
																	lowerCamelCase_	:							int									=CommonSchedulerState.create(self )

										# standard deviation of the initial noise distribution
										lowerCamelCase_	:							Optional[Any]									=jnp.array(1.0       ,	dtype=self.dtype )

										lowerCamelCase_	:							str									=jnp.arange(0       ,	self.config.num_train_timesteps ).round()[::-1]

										return DDPMSchedulerState.create(
										    common=snake_case__       ,	init_noise_sigma=snake_case__       ,	timesteps=snake_case__       ,	)


			def      UpperCAmelCase__						(				self					:  Optional[int]       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  Optional[int] = None ):
										return sample


			def      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  int       ,	snake_case__					:  Tuple = () ):
										lowerCamelCase_	:							Any									=self.config.num_train_timesteps // num_inference_steps
										# creates integer timesteps by multiplying by ratio
										# rounding to avoid issues when num_inference_step is power of 3
										lowerCamelCase_	:							List[str]									=(jnp.arange(0       ,	snake_case__ ) * step_ratio).round()[::-1]

										return state.replace(
										    num_inference_steps=snake_case__       ,	timesteps=snake_case__       ,	)


			def      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  List[Any]=None       ,	snake_case__					:  Any=None ):
										lowerCamelCase_	:							List[str]									=state.common.alphas_cumprod[t]
										lowerCamelCase_	:							Union[str, Any]									=jnp.where(t > 0       ,	state.common.alphas_cumprod[t - 1]       ,	jnp.array(1.0       ,	dtype=self.dtype ) )

										# For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
										# and sample from it to get previous sample
										# x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
										lowerCamelCase_	:							Tuple									=(1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

										if variance_type is None:
																	lowerCamelCase_	:							List[Any]									=self.config.variance_type

										# hacks - were probably added for training stability
										if variance_type == "fixed_small":
																	lowerCamelCase_	:							List[str]									=jnp.clip(snake_case__       ,	a_min=1E-20 )
										# for rl-diffuser https://arxiv.org/abs/2205.09991
										elif variance_type == "fixed_small_log":
																	lowerCamelCase_	:							Dict									=jnp.log(jnp.clip(snake_case__       ,	a_min=1E-20 ) )
										elif variance_type == "fixed_large":
																	lowerCamelCase_	:							Optional[Any]									=state.common.betas[t]
										elif variance_type == "fixed_large_log":
																	# Glide max_log
																	lowerCamelCase_	:							Any									=jnp.log(state.common.betas[t] )
										elif variance_type == "learned":
																	return predicted_variance
										elif variance_type == "learned_range":
																	lowerCamelCase_	:							List[str]									=variance
																	lowerCamelCase_	:							Optional[int]									=state.common.betas[t]
																	lowerCamelCase_	:							Dict									=(predicted_variance + 1) / 2
																	lowerCamelCase_	:							Dict									=frac * max_log + (1 - frac) * min_log

										return variance


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  int       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  Optional[jax.random.KeyArray] = None       ,	snake_case__					:  bool = True       ,	):
										lowerCamelCase_	:							Union[str, Any]									=timestep

										if key is None:
																	lowerCamelCase_	:							Dict									=jax.random.PRNGKey(0 )

										if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
																	lowerCamelCase_							,						lowerCamelCase_	:							Optional[Any]									=jnp.split(snake_case__       ,	sample.shape[1]       ,	axis=1 )
										else:
																	lowerCamelCase_	:							List[str]									=None

										# 1. compute alphas, betas
										lowerCamelCase_	:							Union[str, Any]									=state.common.alphas_cumprod[t]
										lowerCamelCase_	:							Dict									=jnp.where(t > 0       ,	state.common.alphas_cumprod[t - 1]       ,	jnp.array(1.0       ,	dtype=self.dtype ) )
										lowerCamelCase_	:							Any									=1 - alpha_prod_t
										lowerCamelCase_	:							List[str]									=1 - alpha_prod_t_prev

										# 2. compute predicted original sample from predicted noise also called
										# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
										if self.config.prediction_type == "epsilon":
																	lowerCamelCase_	:							int									=(sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
										elif self.config.prediction_type == "sample":
																	lowerCamelCase_	:							List[Any]									=model_output
										elif self.config.prediction_type == "v_prediction":
																	lowerCamelCase_	:							Tuple									=(alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
										else:
																	raise ValueError(
																	    F"""prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` """
																	    " for the FlaxDDPMScheduler." )

										# 3. Clip "predicted x_0"
										if self.config.clip_sample:
																	lowerCamelCase_	:							List[Any]									=jnp.clip(snake_case__       ,	-1       ,	1 )

										# 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
										# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
										lowerCamelCase_	:							int									=(alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
										lowerCamelCase_	:							Optional[Any]									=state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

										# 5. Compute predicted previous sample µ_t
										# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
										lowerCamelCase_	:							Any									=pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

										# 6. Add noise
										def random_variance():
																	lowerCamelCase_	:							Union[str, Any]									=jax.random.split(snake_case__       ,	num=1 )
																	lowerCamelCase_	:							List[Any]									=jax.random.normal(snake_case__       ,	shape=model_output.shape       ,	dtype=self.dtype )
																	return (self._get_variance(snake_case__       ,	snake_case__       ,	predicted_variance=snake_case__ ) ** 0.5) * noise

										lowerCamelCase_	:							Tuple									=jnp.where(t > 0       ,	random_variance()       ,	jnp.zeros(model_output.shape       ,	dtype=self.dtype ) )

										lowerCamelCase_	:							str									=pred_prev_sample + variance

										if not return_dict:
																	return (pred_prev_sample, state)

										return FlaxDDPMSchedulerOutput(prev_sample=snake_case__       ,	state=snake_case__ )


			def      UpperCAmelCase__						(				self					:  Dict       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  jnp.ndarray       ,	):
										return add_noise_common(state.common       ,	snake_case__       ,	snake_case__       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  DDPMSchedulerState       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  jnp.ndarray       ,	snake_case__					:  jnp.ndarray       ,	):
										return get_velocity_common(state.common       ,	snake_case__       ,	snake_case__       ,	snake_case__ )


			def __len__(				self					:  Tuple ):
										return self.config.num_train_timesteps



