




"""simple docstring"""






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





						def       __UpperCamelCase     (							self      :	Tuple)						->       Optional[int]:




											"""simple docstring"""







											with tempfile.TemporaryDirectory() as tmpdirname:
																# pipeline has Flax weights
																_lowerCAmelCase:Optional[Any]									=  FlaxDiffusionPipeline.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-pipe'''     ,safety_checker=a__     ,cache_dir=a__)

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





						def       __UpperCamelCase     (							self      :	str)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:Optional[Any]									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''hf-internal-testing/tiny-stable-diffusion-pipe'''     ,safety_checker=a__)

											_lowerCAmelCase:Optional[int]									=  (
											    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
											    ''' field, close up, split lighting, cinematic'''
											)

											_lowerCAmelCase:Union[str, Any]									=  jax.random.PRNGKey(0)
											_lowerCAmelCase:Dict									=  4

											_lowerCAmelCase:str									=  jax.device_count()
											_lowerCAmelCase:str									=  num_samples * [prompt]
											_lowerCAmelCase:Tuple									=  pipeline.prepare_inputs(a__)

											# shard inputs and rng
											_lowerCAmelCase:int									=  replicate(a__)
											_lowerCAmelCase:Dict									=  jax.random.split(a__     ,a__)
											_lowerCAmelCase:Dict									=  shard(a__)

											_lowerCAmelCase:Union[str, Any]									=  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.1514745) < 1E-3
																assert np.abs(np.abs(a__     ,dtype=np.floataa).sum() - 49947.875) < 5E-1

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





						def       __UpperCamelCase     (							self      :	Any)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:Any									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''     ,revision='''flax'''     ,safety_checker=a__)

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

											_lowerCAmelCase:int									=  jax.random.PRNGKey(0)
											_lowerCAmelCase:Optional[int]									=  50

											_lowerCAmelCase:str									=  jax.device_count()
											_lowerCAmelCase:str									=  num_samples * [prompt]
											_lowerCAmelCase:Union[str, Any]									=  pipeline.prepare_inputs(a__)

											# shard inputs and rng
											_lowerCAmelCase:Tuple									=  replicate(a__)
											_lowerCAmelCase:int									=  jax.random.split(a__     ,a__)
											_lowerCAmelCase:Union[str, Any]									=  shard(a__)

											_lowerCAmelCase:List[str]									=  pipeline(a__     ,a__     ,a__     ,a__     ,jit=a__).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(a__     ,dtype=np.floataa).sum() - 2383808.2)) < 5E-1





						def       __UpperCamelCase     (							self      :	Union[str, Any])						->       Union[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''     ,revision='''bf16'''     ,dtype=jnp.bfloataa     ,safety_checker=a__)

											_lowerCAmelCase:List[Any]									=  (
											    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
											    ''' field, close up, split lighting, cinematic'''
											)

											_lowerCAmelCase:List[Any]									=  jax.random.PRNGKey(0)
											_lowerCAmelCase:int									=  50

											_lowerCAmelCase:int									=  jax.device_count()
											_lowerCAmelCase:List[Any]									=  num_samples * [prompt]
											_lowerCAmelCase:str									=  pipeline.prepare_inputs(a__)

											# shard inputs and rng
											_lowerCAmelCase:List[Any]									=  replicate(a__)
											_lowerCAmelCase:Any									=  jax.random.split(a__     ,a__)
											_lowerCAmelCase:str									=  shard(a__)

											_lowerCAmelCase:int									=  pipeline(a__     ,a__     ,a__     ,a__     ,jit=a__).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(a__     ,dtype=np.floataa).sum() - 2373516.75)) < 5E-1





						def       __UpperCamelCase     (							self      :	Dict)						->       List[Any]:




											"""simple docstring"""







											_lowerCAmelCase:List[str]									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''     ,revision='''bf16'''     ,dtype=jnp.bfloataa)

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

											_lowerCAmelCase:int									=  jax.random.PRNGKey(0)
											_lowerCAmelCase:Dict									=  50

											_lowerCAmelCase:int									=  jax.device_count()
											_lowerCAmelCase:Optional[Any]									=  num_samples * [prompt]
											_lowerCAmelCase:List[str]									=  pipeline.prepare_inputs(a__)

											# shard inputs and rng
											_lowerCAmelCase:List[Any]									=  replicate(a__)
											_lowerCAmelCase:Optional[int]									=  jax.random.split(a__     ,a__)
											_lowerCAmelCase:List[str]									=  shard(a__)

											_lowerCAmelCase:Tuple									=  pipeline(a__     ,a__     ,a__     ,a__     ,jit=a__).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(a__     ,dtype=np.floataa).sum() - 2373516.75)) < 5E-1





						def       __UpperCamelCase     (							self      :	List[str])						->       int:




											"""simple docstring"""







											_lowerCAmelCase:str									=  FlaxDDIMScheduler(
											    beta_start=0.00085     ,beta_end=0.012     ,beta_schedule='''scaled_linear'''     ,set_alpha_to_one=a__     ,steps_offset=1     ,)

											_lowerCAmelCase:List[Any]									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''     ,revision='''bf16'''     ,dtype=jnp.bfloataa     ,scheduler=a__     ,safety_checker=a__     ,)
											_lowerCAmelCase:List[Any]									=  scheduler.create_state()

											_lowerCAmelCase:Union[str, Any]									=  scheduler_state

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

											_lowerCAmelCase:Optional[Any]									=  jax.random.PRNGKey(0)
											_lowerCAmelCase:str									=  50

											_lowerCAmelCase:Tuple									=  jax.device_count()
											_lowerCAmelCase:str									=  num_samples * [prompt]
											_lowerCAmelCase:str									=  pipeline.prepare_inputs(a__)

											# shard inputs and rng
											_lowerCAmelCase:Tuple									=  replicate(a__)
											_lowerCAmelCase:List[str]									=  jax.random.split(a__     ,a__)
											_lowerCAmelCase:str									=  shard(a__)

											_lowerCAmelCase:Union[str, Any]									=  pipeline(a__     ,a__     ,a__     ,a__     ,jit=a__).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(a__     ,dtype=np.floataa).sum() - 2347693.5)) < 5E-1







						def       __UpperCamelCase     (							self      :	List[Any])						->       Optional[Any]:




											"""simple docstring"""







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

											_lowerCAmelCase:Any									=  jax.device_count()
											_lowerCAmelCase:int									=  num_samples * [prompt]
											_lowerCAmelCase:List[Any]									=  jax.random.split(jax.random.PRNGKey(0)     ,a__)

											_lowerCAmelCase:Any									=  FlaxStableDiffusionPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''     ,revision='''bf16'''     ,dtype=jnp.bfloataa     ,safety_checker=a__     ,)

											_lowerCAmelCase:Optional[int]									=  replicate(a__)
											_lowerCAmelCase:Dict									=  pipeline.prepare_inputs(a__)
											_lowerCAmelCase:Dict									=  shard(a__)
											_lowerCAmelCase:Tuple									=  pipeline(a__     ,a__     ,a__     ,jit=a__).images
											assert images.shape == (num_samples, 1, 512, 512, 3)
											_lowerCAmelCase:Optional[Any]									=  images[2, 0, 256, 10:17, 1]

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

											_lowerCAmelCase:Optional[Any]									=  replicate(a__)
											_lowerCAmelCase:Optional[Any]									=  pipeline.prepare_inputs(a__)
											_lowerCAmelCase:int									=  shard(a__)
											_lowerCAmelCase:List[Any]									=  pipeline(a__     ,a__     ,a__     ,jit=a__).images
											assert images_eff.shape == (num_samples, 1, 512, 512, 3)
											_lowerCAmelCase:Tuple									=  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












"""simple docstring"""






import warnings

from ...utils import logging
from .image_processing_poolformer import PoolFormerImageProcessor


UpperCamelCase__						=			logging.get_logger(__name__)




class   a__			(       UpperCamelCase_  ):





						def __init__(							self      :	Dict     ,*a__      :	List[str]     ,**a__      :	str)						->       None:




											"""simple docstring"""







											warnings.warn(
											    '''The class PoolFormerFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'''
											    ''' Please use PoolFormerImageProcessor instead.'''     ,a__     ,)
											super().__init__(*a__     ,**a__)






