






import unittest

import numpy as np

from diffusers import LMSDiscreteScheduler, OnnxStableDiffusionInpaintPipeline
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 					_snake_case						(			_A       ,    unittest.TestCase ):
	# FIXME: add fast tests
	pass




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






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






	@property
	def    lowerCAmelCase_							(		self       )						->    Tuple:
					snake_case__      :Optional[int]         =   ort.SessionOptions()
					snake_case__      :Any         =   False
					return options






	def    lowerCAmelCase_							(		self       )						->    Tuple:
					snake_case__      :Union[str, Any]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
					    "/in_paint/overture-creations-5sI6fQgYIuo.png"       )
					snake_case__      :Optional[Any]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
					    "/in_paint/overture-creations-5sI6fQgYIuo_mask.png"       )
					snake_case__      :List[Any]         =   OnnxStableDiffusionInpaintPipeline.from_pretrained(
					    "runwayml/stable-diffusion-inpainting"					,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__      :Tuple         =   "A red cat sitting on a park bench"

					snake_case__      :Optional[Any]         =   np.random.RandomState(0       )
					snake_case__      :List[str]         =   pipe(
					    prompt=UpperCamelCase					,image=UpperCamelCase					,mask_image=UpperCamelCase					,guidance_scale=7.5					,num_inference_steps=10					,generator=UpperCamelCase					,output_type="np"					,)
					snake_case__      :int         =   output.images
					snake_case__      :Optional[int]         =   images[0, 255:258, 255:258, -1]

					assert images.shape == (1, 512, 512, 3)
					snake_case__      :Any         =   np.array([0.2514, 0.3007, 0.3517, 0.1790, 0.2382, 0.3167, 0.1944, 0.2273, 0.2464]       )

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






	def    lowerCAmelCase_							(		self       )						->    List[Any]:
					snake_case__      :Optional[Any]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
					    "/in_paint/overture-creations-5sI6fQgYIuo.png"       )
					snake_case__      :List[str]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
					    "/in_paint/overture-creations-5sI6fQgYIuo_mask.png"       )
					snake_case__      :Optional[Any]         =   LMSDiscreteScheduler.from_pretrained(
					    "runwayml/stable-diffusion-inpainting"					,subfolder="scheduler"					,revision="onnx"       )
					snake_case__      :Union[str, Any]         =   OnnxStableDiffusionInpaintPipeline.from_pretrained(
					    "runwayml/stable-diffusion-inpainting"					,revision="onnx"					,scheduler=UpperCamelCase					,safety_checker=UpperCamelCase					,feature_extractor=UpperCamelCase					,provider=self.gpu_provider					,sess_options=self.gpu_options					,)
					pipe.set_progress_bar_config(disable=UpperCamelCase       )

					snake_case__      :Tuple         =   "A red cat sitting on a park bench"

					snake_case__      :Dict         =   np.random.RandomState(0       )
					snake_case__      :Optional[Any]         =   pipe(
					    prompt=UpperCamelCase					,image=UpperCamelCase					,mask_image=UpperCamelCase					,guidance_scale=7.5					,num_inference_steps=20					,generator=UpperCamelCase					,output_type="np"					,)
					snake_case__      :Tuple         =   output.images
					snake_case__      :str         =   images[0, 255:258, 255:258, -1]

					assert images.shape == (1, 512, 512, 3)
					snake_case__      :List[Any]         =   np.array([0.0086, 0.0077, 0.0083, 0.0093, 0.0107, 0.0139, 0.0094, 0.0097, 0.0125]       )

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







import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
	import jax
	import jax.numpy as jnp
	from flax.jax_utils import replicate
	from flax.training.common_utils import shard


@slow
@require_flax
class 					_snake_case						(			unittest.TestCase ):






	def    lowerCAmelCase_							(		self       )						->    int:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()






	def    lowerCAmelCase_							(		self       )						->    str:
					snake_case__    ,						snake_case__      :Tuple         =   FlaxControlNetModel.from_pretrained(
					    "lllyasviel/sd-controlnet-canny"					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__    ,						snake_case__      :Any         =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    "runwayml/stable-diffusion-v1-5"					,controlnet=UpperCamelCase					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__      :List[str]         =   controlnet_params

					snake_case__      :Union[str, Any]         =   "bird"
					snake_case__      :Optional[int]         =   jax.device_count()
					snake_case__      :Tuple         =   pipe.prepare_text_inputs([prompts] * num_samples       )

					snake_case__      :Union[str, Any]         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"       )
					snake_case__      :str         =   pipe.prepare_image_inputs([canny_image] * num_samples       )

					snake_case__      :List[str]         =   jax.random.PRNGKey(0       )
					snake_case__      :str         =   jax.random.split(UpperCamelCase					,jax.device_count()       )

					snake_case__      :int         =   replicate(UpperCamelCase       )
					snake_case__      :Any         =   shard(UpperCamelCase       )
					snake_case__      :Any         =   shard(UpperCamelCase       )

					snake_case__      :str         =   pipe(
					    prompt_ids=UpperCamelCase					,image=UpperCamelCase					,params=UpperCamelCase					,prng_seed=UpperCamelCase					,num_inference_steps=50					,jit=UpperCamelCase					,).images
					assert images.shape == (jax.device_count(), 1, 768, 512, 3)

					snake_case__      :List[str]         =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
					snake_case__      :Any         =   images[0, 253:256, 253:256, -1]

					snake_case__      :Union[str, Any]         =   jnp.asarray(jax.device_get(image_slice.flatten()       )       )
					snake_case__      :List[Any]         =   jnp.array(
					    [0.167969, 0.116699, 0.081543, 0.154297, 0.132812, 0.108887, 0.169922, 0.169922, 0.205078]       )
					print(f'output_slice: {output_slice}'       )
					assert jnp.abs(output_slice - expected_slice       ).max() < 1E-2






	def    lowerCAmelCase_							(		self       )						->    Optional[int]:
					snake_case__    ,						snake_case__      :List[str]         =   FlaxControlNetModel.from_pretrained(
					    "lllyasviel/sd-controlnet-openpose"					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__    ,						snake_case__      :Optional[Any]         =   FlaxStableDiffusionControlNetPipeline.from_pretrained(
					    "runwayml/stable-diffusion-v1-5"					,controlnet=UpperCamelCase					,from_pt=UpperCamelCase					,dtype=jnp.bfloataa       )
					snake_case__      :str         =   controlnet_params

					snake_case__      :int         =   "Chef in the kitchen"
					snake_case__      :List[Any]         =   jax.device_count()
					snake_case__      :Dict         =   pipe.prepare_text_inputs([prompts] * num_samples       )

					snake_case__      :Any         =   load_image(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png"       )
					snake_case__      :Optional[int]         =   pipe.prepare_image_inputs([pose_image] * num_samples       )

					snake_case__      :List[str]         =   jax.random.PRNGKey(0       )
					snake_case__      :Any         =   jax.random.split(UpperCamelCase					,jax.device_count()       )

					snake_case__      :Dict         =   replicate(UpperCamelCase       )
					snake_case__      :Tuple         =   shard(UpperCamelCase       )
					snake_case__      :Optional[int]         =   shard(UpperCamelCase       )

					snake_case__      :Optional[Any]         =   pipe(
					    prompt_ids=UpperCamelCase					,image=UpperCamelCase					,params=UpperCamelCase					,prng_seed=UpperCamelCase					,num_inference_steps=50					,jit=UpperCamelCase					,).images
					assert images.shape == (jax.device_count(), 1, 768, 512, 3)

					snake_case__      :int         =   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
					snake_case__      :List[str]         =   images[0, 253:256, 253:256, -1]

					snake_case__      :Tuple         =   jnp.asarray(jax.device_get(image_slice.flatten()       )       )
					snake_case__      :List[str]         =   jnp.array(
					    [[0.271484, 0.261719, 0.275391, 0.277344, 0.279297, 0.291016, 0.294922, 0.302734, 0.302734]]       )
					print(f'output_slice: {output_slice}'       )
					assert jnp.abs(output_slice - expected_slice       ).max() < 1E-2