






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  ):
								"""simple docstring"""
								with tempfile.TemporaryDirectory() as tmpdirname:
												# pipeline has Flax weights
												__magic_name__		:Optional[Any]   =   FlaxDiffusionPipeline.from_pretrained(
												    '''hf-internal-testing/tiny-stable-diffusion-pipe'''							,       safety_checker=__lowerCAmelCase							,       cache_dir=__lowerCAmelCase  )

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







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Tuple   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''hf-internal-testing/tiny-stable-diffusion-pipe'''							,       safety_checker=__lowerCAmelCase  )

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

								__magic_name__		:int   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   4

								__magic_name__		:str   =   jax.device_count()
								__magic_name__		:List[str]   =   num_samples * [prompt]
								__magic_name__		:str   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:List[Any]   =   replicate(__lowerCAmelCase  )
								__magic_name__		:int   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:int   =   shard(__lowerCAmelCase  )

								__magic_name__		:Any   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images

								assert images.shape == (num_samples, 1, 6_4, 6_4, 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(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 49947.875  ) < 5E-1

								__magic_name__		:List[Any]   =   pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]  )  )  )
								assert len(__lowerCAmelCase  ) == num_samples







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:str   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''flax'''							,       safety_checker=__lowerCAmelCase  )

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

								__magic_name__		:List[str]   =   jax.random.PRNGKey(0  )
								__magic_name__		:str   =   5_0

								__magic_name__		:Tuple   =   jax.device_count()
								__magic_name__		:Any   =   num_samples * [prompt]
								__magic_name__		:Dict   =   pipeline.prepare_inputs(__lowerCAmelCase  )

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

								__magic_name__		:Optional[int]   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2383808.2)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Union[str, Any]   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa							,       safety_checker=__lowerCAmelCase  )

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

								__magic_name__		:Tuple   =   jax.random.PRNGKey(0  )
								__magic_name__		:str   =   5_0

								__magic_name__		:Optional[Any]   =   jax.device_count()
								__magic_name__		:List[Any]   =   num_samples * [prompt]
								__magic_name__		:List[Any]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:Dict   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Tuple   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:List[Any]   =   shard(__lowerCAmelCase  )

								__magic_name__		:List[str]   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2373516.75)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Dict   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa  )

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

								__magic_name__		:Any   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   5_0

								__magic_name__		:Optional[Any]   =   jax.device_count()
								__magic_name__		:Optional[Any]   =   num_samples * [prompt]
								__magic_name__		:Optional[Any]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

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

								__magic_name__		:Tuple   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2373516.75)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   FlaxDDIMScheduler(
								    beta_start=0.00085							,       beta_end=0.012							,       beta_schedule='''scaled_linear'''							,       set_alpha_to_one=__lowerCAmelCase							,       steps_offset=1							,       )

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

								__magic_name__		:Optional[int]   =   scheduler_state

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

								__magic_name__		:Tuple   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   5_0

								__magic_name__		:Dict   =   jax.device_count()
								__magic_name__		:List[Any]   =   num_samples * [prompt]
								__magic_name__		:Optional[int]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

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

								__magic_name__		:Tuple   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images

								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2347693.5)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		: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'''
								)

								__magic_name__		:str   =   jax.device_count()
								__magic_name__		:str   =   num_samples * [prompt]
								__magic_name__		:Dict   =   jax.random.split(jax.random.PRNGKey(0  )							,       __lowerCAmelCase  )

								__magic_name__  ,					__magic_name__		:Union[str, Any]   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa							,       safety_checker=__lowerCAmelCase							,       )

								__magic_name__		:Any   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Optional[int]   =   pipeline.prepare_inputs(__lowerCAmelCase  )
								__magic_name__		:Any   =   shard(__lowerCAmelCase  )
								__magic_name__		:int   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images
								assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								__magic_name__		:List[Any]   =   images[2, 0, 2_5_6, 1_0:1_7, 1]

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

								__magic_name__		:List[Any]   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Any   =   pipeline.prepare_inputs(__lowerCAmelCase  )
								__magic_name__		:Any   =   shard(__lowerCAmelCase  )
								__magic_name__		:Any   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images
								assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								__magic_name__		:List[Any]   =   images[2, 0, 2_5_6, 1_0:1_7, 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 pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				assert isinstance(snake_case,	snake_case			)
				assert dataset.num_rows == 4
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_1", "col_2", "col_3"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize('''keep_in_memory''',	[False, True]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Tuple   =   tmp_path / '''cache'''
				__magic_name__		:int   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								__magic_name__		:Optional[Any]   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case,	keep_in_memory=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:List[str]   =   tmp_path / '''cache'''
				__magic_name__		:int   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:Tuple   =   features.copy() if features else default_expected_features
				__magic_name__		:Union[str, Any]   =   (
				    Features({feature: Value(snake_case			) for feature, dtype in features.items()}			) if features is not None else None
				)
				__magic_name__		:int   =   ParquetDatasetReader(snake_case,	features=snake_case,	cache_dir=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''			), '''train''', '''test''']			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:str   =   tmp_path / '''cache'''
				__magic_name__		:List[Any]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:int   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case,	split=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)
				assert dataset.split == split if split else "train"

@pytest.mark.parametrize('''path_type''',	[str, list]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				if issubclass(snake_case,	snake_case			):
								__magic_name__		:Union[str, Any]   =   parquet_path
				elif issubclass(snake_case,	snake_case			):
								__magic_name__		:Union[str, Any]   =   [parquet_path]
				__magic_name__		:Optional[int]   =   tmp_path / '''cache'''
				__magic_name__		:Optional[int]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:str   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case			).read()
				_check_parquet_dataset(snake_case,	snake_case			)

def   __lowercase       (			snake_case,	snake_case,	snake_case=("train",)			):







				"""simple docstring"""






				assert isinstance(snake_case,	snake_case			)
				for split in splits:
								__magic_name__		:Optional[Any]   =   dataset_dict[split]
								assert dataset.num_rows == 4
								assert dataset.num_columns == 3
								assert dataset.column_names == ["col_1", "col_2", "col_3"]
								for feature, expected_dtype in expected_features.items():
												assert dataset.features[feature].dtype == expected_dtype

@pytest.mark.parametrize('''keep_in_memory''',	[False, True]			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Any   =   tmp_path / '''cache'''
				__magic_name__		:Optional[int]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								__magic_name__		:Tuple   =   ParquetDatasetReader(
								    {'''train''': parquet_path},	cache_dir=snake_case,	keep_in_memory=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case			)

@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:Optional[Any]   =   tmp_path / '''cache'''
				__magic_name__		:Dict   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:int   =   features.copy() if features else default_expected_features
				__magic_name__		:List[Any]   =   (
				    Features({feature: Value(snake_case			) for feature, dtype in features.items()}			) if features is not None else None
				)
				__magic_name__		:Optional[int]   =   ParquetDatasetReader({'''train''': parquet_path},	features=snake_case,	cache_dir=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case			)

@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''			), '''train''', '''test''']			)
def   __lowercase       (			snake_case,	snake_case,	snake_case			):







				"""simple docstring"""






				if split:
								__magic_name__		:Dict   =   {split: parquet_path}
				else:
								__magic_name__		:Optional[int]   =   '''train'''
								__magic_name__		:Dict   =   {'''train''': parquet_path, '''test''': parquet_path}
				__magic_name__		:List[Any]   =   tmp_path / '''cache'''
				__magic_name__		:Optional[Any]   =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				__magic_name__		:Optional[Any]   =   ParquetDatasetReader(snake_case,	cache_dir=snake_case			).read()
				_check_parquet_datasetdict(snake_case,	snake_case,	splits=list(path.keys()			)			)
				assert all(dataset[split].split == split for split in path.keys()			)

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:str   =   ParquetDatasetWriter(snake_case,	tmp_path / '''foo.parquet'''			)
				assert writer.write() > 0
				__magic_name__		:List[Any]   =   pq.ParquetFile(tmp_path / '''foo.parquet'''			)
				__magic_name__		:List[Any]   =   pf.read()
				assert dataset.data.table == output_table

def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				__magic_name__		:List[str]   =   str(shared_datadir / '''test_image_rgb.jpg'''			)
				__magic_name__		:Tuple   =   {'''image''': [image_path]}
				__magic_name__		:List[Any]   =   Features({'''image''': Image()}			)
				__magic_name__		:Tuple   =   Dataset.from_dict(snake_case,	features=snake_case			)
				__magic_name__		:Union[str, Any]   =   ParquetDatasetWriter(snake_case,	tmp_path / '''foo.parquet'''			)
				assert writer.write() > 0

				__magic_name__		:List[str]   =   Dataset.from_parquet(str(tmp_path / '''foo.parquet'''			)			)
				assert dataset.features == reloaded_dataset.features

				__magic_name__		:List[str]   =   ParquetDatasetReader(str(tmp_path / '''foo.parquet'''			),	streaming=snake_case			).read()
				assert dataset.features == reloaded_iterable_dataset.features

@pytest.mark.parametrize(
    '''feature, expected''',	[
        (Features({'''foo''': Value('''int32'''			)}			), None),
        (Features({'''image''': Image(), '''foo''': Value('''int32'''			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({'''nested''': Sequence(Audio()			)}			), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ],	)
def   __lowercase       (			snake_case,	snake_case			):







				"""simple docstring"""






				assert get_writer_batch_size(snake_case			) == expected






