
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			):




				def 	__a		(   self	)							-> Tuple:
							with tempfile.TemporaryDirectory() as tmpdirname:
										# pipeline has Flax weights
										lowerCAmelCase_         =      FlaxDiffusionPipeline.from_pretrained(
										    "hf-internal-testing/tiny-stable-diffusion-pipe"							,							safety_checker=_UpperCamelCase							,							cache_dir=_UpperCamelCase	)

										lowerCAmelCase_         =      [t[-1] for t in os.walk(os.path.join(_UpperCamelCase							,							os.listdir(_UpperCamelCase	)[0]							,							"snapshots"	)	)]
										lowerCAmelCase_         =      [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			):




				def 	__a		(   self	)							-> Any:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "hf-internal-testing/tiny-stable-diffusion-pipe"							,							safety_checker=_UpperCamelCase	)

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

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

							lowerCAmelCase_         =      jax.device_count()
							lowerCAmelCase_         =      num_samples * [prompt]
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)

							# shard inputs and rng
							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      jax.random.split(_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)

							lowerCAmelCase_         =      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.1514745	) < 1e-3
										assert np.abs(np.abs(_UpperCamelCase							,							dtype=np.floataa	).sum() - 49947.875	) < 5e-1

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




				def 	__a		(   self	)							-> Tuple:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"							,							revision="flax"							,							safety_checker=_UpperCamelCase	)

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

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

							lowerCAmelCase_         =      jax.device_count()
							lowerCAmelCase_         =      num_samples * [prompt]
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)

							# shard inputs and rng
							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      jax.random.split(_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)

							lowerCAmelCase_         =      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.05652401)	) < 1e-3
										assert np.abs((np.abs(_UpperCamelCase							,							dtype=np.floataa	).sum() - 2383808.2)	) < 5e-1




				def 	__a		(   self	)							-> Tuple:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"							,							revision="bf16"							,							dtype=jnp.bfloataa							,							safety_checker=_UpperCamelCase	)

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

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

							lowerCAmelCase_         =      jax.device_count()
							lowerCAmelCase_         =      num_samples * [prompt]
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)

							# shard inputs and rng
							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      jax.random.split(_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)

							lowerCAmelCase_         =      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.04003906)	) < 1e-3
										assert np.abs((np.abs(_UpperCamelCase							,							dtype=np.floataa	).sum() - 2373516.75)	) < 5e-1




				def 	__a		(   self	)							-> List[Any]:
							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"							,							revision="bf16"							,							dtype=jnp.bfloataa	)

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

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

							lowerCAmelCase_         =      jax.device_count()
							lowerCAmelCase_         =      num_samples * [prompt]
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)

							# shard inputs and rng
							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      jax.random.split(_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)

							lowerCAmelCase_         =      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.04003906)	) < 1e-3
										assert np.abs((np.abs(_UpperCamelCase							,							dtype=np.floataa	).sum() - 2373516.75)	) < 5e-1




				def 	__a		(   self	)							-> Any:
							lowerCAmelCase_         =      FlaxDDIMScheduler(
							    beta_start=0.00085							,							beta_end=0.012							,							beta_schedule="scaled_linear"							,							set_alpha_to_one=_UpperCamelCase							,							steps_offset=1							,							)

							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"							,							revision="bf16"							,							dtype=jnp.bfloataa							,							scheduler=_UpperCamelCase							,							safety_checker=_UpperCamelCase							,							)
							lowerCAmelCase_         =      scheduler.create_state()

							lowerCAmelCase_         =      scheduler_state

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

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

							lowerCAmelCase_         =      jax.device_count()
							lowerCAmelCase_         =      num_samples * [prompt]
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)

							# shard inputs and rng
							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      jax.random.split(_UpperCamelCase							,							_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)

							lowerCAmelCase_         =      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.045043945)	) < 1e-3
										assert np.abs((np.abs(_UpperCamelCase							,							dtype=np.floataa	).sum() - 2347693.5)	) < 5e-1




				def 	__a		(   self	)							-> Tuple:
							lowerCAmelCase_         =      (
							    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
							    " field, close up, split lighting, cinematic"
							)

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

							lowerCAmelCase_      ,    lowerCAmelCase_         =      FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"							,							revision="bf16"							,							dtype=jnp.bfloataa							,							safety_checker=_UpperCamelCase							,							)

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

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

							lowerCAmelCase_         =      replicate(_UpperCamelCase	)
							lowerCAmelCase_         =      pipeline.prepare_inputs(_UpperCamelCase	)
							lowerCAmelCase_         =      shard(_UpperCamelCase	)
							lowerCAmelCase_         =      pipeline(_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							jit=_UpperCamelCase	).images
							assert images_eff.shape == (num_samples, 1, 512, 512, 3)
							lowerCAmelCase_         =      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





from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional

from packaging import version


if TYPE_CHECKING:
			from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


_A  =	logging.get_logger(__name__)

_A  =	{
    "bigscience/bloom": "https://huggingface.co/bigscience/bloom/resolve/main/config.json",
    "bigscience/bloom-560m": "https://huggingface.co/bigscience/bloom-560m/blob/main/config.json",
    "bigscience/bloom-1b1": "https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json",
    "bigscience/bloom-1b7": "https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json",
    "bigscience/bloom-3b": "https://huggingface.co/bigscience/bloom-3b/blob/main/config.json",
    "bigscience/bloom-7b1": "https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json",
}





class       _lowerCAmelCase   (      __a			):
				_lowercase     						='''bloom'''
				_lowercase     						=['''past_key_values''']
				_lowercase     						={
				    '''num_hidden_layers''': '''n_layer''',
				    '''num_attention_heads''': '''n_head''',
				}




				def __init__(   self							,							_UpperCamelCase=250_880							,							_UpperCamelCase=64							,							_UpperCamelCase=2							,							_UpperCamelCase=8							,							_UpperCamelCase=1e-5							,							_UpperCamelCase=0.02							,							_UpperCamelCase=True							,							_UpperCamelCase=1							,							_UpperCamelCase=2							,							_UpperCamelCase=False							,							_UpperCamelCase=0.0							,							_UpperCamelCase=0.0							,							_UpperCamelCase=1							,							_UpperCamelCase=False							,							**_UpperCamelCase							,							)							-> str:
							lowerCAmelCase_         =      vocab_size
							# Backward compatibility with n_embed kwarg
							lowerCAmelCase_         =      kwargs.pop("n_embed"							,							_UpperCamelCase	)
							lowerCAmelCase_         =      hidden_size if n_embed is None else n_embed
							lowerCAmelCase_         =      n_layer
							lowerCAmelCase_         =      n_head
							lowerCAmelCase_         =      layer_norm_epsilon
							lowerCAmelCase_         =      initializer_range
							lowerCAmelCase_         =      use_cache
							lowerCAmelCase_         =      pretraining_tp
							lowerCAmelCase_         =      apply_residual_connection_post_layernorm
							lowerCAmelCase_         =      hidden_dropout
							lowerCAmelCase_         =      attention_dropout

							lowerCAmelCase_         =      bos_token_id
							lowerCAmelCase_         =      eos_token_id
							lowerCAmelCase_         =      slow_but_exact

							super().__init__(bos_token_id=_UpperCamelCase							,							eos_token_id=_UpperCamelCase							,							**_UpperCamelCase	)





class       _lowerCAmelCase   (      __a			):
				_lowercase     						=version.parse('''1.12'''			)




				def __init__(   self							,							_UpperCamelCase							,							_UpperCamelCase = "default"							,							_UpperCamelCase = None							,							_UpperCamelCase = False							,							)							-> int:
							super().__init__(_UpperCamelCase							,							task=_UpperCamelCase							,							patching_specs=_UpperCamelCase							,							use_past=_UpperCamelCase	)
							if not getattr(self._config							,							"pad_token_id"							,							_UpperCamelCase	):
										# TODO: how to do that better?
										lowerCAmelCase_         =      0




				@property
				def 	__a		(   self	)							-> Mapping[str, Mapping[int, str]]:
							lowerCAmelCase_         =      OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}	)
							if self.use_past:
										# BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
										self.fill_with_past_key_values_(_UpperCamelCase							,							direction="inputs"							,							inverted_values_shape=_UpperCamelCase	)
										lowerCAmelCase_         =      {0: "batch", 1: "past_sequence + sequence"}
							else:
										lowerCAmelCase_         =      {0: "batch", 1: "sequence"}

							return common_inputs




				@property
				def 	__a		(   self	)							-> int:
							return self._config.n_layer




				@property
				def 	__a		(   self	)							-> int:
							return self._config.n_head




				@property
				def 	__a		(   self	)							-> float:
							return 1e-3




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase = -1							,							_UpperCamelCase = -1							,							_UpperCamelCase = False							,							_UpperCamelCase = None							,							)							-> Mapping[str, Any]:
							lowerCAmelCase_         =      super(_UpperCamelCase							,							self	).generate_dummy_inputs(
							    _UpperCamelCase							,							batch_size=_UpperCamelCase							,							seq_length=_UpperCamelCase							,							is_pair=_UpperCamelCase							,							framework=_UpperCamelCase	)

							# We need to order the input in the way they appears in the forward()
							lowerCAmelCase_         =      OrderedDict({"input_ids": common_inputs["input_ids"]}	)

							# Need to add the past_keys
							if self.use_past:
										if not is_torch_available():
													raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."	)
										else:
													import torch

													lowerCAmelCase_      ,    lowerCAmelCase_         =      common_inputs["input_ids"].shape
													# Not using the same length for past_key_values
													lowerCAmelCase_         =      seqlen + 2
													lowerCAmelCase_         =      self._config.hidden_size // self.num_attention_heads
													lowerCAmelCase_         =      (
													    batch * self.num_attention_heads,
													    head_dim,
													    past_key_values_length,
													)
													lowerCAmelCase_         =      (
													    batch * self.num_attention_heads,
													    past_key_values_length,
													    head_dim,
													)
													lowerCAmelCase_         =      [
													    (torch.zeros(_UpperCamelCase	), torch.zeros(_UpperCamelCase	)) for _ in range(self.num_layers	)
													]

							lowerCAmelCase_         =      common_inputs["attention_mask"]
							if self.use_past:
										lowerCAmelCase_         =      ordered_inputs["attention_mask"].dtype
										lowerCAmelCase_         =      torch.cat(
										    [ordered_inputs["attention_mask"], torch.ones(_UpperCamelCase							,							_UpperCamelCase							,							dtype=_UpperCamelCase	)]							,							dim=1	)

							return ordered_inputs




				@property
				def 	__a		(   self	)							-> int:
							return 13



