
"""simple docstring"""


import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


SCREAMING_SNAKE_CASE      =     False



class   __a  (      unittest.TestCase	):
       pass





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






       def 	_SCREAMING_SNAKE_CASE       (  self :							str			)->  Tuple:







              """simple docstring"""
              # clean up the VRAM after each test
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[str]			)->  Any:







              """simple docstring"""
              UpperCamelCase					 =		VersatileDiffusionTextToImagePipeline.from_pretrained("shi-labs/versatile-diffusion"			)
              # remove text_unet
              pipe.remove_unused_weights()
              pipe.to(UpperCAmelCase_			)
              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)

              UpperCamelCase					 =		"A painting of a squirrel eating a burger "
              UpperCamelCase					 =		torch.manual_seed(0			)
              UpperCamelCase					 =		pipe(
                  prompt=UpperCAmelCase_       ,      generator=UpperCAmelCase_       ,      guidance_scale=7.5       ,      num_inference_steps=2       ,      output_type="numpy"			).images

              with tempfile.TemporaryDirectory() as tmpdirname:
                     pipe.save_pretrained(UpperCAmelCase_			)
                     UpperCamelCase					 =		VersatileDiffusionTextToImagePipeline.from_pretrained(UpperCAmelCase_			)
              pipe.to(UpperCAmelCase_			)
              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)

              UpperCamelCase					 =		generator.manual_seed(0			)
              UpperCamelCase					 =		pipe(
                  prompt=UpperCAmelCase_       ,      generator=UpperCAmelCase_       ,      guidance_scale=7.5       ,      num_inference_steps=2       ,      output_type="numpy"			).images

              assert np.abs(image - new_image			).sum() < 1e-5, "Models don't have the same forward pass"






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







              """simple docstring"""
              UpperCamelCase					 =		VersatileDiffusionTextToImagePipeline.from_pretrained(
                  "shi-labs/versatile-diffusion"       ,      torch_dtype=torch.floataa			)
              pipe.to(UpperCAmelCase_			)
              pipe.set_progress_bar_config(disable=UpperCAmelCase_			)

              UpperCamelCase					 =		"A painting of a squirrel eating a burger "
              UpperCamelCase					 =		torch.manual_seed(0			)
              UpperCamelCase					 =		pipe(
                  prompt=UpperCAmelCase_       ,      generator=UpperCAmelCase_       ,      guidance_scale=7.5       ,      num_inference_steps=50       ,      output_type="numpy"			).images

              UpperCamelCase					 =		image[0, 253:256, 253:256, -1]

              assert image.shape == (1, 512, 512, 3)
              UpperCamelCase					 =		np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]			)

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







"""simple docstring"""


import torch

from diffusers import EulerDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class   __a  (      _lowerCAmelCase	):
       UpperCamelCase_		:					Any   =			(EulerDiscreteScheduler,)
       UpperCamelCase_		:					Dict   =			10






       def 	_SCREAMING_SNAKE_CASE       (  self :							Any       ,      **UpperCAmelCase_ :							str			)->  Union[str, Any]:







              """simple docstring"""
              UpperCamelCase					 =		{
                  "num_train_timesteps": 1_100,
                  "beta_start": 0.0001,
                  "beta_end": 0.02,
                  "beta_schedule": "linear",
              }

              config.update(**UpperCAmelCase_			)
              return config






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[Any]			)->  Any:







              """simple docstring"""
              for timesteps in [10, 50, 100, 1_000]:
                     self.check_over_configs(num_train_timesteps=UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							int			)->  Any:







              """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=UpperCAmelCase_       ,      beta_end=UpperCAmelCase_			)






       def 	_SCREAMING_SNAKE_CASE       (  self :							str			)->  Dict:







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






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[Any]			)->  int:







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






       def 	_SCREAMING_SNAKE_CASE       (  self :							str			)->  Any:







              """simple docstring"""
              UpperCamelCase					 =		self.scheduler_classes[0]
              UpperCamelCase					 =		self.get_scheduler_config()
              UpperCamelCase					 =		scheduler_class(**UpperCAmelCase_			)

              scheduler.set_timesteps(self.num_inference_steps			)

              UpperCamelCase					 =		torch.manual_seed(0			)

              UpperCamelCase					 =		self.dummy_model()
              UpperCamelCase					 =		self.dummy_sample_deter * scheduler.init_noise_sigma
              UpperCamelCase					 =		sample.to(UpperCAmelCase_			)

              for i, t in enumerate(scheduler.timesteps			):
                     UpperCamelCase					 =		scheduler.scale_model_input(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		model(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		scheduler.step(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      generator=UpperCAmelCase_			)
                     UpperCamelCase					 =		output.prev_sample

              UpperCamelCase					 =		torch.sum(torch.abs(UpperCAmelCase_			)			)
              UpperCamelCase					 =		torch.mean(torch.abs(UpperCAmelCase_			)			)

              assert abs(result_sum.item() - 10.0807			) < 1e-2
              assert abs(result_mean.item() - 0.0131			) < 1e-3






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[str]			)->  str:







              """simple docstring"""
              UpperCamelCase					 =		self.scheduler_classes[0]
              UpperCamelCase					 =		self.get_scheduler_config(prediction_type="v_prediction"			)
              UpperCamelCase					 =		scheduler_class(**UpperCAmelCase_			)

              scheduler.set_timesteps(self.num_inference_steps			)

              UpperCamelCase					 =		torch.manual_seed(0			)

              UpperCamelCase					 =		self.dummy_model()
              UpperCamelCase					 =		self.dummy_sample_deter * scheduler.init_noise_sigma
              UpperCamelCase					 =		sample.to(UpperCAmelCase_			)

              for i, t in enumerate(scheduler.timesteps			):
                     UpperCamelCase					 =		scheduler.scale_model_input(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		model(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		scheduler.step(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      generator=UpperCAmelCase_			)
                     UpperCamelCase					 =		output.prev_sample

              UpperCamelCase					 =		torch.sum(torch.abs(UpperCAmelCase_			)			)
              UpperCamelCase					 =		torch.mean(torch.abs(UpperCAmelCase_			)			)

              assert abs(result_sum.item() - 0.0002			) < 1e-2
              assert abs(result_mean.item() - 2.2_676e-06			) < 1e-3






       def 	_SCREAMING_SNAKE_CASE       (  self :							List[str]			)->  Tuple:







              """simple docstring"""
              UpperCamelCase					 =		self.scheduler_classes[0]
              UpperCamelCase					 =		self.get_scheduler_config()
              UpperCamelCase					 =		scheduler_class(**UpperCAmelCase_			)

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

              UpperCamelCase					 =		torch.manual_seed(0			)

              UpperCamelCase					 =		self.dummy_model()
              UpperCamelCase					 =		self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
              UpperCamelCase					 =		sample.to(UpperCAmelCase_			)

              for t in scheduler.timesteps:
                     UpperCamelCase					 =		scheduler.scale_model_input(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		model(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		scheduler.step(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      generator=UpperCAmelCase_			)
                     UpperCamelCase					 =		output.prev_sample

              UpperCamelCase					 =		torch.sum(torch.abs(UpperCAmelCase_			)			)
              UpperCamelCase					 =		torch.mean(torch.abs(UpperCAmelCase_			)			)

              assert abs(result_sum.item() - 10.0807			) < 1e-2
              assert abs(result_mean.item() - 0.0131			) < 1e-3






       def 	_SCREAMING_SNAKE_CASE       (  self :							Optional[int]			)->  Any:







              """simple docstring"""
              UpperCamelCase					 =		self.scheduler_classes[0]
              UpperCamelCase					 =		self.get_scheduler_config()
              UpperCamelCase					 =		scheduler_class(**UpperCAmelCase_       ,      use_karras_sigmas=UpperCAmelCase_			)

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

              UpperCamelCase					 =		torch.manual_seed(0			)

              UpperCamelCase					 =		self.dummy_model()
              UpperCamelCase					 =		self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
              UpperCamelCase					 =		sample.to(UpperCAmelCase_			)

              for t in scheduler.timesteps:
                     UpperCamelCase					 =		scheduler.scale_model_input(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		model(UpperCAmelCase_       ,      UpperCAmelCase_			)

                     UpperCamelCase					 =		scheduler.step(UpperCAmelCase_       ,      UpperCAmelCase_       ,      UpperCAmelCase_       ,      generator=UpperCAmelCase_			)
                     UpperCamelCase					 =		output.prev_sample

              UpperCamelCase					 =		torch.sum(torch.abs(UpperCAmelCase_			)			)
              UpperCamelCase					 =		torch.mean(torch.abs(UpperCAmelCase_			)			)

              assert abs(result_sum.item() - 124.52299499511719			) < 1e-2
              assert abs(result_mean.item() - 0.16213932633399963			) < 1e-3





