From b409ef399af45decd1a31b0ebe7ec7d09286c639 Mon Sep 17 00:00:00 2001 From: qqingzheng <2533221180@qq.com> Date: Wed, 6 Mar 2024 17:23:40 +0000 Subject: [PATCH] [refactor] added more methods in abstract classes Former-commit-id: a535e3a06eb766bebe7baa3d9df64b04d04d40e6 --- examples/rec_video.py | 12 +- .../ae/videobase/configuration_videobase.py | 20 +- .../models/ae/videobase/modeling_videobase.py | 8 +- .../ae/videobase/vqvae/configuration_vqvae.py | 44 ++- .../ae/videobase/vqvae/modeling_vqvae.py | 340 +++++++++++------- .../ae/videobase/vqvae/trainer_vqvae.py | 22 +- .../train}/train_videogpt.py | 17 +- scripts/train_videogpt.sh | 4 +- train_scripts/train.py | 294 --------------- 9 files changed, 300 insertions(+), 461 deletions(-) rename {train_scripts => opensora/train}/train_videogpt.py (70%) delete mode 100644 train_scripts/train.py diff --git a/examples/rec_video.py b/examples/rec_video.py index 160c8e337..72587c1f0 100644 --- a/examples/rec_video.py +++ b/examples/rec_video.py @@ -1,3 +1,6 @@ +import sys +sys.path.append(".") + import random import cv2 import numpy as np @@ -81,8 +84,10 @@ def main(args): sample_fps = args.sample_fps sample_rate = args.sample_rate device = torch.device('cuda') - - vqvae = VQVAEModel.download_and_load_model(args.ckpt) + if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']: + vqvae = VQVAEModel.download_and_load_model(args.ckpt) + else: + vqvae = VQVAEModel.load_from_checkpoint(args.ckpt) vqvae.eval() vqvae = vqvae.to(device) @@ -99,8 +104,7 @@ if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--video-path', type=str, default='') parser.add_argument('--rec-path', type=str, default='') - parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4', - choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']) + parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4') parser.add_argument('--sample-fps', type=int, default=30) parser.add_argument('--resolution', type=int, default=336) parser.add_argument('--crop-size', type=int, default=None) diff --git a/opensora/models/ae/videobase/configuration_videobase.py b/opensora/models/ae/videobase/configuration_videobase.py index bc13717fb..cce98be87 100644 --- a/opensora/models/ae/videobase/configuration_videobase.py +++ b/opensora/models/ae/videobase/configuration_videobase.py @@ -1,5 +1,4 @@ -from dataclasses import dataclass - +import json """ We are currently working on the abstraction of the code structure, @@ -7,6 +6,19 @@ so there is no content here for the moment. It will be gradually optimized over time. Contributions from the open-source community are welcome. """ -@dataclass class VideoBaseConfiguration: - ... \ No newline at end of file + def __init__(self, **kwargs): + pass + + def to_json_string(self): + json_string = json.dumps(vars(self)) + return json_string + + def to_dict(self): + return vars(self) + + @classmethod + def load_from_file(cls, config_path): + with open(config_path, 'r') as json_file: + config_dict = json.load(json_file) + return cls(**config_dict) \ No newline at end of file diff --git a/opensora/models/ae/videobase/modeling_videobase.py b/opensora/models/ae/videobase/modeling_videobase.py index 022470893..2e8bf7e80 100644 --- a/opensora/models/ae/videobase/modeling_videobase.py +++ b/opensora/models/ae/videobase/modeling_videobase.py @@ -10,8 +10,12 @@ class VideoBaseAE(nn.Module): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) - def download_and_load_model(self): - return None + @classmethod + def load_from_checkpoint(cls, model_path): + pass + @classmethod + def download_and_load_model(cls, model_name, cache_dir=None): + pass def encode(self, x, *args, **kwargs): pass diff --git a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py index b472101d4..450735aeb 100644 --- a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py +++ b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py @@ -1,17 +1,31 @@ from ..configuration_videobase import VideoBaseConfiguration -from dataclasses import dataclass, field -from typing import Tuple, Union +from typing import Union, Tuple -@dataclass -class VQVAEConfiguration: - embedding_dim: int = field(default=256) - n_codes: int = field(default=2048) - n_hiddens: int = field(default=240) - n_res_layers: int = field(default=4) - resolution: int = field(default=128) - sequence_length: int = field(default=16) - downsample: Union[str, Tuple[int, int, int]] = field(default=(4,4,4)) - no_pos_embd: bool = field(default=True) - def __post_init__(self): - if isinstance(self.downsample, str): - self.downsample = tuple(map(int, self.downsample.split(','))) \ No newline at end of file +class VQVAEConfiguration(VideoBaseConfiguration): + def __init__( + self, + embedding_dim: int = 256, + n_codes: int = 2048, + n_hiddens: int = 240, + n_res_layers: int = 4, + resolution: int = 128, + sequence_length: int = 16, + downsample: Union[Tuple[int, int, int], str] = (4, 4, 4), + no_pos_embd: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + + self.embedding_dim = embedding_dim + self.n_codes = n_codes + self.n_hiddens = n_hiddens + self.n_res_layers = n_res_layers + self.resolution = resolution + self.sequence_length = sequence_length + + if isinstance(downsample, str): + self.downsample = tuple(map(int, downsample.split(","))) + else: + self.downsample = downsample + + self.no_pos_embd = no_pos_embd diff --git a/opensora/models/ae/videobase/vqvae/modeling_vqvae.py b/opensora/models/ae/videobase/vqvae/modeling_vqvae.py index 68f8d0a9f..2ff9e29fc 100644 --- a/opensora/models/ae/videobase/vqvae/modeling_vqvae.py +++ b/opensora/models/ae/videobase/vqvae/modeling_vqvae.py @@ -5,8 +5,11 @@ import numpy as np import torch.distributed as dist import torch.nn.functional as F import math +import os +import json from .configuration_vqvae import VQVAEConfiguration + # Copied from https://github.com/wilson1yan/VideoGPT def view_range(x, i, j, shape): shape = tuple(shape) @@ -26,6 +29,7 @@ def view_range(x, i, j, shape): target_shape = x_shape[:i] + shape + x_shape[j:] return x.view(target_shape) + # Copied from https://github.com/wilson1yan/VideoGPT def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True): n_dims = len(x.shape) @@ -49,36 +53,41 @@ def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True): x = x.contiguous() return x + # Copied from https://github.com/wilson1yan/VideoGPT -def scaled_dot_product_attention(q, k, v, mask=None, attn_dropout=0., training=True): +def scaled_dot_product_attention(q, k, v, mask=None, attn_dropout=0.0, training=True): # Performs scaled dot-product attention over the second to last dimension dn # (b, n_head, d1, ..., dn, d) attn = torch.matmul(q, k.transpose(-1, -2)) attn = attn / np.sqrt(q.shape[-1]) if mask is not None: - attn = attn.masked_fill(mask == 0, float('-inf')) + attn = attn.masked_fill(mask == 0, float("-inf")) attn_float = F.softmax(attn, dim=-1) - attn = attn_float.type_as(attn) # b x n_head x d1 x ... x dn x d + attn = attn_float.type_as(attn) # b x n_head x d1 x ... x dn x d attn = F.dropout(attn, p=attn_dropout, training=training) - a = torch.matmul(attn, v) # b x n_head x d1 x ... x dn x d + a = torch.matmul(attn, v) # b x n_head x d1 x ... x dn x d return a + # Copied from https://github.com/wilson1yan/VideoGPT class AxialBlock(nn.Module): def __init__(self, n_hiddens, n_head): super().__init__() - kwargs = dict(shape=(0,) * 3, dim_q=n_hiddens, - dim_kv=n_hiddens, n_head=n_head, - n_layer=1, causal=False, attn_type='axial') - self.attn_w = MultiHeadAttention(attn_kwargs=dict(axial_dim=-2), - **kwargs) - self.attn_h = MultiHeadAttention(attn_kwargs=dict(axial_dim=-3), - **kwargs) - self.attn_t = MultiHeadAttention(attn_kwargs=dict(axial_dim=-4), - **kwargs) + kwargs = dict( + shape=(0,) * 3, + dim_q=n_hiddens, + dim_kv=n_hiddens, + n_head=n_head, + n_layer=1, + causal=False, + attn_type="axial", + ) + self.attn_w = MultiHeadAttention(attn_kwargs=dict(axial_dim=-2), **kwargs) + self.attn_h = MultiHeadAttention(attn_kwargs=dict(axial_dim=-3), **kwargs) + self.attn_t = MultiHeadAttention(attn_kwargs=dict(axial_dim=-4), **kwargs) def forward(self, x): x = shift_dim(x, 1, -1) @@ -86,6 +95,7 @@ class AxialBlock(nn.Module): x = shift_dim(x, -1, 1) return x + # Copied from https://github.com/wilson1yan/VideoGPT class AttentionResidualBlock(nn.Module): def __init__(self, n_hiddens): @@ -99,19 +109,20 @@ class AttentionResidualBlock(nn.Module): SamePadConv3d(n_hiddens // 2, n_hiddens, 1, bias=False), nn.BatchNorm3d(n_hiddens), nn.ReLU(), - AxialBlock(n_hiddens, 2) + AxialBlock(n_hiddens, 2), ) def forward(self, x): return x + self.block(x) + # Copied from https://github.com/wilson1yan/VideoGPT class Codebook(nn.Module): def __init__(self, n_codes, embedding_dim): super().__init__() - self.register_buffer('embeddings', torch.randn(n_codes, embedding_dim)) - self.register_buffer('N', torch.zeros(n_codes)) - self.register_buffer('z_avg', self.embeddings.data.clone()) + self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim)) + self.register_buffer("N", torch.zeros(n_codes)) + self.register_buffer("z_avg", self.embeddings.data.clone()) self.n_codes = n_codes self.embedding_dim = embedding_dim @@ -133,7 +144,7 @@ class Codebook(nn.Module): y = self._tile(flat_inputs) d = y.shape[0] - _k_rand = y[torch.randperm(y.shape[0])][:self.n_codes] + _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes] if dist.is_initialized(): dist.broadcast(_k_rand, 0) self.embeddings.data.copy_(_k_rand) @@ -145,9 +156,11 @@ class Codebook(nn.Module): if self._need_init and self.training: self._init_embeddings(z) flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2) - distances = (flat_inputs ** 2).sum(dim=1, keepdim=True) \ - - 2 * flat_inputs @ self.embeddings.t() \ - + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True) + distances = ( + (flat_inputs**2).sum(dim=1, keepdim=True) + - 2 * flat_inputs @ self.embeddings.t() + + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True) + ) encoding_indices = torch.argmin(distances, dim=1) encode_onehot = F.one_hot(encoding_indices, self.n_codes).type_as(flat_inputs) @@ -175,7 +188,7 @@ class Codebook(nn.Module): self.embeddings.data.copy_(encode_normalized) y = self._tile(flat_inputs) - _k_rand = y[torch.randperm(y.shape[0])][:self.n_codes] + _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes] if dist.is_initialized(): dist.broadcast(_k_rand, 0) @@ -187,13 +200,18 @@ class Codebook(nn.Module): avg_probs = torch.mean(encode_onehot, dim=0) perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) - return dict(embeddings=embeddings_st, encodings=encoding_indices, - commitment_loss=commitment_loss, perplexity=perplexity) + return dict( + embeddings=embeddings_st, + encodings=encoding_indices, + commitment_loss=commitment_loss, + perplexity=perplexity, + ) def dictionary_lookup(self, encodings): embeddings = F.embedding(encodings, self.embeddings) return embeddings + # Copied from https://github.com/wilson1yan/VideoGPT class Encoder(nn.Module): def __init__(self, n_hiddens, n_res_layers, downsample): @@ -210,10 +228,9 @@ class Encoder(nn.Module): self.conv_last = SamePadConv3d(in_channels, n_hiddens, kernel_size=3) self.res_stack = nn.Sequential( - *[AttentionResidualBlock(n_hiddens) - for _ in range(n_res_layers)], + *[AttentionResidualBlock(n_hiddens) for _ in range(n_res_layers)], nn.BatchNorm3d(n_hiddens), - nn.ReLU() + nn.ReLU(), ) def forward(self, x): @@ -224,10 +241,12 @@ class Encoder(nn.Module): h = self.res_stack(h) return h + # Copied from https://github.com/wilson1yan/VideoGPT class MultiHeadAttention(nn.Module): - def __init__(self, shape, dim_q, dim_kv, n_head, n_layer, - causal, attn_type, attn_kwargs): + def __init__( + self, shape, dim_q, dim_kv, n_head, n_layer, causal, attn_type, attn_kwargs + ): super().__init__() self.causal = causal self.shape = shape @@ -236,30 +255,30 @@ class MultiHeadAttention(nn.Module): self.d_v = dim_kv // n_head self.n_head = n_head - self.w_qs = nn.Linear(dim_q, n_head * self.d_k, bias=False) # q + self.w_qs = nn.Linear(dim_q, n_head * self.d_k, bias=False) # q self.w_qs.weight.data.normal_(std=1.0 / np.sqrt(dim_q)) - self.w_ks = nn.Linear(dim_kv, n_head * self.d_k, bias=False) # k + self.w_ks = nn.Linear(dim_kv, n_head * self.d_k, bias=False) # k self.w_ks.weight.data.normal_(std=1.0 / np.sqrt(dim_kv)) - self.w_vs = nn.Linear(dim_kv, n_head * self.d_v, bias=False) # v + self.w_vs = nn.Linear(dim_kv, n_head * self.d_v, bias=False) # v self.w_vs.weight.data.normal_(std=1.0 / np.sqrt(dim_kv)) - self.fc = nn.Linear(n_head * self.d_v, dim_q, bias=True) # c + self.fc = nn.Linear(n_head * self.d_v, dim_q, bias=True) # c self.fc.weight.data.normal_(std=1.0 / np.sqrt(dim_q * n_layer)) - if attn_type == 'full': + if attn_type == "full": self.attn = FullAttention(shape, causal, **attn_kwargs) - elif attn_type == 'axial': - assert not causal, 'causal axial attention is not supported' + elif attn_type == "axial": + assert not causal, "causal axial attention is not supported" self.attn = AxialAttention(len(shape), **attn_kwargs) - elif attn_type == 'sparse': + elif attn_type == "sparse": self.attn = SparseAttention(shape, n_head, causal, **attn_kwargs) self.cache = None def forward(self, q, k, v, decode_step=None, decode_idx=None): - """ Compute multi-head attention + """Compute multi-head attention Args q, k, v: a [b, d1, ..., dn, c] tensor or a [b, 1, ..., 1, c] tensor if decode_step is not None @@ -286,34 +305,40 @@ class MultiHeadAttention(nn.Module): if self.causal: k_shape = (q.shape[0], n_head, *self.shape, self.d_k) v_shape = (q.shape[0], n_head, *self.shape, self.d_v) - self.cache = dict(k=torch.zeros(k_shape, dtype=k.dtype, device=q.device), - v=torch.zeros(v_shape, dtype=v.dtype, device=q.device)) + self.cache = dict( + k=torch.zeros(k_shape, dtype=k.dtype, device=q.device), + v=torch.zeros(v_shape, dtype=v.dtype, device=q.device), + ) else: # cache only once in the non-causal case self.cache = dict(k=k.clone(), v=v.clone()) if self.causal: - idx = (slice(None, None), slice(None, None), *[slice(i, i+ 1) for i in decode_idx]) - self.cache['k'][idx] = k - self.cache['v'][idx] = v - k, v = self.cache['k'], self.cache['v'] + idx = ( + slice(None, None), + slice(None, None), + *[slice(i, i + 1) for i in decode_idx], + ) + self.cache["k"][idx] = k + self.cache["v"][idx] = v + k, v = self.cache["k"], self.cache["v"] a = self.attn(q, k, v, decode_step, decode_idx) # (b, *d_shape, n_head, d) -> (b, *d_shape, n_head * d) a = shift_dim(a, 1, -2).flatten(start_dim=-2) - a = self.fc(a) # (b x seq_len x embd_dim) + a = self.fc(a) # (b x seq_len x embd_dim) return a + # Copied from https://github.com/wilson1yan/VideoGPT class Decoder(nn.Module): def __init__(self, n_hiddens, n_res_layers, upsample): super().__init__() self.res_stack = nn.Sequential( - *[AttentionResidualBlock(n_hiddens) - for _ in range(n_res_layers)], + *[AttentionResidualBlock(n_hiddens) for _ in range(n_res_layers)], nn.BatchNorm3d(n_hiddens), - nn.ReLU() + nn.ReLU(), ) n_times_upsample = np.array([int(math.log2(d)) for d in upsample]) @@ -322,8 +347,7 @@ class Decoder(nn.Module): for i in range(max_us): out_channels = 3 if i == max_us - 1 else n_hiddens us = tuple([2 if d > 0 else 1 for d in n_times_upsample]) - convt = SamePadConvTranspose3d(n_hiddens, out_channels, 4, - stride=us) + convt = SamePadConvTranspose3d(n_hiddens, out_channels, 4, stride=us) self.convts.append(convt) n_times_upsample -= 1 @@ -335,6 +359,7 @@ class Decoder(nn.Module): h = F.relu(h) return h + # Copied from https://github.com/wilson1yan/VideoGPT class SamePadConv3d(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True): @@ -347,17 +372,19 @@ class SamePadConv3d(nn.Module): # assumes that the input shape is divisible by stride total_pad = tuple([k - s for k, s in zip(kernel_size, stride)]) pad_input = [] - for p in total_pad[::-1]: # reverse since F.pad starts from last dim + for p in total_pad[::-1]: # reverse since F.pad starts from last dim pad_input.append((p // 2 + p % 2, p // 2)) pad_input = sum(pad_input, tuple()) self.pad_input = pad_input - self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, - stride=stride, padding=0, bias=bias) + self.conv = nn.Conv3d( + in_channels, out_channels, kernel_size, stride=stride, padding=0, bias=bias + ) def forward(self, x): return self.conv(F.pad(x, self.pad_input)) + # Copied from https://github.com/wilson1yan/VideoGPT class SamePadConvTranspose3d(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True): @@ -369,18 +396,24 @@ class SamePadConvTranspose3d(nn.Module): total_pad = tuple([k - s for k, s in zip(kernel_size, stride)]) pad_input = [] - for p in total_pad[::-1]: # reverse since F.pad starts from last dim + for p in total_pad[::-1]: # reverse since F.pad starts from last dim pad_input.append((p // 2 + p % 2, p // 2)) pad_input = sum(pad_input, tuple()) self.pad_input = pad_input - self.convt = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, - stride=stride, bias=bias, - padding=tuple([k - 1 for k in kernel_size])) + self.convt = nn.ConvTranspose3d( + in_channels, + out_channels, + kernel_size, + stride=stride, + bias=bias, + padding=tuple([k - 1 for k in kernel_size]), + ) def forward(self, x): return self.convt(F.pad(x, self.pad_input)) + # Copied from https://github.com/wilson1yan/VideoGPT class FullAttention(nn.Module): def __init__(self, shape, causal, attn_dropout): @@ -390,7 +423,7 @@ class FullAttention(nn.Module): seq_len = np.prod(shape) if self.causal: - self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len))) + self.register_buffer("mask", torch.tril(torch.ones(seq_len, seq_len))) def forward(self, q, k, v, decode_step, decode_idx): mask = self.mask if self.causal else None @@ -402,12 +435,13 @@ class FullAttention(nn.Module): k = k.flatten(start_dim=2, end_dim=-2) v = v.flatten(start_dim=2, end_dim=-2) - out = scaled_dot_product_attention(q, k, v, mask=mask, - attn_dropout=self.attn_dropout, - training=self.training) + out = scaled_dot_product_attention( + q, k, v, mask=mask, attn_dropout=self.attn_dropout, training=self.training + ) return view_range(out, 2, 3, old_shape) + # Copied from https://github.com/wilson1yan/VideoGPT class AxialAttention(nn.Module): def __init__(self, n_dim, axial_dim): @@ -415,7 +449,7 @@ class AxialAttention(nn.Module): if axial_dim < 0: axial_dim = 2 + n_dim + 1 + axial_dim else: - axial_dim += 2 # account for batch, head, dim + axial_dim += 2 # account for batch, head, dim self.axial_dim = axial_dim def forward(self, q, k, v, decode_step, decode_idx): @@ -430,12 +464,14 @@ class AxialAttention(nn.Module): out = shift_dim(out, -2, self.axial_dim) return out + # Copied from https://github.com/wilson1yan/VideoGPT class StridedSparsityConfig(object): """ Strided Sparse configuration specified in https://arxiv.org/abs/1904.10509 that generalizes to arbitrary dimensions """ + def __init__(self, shape, n_head, causal, block, num_local_blocks): self.n_head = n_head self.shape = shape @@ -443,8 +479,8 @@ class StridedSparsityConfig(object): self.block = block self.num_local_blocks = num_local_blocks - assert self.num_local_blocks >= 1, 'Must have at least 1 local block' - assert self.seq_len % self.block == 0, 'seq len must be divisible by block size' + assert self.num_local_blocks >= 1, "Must have at least 1 local block" + assert self.seq_len % self.block == 0, "seq len must be divisible by block size" self._block_shape = self._compute_block_shape() self._block_shape_cum = self._block_shape_cum_sizes() @@ -462,8 +498,8 @@ class StridedSparsityConfig(object): for row in range(0, num_blocks): end = min(row + self.num_local_blocks, num_blocks) for col in range( - max(0, row - self.num_local_blocks), - (row + 1 if self.causal else end)): + max(0, row - self.num_local_blocks), (row + 1 if self.causal else end) + ): layout[:, row, col] = 1 return layout @@ -487,7 +523,9 @@ class StridedSparsityConfig(object): return layout def make_layout(self): - layout = torch.zeros((self.n_head, self.num_blocks, self.num_blocks), dtype=torch.int64) + layout = torch.zeros( + (self.n_head, self.num_blocks, self.num_blocks), dtype=torch.int64 + ) layout = self.set_local_layout(layout) layout = self.set_global_layout(layout) return layout @@ -505,7 +543,7 @@ class StridedSparsityConfig(object): elem = block_layout[h, i, j].item() if elem == 1: assert i >= j - if i == j: # need to mask within block on diagonals + if i == j: # need to mask within block on diagonals attn_mask[counter] = torch.tril(attn_mask[counter]) counter += 1 assert counter == num_dense_blocks @@ -514,9 +552,9 @@ class StridedSparsityConfig(object): def get_non_block_layout_row(self, block_layout, row): block_row = row // self.block - block_row = block_layout[:, [block_row]] # n_head x 1 x n_blocks + block_row = block_layout[:, [block_row]] # n_head x 1 x n_blocks block_row = block_row.repeat_interleave(self.block, dim=-1) - block_row[:, :, row + 1:] = 0. + block_row[:, :, row + 1 :] = 0.0 return block_row ############# Helper functions ########################## @@ -540,7 +578,9 @@ class StridedSparsityConfig(object): return tuple(np.flip(np.cumprod(bs)[:-1])) + (1,) def _to_flattened_idx(self, idx): - assert len(idx) == len(self._block_shape), f"{len(idx)} != {len(self._block_shape)}" + assert len(idx) == len( + self._block_shape + ), f"{len(idx)} != {len(self._block_shape)}" flat_idx = 0 for i in range(len(self._block_shape)): flat_idx += idx[i] * self._block_shape_cum[i] @@ -554,60 +594,82 @@ class StridedSparsityConfig(object): flat_idx %= self._block_shape_cum[i] return tuple(idx) + # Copied from https://github.com/wilson1yan/VideoGPT class SparseAttention(nn.Module): ops = dict() attn_mask = dict() block_layout = dict() - def __init__(self, shape, n_head, causal, num_local_blocks=4, block=32, - attn_dropout=0.): # does not use attn_dropout + def __init__( + self, shape, n_head, causal, num_local_blocks=4, block=32, attn_dropout=0.0 + ): # does not use attn_dropout super().__init__() self.causal = causal self.shape = shape - self.sparsity_config = StridedSparsityConfig(shape=shape, n_head=n_head, - causal=causal, block=block, - num_local_blocks=num_local_blocks) + self.sparsity_config = StridedSparsityConfig( + shape=shape, + n_head=n_head, + causal=causal, + block=block, + num_local_blocks=num_local_blocks, + ) if self.shape not in SparseAttention.block_layout: - SparseAttention.block_layout[self.shape] = self.sparsity_config.make_layout() + SparseAttention.block_layout[self.shape] = ( + self.sparsity_config.make_layout() + ) if causal and self.shape not in SparseAttention.attn_mask: - SparseAttention.attn_mask[self.shape] = self.sparsity_config.make_sparse_attn_mask() + SparseAttention.attn_mask[self.shape] = ( + self.sparsity_config.make_sparse_attn_mask() + ) def get_ops(self): try: from deepspeed.ops.sparse_attention import MatMul, Softmax except: - raise Exception('Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`') + raise Exception( + "Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`" + ) if self.shape not in SparseAttention.ops: sparsity_layout = self.sparsity_config.make_layout() - sparse_dot_sdd_nt = MatMul(sparsity_layout, - self.sparsity_config.block, - 'sdd', - trans_a=False, - trans_b=True) + sparse_dot_sdd_nt = MatMul( + sparsity_layout, + self.sparsity_config.block, + "sdd", + trans_a=False, + trans_b=True, + ) - sparse_dot_dsd_nn = MatMul(sparsity_layout, - self.sparsity_config.block, - 'dsd', - trans_a=False, - trans_b=False) + sparse_dot_dsd_nn = MatMul( + sparsity_layout, + self.sparsity_config.block, + "dsd", + trans_a=False, + trans_b=False, + ) sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block) - SparseAttention.ops[self.shape] = (sparse_dot_sdd_nt, - sparse_dot_dsd_nn, - sparse_softmax) + SparseAttention.ops[self.shape] = ( + sparse_dot_sdd_nt, + sparse_dot_dsd_nn, + sparse_softmax, + ) return SparseAttention.ops[self.shape] def forward(self, q, k, v, decode_step, decode_idx): if self.training and self.shape not in SparseAttention.ops: self.get_ops() - SparseAttention.block_layout[self.shape] = SparseAttention.block_layout[self.shape].to(q) + SparseAttention.block_layout[self.shape] = SparseAttention.block_layout[ + self.shape + ].to(q) if self.causal: - SparseAttention.attn_mask[self.shape] = SparseAttention.attn_mask[self.shape].to(q).type_as(q) + SparseAttention.attn_mask[self.shape] = ( + SparseAttention.attn_mask[self.shape].to(q).type_as(q) + ) attn_mask = SparseAttention.attn_mask[self.shape] if self.causal else None old_shape = q.shape[2:-1] @@ -616,39 +678,43 @@ class SparseAttention(nn.Module): v = v.flatten(start_dim=2, end_dim=-2) if decode_step is not None: - mask = self.sparsity_config.get_non_block_layout_row(SparseAttention.block_layout[self.shape], decode_step) - out = scaled_dot_product_attention(q, k, v, mask=mask, training=self.training) + mask = self.sparsity_config.get_non_block_layout_row( + SparseAttention.block_layout[self.shape], decode_step + ) + out = scaled_dot_product_attention( + q, k, v, mask=mask, training=self.training + ) else: if q.shape != k.shape or k.shape != v.shape: - raise Exception('SparseAttention only support self-attention') + raise Exception("SparseAttention only support self-attention") sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops() scaling = float(q.shape[-1]) ** -0.5 attn_output_weights = sparse_dot_sdd_nt(q, k) if attn_mask is not None: - attn_output_weights = attn_output_weights.masked_fill(attn_mask == 0, - float('-inf')) - attn_output_weights = sparse_softmax( - attn_output_weights, - scale=scaling - ) + attn_output_weights = attn_output_weights.masked_fill( + attn_mask == 0, float("-inf") + ) + attn_output_weights = sparse_softmax(attn_output_weights, scale=scaling) out = sparse_dot_dsd_nn(attn_output_weights, v) return view_range(out, 2, 3, old_shape) - + + # Modified from https://github.com/wilson1yan/VideoGPT class VQVAEModel(VideoBaseAE): - + DOWNLOADED_VQVAE = { - 'bair_stride4x2x2': '1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L', - 'ucf101_stride4x4x4': '1uuB_8WzHP_bbBmfuaIV7PK_Itl3DyHY5', - 'kinetics_stride4x4x4': '1DOvOZnFAIQmux6hG7pN_HkyJZy3lXbCB', - 'kinetics_stride2x4x4': '1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB' + "bair_stride4x2x2": "1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L", + "ucf101_stride4x4x4": "1uuB_8WzHP_bbBmfuaIV7PK_Itl3DyHY5", + "kinetics_stride4x4x4": "1DOvOZnFAIQmux6hG7pN_HkyJZy3lXbCB", + "kinetics_stride2x4x4": "1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB", } def __init__(self, config: VQVAEConfiguration): super().__init__() + self.config = config self.embedding_dim = config.embedding_dim self.n_codes = config.n_codes self.encoder = Encoder(config.n_hiddens, config.n_res_layers, config.downsample) @@ -656,45 +722,53 @@ class VQVAEModel(VideoBaseAE): self.pre_vq_conv = SamePadConv3d(config.n_hiddens, config.embedding_dim, 1) self.post_vq_conv = SamePadConv3d(config.embedding_dim, config.n_hiddens, 1) self.codebook = Codebook(config.n_codes, config.embedding_dim) - + def forward(self, x): z = self.pre_vq_conv(self.encoder(x)) vq_output = self.codebook(z) - x_recon = self.decoder(self.post_vq_conv(vq_output['embeddings'])) + x_recon = self.decoder(self.post_vq_conv(vq_output["embeddings"])) recon_loss = F.mse_loss(x_recon, x) / 0.06 return recon_loss, x_recon, vq_output - + def encode(self, x, include_embeddings=False): h = self.pre_vq_conv(self.encoder(x)) vq_output = self.codebook(h) if include_embeddings: - return vq_output['encodings'], vq_output['embeddings'] + return vq_output["encodings"], vq_output["embeddings"] else: - return vq_output['encodings'] + return vq_output["encodings"] def decode(self, encodings): h = F.embedding(encodings, self.codebook.embeddings) h = self.post_vq_conv(shift_dim(h, -1, 1)) return self.decoder(h) - - @staticmethod - def load_from_checkpoint(model_path): - model_cpkt = torch.load(model_path) - - # Compatible with old videogpt model formats. - if 'hyper_parameters' in model_cpkt: - hyper_parameters = vars(model_cpkt.get("hyper_parameters").get("args")) - print(hyper_parameters) - state_dict = model_cpkt.get("state_dict") - model = VQVAEModel(config=VQVAEConfiguration(**hyper_parameters)) + + @classmethod + def load_from_checkpoint(cls, model_path): + if not os.path.isdir(model_path): + """model downloaded from internet""" + model_cpkt = torch.load(model_path) + # Compatible with old videogpt model formats. + if "hyper_parameters" in model_cpkt: + hyper_parameters = vars(model_cpkt.get("hyper_parameters").get("args")) + state_dict = model_cpkt.get("state_dict") + model = cls(config=VQVAEConfiguration(**hyper_parameters)) + model.load_state_dict(state_dict) + return model + else: + raise RuntimeError("Model checkpoint has a wrong format.") + else: + with open(os.path.join(model_path, "config.json"), "r") as file: + config = json.load(file) + state_dict = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu") + model = cls(config=VQVAEConfiguration(**config)) model.load_state_dict(state_dict) return model - else: - raise RuntimeError("Model checkpoint has a wrong format.") - - @staticmethod - def download_and_load_model(model_name, cache_dir=None): + + @classmethod + def download_and_load_model(cls, model_name, cache_dir=None): from .....utils.downloader import gdown_download - path = gdown_download(VQVAEModel.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir) - return VQVAEModel.load_from_checkpoint(path) - \ No newline at end of file + path = gdown_download( + cls.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir + ) + return cls.load_from_checkpoint(path) diff --git a/opensora/models/ae/videobase/vqvae/trainer_vqvae.py b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py index c2ede2ed0..5b465c7f6 100644 --- a/opensora/models/ae/videobase/vqvae/trainer_vqvae.py +++ b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py @@ -1,15 +1,29 @@ -from torch import nn from transformers import Trainer import torch.nn.functional as F +from typing import Optional +import os +import torch +from transformers.utils import WEIGHTS_NAME +import json class VQVAETrainer(Trainer): - + def compute_loss(self, model, inputs, return_outputs=False): model = model.module x = inputs.get("video") z = model.pre_vq_conv(model.encoder(x)) vq_output = model.codebook(z) - x_recon = model.decoder(model.post_vq_conv(vq_output['embeddings'])) + x_recon = model.decoder(model.post_vq_conv(vq_output["embeddings"])) recon_loss = F.mse_loss(x_recon, x) / 0.06 - return recon_loss \ No newline at end of file + return recon_loss + + def _save(self, output_dir: Optional[str] = None, state_dict=None): + output_dir = output_dir if output_dir is not None else self.args.output_dir + os.makedirs(output_dir, exist_ok=True) + if state_dict is None: + state_dict = self.model.state_dict() + with open(os.path.join(output_dir, "config.json"), "w") as file: + json.dump(self.model.config.to_dict(), file) + torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + torch.save(self.args, os.path.join(output_dir, "training_args.bin")) diff --git a/train_scripts/train_videogpt.py b/opensora/train/train_videogpt.py similarity index 70% rename from train_scripts/train_videogpt.py rename to opensora/train/train_videogpt.py index 532c94995..a9c102b43 100644 --- a/train_scripts/train_videogpt.py +++ b/opensora/train/train_videogpt.py @@ -1,5 +1,6 @@ import sys sys.path.append(".") + from opensora.models.ae.videobase import ( VQVAEModel, VQVAEConfiguration, @@ -10,11 +11,19 @@ import argparse from typing import Optional from accelerate.utils import set_seed from transformers import HfArgumentParser, TrainingArguments -from dataclasses import dataclass, field +from dataclasses import dataclass, field, asdict @dataclass -class VQVAEArgument(VQVAEConfiguration): +class VQVAEArgument: + embedding_dim: int = field(default=256), + n_codes: int = field(default=2048), + n_hiddens: int = field(default=240), + n_res_layers: int = field(default=4), + resolution: int = field(default=128), + sequence_length: int = field(default=16), + downsample: str = field(default="4,4,4"), + no_pos_embd: bool = True, data_path: str = field(default=None, metadata={"help": "data path"}) @dataclass @@ -24,8 +33,10 @@ class VQVAETrainingArgument(TrainingArguments): ) def train(args, vqvae_args, training_args): + # Load Config + config = VQVAEConfiguration(**asdict(vqvae_args)) # Load Model - model = VQVAEModel(vqvae_args) + model = VQVAEModel(config) # Load Dataset dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length) diff --git a/scripts/train_videogpt.sh b/scripts/train_videogpt.sh index e4498bfcd..45a9558fc 100644 --- a/scripts/train_videogpt.sh +++ b/scripts/train_videogpt.sh @@ -1,6 +1,6 @@ accelerate launch \ --config_file scripts/accelerate_configs/default_config.yaml \ - train_scripts/train_videogpt.py \ + opensora/train/train_videogpt.py \ --do_train \ --seed 1234 \ --data_path UCF-101 \ @@ -15,7 +15,7 @@ accelerate launch \ --save_strategy steps \ --save_total_limit 5 \ --logging_steps 5 \ - --save_steps 500 \ + --save_steps 5 \ --n_codes 2048 \ --n_hiddens 240 \ --n_res_layers 4 \ diff --git a/train_scripts/train.py b/train_scripts/train.py deleted file mode 100644 index 625355d07..000000000 --- a/train_scripts/train.py +++ /dev/null @@ -1,294 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -""" -A minimal training script for DiT using PyTorch DDP. -""" -import torch -# the first flag below was False when we tested this script but True makes A100 training a lot faster: -from torch import nn - -torch.backends.cuda.matmul.allow_tf32 = True -torch.backends.cudnn.allow_tf32 = True -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler -from collections import OrderedDict -from copy import deepcopy -from glob import glob -from time import time -import argparse -import logging -import os -from opensora.dataset import getdataset -from opensora.models.ae import getae -from opensora.utils.dataset_utils import Collate -from opensora.models.diffusion.dit.models import DiT_models as Diffusion_models -from opensora.models.diffusion.dit.diffusion import create_diffusion - -################################################################################# -# Training Helper Functions # -################################################################################# - -@torch.no_grad() -def update_ema(ema_model, model, decay=0.9999): - """ - Step the EMA model towards the current model. - """ - ema_params = OrderedDict(ema_model.named_parameters()) - model_params = OrderedDict(model.named_parameters()) - - for name, param in model_params.items(): - # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed - ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) - - -def requires_grad(model, flag=True): - """ - Set requires_grad flag for all parameters in a model. - """ - for p in model.parameters(): - p.requires_grad = flag - - -def cleanup(): - """ - End DDP training. - """ - dist.destroy_process_group() - - -def create_logger(logging_dir): - """ - Create a logger that writes to a log file and stdout. - """ - if dist.get_rank() == 0: # real logger - logging.basicConfig( - level=logging.INFO, - format='[\033[34m%(asctime)s\033[0m] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] - ) - logger = logging.getLogger(__name__) - else: # dummy logger (does nothing) - logger = logging.getLogger(__name__) - logger.addHandler(logging.NullHandler()) - return logger - -################################################################################# -# Training Loop # -################################################################################# - -def main(args): - """ - Trains a new DiT model. - """ - assert torch.cuda.is_available(), "Training currently requires at least one GPU." - - # Setup DDP: - dist.init_process_group("nccl") - assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size." - rank = dist.get_rank() - device = rank % torch.cuda.device_count() - seed = args.global_seed * dist.get_world_size() + rank - torch.manual_seed(seed) - torch.cuda.set_device(device) - print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.") - - # Setup an experiment folder: - if rank == 0: - print(args) - os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders) - experiment_index = len(glob(f"{args.results_dir}/*")) - model_string_name = args.model.replace("/", "-") # e.g., DiT-XL/2 --> DiT-XL-2 (for naming folders) - experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Create an experiment folder - checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints - os.makedirs(checkpoint_dir, exist_ok=True) - logger = create_logger(experiment_dir) - logger.info(f"Experiment directory created at {experiment_dir}") - else: - logger = create_logger(None) - - # Create model: - vae_stride_t, vae_stride_h, vae_stride_w = [int(i) for i in args.ae[-5:].split('x')] - args.vae_stride_t, args.vae_stride_h, args.vae_stride_w = vae_stride_t, vae_stride_h, vae_stride_w - args.vae_stride = args.vae_stride_h - patch_size = args.model[-3:] - patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) - args.patch_size = patch_size_h - args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w - assert args.num_frames % vae_stride_t == 0, "Image size must be divisible by 8 (for the VAE encoder)." - assert args.max_image_size % vae_stride_h == 0, "Image size must be divisible by 8 (for the VAE encoder)." - assert vae_stride_h == vae_stride_w, "Support now." - assert patch_size_h == patch_size_w, "Support now." - - latent_size = (args.num_frames // vae_stride_t, args.max_image_size // vae_stride_h, args.max_image_size // vae_stride_w) - - model = Diffusion_models[args.model]( - input_size=latent_size, - num_classes=args.num_classes - ) - model.gradient_checkpointing = args.gradient_checkpointing - - - if args.pt_ckpt is not None: - state_dict = torch.load(args.pt_ckpt, map_location='cpu') - if state_dict.get('model', None) is not None: - state_dict = state_dict['model'] - del state_dict['x_embedder.proj.weight'] - del state_dict['final_layer.linear.weight'] - del state_dict['final_layer.linear.bias'] - del state_dict['pos_embed'] - del state_dict['y_embedder.embedding_table.weight'] - missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) - print('missing_keys:', missing_keys, 'unexpected_keys:', unexpected_keys) - assert len(missing_keys) == 4 - logger.info(f"{missing_keys}, {unexpected_keys}") - - # Note that parameter initialization is done within the DiT constructor - ema = deepcopy(model).to(device) # Create an EMA of the model for use after training - requires_grad(ema, False) - model = DDP(model.to(device), device_ids=[rank]) - diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule - ae = getae(args).to(device) - logger.info(f"{model}") - logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}") - for n, p in model.named_parameters(): - if p.requires_grad: - logger.info(f"Training Parameters: {n}") - - # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper): - opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0) - - # Setup data: - dataset = getdataset(args) - sampler = DistributedSampler( - dataset, - num_replicas=dist.get_world_size(), - rank=rank, - shuffle=True, - seed=args.global_seed - ) - loader = DataLoader( - dataset, - batch_size=int(args.global_batch_size // dist.get_world_size()), - shuffle=False, - sampler=sampler, - num_workers=args.num_workers, - pin_memory=True, - drop_last=True, - collate_fn=Collate(args) - ) - logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})") - - # Prepare models for training: - update_ema(ema, model.module, decay=0) # Ensure EMA is initialized with synced weights - model.train() # important! This enables embedding dropout for classifier-free guidance - ema.eval() # EMA model should always be in eval mode - - # Variables for monitoring/logging purposes: - train_steps = 0 - log_steps = 0 - running_loss = 0 - start_time = time() - - logger.info(f"Training for {args.epochs} epochs...") - for epoch in range(args.epochs): - sampler.set_epoch(epoch) - logger.info(f"Beginning epoch {epoch}...") - for x, y, attn_mask in loader: - x = x.to(device) - y = y.to(device) - attn_mask = attn_mask.to(device) - with torch.no_grad(): - x = ae.encode(x) # Map input images to latent space + normalize latents - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) - model_kwargs = dict(y=y, attention_mask=attn_mask) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() - opt.zero_grad() - - loss.backward() - if args.clip_grad_norm is not None: - nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm) - opt.step() - - update_ema(ema, model.module) - - # Log loss values: - running_loss += loss.item() - log_steps += 1 - train_steps += 1 - if train_steps % args.log_every == 0: - # Measure training speed: - torch.cuda.synchronize() - end_time = time() - 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() / dist.get_world_size() - logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}") - # Reset monitoring variables: - running_loss = 0 - log_steps = 0 - start_time = time() - - # Save DiT checkpoint: - if train_steps % args.ckpt_every == 0 and train_steps > 0: - if rank == 0: - checkpoint = { - "model": model.module.state_dict(), - "ema": ema.state_dict(), - "opt": opt.state_dict(), - "args": args - } - checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" - torch.save(checkpoint, checkpoint_path) - logger.info(f"Saved checkpoint to {checkpoint_path}") - dist.barrier() - - model.eval() # important! This disables randomized embedding dropout - # do any sampling/FID calculation/etc. with ema (or model) in eval mode ... - - logger.info("Done!") - cleanup() - - -if __name__ == "__main__": - # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters). - parser = argparse.ArgumentParser() - parser.add_argument("--dataset", type=str, required=True) - parser.add_argument("--data-path", type=str, required=True) - parser.add_argument("--results-dir", type=str, default="results") - parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122") - parser.add_argument("--num-classes", type=int, default=1000) - parser.add_argument("--epochs", type=int, default=1400) - parser.add_argument("--global-batch-size", type=int, default=256) - parser.add_argument("--global-seed", type=int, default=0) - parser.add_argument("--num-workers", type=int, default=8) - parser.add_argument("--log-every", type=int, default=100) - parser.add_argument("--ckpt-every", type=int, default=50_000) - - - # -------------------------------------- - parser.add_argument("--ae", type=str, choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', - 'kinetics_stride4x4x4', 'kinetics_stride2x4x4'], - default="ucf101_stride4x4x4") - parser.add_argument("--pt-ckpt", type=str, default=None) - parser.add_argument("--sample-rate", type=int, default=4) - parser.add_argument("--num-frames", type=int, default=16) - parser.add_argument("--max-image-size", type=int, default=128) - parser.add_argument("--dynamic-frames", action="store_true") - parser.add_argument("--gradient-checkpointing", action="store_true") - parser.add_argument("--lr", type=float, default=1e-4) - parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)") - # -------------------------------------- - - - args = parser.parse_args() - main(args)