<commit_before># Copyright 2022 NVIDIA and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from dataclasses import dataclass
from typing import Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .scheduling_utils import SchedulerMixin


@dataclass
class KarrasVeOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        derivative (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Derivate of predicted original image sample (x_0).
    """

    prev_sample: torch.FloatTensor
    derivative: torch.FloatTensor


class KarrasVeScheduler(SchedulerMixin, ConfigMixin):
    """
    Stochastic sampling from Karras et al. [1] tailored to the Variance-Expanding (VE) models [2]. Use Algorithm 2 and
    the VE column of Table 1 from [1] for reference.

    [1] <NAME>, et al. "Elucidating the Design Space of Diffusion-Based Generative Models."
    https://arxiv.org/abs/2206.00364 [2] <NAME>, et al. "Score-based generative modeling through stochastic
    differential equations." https://arxiv.org/abs/2011.13456
    """

    @register_to_config
    def __init__(
        self,
        sigma_min=0.02,
        sigma_max=100,
        s_noise=1.007,
        s_churn=80,
        s_min=0.05,
        s_max=50,
        tensor_format="pt",
    ):
        """
        For more details on the parameters, see the original paper's Appendix E.: "Elucidating the Design Space of
        Diffusion-Based Generative Models." https://arxiv.org/abs/2206.00364. The grid search values used to find the
        optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper.

        Args:
            sigma_min (`float`): minimum noise magnitude
            sigma_max (`float`): maximum noise magnitude
            s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling.
                A reasonable range is [1.000, 1.011].
            s_churn (`float`): the parameter controlling the overall amount of stochasticity.
                A reasonable range is [0, 100].
            s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity).
                A reasonable range is [0, 10].
            s_max (`float`): the end value of the sigma range where we add noise.
                A reasonable range is [0.2, 80].
        """
        # setable values
        self.num_inference_steps = None
        self.timesteps = None
        self.schedule = None  # sigma(t_i)

        self.tensor_format = tensor_format
        self.set_format(tensor_format=tensor_format)

    def set_timesteps(self, num_inference_steps):
        self.num_inference_steps = num_inference_steps
        self.timesteps = np.arange(0, self.num_inference_steps)[::-1].copy()
        self.schedule = [
            (self.sigma_max * (self.sigma_min**2 / self.sigma_max**2) ** (i / (num_inference_steps - 1)))
            for i in self.timesteps
        ]
        self.schedule = np.array(self.schedule, dtype=np.float32)

        self.set_format(tensor_format=self.tensor_format)

    def add_noise_to_input(self, sample, sigma, generator=None):
        """
        Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
        higher noise level sigma_hat = sigma_i + gamma_i*sigma_i.
        """<commit_msg>[Type hint] scheduling karras ve (#359)<commit_after># Copyright 2022 NVIDIA and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


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

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .scheduling_utils import SchedulerMixin


@dataclass
class KarrasVeOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        derivative (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Derivate of predicted original image sample (x_0).
    """

    prev_sample: torch.FloatTensor
    derivative: torch.FloatTensor


class KarrasVeScheduler(SchedulerMixin, ConfigMixin):
    """
    Stochastic sampling from Karras et al. [1] tailored to the Variance-Expanding (VE) models [2]. Use Algorithm 2 and
    the VE column of Table 1 from [1] for reference.

    [1] <NAME>, et al. "Elucidating the Design Space of Diffusion-Based Generative Models."
    https://arxiv.org/abs/2206.00364 [2] <NAME>, et al. "Score-based generative modeling through stochastic
    differential equations." https://arxiv.org/abs/2011.13456
    """

    @register_to_config
    def __init__(
        self,
        sigma_min: float = 0.02,
        sigma_max: float = 100,
        s_noise: float = 1.007,
        s_churn: float = 80,
        s_min: float = 0.05,
        s_max: float = 50,
        tensor_format: str = "pt",
    ):
        """
        For more details on the parameters, see the original paper's Appendix E.: "Elucidating the Design Space of
        Diffusion-Based Generative Models." https://arxiv.org/abs/2206.00364. The grid search values used to find the
        optimal {s_noise, s_churn, s_min, s_max} for a specific model are described in Table 5 of the paper.

        Args:
            sigma_min (`float`): minimum noise magnitude
            sigma_max (`float`): maximum noise magnitude
            s_noise (`float`): the amount of additional noise to counteract loss of detail during sampling.
                A reasonable range is [1.000, 1.011].
            s_churn (`float`): the parameter controlling the overall amount of stochasticity.
                A reasonable range is [0, 100].
            s_min (`float`): the start value of the sigma range where we add noise (enable stochasticity).
                A reasonable range is [0, 10].
            s_max (`float`): the end value of the sigma range where we add noise.
                A reasonable range is [0.2, 80].
        """
        # setable values
        self.num_inference_steps = None
        self.timesteps = None
        self.schedule = None  # sigma(t_i)

        self.tensor_format = tensor_format
        self.set_format(tensor_format=tensor_format)

    def set_timesteps(self, num_inference_steps: int):
        self.num_inference_steps = num_inference_steps
        self.timesteps = np.arange(0, self.num_inference_steps)[::-1].copy()
        self.schedule = [
            (self.sigma_max * (self.sigma_min**2 / self.sigma_max**2) ** (i / (num_inference_steps - 1)))
            for i in self.timesteps
        ]
        self.schedule = np.array(self.schedule, dtype=np.float32)

        self.set_format(tensor_format=self.tensor_format)

    def add_noise_to_input(
        self, sample: Union[torch.FloatTensor, np.ndarray], sigma: float, generator: Optional[torch.Generator] = None
    ) -> Tuple[Union[torch.FloatTensor, np.ndarray], float]:
        """
        Explicit Langevin-like "churn" step of adding noise to the sample according to a factor gamma_i ≥ 0 to reach a
        higher noise level sigma_hat = sigma_i + gamma_i*sigma_i.
        """
{'fraction_non_alphanumeric': 0.08639333183149359, 'fraction_numerical': 0.020387474656454158, 'mean_word_length': 2.825506247307195, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 8, 'https://': 6, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5556405', 'n_tokens_mistral': 2720, 'n_tokens_neox': 2492, 'n_words': 918}