



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







							'''simple docstring'''






							def _a      (							self		):
											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  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											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_514_745		) < 1E-3
															assert np.abs(np.abs(a_  ,  dtype=np.floataa		).sum() - 49_947.875		) < 5E-1

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


							def _a      (							self		):
											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, 512, 512, 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(a_  ,  dtype=np.floataa		).sum() - 2_383_808.2)		) < 5E-1


							def _a      (							self		):
											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, 512, 512, 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(a_  ,  dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1


							def _a      (							self		):
											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, 512, 512, 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(a_  ,  dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1


							def _a      (							self		):
											a__				      =		FlaxDDIMScheduler(
											    beta_start=0.00_085  ,  beta_end=0.012  ,  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, 512, 512, 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(a_  ,  dtype=np.floataa		).sum() - 2_347_693.5)		) < 5E-1







							def _a      (							self		):
											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, 512, 512, 3)
											a__				      =		images[2, 0, 256, 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, 512, 512, 3)
											a__				      =		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






import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDModel(
											    sample_size=(32, 64)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDConditionModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")  ,  cross_attention_dim=10  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		AutoencoderKL(
											    sample_size=(128, 64)  ,  in_channels=1  ,  out_channels=1  ,  latent_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")  ,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")  ,  )
											a__				      =		UNetaDModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return vqvae, unet







							@slow
							def _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		Mel(
											    x_res=self.dummy_unet.config.sample_size[1]  ,  y_res=self.dummy_unet.config.sample_size[0]  ,  )

											a__				      =		DDPMScheduler()
											a__				      =		AudioDiffusionPipeline(vqvae=a_  ,  unet=self.dummy_unet  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4  ,  return_dict=a_		)
											a__				      =		output[0][0]

											assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
											assert (
											    image.height == self.dummy_unet.config.sample_size[0]
											    and image.width == self.dummy_unet.config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.frombuffer(image_from_tuple.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0
											assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() == 0

											a__				      =		Mel(
											    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]  ,  y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]  ,  )

											a__				      =		DDIMScheduler()
											a__				      =		self.dummy_vqvae_and_unet
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=dummy_vqvae_and_unet[1]  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		np.random.uniform(-1  ,  1  ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(raw_audio=a_  ,  generator=a_  ,  start_step=5  ,  steps=10		)
											a__				      =		output.images[0]

											assert (
											    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
											    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

											a__				      =		self.dummy_unet_condition
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=a_  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0




@slow
@require_torch_gpu
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		torch_device

											a__				      =		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
											assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

