nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[feature] update diffusion pipeline and sample script (#10)

Browse source 
* [feature] diffusion natively support video format

* [feature] add timestamp embedding

* [feature] support learn from raw video

* [feature] update sample script
This commit is contained in:
Hongxin Liu 2024-02-26 11:23:09 +08:00 committed by GitHub
parent 6f887f453b
commit adba00f151
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 146 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
data_utils.py
View file

@ -60,7 +60,7 @@ def col2video(
if y + patch_size > h or x + patch_size > w:
continue
# [T, C, P, P]
patch = patches[:, y * num_x_patches + x]
patch = patches[:, (y // patch_size) * num_x_patches + x // patch_size]
video[:, :, y : y + patch_size, x : x + patch_size].copy_(patch)
return video
@ -157,7 +157,7 @@ def unnormalize_video(video: torch.Tensor) -> torch.Tensor:
@torch.no_grad()
def preprocess_batch(
batch: dict, patch_size: int, vqvae: nn.Module, device=None
batch: dict, patch_size: int, vqvae: Optional[nn.Module] = None, device=None
) -> dict:
if device is None:
device = get_current_device()
@ -165,16 +165,19 @@ def preprocess_batch(
for video in batch.pop("videos"):
video = video.to(device)
video = normalize_video(video)
# [T, H, W, C] -> [B, C, T, H, W]
video = video.permute(3, 0, 1, 2)
video = video.unsqueeze(0)
latent_indices, embeddings = vqvae.encode(video, include_embeddings=True)
# [B, C, T, H, W] -> [T, C, H, W]
embeddings = embeddings.squeeze(0).permute(1, 0, 2, 3)
videos.append(embeddings)
if vqvae is not None:
# [T, H, W, C] -> [B, C, T, H, W]
video = video.permute(3, 0, 1, 2)
video = video.unsqueeze(0)
latent_indices, embeddings = vqvae.encode(video, include_embeddings=True)
# [B, C, T, H, W] -> [T, C, H, W]
embeddings = embeddings.squeeze(0).permute(1, 0, 2, 3)
videos.append(embeddings)
else:
# [T, H, W, C] -> [T, C, H, W]
video = video.permute(0, 3, 1, 2).contiguous()
videos.append(video)
video_latent_states, video_padding_mask = patchify_batch(videos, patch_size)
# hack diffuser, [B, S, C, P, P] -> [B, C, S, P, P]
video_latent_states = video_latent_states.transpose(1, 2)
batch["video_latent_states"] = video_latent_states
batch["video_padding_mask"] = video_padding_mask
text_padding_mask = batch.pop("text_padding_mask").to(device)

29
diffusion/gaussian_diffusion.py
View file

@ -205,7 +205,7 @@ class GaussianDiffusion:
def q_mean_variance(self, x_start, t):
"""
Get the distribution q(x_t | x_0).
:param x_start: the [N x C x ...] tensor of noiseless inputs.
:param x_start: the [N x T x C x ...] tensor of noiseless inputs.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:return: A tuple (mean, variance, log_variance), all of x_start's shape.
"""
@ -266,7 +266,7 @@ class GaussianDiffusion:
the initial x, x_0.
:param model: the model, which takes a signal and a batch of timesteps
as input.
:param x: the [N x C x ...] tensor at time t.
:param x: the [N x T x C x ...] tensor at time t.
:param t: a 1-D Tensor of timesteps.
:param clip_denoised: if True, clip the denoised signal into [-1, 1].
:param denoised_fn: if not None, a function which applies to the
@ -283,7 +283,7 @@ class GaussianDiffusion:
if model_kwargs is None:
model_kwargs = {}
B, C = x.shape[:2]
B, S, C = x.shape[:3]
assert t.shape == (B,)
model_output = model(x, t, **model_kwargs)
if isinstance(model_output, tuple):
@ -292,8 +292,8 @@ class GaussianDiffusion:
extra = None
if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
assert model_output.shape == (B, C * 2, *x.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
assert model_output.shape == (B, S, C * 2, *x.shape[3:])
model_output, model_var_values = th.split(model_output, C, dim=2)
min_log = _extract_into_tensor(
self.posterior_log_variance_clipped, t, x.shape
)
@ -453,7 +453,7 @@ class GaussianDiffusion:
"""
Generate samples from the model.
:param model: the model module.
:param shape: the shape of the samples, (N, C, H, W).
:param shape: the shape of the samples, (N, T, C, H, W).
:param noise: if specified, the noise from the encoder to sample.
Should be of the same shape as `shape`.
:param clip_denoised: if True, clip x_start predictions to [-1, 1].
@ -739,8 +739,7 @@ class GaussianDiffusion:
def _expand_mask(self, mask, ndim: int):
assert mask.ndim == 2
# [B, S] -> [B, 1, S, ...]
mask = mask.unsqueeze(1)
# [B, S] -> [B, S, ...]
mask = mask.view(*mask.shape, *([1] * (ndim - mask.ndim)))
return mask
@ -750,7 +749,7 @@ class GaussianDiffusion:
"""
Compute training losses for a single timestep.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param x_start: the [N x T x C x ...] tensor of inputs.
:param t: a batch of timestep indices.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
@ -788,12 +787,12 @@ class GaussianDiffusion:
ModelVarType.LEARNED,
ModelVarType.LEARNED_RANGE,
]:
B, C = x_t.shape[:2]
assert model_output.shape == (B, C * 2, *x_t.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
B, S, C = x_t.shape[:3]
assert model_output.shape == (B, S, C * 2, *x_t.shape[3:])
model_output, model_var_values = th.split(model_output, C, dim=2)
# Learn the variance using the variational bound, but don't let
# it affect our mean prediction.
frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
frozen_out = th.cat([model_output.detach(), model_var_values], dim=2)
terms["vb"] = self._vb_terms_bpd(
model=lambda *args, r=frozen_out: r,
x_start=x_start,
@ -830,7 +829,7 @@ class GaussianDiffusion:
Get the prior KL term for the variational lower-bound, measured in
bits-per-dim.
This term can't be optimized, as it only depends on the encoder.
:param x_start: the [N x C x ...] tensor of inputs.
:param x_start: the [N x T x C x ...] tensor of inputs.
:return: a batch of [N] KL values (in bits), one per batch element.
"""
batch_size = x_start.shape[0]
@ -846,7 +845,7 @@ class GaussianDiffusion:
Compute the entire variational lower-bound, measured in bits-per-dim,
as well as other related quantities.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param x_start: the [N x T x C x ...] tensor of inputs.
:param clip_denoised: if True, clip denoised samples.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.

44
models.py
View file

@ -150,7 +150,9 @@ class TimestepEmbedder(nn.Module):
return embedding
def forward(self, t):
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(
self.mlp[0].weight.dtype
)
t_emb = self.mlp(t_freq)
return t_emb
@ -211,7 +213,7 @@ class PatchEmbedder(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
# [B, S, C, P, P] -> [B, S, C*P*P]
# FIXME: hack diffusion and use view
x = x.reshape(*x.shape[:2], -1)
x = x.view(*x.shape[:2], -1)
out = F.linear(
x, self.proj.weight.view(self.proj.weight.shape[0], -1), self.proj.bias
)
@ -239,12 +241,10 @@ class PositionEmbedding(nn.Module):
self.max_position_embeddings = max_position_embeddings
self._set_pos_embed_cache(max_position_embeddings)
def _set_pos_embed_cache(self, seq_len: int):
def _set_pos_embed_cache(self, seq_len: int, device="cpu", dtype=torch.float):
self.max_seq_len_cached = seq_len
pos_embed = get_2d_sincos_pos_embed(
self.dim, int(self.max_position_embeddings**0.5)
)
pos_embed = torch.from_numpy(pos_embed).float()
pos_embed = get_2d_sincos_pos_embed(self.dim, math.ceil(seq_len**0.5))
pos_embed = torch.from_numpy(pos_embed).to(device=device, dtype=dtype)
# [S, H]
self.register_buffer("pos_embed_cache", pos_embed, persistent=False)
@ -252,7 +252,7 @@ class PositionEmbedding(nn.Module):
# [B, S, H]
seq_len = x.shape[1]
if seq_len > self.max_seq_len_cached:
self._set_pos_embed_cache(seq_len)
self._set_pos_embed_cache(seq_len, x.device, x.dtype)
pos_embed = self.pos_embed_cache[None, :seq_len]
return pos_embed
@ -328,7 +328,7 @@ class DiT(nn.Module):
def __init__(
self,
patch_size=2,
in_channels=256,
in_channels=3,
text_embed_dim=512,
hidden_size=1152,
depth=28,
@ -348,7 +348,7 @@ class DiT(nn.Module):
self.video_embedder = PatchEmbedder(
patch_size, in_channels, hidden_size, bias=True
)
self.t_embedder = TimestepEmbedder(hidden_size)
self.t_embedder = TimestepEmbedder(text_embed_dim)
self.pos_embed = PositionEmbedding(hidden_size, max_num_embeddings)
self.blocks = nn.ModuleList(
@ -403,15 +403,13 @@ class DiT(nn.Module):
attention_mask=None,
):
"""
video_latent_states: [B, C, S, P, P]
video_latent_states: [B, S, C, P, P]
"""
# [B, C, S, P, P] -> [B, S, C, P, P]
video_latent_states = video_latent_states.transpose(1, 2)
video_latent_states = self.video_embedder(video_latent_states)
pos_embed = self.pos_embed(video_latent_states)
video_latent_states = video_latent_states + pos_embed
# TODO: use timestep embedding
# t = self.t_embedder(t) # (N, D)
t = self.t_embedder(t) # (N, D)
text_latent_states = text_latent_states + t.unsqueeze(1)
attention_mask = self._prepare_mask(attention_mask, video_latent_states.dtype)
for block in self.blocks:
if self.grad_checkpointing and self.training:
@ -423,25 +421,31 @@ class DiT(nn.Module):
video_latent_states, text_latent_states, attention_mask
)
video_latent_states = self.final_layer(video_latent_states)
return video_latent_states.transpose(1, 2)
return video_latent_states
def forward_with_cfg(self, x, t, y, cfg_scale):
def forward_with_cfg(
self, x, t, text_latent_states, cfg_scale, attention_mask=None
):
"""
Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
"""
# https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
half = x[: len(x) // 2]
combined = torch.cat([half, half], dim=0)
model_out = self.forward(combined, t, y)
model_out = self.forward(
combined, t, text_latent_states, attention_mask=attention_mask
)
# For exact reproducibility reasons, we apply classifier-free guidance on only
# three channels by default. The standard approach to cfg applies it to all channels.
# This can be done by uncommenting the following line and commenting-out the line following that.
# eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
eps, rest = model_out[:, :3], model_out[:, 3:]
c = model_out.shape[2]
assert c == 2 * self.in_channels
eps, rest = model_out.chunk(2, dim=2)
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
eps = torch.cat([half_eps, half_eps], dim=0)
return torch.cat([eps, rest], dim=1)
return torch.cat([eps, rest], dim=2)
#################################################################################

118
sample.py
View file

@ -13,11 +13,12 @@ torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import argparse
from diffusers.models import AutoencoderKL
from torchvision.utils import save_image
from colossalai.utils import get_current_device
from torchvision.io import write_video
from transformers import AutoModel, AutoTokenizer, CLIPTextModel
from data_utils import col2video
from diffusion import create_diffusion
from download import find_model
from models import DiT_models
@ -25,63 +26,108 @@ def main(args):
# Setup PyTorch:
torch.manual_seed(args.seed)
torch.set_grad_enabled(False)
device = "cuda" if torch.cuda.is_available() else "cpu"
device = get_current_device()
if len(args.vqvae) > 0:
vqvae = (
AutoModel.from_pretrained(args.vqvae, trust_remote_code=True)
.to(device)
.eval()
)
in_channels = vqvae.embedding_dim
else:
# disable VQ-VAE if not provided, just use raw video frames
vqvae = None
in_channels = 3
text_model = CLIPTextModel.from_pretrained(args.text_model).to(device).eval()
tokenizer = AutoTokenizer.from_pretrained(args.text_model)
if args.ckpt is None:
assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download."
assert args.image_size in [256, 512]
assert args.num_classes == 1000
# Load model:
latent_size = args.image_size // 8
model = DiT_models[args.model](input_size=latent_size, num_classes=args.num_classes).to(device)
# Auto-download a pre-trained model or load a custom DiT checkpoint from train.py:
ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
state_dict = find_model(ckpt_path)
model.load_state_dict(state_dict)
model.eval() # important!
model = DiT_models[args.model](in_channels=in_channels).to(device).eval()
patch_size = model.patch_size
# model.load_state_dict(torch.load(args.ckpt))
diffusion = create_diffusion(str(args.num_sampling_steps))
vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
# Labels to condition the model with (feel free to change):
class_labels = [207, 360, 387, 974, 88, 979, 417, 279]
# Create sampling noise:
n = len(class_labels)
z = torch.randn(n, 4, latent_size, latent_size, device=device)
y = torch.tensor(class_labels, device=device)
text_inputs = tokenizer(args.text, return_tensors="pt")
text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
text_latent_states = text_model(**text_inputs).last_hidden_state
num_frames = args.fps * args.sec
z = torch.randn(
1,
(args.height // patch_size // 4)
* (args.width // patch_size // 4)
* (num_frames // 2),
in_channels,
patch_size,
patch_size,
device=device,
)
# Setup classifier-free guidance:
model_kwargs = {}
z = torch.cat([z, z], 0)
y_null = torch.tensor([1000] * n, device=device)
y = torch.cat([y, y_null], 0)
model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
model_kwargs["text_latent_states"] = torch.cat(
[text_latent_states, text_latent_states], 0
)
model_kwargs["cfg_scale"] = args.cfg_scale
model_kwargs["attention_mask"] = torch.ones(
2, 1, z.shape[1], text_latent_states.shape[1], device=device, dtype=torch.int
)
# Sample images:
samples = diffusion.p_sample_loop(
model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
model.forward_with_cfg,
z.shape,
z,
clip_denoised=False,
model_kwargs=model_kwargs,
progress=True,
device=device,
)
samples, _ = samples.chunk(2, dim=0) # Remove null class samples
samples = vae.decode(samples / 0.18215).sample
samples = col2video(
samples.squeeze(),
(num_frames // 2, in_channels, args.height // 4, args.width // 4),
)
if vqvae is not None:
# [T, C, H, W] -> [B, C, T, H, W]
samples = samples.permute(1, 0, 2, 3).unsqueeze(0)
samples = vqvae.decode_from_embeddings(samples)
# [B, C, T, H, W] -> [T, H, W, C]
samples = samples.squeeze(0).permute(1, 2, 3, 0)
else:
# [T, C, H, W] -> [T, H, W, C]
samples = samples.permute(0, 2, 3, 1)
# Save and display images:
save_image(samples, "sample.png", nrow=4, normalize=True, value_range=(-1, 1))
write_video("sample.mp4", samples.cpu(), args.fps)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2")
parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="mse")
parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
parser.add_argument("--num-classes", type=int, default=1000)
parser.add_argument(
"--model", type=str, choices=list(DiT_models.keys()), default="DiT-S/8"
)
parser.add_argument(
"--text",
type=str,
default="two ladies laughing by seeing some thing another lady throw dresses and keep it back by reverse motion",
)
parser.add_argument("--cfg-scale", type=float, default=4.0)
parser.add_argument("--num-sampling-steps", type=int, default=250)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument(
"--ckpt",
type=str,
default=None,
required=True,
help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).",
)
parser.add_argument("--vqvae", default="hpcai-tech/vqvae")
parser.add_argument(
"--text_model", type=str, default="openai/clip-vit-base-patch32"
)
parser.add_argument("--width", type=int, default=480)
parser.add_argument("--height", type=int, default=320)
parser.add_argument("--fps", type=int, default=15)
parser.add_argument("--sec", type=int, default=8)
args = parser.parse_args()
main(args)

18
train.py
View file

@ -80,12 +80,18 @@ def main(args):
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Setup model
vqvae = (
AutoModel.from_pretrained(args.vqvae, trust_remote_code=True)
.to(get_current_device())
.eval()
)
model = DiT_models[args.model]().to(get_current_device())
if len(args.vqvae) > 0:
vqvae = (
AutoModel.from_pretrained(args.vqvae, trust_remote_code=True)
.to(get_current_device())
.eval()
)
model_kwargs = {"in_channels": vqvae.embedding_dim}
else:
# disable VQ-VAE if not provided, just use raw video frames
vqvae = None
model_kwargs = {}
model = DiT_models[args.model](**model_kwargs).to(get_current_device())
patch_size = model.patch_size
ema = deepcopy(model)
requires_grad(ema, False)