
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   :					Optional[Any]					):

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

									__A       =					[t[-1] for t in os.walk(os.path.join(A				,os.listdir(A					)[0]				,"snapshots"					)					)]
									__A       =					[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   :					Union[str, Any]					):
					__A     ,				__A       =					FlaxStableDiffusionPipeline.from_pretrained(
					    "hf-internal-testing/tiny-stable-diffusion-pipe"				,safety_checker=A					)

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

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

					__A       =					jax.device_count()
					__A       =					num_samples * [prompt]
					__A       =					pipeline.prepare_inputs(A					)

					# shard inputs and rng
					__A       =					replicate(A					)
					__A       =					jax.random.split(A				,A					)
					__A       =					shard(A					)

					__A       =					pipeline(A				,A				,A				,A				,jit=A					).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(A				,dtype=np.floataa					).sum() - 4_99_47.8_75					) < 5E-1

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






	def 							UpperCamelCase_    (			self   :					Optional[int]					):
					__A     ,				__A       =					FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"				,revision="flax"				,safety_checker=A					)

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

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

					__A       =					jax.device_count()
					__A       =					num_samples * [prompt]
					__A       =					pipeline.prepare_inputs(A					)

					# shard inputs and rng
					__A       =					replicate(A					)
					__A       =					jax.random.split(A				,A					)
					__A       =					shard(A					)

					__A       =					pipeline(A				,A				,A				,A				,jit=A					).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.05_65_24_01)					) < 1E-3
									assert np.abs((np.abs(A				,dtype=np.floataa					).sum() - 2_38_38_08.2)					) < 5E-1






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

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

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

					__A       =					jax.device_count()
					__A       =					num_samples * [prompt]
					__A       =					pipeline.prepare_inputs(A					)

					# shard inputs and rng
					__A       =					replicate(A					)
					__A       =					jax.random.split(A				,A					)
					__A       =					shard(A					)

					__A       =					pipeline(A				,A				,A				,A				,jit=A					).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.04_00_39_06)					) < 1E-3
									assert np.abs((np.abs(A				,dtype=np.floataa					).sum() - 2_37_35_16.75)					) < 5E-1






	def 							UpperCamelCase_    (			self   :					List[str]					):
					__A     ,				__A       =					FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"				,revision="bf16"				,dtype=jnp.bfloataa					)

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

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

					__A       =					jax.device_count()
					__A       =					num_samples * [prompt]
					__A       =					pipeline.prepare_inputs(A					)

					# shard inputs and rng
					__A       =					replicate(A					)
					__A       =					jax.random.split(A				,A					)
					__A       =					shard(A					)

					__A       =					pipeline(A				,A				,A				,A				,jit=A					).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.04_00_39_06)					) < 1E-3
									assert np.abs((np.abs(A				,dtype=np.floataa					).sum() - 2_37_35_16.75)					) < 5E-1






	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					__A       =					FlaxDDIMScheduler(
					    beta_start=0.0_00_85				,beta_end=0.0_12				,beta_schedule="scaled_linear"				,set_alpha_to_one=A				,steps_offset=1				,)

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

					__A       =					scheduler_state

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

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

					__A       =					jax.device_count()
					__A       =					num_samples * [prompt]
					__A       =					pipeline.prepare_inputs(A					)

					# shard inputs and rng
					__A       =					replicate(A					)
					__A       =					jax.random.split(A				,A					)
					__A       =					shard(A					)

					__A       =					pipeline(A				,A				,A				,A				,jit=A					).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_45_04_39_45)					) < 1E-3
									assert np.abs((np.abs(A				,dtype=np.floataa					).sum() - 2_34_76_93.5)					) < 5E-1







	def 							UpperCamelCase_    (			self   :					Dict					):
					__A       =					(
					    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
					    " field, close up, split lighting, cinematic"
					)

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

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

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

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

					__A       =					replicate(A					)
					__A       =					pipeline.prepare_inputs(A					)
					__A       =					shard(A					)
					__A       =					pipeline(A				,A				,A				,jit=A					).images
					assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
					__A       =					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





import json
import os
import unittest

from transformers.models.ctrl.tokenization_ctrl import VOCAB_FILES_NAMES, CTRLTokenizer

from ...test_tokenization_common import TokenizerTesterMixin



class    UpperCAmelCase__(				lowerCamelCase  ,				unittest.TestCase						):

   '''simple docstring'''




   A     :	int												=     CTRLTokenizer
   A     :	Tuple												=     False
   A     :	Any												=     False


   def       UpperCAmelCase		(    self	:      int)						->					Any:




         """simple docstring"""

         super().setUp()

         # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
         lowercase__  =	['adapt', 're@@', 'a@@', 'apt', 'c@@', 't', '<unk>']
         lowercase__  =	dict(zip(lowerCAmelCase      ,					range(len(lowerCAmelCase))))
         lowercase__  =	['#version: 0.2', 'a p', 'ap t</w>', 'r e', 'a d', 'ad apt</w>', '']
         lowercase__  =	{'unk_token': '<unk>'}

         lowercase__  =	os.path.join(self.tmpdirname      ,					VOCAB_FILES_NAMES['vocab_file'])
         lowercase__  =	os.path.join(self.tmpdirname      ,					VOCAB_FILES_NAMES['merges_file'])
         with open(self.vocab_file      ,					'w'      ,					encoding='utf-8') as fp:
               fp.write(json.dumps(lowerCAmelCase) + '\n')
         with open(self.merges_file      ,					'w'      ,					encoding='utf-8') as fp:
               fp.write('\n'.join(lowerCAmelCase))


   def       UpperCAmelCase		(    self	:      int      ,					**lowerCAmelCase	:      str)						->					Dict:




         """simple docstring"""

         kwargs.update(self.special_tokens_map)
         return CTRLTokenizer.from_pretrained(self.tmpdirname      ,					**lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      List[str]      ,					lowerCAmelCase	:      Optional[int])						->					Optional[Any]:




         """simple docstring"""

         lowercase__  =	'adapt react readapt apt'
         lowercase__  =	'adapt react readapt apt'
         return input_text, output_text



   def       UpperCAmelCase		(    self	:      Optional[int])						->					str:




         """simple docstring"""

         lowercase__  =	CTRLTokenizer(self.vocab_file      ,					self.merges_file      ,					**self.special_tokens_map)
         lowercase__  =	'adapt react readapt apt'
         lowercase__  =	'adapt re@@ a@@ c@@ t re@@ adapt apt'.split()
         lowercase__  =	tokenizer.tokenize(lowerCAmelCase)
         self.assertListEqual(lowerCAmelCase      ,					lowerCAmelCase)

         lowercase__  =	tokens + [tokenizer.unk_token]

         lowercase__  =	[0, 1, 2, 4, 5, 1, 0, 3, 6]
         self.assertListEqual(tokenizer.convert_tokens_to_ids(lowerCAmelCase)      ,					lowerCAmelCase)
