





'''simple docstring'''




import inspect
from typing import List, Optional, Tuple, Union

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput






def        _lowerCAmelCase			(  __snake_case  : Dict			)     ->     Optional[int]:
     __A ,__A						: Union[str, Any]      =       image.size
     __A ,__A						: List[str]      =       (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
     __A						: Optional[Any]      =       image.resize((w, h) ,							resample=PIL_INTERPOLATION['lanczos']			)
     __A						: Optional[Any]      =       np.array(__snake_case			).astype(np.floataa			) / 255.0
     __A						: int      =       image[None].transpose(0 ,							3 ,							1 ,							2			)
     __A						: Tuple      =       torch.from_numpy(__snake_case			)
     return 2.0 * image - 1.0







class  SCREAMING_SNAKE_CASE    (a__							):


     def __init__(						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	):

          '''simple docstring'''






          super().__init__()
          self.register_modules(vqvae=_UpperCAmelCase						,	unet=_UpperCAmelCase						,	scheduler=_UpperCAmelCase)


     @torch.no_grad()
     def __call__(						self						,	_UpperCAmelCase = None						,	_UpperCAmelCase = 1						,	_UpperCAmelCase = 100						,	_UpperCAmelCase = 0.0						,	_UpperCAmelCase = None						,	_UpperCAmelCase = "pil"						,	_UpperCAmelCase = True						,	):

          '''simple docstring'''






          if isinstance(_UpperCAmelCase						,	PIL.Image.Image):
               __A						: Any      =       1
          elif isinstance(_UpperCAmelCase						,	torch.Tensor):
               __A						: Optional[int]      =       image.shape[0]
          else:
               raise ValueError(F'`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(_UpperCAmelCase)}')

          if isinstance(_UpperCAmelCase						,	PIL.Image.Image):
               __A						: List[str]      =       preprocess(_UpperCAmelCase)

          __A ,__A						: List[str]      =       image.shape[-2:]

          # in_channels should be 6: 3 for latents, 3 for low resolution image
          __A						: List[Any]      =       (batch_size, self.unet.config.in_channels // 2, height, width)
          __A						: Union[str, Any]      =       next(self.unet.parameters()).dtype

          __A						: str      =       randn_tensor(_UpperCAmelCase						,	generator=_UpperCAmelCase						,	device=self.device						,	dtype=_UpperCAmelCase)

          __A						: List[Any]      =       image.to(device=self.device						,	dtype=_UpperCAmelCase)

          # set timesteps and move to the correct device
          self.scheduler.set_timesteps(_UpperCAmelCase						,	device=self.device)
          __A						: List[Any]      =       self.scheduler.timesteps

          # scale the initial noise by the standard deviation required by the scheduler
          __A						: Union[str, Any]      =       latents * self.scheduler.init_noise_sigma

          # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
          # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
          # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
          # and should be between [0, 1]
          __A						: Union[str, Any]      =       'eta' in set(inspect.signature(self.scheduler.step).parameters.keys())
          __A						: Optional[Any]      =       {}
          if accepts_eta:
               __A						: Union[str, Any]      =       eta

          for t in self.progress_bar(_UpperCAmelCase):
               # concat latents and low resolution image in the channel dimension.
               __A						: Any      =       torch.cat([latents, image]						,	dim=1)
               __A						: Dict      =       self.scheduler.scale_model_input(_UpperCAmelCase						,	_UpperCAmelCase)
               # predict the noise residual
               __A						: List[str]      =       self.unet(_UpperCAmelCase						,	_UpperCAmelCase).sample
               # compute the previous noisy sample x_t -> x_t-1
               __A						: int      =       self.scheduler.step(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	**_UpperCAmelCase).prev_sample

          # decode the image latents with the VQVAE
          __A						: List[str]      =       self.vqvae.decode(_UpperCAmelCase).sample
          __A						: List[Any]      =       torch.clamp(_UpperCAmelCase						,	-1.0						,	1.0)
          __A						: Tuple      =       image / 2 + 0.5
          __A						: Tuple      =       image.cpu().permute(0						,	2						,	3						,	1).numpy()

          if output_type == "pil":
               __A						: Dict      =       self.numpy_to_pil(_UpperCAmelCase)

          if not return_dict:
               return (image,)

          return ImagePipelineOutput(images=_UpperCAmelCase)






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          with tempfile.TemporaryDirectory() as tmpdirname:
               # pipeline has Flax weights
               __A						: Union[str, Any]      =       FlaxDiffusionPipeline.from_pretrained(
                   'hf-internal-testing/tiny-stable-diffusion-pipe'						,	safety_checker=_UpperCAmelCase						,	cache_dir=_UpperCAmelCase)

               __A						: Optional[Any]      =       [t[-1] for t in os.walk(os.path.join(_UpperCAmelCase						,	os.listdir(_UpperCAmelCase)[0]						,	'snapshots'))]
               __A						: 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  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: Optional[int]      =       FlaxStableDiffusionPipeline.from_pretrained(
              'hf-internal-testing/tiny-stable-diffusion-pipe'						,	safety_checker=_UpperCAmelCase)

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

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

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

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

          __A						: Union[str, Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).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.1514745) < 1e-3
               assert np.abs(np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 49947.875) < 5e-1

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A ,__A						: str      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='flax'						,	safety_checker=_UpperCAmelCase)

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

          __A						: Tuple      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: str      =       jax.device_count()
          __A						: Union[str, Any]      =       num_samples * [prompt]
          __A						: List[str]      =       pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						: Tuple      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).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.05652401)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2383808.2)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

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

          __A						: str      =       jax.random.PRNGKey(0)
          __A						: Any      =       50

          __A						: Optional[int]      =       jax.device_count()
          __A						: int      =       num_samples * [prompt]
          __A						: Optional[int]      =       pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						: Optional[int]      =       replicate(_UpperCAmelCase)
          __A						: List[str]      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Dict      =       shard(_UpperCAmelCase)

          __A						: str      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

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

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

          __A						: Optional[int]      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						: int      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[Any]      =       FlaxDDIMScheduler(
              beta_start=0.00085						,	beta_end=0.012						,	beta_schedule='scaled_linear'						,	set_alpha_to_one=_UpperCAmelCase						,	steps_offset=1						,	)

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

          __A						: Any      =       scheduler_state

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

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

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

          # shard inputs and rng
          __A						: int      =       replicate(_UpperCAmelCase)
          __A						: Any      =       jax.random.split(_UpperCAmelCase						,	_UpperCAmelCase)
          __A						: Tuple      =       shard(_UpperCAmelCase)

          __A						: Union[str, Any]      =       pipeline(_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	jit=_UpperCAmelCase).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.045043945)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,	dtype=np.floataa).sum() - 2347693.5)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						: int      =       jax.device_count()
          __A						: List[Any]      =       num_samples * [prompt]
          __A						: List[Any]      =       jax.random.split(jax.random.PRNGKey(0)						,	_UpperCAmelCase)

          __A ,__A						: Tuple      =       FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,	revision='bf16'						,	dtype=jnp.bfloataa						,	safety_checker=_UpperCAmelCase						,	)

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

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

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