"""simple docstring"""




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

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)





@dataclass
class UpperCAmelCase						( snake_case_  ):


				'''simple docstring'''






				lowerCamelCase_    = 4_2
				lowerCamelCase_    = 4_2


try:
	if not (is_transformers_available() and is_torch_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
	from .pipeline_cycle_diffusion import CycleDiffusionPipeline
	from .pipeline_stable_diffusion import StableDiffusionPipeline
	from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
	from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
	from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
	from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
	from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
	from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
	from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
	from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
	from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
	from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
	from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
	from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
	from .pipeline_stable_unclip import StableUnCLIPPipeline
	from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
	from .safety_checker import StableDiffusionSafetyChecker
	from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
	if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
	from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
	if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.26.0""")):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_objects import (
	    StableDiffusionDepthaImgPipeline,
	    StableDiffusionDiffEditPipeline,
	    StableDiffusionPixaPixZeroPipeline,
	)
else:
	from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
	from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
	from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
	if not (
	    is_torch_available()
	    and is_transformers_available()
	    and is_k_diffusion_available()
	    and is_k_diffusion_version(""">=""", """0.0.12""")
	):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
	from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
	if not (is_transformers_available() and is_onnx_available()):
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	from ...utils.dummy_onnx_objects import *  # noqa F403
else:
	from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
	from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
	from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
	from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
	from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
	import flax





	@flax.struct.dataclass
	class UpperCAmelCase						( snake_case_  ):


					'''simple docstring'''






					lowerCamelCase_    = 4_2
					lowerCamelCase_    = 4_2

	from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
	from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
	from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
	from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
	from .safety_checker_flax import FlaxStableDiffusionSafetyChecker







import os
import shutil
from pathlib import Path
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ..utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, is_onnx_available, logging


if is_onnx_available():
	import onnxruntime as ort


_UpperCAmelCase            =		logging.get_logger(__name__)

_UpperCAmelCase            =		{
    """tensor(bool)""": np.bool_,
    """tensor(int8)""": np.inta,
    """tensor(uint8)""": np.uinta,
    """tensor(int16)""": np.intaa,
    """tensor(uint16)""": np.uintaa,
    """tensor(int32)""": np.intaa,
    """tensor(uint32)""": np.uintaa,
    """tensor(int64)""": np.intaa,
    """tensor(uint64)""": np.uintaa,
    """tensor(float16)""": np.floataa,
    """tensor(float)""": np.floataa,
    """tensor(double)""": np.floataa,
}





class UpperCAmelCase						:


				'''simple docstring'''






				def __init__(						self		,     lowercase=None		,     **lowercase ):



						"""simple docstring"""





						logger.info('`diffusers.OnnxRuntimeModel` is experimental and might change in the future.' )
						A_		:    List[Any]         =				model
						A_		:    Dict         =				kwargs.get('model_save_dir'		,     lowercase )
						A_		:    List[str]         =				kwargs.get('latest_model_name'		,     lowercase )
				def __call__(						self		,     **lowercase ):



						"""simple docstring"""





						A_		:    str         =				{k: np.array(lowercase ) for k, v in kwargs.items()}
						return self.model.run(lowercase		,     lowercase )
				@staticmethod
				def   lowerCAmelCase_       (						lowercase		,     lowercase=None		,     lowercase=None ):



						"""simple docstring"""





						if provider is None:
								logger.info('No onnxruntime provider specified, using CPUExecutionProvider' )
								A_		:    List[Any]         =				'CPUExecutionProvider'

						return ort.InferenceSession(lowercase		,     providers=[provider]		,     sess_options=lowercase )
				def   lowerCAmelCase_       (						self		,     lowercase		,     lowercase = None		,     **lowercase ):



						"""simple docstring"""





						A_		:    str         =				file_name if file_name is not None else ONNX_WEIGHTS_NAME

						A_		:    Optional[int]         =				self.model_save_dir.joinpath(self.latest_model_name )
						A_		:    int         =				Path(lowercase ).joinpath(lowercase )
						try:
								shutil.copyfile(lowercase		,     lowercase )
						except shutil.SameFileError:
								pass

						# copy external weights (for models >2GB)
						A_		:    Optional[Any]         =				self.model_save_dir.joinpath(lowercase )
						if src_path.exists():
								A_		:    int         =				Path(lowercase ).joinpath(lowercase )
								try:
										shutil.copyfile(lowercase		,     lowercase )
								except shutil.SameFileError:
										pass
				def   lowerCAmelCase_       (						self		,     lowercase		,     **lowercase		,     ):



						"""simple docstring"""





						if os.path.isfile(lowercase ):
								logger.error(F'''Provided path ({save_directory}) should be a directory, not a file''' )
								return

						os.makedirs(lowercase		,     exist_ok=lowercase )

						# saving model weights/files
						self._save_pretrained(lowercase		,     **lowercase )
				@classmethod
				def   lowerCAmelCase_       (						cls		,     lowercase		,     lowercase = None		,     lowercase = None		,     lowercase = False		,     lowercase = None		,     lowercase = None		,     lowercase = None		,     lowercase = None		,     **lowercase		,     ):



						"""simple docstring"""





						A_		:    Any         =				file_name if file_name is not None else ONNX_WEIGHTS_NAME
						# load model from local directory
						if os.path.isdir(lowercase ):
								A_		:    Optional[int]         =				OnnxRuntimeModel.load_model(
								    os.path.join(lowercase		,     lowercase )		,     provider=lowercase		,     sess_options=lowercase )
								A_		:    Dict         =				Path(lowercase )
						# load model from hub
						else:
								# download model
								A_		:    List[str]         =				hf_hub_download(
								    repo_id=lowercase		,     filename=lowercase		,     use_auth_token=lowercase		,     revision=lowercase		,     cache_dir=lowercase		,     force_download=lowercase		,     )
								A_		:    int         =				Path(lowercase ).parent
								A_		:    Optional[Any]         =				Path(lowercase ).name
								A_		:    Any         =				OnnxRuntimeModel.load_model(lowercase		,     provider=lowercase		,     sess_options=lowercase )
						return cls(model=lowercase		,     **lowercase )







				@classmethod
				def   lowerCAmelCase_       (						cls		,     lowercase		,     lowercase = True		,     lowercase = None		,     lowercase = None		,     **lowercase		,     ):



						"""simple docstring"""





						A_		:    List[Any]         =				None
						if len(str(lowercase ).split('@' ) ) == 2:
								A_       ,      A_		:    int         =				model_id.split('@' )

						return cls._from_pretrained(
						    model_id=lowercase		,     revision=lowercase		,     cache_dir=lowercase		,     force_download=lowercase		,     use_auth_token=lowercase		,     **lowercase		,     )






