






"""simple docstring"""
import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest



class 						a_ (      snake_case_					):







      '''simple docstring'''

      lowerCamelCase__			:							Dict					=      (DPMSolverMultistepScheduler,)
      lowerCamelCase__			:							List[str]					=      (('num_inference_steps', 25),)




      def      a__				(self,	**lowerCamelCase_   ):

       '''simple docstring'''



       lowerCamelCase__	:     Tuple								=		{
           'num_train_timesteps': 1_0_0_0,
           'beta_start': 0.0_001,
           'beta_end': 0.02,
           'beta_schedule': 'linear',
           'solver_order': 2,
           'prediction_type': 'epsilon',
           'thresholding': False,
           'sample_max_value': 1.0,
           'algorithm_type': 'dpmsolver++',
           'solver_type': 'midpoint',
           'lower_order_final': False,
           'lambda_min_clipped': -float('inf'   ),
           'variance_type': None,
       }

       config.update(**lowerCamelCase_   )
       return config




      def      a__				(self,	lowerCamelCase_=0,	**lowerCamelCase_   ):

       '''simple docstring'''



       lowerCamelCase__	:     Any								=		dict(self.forward_default_kwargs   )
       lowerCamelCase__	:     Tuple								=		kwargs.pop('num_inference_steps',	lowerCamelCase_   )
       lowerCamelCase__	:     Optional[int]								=		self.dummy_sample
       lowerCamelCase__	:     int								=		0.1 * sample
       lowerCamelCase__	:     str								=		[residual + 0.2, residual + 0.15, residual + 0.10]

       for scheduler_class in self.scheduler_classes:
        lowerCamelCase__	:     List[str]								=		self.get_scheduler_config(**lowerCamelCase_   )
        lowerCamelCase__	:     Dict								=		scheduler_class(**lowerCamelCase_   )
        scheduler.set_timesteps(lowerCamelCase_   )
        # copy over dummy past residuals
        lowerCamelCase__	:     Dict								=		dummy_past_residuals[: scheduler.config.solver_order]

        with tempfile.TemporaryDirectory() as tmpdirname:
         scheduler.save_config(lowerCamelCase_   )
         lowerCamelCase__	:     int								=		scheduler_class.from_pretrained(lowerCamelCase_   )
         new_scheduler.set_timesteps(lowerCamelCase_   )
         # copy over dummy past residuals
         lowerCamelCase__	:     List[Any]								=		dummy_past_residuals[: new_scheduler.config.solver_order]

        lowerCamelCase__							,	lowerCamelCase__	:     Tuple								=		sample, sample
        for t in range(lowerCamelCase_,	time_step + scheduler.config.solver_order + 1   ):
         lowerCamelCase__	:     int								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).prev_sample
         lowerCamelCase__	:     Dict								=		new_scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).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,	lowerCamelCase_=0,	**lowerCamelCase_   ):

       '''simple docstring'''



       lowerCamelCase__	:     Optional[int]								=		dict(self.forward_default_kwargs   )
       lowerCamelCase__	:     int								=		kwargs.pop('num_inference_steps',	lowerCamelCase_   )
       lowerCamelCase__	:     int								=		self.dummy_sample
       lowerCamelCase__	:     List[Any]								=		0.1 * sample
       lowerCamelCase__	:     Any								=		[residual + 0.2, residual + 0.15, residual + 0.10]

       for scheduler_class in self.scheduler_classes:
        lowerCamelCase__	:     Dict								=		self.get_scheduler_config()
        lowerCamelCase__	:     Union[str, Any]								=		scheduler_class(**lowerCamelCase_   )
        scheduler.set_timesteps(lowerCamelCase_   )

        # copy over dummy past residuals (must be after setting timesteps)
        lowerCamelCase__	:     int								=		dummy_past_residuals[: scheduler.config.solver_order]

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

         # copy over dummy past residual (must be after setting timesteps)
         lowerCamelCase__	:     List[str]								=		dummy_past_residuals[: new_scheduler.config.solver_order]

        lowerCamelCase__	:     Optional[Any]								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).prev_sample
        lowerCamelCase__	:     int								=		new_scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).prev_sample

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




      def      a__				(self,	lowerCamelCase_=None,	**lowerCamelCase_   ):

       '''simple docstring'''



       if scheduler is None:
        lowerCamelCase__	:     Dict								=		self.scheduler_classes[0]
        lowerCamelCase__	:     Optional[int]								=		self.get_scheduler_config(**lowerCamelCase_   )
        lowerCamelCase__	:     Any								=		scheduler_class(**lowerCamelCase_   )

       lowerCamelCase__	:     List[Any]								=		1_0
       lowerCamelCase__	:     Tuple								=		self.dummy_model()
       lowerCamelCase__	:     Dict								=		self.dummy_sample_deter
       scheduler.set_timesteps(lowerCamelCase_   )

       for i, t in enumerate(scheduler.timesteps   ):
        lowerCamelCase__	:     str								=		model(lowerCamelCase_,	lowerCamelCase_   )
        lowerCamelCase__	:     Optional[Any]								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_   ).prev_sample

       return sample




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     str								=		dict(self.forward_default_kwargs   )

       lowerCamelCase__	:     Optional[Any]								=		kwargs.pop('num_inference_steps',	lowerCamelCase_   )

       for scheduler_class in self.scheduler_classes:
        lowerCamelCase__	:     Any								=		self.get_scheduler_config()
        lowerCamelCase__	:     List[Any]								=		scheduler_class(**lowerCamelCase_   )

        lowerCamelCase__	:     Dict								=		self.dummy_sample
        lowerCamelCase__	:     str								=		0.1 * sample

        if num_inference_steps is not None and hasattr(lowerCamelCase_,	'set_timesteps'   ):
         scheduler.set_timesteps(lowerCamelCase_   )
        elif num_inference_steps is not None and not hasattr(lowerCamelCase_,	'set_timesteps'   ):
         lowerCamelCase__	:     Optional[int]								=		num_inference_steps

        # copy over dummy past residuals (must be done after set_timesteps)
        lowerCamelCase__	:     str								=		[residual + 0.2, residual + 0.15, residual + 0.10]
        lowerCamelCase__	:     Optional[int]								=		dummy_past_residuals[: scheduler.config.solver_order]

        lowerCamelCase__	:     Dict								=		scheduler.timesteps[5]
        lowerCamelCase__	:     Optional[Any]								=		scheduler.timesteps[6]

        lowerCamelCase__	:     Any								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).prev_sample
        lowerCamelCase__	:     List[str]								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_,	**lowerCamelCase_   ).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 [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
        self.check_over_configs(num_train_timesteps=lowerCamelCase_   )




      def      a__				(self   ):

       '''simple docstring'''



       self.check_over_configs(thresholding=lowerCamelCase_   )
       for order in [1, 2, 3]:
        for solver_type in ["midpoint", "heun"]:
         for threshold in [0.5, 1.0, 2.0]:
          for prediction_type in ["epsilon", "sample"]:
           self.check_over_configs(
               thresholding=lowerCamelCase_,	prediction_type=lowerCamelCase_,	sample_max_value=lowerCamelCase_,	algorithm_type='dpmsolver++',	solver_order=lowerCamelCase_,	solver_type=lowerCamelCase_,	)




      def      a__				(self   ):

       '''simple docstring'''



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




      def      a__				(self   ):

       '''simple docstring'''



       for algorithm_type in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
        for solver_type in ["midpoint", "heun"]:
         for order in [1, 2, 3]:
          for prediction_type in ["epsilon", "sample"]:
           if algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
            if order == 3:
             continue
           else:
            self.check_over_configs(
                solver_order=lowerCamelCase_,	solver_type=lowerCamelCase_,	prediction_type=lowerCamelCase_,	algorithm_type=lowerCamelCase_,	)
           lowerCamelCase__	:     List[str]								=		self.full_loop(
               solver_order=lowerCamelCase_,	solver_type=lowerCamelCase_,	prediction_type=lowerCamelCase_,	algorithm_type=lowerCamelCase_,	)
           assert not torch.isnan(lowerCamelCase_   ).any(), "Samples have nan numbers"




      def      a__				(self   ):

       '''simple docstring'''



       self.check_over_configs(lower_order_final=lowerCamelCase_   )
       self.check_over_configs(lower_order_final=lowerCamelCase_   )




      def      a__				(self   ):

       '''simple docstring'''



       self.check_over_configs(lambda_min_clipped=-float('inf'   )   )
       self.check_over_configs(lambda_min_clipped=-5.1   )




      def      a__				(self   ):

       '''simple docstring'''



       self.check_over_configs(variance_type=lowerCamelCase_   )
       self.check_over_configs(variance_type='learned_range'   )




      def      a__				(self   ):

       '''simple docstring'''



       for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
        self.check_over_forward(num_inference_steps=lowerCamelCase_,	time_step=0   )




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Optional[Any]								=		self.full_loop()
       lowerCamelCase__	:     Tuple								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 0.3_301   ) < 1e-3




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     int								=		self.full_loop(thresholding=lowerCamelCase_,	dynamic_thresholding_ratio=0.87,	sample_max_value=0.5   )
       lowerCamelCase__	:     List[Any]								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 1.1_364   ) < 1e-3




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Tuple								=		self.full_loop(prediction_type='v_prediction'   )
       lowerCamelCase__	:     Tuple								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 0.2_251   ) < 1e-3




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     int								=		self.full_loop(prediction_type='v_prediction',	use_karras_sigmas=lowerCamelCase_   )
       lowerCamelCase__	:     str								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 0.2_096   ) < 1e-3




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Any								=		DPMSolverMultistepScheduler(**self.get_scheduler_config()   )
       lowerCamelCase__	:     Tuple								=		self.full_loop(scheduler=lowerCamelCase_   )
       lowerCamelCase__	:     List[str]								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 0.3_301   ) < 1e-3

       lowerCamelCase__	:     Optional[int]								=		DPMSolverSinglestepScheduler.from_config(scheduler.config   )
       lowerCamelCase__	:     List[Any]								=		UniPCMultistepScheduler.from_config(scheduler.config   )
       lowerCamelCase__	:     Optional[int]								=		DEISMultistepScheduler.from_config(scheduler.config   )
       lowerCamelCase__	:     Optional[Any]								=		DPMSolverMultistepScheduler.from_config(scheduler.config   )

       lowerCamelCase__	:     Optional[Any]								=		self.full_loop(scheduler=lowerCamelCase_   )
       lowerCamelCase__	:     Optional[Any]								=		torch.mean(torch.abs(lowerCamelCase_   )   )

       assert abs(result_mean.item() - 0.3_301   ) < 1e-3




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Optional[int]								=		self.scheduler_classes[0]
       lowerCamelCase__	:     Union[str, Any]								=		self.get_scheduler_config(thresholding=lowerCamelCase_,	dynamic_thresholding_ratio=0   )
       lowerCamelCase__	:     Dict								=		scheduler_class(**lowerCamelCase_   )

       lowerCamelCase__	:     Optional[Any]								=		1_0
       lowerCamelCase__	:     Any								=		self.dummy_model()
       lowerCamelCase__	:     str								=		self.dummy_sample_deter.half()
       scheduler.set_timesteps(lowerCamelCase_   )

       for i, t in enumerate(scheduler.timesteps   ):
        lowerCamelCase__	:     int								=		model(lowerCamelCase_,	lowerCamelCase_   )
        lowerCamelCase__	:     List[str]								=		scheduler.step(lowerCamelCase_,	lowerCamelCase_,	lowerCamelCase_   ).prev_sample

       assert sample.dtype == torch.floataa





      def      a__				(self,	**lowerCamelCase_   ):

       '''simple docstring'''



       for scheduler_class in self.scheduler_classes:
        lowerCamelCase__	:     List[str]								=		self.get_scheduler_config(**lowerCamelCase_   )
        lowerCamelCase__	:     List[Any]								=		scheduler_class(**lowerCamelCase_   )

        scheduler.set_timesteps(scheduler.config.num_train_timesteps   )
        assert len(scheduler.timesteps.unique()   ) == scheduler.num_inference_steps









"""simple docstring"""
import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
      import transformer_engine.pytorch as te






def 			lowerCamelCase_    (		_lowerCamelCase     ,  _lowerCamelCase=True     ,  _lowerCamelCase=True     ,  _lowerCamelCase=True							):
 if not is_fpa_available():
  raise ImportError('Using `convert_model` requires transformer_engine to be installed.'							)
 for name, module in model.named_children():
  if isinstance(_lowerCamelCase     ,  nn.Linear							) and to_transformer_engine and _convert_linear:
   # Return early if the linear layer weights are not multiples of 16
   if any(p % 16 != 0 for p in module.weight.shape							):
    return
   lowerCamelCase__	:     List[Any]								=		module.bias is not None
   lowerCamelCase__	:     Tuple								=		te.Linear(
       module.in_features     ,  module.out_features     ,  bias=_lowerCamelCase     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Optional[Any]								=		module.weight.data.clone()
   if has_bias:
    lowerCamelCase__	:     List[str]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  nn.LayerNorm							) and to_transformer_engine and _convert_ln:
   lowerCamelCase__	:     Tuple								=		te.LayerNorm(module.normalized_shape[0]     ,  eps=module.eps     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Tuple								=		module.weight.data.clone()
   lowerCamelCase__	:     Optional[Any]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  te.Linear							) and not to_transformer_engine and _convert_linear:
   lowerCamelCase__	:     Optional[Any]								=		module.bias is not None
   lowerCamelCase__	:     Union[str, Any]								=		nn.Linear(
       module.in_features     ,  module.out_features     ,  bias=_lowerCamelCase     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     str								=		module.weight.data.clone()
   if has_bias:
    lowerCamelCase__	:     Optional[int]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  te.LayerNorm							) and not to_transformer_engine and _convert_ln:
   lowerCamelCase__	:     str								=		nn.LayerNorm(module.normalized_shape[0]     ,  eps=module.eps     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Dict								=		module.weight.data.clone()
   lowerCamelCase__	:     Dict								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  else:
   convert_model(
       _lowerCamelCase     ,  to_transformer_engine=_lowerCamelCase     ,  _convert_linear=_convert_linear     ,  _convert_ln=_convert_ln     ,  )






def 			lowerCamelCase_    (		_lowerCamelCase							):
 if not is_fpa_available():
  raise ImportError('Using `has_transformer_engine_layers` requires transformer_engine to be installed.'							)
 for m in model.modules():
  if isinstance(_lowerCamelCase     ,  (te.LayerNorm, te.Linear)							):
   return True
 return False

