

from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import torch

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available

@dataclass
class 			a					(		UpperCAmelCase			):
						_lowercase        =     4_2


try:
					if not (is_transformers_available() and is_torch_available()):
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
					from .pipeline_text_to_video_synth import TextToVideoSDPipeline
					from .pipeline_text_to_video_synth_imgaimg import VideoToVideoSDPipeline  # noqa: F401
					from .pipeline_text_to_video_zero import TextToVideoZeroPipeline





import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest

class 			a					(		UpperCAmelCase			):
						_lowercase        =     (PNDMScheduler,)
						_lowercase        =     (("num_inference_steps", 5_0),)


						def     _UpperCAmelCase			(    self  ,   **A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      {
										    "num_train_timesteps": 1000,
										    "beta_start": 0.00_01,
										    "beta_end": 0.02,
										    "beta_schedule": "linear",
										}

										config.update(**A_		)
										return config


						def     _UpperCAmelCase			(    self  ,   A_=0  ,   **A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      dict(self.forward_default_kwargs		)
										_UpperCAmelCase    :						Optional[Any]							=      kwargs.pop("num_inference_steps"  ,   A_		)
										_UpperCAmelCase    :						int							=      self.dummy_sample
										_UpperCAmelCase    :						Dict							=      0.1 * sample
										_UpperCAmelCase    :						Optional[int]							=      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

										for scheduler_class in self.scheduler_classes:
														_UpperCAmelCase    :						Dict							=      self.get_scheduler_config(**A_		)
														_UpperCAmelCase    :						Any							=      scheduler_class(**A_		)
														scheduler.set_timesteps(A_		)
														# copy over dummy past residuals
														_UpperCAmelCase    :						Union[str, Any]							=      dummy_past_residuals[:]

														with tempfile.TemporaryDirectory() as tmpdirname:
																		scheduler.save_config(A_		)
																		_UpperCAmelCase    :						str							=      scheduler_class.from_pretrained(A_		)
																		new_scheduler.set_timesteps(A_		)
																		# copy over dummy past residuals
																		_UpperCAmelCase    :						Union[str, Any]							=      dummy_past_residuals[:]

														_UpperCAmelCase    :						Tuple							=      scheduler.step_prk(A_  ,   A_  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						Any							=      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"

														_UpperCAmelCase    :						Union[str, Any]							=      scheduler.step_plms(A_  ,   A_  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						Optional[Any]							=      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     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										pass


						def     _UpperCAmelCase			(    self  ,   A_=0  ,   **A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						Any							=      dict(self.forward_default_kwargs		)
										_UpperCAmelCase    :						int							=      kwargs.pop("num_inference_steps"  ,   A_		)
										_UpperCAmelCase    :						int							=      self.dummy_sample
										_UpperCAmelCase    :						Any							=      0.1 * sample
										_UpperCAmelCase    :						Tuple							=      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

										for scheduler_class in self.scheduler_classes:
														_UpperCAmelCase    :						Optional[Any]							=      self.get_scheduler_config()
														_UpperCAmelCase    :						Dict							=      scheduler_class(**A_		)
														scheduler.set_timesteps(A_		)

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

														with tempfile.TemporaryDirectory() as tmpdirname:
																		scheduler.save_config(A_		)
																		_UpperCAmelCase    :						str							=      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)
																		_UpperCAmelCase    :						List[Any]							=      dummy_past_residuals[:]

														_UpperCAmelCase    :						Tuple							=      scheduler.step_prk(A_  ,   A_  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						Union[str, Any]							=      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"

														_UpperCAmelCase    :						Optional[Any]							=      scheduler.step_plms(A_  ,   A_  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						List[str]							=      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     _UpperCAmelCase			(    self  ,   **A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[Any]							=      self.scheduler_classes[0]
										_UpperCAmelCase    :						Union[str, Any]							=      self.get_scheduler_config(**A_		)
										_UpperCAmelCase    :						List[str]							=      scheduler_class(**A_		)

										_UpperCAmelCase    :						List[str]							=      10
										_UpperCAmelCase    :						Optional[int]							=      self.dummy_model()
										_UpperCAmelCase    :						List[Any]							=      self.dummy_sample_deter
										scheduler.set_timesteps(A_		)

										for i, t in enumerate(scheduler.prk_timesteps		):
														_UpperCAmelCase    :						Dict							=      model(A_  ,   A_		)
														_UpperCAmelCase    :						int							=      scheduler.step_prk(A_  ,   A_  ,   A_		).prev_sample

										for i, t in enumerate(scheduler.plms_timesteps		):
														_UpperCAmelCase    :						Any							=      model(A_  ,   A_		)
														_UpperCAmelCase    :						Any							=      scheduler.step_plms(A_  ,   A_  ,   A_		).prev_sample

										return sample


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      dict(self.forward_default_kwargs		)

										_UpperCAmelCase    :						str							=      kwargs.pop("num_inference_steps"  ,   A_		)

										for scheduler_class in self.scheduler_classes:
														_UpperCAmelCase    :						int							=      self.get_scheduler_config()
														_UpperCAmelCase    :						Optional[Any]							=      scheduler_class(**A_		)

														_UpperCAmelCase    :						int							=      self.dummy_sample
														_UpperCAmelCase    :						str							=      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"		):
																		_UpperCAmelCase    :						Union[str, Any]							=      num_inference_steps

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

														_UpperCAmelCase    :						Any							=      scheduler.step_prk(A_  ,   0  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						Optional[Any]							=      scheduler.step_prk(A_  ,   1  ,   A_  ,   **A_		).prev_sample

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

														_UpperCAmelCase    :						Optional[int]							=      scheduler.step_plms(A_  ,   0  ,   A_  ,   **A_		).prev_sample
														_UpperCAmelCase    :						Optional[Any]							=      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     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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

										_UpperCAmelCase    :						List[str]							=      self.scheduler_classes[0]
										_UpperCAmelCase    :						Optional[int]							=      self.get_scheduler_config(steps_offset=1		)
										_UpperCAmelCase    :						str							=      scheduler_class(**A_		)
										scheduler.set_timesteps(10		)
										assert torch.equal(
										    scheduler.timesteps  ,   torch.LongTensor(
										        [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]		)  ,   )


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										for beta_start, beta_end in zip([0.00_01, 0.0_01]  ,   [0.0_02, 0.02]		):
														self.check_over_configs(beta_start=A_  ,   beta_end=A_		)


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      27

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

														_UpperCAmelCase    :						str							=      self.get_scheduler_config()
														_UpperCAmelCase    :						Tuple							=      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]		):
																		_UpperCAmelCase    :						Dict							=      scheduler.step_prk(A_  ,   A_  ,   A_		).prev_sample


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										with self.assertRaises(A_		):
														_UpperCAmelCase    :						Union[str, Any]							=      self.scheduler_classes[0]
														_UpperCAmelCase    :						Dict							=      self.get_scheduler_config()
														_UpperCAmelCase    :						Union[str, Any]							=      scheduler_class(**A_		)

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


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Dict							=      self.full_loop()
										_UpperCAmelCase    :						int							=      torch.sum(torch.abs(A_		)		)
										_UpperCAmelCase    :						Optional[Any]							=      torch.mean(torch.abs(A_		)		)

										assert abs(result_sum.item() - 1_98.13_18		) < 1e-2
										assert abs(result_mean.item() - 0.25_80		) < 1e-3


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[Any]							=      self.full_loop(prediction_type="v_prediction"		)
										_UpperCAmelCase    :						Union[str, Any]							=      torch.sum(torch.abs(A_		)		)
										_UpperCAmelCase    :						List[Any]							=      torch.mean(torch.abs(A_		)		)

										assert abs(result_sum.item() - 67.39_86		) < 1e-2
										assert abs(result_mean.item() - 0.08_78		) < 1e-3


						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      self.full_loop(set_alpha_to_one=A_  ,   beta_start=0.01		)
										_UpperCAmelCase    :						Dict							=      torch.sum(torch.abs(A_		)		)
										_UpperCAmelCase    :						Optional[int]							=      torch.mean(torch.abs(A_		)		)

										assert abs(result_sum.item() - 2_30.03_99		) < 1e-2
										assert abs(result_mean.item() - 0.29_95		) < 1e-3





						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      self.full_loop(set_alpha_to_one=A_  ,   beta_start=0.01		)
										_UpperCAmelCase    :						Optional[Any]							=      torch.sum(torch.abs(A_		)		)
										_UpperCAmelCase    :						Dict							=      torch.mean(torch.abs(A_		)		)

										assert abs(result_sum.item() - 1_86.94_82		) < 1e-2
										assert abs(result_mean.item() - 0.24_34		) < 1e-3


