
'''simple docstring'''

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.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 	__lowerCAmelCase    (  unittest.TestCase				):



     def     snake_case_ (self	):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



     @property
     def     snake_case_ (self	):
          _UpperCAmelCase					:  Optional[Any]         =    1
          _UpperCAmelCase					:  Union[str, Any]         =    3
          _UpperCAmelCase					:  List[str]         =    (3_2, 3_2)

          _UpperCAmelCase					:  Any         =    floats_tensor((batch_size, num_channels) + sizes						,	rng=random.Random(0	)	).to(lowerCAmelCase__	)
          return image



     @property
     def     snake_case_ (self	):
          torch.manual_seed(0	)
          _UpperCAmelCase					:  Optional[Any]         =    UNetaDConditionModel(
              block_out_channels=(3_2, 6_4)						,	layers_per_block=2						,	sample_size=3_2						,	in_channels=4						,	out_channels=4						,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")						,	up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")						,	cross_attention_dim=3_2						,	)
          return model



     @property
     def     snake_case_ (self	):
          torch.manual_seed(0	)
          _UpperCAmelCase					:  int         =    AutoencoderKL(
              block_out_channels=[3_2, 6_4]						,	in_channels=3						,	out_channels=3						,	down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]						,	up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]						,	latent_channels=4						,	)
          return model



     @property
     def     snake_case_ (self	):
          torch.manual_seed(0	)
          _UpperCAmelCase					:  Optional[int]         =    CLIPTextConfig(
              bos_token_id=0						,	eos_token_id=2						,	hidden_size=3_2						,	intermediate_size=3_7						,	layer_norm_eps=1e-05						,	num_attention_heads=4						,	num_hidden_layers=5						,	pad_token_id=1						,	vocab_size=1_0_0_0						,	)
          return CLIPTextModel(lowerCAmelCase__	)



     @property
     def     snake_case_ (self	):
          def extract(*lowerCAmelCase__						,	**lowerCAmelCase__	):







               class 	__lowerCAmelCase    :



                    def __init__(self	):
                         _UpperCAmelCase					:  Tuple         =    torch.ones([0]	)



                    def     snake_case_ (self						,	lowerCAmelCase__	):
                         self.pixel_values.to(lowerCAmelCase__	)
                         return self




               return Out()

          return extract



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase					:  Tuple         =    self.dummy_cond_unet
          _UpperCAmelCase					:  List[str]         =    DDIMScheduler(
              beta_start=0.0_0_0_8_5						,	beta_end=0.0_1_2						,	beta_schedule="""scaled_linear"""						,	clip_sample=lowerCAmelCase__						,	set_alpha_to_one=lowerCAmelCase__						,	)

          _UpperCAmelCase					:  Dict         =    self.dummy_vae
          _UpperCAmelCase					:  Tuple         =    self.dummy_text_encoder
          _UpperCAmelCase					:  Optional[int]         =    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

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

          _UpperCAmelCase					:  Dict         =    """A painting of a squirrel eating a burger"""

          _UpperCAmelCase					:  List[str]         =    torch.Generator(device=lowerCAmelCase__	).manual_seed(0	)
          _UpperCAmelCase					:  List[Any]         =    sd_pipe([prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="""np"""	)
          _UpperCAmelCase					:  Optional[Any]         =    output.images

          _UpperCAmelCase					:  Dict         =    torch.Generator(device=lowerCAmelCase__	).manual_seed(0	)
          _UpperCAmelCase					:  Any         =    sd_pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="""np"""						,	return_dict=lowerCAmelCase__						,	)[0]

          _UpperCAmelCase					:  Union[str, Any]         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  List[str]         =    image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 6_4, 6_4, 3)
          _UpperCAmelCase					:  Union[str, Any]         =    np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7]	)

          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     snake_case_ (self	):
          _UpperCAmelCase					:  Any         =    """cpu"""  # ensure determinism for the device-dependent torch.Generator
          _UpperCAmelCase					:  List[Any]         =    self.dummy_cond_unet
          _UpperCAmelCase					:  Optional[int]         =    PNDMScheduler(skip_prk_steps=lowerCAmelCase__	)
          _UpperCAmelCase					:  Tuple         =    self.dummy_vae
          _UpperCAmelCase					:  Union[str, Any]         =    self.dummy_text_encoder
          _UpperCAmelCase					:  Dict         =    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

          # make sure here that pndm scheduler skips prk
          _UpperCAmelCase					:  int         =    StableDiffusionPipeline(
              unet=lowerCAmelCase__						,	scheduler=lowerCAmelCase__						,	vae=lowerCAmelCase__						,	text_encoder=lowerCAmelCase__						,	tokenizer=lowerCAmelCase__						,	safety_checker=lowerCAmelCase__						,	feature_extractor=self.dummy_extractor						,	)
          _UpperCAmelCase					:  Tuple         =    sd_pipe.to(lowerCAmelCase__	)
          sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  Optional[int]         =    """A painting of a squirrel eating a burger"""
          _UpperCAmelCase					:  List[Any]         =    torch.Generator(device=lowerCAmelCase__	).manual_seed(0	)
          _UpperCAmelCase					:  Tuple         =    sd_pipe([prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="""np"""	)

          _UpperCAmelCase					:  List[Any]         =    output.images

          _UpperCAmelCase					:  Dict         =    torch.Generator(device=lowerCAmelCase__	).manual_seed(0	)
          _UpperCAmelCase					:  Optional[Any]         =    sd_pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=6.0						,	num_inference_steps=2						,	output_type="""np"""						,	return_dict=lowerCAmelCase__						,	)[0]

          _UpperCAmelCase					:  Any         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  int         =    image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 6_4, 6_4, 3)
          _UpperCAmelCase					:  Dict         =    np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4]	)

          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     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    StableDiffusionPipeline.from_pretrained(
              """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""						,	safety_checker=lowerCAmelCase__	)
          assert isinstance(lowerCAmelCase__						,	lowerCAmelCase__	)
          assert isinstance(pipe.scheduler						,	lowerCAmelCase__	)
          assert pipe.safety_checker is None

          _UpperCAmelCase					:  Optional[Any]         =    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(lowerCAmelCase__	)
               _UpperCAmelCase					:  List[str]         =    StableDiffusionPipeline.from_pretrained(lowerCAmelCase__	)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          _UpperCAmelCase					:  Tuple         =    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     snake_case_ (self	):
          _UpperCAmelCase					:  Any         =    self.dummy_cond_unet
          _UpperCAmelCase					:  Union[str, Any]         =    PNDMScheduler(skip_prk_steps=lowerCAmelCase__	)
          _UpperCAmelCase					:  Optional[Any]         =    self.dummy_vae
          _UpperCAmelCase					:  str         =    self.dummy_text_encoder
          _UpperCAmelCase					:  Tuple         =    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

          # put models in fp16
          _UpperCAmelCase					:  Optional[int]         =    unet.half()
          _UpperCAmelCase					:  Tuple         =    vae.half()
          _UpperCAmelCase					:  List[str]         =    bert.half()

          # make sure here that pndm scheduler skips prk
          _UpperCAmelCase					:  Any         =    StableDiffusionPipeline(
              unet=lowerCAmelCase__						,	scheduler=lowerCAmelCase__						,	vae=lowerCAmelCase__						,	text_encoder=lowerCAmelCase__						,	tokenizer=lowerCAmelCase__						,	safety_checker=lowerCAmelCase__						,	feature_extractor=self.dummy_extractor						,	)
          _UpperCAmelCase					:  Optional[Any]         =    sd_pipe.to(lowerCAmelCase__	)
          sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  List[Any]         =    """A painting of a squirrel eating a burger"""
          _UpperCAmelCase					:  str         =    sd_pipe([prompt]						,	num_inference_steps=2						,	output_type="""np"""	).images

          assert image.shape == (1, 6_4, 6_4, 3)







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



     def     snake_case_ (self	):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



     def     snake_case_ (self	):
          _UpperCAmelCase					:  int         =    """cuda"""
          _UpperCAmelCase					:  Tuple         =    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
          _UpperCAmelCase					:  Optional[Any]         =    pipe.to(lowerCAmelCase__	)
          pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  Dict         =    """a photo of a cat"""
          _UpperCAmelCase					:  Optional[int]         =    {
              """editing_prompt""": ["""sunglasses"""],
              """reverse_editing_direction""": [False],
              """edit_warmup_steps""": 1_0,
              """edit_guidance_scale""": 6,
              """edit_threshold""": 0.9_5,
              """edit_momentum_scale""": 0.5,
              """edit_mom_beta""": 0.6,
          }

          _UpperCAmelCase					:  Optional[int]         =    3
          _UpperCAmelCase					:  Optional[int]         =    7

          # no sega enabled
          _UpperCAmelCase					:  int         =    torch.Generator(lowerCAmelCase__	)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Tuple         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	)

          _UpperCAmelCase					:  Tuple         =    output.images
          _UpperCAmelCase					:  str         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  List[Any]         =    [
              0.3_4_6_7_3_1_1_3,
              0.3_8_4_9_2_7_3_3,
              0.3_7_5_9_7_3_5_2,
              0.3_4_0_8_6_3_3_5,
              0.3_5_6_5_0_7_4_8,
              0.3_5_5_7_9_2_0_5,
              0.3_3_8_4_7_6_3,
              0.3_4_3_4_0_2_3_6,
              0.3_5_7_3_2_7_1,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Tuple         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	**lowerCAmelCase__						,	)

          _UpperCAmelCase					:  Union[str, Any]         =    output.images
          _UpperCAmelCase					:  Dict         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  Dict         =    [
              0.4_1_8_8_7_8_2_6,
              0.3_7_7_2_8_7_6_6,
              0.3_0_1_3_8_2_7_2,
              0.4_1_4_1_6_3_3_5,
              0.4_1_6_6_4_9_8_5,
              0.3_6_2_8_3_3_9_2,
              0.3_6_1_9_1_2_4_6,
              0.4_3_3_6_4_4_6_5,
              0.4_3_0_0_1_7_3_2,
          ]

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

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



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Dict         =    """cuda"""
          _UpperCAmelCase					:  List[Any]         =    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
          _UpperCAmelCase					:  str         =    pipe.to(lowerCAmelCase__	)
          pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  List[str]         =    """an image of a crowded boulevard, realistic, 4k"""
          _UpperCAmelCase					:  Union[str, Any]         =    {
              """editing_prompt""": """crowd, crowded, people""",
              """reverse_editing_direction""": True,
              """edit_warmup_steps""": 1_0,
              """edit_guidance_scale""": 8.3,
              """edit_threshold""": 0.9,
              """edit_momentum_scale""": 0.5,
              """edit_mom_beta""": 0.6,
          }

          _UpperCAmelCase					:  Optional[Any]         =    9
          _UpperCAmelCase					:  List[Any]         =    7

          # no sega enabled
          _UpperCAmelCase					:  Optional[int]         =    torch.Generator(lowerCAmelCase__	)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Dict         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	)

          _UpperCAmelCase					:  int         =    output.images
          _UpperCAmelCase					:  List[Any]         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  List[str]         =    [
              0.4_3_4_9_7_9_9_8,
              0.9_1_8_1_4_0_6_5,
              0.7_5_4_0_7_3_9,
              0.5_5_5_8_0_2_0_5,
              0.8_4_6_7_2_6_5,
              0.5_3_8_9_6_9_1,
              0.6_2_5_7_4_5_0_6,
              0.5_8_8_9_7_7_6_3,
              0.5_0_9_2_6_7_5_7,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Union[str, Any]         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	**lowerCAmelCase__						,	)

          _UpperCAmelCase					:  Union[str, Any]         =    output.images
          _UpperCAmelCase					:  Any         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  Optional[Any]         =    [
              0.3_0_8_9_7_1_9,
              0.3_0_5_0_0_1_4_4,
              0.2_9_0_1_6_0_4_2,
              0.3_0_6_3_0_9_6_4,
              0.3_2_5_6_8_7,
              0.2_9_4_1_9_2_2_5,
              0.2_9_0_8_0_9_1,
              0.2_8_7_2_3_5_9_8,
              0.2_7_6_9_6_2_9_4,
          ]

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

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



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Optional[int]         =    """cuda"""
          _UpperCAmelCase					:  Any         =    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
          _UpperCAmelCase					:  Any         =    pipe.to(lowerCAmelCase__	)
          pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  Any         =    """a castle next to a river"""
          _UpperCAmelCase					:  List[Any]         =    {
              """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
              """reverse_editing_direction""": False,
              """edit_warmup_steps""": [1_5, 1_8],
              """edit_guidance_scale""": 6,
              """edit_threshold""": [0.9, 0.8],
              """edit_momentum_scale""": 0.5,
              """edit_mom_beta""": 0.6,
          }

          _UpperCAmelCase					:  List[Any]         =    4_8
          _UpperCAmelCase					:  int         =    7

          # no sega enabled
          _UpperCAmelCase					:  List[Any]         =    torch.Generator(lowerCAmelCase__	)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Dict         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	)

          _UpperCAmelCase					:  int         =    output.images
          _UpperCAmelCase					:  List[Any]         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  List[str]         =    [
              0.7_5_1_6_3_5_5_5,
              0.7_6_0_3_7_1_4_5,
              0.6_1_7_8_5,
              0.9_1_8_9_6_7_3,
              0.8_6_2_7_7_0_1,
              0.8_5_1_8_9_6_9_4,
              0.8_5_1_2_8_1_3,
              0.8_7_0_1_2_0_7_6,
              0.8_3_1_2_8_5_7,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Optional[Any]         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	**lowerCAmelCase__						,	)

          _UpperCAmelCase					:  Optional[int]         =    output.images
          _UpperCAmelCase					:  int         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  int         =    [
              0.7_3_5_5_3_3_6_5,
              0.7_5_3_7_2_7_1,
              0.7_4_3_4_1_9_0_5,
              0.6_6_4_8_0_3_5_6,
              0.6_4_7_2_9_2_5,
              0.6_3_0_3_9_4_1_6,
              0.6_4_8_1_2_9_0_5,
              0.6_7_4_9_7_1_7,
              0.6_5_1_7_1_0_2,
          ]

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

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



     def     snake_case_ (self	):
          _UpperCAmelCase					:  List[Any]         =    """cuda"""
          _UpperCAmelCase					:  List[str]         =    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						,	torch_dtype=torch.floataa	)
          _UpperCAmelCase					:  Optional[Any]         =    pipe.to(lowerCAmelCase__	)
          pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

          _UpperCAmelCase					:  Optional[int]         =    """a photo of a cat"""
          _UpperCAmelCase					:  Tuple         =    {
              """editing_prompt""": ["""sunglasses"""],
              """reverse_editing_direction""": [False],
              """edit_warmup_steps""": 1_0,
              """edit_guidance_scale""": 6,
              """edit_threshold""": 0.9_5,
              """edit_momentum_scale""": 0.5,
              """edit_mom_beta""": 0.6,
          }

          _UpperCAmelCase					:  Union[str, Any]         =    3
          _UpperCAmelCase					:  str         =    7

          # no sega enabled
          _UpperCAmelCase					:  Optional[int]         =    torch.Generator(lowerCAmelCase__	)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  Union[str, Any]         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	)

          _UpperCAmelCase					:  Union[str, Any]         =    output.images
          _UpperCAmelCase					:  List[str]         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  Any         =    [
              0.3_4_8_8_7_6_9_5,
              0.3_8_7_6_9_5_3,
              0.3_7_5,
              0.3_4_4_2_3_8_2_8,
              0.3_5_8_1_5_4_3,
              0.3_5_7_1_7_7_7_3,
              0.3_3_8_3_7_8_9,
              0.3_4_5_7_0_3_1_2,
              0.3_5_9_3_7_5,
          ]

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

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

          # with sega enabled
          # generator = torch.manual_seed(seed)
          generator.manual_seed(lowerCAmelCase__	)
          _UpperCAmelCase					:  int         =    pipe(
              [prompt]						,	generator=lowerCAmelCase__						,	guidance_scale=lowerCAmelCase__						,	num_inference_steps=5_0						,	output_type="""np"""						,	width=5_1_2						,	height=5_1_2						,	**lowerCAmelCase__						,	)

          _UpperCAmelCase					:  Any         =    output.images
          _UpperCAmelCase					:  Optional[int]         =    image[0, -3:, -3:, -1]
          _UpperCAmelCase					:  Optional[int]         =    [
              0.4_2_2_8_5_1_5_6,
              0.3_6_9_1_4_0_6_2,
              0.2_9_0_7_7_1_4_8,
              0.4_2_0_4_1_0_1_6,
              0.4_1_9_1_8_9_4_5,
              0.3_5_4_9_8_0_4_7,
              0.3_6_1_8_1_6_4,
              0.4_4_2_3_8_2_8,
              0.4_3_1_1_5_2_3_4,
          ]

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

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




'''simple docstring'''

import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class 	__lowerCAmelCase    (  __a				):
     snake_case   :    List[str]							      =      (HeunDiscreteScheduler,)
     snake_case   :    Optional[Any]							      =      1_0



     def     snake_case_ (self						,	**lowerCAmelCase__	):
          _UpperCAmelCase					:  Union[str, Any]         =    {
              """num_train_timesteps""": 1_1_0_0,
              """beta_start""": 0.0_0_0_1,
              """beta_end""": 0.0_2,
              """beta_schedule""": """linear""",
          }

          config.update(**lowerCAmelCase__	)
          return config



     def     snake_case_ (self	):
          for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
               self.check_over_configs(num_train_timesteps=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]						,	[0.0_0_0_2, 0.0_0_2, 0.0_2]	):
               self.check_over_configs(beta_start=lowerCAmelCase__						,	beta_end=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for schedule in ["linear", "scaled_linear", "exp"]:
               self.check_over_configs(beta_schedule=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for clip_sample_range in [1.0, 2.0, 3.0]:
               self.check_over_configs(clip_sample_range=lowerCAmelCase__						,	clip_sample=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for prediction_type in ["epsilon", "v_prediction", "sample"]:
               self.check_over_configs(prediction_type=lowerCAmelCase__	)



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Optional[int]         =    self.get_scheduler_config()
          _UpperCAmelCase					:  List[str]         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps	)

          _UpperCAmelCase					:  str         =    self.dummy_model()
          _UpperCAmelCase					:  List[Any]         =    self.dummy_sample_deter * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Tuple         =    sample.to(lowerCAmelCase__	)

          for i, t in enumerate(scheduler.timesteps	):
               _UpperCAmelCase					:  Union[str, Any]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Any         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Tuple         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Optional[Any]         =    output.prev_sample

          _UpperCAmelCase					:  Optional[Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  int         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3
          else:
               # CUDA
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Optional[int]         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Any         =    self.get_scheduler_config(prediction_type="""v_prediction"""	)
          _UpperCAmelCase					:  Any         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps	)

          _UpperCAmelCase					:  Optional[Any]         =    self.dummy_model()
          _UpperCAmelCase					:  Any         =    self.dummy_sample_deter * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Any         =    sample.to(lowerCAmelCase__	)

          for i, t in enumerate(scheduler.timesteps	):
               _UpperCAmelCase					:  List[str]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  List[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Union[str, Any]         =    output.prev_sample

          _UpperCAmelCase					:  Union[str, Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  str         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 4.6_934e-07	) < 1e-2
               assert abs(result_mean.item() - 6.1_112e-10	) < 1e-3
          else:
               # CUDA
               assert abs(result_sum.item() - 4.693_428_650_170_972e-07	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Any         =    self.get_scheduler_config()
          _UpperCAmelCase					:  Optional[int]         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps						,	device=lowerCAmelCase__	)

          _UpperCAmelCase					:  List[Any]         =    self.dummy_model()
          _UpperCAmelCase					:  Dict         =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
               _UpperCAmelCase					:  Union[str, Any]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[int]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  int         =    output.prev_sample

          _UpperCAmelCase					:  Any         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  Optional[Any]         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if str(lowerCAmelCase__	).startswith("""cpu"""	):
               # The following sum varies between 148 and 156 on mps. Why?
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3
          elif str(lowerCAmelCase__	).startswith("""mps"""	):
               # Larger tolerance on mps
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-2
          else:
               # CUDA
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  int         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  str         =    self.get_scheduler_config()
          _UpperCAmelCase					:  Union[str, Any]         =    scheduler_class(**lowerCAmelCase__						,	use_karras_sigmas=lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps						,	device=lowerCAmelCase__	)

          _UpperCAmelCase					:  Optional[int]         =    self.dummy_model()
          _UpperCAmelCase					:  Optional[Any]         =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Union[str, Any]         =    sample.to(lowerCAmelCase__	)

          for t in scheduler.timesteps:
               _UpperCAmelCase					:  Dict         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[int]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Optional[Any]         =    output.prev_sample

          _UpperCAmelCase					:  Union[str, Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  List[Any]         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          assert abs(result_sum.item() - 0.0_0_0_1_5	) < 1e-2
          assert abs(result_mean.item() - 1.9_869_554_535_034_695e-07	) < 1e-2


