Repository: huggingface/diffusers
File: src/diffusers/schedulers/scheduling_dpmsolver_sde.py (Python)

Code before:
        """
        Sorts two float values and returns them along with a sign indicating if they were swapped.

        Args:
            a (`float`):
                The first value.
            b (`float`):
                The second value.

        Returns:
            `Tuple[float, float, float]`:
                A tuple containing the sorted values (min, max) and a sign (1.0 if a < b, -1.0 otherwise).
        """
        return (a, b, 1.0) if a < b else (b, a, -1.0)

    def __call__(self, t0: float, t1: float) -> torch.Tensor:
        t0, t1, sign = self.sort(t0, t1)
        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
        return w if self.batched else w[0]


class BrownianTreeNoiseSampler:
    """A noise sampler backed by a torchsde.BrownianTree.

    Args:
        x (`torch.Tensor`): The tensor whose shape, device and dtype to use to generate random samples.
        sigma_min (`float`): The low end of the valid interval.
        sigma_max (`float`): The high end of the valid interval.
        seed (`int` or `List[int]`): The random seed. If a list of seeds is
            supplied instead of a single integer, then the noise sampler will use one BrownianTree per batch item, each
            with its own seed.
        transform (`callable`): A function that maps sigma to the sampler's
            internal timestep.
    """

    def __init__(
        self,
        x: torch.Tensor,
        sigma_min: float,
        sigma_max: float,
        seed: Optional[Union[int, List[int]]] = None,
        transform: Callable[[float], float] = lambda x: x,
    ):
        self.transform = transform
        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
        self.tree = BatchedBrownianTree(x, t0, t1, seed)

    def __call__(self, sigma: float, sigma_next: float) -> torch.Tensor:
        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()


Review comment (suggestion):
```suggestion
        x (`torch.Tensor`): The tensor whose shape, device and dtype is used to generate random samples.
```

Code after:
        """
        Sorts two float values and returns them along with a sign indicating if they were swapped.

        Args:
            a (`float`):
                The first value.
            b (`float`):
                The second value.

        Returns:
            `Tuple[float, float, float]`:
                A tuple containing the sorted values (min, max) and a sign (1.0 if a < b, -1.0 otherwise).
        """
        return (a, b, 1.0) if a < b else (b, a, -1.0)

    def __call__(self, t0: float, t1: float) -> torch.Tensor:
        t0, t1, sign = self.sort(t0, t1)
        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
        return w if self.batched else w[0]


class BrownianTreeNoiseSampler:
    """A noise sampler backed by a torchsde.BrownianTree.

    Args:
        x (`torch.Tensor`): The tensor whose shape, device and dtype is used to generate random samples.
        sigma_min (`float`): The low end of the valid interval.
        sigma_max (`float`): The high end of the valid interval.
        seed (`int` or `List[int]`): The random seed. If a list of seeds is
            supplied instead of a single integer, then the noise sampler will use one BrownianTree per batch item, each
            with its own seed.
        transform (`callable`): A function that maps sigma to the sampler's
            internal timestep.
    """

    def __init__(
        self,
        x: torch.Tensor,
        sigma_min: float,
        sigma_max: float,
        seed: Optional[Union[int, List[int]]] = None,
        transform: Callable[[float], float] = lambda x: x,
    ):
        self.transform = transform
        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
        self.tree = BatchedBrownianTree(x, t0, t1, seed)

    def __call__(self, sigma: float, sigma_next: float) -> torch.Tensor:
        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()