





import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu





class  lowerCAmelCase							(       unittest.TestCase ):







  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







  @property
  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     lowercase__			       =  1
     lowercase__			       =  3
     lowercase__			       =  (32, 32)

     lowercase__			       =  floats_tensor((batch_size, num_channels) + sizes   ,							rng=random.Random(0 ) ).to(_lowercase )
     return image







  @property
  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     torch.manual_seed(0 )
     lowercase__			       =  UNetaDConditionModel(
         block_out_channels=(32, 64)   ,							layers_per_block=2   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,							up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,							cross_attention_dim=32   ,							)
     return model







  @property
  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     torch.manual_seed(0 )
     lowercase__			       =  AutoencoderKL(
         block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]   ,							up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]   ,							latent_channels=4   ,							)
     return model







  @property
  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     torch.manual_seed(0 )
     lowercase__			       =  CLIPTextConfig(
         bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1e-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=10_00   ,							)
     return CLIPTextModel(_lowercase )







  @property
  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     def extract(*_lowercase  :Optional[Any]   ,							**_lowercase  :Tuple ):





        class  lowerCAmelCase							:







          def __init__(      self  :Union[str, Any] ):






             '''simple docstring'''





             lowercase__			       =  torch.ones([0] )







          def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Optional[int] ):






             '''simple docstring'''





             self.pixel_values.to(_lowercase )
             return self
        return Out()

     return extract







  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  "cpu"  # ensure determinism for the device-dependent torch.Generator
     lowercase__			       =  self.dummy_cond_unet
     lowercase__			       =  DDIMScheduler(
         beta_start=0.00085   ,							beta_end=0.012   ,							beta_schedule="scaled_linear"   ,							clip_sample=_lowercase   ,							set_alpha_to_one=_lowercase   ,							)

     lowercase__			       =  self.dummy_vae
     lowercase__			       =  self.dummy_text_encoder
     lowercase__			       =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

     # make sure here that pndm scheduler skips prk
     lowercase__			       =  StableDiffusionPipeline(
         unet=_lowercase   ,							scheduler=_lowercase   ,							vae=_lowercase   ,							text_encoder=_lowercase   ,							tokenizer=_lowercase   ,							safety_checker=_lowercase   ,							feature_extractor=self.dummy_extractor   ,							)
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  "A painting of a squirrel eating a burger"

     lowercase__			       =  torch.Generator(device=_lowercase ).manual_seed(0 )
     lowercase__			       =  sd_pipe([prompt]   ,							generator=_lowercase   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type="np" )
     lowercase__			       =  output.images

     lowercase__			       =  torch.Generator(device=_lowercase ).manual_seed(0 )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type="np"   ,							return_dict=_lowercase   ,							)[0]

     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)
     lowercase__			       =  np.array([0.5756, 0.6118, 0.5005, 0.5041, 0.5471, 0.4726, 0.4976, 0.4865, 0.4864] )

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1e-2







  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





     lowercase__			       =  "cpu"  # ensure determinism for the device-dependent torch.Generator
     lowercase__			       =  self.dummy_cond_unet
     lowercase__			       =  PNDMScheduler(skip_prk_steps=_lowercase )
     lowercase__			       =  self.dummy_vae
     lowercase__			       =  self.dummy_text_encoder
     lowercase__			       =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

     # make sure here that pndm scheduler skips prk
     lowercase__			       =  StableDiffusionPipeline(
         unet=_lowercase   ,							scheduler=_lowercase   ,							vae=_lowercase   ,							text_encoder=_lowercase   ,							tokenizer=_lowercase   ,							safety_checker=_lowercase   ,							feature_extractor=self.dummy_extractor   ,							)
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  "A painting of a squirrel eating a burger"
     lowercase__			       =  torch.Generator(device=_lowercase ).manual_seed(0 )
     lowercase__			       =  sd_pipe([prompt]   ,							generator=_lowercase   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type="np" )

     lowercase__			       =  output.images

     lowercase__			       =  torch.Generator(device=_lowercase ).manual_seed(0 )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=6.0   ,							num_inference_steps=2   ,							output_type="np"   ,							return_dict=_lowercase   ,							)[0]

     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)
     lowercase__			       =  np.array([0.5125, 0.5716, 0.4828, 0.5060, 0.5650, 0.4768, 0.5185, 0.4895, 0.4993] )

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1e-2







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  StableDiffusionPipeline.from_pretrained(
         "hf-internal-testing/tiny-stable-diffusion-lms-pipe"   ,							safety_checker=_lowercase )
     assert isinstance(_lowercase   ,							_lowercase )
     assert isinstance(pipe.scheduler   ,							_lowercase )
     assert pipe.safety_checker is None

     lowercase__			       =  pipe("example prompt"   ,							num_inference_steps=2 ).images[0]
     assert image is not None

     # check that there's no error when saving a pipeline with one of the models being None
     with tempfile.TemporaryDirectory() as tmpdirname:
        pipe.save_pretrained(_lowercase )
        lowercase__			       =  StableDiffusionPipeline.from_pretrained(_lowercase )

     # sanity check that the pipeline still works
     assert pipe.safety_checker is None
     lowercase__			       =  pipe("example prompt"   ,							num_inference_steps=2 ).images[0]
     assert image is not None







  @unittest.skipIf(torch_device != "cuda"   ,							"This test requires a GPU" )
  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  self.dummy_cond_unet
     lowercase__			       =  PNDMScheduler(skip_prk_steps=_lowercase )
     lowercase__			       =  self.dummy_vae
     lowercase__			       =  self.dummy_text_encoder
     lowercase__			       =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

     # put models in fp16
     lowercase__			       =  unet.half()
     lowercase__			       =  vae.half()
     lowercase__			       =  bert.half()

     # make sure here that pndm scheduler skips prk
     lowercase__			       =  StableDiffusionPipeline(
         unet=_lowercase   ,							scheduler=_lowercase   ,							vae=_lowercase   ,							text_encoder=_lowercase   ,							tokenizer=_lowercase   ,							safety_checker=_lowercase   ,							feature_extractor=self.dummy_extractor   ,							)
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  "A painting of a squirrel eating a burger"
     lowercase__			       =  sd_pipe([prompt]   ,							num_inference_steps=2   ,							output_type="np" ).images

     assert image.shape == (1, 64, 64, 3)





@nightly
@require_torch_gpu
class  lowerCAmelCase							(       unittest.TestCase ):







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"   ,							safety_checker=_lowercase )
     lowercase__			       =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config )
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  (
         "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
         " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
         " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
         " children from bahnhof zoo, detailed "
     )
     lowercase__			       =  40_03_66_03_46
     lowercase__			       =  7

     # without safety guidance (sld_guidance_scale = 0)
     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=0   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  [0.2278, 0.2231, 0.2249, 0.2333, 0.2303, 0.1885, 0.2273, 0.2144, 0.2176]

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2

     # without safety guidance (strong configuration)
     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=20_00   ,							sld_warmup_steps=7   ,							sld_threshold=0.025   ,							sld_momentum_scale=0.5   ,							sld_mom_beta=0.7   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  [0.2383, 0.2276, 0.236, 0.2192, 0.2186, 0.2053, 0.1971, 0.1901, 0.1719]

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2







  def   UpperCAmelCase					(      self  :Union[str, Any] ):






     '''simple docstring'''





     lowercase__			       =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"   ,							safety_checker=_lowercase )
     lowercase__			       =  LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config )
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  "padme amidala taking a bath artwork, safe for work, no nudity"
     lowercase__			       =  27_34_97_17_55
     lowercase__			       =  7

     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=0   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  [0.3502, 0.3622, 0.3396, 0.3642, 0.3478, 0.3318, 0.35, 0.3348, 0.3297]

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2

     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=20_00   ,							sld_warmup_steps=7   ,							sld_threshold=0.025   ,							sld_momentum_scale=0.5   ,							sld_mom_beta=0.7   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  [0.5531, 0.5206, 0.4895, 0.5156, 0.5182, 0.4751, 0.4802, 0.4803, 0.4443]

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2







  def   UpperCAmelCase					(      self  :str ):






     '''simple docstring'''





     lowercase__			       =  StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5" )
     lowercase__			       =  sd_pipe.to(_lowercase )
     sd_pipe.set_progress_bar_config(disable=_lowercase )

     lowercase__			       =  (
         "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
         " leyendecker"
     )
     lowercase__			       =  10_44_35_52_34
     lowercase__			       =  12

     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=0   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] )

     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-7

     lowercase__			       =  torch.manual_seed(_lowercase )
     lowercase__			       =  sd_pipe(
         [prompt]   ,							generator=_lowercase   ,							guidance_scale=_lowercase   ,							num_inference_steps=50   ,							output_type="np"   ,							width=5_12   ,							height=5_12   ,							sld_guidance_scale=20_00   ,							sld_warmup_steps=7   ,							sld_threshold=0.025   ,							sld_momentum_scale=0.5   ,							sld_mom_beta=0.7   ,							)

     lowercase__			       =  output.images
     lowercase__			       =  image[0, -3:, -3:, -1]
     lowercase__			       =  np.array([0.5818, 0.6285, 0.6835, 0.6019, 0.625, 0.6754, 0.6096, 0.6334, 0.6561] )
     assert image.shape == (1, 5_12, 5_12, 3)

     assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2













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

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils import SchedulerMixin, SchedulerOutput





class  lowerCAmelCase							(       lowercase_							,							lowercase_ ):
  __lowerCamelCase		=      1







  @register_to_config
  def __init__(      self  :Dict   ,							_lowercase  :int = 10_00   ,							_lowercase  :Optional[Union[np.ndarray, List[float]]] = None ):






     '''simple docstring'''





     self.set_timesteps(_lowercase )

     # standard deviation of the initial noise distribution
     lowercase__			       =  1.0

     # For now we only support F-PNDM, i.e. the runge-kutta method
     # For more information on the algorithm please take a look at the paper: https://arxiv.org/pdf/2202.09778.pdf
     # mainly at formula (9), (12), (13) and the Algorithm 2.
     lowercase__			       =  4

     # running values
     lowercase__			       =  []







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :int   ,							_lowercase  :Union[str, torch.device] = None ):






     '''simple docstring'''





     lowercase__			       =  num_inference_steps
     lowercase__			       =  torch.linspace(1   ,							0   ,							num_inference_steps + 1 )[:-1]
     lowercase__			       =  torch.cat([steps, torch.tensor([0.0] )] )

     if self.config.trained_betas is not None:
        lowercase__			       =  torch.tensor(self.config.trained_betas   ,							dtype=torch.floataa )
     else:
        lowercase__			       =  torch.sin(steps * math.pi / 2 ) ** 2

     lowercase__			       =  (1.0 - self.betas**2) ** 0.5

     lowercase__			       =  (torch.atana(self.betas   ,							self.alphas ) / math.pi * 2)[:-1]
     lowercase__			       =  timesteps.to(_lowercase )

     lowercase__			       =  []







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :torch.FloatTensor   ,							_lowercase  :int   ,							_lowercase  :torch.FloatTensor   ,							_lowercase  :bool = True   ,							):






     '''simple docstring'''





     if self.num_inference_steps is None:
        raise ValueError(
            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" )

     lowercase__			       =  (self.timesteps == timestep).nonzero().item()
     lowercase__			       =  timestep_index + 1

     lowercase__			       =  sample * self.betas[timestep_index] + model_output * self.alphas[timestep_index]
     self.ets.append(_lowercase )

     if len(self.ets ) == 1:
        lowercase__			       =  self.ets[-1]
     elif len(self.ets ) == 2:
        lowercase__			       =  (3 * self.ets[-1] - self.ets[-2]) / 2
     elif len(self.ets ) == 3:
        lowercase__			       =  (23 * self.ets[-1] - 16 * self.ets[-2] + 5 * self.ets[-3]) / 12
     else:
        lowercase__			       =  (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4])

     lowercase__			       =  self._get_prev_sample(_lowercase   ,							_lowercase   ,							_lowercase   ,							_lowercase )

     if not return_dict:
        return (prev_sample,)

     return SchedulerOutput(prev_sample=_lowercase )







  def   UpperCAmelCase					(      self  :Any   ,							_lowercase  :torch.FloatTensor   ,							*_lowercase  :int   ,							**_lowercase  :int ):






     '''simple docstring'''





     return sample







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :Tuple   ,							_lowercase  :int   ,							_lowercase  :Optional[Any]   ,							_lowercase  :List[str] ):






     '''simple docstring'''





     lowercase__			       =  self.alphas[timestep_index]
     lowercase__			       =  self.betas[timestep_index]

     lowercase__			       =  self.alphas[prev_timestep_index]
     lowercase__			       =  self.betas[prev_timestep_index]

     lowercase__			       =  (sample - sigma * ets) / max(_lowercase   ,							1e-8 )
     lowercase__			       =  next_alpha * pred + ets * next_sigma

     return prev_sample







  def __len__(      self  :Union[str, Any] ):






     '''simple docstring'''





     return self.config.num_train_timesteps






