






UpperCamelCase			    =     '''0.18.2'''

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
		if not is_onnx_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_onnx_objects import *  # noqa F403
else:
		from .pipelines import OnnxRuntimeModel

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_pt_objects import *  # noqa F403
else:
		from .models import (
		    AutoencoderKL,
		    ControlNetModel,
		    ModelMixin,
		    PriorTransformer,
		    TaFilmDecoder,
		    TransformeraDModel,
		    UNetaDModel,
		    UNetaDConditionModel,
		    UNetaDModel,
		    UNetaDConditionModel,
		    VQModel,
		)
		from .optimization import (
		    get_constant_schedule,
		    get_constant_schedule_with_warmup,
		    get_cosine_schedule_with_warmup,
		    get_cosine_with_hard_restarts_schedule_with_warmup,
		    get_linear_schedule_with_warmup,
		    get_polynomial_decay_schedule_with_warmup,
		    get_scheduler,
		)
		from .pipelines import (
		    AudioPipelineOutput,
		    ConsistencyModelPipeline,
		    DanceDiffusionPipeline,
		    DDIMPipeline,
		    DDPMPipeline,
		    DiffusionPipeline,
		    DiTPipeline,
		    ImagePipelineOutput,
		    KarrasVePipeline,
		    LDMPipeline,
		    LDMSuperResolutionPipeline,
		    PNDMPipeline,
		    RePaintPipeline,
		    ScoreSdeVePipeline,
		)
		from .schedulers import (
		    CMStochasticIterativeScheduler,
		    DDIMInverseScheduler,
		    DDIMParallelScheduler,
		    DDIMScheduler,
		    DDPMParallelScheduler,
		    DDPMScheduler,
		    DEISMultistepScheduler,
		    DPMSolverMultistepInverseScheduler,
		    DPMSolverMultistepScheduler,
		    DPMSolverSinglestepScheduler,
		    EulerAncestralDiscreteScheduler,
		    EulerDiscreteScheduler,
		    HeunDiscreteScheduler,
		    IPNDMScheduler,
		    KarrasVeScheduler,
		    KDPMaAncestralDiscreteScheduler,
		    KDPMaDiscreteScheduler,
		    PNDMScheduler,
		    RePaintScheduler,
		    SchedulerMixin,
		    ScoreSdeVeScheduler,
		    UnCLIPScheduler,
		    UniPCMultistepScheduler,
		    VQDiffusionScheduler,
		)
		from .training_utils import EMAModel

try:
		if not (is_torch_available() and is_scipy_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
		from .schedulers import LMSDiscreteScheduler

try:
		if not (is_torch_available() and is_torchsde_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
		from .schedulers import DPMSolverSDEScheduler

try:
		if not (is_torch_available() and is_transformers_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
		from .pipelines import (
		    AltDiffusionImgaImgPipeline,
		    AltDiffusionPipeline,
		    AudioLDMPipeline,
		    CycleDiffusionPipeline,
		    IFImgaImgPipeline,
		    IFImgaImgSuperResolutionPipeline,
		    IFInpaintingPipeline,
		    IFInpaintingSuperResolutionPipeline,
		    IFPipeline,
		    IFSuperResolutionPipeline,
		    ImageTextPipelineOutput,
		    KandinskyImgaImgPipeline,
		    KandinskyInpaintPipeline,
		    KandinskyPipeline,
		    KandinskyPriorPipeline,
		    KandinskyVaaControlnetImgaImgPipeline,
		    KandinskyVaaControlnetPipeline,
		    KandinskyVaaImgaImgPipeline,
		    KandinskyVaaInpaintPipeline,
		    KandinskyVaaPipeline,
		    KandinskyVaaPriorEmbaEmbPipeline,
		    KandinskyVaaPriorPipeline,
		    LDMTextToImagePipeline,
		    PaintByExamplePipeline,
		    SemanticStableDiffusionPipeline,
		    ShapEImgaImgPipeline,
		    ShapEPipeline,
		    StableDiffusionAttendAndExcitePipeline,
		    StableDiffusionControlNetImgaImgPipeline,
		    StableDiffusionControlNetInpaintPipeline,
		    StableDiffusionControlNetPipeline,
		    StableDiffusionDepthaImgPipeline,
		    StableDiffusionDiffEditPipeline,
		    StableDiffusionImageVariationPipeline,
		    StableDiffusionImgaImgPipeline,
		    StableDiffusionInpaintPipeline,
		    StableDiffusionInpaintPipelineLegacy,
		    StableDiffusionInstructPixaPixPipeline,
		    StableDiffusionLatentUpscalePipeline,
		    StableDiffusionLDMaDPipeline,
		    StableDiffusionModelEditingPipeline,
		    StableDiffusionPanoramaPipeline,
		    StableDiffusionParadigmsPipeline,
		    StableDiffusionPipeline,
		    StableDiffusionPipelineSafe,
		    StableDiffusionPixaPixZeroPipeline,
		    StableDiffusionSAGPipeline,
		    StableDiffusionUpscalePipeline,
		    StableUnCLIPImgaImgPipeline,
		    StableUnCLIPPipeline,
		    TextToVideoSDPipeline,
		    TextToVideoZeroPipeline,
		    UnCLIPImageVariationPipeline,
		    UnCLIPPipeline,
		    UniDiffuserModel,
		    UniDiffuserPipeline,
		    UniDiffuserTextDecoder,
		    VersatileDiffusionDualGuidedPipeline,
		    VersatileDiffusionImageVariationPipeline,
		    VersatileDiffusionPipeline,
		    VersatileDiffusionTextToImagePipeline,
		    VideoToVideoSDPipeline,
		    VQDiffusionPipeline,
		)

try:
		if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
		from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

try:
		if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
		from .pipelines import StableDiffusionKDiffusionPipeline

try:
		if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
		from .pipelines import (
		    OnnxStableDiffusionImgaImgPipeline,
		    OnnxStableDiffusionInpaintPipeline,
		    OnnxStableDiffusionInpaintPipelineLegacy,
		    OnnxStableDiffusionPipeline,
		    OnnxStableDiffusionUpscalePipeline,
		    StableDiffusionOnnxPipeline,
		)

try:
		if not (is_torch_available() and is_librosa_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
		from .pipelines import AudioDiffusionPipeline, Mel

try:
		if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
		from .pipelines import SpectrogramDiffusionPipeline

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_flax_objects import *  # noqa F403
else:
		from .models.controlnet_flax import FlaxControlNetModel
		from .models.modeling_flax_utils import FlaxModelMixin
		from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
		from .models.vae_flax import FlaxAutoencoderKL
		from .pipelines import FlaxDiffusionPipeline
		from .schedulers import (
		    FlaxDDIMScheduler,
		    FlaxDDPMScheduler,
		    FlaxDPMSolverMultistepScheduler,
		    FlaxKarrasVeScheduler,
		    FlaxLMSDiscreteScheduler,
		    FlaxPNDMScheduler,
		    FlaxSchedulerMixin,
		    FlaxScoreSdeVeScheduler,
		)


try:
		if not (is_flax_available() and is_transformers_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
		from .pipelines import (
		    FlaxStableDiffusionControlNetPipeline,
		    FlaxStableDiffusionImgaImgPipeline,
		    FlaxStableDiffusionInpaintPipeline,
		    FlaxStableDiffusionPipeline,
		)

try:
		if not (is_note_seq_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from .utils.dummy_note_seq_objects import *  # noqa F403
else:
		from .pipelines import MidiProcessor









import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
		import os

		from flax.core.frozen_dict import unfreeze
		from flax.traverse_util import flatten_dict

		from transformers import FlaxBertModel

		UpperCamelCase			    =     '''0.12'''  # assumed parallelism: 8



@require_flax
@is_staging_test
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@classmethod
			def 							UpperCamelCase				(     cls:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TOKEN
									HfFolder.save_token(UpperCAmelCase_					)







			@classmethod
			def 							UpperCamelCase				(     cls:  Union[str, Any]					):







									'''simple docstring'''




									try:
															delete_repo(token=cls._token				,     repo_id="""test-model-flax"""					)
									except HTTPError:
															pass

									try:
															delete_repo(token=cls._token				,     repo_id="""valid_org/test-model-flax-org"""					)
									except HTTPError:
															pass







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=99				,     hidden_size=32				,     num_hidden_layers=5				,     num_attention_heads=4				,     intermediate_size=37					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)
									model.push_to_hub("""test-model-flax"""				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

									# Reset repo
									delete_repo(token=self._token				,     repo_id="""test-model-flax"""					)

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(UpperCAmelCase_				,     repo_id="""test-model-flax"""				,     push_to_hub=UpperCAmelCase_				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=99				,     hidden_size=32				,     num_hidden_layers=5				,     num_attention_heads=4				,     intermediate_size=37					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)
									model.push_to_hub("""valid_org/test-model-flax-org"""				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

									# Reset repo
									delete_repo(token=self._token				,     repo_id="""valid_org/test-model-flax-org"""					)

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(
															    UpperCAmelCase_				,     repo_id="""valid_org/test-model-flax-org"""				,     push_to_hub=UpperCAmelCase_				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  True
						_SCREAMING_SNAKE_CASE     				=  flatten_dict(modela.params		)
						_SCREAMING_SNAKE_CASE     				=  flatten_dict(modela.params		)
						for key in flat_params_a.keys():
												if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]		)		) > 1e-4:
																		_SCREAMING_SNAKE_CASE     				=  False

						return models_are_equal




@require_flax
class 				__UpperCAmelCase		(unittest.TestCase   ):







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """bert"""
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(UpperCAmelCase_				,     UpperCAmelCase_					)					)

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """bert"""
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(UpperCAmelCase_				,     UpperCAmelCase_					)				,     max_shard_size="""10KB"""					)

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-subfolder"""

									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-sharded-subfolder"""
									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)

