Merge branch 'PKU-YuanGroup:main' into videogpt
Former-commit-id: 339857e1b672a1a688f86c1b8e1d8dae29f10301
This commit is contained in:
Chestnut
@@ -1,6 +1,6 @@
|
||||
# Open-Sora Plan
|
||||
|
||||
[[Project Page]](https://pku-yuangroup.github.io/Open-Sora-Plan/) [[中文主页]](https://pku-yuangroup.github.io/Open-Sora-Plan/blog_cn.html) [[Discord]](https://discord.gg/fqpmStRX) [[Wechat Group]](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/53#issuecomment-1980312563)
|
||||
[[Project Page]](https://pku-yuangroup.github.io/Open-Sora-Plan/) [[中文主页]](https://pku-yuangroup.github.io/Open-Sora-Plan/blog_cn.html) [[Discord]](https://discord.gg/fqpmStRX) [[Wechat Group]](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/53#issuecomment-1980312563)
|
||||
|
||||
## Goal
|
||||
This project aims to create a simple and scalable repo, to reproduce [Sora](https://openai.com/sora) (OpenAI, but we prefer to call it "CloseAI" ) and build knowledge about Video-VQVAE (VideoGPT) + DiT at scale. However, we have limited resources, we deeply wish all open-source community can contribute to this project. Pull request are welcome!!!
|
||||
@@ -38,6 +38,7 @@ Project stages:
|
||||
## Todo
|
||||
|
||||
#### Setup the codebase and train a unconditional model on landscape dataset
|
||||
- [x] Fix typo & Update readme. 🤝 Thanks to [@mio2333](https://github.com/mio2333), [@CreamyLong](https://github.com/CreamyLong), [@chg0901](https://github.com/chg0901)
|
||||
- [x] Setup repo-structure.
|
||||
- [x] Add Video-VQGAN model, which is borrowed from [VideoGPT](https://github.com/wilson1yan/VideoGPT).
|
||||
- [x] Support variable aspect ratios, resolutions, durations training on [DiT](https://github.com/facebookresearch/DiT).
|
||||
@@ -47,16 +48,18 @@ Project stages:
|
||||
- [x] Add VAE model, which is borrowed from [Stable Diffusion](https://github.com/CompVis/latent-diffusion).
|
||||
- [x] Joint dynamic mask input with VAE.
|
||||
- [ ] Make the codebase ready for the cluster training. Add SLURM scripts.
|
||||
- [x] Refactor VideoGPT. 🤝 Thanks to [@qqingzheng](https://github.com/qqingzheng), [@luo3300612](https://github.com/luo3300612)
|
||||
- [x] Add sampling script.
|
||||
- [ ] Incorporating [SiT](https://github.com/willisma/SiT).
|
||||
- [ ] Incorporate [SiT](https://github.com/willisma/SiT). 🤝 Thanks to [@khan-yin](https://github.com/khan-yin)
|
||||
|
||||
#### Train models that boost resolution and duration
|
||||
- [ ] Add [PI](https://arxiv.org/abs/2306.15595) to support out-of-domain size.
|
||||
- [x] Extract offline feature.
|
||||
- [x] Add frame interpolation model.
|
||||
- [x] Add super resolution model.
|
||||
- [x] Add frame interpolation model. 🤝 Thanks to [@yunyangge](https://github.com/yunyangge)
|
||||
- [x] Add super resolution model. 🤝 Thanks to [@Linzy19](https://github.com/Linzy19)
|
||||
- [x] Add accelerate to automatically manage training.
|
||||
- [ ] Joint training with images.
|
||||
- [ ] Incorporate [NaViT](https://arxiv.org/abs/2307.06304).
|
||||
|
||||
#### Conduct text2video experiments on landscape dataset.
|
||||
- [ ] Finish data loading, pre-processing utils.
|
||||
@@ -72,6 +75,7 @@ Project stages:
|
||||
- [x] Add xformers.
|
||||
- [x] Support mixed precision training.
|
||||
- [x] Add gradient checkpoint.
|
||||
- [ ] Support for ReBased and Ring attention. 🤝 Thanks to [@kabachuha](https://github.com/kabachuha)
|
||||
- [ ] Train using the deepspeed engine.
|
||||
|
||||
#### Control model with more condition
|
||||
@@ -195,6 +199,7 @@ We greatly appreciate your contributions to the Open-Sora Plan open-source commu
|
||||
|
||||
For more details, please refer to the [Contribution Guidelines](docs/Contribution_Guidelines.md)
|
||||
|
||||
|
||||
## Acknowledgement
|
||||
* [Latte](https://github.com/Vchitect/Latte): The **main codebase** we built upon and it is an wonderful video gererated model.
|
||||
* [DiT](https://github.com/facebookresearch/DiT): Scalable Diffusion Models with Transformers.
|
||||
@@ -205,8 +210,3 @@ For more details, please refer to the [Contribution Guidelines](docs/Contributio
|
||||
## License
|
||||
* The service is a research preview intended for non-commercial use only. See [LICENSE](LICENSE) for details.
|
||||
|
||||
## Contributors
|
||||
|
||||
<a href="https://github.com/PKU-YuanGroup/Open-Sora-Plan/graphs/contributors">
|
||||
<img src="https://contrib.rocks/image?repo=PKU-YuanGroup/Open-Sora-Plan" />
|
||||
</a>
|
||||
|
||||
@@ -10,19 +10,33 @@ from .transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVi
|
||||
from .ucf101 import UCF101
|
||||
from .sky_datasets import Sky
|
||||
|
||||
ae_norm = {
|
||||
"bair_stride4x2x2": Lambda(lambda x: x - 0.5),
|
||||
"ucf101_stride4x4x4": Lambda(lambda x: x - 0.5),
|
||||
"kinetics_stride4x4x4": Lambda(lambda x: x - 0.5),
|
||||
"kinetics_stride2x4x4": Lambda(lambda x: x - 0.5),
|
||||
'stabilityai/sd-vae-ft-mse': transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
|
||||
'stabilityai/sd-vae-ft-ema': transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
|
||||
'vqgan_imagenet_f16_1024': Lambda(lambda x: 2. * x - 1.),
|
||||
'vqgan_imagenet_f16_16384': Lambda(lambda x: 2. * x - 1.),
|
||||
'vqgan_gumbel_f8': Lambda(lambda x: 2. * x - 1.),
|
||||
|
||||
}
|
||||
ae_denorm = {
|
||||
"bair_stride4x2x2": lambda x: x + 0.5,
|
||||
"ucf101_stride4x4x4": lambda x: x + 0.5,
|
||||
"kinetics_stride4x4x4": lambda x: x + 0.5,
|
||||
"kinetics_stride2x4x4": lambda x: x + 0.5,
|
||||
'stabilityai/sd-vae-ft-mse': lambda x: 0.5 * x + 0.5,
|
||||
'stabilityai/sd-vae-ft-ema': lambda x: 0.5 * x + 0.5,
|
||||
'vqgan_imagenet_f16_1024': lambda x: (x + 1.) / 2.,
|
||||
'vqgan_imagenet_f16_16384': lambda x: (x + 1.) / 2.,
|
||||
'vqgan_gumbel_f8': lambda x: (x + 1.) / 2.,
|
||||
}
|
||||
|
||||
def getdataset(args):
|
||||
temporal_sample = TemporalRandomCrop(args.num_frames * args.sample_rate) # 16 x
|
||||
if args.ae in videovqvae:
|
||||
norm_fun = Lambda(lambda x: x - 0.5)
|
||||
elif args.ae in videovae:
|
||||
raise NotImplementedError
|
||||
elif args.ae in vae:
|
||||
norm_fun = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
|
||||
elif args.ae in vqvae:
|
||||
norm_fun = Lambda(lambda x: 2. * x - 1.)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
norm_fun = ae_norm[args.ae]
|
||||
|
||||
if args.dataset == 'landscope_feature':
|
||||
temporal_sample = TemporalRandomCrop(args.num_frames) # 16 1
|
||||
|
||||
@@ -17,7 +17,13 @@ class HFVAEWrapper(nn.Module):
|
||||
x = rearrange(x, '(b t) c h w -> b t c h w', t=t).contiguous()
|
||||
return x
|
||||
def decode(self, x):
|
||||
t = 0
|
||||
if x.ndim == 5:
|
||||
b, t, c, h, w = x.shape
|
||||
x = rearrange(x, 'b t c h w -> (b t) c h w').contiguous()
|
||||
x = self.vae.decode(x / 0.18215).sample
|
||||
if t != 0:
|
||||
x = rearrange(x, '(b t) c h w -> b t c h w', t=t).contiguous()
|
||||
return x
|
||||
|
||||
class SDVAEWrapper(nn.Module):
|
||||
|
||||
@@ -13,6 +13,8 @@ class VideoGPTVQVAEWrapper(nn.Module):
|
||||
x = rearrange(x, 'b c t h w -> b t c h w').contiguous()
|
||||
return x
|
||||
def decode(self, x):
|
||||
x = rearrange(x, 'b t c h w -> b c t h w').contiguous()
|
||||
vq_output = self.vqvae.codebook(x)
|
||||
x = self.vqvae.decoder(self.vqvae.post_vq_conv(vq_output['embeddings']))
|
||||
x = rearrange(x, 'b c t h w -> b t c h w').contiguous()
|
||||
return x
|
||||
@@ -89,8 +89,9 @@ def main(args):
|
||||
x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size)
|
||||
x_vae = x_vae.to(device)
|
||||
|
||||
encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True)
|
||||
video_recon = vqvae.decode(encodings)
|
||||
with torch.no_grad():
|
||||
encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True)
|
||||
video_recon = vqvae.decode(encodings)
|
||||
|
||||
# custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4')
|
||||
custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)
|
||||
|
||||
@@ -267,7 +267,8 @@ class DiT(nn.Module):
|
||||
nn.init.constant_(self.x_embedder.proj.bias, 0)
|
||||
|
||||
# Initialize label embedding table:
|
||||
nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
|
||||
if self.extras == 2:
|
||||
nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
|
||||
|
||||
# Initialize timestep embedding MLP:
|
||||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||||
@@ -308,7 +309,8 @@ class DiT(nn.Module):
|
||||
attention_mask = attention_mask.flatten(1).unsqueeze(-1) # bs t h w -> bs thw 1
|
||||
attention_mask = attention_mask @ attention_mask.transpose(1, 2) # bs thw 1 @ bs 1 thw = bs thw thw
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(dtype).min)
|
||||
# attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(dtype).min)
|
||||
attention_mask = attention_mask.masked_fill(attention_mask == 0, 1e-8)
|
||||
return attention_mask
|
||||
|
||||
def forward(self, x, t, y, attention_mask):
|
||||
|
||||
@@ -350,7 +350,8 @@ class Latte(nn.Module):
|
||||
attention_mask = attention_mask.flatten(1).unsqueeze(-1) # bs t h w -> bs thw 1
|
||||
attention_mask = attention_mask @ attention_mask.transpose(1, 2) # bs thw 1 @ bs 1 thw = bs thw thw
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(dtype).min)
|
||||
attention_mask = attention_mask.masked_fill(attention_mask == 0, 1e-8)
|
||||
# attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(dtype).min)
|
||||
return attention_mask
|
||||
|
||||
# @torch.cuda.amp.autocast()
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
from .transport import Transport, ModelType, WeightType, PathType, Sampler
|
||||
|
||||
def create_transport(
|
||||
path_type='Linear',
|
||||
prediction="velocity",
|
||||
loss_weight=None,
|
||||
train_eps=None,
|
||||
sample_eps=None,
|
||||
):
|
||||
"""function for creating Transport object
|
||||
**Note**: model prediction defaults to velocity
|
||||
Args:
|
||||
- path_type: type of path to use; default to linear
|
||||
- learn_score: set model prediction to score
|
||||
- learn_noise: set model prediction to noise
|
||||
- velocity_weighted: weight loss by velocity weight
|
||||
- likelihood_weighted: weight loss by likelihood weight
|
||||
- train_eps: small epsilon for avoiding instability during training
|
||||
- sample_eps: small epsilon for avoiding instability during sampling
|
||||
"""
|
||||
|
||||
if prediction == "noise":
|
||||
model_type = ModelType.NOISE
|
||||
elif prediction == "score":
|
||||
model_type = ModelType.SCORE
|
||||
else:
|
||||
model_type = ModelType.VELOCITY
|
||||
|
||||
if loss_weight == "velocity":
|
||||
loss_type = WeightType.VELOCITY
|
||||
elif loss_weight == "likelihood":
|
||||
loss_type = WeightType.LIKELIHOOD
|
||||
else:
|
||||
loss_type = WeightType.NONE
|
||||
|
||||
path_choice = {
|
||||
"Linear": PathType.LINEAR,
|
||||
"GVP": PathType.GVP,
|
||||
"VP": PathType.VP,
|
||||
}
|
||||
|
||||
path_type = path_choice[path_type]
|
||||
|
||||
if (path_type in [PathType.VP]):
|
||||
train_eps = 1e-5 if train_eps is None else train_eps
|
||||
sample_eps = 1e-3 if train_eps is None else sample_eps
|
||||
elif (path_type in [PathType.GVP, PathType.LINEAR] and model_type != ModelType.VELOCITY):
|
||||
train_eps = 1e-3 if train_eps is None else train_eps
|
||||
sample_eps = 1e-3 if train_eps is None else sample_eps
|
||||
else: # velocity & [GVP, LINEAR] is stable everywhere
|
||||
train_eps = 0
|
||||
sample_eps = 0
|
||||
|
||||
# create flow state
|
||||
state = Transport(
|
||||
model_type=model_type,
|
||||
path_type=path_type,
|
||||
loss_type=loss_type,
|
||||
train_eps=train_eps,
|
||||
sample_eps=sample_eps,
|
||||
)
|
||||
|
||||
return state
|
||||
@@ -0,0 +1,117 @@
|
||||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
from torchdiffeq import odeint
|
||||
from functools import partial
|
||||
from tqdm import tqdm
|
||||
|
||||
class sde:
|
||||
"""SDE solver class"""
|
||||
def __init__(
|
||||
self,
|
||||
drift,
|
||||
diffusion,
|
||||
*,
|
||||
t0,
|
||||
t1,
|
||||
num_steps,
|
||||
sampler_type,
|
||||
):
|
||||
assert t0 < t1, "SDE sampler has to be in forward time"
|
||||
|
||||
self.num_timesteps = num_steps
|
||||
self.t = th.linspace(t0, t1, num_steps)
|
||||
self.dt = self.t[1] - self.t[0]
|
||||
self.drift = drift
|
||||
self.diffusion = diffusion
|
||||
self.sampler_type = sampler_type
|
||||
|
||||
def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs):
|
||||
w_cur = th.randn(x.size()).to(x)
|
||||
t = th.ones(x.size(0)).to(x) * t
|
||||
dw = w_cur * th.sqrt(self.dt)
|
||||
drift = self.drift(x, t, model, **model_kwargs)
|
||||
diffusion = self.diffusion(x, t)
|
||||
mean_x = x + drift * self.dt
|
||||
x = mean_x + th.sqrt(2 * diffusion) * dw
|
||||
return x, mean_x
|
||||
|
||||
def __Heun_step(self, x, _, t, model, **model_kwargs):
|
||||
w_cur = th.randn(x.size()).to(x)
|
||||
dw = w_cur * th.sqrt(self.dt)
|
||||
t_cur = th.ones(x.size(0)).to(x) * t
|
||||
diffusion = self.diffusion(x, t_cur)
|
||||
xhat = x + th.sqrt(2 * diffusion) * dw
|
||||
K1 = self.drift(xhat, t_cur, model, **model_kwargs)
|
||||
xp = xhat + self.dt * K1
|
||||
K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs)
|
||||
return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step
|
||||
|
||||
def __forward_fn(self):
|
||||
"""TODO: generalize here by adding all private functions ending with steps to it"""
|
||||
sampler_dict = {
|
||||
"Euler": self.__Euler_Maruyama_step,
|
||||
"Heun": self.__Heun_step,
|
||||
}
|
||||
|
||||
try:
|
||||
sampler = sampler_dict[self.sampler_type]
|
||||
except:
|
||||
raise NotImplementedError("Smapler type not implemented.")
|
||||
|
||||
return sampler
|
||||
|
||||
def sample(self, init, model, **model_kwargs):
|
||||
"""forward loop of sde"""
|
||||
x = init
|
||||
mean_x = init
|
||||
samples = []
|
||||
sampler = self.__forward_fn()
|
||||
for ti in self.t[:-1]:
|
||||
with th.no_grad():
|
||||
x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs)
|
||||
samples.append(x)
|
||||
|
||||
return samples
|
||||
|
||||
class ode:
|
||||
"""ODE solver class"""
|
||||
def __init__(
|
||||
self,
|
||||
drift,
|
||||
*,
|
||||
t0,
|
||||
t1,
|
||||
sampler_type,
|
||||
num_steps,
|
||||
atol,
|
||||
rtol,
|
||||
):
|
||||
assert t0 < t1, "ODE sampler has to be in forward time"
|
||||
|
||||
self.drift = drift
|
||||
self.t = th.linspace(t0, t1, num_steps)
|
||||
self.atol = atol
|
||||
self.rtol = rtol
|
||||
self.sampler_type = sampler_type
|
||||
|
||||
def sample(self, x, model, **model_kwargs):
|
||||
|
||||
device = x[0].device if isinstance(x, tuple) else x.device
|
||||
def _fn(t, x):
|
||||
t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t
|
||||
model_output = self.drift(x, t, model, **model_kwargs)
|
||||
return model_output
|
||||
|
||||
t = self.t.to(device)
|
||||
atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
|
||||
rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
|
||||
samples = odeint(
|
||||
_fn,
|
||||
x,
|
||||
t,
|
||||
method=self.sampler_type,
|
||||
atol=atol,
|
||||
rtol=rtol
|
||||
)
|
||||
return samples
|
||||
@@ -0,0 +1,192 @@
|
||||
import torch as th
|
||||
import numpy as np
|
||||
from functools import partial
|
||||
|
||||
def expand_t_like_x(t, x):
|
||||
"""Function to reshape time t to broadcastable dimension of x
|
||||
Args:
|
||||
t: [batch_dim,], time vector
|
||||
x: [batch_dim,...], data point
|
||||
"""
|
||||
dims = [1] * (len(x.size()) - 1)
|
||||
t = t.view(t.size(0), *dims)
|
||||
return t
|
||||
|
||||
|
||||
#################### Coupling Plans ####################
|
||||
|
||||
class ICPlan:
|
||||
"""Linear Coupling Plan"""
|
||||
def __init__(self, sigma=0.0):
|
||||
self.sigma = sigma
|
||||
|
||||
def compute_alpha_t(self, t):
|
||||
"""Compute the data coefficient along the path"""
|
||||
return t, 1
|
||||
|
||||
def compute_sigma_t(self, t):
|
||||
"""Compute the noise coefficient along the path"""
|
||||
return 1 - t, -1
|
||||
|
||||
def compute_d_alpha_alpha_ratio_t(self, t):
|
||||
"""Compute the ratio between d_alpha and alpha"""
|
||||
return 1 / t
|
||||
|
||||
def compute_drift(self, x, t):
|
||||
"""We always output sde according to score parametrization; """
|
||||
t = expand_t_like_x(t, x)
|
||||
alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
|
||||
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
||||
drift = alpha_ratio * x
|
||||
diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
|
||||
|
||||
return -drift, diffusion
|
||||
|
||||
def compute_diffusion(self, x, t, form="constant", norm=1.0):
|
||||
"""Compute the diffusion term of the SDE
|
||||
Args:
|
||||
x: [batch_dim, ...], data point
|
||||
t: [batch_dim,], time vector
|
||||
form: str, form of the diffusion term
|
||||
norm: float, norm of the diffusion term
|
||||
"""
|
||||
t = expand_t_like_x(t, x)
|
||||
choices = {
|
||||
"constant": norm,
|
||||
"SBDM": norm * self.compute_drift(x, t)[1],
|
||||
"sigma": norm * self.compute_sigma_t(t)[0],
|
||||
"linear": norm * (1 - t),
|
||||
"decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2,
|
||||
"inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2,
|
||||
}
|
||||
|
||||
try:
|
||||
diffusion = choices[form]
|
||||
except KeyError:
|
||||
raise NotImplementedError(f"Diffusion form {form} not implemented")
|
||||
|
||||
return diffusion
|
||||
|
||||
def get_score_from_velocity(self, velocity, x, t):
|
||||
"""Wrapper function: transfrom velocity prediction model to score
|
||||
Args:
|
||||
velocity: [batch_dim, ...] shaped tensor; velocity model output
|
||||
x: [batch_dim, ...] shaped tensor; x_t data point
|
||||
t: [batch_dim,] time tensor
|
||||
"""
|
||||
t = expand_t_like_x(t, x)
|
||||
alpha_t, d_alpha_t = self.compute_alpha_t(t)
|
||||
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
||||
mean = x
|
||||
reverse_alpha_ratio = alpha_t / d_alpha_t
|
||||
var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
|
||||
score = (reverse_alpha_ratio * velocity - mean) / var
|
||||
return score
|
||||
|
||||
def get_noise_from_velocity(self, velocity, x, t):
|
||||
"""Wrapper function: transfrom velocity prediction model to denoiser
|
||||
Args:
|
||||
velocity: [batch_dim, ...] shaped tensor; velocity model output
|
||||
x: [batch_dim, ...] shaped tensor; x_t data point
|
||||
t: [batch_dim,] time tensor
|
||||
"""
|
||||
t = expand_t_like_x(t, x)
|
||||
alpha_t, d_alpha_t = self.compute_alpha_t(t)
|
||||
sigma_t, d_sigma_t = self.compute_sigma_t(t)
|
||||
mean = x
|
||||
reverse_alpha_ratio = alpha_t / d_alpha_t
|
||||
var = reverse_alpha_ratio * d_sigma_t - sigma_t
|
||||
noise = (reverse_alpha_ratio * velocity - mean) / var
|
||||
return noise
|
||||
|
||||
def get_velocity_from_score(self, score, x, t):
|
||||
"""Wrapper function: transfrom score prediction model to velocity
|
||||
Args:
|
||||
score: [batch_dim, ...] shaped tensor; score model output
|
||||
x: [batch_dim, ...] shaped tensor; x_t data point
|
||||
t: [batch_dim,] time tensor
|
||||
"""
|
||||
t = expand_t_like_x(t, x)
|
||||
drift, var = self.compute_drift(x, t)
|
||||
velocity = var * score - drift
|
||||
return velocity
|
||||
|
||||
def compute_mu_t(self, t, x0, x1):
|
||||
"""Compute the mean of time-dependent density p_t"""
|
||||
t = expand_t_like_x(t, x1)
|
||||
alpha_t, _ = self.compute_alpha_t(t)
|
||||
sigma_t, _ = self.compute_sigma_t(t)
|
||||
return alpha_t * x1 + sigma_t * x0
|
||||
|
||||
def compute_xt(self, t, x0, x1):
|
||||
"""Sample xt from time-dependent density p_t; rng is required"""
|
||||
xt = self.compute_mu_t(t, x0, x1)
|
||||
return xt
|
||||
|
||||
def compute_ut(self, t, x0, x1, xt):
|
||||
"""Compute the vector field corresponding to p_t"""
|
||||
t = expand_t_like_x(t, x1)
|
||||
_, d_alpha_t = self.compute_alpha_t(t)
|
||||
_, d_sigma_t = self.compute_sigma_t(t)
|
||||
return d_alpha_t * x1 + d_sigma_t * x0
|
||||
|
||||
def plan(self, t, x0, x1):
|
||||
xt = self.compute_xt(t, x0, x1)
|
||||
ut = self.compute_ut(t, x0, x1, xt)
|
||||
return t, xt, ut
|
||||
|
||||
|
||||
class VPCPlan(ICPlan):
|
||||
"""class for VP path flow matching"""
|
||||
|
||||
def __init__(self, sigma_min=0.1, sigma_max=20.0):
|
||||
self.sigma_min = sigma_min
|
||||
self.sigma_max = sigma_max
|
||||
self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min
|
||||
self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min
|
||||
|
||||
|
||||
def compute_alpha_t(self, t):
|
||||
"""Compute coefficient of x1"""
|
||||
alpha_t = self.log_mean_coeff(t)
|
||||
alpha_t = th.exp(alpha_t)
|
||||
d_alpha_t = alpha_t * self.d_log_mean_coeff(t)
|
||||
return alpha_t, d_alpha_t
|
||||
|
||||
def compute_sigma_t(self, t):
|
||||
"""Compute coefficient of x0"""
|
||||
p_sigma_t = 2 * self.log_mean_coeff(t)
|
||||
sigma_t = th.sqrt(1 - th.exp(p_sigma_t))
|
||||
d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t)
|
||||
return sigma_t, d_sigma_t
|
||||
|
||||
def compute_d_alpha_alpha_ratio_t(self, t):
|
||||
"""Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
|
||||
return self.d_log_mean_coeff(t)
|
||||
|
||||
def compute_drift(self, x, t):
|
||||
"""Compute the drift term of the SDE"""
|
||||
t = expand_t_like_x(t, x)
|
||||
beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
|
||||
return -0.5 * beta_t * x, beta_t / 2
|
||||
|
||||
|
||||
class GVPCPlan(ICPlan):
|
||||
def __init__(self, sigma=0.0):
|
||||
super().__init__(sigma)
|
||||
|
||||
def compute_alpha_t(self, t):
|
||||
"""Compute coefficient of x1"""
|
||||
alpha_t = th.sin(t * np.pi / 2)
|
||||
d_alpha_t = np.pi / 2 * th.cos(t * np.pi / 2)
|
||||
return alpha_t, d_alpha_t
|
||||
|
||||
def compute_sigma_t(self, t):
|
||||
"""Compute coefficient of x0"""
|
||||
sigma_t = th.cos(t * np.pi / 2)
|
||||
d_sigma_t = -np.pi / 2 * th.sin(t * np.pi / 2)
|
||||
return sigma_t, d_sigma_t
|
||||
|
||||
def compute_d_alpha_alpha_ratio_t(self, t):
|
||||
"""Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
|
||||
return np.pi / (2 * th.tan(t * np.pi / 2))
|
||||
@@ -0,0 +1,443 @@
|
||||
import torch as th
|
||||
import numpy as np
|
||||
import logging
|
||||
|
||||
import enum
|
||||
|
||||
from . import path
|
||||
from .utils import EasyDict, log_state, mean_flat
|
||||
from .integrators import ode, sde
|
||||
|
||||
class ModelType(enum.Enum):
|
||||
"""
|
||||
Which type of output the model predicts.
|
||||
"""
|
||||
|
||||
NOISE = enum.auto() # the model predicts epsilon
|
||||
SCORE = enum.auto() # the model predicts \nabla \log p(x)
|
||||
VELOCITY = enum.auto() # the model predicts v(x)
|
||||
|
||||
class PathType(enum.Enum):
|
||||
"""
|
||||
Which type of path to use.
|
||||
"""
|
||||
|
||||
LINEAR = enum.auto()
|
||||
GVP = enum.auto()
|
||||
VP = enum.auto()
|
||||
|
||||
class WeightType(enum.Enum):
|
||||
"""
|
||||
Which type of weighting to use.
|
||||
"""
|
||||
|
||||
NONE = enum.auto()
|
||||
VELOCITY = enum.auto()
|
||||
LIKELIHOOD = enum.auto()
|
||||
|
||||
|
||||
class Transport:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model_type,
|
||||
path_type,
|
||||
loss_type,
|
||||
train_eps,
|
||||
sample_eps,
|
||||
):
|
||||
path_options = {
|
||||
PathType.LINEAR: path.ICPlan,
|
||||
PathType.GVP: path.GVPCPlan,
|
||||
PathType.VP: path.VPCPlan,
|
||||
}
|
||||
|
||||
self.loss_type = loss_type
|
||||
self.model_type = model_type
|
||||
self.path_sampler = path_options[path_type]()
|
||||
self.train_eps = train_eps
|
||||
self.sample_eps = sample_eps
|
||||
|
||||
def prior_logp(self, z):
|
||||
'''
|
||||
Standard multivariate normal prior
|
||||
Assume z is batched
|
||||
'''
|
||||
shape = th.tensor(z.size())
|
||||
N = th.prod(shape[1:])
|
||||
_fn = lambda x: -N / 2. * np.log(2 * np.pi) - th.sum(x ** 2) / 2.
|
||||
return th.vmap(_fn)(z)
|
||||
|
||||
|
||||
def check_interval(
|
||||
self,
|
||||
train_eps,
|
||||
sample_eps,
|
||||
*,
|
||||
diffusion_form="SBDM",
|
||||
sde=False,
|
||||
reverse=False,
|
||||
eval=False,
|
||||
last_step_size=0.0,
|
||||
):
|
||||
t0 = 0
|
||||
t1 = 1
|
||||
eps = train_eps if not eval else sample_eps
|
||||
if (type(self.path_sampler) in [path.VPCPlan]):
|
||||
|
||||
t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
|
||||
|
||||
elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan]) \
|
||||
and (self.model_type != ModelType.VELOCITY or sde): # avoid numerical issue by taking a first semi-implicit step
|
||||
|
||||
t0 = eps if (diffusion_form == "SBDM" and sde) or self.model_type != ModelType.VELOCITY else 0
|
||||
t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
|
||||
|
||||
if reverse:
|
||||
t0, t1 = 1 - t0, 1 - t1
|
||||
|
||||
return t0, t1
|
||||
|
||||
|
||||
def sample(self, x1):
|
||||
"""Sampling x0 & t based on shape of x1 (if needed)
|
||||
Args:
|
||||
x1 - data point; [batch, *dim]
|
||||
"""
|
||||
|
||||
x0 = th.randn_like(x1)
|
||||
t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
|
||||
t = th.rand((x1.shape[0],)) * (t1 - t0) + t0
|
||||
t = t.to(x1)
|
||||
return t, x0, x1
|
||||
|
||||
|
||||
def training_losses(
|
||||
self,
|
||||
model,
|
||||
x1,
|
||||
model_kwargs=None
|
||||
):
|
||||
"""Loss for training the score model
|
||||
Args:
|
||||
- model: backbone model; could be score, noise, or velocity
|
||||
- x1: datapoint
|
||||
- model_kwargs: additional arguments for the model
|
||||
"""
|
||||
if model_kwargs == None:
|
||||
model_kwargs = {}
|
||||
|
||||
t, x0, x1 = self.sample(x1)
|
||||
t, xt, ut = self.path_sampler.plan(t, x0, x1)
|
||||
model_output = model(xt, t, **model_kwargs)
|
||||
B, *_, C = xt.shape
|
||||
assert model_output.size() == (B, *xt.size()[1:-1], C)
|
||||
|
||||
terms = {}
|
||||
terms['pred'] = model_output
|
||||
if self.model_type == ModelType.VELOCITY:
|
||||
terms['loss'] = mean_flat(((model_output - ut) ** 2))
|
||||
else:
|
||||
_, drift_var = self.path_sampler.compute_drift(xt, t)
|
||||
sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
|
||||
if self.loss_type in [WeightType.VELOCITY]:
|
||||
weight = (drift_var / sigma_t) ** 2
|
||||
elif self.loss_type in [WeightType.LIKELIHOOD]:
|
||||
weight = drift_var / (sigma_t ** 2)
|
||||
elif self.loss_type in [WeightType.NONE]:
|
||||
weight = 1
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
if self.model_type == ModelType.NOISE:
|
||||
terms['loss'] = mean_flat(weight * ((model_output - x0) ** 2))
|
||||
else:
|
||||
terms['loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
|
||||
|
||||
return terms
|
||||
|
||||
|
||||
def get_drift(
|
||||
self
|
||||
):
|
||||
"""member function for obtaining the drift of the probability flow ODE"""
|
||||
def score_ode(x, t, model, **model_kwargs):
|
||||
drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
|
||||
model_output = model(x, t, **model_kwargs)
|
||||
return (-drift_mean + drift_var * model_output) # by change of variable
|
||||
|
||||
def noise_ode(x, t, model, **model_kwargs):
|
||||
drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
|
||||
sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
|
||||
model_output = model(x, t, **model_kwargs)
|
||||
score = model_output / -sigma_t
|
||||
return (-drift_mean + drift_var * score)
|
||||
|
||||
def velocity_ode(x, t, model, **model_kwargs):
|
||||
model_output = model(x, t, **model_kwargs)
|
||||
return model_output
|
||||
|
||||
if self.model_type == ModelType.NOISE:
|
||||
drift_fn = noise_ode
|
||||
elif self.model_type == ModelType.SCORE:
|
||||
drift_fn = score_ode
|
||||
else:
|
||||
drift_fn = velocity_ode
|
||||
|
||||
def body_fn(x, t, model, **model_kwargs):
|
||||
model_output = drift_fn(x, t, model, **model_kwargs)
|
||||
assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape"
|
||||
return model_output
|
||||
|
||||
return body_fn
|
||||
|
||||
|
||||
def get_score(
|
||||
self,
|
||||
):
|
||||
"""member function for obtaining score of
|
||||
x_t = alpha_t * x + sigma_t * eps"""
|
||||
if self.model_type == ModelType.NOISE:
|
||||
score_fn = lambda x, t, model, **kwargs: model(x, t, **kwargs) / -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
|
||||
elif self.model_type == ModelType.SCORE:
|
||||
score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
|
||||
elif self.model_type == ModelType.VELOCITY:
|
||||
score_fn = lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(model(x, t, **kwargs), x, t)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
return score_fn
|
||||
|
||||
|
||||
class Sampler:
|
||||
"""Sampler class for the transport model"""
|
||||
def __init__(
|
||||
self,
|
||||
transport,
|
||||
):
|
||||
"""Constructor for a general sampler; supporting different sampling methods
|
||||
Args:
|
||||
- transport: an tranport object specify model prediction & interpolant type
|
||||
"""
|
||||
|
||||
self.transport = transport
|
||||
self.drift = self.transport.get_drift()
|
||||
self.score = self.transport.get_score()
|
||||
|
||||
def __get_sde_diffusion_and_drift(
|
||||
self,
|
||||
*,
|
||||
diffusion_form="SBDM",
|
||||
diffusion_norm=1.0,
|
||||
):
|
||||
|
||||
def diffusion_fn(x, t):
|
||||
diffusion = self.transport.path_sampler.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm)
|
||||
return diffusion
|
||||
|
||||
sde_drift = \
|
||||
lambda x, t, model, **kwargs: \
|
||||
self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
|
||||
|
||||
sde_diffusion = diffusion_fn
|
||||
|
||||
return sde_drift, sde_diffusion
|
||||
|
||||
def __get_last_step(
|
||||
self,
|
||||
sde_drift,
|
||||
*,
|
||||
last_step,
|
||||
last_step_size,
|
||||
):
|
||||
"""Get the last step function of the SDE solver"""
|
||||
|
||||
if last_step is None:
|
||||
last_step_fn = \
|
||||
lambda x, t, model, **model_kwargs: \
|
||||
x
|
||||
elif last_step == "Mean":
|
||||
last_step_fn = \
|
||||
lambda x, t, model, **model_kwargs: \
|
||||
x + sde_drift(x, t, model, **model_kwargs) * last_step_size
|
||||
elif last_step == "Tweedie":
|
||||
alpha = self.transport.path_sampler.compute_alpha_t # simple aliasing; the original name was too long
|
||||
sigma = self.transport.path_sampler.compute_sigma_t
|
||||
last_step_fn = \
|
||||
lambda x, t, model, **model_kwargs: \
|
||||
x / alpha(t)[0][0] + (sigma(t)[0][0] ** 2) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
|
||||
elif last_step == "Euler":
|
||||
last_step_fn = \
|
||||
lambda x, t, model, **model_kwargs: \
|
||||
x + self.drift(x, t, model, **model_kwargs) * last_step_size
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
return last_step_fn
|
||||
|
||||
def sample_sde(
|
||||
self,
|
||||
*,
|
||||
sampling_method="Euler",
|
||||
diffusion_form="SBDM",
|
||||
diffusion_norm=1.0,
|
||||
last_step="Mean",
|
||||
last_step_size=0.04,
|
||||
num_steps=250,
|
||||
):
|
||||
"""returns a sampling function with given SDE settings
|
||||
Args:
|
||||
- sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
|
||||
- diffusion_form: function form of diffusion coefficient; default to be matching SBDM
|
||||
- diffusion_norm: function magnitude of diffusion coefficient; default to 1
|
||||
- last_step: type of the last step; default to identity
|
||||
- last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
|
||||
- num_steps: total integration step of SDE
|
||||
"""
|
||||
|
||||
if last_step is None:
|
||||
last_step_size = 0.0
|
||||
|
||||
sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
|
||||
diffusion_form=diffusion_form,
|
||||
diffusion_norm=diffusion_norm,
|
||||
)
|
||||
|
||||
t0, t1 = self.transport.check_interval(
|
||||
self.transport.train_eps,
|
||||
self.transport.sample_eps,
|
||||
diffusion_form=diffusion_form,
|
||||
sde=True,
|
||||
eval=True,
|
||||
reverse=False,
|
||||
last_step_size=last_step_size,
|
||||
)
|
||||
|
||||
_sde = sde(
|
||||
sde_drift,
|
||||
sde_diffusion,
|
||||
t0=t0,
|
||||
t1=t1,
|
||||
num_steps=num_steps,
|
||||
sampler_type=sampling_method
|
||||
)
|
||||
|
||||
last_step_fn = self.__get_last_step(sde_drift, last_step=last_step, last_step_size=last_step_size)
|
||||
|
||||
|
||||
def _sample(init, model, **model_kwargs):
|
||||
xs = _sde.sample(init, model, **model_kwargs)
|
||||
ts = th.ones(init.size(0), device=init.device) * t1
|
||||
x = last_step_fn(xs[-1], ts, model, **model_kwargs)
|
||||
xs.append(x)
|
||||
|
||||
assert len(xs) == num_steps, "Samples does not match the number of steps"
|
||||
|
||||
return xs
|
||||
|
||||
return _sample
|
||||
|
||||
def sample_ode(
|
||||
self,
|
||||
*,
|
||||
sampling_method="dopri5",
|
||||
num_steps=50,
|
||||
atol=1e-6,
|
||||
rtol=1e-3,
|
||||
reverse=False,
|
||||
):
|
||||
"""returns a sampling function with given ODE settings
|
||||
Args:
|
||||
- sampling_method: type of sampler used in solving the ODE; default to be Dopri5
|
||||
- num_steps:
|
||||
- fixed solver (Euler, Heun): the actual number of integration steps performed
|
||||
- adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
|
||||
- atol: absolute error tolerance for the solver
|
||||
- rtol: relative error tolerance for the solver
|
||||
- reverse: whether solving the ODE in reverse (data to noise); default to False
|
||||
"""
|
||||
if reverse:
|
||||
drift = lambda x, t, model, **kwargs: self.drift(x, th.ones_like(t) * (1 - t), model, **kwargs)
|
||||
else:
|
||||
drift = self.drift
|
||||
|
||||
t0, t1 = self.transport.check_interval(
|
||||
self.transport.train_eps,
|
||||
self.transport.sample_eps,
|
||||
sde=False,
|
||||
eval=True,
|
||||
reverse=reverse,
|
||||
last_step_size=0.0,
|
||||
)
|
||||
|
||||
_ode = ode(
|
||||
drift=drift,
|
||||
t0=t0,
|
||||
t1=t1,
|
||||
sampler_type=sampling_method,
|
||||
num_steps=num_steps,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
return _ode.sample
|
||||
|
||||
def sample_ode_likelihood(
|
||||
self,
|
||||
*,
|
||||
sampling_method="dopri5",
|
||||
num_steps=50,
|
||||
atol=1e-6,
|
||||
rtol=1e-3,
|
||||
):
|
||||
|
||||
"""returns a sampling function for calculating likelihood with given ODE settings
|
||||
Args:
|
||||
- sampling_method: type of sampler used in solving the ODE; default to be Dopri5
|
||||
- num_steps:
|
||||
- fixed solver (Euler, Heun): the actual number of integration steps performed
|
||||
- adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
|
||||
- atol: absolute error tolerance for the solver
|
||||
- rtol: relative error tolerance for the solver
|
||||
"""
|
||||
def _likelihood_drift(x, t, model, **model_kwargs):
|
||||
x, _ = x
|
||||
eps = th.randint(2, x.size(), dtype=th.float, device=x.device) * 2 - 1
|
||||
t = th.ones_like(t) * (1 - t)
|
||||
with th.enable_grad():
|
||||
x.requires_grad = True
|
||||
grad = th.autograd.grad(th.sum(self.drift(x, t, model, **model_kwargs) * eps), x)[0]
|
||||
logp_grad = th.sum(grad * eps, dim=tuple(range(1, len(x.size()))))
|
||||
drift = self.drift(x, t, model, **model_kwargs)
|
||||
return (-drift, logp_grad)
|
||||
|
||||
t0, t1 = self.transport.check_interval(
|
||||
self.transport.train_eps,
|
||||
self.transport.sample_eps,
|
||||
sde=False,
|
||||
eval=True,
|
||||
reverse=False,
|
||||
last_step_size=0.0,
|
||||
)
|
||||
|
||||
_ode = ode(
|
||||
drift=_likelihood_drift,
|
||||
t0=t0,
|
||||
t1=t1,
|
||||
sampler_type=sampling_method,
|
||||
num_steps=num_steps,
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
)
|
||||
|
||||
def _sample_fn(x, model, **model_kwargs):
|
||||
init_logp = th.zeros(x.size(0)).to(x)
|
||||
input = (x, init_logp)
|
||||
drift, delta_logp = _ode.sample(input, model, **model_kwargs)
|
||||
drift, delta_logp = drift[-1], delta_logp[-1]
|
||||
prior_logp = self.transport.prior_logp(drift)
|
||||
logp = prior_logp - delta_logp
|
||||
return logp, drift
|
||||
|
||||
return _sample_fn
|
||||
@@ -0,0 +1,29 @@
|
||||
import torch as th
|
||||
|
||||
class EasyDict:
|
||||
|
||||
def __init__(self, sub_dict):
|
||||
for k, v in sub_dict.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return getattr(self, key)
|
||||
|
||||
def mean_flat(x):
|
||||
"""
|
||||
Take the mean over all non-batch dimensions.
|
||||
"""
|
||||
return th.mean(x, dim=list(range(1, len(x.size()))))
|
||||
|
||||
def log_state(state):
|
||||
result = []
|
||||
|
||||
sorted_state = dict(sorted(state.items()))
|
||||
for key, value in sorted_state.items():
|
||||
# Check if the value is an instance of a class
|
||||
if "<object" in str(value) or "object at" in str(value):
|
||||
result.append(f"{key}: [{value.__class__.__name__}]")
|
||||
else:
|
||||
result.append(f"{key}: {value}")
|
||||
|
||||
return '\n'.join(result)
|
||||
@@ -9,6 +9,7 @@ Sample new images from a pre-trained Latte.
|
||||
import os
|
||||
import sys
|
||||
|
||||
from opensora.dataset import ae_denorm
|
||||
from opensora.models.ae import ae_channel_config, getae, ae_stride_config
|
||||
from opensora.models.diffusion import Diffusion_models
|
||||
from opensora.models.diffusion.diffusion import create_diffusion
|
||||
@@ -62,9 +63,9 @@ def main(args):
|
||||
|
||||
# Create sampling noise:
|
||||
if args.use_fp16:
|
||||
z = torch.randn(1, args.num_frames, model.in_channels, latent_size, latent_size, dtype=torch.float16, device=device) # b c f h w
|
||||
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, dtype=torch.float16, device=device) # b c f h w
|
||||
else:
|
||||
z = torch.randn(1, args.num_frames, model.in_channels, latent_size, latent_size, device=device)
|
||||
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, device=device)
|
||||
|
||||
# Setup classifier-free guidance:
|
||||
if using_cfg:
|
||||
@@ -90,17 +91,14 @@ def main(args):
|
||||
|
||||
if args.use_fp16:
|
||||
samples = samples.to(dtype=torch.float16)
|
||||
b, f, c, h, w = samples.shape
|
||||
samples = rearrange(samples, 'b f c h w -> (b f) c h w')
|
||||
samples = ae.decode(samples)
|
||||
samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b)
|
||||
# Save and display images:
|
||||
|
||||
if not os.path.exists(args.save_video_path):
|
||||
os.makedirs(args.save_video_path)
|
||||
|
||||
|
||||
video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
|
||||
video_ = (ae_denorm[args.ae](samples[0]) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
|
||||
video_save_path = os.path.join(args.save_video_path, 'sample' + '.mp4')
|
||||
print(video_save_path)
|
||||
imageio.mimwrite(video_save_path, video_, fps=args.fps, quality=9)
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""
|
||||
Sample new images from a pre-trained SiT.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
from opensora.dataset import ae_denorm
|
||||
from opensora.models.ae import ae_channel_config, getae, ae_stride_config
|
||||
from opensora.models.diffusion import Diffusion_models
|
||||
from opensora.models.diffusion.transport import create_transport, Sampler
|
||||
from opensora.utils.utils import find_model
|
||||
|
||||
import torch
|
||||
import argparse
|
||||
|
||||
from einops import rearrange
|
||||
import imageio
|
||||
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
|
||||
|
||||
def main(mode, args):
|
||||
# Setup PyTorch:
|
||||
# torch.manual_seed(args.seed)
|
||||
torch.set_grad_enabled(False)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
using_cfg = args.cfg_scale > 1.0
|
||||
|
||||
# Load model:
|
||||
latent_size = args.image_size // ae_stride_config[args.ae][1]
|
||||
args.latent_size = latent_size
|
||||
model = Diffusion_models[args.model](
|
||||
input_size=latent_size,
|
||||
num_classes=args.num_classes,
|
||||
in_channels=ae_channel_config[args.ae],
|
||||
extras=args.extras
|
||||
).to(device)
|
||||
|
||||
if args.use_compile:
|
||||
model = torch.compile(model)
|
||||
|
||||
# a pre-trained model or load a custom Latte checkpoint from train.py:
|
||||
ckpt_path = args.ckpt
|
||||
state_dict = find_model(ckpt_path)
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
model.eval() # important!
|
||||
transport = create_transport(
|
||||
args.path_type,
|
||||
args.prediction,
|
||||
args.loss_weight,
|
||||
args.train_eps,
|
||||
args.sample_eps
|
||||
)
|
||||
sampler = Sampler(transport)
|
||||
if mode == "ODE":
|
||||
if args.likelihood:
|
||||
assert args.cfg_scale == 1, "Likelihood is incompatible with guidance"
|
||||
sample_fn = sampler.sample_ode_likelihood(
|
||||
sampling_method=args.sampling_method,
|
||||
num_steps=args.num_sampling_steps,
|
||||
atol=args.atol,
|
||||
rtol=args.rtol,
|
||||
)
|
||||
else:
|
||||
sample_fn = sampler.sample_ode(
|
||||
sampling_method=args.sampling_method,
|
||||
num_steps=args.num_sampling_steps,
|
||||
atol=args.atol,
|
||||
rtol=args.rtol,
|
||||
reverse=args.reverse
|
||||
)
|
||||
elif mode == "SDE":
|
||||
sample_fn = sampler.sample_sde(
|
||||
sampling_method=args.sampling_method,
|
||||
diffusion_form=args.diffusion_form,
|
||||
diffusion_norm=args.diffusion_norm,
|
||||
last_step=args.last_step,
|
||||
last_step_size=args.last_step_size,
|
||||
num_steps=args.num_sampling_steps,
|
||||
)
|
||||
|
||||
ae = getae(args).to(device)
|
||||
|
||||
if args.use_fp16:
|
||||
print('WARNING: using half percision for inferencing!')
|
||||
ae.to(dtype=torch.float16)
|
||||
model.to(dtype=torch.float16)
|
||||
|
||||
# Labels to condition the model with (feel free to change):
|
||||
|
||||
# Create sampling noise:
|
||||
if args.use_fp16:
|
||||
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, dtype=torch.float16, device=device) # b c f h w
|
||||
else:
|
||||
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, device=device)
|
||||
|
||||
# Setup classifier-free guidance:
|
||||
if using_cfg:
|
||||
z = torch.cat([z, z], 0)
|
||||
y = torch.randint(0, args.num_classes, (1,), device=device)
|
||||
y_null = torch.tensor([args.num_classes] * 1, device=device)
|
||||
y = torch.cat([y, y_null], dim=0)
|
||||
model_kwargs = dict(y=y, cfg_scale=args.cfg_scale, use_fp16=args.use_fp16)
|
||||
forward_fn = model.forward_with_cfg
|
||||
else:
|
||||
forward_fn = model.forward
|
||||
model_kwargs = dict(y=None, use_fp16=args.use_fp16)
|
||||
|
||||
# Sample images:
|
||||
samples = sample_fn(z, forward_fn, **model_kwargs)[-1]
|
||||
|
||||
if args.use_fp16:
|
||||
samples = samples.to(dtype=torch.float16)
|
||||
samples = ae.decode(samples)
|
||||
|
||||
# Save and display images:
|
||||
if not os.path.exists(args.save_video_path):
|
||||
os.makedirs(args.save_video_path)
|
||||
|
||||
|
||||
video_ = (ae_denorm[args.ae](samples[0]) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
|
||||
video_save_path = os.path.join(args.save_video_path, 'sample' + '.mp4')
|
||||
print(video_save_path)
|
||||
imageio.mimwrite(video_save_path, video_, fps=args.fps, quality=9)
|
||||
print('save path {}'.format(args.save_video_path))
|
||||
|
||||
|
||||
def none_or_str(value):
|
||||
if value == 'None':
|
||||
return None
|
||||
return value
|
||||
|
||||
def parse_transport_args(parser):
|
||||
group = parser.add_argument_group("Transport arguments")
|
||||
group.add_argument("--path-type", type=str, default="Linear", choices=["Linear", "GVP", "VP"])
|
||||
group.add_argument("--prediction", type=str, default="velocity", choices=["velocity", "score", "noise"])
|
||||
group.add_argument("--loss-weight", type=none_or_str, default=None, choices=[None, "velocity", "likelihood"])
|
||||
group.add_argument("--sample-eps", type=float)
|
||||
group.add_argument("--train-eps", type=float)
|
||||
|
||||
def parse_ode_args(parser):
|
||||
group = parser.add_argument_group("ODE arguments")
|
||||
group.add_argument("--sampling-method", type=str, default="dopri5", help="blackbox ODE solver methods; for full list check https://github.com/rtqichen/torchdiffeq")
|
||||
group.add_argument("--atol", type=float, default=1e-6, help="Absolute tolerance")
|
||||
group.add_argument("--rtol", type=float, default=1e-3, help="Relative tolerance")
|
||||
group.add_argument("--reverse", action="store_true")
|
||||
group.add_argument("--likelihood", action="store_true")
|
||||
|
||||
def parse_sde_args(parser):
|
||||
group = parser.add_argument_group("SDE arguments")
|
||||
group.add_argument("--sampling-method", type=str, default="Euler", choices=["Euler", "Heun"])
|
||||
group.add_argument("--diffusion-form", type=str, default="sigma", \
|
||||
choices=["constant", "SBDM", "sigma", "linear", "decreasing", "increasing-decreasing"],\
|
||||
help="form of diffusion coefficient in the SDE")
|
||||
group.add_argument("--diffusion-norm", type=float, default=1.0)
|
||||
group.add_argument("--last-step", type=none_or_str, default="Mean", choices=[None, "Mean", "Tweedie", "Euler"],\
|
||||
help="form of last step taken in the SDE")
|
||||
group.add_argument("--last-step-size", type=float, default=0.04, \
|
||||
help="size of the last step taken")
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: program.py <mode> [options]")
|
||||
sys.exit(1)
|
||||
|
||||
mode = sys.argv[1]
|
||||
|
||||
assert mode[:2] != "--", "Usage: program.py <mode> [options]"
|
||||
assert mode in ["ODE", "SDE"], "Invalid mode. Please choose 'ODE' or 'SDE'"
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--ckpt", type=str, default="")
|
||||
parser.add_argument("--model", type=str, default='Latte-XL/122')
|
||||
parser.add_argument("--ae", type=str, default='stabilityai/sd-vae-ft-mse')
|
||||
parser.add_argument("--save-video-path", type=str, default="./sample_videos/")
|
||||
parser.add_argument("--fps", type=int, default=10)
|
||||
parser.add_argument("--num-classes", type=int, default=101)
|
||||
parser.add_argument("--num-frames", type=int, default=16)
|
||||
parser.add_argument("--image-size", type=int, default=256, choices=[256, 512])
|
||||
parser.add_argument("--extras", type=int, default=1)
|
||||
parser.add_argument("--num-sampling-steps", type=int, default=250)
|
||||
parser.add_argument("--cfg-scale", type=float, default=1.0)
|
||||
parser.add_argument("--use-fp16", action="store_true")
|
||||
parser.add_argument("--use-compile", action="store_true")
|
||||
parser.add_argument("--sample-method", type=str, default='ddpm')
|
||||
|
||||
parse_transport_args(parser)
|
||||
if mode == "ODE":
|
||||
parse_ode_args(parser)
|
||||
# Further processing for ODE
|
||||
elif mode == "SDE":
|
||||
parse_sde_args(parser)
|
||||
# Further processing for SDE
|
||||
|
||||
args = parser.parse_known_args()[0]
|
||||
main(mode, args)
|
||||
@@ -249,6 +249,7 @@ def main(args):
|
||||
steps_per_sec = log_steps / (end_time - start_time)
|
||||
# Reduce loss history over all processes:
|
||||
avg_loss = torch.tensor(running_loss / log_steps, device=device)
|
||||
dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
|
||||
avg_loss = avg_loss.item() / accelerator.num_processes
|
||||
# logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
|
||||
if accelerator.is_main_process:
|
||||
|
||||
@@ -1,66 +0,0 @@
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def crop_video(input_file, start_and_end):
|
||||
# 打开视频文件
|
||||
video = cv2.VideoCapture(input_file)
|
||||
|
||||
# 获取视频的帧速率和总帧数
|
||||
fps = video.get(cv2.CAP_PROP_FPS)
|
||||
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
|
||||
current_frame = 0
|
||||
# 设置当前帧的索引
|
||||
for s, e in tqdm(start_and_end):
|
||||
# 计算裁剪的起始帧和结束帧
|
||||
start_frame = int(s * fps)
|
||||
end_frame = int(e * fps)
|
||||
# 设置视频编解码器并创建输出视频对象
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
output = cv2.VideoWriter(f'clip/{s}-{e}.mp4', fourcc, fps, (int(video.get(3)), int(video.get(4))))
|
||||
|
||||
|
||||
# 读取并写入视频帧，直到达到结束帧
|
||||
while current_frame < end_frame:
|
||||
ret, frame = video.read()
|
||||
if not ret:
|
||||
break
|
||||
if current_frame >= start_frame:
|
||||
output.write(frame)
|
||||
current_frame += 1
|
||||
|
||||
# 释放视频对象和输出对象
|
||||
video.release()
|
||||
output.release()
|
||||
|
||||
print("视频裁剪完成！")
|
||||
|
||||
# 使用示例
|
||||
input_file = r"C:\Users\ABin\Desktop\S01E03 热带宁静Tropical Serenity_超清 4K.mp4"
|
||||
start_and_end = [
|
||||
[24, 71],
|
||||
[76, 178],
|
||||
[182, 265],
|
||||
[268, 315],
|
||||
[322, 366],
|
||||
[372, 467],
|
||||
[473, 535],
|
||||
[550, 610],
|
||||
[618, 684],
|
||||
[689, 792],
|
||||
[800, 880],
|
||||
[884, 945],
|
||||
[951, 1002],
|
||||
[1009, 1116],
|
||||
[1127, 1196],
|
||||
[1205, 1301],
|
||||
[1306, 1381],
|
||||
[1386, 1421],
|
||||
[1430, 1505],
|
||||
[1571, 1628],
|
||||
[1636, 1724],
|
||||
[1731, 1784],
|
||||
]
|
||||
crop_video(input_file, start_and_end)
|
||||
Reference in New Issue
Block a user