

import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class SCREAMING_SNAKE_CASE	(  a_   ,							a_   ,							unittest.TestCase ):
  """simple docstring"""






  lowerCamelCase    :			List[Any]						      =StableDiffusionParadigmsPipeline
  lowerCamelCase    :			Tuple						      =TEXT_TO_IMAGE_PARAMS
  lowerCamelCase    :			List[Any]						      =TEXT_TO_IMAGE_BATCH_PARAMS
  lowerCamelCase    :			Union[str, Any]						      =TEXT_TO_IMAGE_IMAGE_PARAMS
  lowerCamelCase    :			Optional[Any]						      =TEXT_TO_IMAGE_IMAGE_PARAMS







  def        SCREAMING_SNAKE_CASE    (			self :    Any    )							->	str:






     """simple docstring"""







     torch.manual_seed(0    )
     __lowerCAmelCase  :			Optional[int]					    =						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       ,							attention_head_dim=(2, 4)       ,							use_linear_projection=lowerCAmelCase       ,							)
     __lowerCAmelCase  :			Union[str, Any]					    =						DDIMParallelScheduler(
         beta_start=0.0_0085       ,							beta_end=0.012       ,							beta_schedule="""scaled_linear"""       ,							clip_sample=lowerCAmelCase       ,							set_alpha_to_one=lowerCAmelCase       ,							)
     torch.manual_seed(0    )
     __lowerCAmelCase  :			Optional[int]					    =						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       ,							sample_size=1_28       ,							)
     torch.manual_seed(0    )
     __lowerCAmelCase  :			List[Any]					    =						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       ,							hidden_act="""gelu"""       ,							projection_dim=5_12       ,							)
     __lowerCAmelCase  :			Dict					    =						CLIPTextModel(lowerCAmelCase    )
     __lowerCAmelCase  :			Optional[int]					    =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

     __lowerCAmelCase  :			List[Any]					    =						{
         """unet""": unet,
         """scheduler""": scheduler,
         """vae""": vae,
         """text_encoder""": text_encoder,
         """tokenizer""": tokenizer,
         """safety_checker""": None,
         """feature_extractor""": None,
     }
     return components







  def        SCREAMING_SNAKE_CASE    (			self :    List[str]       ,							lowerCAmelCase :    List[str]       ,							lowerCAmelCase :    str=0    )							->	str:






     """simple docstring"""







     if str(lowerCAmelCase    ).startswith("""mps"""    ):
        __lowerCAmelCase  :			List[Any]					    =						torch.manual_seed(lowerCAmelCase    )
     else:
        __lowerCAmelCase  :			Optional[int]					    =						torch.Generator(device=lowerCAmelCase    ).manual_seed(lowerCAmelCase    )
     __lowerCAmelCase  :			Any					    =						{
         """prompt""": """a photograph of an astronaut riding a horse""",
         """generator""": generator,
         """num_inference_steps""": 10,
         """guidance_scale""": 6.0,
         """output_type""": """numpy""",
         """parallel""": 3,
         """debug""": True,
     }
     return inputs







  def        SCREAMING_SNAKE_CASE    (			self :    List[str]    )							->	int:






     """simple docstring"""







     __lowerCAmelCase  :			Union[str, Any]					    =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase  :			Union[str, Any]					    =						self.get_dummy_components()
     __lowerCAmelCase  :			List[str]					    =						StableDiffusionParadigmsPipeline(**lowerCAmelCase    )
     __lowerCAmelCase  :			Optional[Any]					    =						sd_pipe.to(lowerCAmelCase    )
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase    )

     __lowerCAmelCase  :			Optional[Any]					    =						self.get_dummy_inputs(lowerCAmelCase    )
     __lowerCAmelCase  :			Union[str, Any]					    =						sd_pipe(**lowerCAmelCase    ).images
     __lowerCAmelCase  :			Union[str, Any]					    =						image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     __lowerCAmelCase  :			int					    =						np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]    )

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







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[Any]    )							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			Optional[Any]					    =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase  :			str					    =						self.get_dummy_components()
     torch.manual_seed(0    )
     __lowerCAmelCase  :			List[Any]					    =						DDPMParallelScheduler()
     torch.manual_seed(0    )
     __lowerCAmelCase  :			Any					    =						StableDiffusionParadigmsPipeline(**lowerCAmelCase    )
     __lowerCAmelCase  :			Any					    =						sd_pipe.to(lowerCAmelCase    )
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase    )

     __lowerCAmelCase  :			Dict					    =						self.get_dummy_inputs(lowerCAmelCase    )
     __lowerCAmelCase  :			Dict					    =						sd_pipe(**lowerCAmelCase    ).images
     __lowerCAmelCase  :			str					    =						image[0, -3:, -3:, -1]
     assert image.shape == (1, 64, 64, 3)

     __lowerCAmelCase  :			List[str]					    =						np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]    )

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







  def        SCREAMING_SNAKE_CASE    (			self :    List[Any]    )							->	Optional[Any]:






     """simple docstring"""







     super().test_inference_batch_consistent(batch_sizes=[1, 2]    )







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[Any]    )							->	Any:






     """simple docstring"""







     super().test_inference_batch_single_identical(batch_size=2       ,							expected_max_diff=3e-3    )







  def        SCREAMING_SNAKE_CASE    (			self :    Dict    )							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			List[str]					    =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase  :			Union[str, Any]					    =						self.get_dummy_components()
     __lowerCAmelCase  :			str					    =						StableDiffusionParadigmsPipeline(**lowerCAmelCase    )
     __lowerCAmelCase  :			str					    =						sd_pipe.to(lowerCAmelCase    )
     sd_pipe.set_progress_bar_config(disable=lowerCAmelCase    )

     __lowerCAmelCase  :			List[Any]					    =						self.get_dummy_inputs(lowerCAmelCase    )
     __lowerCAmelCase  :			Optional[int]					    =						"""french fries"""
     __lowerCAmelCase  :			Union[str, Any]					    =						sd_pipe(**lowerCAmelCase       ,							negative_prompt=lowerCAmelCase    )
     __lowerCAmelCase  :			Tuple					    =						output.images
     __lowerCAmelCase  :			int					    =						image[0, -3:, -3:, -1]

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

     __lowerCAmelCase  :			List[Any]					    =						np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]    )

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






@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE	(  unittest.TestCase ):
  """simple docstring"""






  def        SCREAMING_SNAKE_CASE    (			self :    Dict    )							->	Any:






     """simple docstring"""







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







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[int]       ,							lowerCAmelCase :    List[Any]=0    )							->	List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :			Optional[int]					    =						torch.Generator(device=lowerCAmelCase    ).manual_seed(lowerCAmelCase    )
     __lowerCAmelCase  :			Any					    =						{
         """prompt""": """a photograph of an astronaut riding a horse""",
         """generator""": generator,
         """num_inference_steps""": 10,
         """guidance_scale""": 7.5,
         """output_type""": """numpy""",
         """parallel""": 3,
         """debug""": True,
     }
     return inputs







  def        SCREAMING_SNAKE_CASE    (			self :    Tuple    )							->	List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :			Optional[int]					    =						"""stabilityai/stable-diffusion-2-base"""
     __lowerCAmelCase  :			Tuple					    =						DDIMParallelScheduler.from_pretrained(lowerCAmelCase       ,							subfolder="""scheduler"""    )
     __lowerCAmelCase  :			Union[str, Any]					    =						StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase       ,							scheduler=lowerCAmelCase       ,							safety_checker=lowerCAmelCase    )
     pipe.to(lowerCAmelCase    )
     pipe.set_progress_bar_config(disable=lowerCAmelCase    )
     pipe.enable_attention_slicing()

     __lowerCAmelCase  :			Union[str, Any]					    =						self.get_inputs()
     __lowerCAmelCase  :			List[Any]					    =						pipe(**lowerCAmelCase    ).images
     __lowerCAmelCase  :			List[str]					    =						image[0, -3:, -3:, -1].flatten()

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

     __lowerCAmelCase  :			List[str]					    =						np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]    )

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





'''simple docstring'''






import torch

from diffusers import EulerAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class 	SCREAMING_SNAKE_CASE__			(				_UpperCAmelCase      ):
    _A			    =							(EulerAncestralDiscreteScheduler,)
    _A			    =							10




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[Any]     =       {
              "num_train_timesteps": 1100,
              "beta_start": 0.0001,
              "beta_end": 0.02,
              "beta_schedule": "linear",
          }

          config.update(**lowercase__			)
          return config




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for timesteps in [10, 50, 100, 1000]:
                self.check_over_configs(num_train_timesteps=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for beta_start, beta_end in zip([0.00001, 0.0001, 0.001]				,     [0.0002, 0.002, 0.02]			):
                self.check_over_configs(beta_start=lowercase__				,     beta_end=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for schedule in ["linear", "scaled_linear"]:
                self.check_over_configs(beta_schedule=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          for prediction_type in ["epsilon", "v_prediction"]:
                self.check_over_configs(prediction_type=lowercase__			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			int     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			Any     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			int     =       torch.manual_seed(0			)

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       sample.to(lowercase__			)

          for i, t in enumerate(scheduler.timesteps			):
                SCREAMING_SNAKE_CASE_      :			Tuple     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			List[str]     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Dict     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__				,     generator=lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Tuple     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			Tuple     =       torch.mean(torch.abs(lowercase__			)			)

          assert abs(result_sum.item() - 152.3192			) < 1e-2
          assert abs(result_mean.item() - 0.1983			) < 1e-3




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			int     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.get_scheduler_config(prediction_type="v_prediction"			)
          SCREAMING_SNAKE_CASE_      :			str     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			Dict     =       torch.manual_seed(0			)

          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       self.dummy_sample_deter * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       sample.to(lowercase__			)

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

                SCREAMING_SNAKE_CASE_      :			List[Any]     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__				,     generator=lowercase__			)
                SCREAMING_SNAKE_CASE_      :			List[Any]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			Dict     =       torch.mean(torch.abs(lowercase__			)			)

          assert abs(result_sum.item() - 108.4439			) < 1e-2
          assert abs(result_mean.item() - 0.1412			) < 1e-3







    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			Tuple     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps				,     device=lowercase__			)
          SCREAMING_SNAKE_CASE_      :			Any     =       torch.manual_seed(0			)

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       sample.to(lowercase__			)

          for t in scheduler.timesteps:
                SCREAMING_SNAKE_CASE_      :			Dict     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__				,     generator=lowercase__			)
                SCREAMING_SNAKE_CASE_      :			List[str]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			Any     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.mean(torch.abs(lowercase__			)			)

          assert abs(result_sum.item() - 152.3192			) < 1e-2
          assert abs(result_mean.item() - 0.1983			) < 1e-3

