147 lines
4.7 KiB
Python
147 lines
4.7 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# This source code is licensed under the CC BY-NC 4.0 license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
#
|
|
# This file is derived from https://github.com/facebookresearch/flow_matching
|
|
# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
|
|
|
|
# pyre-unsafe
|
|
|
|
import math
|
|
from dataclasses import dataclass, field
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.distributions import LogisticNormal, Uniform
|
|
|
|
|
|
def get_lin_function(
|
|
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
|
):
|
|
m = (y2 - y1) / (x2 - x1)
|
|
b = y1 - m * x1
|
|
return lambda x: m * x + b
|
|
|
|
|
|
def time_shift(mu: float, sigma: float, t: Tensor):
|
|
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
|
|
|
|
|
@dataclass
|
|
class SchedulerOutput:
|
|
r"""Represents a sample of a conditional-flow generated probability path.
|
|
|
|
Attributes:
|
|
alpha_t (Tensor): :math:`\alpha_t`, shape (...).
|
|
sigma_t (Tensor): :math:`\sigma_t`, shape (...).
|
|
d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
|
|
d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
|
|
|
|
"""
|
|
|
|
alpha_t: Tensor = field(metadata={"help": "alpha_t"})
|
|
sigma_t: Tensor = field(metadata={"help": "sigma_t"})
|
|
d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
|
|
d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
|
|
|
|
|
|
class CondOTScheduler:
|
|
"""CondOT Scheduler."""
|
|
|
|
def __call__(self, t: Tensor) -> SchedulerOutput:
|
|
r"""
|
|
Args:
|
|
t (Tensor): times in [0,1], shape (...).
|
|
|
|
Returns:
|
|
SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
|
|
"""
|
|
return SchedulerOutput(
|
|
alpha_t=t,
|
|
sigma_t=1 - t,
|
|
d_alpha_t=torch.ones_like(t),
|
|
d_sigma_t=-torch.ones_like(t),
|
|
)
|
|
|
|
def snr_inverse(self, snr: Tensor) -> Tensor:
|
|
r"""
|
|
Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
|
|
|
|
Args:
|
|
snr (Tensor): The signal-to-noise, shape (...)
|
|
|
|
Returns:
|
|
Tensor: t, shape (...)
|
|
"""
|
|
kappa_t = snr / (1.0 + snr)
|
|
|
|
return self.kappa_inverse(kappa=kappa_t)
|
|
|
|
def kappa_inverse(self, kappa: Tensor) -> Tensor:
|
|
"""
|
|
Computes :math:`t` from :math:`kappa_t`.
|
|
|
|
Args:
|
|
kappa (Tensor): :math:`kappa`, shape (...)
|
|
|
|
Returns:
|
|
Tensor: t, shape (...)
|
|
"""
|
|
return kappa
|
|
|
|
|
|
class TimeSampler:
|
|
def __init__(self, dist):
|
|
self.dist_name = dist
|
|
if dist == "uniform":
|
|
self.distribution = Uniform(0.0, 1.0)
|
|
elif dist == "lognormal":
|
|
self.distribution = LogisticNormal(0.0, 1.0)
|
|
elif dist == "flux":
|
|
self.distribution = FluxTimeSampler(mode="train")
|
|
else:
|
|
raise ValueError(f"Unknown distribution {dist}")
|
|
|
|
def __call__(self, shape, num_tokens, device):
|
|
if self.dist_name == "uniform":
|
|
return self.distribution.sample(sample_shape=(shape,)).to(device).squeeze(0)
|
|
elif self.dist_name == "lognormal":
|
|
return self.distribution.sample(sample_shape=(shape,))[:, 0].to(device)
|
|
elif self.dist_name == "flux":
|
|
return self.distribution(shape, num_tokens, device)
|
|
|
|
|
|
class FluxTimeSampler:
|
|
def __init__(self, **kwargs):
|
|
self.mode = kwargs.get("mode", "train")
|
|
|
|
def __call__(self, bs, num_tokens, device):
|
|
if self.mode == "train":
|
|
logits_norm = torch.randn(bs, device=device)
|
|
timesteps = logits_norm.sigmoid()
|
|
mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
|
|
timesteps = time_shift(mu, 1.0, timesteps)
|
|
timesteps = 1 - timesteps
|
|
elif self.mode == "inference":
|
|
timesteps = torch.linspace(1, 0, bs + 1).to(device)
|
|
mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
|
|
timesteps = time_shift(mu, 1.0, timesteps)
|
|
timesteps = 1 - timesteps
|
|
return timesteps
|
|
|
|
|
|
def visualize_sampling():
|
|
import matplotlib.pyplot as plt
|
|
|
|
flux_sampler = FluxTimeSampler(mode="inference")
|
|
timesteps = flux_sampler(25, 2048, torch.device("cuda"))
|
|
x = torch.linspace(0, 1, 25 + 1)
|
|
plt.scatter(x, y=timesteps.cpu().numpy())
|
|
plt.savefig("scatter_plot.png")
|
|
|
|
flux_sampler = FluxTimeSampler(mode="train")
|
|
timesteps = flux_sampler(5000, 4096, torch.device("cuda"))
|
|
plt.hist(timesteps.cpu().numpy(), bins=50, edgecolor="black")
|
|
plt.savefig("scatter_plot_4096.png")
|
|
print(timesteps)
|