






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




						"""simple docstring"""




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												with tempfile.TemporaryDirectory() as tmpdirname:
																		# pipeline has Flax weights
																		lowercase__    :				Dict=	FlaxDiffusionPipeline.from_pretrained(
																		    "hf-internal-testing/tiny-stable-diffusion-pipe"	,  safety_checker=snake_case__	,  cache_dir=snake_case__      )

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




						"""simple docstring"""




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__,     lowercase__    :				List[Any]=	FlaxStableDiffusionPipeline.from_pretrained(
												    "hf-internal-testing/tiny-stable-diffusion-pipe"	,  safety_checker=snake_case__      )

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

												lowercase__    :				Dict=	jax.random.PRNGKey(0      )
												lowercase__    :				List[str]=	4

												lowercase__    :				List[Any]=	jax.device_count()
												lowercase__    :				Dict=	num_samples * [prompt]
												lowercase__    :				List[str]=	pipeline.prepare_inputs(snake_case__      )

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

												lowercase__    :				Tuple=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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_51_47_45      ) < 1e-3
																		assert np.abs(np.abs(snake_case__	,  dtype=np.floataa      ).sum() - 4_99_47.8_75      ) < 5e-1

												lowercase__    :				Optional[Any]=	pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]      )      )      )
												assert len(snake_case__      ) == num_samples




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__,     lowercase__    :				Any=	FlaxStableDiffusionPipeline.from_pretrained(
												    "CompVis/stable-diffusion-v1-4"	,  revision="flax"	,  safety_checker=snake_case__      )

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

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

												lowercase__    :				List[str]=	jax.device_count()
												lowercase__    :				Optional[int]=	num_samples * [prompt]
												lowercase__    :				List[str]=	pipeline.prepare_inputs(snake_case__      )

												# shard inputs and rng
												lowercase__    :				Any=	replicate(snake_case__      )
												lowercase__    :				Optional[Any]=	jax.random.split(snake_case__	,  snake_case__      )
												lowercase__    :				str=	shard(snake_case__      )

												lowercase__    :				Optional[Any]=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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_65_24_01)      ) < 1e-3
																		assert np.abs((np.abs(snake_case__	,  dtype=np.floataa      ).sum() - 2_38_38_08.2)      ) < 5e-1




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__,     lowercase__    :				List[Any]=	FlaxStableDiffusionPipeline.from_pretrained(
												    "CompVis/stable-diffusion-v1-4"	,  revision="bf16"	,  dtype=jnp.bfloataa	,  safety_checker=snake_case__      )

												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__    :				Dict=	jax.random.PRNGKey(0      )
												lowercase__    :				str=	50

												lowercase__    :				Union[str, Any]=	jax.device_count()
												lowercase__    :				List[Any]=	num_samples * [prompt]
												lowercase__    :				Optional[Any]=	pipeline.prepare_inputs(snake_case__      )

												# shard inputs and rng
												lowercase__    :				Optional[Any]=	replicate(snake_case__      )
												lowercase__    :				str=	jax.random.split(snake_case__	,  snake_case__      )
												lowercase__    :				Any=	shard(snake_case__      )

												lowercase__    :				Optional[int]=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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_00_39_06)      ) < 1e-3
																		assert np.abs((np.abs(snake_case__	,  dtype=np.floataa      ).sum() - 2_37_35_16.75)      ) < 5e-1




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__,     lowercase__    :				int=	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__    :				List[str]=	jax.device_count()
												lowercase__    :				List[str]=	num_samples * [prompt]
												lowercase__    :				Any=	pipeline.prepare_inputs(snake_case__      )

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

												lowercase__    :				int=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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_00_39_06)      ) < 1e-3
																		assert np.abs((np.abs(snake_case__	,  dtype=np.floataa      ).sum() - 2_37_35_16.75)      ) < 5e-1




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__    :				str=	FlaxDDIMScheduler(
												    beta_start=0.0_00_85	,  beta_end=0.0_12	,  beta_schedule="scaled_linear"	,  set_alpha_to_one=snake_case__	,  steps_offset=1	,  )

												lowercase__,     lowercase__    :				str=	FlaxStableDiffusionPipeline.from_pretrained(
												    "CompVis/stable-diffusion-v1-4"	,  revision="bf16"	,  dtype=jnp.bfloataa	,  scheduler=snake_case__	,  safety_checker=snake_case__	,  )
												lowercase__    :				List[str]=	scheduler.create_state()

												lowercase__    :				List[Any]=	scheduler_state

												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__    :				str=	jax.random.PRNGKey(0      )
												lowercase__    :				List[Any]=	50

												lowercase__    :				int=	jax.device_count()
												lowercase__    :				int=	num_samples * [prompt]
												lowercase__    :				List[str]=	pipeline.prepare_inputs(snake_case__      )

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

												lowercase__    :				Optional[int]=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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_45_04_39_45)      ) < 1e-3
																		assert np.abs((np.abs(snake_case__	,  dtype=np.floataa      ).sum() - 2_34_76_93.5)      ) < 5e-1




						def 		UpperCAmelCase_	(       self      ):
												'''simple docstring'''
												lowercase__    :				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.device_count()
												lowercase__    :				Optional[Any]=	num_samples * [prompt]
												lowercase__    :				Dict=	jax.random.split(jax.random.PRNGKey(0      )	,  snake_case__      )

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

												lowercase__    :				Tuple=	replicate(snake_case__      )
												lowercase__    :				str=	pipeline.prepare_inputs(snake_case__      )
												lowercase__    :				Tuple=	shard(snake_case__      )
												lowercase__    :				Union[str, Any]=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).images
												assert images.shape == (num_samples, 1, 512, 512, 3)
												lowercase__    :				List[str]=	images[2, 0, 256, 10:17, 1]

												# With memory efficient attention
												lowercase__,     lowercase__    :				Optional[int]=	FlaxStableDiffusionPipeline.from_pretrained(
												    "CompVis/stable-diffusion-v1-4"	,  revision="bf16"	,  dtype=jnp.bfloataa	,  safety_checker=snake_case__	,  use_memory_efficient_attention=snake_case__	,  )

												lowercase__    :				Any=	replicate(snake_case__      )
												lowercase__    :				int=	pipeline.prepare_inputs(snake_case__      )
												lowercase__    :				Dict=	shard(snake_case__      )
												lowercase__    :				List[str]=	pipeline(snake_case__	,  snake_case__	,  snake_case__	,  jit=snake_case__      ).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 Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class    __UpperCAmelCase(       SCREAMING_SNAKE_CASE__							):




						"""simple docstring"""




						def __init__(       self	,  snake_case__	,  snake_case__	,  snake_case__	,  snake_case__ = None	,  ):
												'''simple docstring'''
												super().__init__()
												self.register_modules(transformer=snake_case__	,  vae=snake_case__	,  scheduler=snake_case__      )

												# create a imagenet -> id dictionary for easier use
												lowercase__    :				int=	{}
												if idalabel is not None:
																		for key, value in idalabel.items():
																								for label in value.split(","      ):
																														lowercase__    :				Tuple=	int(snake_case__      )
																		lowercase__    :				Union[str, Any]=	dict(sorted(self.labels.items()      )      )




						def 		UpperCAmelCase_	(       self	,  snake_case__      ):
												'''simple docstring'''

												if not isinstance(snake_case__	,  snake_case__      ):
																		lowercase__    :				List[Any]=	list(snake_case__      )

												for l in label:
																		if l not in self.labels:
																								raise ValueError(
																								    F'''{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}.'''      )

												return [self.labels[l] for l in label]




						@torch.no_grad()
						def __call__(       self	,  snake_case__	,  snake_case__ = 4.0	,  snake_case__ = None	,  snake_case__ = 50	,  snake_case__ = "pil"	,  snake_case__ = True	,  ):
												'''simple docstring'''
												lowercase__    :				List[Any]=	len(snake_case__      )
												lowercase__    :				Optional[int]=	self.transformer.config.sample_size
												lowercase__    :				List[str]=	self.transformer.config.in_channels

												lowercase__    :				Any=	randn_tensor(
												    shape=(batch_size, latent_channels, latent_size, latent_size)	,  generator=snake_case__	,  device=self.device	,  dtype=self.transformer.dtype	,  )
												lowercase__    :				Any=	torch.cat([latents] * 2      ) if guidance_scale > 1 else latents

												lowercase__    :				Tuple=	torch.tensor(snake_case__	,  device=self.device      ).reshape(-1      )
												lowercase__    :				Any=	torch.tensor([1000] * batch_size	,  device=self.device      )
												lowercase__    :				Tuple=	torch.cat([class_labels, class_null]	,  0      ) if guidance_scale > 1 else class_labels

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

												for t in self.progress_bar(self.scheduler.timesteps      ):
																		if guidance_scale > 1:
																								lowercase__    :				List[str]=	latent_model_input[: len(snake_case__      ) // 2]
																								lowercase__    :				int=	torch.cat([half, half]	,  dim=0      )
																		lowercase__    :				Union[str, Any]=	self.scheduler.scale_model_input(snake_case__	,  snake_case__      )

																		lowercase__    :				Optional[int]=	t
																		if not torch.is_tensor(snake_case__      ):
																								# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
																								# This would be a good case for the `match` statement (Python 3.10+)
																								lowercase__    :				List[str]=	latent_model_input.device.type == "mps"
																								if isinstance(snake_case__	,  snake_case__      ):
																														lowercase__    :				int=	torch.floataa if is_mps else torch.floataa
																								else:
																														lowercase__    :				Dict=	torch.intaa if is_mps else torch.intaa
																								lowercase__    :				Tuple=	torch.tensor([timesteps]	,  dtype=snake_case__	,  device=latent_model_input.device      )
																		elif len(timesteps.shape      ) == 0:
																								lowercase__    :				Dict=	timesteps[None].to(latent_model_input.device      )
																		# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
																		lowercase__    :				int=	timesteps.expand(latent_model_input.shape[0]      )
																		# predict noise model_output
																		lowercase__    :				Union[str, Any]=	self.transformer(
																		    snake_case__	,  timestep=snake_case__	,  class_labels=snake_case__      ).sample

																		# perform guidance
																		if guidance_scale > 1:
																								lowercase__,     lowercase__    :				Tuple=	noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
																								lowercase__,     lowercase__    :				Union[str, Any]=	torch.split(snake_case__	,  len(snake_case__      ) // 2	,  dim=0      )

																								lowercase__    :				str=	uncond_eps + guidance_scale * (cond_eps - uncond_eps)
																								lowercase__    :				Dict=	torch.cat([half_eps, half_eps]	,  dim=0      )

																								lowercase__    :				Optional[int]=	torch.cat([eps, rest]	,  dim=1      )

																		# learned sigma
																		if self.transformer.config.out_channels // 2 == latent_channels:
																								lowercase__,     lowercase__    :				Union[str, Any]=	torch.split(snake_case__	,  snake_case__	,  dim=1      )
																		else:
																								lowercase__    :				int=	noise_pred

																		# compute previous image: x_t -> x_t-1
																		lowercase__    :				List[Any]=	self.scheduler.step(snake_case__	,  snake_case__	,  snake_case__      ).prev_sample

												if guidance_scale > 1:
																		lowercase__,     lowercase__    :				Any=	latent_model_input.chunk(2	,  dim=0      )
												else:
																		lowercase__    :				str=	latent_model_input

												lowercase__    :				Dict=	1 / self.vae.config.scaling_factor * latents
												lowercase__    :				Any=	self.vae.decode(snake_case__      ).sample

												lowercase__    :				Tuple=	(samples / 2 + 0.5).clamp(0	,  1      )

												# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
												lowercase__    :				List[Any]=	samples.cpu().permute(0	,  2	,  3	,  1      ).float().numpy()

												if output_type == "pil":
																		lowercase__    :				Optional[Any]=	self.numpy_to_pil(snake_case__      )

												if not return_dict:
																		return (samples,)

												return ImagePipelineOutput(images=snake_case__      )






