diff --git a/README.md b/README.md
index 0de40a1aa..757b80c77 100644
--- a/README.md
+++ b/README.md
@@ -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
-
-
- 
-
diff --git a/opensora/dataset/__init__.py b/opensora/dataset/__init__.py
index 8f9382958..0acdd345c 100644
--- a/opensora/dataset/__init__.py
+++ b/opensora/dataset/__init__.py
@@ -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
diff --git a/opensora/models/ae/imagebase/vae/vae.py b/opensora/models/ae/imagebase/vae/vae.py
index 3ba4488ef..3a19a5239 100644
--- a/opensora/models/ae/imagebase/vae/vae.py
+++ b/opensora/models/ae/imagebase/vae/vae.py
@@ -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):
diff --git a/opensora/models/ae/videobase/vqvae/videogpt/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/__init__.py
index 722457054..f42841f54 100644
--- a/opensora/models/ae/videobase/vqvae/videogpt/__init__.py
+++ b/opensora/models/ae/videobase/vqvae/videogpt/__init__.py
@@ -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
\ No newline at end of file
diff --git a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py
index 0353334d2..5bc2d2f19 100644
--- a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py
+++ b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py
@@ -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)
diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py
index 676f2cd6a..d7611b2b5 100644
--- a/opensora/models/diffusion/dit/dit.py
+++ b/opensora/models/diffusion/dit/dit.py
@@ -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):
diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py
index ad8ac19a2..8824d2a5c 100644
--- a/opensora/models/diffusion/latte/latte.py
+++ b/opensora/models/diffusion/latte/latte.py
@@ -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()
diff --git a/opensora/models/diffusion/transport/__init__.py b/opensora/models/diffusion/transport/__init__.py
new file mode 100644
index 000000000..db68edd20
--- /dev/null
+++ b/opensora/models/diffusion/transport/__init__.py
@@ -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
\ No newline at end of file
diff --git a/opensora/models/diffusion/transport/integrators.py b/opensora/models/diffusion/transport/integrators.py
new file mode 100644
index 000000000..adf7c7b4c
--- /dev/null
+++ b/opensora/models/diffusion/transport/integrators.py
@@ -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
\ No newline at end of file
diff --git a/opensora/models/diffusion/transport/path.py b/opensora/models/diffusion/transport/path.py
new file mode 100644
index 000000000..156a7b0de
--- /dev/null
+++ b/opensora/models/diffusion/transport/path.py
@@ -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))
\ No newline at end of file
diff --git a/opensora/models/diffusion/transport/transport.py b/opensora/models/diffusion/transport/transport.py
new file mode 100644
index 000000000..396c516cf
--- /dev/null
+++ b/opensora/models/diffusion/transport/transport.py
@@ -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
\ No newline at end of file
diff --git a/opensora/models/diffusion/transport/utils.py b/opensora/models/diffusion/transport/utils.py
new file mode 100644
index 000000000..446460355
--- /dev/null
+++ b/opensora/models/diffusion/transport/utils.py
@@ -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 " (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)
diff --git a/opensora/sample/transport_sample.py b/opensora/sample/transport_sample.py
new file mode 100644
index 000000000..b88b55206
--- /dev/null
+++ b/opensora/sample/transport_sample.py
@@ -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 [options]")
+ sys.exit(1)
+
+ mode = sys.argv[1]
+
+ assert mode[:2] != "--", "Usage: program.py [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)
diff --git a/opensora/train/train.py b/opensora/train/train.py
index d885d1236..f027a04cf 100644
--- a/opensora/train/train.py
+++ b/opensora/train/train.py
@@ -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:
diff --git a/opensora/utils/movie2clip.py b/opensora/utils/movie2clip.py
deleted file mode 100644
index 2da7a1d18..000000000
--- a/opensora/utils/movie2clip.py
+++ /dev/null
@@ -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)
\ No newline at end of file