




'''simple docstring'''



import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest





class     lowerCAmelCase   ( __lowerCAmelCase):
       __lowercase  :      List[str]        =			(DPMSolverMultistepScheduler,)
       __lowercase  :      Tuple        =			(('''num_inference_steps''', 25),)


       def       lowerCAmelCase      (   self    ,  **__SCREAMING_SNAKE_CASE  )      ->  Dict:



            '''simple docstring'''




            __snake_case    					=  {
                '''num_train_timesteps''': 1000,
                '''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(**__SCREAMING_SNAKE_CASE  )
            return config


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE=0    ,  **__SCREAMING_SNAKE_CASE  )      ->  Dict:



            '''simple docstring'''




            __snake_case    					=  dict(self.forward_default_kwargs  )
            __snake_case    					=  kwargs.pop('''num_inference_steps'''    ,  __SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.dummy_sample
            __snake_case    					=  0.1 * sample
            __snake_case    					=  [residual + 0.2, residual + 0.15, residual + 0.10]

            for scheduler_class in self.scheduler_classes:
                 __snake_case    					=  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE  )
                 __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )
                 scheduler.set_timesteps(__SCREAMING_SNAKE_CASE  )
                 # copy over dummy past residuals
                 __snake_case    					=  dummy_past_residuals[: scheduler.config.solver_order]

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

                 __snake_case					,    __snake_case    					=  sample, sample
                 for t in range(__SCREAMING_SNAKE_CASE    ,  time_step + scheduler.config.solver_order + 1  ):
                      __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample
                      __snake_case    					=  new_scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample

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


       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            pass


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE=0    ,  **__SCREAMING_SNAKE_CASE  )      ->  str:



            '''simple docstring'''




            __snake_case    					=  dict(self.forward_default_kwargs  )
            __snake_case    					=  kwargs.pop('''num_inference_steps'''    ,  __SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.dummy_sample
            __snake_case    					=  0.1 * sample
            __snake_case    					=  [residual + 0.2, residual + 0.15, residual + 0.10]

            for scheduler_class in self.scheduler_classes:
                 __snake_case    					=  self.get_scheduler_config()
                 __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )
                 scheduler.set_timesteps(__SCREAMING_SNAKE_CASE  )

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

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

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

                 __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample
                 __snake_case    					=  new_scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample

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


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE=None    ,  **__SCREAMING_SNAKE_CASE  )      ->  List[str]:



            '''simple docstring'''




            if scheduler is None:
                 __snake_case    					=  self.scheduler_classes[0]
                 __snake_case    					=  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE  )
                 __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )

            __snake_case    					=  10
            __snake_case    					=  self.dummy_model()
            __snake_case    					=  self.dummy_sample_deter
            scheduler.set_timesteps(__SCREAMING_SNAKE_CASE  )

            for i, t in enumerate(scheduler.timesteps  ):
                 __snake_case    					=  model(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
                 __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ).prev_sample

            return sample


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



            '''simple docstring'''




            __snake_case    					=  dict(self.forward_default_kwargs  )

            __snake_case    					=  kwargs.pop('''num_inference_steps'''    ,  __SCREAMING_SNAKE_CASE  )

            for scheduler_class in self.scheduler_classes:
                 __snake_case    					=  self.get_scheduler_config()
                 __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )

                 __snake_case    					=  self.dummy_sample
                 __snake_case    					=  0.1 * sample

                 if num_inference_steps is not None and hasattr(__SCREAMING_SNAKE_CASE    ,  '''set_timesteps'''  ):
                      scheduler.set_timesteps(__SCREAMING_SNAKE_CASE  )
                 elif num_inference_steps is not None and not hasattr(__SCREAMING_SNAKE_CASE    ,  '''set_timesteps'''  ):
                      __snake_case    					=  num_inference_steps

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

                 __snake_case    					=  scheduler.timesteps[5]
                 __snake_case    					=  scheduler.timesteps[6]

                 __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample
                 __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  **__SCREAMING_SNAKE_CASE  ).prev_sample

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


       def       lowerCAmelCase      (   self  )      ->  Tuple:



            '''simple docstring'''




            for timesteps in [25, 50, 100, 999, 1000]:
                 self.check_over_configs(num_train_timesteps=__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self  )      ->  Optional[int]:



            '''simple docstring'''




            self.check_over_configs(thresholding=__SCREAMING_SNAKE_CASE  )
            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=__SCREAMING_SNAKE_CASE    ,  prediction_type=__SCREAMING_SNAKE_CASE    ,  sample_max_value=__SCREAMING_SNAKE_CASE    ,  algorithm_type='''dpmsolver++'''    ,  solver_order=__SCREAMING_SNAKE_CASE    ,  solver_type=__SCREAMING_SNAKE_CASE    ,  )


       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




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


       def       lowerCAmelCase      (   self  )      ->  int:



            '''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=__SCREAMING_SNAKE_CASE    ,  solver_type=__SCREAMING_SNAKE_CASE    ,  prediction_type=__SCREAMING_SNAKE_CASE    ,  algorithm_type=__SCREAMING_SNAKE_CASE    ,  )
                                __snake_case    					=  self.full_loop(
                                    solver_order=__SCREAMING_SNAKE_CASE    ,  solver_type=__SCREAMING_SNAKE_CASE    ,  prediction_type=__SCREAMING_SNAKE_CASE    ,  algorithm_type=__SCREAMING_SNAKE_CASE    ,  )
                                assert not torch.isnan(__SCREAMING_SNAKE_CASE  ).any(), "Samples have nan numbers"


       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            self.check_over_configs(lower_order_final=__SCREAMING_SNAKE_CASE  )
            self.check_over_configs(lower_order_final=__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self  )      ->  Tuple:



            '''simple docstring'''




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


       def       lowerCAmelCase      (   self  )      ->  Union[str, Any]:



            '''simple docstring'''




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


       def       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




            for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
                 self.check_over_forward(num_inference_steps=__SCREAMING_SNAKE_CASE    ,  time_step=0  )


       def       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




            __snake_case    					=  self.full_loop()
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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


       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            __snake_case    					=  self.full_loop(thresholding=__SCREAMING_SNAKE_CASE    ,  dynamic_thresholding_ratio=0.87    ,  sample_max_value=0.5  )
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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


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



            '''simple docstring'''




            __snake_case    					=  self.full_loop(prediction_type='''v_prediction'''  )
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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


       def       lowerCAmelCase      (   self  )      ->  Any:



            '''simple docstring'''




            __snake_case    					=  self.full_loop(prediction_type='''v_prediction'''    ,  use_karras_sigmas=__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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


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



            '''simple docstring'''




            __snake_case    					=  DPMSolverMultistepScheduler(**self.get_scheduler_config()  )
            __snake_case    					=  self.full_loop(scheduler=__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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

            __snake_case    					=  DPMSolverSinglestepScheduler.from_config(scheduler.config  )
            __snake_case    					=  UniPCMultistepScheduler.from_config(scheduler.config  )
            __snake_case    					=  DEISMultistepScheduler.from_config(scheduler.config  )
            __snake_case    					=  DPMSolverMultistepScheduler.from_config(scheduler.config  )

            __snake_case    					=  self.full_loop(scheduler=__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  torch.mean(torch.abs(__SCREAMING_SNAKE_CASE  )  )

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


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



            '''simple docstring'''




            __snake_case    					=  self.scheduler_classes[0]
            __snake_case    					=  self.get_scheduler_config(thresholding=__SCREAMING_SNAKE_CASE    ,  dynamic_thresholding_ratio=0  )
            __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )

            __snake_case    					=  10
            __snake_case    					=  self.dummy_model()
            __snake_case    					=  self.dummy_sample_deter.half()
            scheduler.set_timesteps(__SCREAMING_SNAKE_CASE  )

            for i, t in enumerate(scheduler.timesteps  ):
                 __snake_case    					=  model(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
                 __snake_case    					=  scheduler.step(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ).prev_sample

            assert sample.dtype == torch.floataa




       def       lowerCAmelCase      (   self    ,  **__SCREAMING_SNAKE_CASE  )      ->  int:



            '''simple docstring'''




            for scheduler_class in self.scheduler_classes:
                 __snake_case    					=  self.get_scheduler_config(**__SCREAMING_SNAKE_CASE  )
                 __snake_case    					=  scheduler_class(**__SCREAMING_SNAKE_CASE  )

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












'''simple docstring'''



import itertools
from dataclasses import dataclass
from io import StringIO
from typing import Optional

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


UpperCAmelCase_       :  str      =		datasets.utils.logging.get_logger(__name__)





@dataclass
class     lowerCAmelCase   ( datasets.BuilderConfig):
       __lowercase  :      Optional[datasets.Features]        =			None
       __lowercase  :      str        =			"utf-8"
       __lowercase  :      Optional[str]        =			None
       __lowercase  :      int        =			10 << 20  # 10MB
       __lowercase  :      bool        =			False
       __lowercase  :      str        =			"line"





class     lowerCAmelCase   ( datasets.ArrowBasedBuilder):
       __lowercase  :      Dict        =			TextConfig


       def       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




            return datasets.DatasetInfo(features=self.config.features  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            if not self.config.data_files:
                 raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''  )
            __snake_case    					=  dl_manager.download_and_extract(self.config.data_files  )
            if isinstance(__SCREAMING_SNAKE_CASE    ,  (str, list, tuple)  ):
                 __snake_case    					=  data_files
                 if isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ):
                      __snake_case    					=  [files]
                 __snake_case    					=  [dl_manager.iter_files(__SCREAMING_SNAKE_CASE  ) for file in files]
                 return [datasets.SplitGenerator(name=datasets.Split.TRAIN    ,  gen_kwargs={'''files''': files}  )]
            __snake_case    					=  []
            for split_name, files in data_files.items():
                 if isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ):
                      __snake_case    					=  [files]
                 __snake_case    					=  [dl_manager.iter_files(__SCREAMING_SNAKE_CASE  ) for file in files]
                 splits.append(datasets.SplitGenerator(name=__SCREAMING_SNAKE_CASE    ,  gen_kwargs={'''files''': files}  )  )
            return splits


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  pa.Table:



            '''simple docstring'''




            if self.config.features is not None:
                 __snake_case    					=  self.config.features.arrow_schema
                 if all(not require_storage_cast(__SCREAMING_SNAKE_CASE  ) for feature in self.config.features.values()  ):
                      # cheaper cast
                      __snake_case    					=  pa_table.cast(__SCREAMING_SNAKE_CASE  )
                 else:
                      # more expensive cast; allows str <-> int/float or str to Audio for example
                      __snake_case    					=  table_cast(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  )
                 return pa_table
            else:
                 return pa_table.cast(pa.schema({'''text''': pa.string()}  )  )




       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            __snake_case    					=  list(self.config.features  ) if self.config.features is not None else ['''text''']
            for file_idx, file in enumerate(itertools.chain.from_iterable(__SCREAMING_SNAKE_CASE  )  ):
                 # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
                 with open(__SCREAMING_SNAKE_CASE    ,  encoding=self.config.encoding    ,  errors=self.config.errors  ) as f:
                      if self.config.sample_by == "line":
                           __snake_case    					=  0
                           while True:
                                __snake_case    					=  f.read(self.config.chunksize  )
                                if not batch:
                                     break
                                batch += f.readline()  # finish current line
                                # StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
                                __snake_case    					=  StringIO(__SCREAMING_SNAKE_CASE  ).readlines()
                                if not self.config.keep_linebreaks:
                                     __snake_case    					=  [line.rstrip('''\n'''  ) for line in batch]
                                __snake_case    					=  pa.Table.from_arrays([pa.array(__SCREAMING_SNAKE_CASE  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                # Uncomment for debugging (will print the Arrow table size and elements)
                                # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                                batch_idx += 1
                      elif self.config.sample_by == "paragraph":
                           __snake_case    					=  0
                           __snake_case    					=  ''''''
                           while True:
                                __snake_case    					=  f.read(self.config.chunksize  )
                                if not new_batch:
                                     break
                                batch += new_batch
                                batch += f.readline()  # finish current line
                                __snake_case    					=  batch.split('''\n\n'''  )
                                __snake_case    					=  pa.Table.from_arrays(
                                    [pa.array([example for example in batch[:-1] if example]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                # Uncomment for debugging (will print the Arrow table size and elements)
                                # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                                batch_idx += 1
                                __snake_case    					=  batch[-1]
                           if batch:
                                __snake_case    					=  pa.Table.from_arrays([pa.array([batch]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                                yield (file_idx, batch_idx), self._cast_table(__SCREAMING_SNAKE_CASE  )
                      elif self.config.sample_by == "document":
                           __snake_case    					=  f.read()
                           __snake_case    					=  pa.Table.from_arrays([pa.array([text]  )]    ,  names=__SCREAMING_SNAKE_CASE  )
                           yield file_idx, self._cast_table(__SCREAMING_SNAKE_CASE  )






