





import copy
from collections import OrderedDict
from typing import Dict, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


lowerCAmelCase_    =				logging.get_logger(__name__)

lowerCAmelCase_    =				{
    '''facebook/detr-resnet-50''': '''https://huggingface.co/facebook/detr-resnet-50/resolve/main/config.json''',
    # See all DETR models at https://huggingface.co/models?filter=detr
}



class 					__lowerCAmelCase       (			snake_case__     ):
			lowerCamelCase_     :     Any       			=      '''detr'''
			lowerCamelCase_     :     Any       			=      ['''past_key_values''']
			lowerCamelCase_     :     Optional[int]       			=      {
			    '''hidden_size''': '''d_model''',
			    '''num_attention_heads''': '''encoder_attention_heads''',
			}






			def __init__(self					,							__magic_name__=True					,							__magic_name__=None					,							__magic_name__=3					,							__magic_name__=100					,							__magic_name__=6					,							__magic_name__=2048					,							__magic_name__=8					,							__magic_name__=6					,							__magic_name__=2048					,							__magic_name__=8					,							__magic_name__=0.0					,							__magic_name__=0.0					,							__magic_name__=True					,							__magic_name__="relu"					,							__magic_name__=256					,							__magic_name__=0.1					,							__magic_name__=0.0					,							__magic_name__=0.0					,							__magic_name__=0.02					,							__magic_name__=1.0					,							__magic_name__=False					,							__magic_name__="sine"					,							__magic_name__="resnet50"					,							__magic_name__=True					,							__magic_name__=False					,							__magic_name__=1					,							__magic_name__=5					,							__magic_name__=2					,							__magic_name__=1					,							__magic_name__=1					,							__magic_name__=5					,							__magic_name__=2					,							__magic_name__=0.1					,							**__magic_name__					,							)     ->   Any:


							'''simple docstring'''



							if backbone_config is not None and use_timm_backbone:
											raise ValueError('''You can\'t specify both `backbone_config` and `use_timm_backbone`.'''    )

							if not use_timm_backbone:
											if backbone_config is None:
															logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''    )
															snake_case_						:    Tuple				=  CONFIG_MAPPING['''resnet'''](out_features=['''stage4''']    )
											elif isinstance(__magic_name__					,							__magic_name__    ):
															snake_case_						:    int				=  backbone_config.get('''model_type'''    )
															snake_case_						:    Dict				=  CONFIG_MAPPING[backbone_model_type]
															snake_case_						:    Any				=  config_class.from_dict(__magic_name__    )
											# set timm attributes to None
											snake_case_			,							snake_case_			,							snake_case_						:    Optional[Any]				=  None, None, None

							snake_case_						:    Tuple				=  use_timm_backbone
							snake_case_						:    List[Any]				=  backbone_config
							snake_case_						:    Tuple				=  num_channels
							snake_case_						:    Tuple				=  num_queries
							snake_case_						:    int				=  d_model
							snake_case_						:    str				=  encoder_ffn_dim
							snake_case_						:    Tuple				=  encoder_layers
							snake_case_						:    Union[str, Any]				=  encoder_attention_heads
							snake_case_						:    Tuple				=  decoder_ffn_dim
							snake_case_						:    Tuple				=  decoder_layers
							snake_case_						:    Union[str, Any]				=  decoder_attention_heads
							snake_case_						:    Any				=  dropout
							snake_case_						:    Dict				=  attention_dropout
							snake_case_						:    Optional[Any]				=  activation_dropout
							snake_case_						:    int				=  activation_function
							snake_case_						:    Any				=  init_std
							snake_case_						:    List[Any]				=  init_xavier_std
							snake_case_						:    Optional[int]				=  encoder_layerdrop
							snake_case_						:    List[str]				=  decoder_layerdrop
							snake_case_						:    str				=  encoder_layers
							snake_case_						:    Union[str, Any]				=  auxiliary_loss
							snake_case_						:    Optional[Any]				=  position_embedding_type
							snake_case_						:    int				=  backbone
							snake_case_						:    List[str]				=  use_pretrained_backbone
							snake_case_						:    List[Any]				=  dilation
							# Hungarian matcher
							snake_case_						:    Optional[Any]				=  class_cost
							snake_case_						:    List[Any]				=  bbox_cost
							snake_case_						:    Any				=  giou_cost
							# Loss coefficients
							snake_case_						:    List[Any]				=  mask_loss_coefficient
							snake_case_						:    Union[str, Any]				=  dice_loss_coefficient
							snake_case_						:    List[str]				=  bbox_loss_coefficient
							snake_case_						:    str				=  giou_loss_coefficient
							snake_case_						:    Any				=  eos_coefficient
							super().__init__(is_encoder_decoder=__magic_name__					,							**__magic_name__    )






			@property
			def 			lowerCamelCase  (self    )     ->   Optional[int]:


							'''simple docstring'''



							return self.encoder_attention_heads






			@property
			def 			lowerCamelCase  (self    )     ->   Optional[int]:


							'''simple docstring'''



							return self.d_model






			@classmethod
			def 			lowerCamelCase  (cls					,							__magic_name__					,							**__magic_name__    )     ->   Tuple:


							'''simple docstring'''



							return cls(backbone_config=__magic_name__					,							**__magic_name__    )






			def 			lowerCamelCase  (self    )     ->   List[Any]:


							'''simple docstring'''



							snake_case_						:    Dict				=  copy.deepcopy(self.__dict__    )
							if output["backbone_config"] is not None:
											snake_case_						:    Any				=  self.backbone_config.to_dict()
							snake_case_						:    Any				=  self.__class__.model_type
							return output




class 					__lowerCAmelCase       (			snake_case__     ):
			lowerCamelCase_     :     Union[str, Any]       			=      version.parse('''1.11'''     )






			@property
			def 			lowerCamelCase  (self    )     ->   Optional[Any]:


							'''simple docstring'''



							return OrderedDict(
							    [
							        ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
							        ('''pixel_mask''', {0: '''batch'''}),
							    ]    )






			@property
			def 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							return 1e-5






			@property
			def 			lowerCamelCase  (self    )     ->   str:


							'''simple docstring'''



							return 12








import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


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

		from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline





@require_flax
class __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							with tempfile.TemporaryDirectory() as tmpdirname:
													# pipeline has Flax weights
													_UpperCAmelCase									=			FlaxDiffusionPipeline.from_pretrained(
													    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,    safety_checker=UpperCAmelCase       ,    cache_dir=UpperCAmelCase    )

													_UpperCAmelCase									=			[t[-1] for t in os.walk(os.path.join(UpperCAmelCase       ,    os.listdir(UpperCAmelCase    )[0]       ,    'snapshots'    )    )]
													_UpperCAmelCase									=			[item for sublist in all_root_files for item in sublist]

													# None of the downloaded files should be a PyTorch file even if we have some here:
													# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
													assert not any(f.endswith('.bin'    ) for f in files    )





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





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,    safety_checker=UpperCAmelCase    )

							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			4

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 64, 64, 3)
							if jax.device_count() == 8:
													assert np.abs(np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 4.1_51_47_45    ) < 1e-3
													assert np.abs(np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 4_99_47.8_75    ) < 5e-1

							_UpperCAmelCase									=			pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]    )    )    )
							assert len(UpperCAmelCase    ) == num_samples

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='flax'       ,    safety_checker=UpperCAmelCase    )

							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.05_65_24_01)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_38_38_08.2)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    safety_checker=UpperCAmelCase    )

							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa    )

							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.04_00_39_06)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_37_35_16.75)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			FlaxDDIMScheduler(
							    beta_start=0.0_00_85       ,    beta_end=0.0_12       ,    beta_schedule='scaled_linear'       ,    set_alpha_to_one=UpperCAmelCase       ,    steps_offset=1       ,    )

							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    scheduler=UpperCAmelCase       ,    safety_checker=UpperCAmelCase       ,    )
							_UpperCAmelCase									=			scheduler.create_state()

							_UpperCAmelCase									=			scheduler_state

							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.random.PRNGKey(0    )
							_UpperCAmelCase									=			50

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images

							assert images.shape == (num_samples, 1, 512, 512, 3)
							if jax.device_count() == 8:
													assert np.abs((np.abs(images[0, 0, :2, :2, -2:]       ,    dtype=np.floataa    ).sum() - 0.0_45_04_39_45)    ) < 1e-3
													assert np.abs((np.abs(UpperCAmelCase       ,    dtype=np.floataa    ).sum() - 2_34_76_93.5)    ) < 5e-1

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			(
							    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
							    ' field, close up, split lighting, cinematic'
							)

							_UpperCAmelCase									=			jax.device_count()
							_UpperCAmelCase									=			num_samples * [prompt]
							_UpperCAmelCase									=			jax.random.split(jax.random.PRNGKey(0    )       ,    UpperCAmelCase    )

							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    safety_checker=UpperCAmelCase       ,    )

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			images[2, 0, 256, 10:17, 1]

							# With memory efficient attention
							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    safety_checker=UpperCAmelCase       ,    use_memory_efficient_attention=UpperCAmelCase       ,    )

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							_UpperCAmelCase									=			images[2, 0, 256, 10:17, 1]

							# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
							# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
							assert abs(slice_eff - slice    ).max() < 1e-2
