

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







 """simple docstring"""







 def 					_lowercase (			self							:			Optional[int]	):
    with tempfile.TemporaryDirectory() as tmpdirname:
       # pipeline has Flax weights
       __lowercase								=						FlaxDiffusionPipeline.from_pretrained(
           "hf-internal-testing/tiny-stable-diffusion-pipe",					safety_checker=UpperCAmelCase__,					cache_dir=UpperCAmelCase__	)

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







 """simple docstring"""







 def 					_lowercase (			self							:			List[Any]	):
    __lowercase		,__lowercase								=						FlaxStableDiffusionPipeline.from_pretrained(
        "hf-internal-testing/tiny-stable-diffusion-pipe",					safety_checker=UpperCAmelCase__	)

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

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

    __lowercase								=						jax.device_count()
    __lowercase								=						num_samples * [prompt]
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)

    # shard inputs and rng
    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						jax.random.split(UpperCAmelCase__,					UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)

    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).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.1_514_745	) < 1E-3
       assert np.abs(np.abs(UpperCAmelCase__,					dtype=np.floataa	).sum() - 49_947.875	) < 5E-1

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


 def 					_lowercase (			self							:			int	):
    __lowercase		,__lowercase								=						FlaxStableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",					revision="flax",					safety_checker=UpperCAmelCase__	)

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

    __lowercase								=						jax.random.PRNGKey(0	)
    __lowercase								=						5_0

    __lowercase								=						jax.device_count()
    __lowercase								=						num_samples * [prompt]
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)

    # shard inputs and rng
    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						jax.random.split(UpperCAmelCase__,					UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)

    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).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.05_652_401)	) < 1E-3
       assert np.abs((np.abs(UpperCAmelCase__,					dtype=np.floataa	).sum() - 2_383_808.2)	) < 5E-1


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase		,__lowercase								=						FlaxStableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",					revision="bf16",					dtype=jnp.bfloataa,					safety_checker=UpperCAmelCase__	)

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

    __lowercase								=						jax.random.PRNGKey(0	)
    __lowercase								=						5_0

    __lowercase								=						jax.device_count()
    __lowercase								=						num_samples * [prompt]
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)

    # shard inputs and rng
    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						jax.random.split(UpperCAmelCase__,					UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)

    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).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.04_003_906)	) < 1E-3
       assert np.abs((np.abs(UpperCAmelCase__,					dtype=np.floataa	).sum() - 2_373_516.75)	) < 5E-1


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase		,__lowercase								=						FlaxStableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",					revision="bf16",					dtype=jnp.bfloataa	)

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

    __lowercase								=						jax.random.PRNGKey(0	)
    __lowercase								=						5_0

    __lowercase								=						jax.device_count()
    __lowercase								=						num_samples * [prompt]
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)

    # shard inputs and rng
    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						jax.random.split(UpperCAmelCase__,					UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)

    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).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.04_003_906)	) < 1E-3
       assert np.abs((np.abs(UpperCAmelCase__,					dtype=np.floataa	).sum() - 2_373_516.75)	) < 5E-1


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase								=						FlaxDDIMScheduler(
        beta_start=0.00_085,					beta_end=0.012,					beta_schedule="scaled_linear",					set_alpha_to_one=UpperCAmelCase__,					steps_offset=1,					)

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

    __lowercase								=						scheduler_state

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

    __lowercase								=						jax.random.PRNGKey(0	)
    __lowercase								=						5_0

    __lowercase								=						jax.device_count()
    __lowercase								=						num_samples * [prompt]
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)

    # shard inputs and rng
    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						jax.random.split(UpperCAmelCase__,					UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)

    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).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.045_043_945)	) < 1E-3
       assert np.abs((np.abs(UpperCAmelCase__,					dtype=np.floataa	).sum() - 2_347_693.5)	) < 5E-1



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

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

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

    __lowercase								=						replicate(UpperCAmelCase__	)
    __lowercase								=						pipeline.prepare_inputs(UpperCAmelCase__	)
    __lowercase								=						shard(UpperCAmelCase__	)
    __lowercase								=						pipeline(UpperCAmelCase__,					UpperCAmelCase__,					UpperCAmelCase__,					jit=UpperCAmelCase__	).images
    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
    __lowercase								=						images[2, 0, 2_5_6, 1_0:1_7, 1]

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

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








'''simple docstring'''







import contextlib
import os
import sqlitea

import pytest

from datasets import Dataset, Features, Value
from datasets.io.sql import SqlDatasetReader, SqlDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases, require_sqlalchemy


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Any    ,_UpperCAmelCase	:	str  )  ->			Dict:
       assert isinstance(_UpperCAmelCase    ,_UpperCAmelCase  )
       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


@require_sqlalchemy
@pytest.mark.parametrize("""keep_in_memory"""    ,[False, True]  )
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	Optional[int]    ,_UpperCAmelCase	:	Optional[int]    ,_UpperCAmelCase	:	str  )  ->			Optional[Any]:
       _a    : Any     					=tmp_path / """cache"""
       _a    : 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():
              _a    : Tuple     					=SqlDatasetReader(
                  """dataset"""    ,"""sqlite:///""" + sqlite_path    ,cache_dir=_UpperCAmelCase    ,keep_in_memory=_UpperCAmelCase  ).read()
       _check_sql_dataset(_UpperCAmelCase    ,_UpperCAmelCase  )


@require_sqlalchemy
@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 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	List[Any]    ,_UpperCAmelCase	:	Dict    ,_UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	int  )  ->			List[Any]:
       _a    : Union[str, Any]     					=tmp_path / """cache"""
       _a    : str     					={"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
       _a    : Optional[int]     					=features.copy() if features else default_expected_features
       _a    : Union[str, Any]     					=(
           Features({feature: Value(_UpperCAmelCase  ) for feature, dtype in features.items()}  ) if features is not None else None
       )
       _a    : Optional[Any]     					=SqlDatasetReader("""dataset"""    ,"""sqlite:///""" + sqlite_path    ,features=_UpperCAmelCase    ,cache_dir=_UpperCAmelCase  ).read()
       _check_sql_dataset(_UpperCAmelCase    ,_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]  )  ->			List[str]:
       with contextlib.closing(sqlitea.connect(_UpperCAmelCase  )  ) as con:
              _a    : Any     					=con.cursor()
              cur.execute("""SELECT * FROM dataset"""  )
              for row in cur:
                     yield row


@require_sqlalchemy
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Dict    ,_UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	List[str]  )  ->			Union[str, Any]:
       _a    : Union[str, Any]     					=tmp_path / """cache"""
       _a    : Union[str, Any]     					=os.path.join(_UpperCAmelCase    ,"""tmp.sql"""  )
       _a    : Tuple     					=SqlDatasetReader("""dataset"""    ,"""sqlite:///""" + sqlite_path    ,cache_dir=_UpperCAmelCase  ).read()
       SqlDatasetWriter(_UpperCAmelCase    ,"""dataset"""    ,"""sqlite:///""" + output_sqlite_path    ,num_proc=1  ).write()

       _a    : Tuple     					=iter_sql_file(_UpperCAmelCase  )
       _a    : List[Any]     					=iter_sql_file(_UpperCAmelCase  )

       for rowa, rowa in zip(_UpperCAmelCase    ,_UpperCAmelCase  ):
              assert rowa == rowa


@require_sqlalchemy
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[int]    ,_UpperCAmelCase	:	Any    ,_UpperCAmelCase	:	List[Any]  )  ->			Optional[int]:
       _a    : int     					=tmp_path / """cache"""
       _a    : Any     					=os.path.join(_UpperCAmelCase    ,"""tmp.sql"""  )
       _a    : Union[str, Any]     					=SqlDatasetReader("""dataset"""    ,"""sqlite:///""" + sqlite_path    ,cache_dir=_UpperCAmelCase  ).read()
       SqlDatasetWriter(_UpperCAmelCase    ,"""dataset"""    ,"""sqlite:///""" + output_sqlite_path    ,num_proc=2  ).write()

       _a    : List[Any]     					=iter_sql_file(_UpperCAmelCase  )
       _a    : str     					=iter_sql_file(_UpperCAmelCase  )

       for rowa, rowa in zip(_UpperCAmelCase    ,_UpperCAmelCase  ):
              assert rowa == rowa


@require_sqlalchemy
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Union[str, Any]    ,_UpperCAmelCase	:	str    ,_UpperCAmelCase	:	List[Any]  )  ->			List[str]:
       _a    : List[str]     					=tmp_path / """cache"""
       _a    : Dict     					=os.path.join(_UpperCAmelCase    ,"""tmp.sql"""  )
       _a    : Optional[Any]     					=SqlDatasetReader("""dataset"""    ,"""sqlite:///""" + sqlite_path    ,cache_dir=_UpperCAmelCase  ).read()
       with pytest.raises(_UpperCAmelCase  ):
              SqlDatasetWriter(_UpperCAmelCase    ,"""dataset"""    ,"""sqlite:///""" + output_sqlite_path    ,num_proc=0  ).write()


