






# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()


class 			SCREAMING_SNAKE_CASE      (      lowerCAmelCase   ,		lowerCAmelCase   ,		lowerCAmelCase   ,		unittest.TestCase       ):
     '''simple docstring'''


     UpperCamelCase_				:		int								=						StableDiffusionControlNetImgaImgPipeline
     UpperCamelCase_				:		List[Any]								=						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
     UpperCamelCase_				:		Optional[Any]								=						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     UpperCamelCase_				:		Tuple								=						IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}       )
     UpperCamelCase_				:		Any								=						IMAGE_TO_IMAGE_IMAGE_PARAMS


     def 					_A							(					self							:						List[str]   ):
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   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       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   int						    =     ControlNetModel(
           block_out_channels=(32, 64)       ,			layers_per_block=2       ,			in_channels=4       ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,			cross_attention_dim=32       ,			conditioning_embedding_out_channels=(16, 32)       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     DDIMScheduler(
           beta_start=0.00_085       ,			beta_end=0.012       ,			beta_schedule="scaled_linear"       ,			clip_sample=UpperCAmelCase_       ,			set_alpha_to_one=UpperCAmelCase_       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Any						    =     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       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   int						    =     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=1000       ,			)
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     CLIPTextModel(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

       SCREAMING_SNAKE_CASE							:   str						    =     {
           "unet": unet,
           "controlnet": controlnet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       return components


     def 					_A							(					self							:						str       ,			UpperCAmelCase_							:						List[str]       ,			UpperCAmelCase_							:						int=0   ):
       if str(UpperCAmelCase_   ).startswith("mps"   ):
         SCREAMING_SNAKE_CASE							:   Any						    =     torch.manual_seed(UpperCAmelCase_   )
       else:
         SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     torch.Generator(device=UpperCAmelCase_   ).manual_seed(UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     2
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,			generator=UpperCAmelCase_       ,			device=torch.device(UpperCAmelCase_   )       ,			)
       SCREAMING_SNAKE_CASE							:   Tuple						    =     floats_tensor(control_image.shape       ,			rng=random.Random(UpperCAmelCase_   )   ).to(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Dict						    =     image.cpu().permute(0       ,			2       ,			3       ,			1   )[0]
       SCREAMING_SNAKE_CASE							:   str						    =     Image.fromarray(np.uinta(UpperCAmelCase_   )   ).convert("RGB"   ).resize((64, 64)   )
       SCREAMING_SNAKE_CASE							:   List[str]						    =     {
           "prompt": "A painting of a squirrel eating a burger",
           "generator": generator,
           "num_inference_steps": 2,
           "guidance_scale": 6.0,
           "output_type": "numpy",
           "image": image,
           "control_image": control_image,
       }

       return inputs


     def 					_A							(					self							:						int   ):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3   )


     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available()       ,			reason="XFormers attention is only available with CUDA and `xformers` installed"       ,			)
     def 					_A							(					self							:						str   ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3   )






     def 					_A							(					self							:						Union[str, Any]   ):
       self._test_inference_batch_single_identical(expected_max_diff=2E-3   )





class 			SCREAMING_SNAKE_CASE      (      lowerCAmelCase   ,		lowerCAmelCase   ,		unittest.TestCase       ):
     '''simple docstring'''


     UpperCamelCase_				:		List[str]								=						StableDiffusionControlNetImgaImgPipeline
     UpperCamelCase_				:		str								=						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
     UpperCamelCase_				:		Optional[int]								=						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     UpperCamelCase_				:		Dict								=						frozenset([]       )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess


     def 					_A							(					self							:						Optional[Any]   ):
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     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       ,			)
       torch.manual_seed(0   )

       def init_weights(UpperCAmelCase_							:						List[Any]   ):
         if isinstance(UpperCAmelCase_       ,			torch.nn.Convad   ):
           torch.nn.init.normal(m.weight   )
           m.bias.data.fill_(1.0   )

       SCREAMING_SNAKE_CASE							:   List[str]						    =     ControlNetModel(
           block_out_channels=(32, 64)       ,			layers_per_block=2       ,			in_channels=4       ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,			cross_attention_dim=32       ,			conditioning_embedding_out_channels=(16, 32)       ,			)
       controlneta.controlnet_down_blocks.apply(UpperCAmelCase_   )

       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Any						    =     ControlNetModel(
           block_out_channels=(32, 64)       ,			layers_per_block=2       ,			in_channels=4       ,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,			cross_attention_dim=32       ,			conditioning_embedding_out_channels=(16, 32)       ,			)
       controlneta.controlnet_down_blocks.apply(UpperCAmelCase_   )

       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Any						    =     DDIMScheduler(
           beta_start=0.00_085       ,			beta_end=0.012       ,			beta_schedule="scaled_linear"       ,			clip_sample=UpperCAmelCase_       ,			set_alpha_to_one=UpperCAmelCase_       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Dict						    =     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       ,			)
       torch.manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   Tuple						    =     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=1000       ,			)
       SCREAMING_SNAKE_CASE							:   Any						    =     CLIPTextModel(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

       SCREAMING_SNAKE_CASE							:   Tuple						    =     MultiControlNetModel([controlneta, controlneta]   )

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     {
           "unet": unet,
           "controlnet": controlnet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       return components


     def 					_A							(					self							:						List[str]       ,			UpperCAmelCase_							:						Optional[Any]       ,			UpperCAmelCase_							:						Any=0   ):
       if str(UpperCAmelCase_   ).startswith("mps"   ):
         SCREAMING_SNAKE_CASE							:   Dict						    =     torch.manual_seed(UpperCAmelCase_   )
       else:
         SCREAMING_SNAKE_CASE							:   str						    =     torch.Generator(device=UpperCAmelCase_   ).manual_seed(UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   Any						    =     2

       SCREAMING_SNAKE_CASE							:   Tuple						    =     [
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,			generator=UpperCAmelCase_       ,			device=torch.device(UpperCAmelCase_   )       ,			),
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,			generator=UpperCAmelCase_       ,			device=torch.device(UpperCAmelCase_   )       ,			),
       ]

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     floats_tensor(control_image[0].shape       ,			rng=random.Random(UpperCAmelCase_   )   ).to(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   int						    =     image.cpu().permute(0       ,			2       ,			3       ,			1   )[0]
       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     Image.fromarray(np.uinta(UpperCAmelCase_   )   ).convert("RGB"   ).resize((64, 64)   )
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     {
           "prompt": "A painting of a squirrel eating a burger",
           "generator": generator,
           "num_inference_steps": 2,
           "guidance_scale": 6.0,
           "output_type": "numpy",
           "image": image,
           "control_image": control_image,
       }

       return inputs


     def 					_A							(					self							:						Tuple   ):
       SCREAMING_SNAKE_CASE							:   Any						    =     self.get_dummy_components()
       SCREAMING_SNAKE_CASE							:   str						    =     self.pipeline_class(**UpperCAmelCase_   )
       pipe.to(UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   str						    =     10.0
       SCREAMING_SNAKE_CASE							:   Any						    =     4

       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     self.get_dummy_inputs(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Tuple						    =     steps
       SCREAMING_SNAKE_CASE							:   int						    =     scale
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     pipe(**UpperCAmelCase_   )[0]

       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     self.get_dummy_inputs(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   Tuple						    =     steps
       SCREAMING_SNAKE_CASE							:   Any						    =     scale
       SCREAMING_SNAKE_CASE							:   List[str]						    =     pipe(**UpperCAmelCase_       ,			control_guidance_start=0.1       ,			control_guidance_end=0.2   )[0]

       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     self.get_dummy_inputs(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   List[str]						    =     steps
       SCREAMING_SNAKE_CASE							:   int						    =     scale
       SCREAMING_SNAKE_CASE							:   List[Any]						    =     pipe(**UpperCAmelCase_       ,			control_guidance_start=[0.1, 0.3]       ,			control_guidance_end=[0.2, 0.7]   )[0]

       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     self.get_dummy_inputs(UpperCAmelCase_   )
       SCREAMING_SNAKE_CASE							:   str						    =     steps
       SCREAMING_SNAKE_CASE							:   Dict						    =     scale
       SCREAMING_SNAKE_CASE							:   Dict						    =     pipe(**UpperCAmelCase_       ,			control_guidance_start=0.4       ,			control_guidance_end=[0.5, 0.8]   )[0]

       # make sure that all outputs are different
       assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3
       assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3
       assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3


     def 					_A							(					self							:						Union[str, Any]   ):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3   )


     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available()       ,			reason="XFormers attention is only available with CUDA and `xformers` installed"       ,			)
     def 					_A							(					self							:						str   ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3   )


     def 					_A							(					self							:						List[Any]   ):
       self._test_inference_batch_single_identical(expected_max_diff=2E-3   )






     def 					_A							(					self							:						Any   ):
       SCREAMING_SNAKE_CASE							:   Dict						    =     self.get_dummy_components()
       SCREAMING_SNAKE_CASE							:   Dict						    =     self.pipeline_class(**UpperCAmelCase_   )
       pipe.to(UpperCAmelCase_   )
       pipe.set_progress_bar_config(disable=UpperCAmelCase_   )
       with tempfile.TemporaryDirectory() as tmpdir:
         try:
           # save_pretrained is not implemented for Multi-ControlNet
           pipe.save_pretrained(UpperCAmelCase_   )
         except NotImplementedError:
           pass





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


     def 					_A							(					self							:						Optional[Any]   ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






     def 					_A							(					self							:						Optional[Any]   ):
       SCREAMING_SNAKE_CASE							:   str						    =     ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"   )

       SCREAMING_SNAKE_CASE							:   Union[str, Any]						    =     StableDiffusionControlNetImgaImgPipeline.from_pretrained(
           "runwayml/stable-diffusion-v1-5"       ,			safety_checker=UpperCAmelCase_       ,			controlnet=UpperCAmelCase_   )
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=UpperCAmelCase_   )

       SCREAMING_SNAKE_CASE							:   str						    =     torch.Generator(device="cpu"   ).manual_seed(0   )
       SCREAMING_SNAKE_CASE							:   str						    =     "evil space-punk bird"
       SCREAMING_SNAKE_CASE							:   Optional[Any]						    =     load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"   ).resize((512, 512)   )
       SCREAMING_SNAKE_CASE							:   Optional[int]						    =     load_image(
           "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"   ).resize((512, 512)   )

       SCREAMING_SNAKE_CASE							:   str						    =     pipe(
           UpperCAmelCase_       ,			UpperCAmelCase_       ,			control_image=UpperCAmelCase_       ,			generator=UpperCAmelCase_       ,			output_type="np"       ,			num_inference_steps=50       ,			strength=0.6       ,			)

       SCREAMING_SNAKE_CASE							:   int						    =     output.images[0]

       assert image.shape == (512, 512, 3)

       SCREAMING_SNAKE_CASE							:   Dict						    =     load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy"   )

       assert np.abs(expected_image - image   ).max() < 9E-2











import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest





class  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     _A      =      (PNDMScheduler,)
     _A      =      (('num_inference_steps', 50),)







     def     _a   (self					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		{
               "num_train_timesteps": 10_00,
               "beta_start": 0.0001,
               "beta_end": 0.02,
               "beta_schedule": "linear",
           }

           config.update(**__a     )
           return config







     def     _a   (self					,      __a=0					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )
           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )
           lowerCamelCase             =		self.dummy_sample
           lowerCamelCase             =		0.1 * sample
           lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config(**__a     )
                 lowerCamelCase             =		scheduler_class(**__a     )
                 scheduler.set_timesteps(__a     )
                 # copy over dummy past residuals
                 lowerCamelCase             =		dummy_past_residuals[:]

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       scheduler.save_config(__a     )
                       lowerCamelCase             =		scheduler_class.from_pretrained(__a     )
                       new_scheduler.set_timesteps(__a     )
                       # copy over dummy past residuals
                       lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"

                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"







     def     _a   (self     ):






           '''simple docstring'''



           pass







     def     _a   (self					,      __a=0					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )
           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )
           lowerCamelCase             =		self.dummy_sample
           lowerCamelCase             =		0.1 * sample
           lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )
                 scheduler.set_timesteps(__a     )

                 # copy over dummy past residuals (must be after setting timesteps)
                 lowerCamelCase             =		dummy_past_residuals[:]

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       scheduler.save_config(__a     )
                       lowerCamelCase             =		scheduler_class.from_pretrained(__a     )
                       # copy over dummy past residuals
                       new_scheduler.set_timesteps(__a     )

                       # copy over dummy past residual (must be after setting timesteps)
                       lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_prk(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"

                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		new_scheduler.step_plms(__a					,      __a					,      __a					,      **__a     ).prev_sample

                 assert torch.sum(torch.abs(output - new_output     )     ) < 1E-5, "Scheduler outputs are not identical"







     def     _a   (self					,      **__a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.scheduler_classes[0]
           lowerCamelCase             =		self.get_scheduler_config(**__a     )
           lowerCamelCase             =		scheduler_class(**__a     )

           lowerCamelCase             =		10
           lowerCamelCase             =		self.dummy_model()
           lowerCamelCase             =		self.dummy_sample_deter
           scheduler.set_timesteps(__a     )

           for i, t in enumerate(scheduler.prk_timesteps     ):
                 lowerCamelCase             =		model(__a					,      __a     )
                 lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a     ).prev_sample

           for i, t in enumerate(scheduler.plms_timesteps     ):
                 lowerCamelCase             =		model(__a					,      __a     )
                 lowerCamelCase             =		scheduler.step_plms(__a					,      __a					,      __a     ).prev_sample

           return sample







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		dict(self.forward_default_kwargs     )

           lowerCamelCase             =		kwargs.pop("num_inference_steps"					,      __a     )

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 lowerCamelCase             =		self.dummy_sample
                 lowerCamelCase             =		0.1 * sample

                 if num_inference_steps is not None and hasattr(__a					,      "set_timesteps"     ):
                       scheduler.set_timesteps(__a     )
                 elif num_inference_steps is not None and not hasattr(__a					,      "set_timesteps"     ):
                       lowerCamelCase             =		num_inference_steps

                 # copy over dummy past residuals (must be done after set_timesteps)
                 lowerCamelCase             =		[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
                 lowerCamelCase             =		dummy_past_residuals[:]

                 lowerCamelCase             =		scheduler.step_prk(__a					,      0					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		scheduler.step_prk(__a					,      1					,      __a					,      **__a     ).prev_sample

                 self.assertEqual(output_a.shape					,      sample.shape     )
                 self.assertEqual(output_a.shape					,      output_a.shape     )

                 lowerCamelCase             =		scheduler.step_plms(__a					,      0					,      __a					,      **__a     ).prev_sample
                 lowerCamelCase             =		scheduler.step_plms(__a					,      1					,      __a					,      **__a     ).prev_sample

                 self.assertEqual(output_a.shape					,      sample.shape     )
                 self.assertEqual(output_a.shape					,      output_a.shape     )







     def     _a   (self     ):






           '''simple docstring'''



           for timesteps in [1_00, 10_00]:
                 self.check_over_configs(num_train_timesteps=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for steps_offset in [0, 1]:
                 self.check_over_configs(steps_offset=__a     )

           lowerCamelCase             =		self.scheduler_classes[0]
           lowerCamelCase             =		self.get_scheduler_config(steps_offset=1     )
           lowerCamelCase             =		scheduler_class(**__a     )
           scheduler.set_timesteps(10     )
           assert torch.equal(
               scheduler.timesteps					,      torch.LongTensor(
                   [9_01, 8_51, 8_51, 8_01, 8_01, 7_51, 7_51, 7_01, 7_01, 6_51, 6_51, 6_01, 6_01, 5_01, 4_01, 3_01, 2_01, 1_01, 1]     )					,      )







     def     _a   (self     ):






           '''simple docstring'''



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







     def     _a   (self     ):






           '''simple docstring'''



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







     def     _a   (self     ):






           '''simple docstring'''



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







     def     _a   (self     ):






           '''simple docstring'''



           for t in [1, 5, 10]:
                 self.check_over_forward(time_step=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           for t, num_inference_steps in zip([1, 5, 10]					,      [10, 50, 1_00]     ):
                 self.check_over_forward(num_inference_steps=__a     )







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		27

           for scheduler_class in self.scheduler_classes:
                 lowerCamelCase             =		self.dummy_sample
                 lowerCamelCase             =		0.1 * sample

                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 scheduler.set_timesteps(__a     )

                 # before power of 3 fix, would error on first step, so we only need to do two
                 for i, t in enumerate(scheduler.prk_timesteps[:2]     ):
                       lowerCamelCase             =		scheduler.step_prk(__a					,      __a					,      __a     ).prev_sample







     def     _a   (self     ):






           '''simple docstring'''



           with self.assertRaises(__a     ):
                 lowerCamelCase             =		self.scheduler_classes[0]
                 lowerCamelCase             =		self.get_scheduler_config()
                 lowerCamelCase             =		scheduler_class(**__a     )

                 scheduler.step_plms(self.dummy_sample					,      1					,      self.dummy_sample     ).prev_sample







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop()
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 198.1318     ) < 1E-2
           assert abs(result_mean.item() - 0.2580     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(prediction_type="v_prediction"     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 67.3986     ) < 1E-2
           assert abs(result_mean.item() - 0.0878     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(set_alpha_to_one=__a					,      beta_start=0.01     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 230.0399     ) < 1E-2
           assert abs(result_mean.item() - 0.2995     ) < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.full_loop(set_alpha_to_one=__a					,      beta_start=0.01     )
           lowerCamelCase             =		torch.sum(torch.abs(__a     )     )
           lowerCamelCase             =		torch.mean(torch.abs(__a     )     )

           assert abs(result_sum.item() - 186.9482     ) < 1E-2
           assert abs(result_mean.item() - 0.2434     ) < 1E-3