File size: 1,750 Bytes
7344bef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | from typing import Any
import torch
def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6, in_place = False) -> torch.Tensor:
# deepbeepmeep RMS Norm
if in_place:
scale = torch.linalg.vector_norm(x, ord=2, dim=-1, keepdim=True, dtype=torch.float32)
scale.square_().div_(x.shape[-1]).add_(eps).rsqrt_()
x.mul_(scale)
if weight is not None:
x.mul_(weight)
return x
return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401
actual = config.get(key)
if actual != expected:
raise ValueError(f"Config value {key} is {actual}, expected {expected}")
def to_velocity(
sample: torch.Tensor,
sigma: float | torch.Tensor,
denoised_sample: torch.Tensor,
calc_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Convert the sample and its denoised version to velocity.
Returns:
Velocity
"""
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_dtype).item()
if sigma == 0:
raise ValueError("Sigma can't be 0.0")
return ((sample.to(calc_dtype) - denoised_sample.to(calc_dtype)) / sigma).to(sample.dtype)
def to_denoised(
sample: torch.Tensor,
velocity: torch.Tensor,
sigma: float | torch.Tensor,
calc_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Convert the sample and its denoising velocity to denoised sample.
Returns:
Denoised sample
"""
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_dtype)
return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)
|