






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 		UpperCamelCase__    (unittest.TestCase			):




		'''simple docstring'''





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

														lowerCamelCase__         =  [t[-1] for t in os.walk(os.path.join(_lowerCAmelCase   ,os.listdir(_lowerCAmelCase       )[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 		UpperCamelCase__    (unittest.TestCase			):




		'''simple docstring'''





		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    """hf-internal-testing/tiny-stable-diffusion-pipe"""   ,safety_checker=_lowerCAmelCase       )

								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(_lowerCAmelCase       )

								# shard inputs and rng
								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).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.151_4745       ) < 1E-3
														assert np.abs(np.abs(_lowerCAmelCase   ,dtype=np.floataa       ).sum() - 4_99_47.8_75       ) < 5E-1

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






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""   ,revision="""flax"""   ,safety_checker=_lowerCAmelCase       )

								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(_lowerCAmelCase       )

								# shard inputs and rng
								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images

								assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,dtype=np.floataa       ).sum() - 0.0565_2401)       ) < 1E-3
														assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa       ).sum() - 2_38_38_08.2)       ) < 5E-1






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""   ,revision="""bf16"""   ,dtype=jnp.bfloataa   ,safety_checker=_lowerCAmelCase       )

								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(_lowerCAmelCase       )

								# shard inputs and rng
								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images

								assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,dtype=np.floataa       ).sum() - 0.0400_3906)       ) < 1E-3
														assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5E-1






		def 			UpperCamelCase_					(      self       ):
								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(_lowerCAmelCase       )

								# shard inputs and rng
								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images

								assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,dtype=np.floataa       ).sum() - 0.0400_3906)       ) < 1E-3
														assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5E-1






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  FlaxDDIMScheduler(
								    beta_start=0.0_0085   ,beta_end=0.012   ,beta_schedule="""scaled_linear"""   ,set_alpha_to_one=_lowerCAmelCase   ,steps_offset=1   ,)

								lowerCamelCase__       ,			lowerCamelCase__         =  FlaxStableDiffusionPipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""   ,revision="""bf16"""   ,dtype=jnp.bfloataa   ,scheduler=_lowerCAmelCase   ,safety_checker=_lowerCAmelCase   ,)
								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(_lowerCAmelCase       )

								# shard inputs and rng
								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )

								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images

								assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
								if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]   ,dtype=np.floataa       ).sum() - 0.0_4504_3945)       ) < 1E-3
														assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa       ).sum() - 2_34_76_93.5)       ) < 5E-1






		def 			UpperCamelCase_					(      self       ):
								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       )   ,_lowerCAmelCase       )

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

								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  pipeline.prepare_inputs(_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )
								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images
								assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
								lowerCamelCase__         =  images[2, 0, 2_56, 10:17, 1]

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

								lowerCamelCase__         =  replicate(_lowerCAmelCase       )
								lowerCamelCase__         =  pipeline.prepare_inputs(_lowerCAmelCase       )
								lowerCamelCase__         =  shard(_lowerCAmelCase       )
								lowerCamelCase__         =  pipeline(_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,jit=_lowerCAmelCase       ).images
								assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
								lowerCamelCase__         =  images[2, 0, 2_56, 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









'''simple docstring'''







import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing




class 		UpperCamelCase__    (a			):




		'''simple docstring'''





		def __init__(      self   ,_lowerCAmelCase = "▁"   ,_lowerCAmelCase = True   ,_lowerCAmelCase = "<unk>"   ,_lowerCAmelCase = "</s>"   ,_lowerCAmelCase = "<pad>"   ,):
								lowerCamelCase__         =  {
								    """pad""": {"""id""": 0, """token""": pad_token},
								    """eos""": {"""id""": 1, """token""": eos_token},
								    """unk""": {"""id""": 2, """token""": unk_token},
								}

								lowerCamelCase__         =  [None] * len(self.special_tokens       )
								for token_dict in self.special_tokens.values():
														lowerCamelCase__         =  token_dict["""token"""]

								lowerCamelCase__         =  Tokenizer(Unigram()       )

								lowerCamelCase__         =  normalizers.Sequence(
								    [
								        normalizers.Nmt(),
								        normalizers.NFKC(),
								        normalizers.Replace(Regex(""" {2,}"""       )   ,""" """       ),
								        normalizers.Lowercase(),
								    ]       )
								lowerCamelCase__         =  pre_tokenizers.Sequence(
								    [
								        pre_tokenizers.Metaspace(replacement=_lowerCAmelCase   ,add_prefix_space=_lowerCAmelCase       ),
								        pre_tokenizers.Digits(individual_digits=_lowerCAmelCase       ),
								        pre_tokenizers.Punctuation(),
								    ]       )
								lowerCamelCase__         =  decoders.Metaspace(replacement=_lowerCAmelCase   ,add_prefix_space=_lowerCAmelCase       )

								lowerCamelCase__         =  TemplateProcessing(
								    single=F'''$A {self.special_tokens["eos"]["token"]}'''   ,special_tokens=[(self.special_tokens["""eos"""]["""token"""], self.special_tokens["""eos"""]["""id"""])]   ,)

								lowerCamelCase__         =  {
								    """model""": """SentencePieceUnigram""",
								    """replacement""": replacement,
								    """add_prefix_space""": add_prefix_space,
								}

								super().__init__(_lowerCAmelCase   ,_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase = 80_00   ,_lowerCAmelCase = True   ,):
								lowerCamelCase__         =  trainers.UnigramTrainer(
								    vocab_size=_lowerCAmelCase   ,special_tokens=self.special_tokens_list   ,show_progress=_lowerCAmelCase   ,)

								if isinstance(_lowerCAmelCase   ,_lowerCAmelCase       ):
														lowerCamelCase__         =  [files]
								self._tokenizer.train(_lowerCAmelCase   ,trainer=_lowerCAmelCase       )

								self.add_unk_id()






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase = 80_00   ,_lowerCAmelCase = True   ,):
								lowerCamelCase__         =  trainers.UnigramTrainer(
								    vocab_size=_lowerCAmelCase   ,special_tokens=self.special_tokens_list   ,show_progress=_lowerCAmelCase   ,)

								self._tokenizer.train_from_iterator(_lowerCAmelCase   ,trainer=_lowerCAmelCase       )

								self.add_unk_id()






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  json.loads(self._tokenizer.to_str()       )

								lowerCamelCase__         =  self.special_tokens["""unk"""]["""id"""]

								lowerCamelCase__         =  Tokenizer.from_str(json.dumps(_lowerCAmelCase       )       )

