






"""simple docstring"""



import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)
def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  str		,					__snake_case			:  Tuple		,					__snake_case			:  int    ):





  '''simple docstring'''





  if isinstance(__snake_case		,					torch.Tensor    ):
    return image
  elif isinstance(__snake_case		,					PIL.Image.Image    ):
    lowercase							    =       [image]

  if isinstance(image[0]		,					PIL.Image.Image    ):
    lowercase							    =       [np.array(i.resize((w, h)		,					resample=PIL_INTERPOLATION['lanczos']    )    )[None, :] for i in image]
    lowercase							    =       np.concatenate(__snake_case		,					axis=0    )
    lowercase							    =       np.array(__snake_case    ).astype(np.floataa    ) / 255.0
    lowercase							    =       image.transpose(0		,					3		,					1		,					2    )
    lowercase							    =       2.0 * image - 1.0
    lowercase							    =       torch.from_numpy(__snake_case    )
  elif isinstance(image[0]		,					torch.Tensor    ):
    lowercase							    =       torch.cat(__snake_case		,					dim=0    )
  return image
def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  Optional[int]		,					__snake_case			:  Optional[int]		,					__snake_case			:  Optional[int]		,					__snake_case			:  Any=0.9995    ):





  '''simple docstring'''





  if not isinstance(__snake_case		,					np.ndarray    ):
    lowercase							    =       True
    lowercase							    =       va.device
    lowercase							    =       va.cpu().numpy()
    lowercase							    =       va.cpu().numpy()

  lowercase							    =       np.sum(va * va / (np.linalg.norm(__snake_case    ) * np.linalg.norm(__snake_case    ))    )
  if np.abs(__snake_case    ) > DOT_THRESHOLD:
    lowercase							    =       (1 - t) * va + t * va
  else:
    lowercase							    =       np.arccos(__snake_case    )
    lowercase							    =       np.sin(__snake_case    )
    lowercase							    =       theta_a * t
    lowercase							    =       np.sin(__snake_case    )
    lowercase							    =       np.sin(theta_a - theta_t    ) / sin_theta_a
    lowercase							    =       sin_theta_t / sin_theta_a
    lowercase							    =       sa * va + sa * va

  if inputs_are_torch:
    lowercase							    =       torch.from_numpy(__snake_case    ).to(__snake_case    )

  return va
def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  Optional[int]		,					__snake_case			:  List[Any]    ):





  '''simple docstring'''





  lowercase							    =       F.normalize(__snake_case		,					dim=-1    )
  lowercase							    =       F.normalize(__snake_case		,					dim=-1    )
  return (x - y).norm(dim=-1    ).div(2    ).arcsin().pow(2    ).mul(2    )
def 						_SCREAMING_SNAKE_CASE (    __snake_case			:  Union[str, Any]		,					__snake_case			:  List[Any]    ):





  '''simple docstring'''





  for param in model.parameters():
    lowercase							    =       value



class        a     (   a_	):







    def __init__(			self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase=None							,		_lowerCamelCase=None							,		_lowerCamelCase=None							,		):
      super().__init__()
      self.register_modules(
          vae=_lowerCamelCase							,		text_encoder=_lowerCamelCase							,		clip_model=_lowerCamelCase							,		tokenizer=_lowerCamelCase							,		unet=_lowerCamelCase							,		scheduler=_lowerCamelCase							,		feature_extractor=_lowerCamelCase							,		coca_model=_lowerCamelCase							,		coca_tokenizer=_lowerCamelCase							,		coca_transform=_lowerCamelCase							,		)
      lowercase							    =       (
          feature_extractor.size
          if isinstance(feature_extractor.size							,		_lowerCamelCase       )
          else feature_extractor.size['shortest_edge']
      )
      lowercase							    =       transforms.Normalize(mean=feature_extractor.image_mean							,		std=feature_extractor.image_std       )
      set_requires_grad(self.text_encoder							,		_lowerCamelCase       )
      set_requires_grad(self.clip_model							,		_lowerCamelCase       )







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase = "auto"       ):
      if slice_size == "auto":
        # half the attention head size is usually a good trade-off between
        # speed and memory
        lowercase							    =       self.unet.config.attention_head_dim // 2
      self.unet.set_attention_slice(_lowerCamelCase       )







    def 			UpperCamelCase_      (			self       ):
      self.enable_attention_slicing(_lowerCamelCase       )







    def 			UpperCamelCase_      (			self       ):
      set_requires_grad(self.vae							,		_lowerCamelCase       )







    def 			UpperCamelCase_      (			self       ):
      set_requires_grad(self.vae							,		_lowerCamelCase       )







    def 			UpperCamelCase_      (			self       ):
      set_requires_grad(self.unet							,		_lowerCamelCase       )







    def 			UpperCamelCase_      (			self       ):
      set_requires_grad(self.unet							,		_lowerCamelCase       )







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       ):
      # get the original timestep using init_timestep
      lowercase							    =       min(int(num_inference_steps * strength       )							,		_lowerCamelCase       )

      lowercase							    =       max(num_inference_steps - init_timestep							,		0       )
      lowercase							    =       self.scheduler.timesteps[t_start:]

      return timesteps, num_inference_steps - t_start







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase=None       ):
      if not isinstance(_lowerCamelCase							,		torch.Tensor       ):
        raise ValueError(F'`image` has to be of type `torch.Tensor` but is {type(_lowerCamelCase       )}'       )

      lowercase							    =       image.to(device=_lowerCamelCase							,		dtype=_lowerCamelCase       )

      if isinstance(_lowerCamelCase							,		_lowerCamelCase       ):
        lowercase							    =       [
            self.vae.encode(image[i : i + 1]       ).latent_dist.sample(generator[i]       ) for i in range(_lowerCamelCase       )
        ]
        lowercase							    =       torch.cat(_lowerCamelCase							,		dim=0       )
      else:
        lowercase							    =       self.vae.encode(_lowerCamelCase       ).latent_dist.sample(_lowerCamelCase       )

      # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
      lowercase							    =       0.1_8_2_1_5 * init_latents
      lowercase							    =       init_latents.repeat_interleave(_lowerCamelCase							,		dim=0       )

      lowercase							    =       randn_tensor(init_latents.shape							,		generator=_lowerCamelCase							,		device=_lowerCamelCase							,		dtype=_lowerCamelCase       )

      # get latents
      lowercase							    =       self.scheduler.add_noise(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )
      lowercase							    =       init_latents

      return latents







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase       ):
      lowercase							    =       self.coca_transform(_lowerCamelCase       ).unsqueeze(0       )
      with torch.no_grad(), torch.cuda.amp.autocast():
        lowercase							    =       self.coca_model.generate(transformed_image.to(device=self.device							,		dtype=self.coca_model.dtype       )       )
      lowercase							    =       self.coca_tokenizer.decode(generated[0].cpu().numpy()       )
      return generated.split('<end_of_text>'       )[0].replace('<start_of_text>'							,		''       ).rstrip(' .,'       )







    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase       ):
      lowercase							    =       self.feature_extractor.preprocess(_lowerCamelCase       )
      lowercase							    =       torch.from_numpy(clip_image_input['pixel_values'][0]       ).unsqueeze(0       ).to(self.device       ).half()
      lowercase							    =       self.clip_model.get_image_features(_lowerCamelCase       )
      lowercase							    =       image_embeddings_clip / image_embeddings_clip.norm(p=2							,		dim=-1							,		keepdim=_lowerCamelCase       )
      lowercase							    =       image_embeddings_clip.repeat_interleave(_lowerCamelCase							,		dim=0       )
      return image_embeddings_clip







    @torch.enable_grad()
    def 			UpperCamelCase_      (			self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		):
      lowercase							    =       latents.detach().requires_grad_()

      lowercase							    =       self.scheduler.scale_model_input(_lowerCamelCase							,		_lowerCamelCase       )

      # predict the noise residual
      lowercase							    =       self.unet(_lowerCamelCase							,		_lowerCamelCase							,		encoder_hidden_states=_lowerCamelCase       ).sample

      if isinstance(self.scheduler							,		(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)       ):
        lowercase							    =       self.scheduler.alphas_cumprod[timestep]
        lowercase							    =       1 - alpha_prod_t
        # compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        lowercase							    =       (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

        lowercase							    =       torch.sqrt(_lowerCamelCase       )
        lowercase							    =       pred_original_sample * (fac) + latents * (1 - fac)
      elif isinstance(self.scheduler							,		_lowerCamelCase       ):
        lowercase							    =       self.scheduler.sigmas[index]
        lowercase							    =       latents - sigma * noise_pred
      else:
        raise ValueError(F'scheduler type {type(self.scheduler       )} not supported'       )

      # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
      lowercase							    =       1 / 0.1_8_2_1_5 * sample
      lowercase							    =       self.vae.decode(_lowerCamelCase       ).sample
      lowercase							    =       (image / 2 + 0.5).clamp(0							,		1       )

      lowercase							    =       transforms.Resize(self.feature_extractor_size       )(_lowerCamelCase       )
      lowercase							    =       self.normalize(_lowerCamelCase       ).to(latents.dtype       )

      lowercase							    =       self.clip_model.get_image_features(_lowerCamelCase       )
      lowercase							    =       image_embeddings_clip / image_embeddings_clip.norm(p=2							,		dim=-1							,		keepdim=_lowerCamelCase       )

      lowercase							    =       spherical_dist_loss(_lowerCamelCase							,		_lowerCamelCase       ).mean() * clip_guidance_scale

      lowercase							    =       -torch.autograd.grad(_lowerCamelCase							,		_lowerCamelCase       )[0]

      if isinstance(self.scheduler							,		_lowerCamelCase       ):
        lowercase							    =       latents.detach() + grads * (sigma**2)
        lowercase							    =       noise_pred_original
      else:
        lowercase							    =       noise_pred_original - torch.sqrt(_lowerCamelCase       ) * grads
      return noise_pred, latents







    @torch.no_grad()
    def __call__(			self							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase = None							,		_lowerCamelCase = None							,		_lowerCamelCase = 5_1_2							,		_lowerCamelCase = 5_1_2							,		_lowerCamelCase = 0.6							,		_lowerCamelCase = 5_0							,		_lowerCamelCase = 7.5							,		_lowerCamelCase = 1							,		_lowerCamelCase = 0.0							,		_lowerCamelCase = 1_0_0							,		_lowerCamelCase = None							,		_lowerCamelCase = "pil"							,		_lowerCamelCase = True							,		_lowerCamelCase = 0.8							,		_lowerCamelCase = 0.1							,		_lowerCamelCase = 0.1							,		):
      if isinstance(_lowerCamelCase							,		_lowerCamelCase       ) and len(_lowerCamelCase       ) != batch_size:
        raise ValueError(F'You have passed {batch_size} batch_size, but only {len(_lowerCamelCase       )} generators.'       )

      if height % 8 != 0 or width % 8 != 0:
        raise ValueError(F'`height` and `width` have to be divisible by 8 but are {height} and {width}.'       )

      if isinstance(_lowerCamelCase							,		torch.Generator       ) and batch_size > 1:
        lowercase							    =       [generator] + [None] * (batch_size - 1)

      lowercase							    =       [
          ('model', self.coca_model is None),
          ('tokenizer', self.coca_tokenizer is None),
          ('transform', self.coca_transform is None),
      ]
      lowercase							    =       [x[0] for x in coca_is_none if x[1]]
      lowercase							    =       ', '.join(_lowerCamelCase       )
      # generate prompts with coca model if prompt is None
      if content_prompt is None:
        if len(_lowerCamelCase       ):
          raise ValueError(
              F'Content prompt is None and CoCa [{coca_is_none_str}] is None.'
              F'Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'       )
        lowercase							    =       self.get_image_description(_lowerCamelCase       )
      if style_prompt is None:
        if len(_lowerCamelCase       ):
          raise ValueError(
              F'Style prompt is None and CoCa [{coca_is_none_str}] is None.'
              F' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'       )
        lowercase							    =       self.get_image_description(_lowerCamelCase       )

      # get prompt text embeddings for content and style
      lowercase							    =       self.tokenizer(
          _lowerCamelCase							,		padding='max_length'							,		max_length=self.tokenizer.model_max_length							,		truncation=_lowerCamelCase							,		return_tensors='pt'							,		)
      lowercase							    =       self.text_encoder(content_text_input.input_ids.to(self.device       )       )[0]

      lowercase							    =       self.tokenizer(
          _lowerCamelCase							,		padding='max_length'							,		max_length=self.tokenizer.model_max_length							,		truncation=_lowerCamelCase							,		return_tensors='pt'							,		)
      lowercase							    =       self.text_encoder(style_text_input.input_ids.to(self.device       )       )[0]

      lowercase							    =       slerp(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )

      # duplicate text embeddings for each generation per prompt
      lowercase							    =       text_embeddings.repeat_interleave(_lowerCamelCase							,		dim=0       )

      # set timesteps
      lowercase							    =       'offset' in set(inspect.signature(self.scheduler.set_timesteps       ).parameters.keys()       )
      lowercase							    =       {}
      if accepts_offset:
        lowercase							    =       1

      self.scheduler.set_timesteps(_lowerCamelCase							,		**_lowerCamelCase       )
      # Some schedulers like PNDM have timesteps as arrays
      # It's more optimized to move all timesteps to correct device beforehand
      self.scheduler.timesteps.to(self.device       )

      lowercase    ,     lowercase							    =       self.get_timesteps(_lowerCamelCase							,		_lowerCamelCase							,		self.device       )
      lowercase							    =       timesteps[:1].repeat(_lowerCamelCase       )

      # Preprocess image
      lowercase							    =       preprocess(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )
      lowercase							    =       self.prepare_latents(
          _lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		text_embeddings.dtype							,		self.device							,		_lowerCamelCase       )

      lowercase							    =       preprocess(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )
      lowercase							    =       self.prepare_latents(
          _lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		text_embeddings.dtype							,		self.device							,		_lowerCamelCase       )

      lowercase							    =       slerp(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )

      if clip_guidance_scale > 0:
        lowercase							    =       self.get_clip_image_embeddings(_lowerCamelCase							,		_lowerCamelCase       )
        lowercase							    =       self.get_clip_image_embeddings(_lowerCamelCase							,		_lowerCamelCase       )
        lowercase							    =       slerp(
            _lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase       )

      # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
      # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
      # corresponds to doing no classifier free guidance.
      lowercase							    =       guidance_scale > 1.0
      # get unconditional embeddings for classifier free guidance
      if do_classifier_free_guidance:
        lowercase							    =       content_text_input.input_ids.shape[-1]
        lowercase							    =       self.tokenizer(['']							,		padding='max_length'							,		max_length=_lowerCamelCase							,		return_tensors='pt'       )
        lowercase							    =       self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]
        # duplicate unconditional embeddings for each generation per prompt
        lowercase							    =       uncond_embeddings.repeat_interleave(_lowerCamelCase							,		dim=0       )

        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        lowercase							    =       torch.cat([uncond_embeddings, text_embeddings]       )

      # get the initial random noise unless the user supplied it

      # Unlike in other pipelines, latents need to be generated in the target device
      # for 1-to-1 results reproducibility with the CompVis implementation.
      # However this currently doesn't work in `mps`.
      lowercase							    =       (batch_size, self.unet.config.in_channels, height // 8, width // 8)
      lowercase							    =       text_embeddings.dtype
      if latents is None:
        if self.device.type == "mps":
          # randn does not work reproducibly on mps
          lowercase							    =       torch.randn(_lowerCamelCase							,		generator=_lowerCamelCase							,		device='cpu'							,		dtype=_lowerCamelCase       ).to(
              self.device       )
        else:
          lowercase							    =       torch.randn(_lowerCamelCase							,		generator=_lowerCamelCase							,		device=self.device							,		dtype=_lowerCamelCase       )
      else:
        if latents.shape != latents_shape:
          raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'       )
        lowercase							    =       latents.to(self.device       )

      # scale the initial noise by the standard deviation required by the scheduler
      lowercase							    =       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]
      lowercase							    =       'eta' in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
      lowercase							    =       {}
      if accepts_eta:
        lowercase							    =       eta

      # check if the scheduler accepts generator
      lowercase							    =       'generator' in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
      if accepts_generator:
        lowercase							    =       generator

      with self.progress_bar(total=_lowerCamelCase       ):
        for i, t in enumerate(_lowerCamelCase       ):
          # expand the latents if we are doing classifier free guidance
          lowercase							    =       torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents
          lowercase							    =       self.scheduler.scale_model_input(_lowerCamelCase							,		_lowerCamelCase       )

          # predict the noise residual
          lowercase							    =       self.unet(_lowerCamelCase							,		_lowerCamelCase							,		encoder_hidden_states=_lowerCamelCase       ).sample

          # perform classifier free guidance
          if do_classifier_free_guidance:
            lowercase    ,     lowercase							    =       noise_pred.chunk(2       )
            lowercase							    =       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

          # perform clip guidance
          if clip_guidance_scale > 0:
            lowercase							    =       (
                text_embeddings.chunk(2       )[1] if do_classifier_free_guidance else text_embeddings
            )
            lowercase    ,     lowercase							    =       self.cond_fn(
                _lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		)

          # compute the previous noisy sample x_t -> x_t-1
          lowercase							    =       self.scheduler.step(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		**_lowerCamelCase       ).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
      lowercase							    =       1 / 0.1_8_2_1_5 * latents
      lowercase							    =       self.vae.decode(_lowerCamelCase       ).sample

      lowercase							    =       (image / 2 + 0.5).clamp(0							,		1       )
      lowercase							    =       image.cpu().permute(0							,		2							,		3							,		1       ).numpy()

      if output_type == "pil":
        lowercase							    =       self.numpy_to_pil(_lowerCamelCase       )

      if not return_dict:
        return (image, None)

      return StableDiffusionPipelineOutput(images=_lowerCamelCase							,		nsfw_content_detected=_lowerCamelCase       )










"""simple docstring"""



import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
    import jax
    import jax.numpy as jnp
    from flax.jax_utils import replicate
    from flax.training.common_utils import shard

@slow
@require_flax
class        a     (   unittest.TestCase	):







    def 			UpperCamelCase_      (			self       ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()







    def 			UpperCamelCase_      (			self       ):
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          'stabilityai/stable-diffusion-2'							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

      lowercase							    =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
      lowercase							    =       images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       'stabilityai/stable-diffusion-2'
      lowercase    ,     lowercase							    =       FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCamelCase							,		subfolder='scheduler'       )
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          _lowerCamelCase							,		scheduler=_lowerCamelCase							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)
      lowercase							    =       scheduler_params

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       sd_pipe(_lowerCamelCase							,		_lowerCamelCase							,		_lowerCamelCase							,		num_inference_steps=2_5							,		jit=_lowerCamelCase       )[0]
      assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

      lowercase							    =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]       )
      lowercase							    =       images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2


