59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
"""
|
|
-----------------------------------------------------------------------------
|
|
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
and proprietary rights in and to this software, related documentation
|
|
and any modifications thereto. Any use, reproduction, disclosure or
|
|
distribution of this software and related documentation without an express
|
|
license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
-----------------------------------------------------------------------------
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
class FlowMatchingScheduler:
|
|
def __init__(self, num_train_timesteps: int = 1000, shift: float = 1):
|
|
# set timesteps
|
|
self.num_train_timesteps = num_train_timesteps
|
|
self.shift = shift
|
|
|
|
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
|
|
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
|
|
|
|
sigmas = timesteps / num_train_timesteps
|
|
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
|
|
|
self.sigmas = sigmas # 1 --> 0
|
|
self.timesteps = sigmas * num_train_timesteps # num_train_timesteps --> 1
|
|
|
|
# set device
|
|
def to(self, device):
|
|
self.sigmas = self.sigmas.to(device=device)
|
|
self.timesteps = self.timesteps.to(device=device)
|
|
|
|
# add random noise to latent during training
|
|
def add_noise(self, latent: torch.Tensor, logit_mean: float = 1.0, logit_std: float = 1.0):
|
|
# latent: [B, ...]
|
|
# timesteps: [B]
|
|
# return: [B, ...] noisy_latent, [B, ...] noise, [B] timesteps
|
|
|
|
# logit-normal sampling
|
|
u = torch.normal(mean=logit_mean, std=logit_std, size=(latent.shape[0],), device=self.sigmas.device)
|
|
u = torch.nn.functional.sigmoid(u)
|
|
|
|
step_indices = (u * self.num_train_timesteps).long()
|
|
timesteps = self.timesteps[step_indices]
|
|
|
|
sigmas = self.sigmas[step_indices].flatten()
|
|
|
|
while len(sigmas.shape) < latent.ndim:
|
|
sigmas = sigmas.unsqueeze(-1)
|
|
|
|
noise = torch.randn_like(latent)
|
|
noisy_latent = (1.0 - sigmas) * latent + sigmas * noise
|
|
|
|
return noisy_latent, noise, timesteps
|