
"""simple docstring"""





import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD



class __lowerCAmelCase					(       nn.Module   ):




			"""simple docstring"""







			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					0.0
			snake_case						   =					1
			snake_case						   =					1
			snake_case						   =					True
			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					jnp.floataa




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			int:







								"""simple docstring"""
								A_								=    []
								A_								=    []

								for i in range(self.num_layers						):
													A_								=    self.in_channels if i == 0 else self.out_channels

													A_								=    FlaxResnetBlockaD(
													    in_channels=_snake_case							, out_channels=self.out_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
													resnets.append(_snake_case						)

													A_								=    FlaxTransformeraDModel(
													    in_channels=self.out_channels							, n_heads=self.num_attention_heads							, d_head=self.out_channels // self.num_attention_heads							, depth=1							, use_linear_projection=self.use_linear_projection							, only_cross_attention=self.only_cross_attention							, use_memory_efficient_attention=self.use_memory_efficient_attention							, dtype=self.dtype							, )
													attentions.append(_snake_case						)

								A_								=    resnets
								A_								=    attentions

								if self.add_downsample:
													A_								=    FlaxDownsampleaD(self.out_channels							, dtype=self.dtype						)






			def __call__(					self    :       str							, _snake_case    :       Optional[int]							, _snake_case    :       Tuple							, _snake_case    :       List[Any]							, _snake_case    :       int=True						)   ->			Optional[Any]:







								"""simple docstring"""
								A_								=    ()

								for resnet, attn in zip(self.resnets							, self.attentions						):
													A_								=    resnet(_snake_case							, _snake_case							, deterministic=_snake_case						)
													A_								=    attn(_snake_case							, _snake_case							, deterministic=_snake_case						)
													output_states += (hidden_states,)

								if self.add_downsample:
													A_								=    self.downsamplers_a(_snake_case						)
													output_states += (hidden_states,)

								return hidden_states, output_states





class __lowerCAmelCase					(       nn.Module   ):




			"""simple docstring"""







			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					0.0
			snake_case						   =					1
			snake_case						   =					True
			snake_case						   =					jnp.floataa




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_								=    []

								for i in range(self.num_layers						):
													A_								=    self.in_channels if i == 0 else self.out_channels

													A_								=    FlaxResnetBlockaD(
													    in_channels=_snake_case							, out_channels=self.out_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
													resnets.append(_snake_case						)
								A_								=    resnets

								if self.add_downsample:
													A_								=    FlaxDownsampleaD(self.out_channels							, dtype=self.dtype						)






			def __call__(					self    :       Optional[Any]							, _snake_case    :       Dict							, _snake_case    :       Union[str, Any]							, _snake_case    :       Optional[int]=True						)   ->			str:







								"""simple docstring"""
								A_								=    ()

								for resnet in self.resnets:
													A_								=    resnet(_snake_case							, _snake_case							, deterministic=_snake_case						)
													output_states += (hidden_states,)

								if self.add_downsample:
													A_								=    self.downsamplers_a(_snake_case						)
													output_states += (hidden_states,)

								return hidden_states, output_states





class __lowerCAmelCase					(       nn.Module   ):




			"""simple docstring"""







			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					0.0
			snake_case						   =					1
			snake_case						   =					1
			snake_case						   =					True
			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					jnp.floataa




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			int:







								"""simple docstring"""
								A_								=    []
								A_								=    []

								for i in range(self.num_layers						):
													A_								=    self.in_channels if (i == self.num_layers - 1) else self.out_channels
													A_								=    self.prev_output_channel if i == 0 else self.out_channels

													A_								=    FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels							, out_channels=self.out_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
													resnets.append(_snake_case						)

													A_								=    FlaxTransformeraDModel(
													    in_channels=self.out_channels							, n_heads=self.num_attention_heads							, d_head=self.out_channels // self.num_attention_heads							, depth=1							, use_linear_projection=self.use_linear_projection							, only_cross_attention=self.only_cross_attention							, use_memory_efficient_attention=self.use_memory_efficient_attention							, dtype=self.dtype							, )
													attentions.append(_snake_case						)

								A_								=    resnets
								A_								=    attentions

								if self.add_upsample:
													A_								=    FlaxUpsampleaD(self.out_channels							, dtype=self.dtype						)






			def __call__(					self    :       Optional[Any]							, _snake_case    :       List[Any]							, _snake_case    :       Any							, _snake_case    :       int							, _snake_case    :       Dict							, _snake_case    :       Optional[Any]=True						)   ->			Tuple:







								"""simple docstring"""
								for resnet, attn in zip(self.resnets							, self.attentions						):
													# pop res hidden states
													A_								=    res_hidden_states_tuple[-1]
													A_								=    res_hidden_states_tuple[:-1]
													A_								=    jnp.concatenate((hidden_states, res_hidden_states)							, axis=-1						)

													A_								=    resnet(_snake_case							, _snake_case							, deterministic=_snake_case						)
													A_								=    attn(_snake_case							, _snake_case							, deterministic=_snake_case						)

								if self.add_upsample:
													A_								=    self.upsamplers_a(_snake_case						)

								return hidden_states





class __lowerCAmelCase					(       nn.Module   ):




			"""simple docstring"""







			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					42
			snake_case						   =					0.0
			snake_case						   =					1
			snake_case						   =					True
			snake_case						   =					jnp.floataa




			def   lowerCamelCase__  (					self    :       str						)   ->			str:







								"""simple docstring"""
								A_								=    []

								for i in range(self.num_layers						):
													A_								=    self.in_channels if (i == self.num_layers - 1) else self.out_channels
													A_								=    self.prev_output_channel if i == 0 else self.out_channels

													A_								=    FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels							, out_channels=self.out_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
													resnets.append(_snake_case						)

								A_								=    resnets

								if self.add_upsample:
													A_								=    FlaxUpsampleaD(self.out_channels							, dtype=self.dtype						)






			def __call__(					self    :       List[Any]							, _snake_case    :       Union[str, Any]							, _snake_case    :       List[str]							, _snake_case    :       Dict							, _snake_case    :       Union[str, Any]=True						)   ->			Optional[int]:







								"""simple docstring"""
								for resnet in self.resnets:
													# pop res hidden states
													A_								=    res_hidden_states_tuple[-1]
													A_								=    res_hidden_states_tuple[:-1]
													A_								=    jnp.concatenate((hidden_states, res_hidden_states)							, axis=-1						)

													A_								=    resnet(_snake_case							, _snake_case							, deterministic=_snake_case						)

								if self.add_upsample:
													A_								=    self.upsamplers_a(_snake_case						)

								return hidden_states





class __lowerCAmelCase					(       nn.Module   ):




			"""simple docstring"""







			snake_case						   =					42
			snake_case						   =					0.0
			snake_case						   =					1
			snake_case						   =					1
			snake_case						   =					False
			snake_case						   =					False
			snake_case						   =					jnp.floataa




			def   lowerCamelCase__  (					self    :       int						)   ->			List[Any]:







								"""simple docstring"""
								# there is always at least one resnet
								A_								=    [
								    FlaxResnetBlockaD(
								        in_channels=self.in_channels							, out_channels=self.in_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
								]

								A_								=    []

								for _ in range(self.num_layers						):
													A_								=    FlaxTransformeraDModel(
													    in_channels=self.in_channels							, n_heads=self.num_attention_heads							, d_head=self.in_channels // self.num_attention_heads							, depth=1							, use_linear_projection=self.use_linear_projection							, use_memory_efficient_attention=self.use_memory_efficient_attention							, dtype=self.dtype							, )
													attentions.append(_snake_case						)

													A_								=    FlaxResnetBlockaD(
													    in_channels=self.in_channels							, out_channels=self.in_channels							, dropout_prob=self.dropout							, dtype=self.dtype							, )
													resnets.append(_snake_case						)

								A_								=    resnets
								A_								=    attentions






			def __call__(					self    :       Optional[Any]							, _snake_case    :       Optional[Any]							, _snake_case    :       int							, _snake_case    :       Tuple							, _snake_case    :       List[str]=True						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_								=    self.resnets[0](_snake_case							, _snake_case						)
								for attn, resnet in zip(self.attentions							, self.resnets[1:]						):
													A_								=    attn(_snake_case							, _snake_case							, deterministic=_snake_case						)
													A_								=    resnet(_snake_case							, _snake_case							, deterministic=_snake_case						)

								return hidden_states






"""simple docstring"""





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 __lowerCAmelCase					(       _lowercase  ,					unittest.TestCase   ):




			"""simple docstring"""







			pass





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




			"""simple docstring"""







			@property
			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Dict:







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




			@property
			def   lowerCamelCase__  (					self    :       List[str]						)   ->			List[Any]:







								"""simple docstring"""
								A_								=    ort.SessionOptions()
								A_								=    False
								return options




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			List[Any]:







								"""simple docstring"""
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/in_paint/overture-creations-5sI6fQgYIuo.png"						)
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/in_paint/overture-creations-5sI6fQgYIuo_mask.png"						)
								A_								=    OnnxStableDiffusionInpaintPipeline.from_pretrained(
								    "runwayml/stable-diffusion-inpainting"							, 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						)

								A_								=    "A red cat sitting on a park bench"

								A_								=    np.random.RandomState(0						)
								A_								=    pipe(
								    prompt=_snake_case							, image=_snake_case							, mask_image=_snake_case							, guidance_scale=7.5							, num_inference_steps=10							, generator=_snake_case							, output_type="np"							, )
								A_								=    output.images
								A_								=    images[0, 255:258, 255:258, -1]

								assert images.shape == (1, 512, 512, 3)
								A_								=    np.array([0.2_5_1_4, 0.3_0_0_7, 0.3_5_1_7, 0.1_7_9_0, 0.2_3_8_2, 0.3_1_6_7, 0.1_9_4_4, 0.2_2_7_3, 0.2_4_6_4]						)

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






			def   lowerCamelCase__  (					self    :       Optional[int]						)   ->			Tuple:







								"""simple docstring"""
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/in_paint/overture-creations-5sI6fQgYIuo.png"						)
								A_								=    load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/in_paint/overture-creations-5sI6fQgYIuo_mask.png"						)
								A_								=    LMSDiscreteScheduler.from_pretrained(
								    "runwayml/stable-diffusion-inpainting"							, subfolder="scheduler"							, revision="onnx"						)
								A_								=    OnnxStableDiffusionInpaintPipeline.from_pretrained(
								    "runwayml/stable-diffusion-inpainting"							, 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						)

								A_								=    "A red cat sitting on a park bench"

								A_								=    np.random.RandomState(0						)
								A_								=    pipe(
								    prompt=_snake_case							, image=_snake_case							, mask_image=_snake_case							, guidance_scale=7.5							, num_inference_steps=20							, generator=_snake_case							, output_type="np"							, )
								A_								=    output.images
								A_								=    images[0, 255:258, 255:258, -1]

								assert images.shape == (1, 512, 512, 3)
								A_								=    np.array([0.0_0_8_6, 0.0_0_7_7, 0.0_0_8_3, 0.0_0_9_3, 0.0_1_0_7, 0.0_1_3_9, 0.0_0_9_4, 0.0_0_9_7, 0.0_1_2_5]						)

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




