






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






						def      __snake_case      (		self  :List[Any]				)      ->Optional[int]:
										with tempfile.TemporaryDirectory() as tmpdirname:
														# pipeline has Flax weights
														lowercase    :      str       = FlaxDiffusionPipeline.from_pretrained(
														    """hf-internal-testing/tiny-stable-diffusion-pipe"""			, safety_checker=__magic_name__			, cache_dir=__magic_name__				)

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






						def      __snake_case      (		self  :Union[str, Any]				)      ->Any:
										lowercase      ,      lowercase    :      Tuple       = FlaxStableDiffusionPipeline.from_pretrained(
										    """hf-internal-testing/tiny-stable-diffusion-pipe"""			, safety_checker=__magic_name__				)

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

										lowercase    :      Optional[int]       = jax.random.PRNGKey(0				)
										lowercase    :      Optional[Any]       = 4

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

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

										lowercase    :      Union[str, Any]       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).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.151_4745				) < 1E-3
														assert np.abs(np.abs(__magic_name__			, dtype=np.floataa				).sum() - 4_9947.875				) < 5E-1

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






						def      __snake_case      (		self  :int				)      ->Union[str, Any]:
										lowercase      ,      lowercase    :      Optional[int]       = FlaxStableDiffusionPipeline.from_pretrained(
										    """CompVis/stable-diffusion-v1-4"""			, revision="""flax"""			, safety_checker=__magic_name__				)

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

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

										lowercase    :      Any       = jax.device_count()
										lowercase    :      Tuple       = num_samples * [prompt]
										lowercase    :      Union[str, Any]       = pipeline.prepare_inputs(__magic_name__				)

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

										lowercase    :      Tuple       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).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.0565_2401)				) < 1E-3
														assert np.abs((np.abs(__magic_name__			, dtype=np.floataa				).sum() - 238_3808.2)				) < 5E-1






						def      __snake_case      (		self  :int				)      ->List[str]:
										lowercase      ,      lowercase    :      List[str]       = FlaxStableDiffusionPipeline.from_pretrained(
										    """CompVis/stable-diffusion-v1-4"""			, revision="""bf16"""			, dtype=jnp.bfloataa			, safety_checker=__magic_name__				)

										lowercase    :      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"""
										)

										lowercase    :      Tuple       = jax.random.PRNGKey(0				)
										lowercase    :      str       = 50

										lowercase    :      Optional[int]       = jax.device_count()
										lowercase    :      Union[str, Any]       = num_samples * [prompt]
										lowercase    :      Dict       = pipeline.prepare_inputs(__magic_name__				)

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

										lowercase    :      str       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).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.0400_3906)				) < 1E-3
														assert np.abs((np.abs(__magic_name__			, dtype=np.floataa				).sum() - 237_3516.75)				) < 5E-1






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

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

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

										lowercase    :      Union[str, Any]       = jax.device_count()
										lowercase    :      Optional[int]       = num_samples * [prompt]
										lowercase    :      Optional[int]       = pipeline.prepare_inputs(__magic_name__				)

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

										lowercase    :      Optional[Any]       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).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.0400_3906)				) < 1E-3
														assert np.abs((np.abs(__magic_name__			, dtype=np.floataa				).sum() - 237_3516.75)				) < 5E-1






						def      __snake_case      (		self  :str				)      ->Dict:
										lowercase    :      Union[str, Any]       = FlaxDDIMScheduler(
										    beta_start=0.0_0085			, beta_end=0.012			, beta_schedule="""scaled_linear"""			, set_alpha_to_one=__magic_name__			, steps_offset=1			, )

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

										lowercase    :      List[str]       = scheduler_state

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

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

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

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

										lowercase    :      Any       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).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.0_4504_3945)				) < 1E-3
														assert np.abs((np.abs(__magic_name__			, dtype=np.floataa				).sum() - 234_7693.5)				) < 5E-1






						def      __snake_case      (		self  :Tuple				)      ->Any:
										lowercase    :      str       = (
										    """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
										    """ field, close up, split lighting, cinematic"""
										)

										lowercase    :      Optional[Any]       = jax.device_count()
										lowercase    :      Any       = num_samples * [prompt]
										lowercase    :      List[Any]       = jax.random.split(jax.random.PRNGKey(0				)			, __magic_name__				)

										lowercase      ,      lowercase    :      Any       = FlaxStableDiffusionPipeline.from_pretrained(
										    """CompVis/stable-diffusion-v1-4"""			, revision="""bf16"""			, dtype=jnp.bfloataa			, safety_checker=__magic_name__			, )

										lowercase    :      Optional[Any]       = replicate(__magic_name__				)
										lowercase    :      Optional[int]       = pipeline.prepare_inputs(__magic_name__				)
										lowercase    :      List[str]       = shard(__magic_name__				)
										lowercase    :      List[str]       = pipeline(__magic_name__			, __magic_name__			, __magic_name__			, jit=__magic_name__				).images
										assert images.shape == (num_samples, 1, 512, 512, 3)
										lowercase    :      Optional[Any]       = images[2, 0, 256, 10:17, 1]

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

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







from typing import List, Optional, Tuple, Union

import torch

from ...schedulers import DDIMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class      UpperCamelCase		(__snake_case     ):






						def __init__(		self  :Any			, __magic_name__  :Union[str, Any]			, __magic_name__  :Optional[Any]				)      ->Dict:
										super().__init__()

										# make sure scheduler can always be converted to DDIM
										lowercase    :      Optional[Any]       = DDIMScheduler.from_config(scheduler.config				)

										self.register_modules(unet=__magic_name__			, scheduler=__magic_name__				)






						@torch.no_grad()
						def __call__(		self  :Optional[int]			, __magic_name__  :int = 1			, __magic_name__  :Optional[Union[torch.Generator, List[torch.Generator]]] = None			, __magic_name__  :float = 0.0			, __magic_name__  :int = 50			, __magic_name__  :Optional[bool] = None			, __magic_name__  :Optional[str] = "pil"			, __magic_name__  :bool = True			, )      ->Union[ImagePipelineOutput, Tuple]:

										# Sample gaussian noise to begin loop
										if isinstance(self.unet.config.sample_size			, __magic_name__				):
														lowercase    :      int       = (
														    batch_size,
														    self.unet.config.in_channels,
														    self.unet.config.sample_size,
														    self.unet.config.sample_size,
														)
										else:
														lowercase    :      Tuple       = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)

										if isinstance(__magic_name__			, __magic_name__				) and len(__magic_name__				) != batch_size:
														raise ValueError(
														    f"""You have passed a list of generators of length {len(__magic_name__				)}, but requested an effective batch"""
														    f""" size of {batch_size}. Make sure the batch size matches the length of the generators."""				)

										lowercase    :      Dict       = randn_tensor(__magic_name__			, generator=__magic_name__			, device=self.device			, dtype=self.unet.dtype				)

										# set step values
										self.scheduler.set_timesteps(__magic_name__				)

										for t in self.progress_bar(self.scheduler.timesteps				):
														# 1. predict noise model_output
														lowercase    :      Union[str, Any]       = self.unet(__magic_name__			, __magic_name__				).sample

														# 2. predict previous mean of image x_t-1 and add variance depending on eta
														# eta corresponds to η in paper and should be between [0, 1]
														# do x_t -> x_t-1
														lowercase    :      int       = self.scheduler.step(
														    __magic_name__			, __magic_name__			, __magic_name__			, eta=__magic_name__			, use_clipped_model_output=__magic_name__			, generator=__magic_name__				).prev_sample

										lowercase    :      Optional[Any]       = (image / 2 + 0.5).clamp(0			, 1				)
										lowercase    :      Tuple       = image.cpu().permute(0			, 2			, 3			, 1				).numpy()
										if output_type == "pil":
														lowercase    :      Any       = self.numpy_to_pil(__magic_name__				)

										if not return_dict:
														return (image,)

										return ImagePipelineOutput(images=__magic_name__				)
