





"""simple docstring"""




import math
from collections import defaultdict
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput





def 		UpperCAmelCase__						(	_UpperCAmelCase					,   _UpperCAmelCase=0.999					,   _UpperCAmelCase="cosine"					,   ):






  """simple docstring"""





  if alpha_transform_type == "cosine":

    def alpha_bar_fn(_UpperCAmelCase   ):
      return math.cos((t + 0.008) / 1.008 * math.pi / 2   ) ** 2

  elif alpha_transform_type == "exp":

    def alpha_bar_fn(_UpperCAmelCase   ):
      return math.exp(t * -12.0   )

  else:
    raise ValueError(f"""Unsupported alpha_tranform_type: {alpha_transform_type}"""   )

  A_    :			List[Any]									=    []
  for i in range(_UpperCAmelCase   ):
    A_    :			Dict									=    i / num_diffusion_timesteps
    A_    :			Union[str, Any]									=    (i + 1) / num_diffusion_timesteps
    betas.append(min(1 - alpha_bar_fn(_UpperCAmelCase   ) / alpha_bar_fn(_UpperCAmelCase   )					,   _UpperCAmelCase   )   )
  return torch.tensor(_UpperCAmelCase					,   dtype=torch.floataa   )





class 					_UpperCAmelCase    (  UpperCAmelCase__     ,  UpperCAmelCase__						):



      '''simple docstring'''
      lowercase_ :					List[str]             =  [e.name for e in KarrasDiffusionSchedulers]
      lowercase_ :					Dict             =  2

      @register_to_config
      def __init__(	self   , snake_case_ = 1_0_0_0   , snake_case_ = 0.0_00_85   , snake_case_ = 0.0_12   , snake_case_ = "linear"   , snake_case_ = None   , snake_case_ = "epsilon"   , snake_case_ = "linspace"   , snake_case_ = 0   , ):

        """simple docstring"""

        if trained_betas is not None:
          A_    :			Tuple									=    torch.tensor(snake_case_   , dtype=torch.floataa    )
        elif beta_schedule == "linear":
          A_    :			int									=    torch.linspace(snake_case_   , snake_case_   , snake_case_   , dtype=torch.floataa    )
        elif beta_schedule == "scaled_linear":
          # this schedule is very specific to the latent diffusion model.
          A_    :			Any									=    (
              torch.linspace(beta_start**0.5   , beta_end**0.5   , snake_case_   , dtype=torch.floataa    ) ** 2
          )
        elif beta_schedule == "squaredcos_cap_v2":
          # Glide cosine schedule
          A_    :			Dict									=    betas_for_alpha_bar(snake_case_    )
        else:
          raise NotImplementedError(F"""{beta_schedule} does is not implemented for {self.__class__}"""    )

        A_    :			Union[str, Any]									=    1.0 - self.betas
        A_    :			List[Any]									=    torch.cumprod(self.alphas   , dim=0    )

        #  set all values
        self.set_timesteps(snake_case_   , snake_case_   , snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_=None    ):

        """simple docstring"""

        if schedule_timesteps is None:
          A_    :			List[str]									=    self.timesteps

        A_    :			Any									=    (schedule_timesteps == timestep).nonzero()

        # The sigma index that is taken for the **very** first `step`
        # is always the second index (or the last index if there is only 1)
        # This way we can ensure we don't accidentally skip a sigma in
        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
        if len(self._index_counter    ) == 0:
          A_    :			Optional[int]									=    1 if len(snake_case_    ) > 1 else 0
        else:
          A_    :			Union[str, Any]									=    timestep.cpu().item() if torch.is_tensor(snake_case_    ) else timestep
          A_    :			Tuple									=    self._index_counter[timestep_int]

        return indices[pos].item()

      @property
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        if self.config.timestep_spacing in ["linspace", "trailing"]:
          return self.sigmas.max()

        return (self.sigmas.max() ** 2 + 1) ** 0.5

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , ):

        """simple docstring"""

        A_    :			Any									=    self.index_for_timestep(snake_case_    )

        if self.state_in_first_order:
          A_    :			int									=    self.sigmas[step_index]
        else:
          A_    :			Optional[Any]									=    self.sigmas_interpol[step_index - 1]

        A_    :			int									=    sample / ((sigma**2 + 1) ** 0.5)
        return sample

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_ = None   , snake_case_ = None   , ):

        """simple docstring"""

        A_    :			Tuple									=    num_inference_steps

        A_    :			Any									=    num_train_timesteps or self.config.num_train_timesteps

        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
        if self.config.timestep_spacing == "linspace":
          A_    :			List[str]									=    np.linspace(0   , num_train_timesteps - 1   , snake_case_   , dtype=snake_case_    )[::-1].copy()
        elif self.config.timestep_spacing == "leading":
          A_    :			str									=    num_train_timesteps // self.num_inference_steps
          # creates integer timesteps by multiplying by ratio
          # casting to int to avoid issues when num_inference_step is power of 3
          A_    :			List[Any]									=    (np.arange(0   , snake_case_    ) * step_ratio).round()[::-1].copy().astype(snake_case_    )
          timesteps += self.config.steps_offset
        elif self.config.timestep_spacing == "trailing":
          A_    :			Optional[Any]									=    num_train_timesteps / self.num_inference_steps
          # creates integer timesteps by multiplying by ratio
          # casting to int to avoid issues when num_inference_step is power of 3
          A_    :			Any									=    (np.arange(snake_case_   , 0   , -step_ratio    )).round().copy().astype(snake_case_    )
          timesteps -= 1
        else:
          raise ValueError(
              F"""{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."""    )

        A_    :			Optional[Any]									=    np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5    )
        A_    :			int									=    torch.from_numpy(np.log(snake_case_    )    ).to(snake_case_    )

        A_    :			Optional[int]									=    np.interp(snake_case_   , np.arange(0   , len(snake_case_    )    )   , snake_case_    )
        A_    :			Dict									=    np.concatenate([sigmas, [0.0]]    ).astype(np.floataa    )
        A_    :			Tuple									=    torch.from_numpy(snake_case_    ).to(device=snake_case_    )

        # compute up and down sigmas
        A_    :			Optional[Any]									=    sigmas.roll(-1    )
        A_    :			List[Any]									=    0.0
        A_    :			Optional[Any]									=    (sigmas_next**2 * (sigmas**2 - sigmas_next**2) / sigmas**2) ** 0.5
        A_    :			List[str]									=    (sigmas_next**2 - sigmas_up**2) ** 0.5
        A_    :			Tuple									=    0.0

        # compute interpolated sigmas
        A_    :			Optional[Any]									=    sigmas.log().lerp(sigmas_down.log()   , 0.5    ).exp()
        A_    :			int									=    0.0

        # set sigmas
        A_    :			Union[str, Any]									=    torch.cat([sigmas[:1], sigmas[1:].repeat_interleave(2    ), sigmas[-1:]]    )
        A_    :			Optional[int]									=    torch.cat(
            [sigmas_interpol[:1], sigmas_interpol[1:].repeat_interleave(2    ), sigmas_interpol[-1:]]    )
        A_    :			List[Any]									=    torch.cat([sigmas_up[:1], sigmas_up[1:].repeat_interleave(2    ), sigmas_up[-1:]]    )
        A_    :			Any									=    torch.cat([sigmas_down[:1], sigmas_down[1:].repeat_interleave(2    ), sigmas_down[-1:]]    )

        if str(snake_case_    ).startswith('mps'    ):
          # mps does not support float64
          A_    :			Union[str, Any]									=    torch.from_numpy(snake_case_    ).to(snake_case_   , dtype=torch.floataa    )
        else:
          A_    :			Dict									=    torch.from_numpy(snake_case_    ).to(snake_case_    )

        A_    :			str									=    self.sigma_to_t(snake_case_    ).to(snake_case_   , dtype=timesteps.dtype    )
        A_    :			int									=    torch.stack((timesteps_interpol[:-2, None], timesteps[1:, None])   , dim=-1    ).flatten()

        A_    :			Dict									=    torch.cat([timesteps[:1], interleaved_timesteps]    )

        A_    :			int									=    None

        # for exp beta schedules, such as the one for `pipeline_shap_e.py`
        # we need an index counter
        A_    :			Union[str, Any]									=    defaultdict(snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_    :			int									=    sigma.log()

        # get distribution
        A_    :			int									=    log_sigma - self.log_sigmas[:, None]

        # get sigmas range
        A_    :			str									=    dists.ge(0    ).cumsum(dim=0    ).argmax(dim=0    ).clamp(max=self.log_sigmas.shape[0] - 2    )
        A_    :			str									=    low_idx + 1

        A_    :			Optional[int]									=    self.log_sigmas[low_idx]
        A_    :			Optional[Any]									=    self.log_sigmas[high_idx]

        # interpolate sigmas
        A_    :			int									=    (low - log_sigma) / (low - high)
        A_    :			Any									=    w.clamp(0   , 1    )

        # transform interpolation to time range
        A_    :			int									=    (1 - w) * low_idx + w * high_idx
        A_    :			Optional[int]									=    t.view(sigma.shape    )
        return t

      @property
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        return self.sample is None

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_   , snake_case_ = None   , snake_case_ = True   , ):

        """simple docstring"""

        A_    :			str									=    self.index_for_timestep(snake_case_    )

        # advance index counter by 1
        A_    :			Tuple									=    timestep.cpu().item() if torch.is_tensor(snake_case_    ) else timestep
        self._index_counter[timestep_int] += 1

        if self.state_in_first_order:
          A_    :			Optional[int]									=    self.sigmas[step_index]
          A_    :			Optional[Any]									=    self.sigmas_interpol[step_index]
          A_    :			Any									=    self.sigmas_up[step_index]
          A_    :			Tuple									=    self.sigmas_down[step_index - 1]
        else:
          # 2nd order / KPDM2's method
          A_    :			int									=    self.sigmas[step_index - 1]
          A_    :			str									=    self.sigmas_interpol[step_index - 1]
          A_    :			List[Any]									=    self.sigmas_up[step_index - 1]
          A_    :			List[Any]									=    self.sigmas_down[step_index - 1]

        # currently only gamma=0 is supported. This usually works best anyways.
        # We can support gamma in the future but then need to scale the timestep before
        # passing it to the model which requires a change in API
        A_    :			Dict									=    0
        A_    :			str									=    sigma * (gamma + 1)  # Note: sigma_hat == sigma for now

        A_    :			Optional[Any]									=    model_output.device
        A_    :			List[str]									=    randn_tensor(model_output.shape   , dtype=model_output.dtype   , device=snake_case_   , generator=snake_case_    )

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        if self.config.prediction_type == "epsilon":
          A_    :			List[str]									=    sigma_hat if self.state_in_first_order else sigma_interpol
          A_    :			int									=    sample - sigma_input * model_output
        elif self.config.prediction_type == "v_prediction":
          A_    :			Any									=    sigma_hat if self.state_in_first_order else sigma_interpol
          A_    :			Dict									=    model_output * (-sigma_input / (sigma_input**2 + 1) ** 0.5) + (
              sample / (sigma_input**2 + 1)
          )
        elif self.config.prediction_type == "sample":
          raise NotImplementedError('prediction_type not implemented yet: sample'    )
        else:
          raise ValueError(
              F"""prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"""    )

        if self.state_in_first_order:
          # 2. Convert to an ODE derivative for 1st order
          A_    :			List[Any]									=    (sample - pred_original_sample) / sigma_hat
          # 3. delta timestep
          A_    :			Optional[Any]									=    sigma_interpol - sigma_hat

          # store for 2nd order step
          A_    :			Optional[int]									=    sample
          A_    :			Union[str, Any]									=    dt
          A_    :			int									=    sample + derivative * dt
        else:
          # DPM-Solver-2
          # 2. Convert to an ODE derivative for 2nd order
          A_    :			List[Any]									=    (sample - pred_original_sample) / sigma_interpol
          # 3. delta timestep
          A_    :			List[str]									=    sigma_down - sigma_hat

          A_    :			Optional[int]									=    self.sample
          A_    :			Any									=    None

          A_    :			str									=    sample + derivative * dt
          A_    :			Optional[Any]									=    prev_sample + noise * sigma_up

        if not return_dict:
          return (prev_sample,)

        return SchedulerOutput(prev_sample=snake_case_    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_   , ):

        """simple docstring"""

        A_    :			Dict									=    self.sigmas.to(device=original_samples.device   , dtype=original_samples.dtype    )
        if original_samples.device.type == "mps" and torch.is_floating_point(snake_case_    ):
          # mps does not support float64
          A_    :			Optional[int]									=    self.timesteps.to(original_samples.device   , dtype=torch.floataa    )
          A_    :			Any									=    timesteps.to(original_samples.device   , dtype=torch.floataa    )
        else:
          A_    :			Optional[int]									=    self.timesteps.to(original_samples.device    )
          A_    :			Optional[Any]									=    timesteps.to(original_samples.device    )

        A_    :			str									=    [self.index_for_timestep(snake_case_   , snake_case_    ) for t in timesteps]

        A_    :			str									=    sigmas[step_indices].flatten()
        while len(sigma.shape    ) < len(original_samples.shape    ):
          A_    :			Optional[Any]									=    sigma.unsqueeze(-1    )

        A_    :			Dict									=    original_samples + noise * sigma
        return noisy_samples

      def __len__(	self    ):

        """simple docstring"""

        return self.config.num_train_timesteps






"""simple docstring"""




import multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class 					_UpperCAmelCase    (  UpperCAmelCase__						):



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_ = None   , snake_case_ = None   , snake_case_ = None   , snake_case_ = False   , snake_case_ = False   , snake_case_ = None   , **snake_case_   , ):

        """simple docstring"""

        super().__init__(
            snake_case_   , split=snake_case_   , features=snake_case_   , cache_dir=snake_case_   , keep_in_memory=snake_case_   , streaming=snake_case_   , num_proc=snake_case_   , **snake_case_   , )
        A_    :			int									=    path_or_paths if isinstance(snake_case_   , snake_case_    ) else {self.split: path_or_paths}
        A_    :			List[str]									=    Csv(
            cache_dir=snake_case_   , data_files=snake_case_   , features=snake_case_   , **snake_case_   , )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        if self.streaming:
          A_    :			List[str]									=    self.builder.as_streaming_dataset(split=self.split    )
        # Build regular (map-style) dataset
        else:
          A_    :			List[Any]									=    None
          A_    :			Optional[int]									=    None
          A_    :			Optional[int]									=    None
          A_    :			List[str]									=    None

          self.builder.download_and_prepare(
              download_config=snake_case_   , download_mode=snake_case_   , verification_mode=snake_case_   , base_path=snake_case_   , num_proc=self.num_proc   , )
          A_    :			List[Any]									=    self.builder.as_dataset(
              split=self.split   , verification_mode=snake_case_   , in_memory=self.keep_in_memory    )
        return dataset


class 					_UpperCAmelCase    :



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_   , snake_case_ = None   , snake_case_ = None   , **snake_case_   , ):

        """simple docstring"""

        if num_proc is not None and num_proc <= 0:
          raise ValueError(F"""num_proc {num_proc} must be an integer > 0."""    )

        A_    :			Optional[int]									=    dataset
        A_    :			List[str]									=    path_or_buf
        A_    :			Optional[int]									=    batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        A_    :			Optional[int]									=    num_proc
        A_    :			Optional[int]									=    'utf-8'
        A_    :			int									=    to_csv_kwargs

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    self.to_csv_kwargs.pop('path_or_buf'   , snake_case_    )
        A_    :			str									=    self.to_csv_kwargs.pop('header'   , snake_case_    )
        A_    :			Optional[Any]									=    self.to_csv_kwargs.pop('index'   , snake_case_    )

        if isinstance(self.path_or_buf   , (str, bytes, os.PathLike)    ):
          with open(self.path_or_buf   , 'wb+'    ) as buffer:
            A_    :			Optional[Any]									=    self._write(file_obj=snake_case_   , header=snake_case_   , index=snake_case_   , **self.to_csv_kwargs    )
        else:
          A_    :			int									=    self._write(file_obj=self.path_or_buf   , header=snake_case_   , index=snake_case_   , **self.to_csv_kwargs    )
        return written

      def    lowerCamelCase_     (	self   , snake_case_    ):

        """simple docstring"""

        A_		,			A_		,			A_		,			A_    :			Union[str, Any]									=    args

        A_    :			Dict									=    query_table(
            table=self.dataset.data   , key=slice(snake_case_   , offset + self.batch_size    )   , indices=self.dataset._indices   , )
        A_    :			List[str]									=    batch.to_pandas().to_csv(
            path_or_buf=snake_case_   , header=header if (offset == 0) else False   , index=snake_case_   , **snake_case_    )
        return csv_str.encode(self.encoding    )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_   , **snake_case_    ):

        """simple docstring"""

        A_    :			Any									=    0

        if self.num_proc is None or self.num_proc == 1:
          for offset in logging.tqdm(
              range(0   , len(self.dataset    )   , self.batch_size    )   , unit='ba'   , disable=not logging.is_progress_bar_enabled()   , desc='Creating CSV from Arrow format'   , ):
            A_    :			Optional[int]									=    self._batch_csv((offset, header, index, to_csv_kwargs)    )
            written += file_obj.write(snake_case_    )

        else:
          A_		,			A_    :			Optional[int]									=    len(self.dataset    ), self.batch_size
          with multiprocessing.Pool(self.num_proc    ) as pool:
            for csv_str in logging.tqdm(
                pool.imap(
                    self._batch_csv   , [(offset, header, index, to_csv_kwargs) for offset in range(0   , snake_case_   , snake_case_    )]   , )   , total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size   , unit='ba'   , disable=not logging.is_progress_bar_enabled()   , desc='Creating CSV from Arrow format'   , ):
              written += file_obj.write(snake_case_    )

        return written