From 29e167fecfbfff51d752a3a60873afd3e2bfe6e9 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sat, 9 Mar 2024 18:18:26 +0800 Subject: [PATCH 01/28] Fix einops bug: module 'keras.backend' has no attribute 'is_tensor' Former-commit-id: 849f8ed815295bac78c3c10d319e6731aa26ca02 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index adaa474a8..fa4e2648c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ classifiers = [ ] dependencies = [ "torch==2.1.2", "torchvision==0.16.2", - "transformers==4.32.0", "albumentations==1.4.0", "av==11.0.0", "decord==0.6.0", "einops==0.3.0", "fastapi==0.110.0", + "transformers==4.32.0", "albumentations==1.4.0", "av==11.0.0", "decord==0.6.0", "einops==0.7.0", "fastapi==0.110.0", "accelerate==0.21.0", "gdown==5.1.0", "h5py==3.10.0", "idna==3.6", 'imageio==2.34.0', "matplotlib==3.7.5", "numpy==1.24.4", "omegaconf==2.1.1", "opencv-python==4.9.0.80", "opencv-python-headless==4.9.0.80", "pandas==2.0.3", "pillow==10.2.0", "pydub==0.25.1", "pytorch-lightning==1.4.2", "pytorchvideo==0.1.5", "PyYAML==6.0.1", "regex==2023.12.25", From aef95b043c569a006434b76a11017ec2e2097b86 Mon Sep 17 00:00:00 2001 From: jpthu17 <1003901354@qq.com> Date: Sat, 9 Mar 2024 19:53:48 +0800 Subject: [PATCH 02/28] [feat] Add 2D RoPE Former-commit-id: ccae80d339e34fb23706827565b5095ae2c2f320 --- opensora/models/diffusion/dit/VisionRoPE.py | 89 +++++++++++++++++++ opensora/models/diffusion/dit/dit.py | 36 +++++++- opensora/models/diffusion/latte/VisionRoPE.py | 89 +++++++++++++++++++ opensora/models/diffusion/latte/latte.py | 52 +++++++++-- 4 files changed, 257 insertions(+), 9 deletions(-) create mode 100644 opensora/models/diffusion/dit/VisionRoPE.py create mode 100644 opensora/models/diffusion/latte/VisionRoPE.py diff --git a/opensora/models/diffusion/dit/VisionRoPE.py b/opensora/models/diffusion/dit/VisionRoPE.py new file mode 100644 index 000000000..c6b01f572 --- /dev/null +++ b/opensora/models/diffusion/dit/VisionRoPE.py @@ -0,0 +1,89 @@ +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# EVA: https://github.com/baaivision/EVA +# Transformer升级之路：2、博采众长的旋转式位置编码: https://spaces.ac.cn/archives/8265 +# Transformer升级之路：4、二维位置的旋转式位置编码: https://spaces.ac.cn/archives/8397 +# -------------------------------------------------------- +import torch +import torch.nn as nn +from math import pi +from einops import rearrange, repeat + + +def broadcat(tensors, dim=-1): + num_tensors = len(tensors) + shape_lens = set(list(map(lambda t: len(t.shape), tensors))) + assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions' + shape_len = list(shape_lens)[0] + dim = (dim + shape_len) if dim < 0 else dim + dims = list(zip(*map(lambda t: list(t.shape), tensors))) + expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] + assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation' + max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) + expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) + expanded_dims.insert(dim, (dim, dims[dim])) + expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) + tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) + return torch.cat(tensors, dim=dim) + + +def rotate_half(x): + x = rearrange(x, '... (d r) -> ... d r', r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, '... d r -> ... (d r)') + + +class VisionRotaryEmbeddingFast(nn.Module): + def __init__( + self, + dim, + pt_hw=(int, int), # (H, W) + ft_hw=None, + custom_freqs = None, + freqs_for = 'lang', + theta = 10000, + max_freq = 10, + num_freqs = 1, + ): + super().__init__() + # Unlike a 1d RoPE, a 2d RoPE requires splitting the dimension into four parts + # References: https://spaces.ac.cn/archives/8397 + + if custom_freqs: + freqs = custom_freqs + elif freqs_for == 'lang': + freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) + elif freqs_for == 'pixel': + freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi + elif freqs_for == 'constant': + freqs = torch.ones(num_freqs).float() + else: + raise ValueError(f'unknown modality {freqs_for}') + + if ft_hw is None: ft_hw = pt_hw + h_t = torch.arange(ft_hw[0]) / ft_hw[0] * pt_hw[0] + w_t = torch.arange(ft_hw[1]) / ft_hw[1] * pt_hw[1] + + h_freqs = torch.einsum('..., f -> ... f', h_t, freqs) + w_freqs = torch.einsum('..., f -> ... f', w_t, freqs) + + h_freqs = repeat(h_freqs, '... n -> ... (n r)', r=2) + w_freqs = repeat(w_freqs, '... n -> ... (n r)', r=2) + + freqs = broadcat((h_freqs[:, None, :], w_freqs[None, :, :]), dim=-1) + freqs_cos = freqs.cos().view(-1, freqs.shape[-1]) + freqs_sin = freqs.sin().view(-1, freqs.shape[-1]) + + self.register_buffer("freqs_cos", freqs_cos) + self.register_buffer("freqs_sin", freqs_sin) + + def forward(self, t): + # 2d RoPE: [[cos(h*theta), -sin(h*theta), 0, 0 ], + # [sin(h*theta), cos(h*theta), 0, 0 ], + # [0, 0, cos(w*theta), -sin(w*theta)], + # [0, 0, sin(w*theta), cos(w*theta) ],] + + return t * self.freqs_cos + rotate_half(t) * self.freqs_sin \ No newline at end of file diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index ddc6b4894..8b295c6d8 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -5,7 +5,7 @@ # GLIDE: https://github.com/openai/glide-text2im # MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py # -------------------------------------------------------- -from typing import Final, Optional +from typing import Final, Optional, Tuple, Union import torch import torch.nn as nn @@ -17,7 +17,7 @@ from timm.layers import use_fused_attn, to_2tuple # from timm.models.vision_transformer import PatchEmbed, Attention, Mlp from timm.models.vision_transformer import PatchEmbed, Mlp from torch.nn import functional as F - +from .VisionRoPE import VisionRotaryEmbeddingFast # the xformers lib allows less memory, faster training and inference try: import xformers @@ -50,7 +50,11 @@ class Attention(nn.Module): attn_drop: float = 0., proj_drop: float = 0., norm_layer: nn.Module = nn.LayerNorm, - attention_mode='math' + attention_mode='math', + attention_pe_mode='2d_rope', + hw: Union[int, Tuple[int, int]] = 16, # (h, w) + pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) + intp_vfreq: bool = False, # vision position interpolation ) -> None: super().__init__() assert dim % num_heads == 0, 'dim should be divisible by num_heads' @@ -58,6 +62,7 @@ class Attention(nn.Module): self.head_dim = dim // num_heads self.scale = self.head_dim ** -0.5 self.attention_mode = attention_mode + self.attention_pe_mode = attention_pe_mode self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() @@ -66,11 +71,29 @@ class Attention(nn.Module): self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) + if self.attention_pe_mode == '2d_rope': + half_head_dim = dim // num_heads // 2 + self.hw = to_2tuple(hw) + self.rope = VisionRotaryEmbeddingFast( + dim=half_head_dim, + pt_hw=to_2tuple(pt_hw), + ft_hw=self.hw if intp_vfreq else None, + ) + def forward(self, x: torch.Tensor, attention_mask) -> torch.Tensor: B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c + if self.attention_pe_mode == '2d_rope': + q_t = q.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_q_t = self.rope(q_t) + q = ro_q_t.view(B, self.num_heads, N, C // self.num_heads) + + k_t = k.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_k_t = self.rope(k_t) + k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) + if self.attention_mode == 'xformers': # cause loss nan while using with amp assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool())) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) @@ -256,6 +279,9 @@ class DiT(nn.Module): learn_sigma=True, extras=1, attention_mode='math', + attention_pe_mode='2d_rope', + pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) + intp_vfreq: bool = False, # vision position interpolation ): super().__init__() self.gradient_checkpointing = False @@ -280,7 +306,9 @@ class DiT(nn.Module): self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False) self.blocks = nn.ModuleList([ - DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode) for _ in range(depth) + DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode, + attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size, + intp_vfreq=intp_vfreq) for _ in range(depth) ]) self.final_layer = FinalLayer(hidden_size, patch_size_t, patch_size, self.out_channels) self.initialize_weights() diff --git a/opensora/models/diffusion/latte/VisionRoPE.py b/opensora/models/diffusion/latte/VisionRoPE.py new file mode 100644 index 000000000..c6b01f572 --- /dev/null +++ b/opensora/models/diffusion/latte/VisionRoPE.py @@ -0,0 +1,89 @@ +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# EVA: https://github.com/baaivision/EVA +# Transformer升级之路：2、博采众长的旋转式位置编码: https://spaces.ac.cn/archives/8265 +# Transformer升级之路：4、二维位置的旋转式位置编码: https://spaces.ac.cn/archives/8397 +# -------------------------------------------------------- +import torch +import torch.nn as nn +from math import pi +from einops import rearrange, repeat + + +def broadcat(tensors, dim=-1): + num_tensors = len(tensors) + shape_lens = set(list(map(lambda t: len(t.shape), tensors))) + assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions' + shape_len = list(shape_lens)[0] + dim = (dim + shape_len) if dim < 0 else dim + dims = list(zip(*map(lambda t: list(t.shape), tensors))) + expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] + assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation' + max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) + expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) + expanded_dims.insert(dim, (dim, dims[dim])) + expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) + tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) + return torch.cat(tensors, dim=dim) + + +def rotate_half(x): + x = rearrange(x, '... (d r) -> ... d r', r=2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return rearrange(x, '... d r -> ... (d r)') + + +class VisionRotaryEmbeddingFast(nn.Module): + def __init__( + self, + dim, + pt_hw=(int, int), # (H, W) + ft_hw=None, + custom_freqs = None, + freqs_for = 'lang', + theta = 10000, + max_freq = 10, + num_freqs = 1, + ): + super().__init__() + # Unlike a 1d RoPE, a 2d RoPE requires splitting the dimension into four parts + # References: https://spaces.ac.cn/archives/8397 + + if custom_freqs: + freqs = custom_freqs + elif freqs_for == 'lang': + freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) + elif freqs_for == 'pixel': + freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi + elif freqs_for == 'constant': + freqs = torch.ones(num_freqs).float() + else: + raise ValueError(f'unknown modality {freqs_for}') + + if ft_hw is None: ft_hw = pt_hw + h_t = torch.arange(ft_hw[0]) / ft_hw[0] * pt_hw[0] + w_t = torch.arange(ft_hw[1]) / ft_hw[1] * pt_hw[1] + + h_freqs = torch.einsum('..., f -> ... f', h_t, freqs) + w_freqs = torch.einsum('..., f -> ... f', w_t, freqs) + + h_freqs = repeat(h_freqs, '... n -> ... (n r)', r=2) + w_freqs = repeat(w_freqs, '... n -> ... (n r)', r=2) + + freqs = broadcat((h_freqs[:, None, :], w_freqs[None, :, :]), dim=-1) + freqs_cos = freqs.cos().view(-1, freqs.shape[-1]) + freqs_sin = freqs.sin().view(-1, freqs.shape[-1]) + + self.register_buffer("freqs_cos", freqs_cos) + self.register_buffer("freqs_sin", freqs_sin) + + def forward(self, t): + # 2d RoPE: [[cos(h*theta), -sin(h*theta), 0, 0 ], + # [sin(h*theta), cos(h*theta), 0, 0 ], + # [0, 0, cos(w*theta), -sin(w*theta)], + # [0, 0, sin(w*theta), cos(w*theta) ],] + + return t * self.freqs_cos + rotate_half(t) * self.freqs_sin \ No newline at end of file diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index ed6824f5b..9b069f64d 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -11,10 +11,11 @@ import math import torch import torch.nn as nn import numpy as np - +from .VisionRoPE import VisionRotaryEmbeddingFast from einops import rearrange, repeat from timm.models.vision_transformer import Mlp, PatchEmbed - +from typing import Tuple, Union +from timm.layers import to_2tuple import torch.utils.checkpoint as cp # for i in sys.path: @@ -50,13 +51,30 @@ def modulate(x, shift, scale): ################################################################################# class Attention(nn.Module): - def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math', eps=1e-12, causal=True, ring_bucket_size=1024): + def __init__(self, + dim, + num_heads=8, + qkv_bias=False, + attn_drop=0., + proj_drop=0., + use_lora=False, + attention_mode='math', + eps=1e-12, + causal=True, + ring_bucket_size=1024, + attention_pe_mode='2d_rope', + hw: Union[int, Tuple[int, int]] = 16, # (h, w) + pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) + intp_vfreq: bool = False, # vision position interpolation + ): super().__init__() assert dim % num_heads == 0, 'dim should be divisible by num_heads' self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim ** -0.5 self.attention_mode = attention_mode + self.attention_pe_mode = attention_pe_mode + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) @@ -65,11 +83,29 @@ class Attention(nn.Module): self.causal = causal self.ring_bucket_size = ring_bucket_size + if self.attention_pe_mode == '2d_rope': + half_head_dim = dim // num_heads // 2 + self.hw = to_2tuple(hw) + self.rope = VisionRotaryEmbeddingFast( + dim=half_head_dim, + pt_hw=to_2tuple(pt_hw), + ft_hw=self.hw if intp_vfreq else None, + ) + def forward(self, x, attention_mask): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c - + + if self.attention_pe_mode == '2d_rope': + q_t = q.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_q_t = self.rope(q_t) + q = ro_q_t.view(B, self.num_heads, N, C // self.num_heads) + + k_t = k.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_k_t = self.rope(k_t) + k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) + if self.attention_mode == 'xformers': # cause loss nan while using with amp assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool())) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) @@ -249,6 +285,9 @@ class Latte(nn.Module): learn_sigma=True, extras=1, attention_mode='math', + attention_pe_mode='2d_rope', + pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) + intp_vfreq: bool = False, # vision position interpolation ): super().__init__() self.learn_sigma = learn_sigma @@ -280,7 +319,10 @@ class Latte(nn.Module): self.hidden_size = hidden_size self.blocks = nn.ModuleList([ - TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode) for _ in range(depth) + TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode, + attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size, + intp_vfreq=intp_vfreq + ) for _ in range(depth) ]) self.final_layer = FinalLayer(hidden_size, patch_size_t, patch_size, self.out_channels) From c08f985c8c441700cf6ddbcef5942fd4f7e76ec9 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sat, 9 Mar 2024 21:19:47 +0800 Subject: [PATCH 03/28] Update README.md Former-commit-id: 6630c22baa2340a8b55badd61472daab7177bc63 --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 8e850a26e..88bb5b06b 100644 --- a/README.md +++ b/README.md @@ -78,9 +78,9 @@ Project stages: - [x] Add evaluation scripts (FVD, CLIP score). 🤝 Thanks to [@rain305f](https://github.com/rain305f) #### Train models that boost resolution and duration -- [ ] Add [PI](https://arxiv.org/abs/2306.15595) to support out-of-domain size. 🙏 **[Need your contribution]** -- [ ] Add 2D RoPE to improve generalization ability as [FiT](https://github.com/whlzy/FiT). 🙏 **[Need your contribution]** -- [ ] Train a **low dimension** Video-AE, whether it is VAE or VQVAE. 🚀 **[Require more computation]** +- [x] Add [PI](https://arxiv.org/abs/2306.15595) to support out-of-domain size. 🤝 Thanks to [@jpthu17](https://github.com/jpthu17) +- [x] Add 2D RoPE to improve generalization ability as [FiT](https://github.com/whlzy/FiT). 🤝 Thanks to [@jpthu17](https://github.com/jpthu17) +- [ ] Train a **low dimension** Video-AE, whether it is VAE or VQVAE. ⌛ [WIP] 🚀 **[Require more computation]** - [x] Extract offline feature. - [x] Train with offline feature. - [x] Add frame interpolation model. 🤝 Thanks to [@yunyangge](https://github.com/yunyangge) From 5ebb382d85a7595970530c693ba4a8292b47ba41 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sat, 9 Mar 2024 21:52:03 +0800 Subject: [PATCH 04/28] xformers and flashattn Former-commit-id: 002cad95253f7da36344237ce8c8410ffb517711 --- opensora/models/diffusion/dit/dit.py | 42 ++++++++++---------- opensora/models/diffusion/latte/latte.py | 50 +++++++++++------------- 2 files changed, 43 insertions(+), 49 deletions(-) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index 8b295c6d8..8e5832920 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -18,12 +18,7 @@ from timm.layers import use_fused_attn, to_2tuple from timm.models.vision_transformer import PatchEmbed, Mlp from torch.nn import functional as F from .VisionRoPE import VisionRotaryEmbeddingFast -# the xformers lib allows less memory, faster training and inference -try: - import xformers - import xformers.ops -except: - XFORMERS_IS_AVAILBLE = False +from torch.nn import functional as F try: # needs to have https://github.com/corl-team/rebased/ installed @@ -80,10 +75,12 @@ class Attention(nn.Module): ft_hw=self.hw if intp_vfreq else None, ) - def forward(self, x: torch.Tensor, attention_mask) -> torch.Tensor: + def forward(self, x: torch.Tensor, attn_mask) -> torch.Tensor: B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c + if attn_mask is not None: + attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype) if self.attention_pe_mode == '2d_rope': q_t = q.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) @@ -94,24 +91,23 @@ class Attention(nn.Module): ro_k_t = self.rope(k_t) k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) - if self.attention_mode == 'xformers': # cause loss nan while using with amp - assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool())) - q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + if self.attention_mode == 'xformers': # cause loss nan while using with amp + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 elif self.attention_mode == 'flash': - assert attention_mask is None or torch.all(attention_mask.bool()) # cause loss nan while using with amp # Optionally use the context manager to ensure one of the fused kerenels is run with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p, - scale=self.scale).reshape(B, N, - C) # require pytorch 2.0 + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 elif self.attention_mode == 'math': attn = (q @ k.transpose(-2, -1)) * self.scale - if attention_mask is not None: - attn = attn + attention_mask + if attn_mask is not None: + attn_bias = self.make_attn_bias(attn_mask) + attn = attn + attn_bias attn = attn.softmax(dim=-1) if torch.any(torch.isnan(attn)): print('torch.any(torch.isnan(attn))') @@ -132,6 +128,10 @@ class Attention(nn.Module): x = self.proj_drop(x) return x + def make_attn_bias(self, attn_mask): + attn_bias = torch.where(attn_mask == 0, -1e8, attn_mask) + attn_bias = torch.where(attn_mask == 1, 0., attn_bias) + return attn_bias def modulate(x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) @@ -374,15 +374,13 @@ class DiT(nn.Module): return outputs return ckpt_forward - def make_mask(self, attention_mask, dtype): + def make_mask(self, attention_mask): 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) return attention_mask - def forward(self, x, t, y, attention_mask=None, use_fp16=False): + def forward(self, x, t, y, attention_mask=None): """ Forward pass of DiT. x: (B, T, C, H, W) tensor of spatial inputs (images or latent representations of images) @@ -395,7 +393,7 @@ class DiT(nn.Module): B, T, C, H, W = x.shape self.t = T // self.patch_size_t # 4 // 1 if attention_mask is not None: - attention_mask = self.make_mask(attention_mask, x.dtype) + attention_mask = self.make_mask(attention_mask) diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index 9b069f64d..42fdb5922 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -17,16 +17,7 @@ from timm.models.vision_transformer import Mlp, PatchEmbed from typing import Tuple, Union from timm.layers import to_2tuple import torch.utils.checkpoint as cp - -# for i in sys.path: -# print(i) - -# the xformers lib allows less memory, faster training and inference -try: - import xformers - import xformers.ops -except: - XFORMERS_IS_AVAILBLE = False +from torch.nn import functional as F try: # needs to have https://github.com/corl-team/rebased/ installed @@ -62,7 +53,7 @@ class Attention(nn.Module): eps=1e-12, causal=True, ring_bucket_size=1024, - attention_pe_mode='2d_rope', + attention_pe_mode=None, hw: Union[int, Tuple[int, int]] = 16, # (h, w) pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) intp_vfreq: bool = False, # vision position interpolation @@ -92,10 +83,12 @@ class Attention(nn.Module): ft_hw=self.hw if intp_vfreq else None, ) - def forward(self, x, attention_mask): + def forward(self, x, attn_mask): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c + if attn_mask is not None: + attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype) if self.attention_pe_mode == '2d_rope': q_t = q.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) @@ -107,21 +100,22 @@ class Attention(nn.Module): k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) if self.attention_mode == 'xformers': # cause loss nan while using with amp - assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool())) - q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 elif self.attention_mode == 'flash': - assert attention_mask is None or torch.all(attention_mask.bool()) # cause loss nan while using with amp # Optionally use the context manager to ensure one of the fused kerenels is run with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 elif self.attention_mode == 'math': attn = (q @ k.transpose(-2, -1)) * self.scale - if attention_mask is not None: - attn = attn + attention_mask + if attn_mask is not None: + attn_bias = self.make_attn_bias(attn_mask) + attn = attn + attn_bias attn = attn.softmax(dim=-1) if torch.any(torch.isnan(attn)): print('torch.any(torch.isnan(attn))') @@ -142,6 +136,10 @@ class Attention(nn.Module): x = self.proj_drop(x) return x + def make_attn_bias(self, attn_mask): + attn_bias = torch.where(attn_mask == 0, -1e8, attn_mask) + attn_bias = torch.where(attn_mask == 1, 0., attn_bias) + return attn_bias ################################################################################# # Embedding Layers for Timesteps and Class Labels # @@ -238,9 +236,9 @@ class TransformerBlock(nn.Module): nn.Linear(hidden_size, 6 * hidden_size, bias=True) ) - def forward(self, x, c, attention_mask): + def forward(self, x, c, attn_bias): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) - x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), attention_mask) + x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), attn_bias) x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) return x @@ -285,7 +283,7 @@ class Latte(nn.Module): learn_sigma=True, extras=1, attention_mode='math', - attention_pe_mode='2d_rope', + attention_pe_mode=None, pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) intp_vfreq: bool = False, # vision position interpolation ): @@ -389,12 +387,10 @@ class Latte(nn.Module): return outputs return ckpt_forward - def make_mask(self, attention_mask, dtype): + def make_mask(self, attention_mask): 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, 1e-8) - # attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(dtype).min) return attention_mask # @torch.cuda.amp.autocast() @@ -415,10 +411,10 @@ class Latte(nn.Module): attention_mask_temproal, attention_mask_spatial = None, None if attention_mask is not None: attention_mask_spatial = rearrange(attention_mask, 'b t h w -> (b t) h w') - attention_mask_spatial = self.make_mask(attention_mask_spatial, x.dtype) + attention_mask_spatial = self.make_mask(attention_mask_spatial) attention_mask_temproal = rearrange(attention_mask, 'b t h w -> (b h w) t') - attention_mask_temproal = self.make_mask(attention_mask_temproal, x.dtype) + attention_mask_temproal = self.make_mask(attention_mask_temproal) batches, frames, channels, high, weight = x.shape From 9a0d0bc843ab05218f23ea324c53b561f6b49cfc Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sat, 9 Mar 2024 21:53:43 +0800 Subject: [PATCH 05/28] safe dit Former-commit-id: f407cab7dc0e53257972d430e873f4fd77bb039d --- opensora/models/diffusion/dit/dit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index 8e5832920..b6f3df4f5 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -46,7 +46,7 @@ class Attention(nn.Module): proj_drop: float = 0., norm_layer: nn.Module = nn.LayerNorm, attention_mode='math', - attention_pe_mode='2d_rope', + attention_pe_mode=None, hw: Union[int, Tuple[int, int]] = 16, # (h, w) pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) intp_vfreq: bool = False, # vision position interpolation @@ -279,7 +279,7 @@ class DiT(nn.Module): learn_sigma=True, extras=1, attention_mode='math', - attention_pe_mode='2d_rope', + attention_pe_mode=None, pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) intp_vfreq: bool = False, # vision position interpolation ): From 75938cfaeb1fe0d95e073848af579a04fd20f547 Mon Sep 17 00:00:00 2001 From: jpthu17 <1003901354@qq.com> Date: Sat, 9 Mar 2024 22:04:15 +0800 Subject: [PATCH 06/28] [fix] 2d RoPE init Former-commit-id: ff5931bae6efa94675b9ef6bc25f1b76217d128e --- opensora/models/diffusion/dit/dit.py | 8 +++++--- opensora/models/diffusion/latte/latte.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index b6f3df4f5..e514a213a 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -49,7 +49,7 @@ class Attention(nn.Module): attention_pe_mode=None, hw: Union[int, Tuple[int, int]] = 16, # (h, w) pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) - intp_vfreq: bool = False, # vision position interpolation + intp_vfreq: bool = True, # vision position interpolation ) -> None: super().__init__() assert dim % num_heads == 0, 'dim should be divisible by num_heads' @@ -280,8 +280,8 @@ class DiT(nn.Module): extras=1, attention_mode='math', attention_pe_mode=None, - pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) - intp_vfreq: bool = False, # vision position interpolation + pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w) + intp_vfreq: bool = True, # vision position interpolation ): super().__init__() self.gradient_checkpointing = False @@ -305,6 +305,8 @@ class DiT(nn.Module): self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False) self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False) + if pt_input_size is None: + pt_input_size = input_size self.blocks = nn.ModuleList([ DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode, attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size, diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index 42fdb5922..e89a87e57 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -56,7 +56,7 @@ class Attention(nn.Module): attention_pe_mode=None, hw: Union[int, Tuple[int, int]] = 16, # (h, w) pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) - intp_vfreq: bool = False, # vision position interpolation + intp_vfreq: bool = True, # vision position interpolation ): super().__init__() assert dim % num_heads == 0, 'dim should be divisible by num_heads' @@ -284,8 +284,8 @@ class Latte(nn.Module): extras=1, attention_mode='math', attention_pe_mode=None, - pt_input_size: Union[int, Tuple[int, int]] = 16, # (h, w) - intp_vfreq: bool = False, # vision position interpolation + pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w) + intp_vfreq: bool = True, # vision position interpolation ): super().__init__() self.learn_sigma = learn_sigma @@ -316,6 +316,8 @@ class Latte(nn.Module): self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False) self.hidden_size = hidden_size + if pt_input_size is None: + pt_input_size = input_size self.blocks = nn.ModuleList([ TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode, attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size, From f034826c6d0d0d148d102455a3000122b8fa6dee Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sun, 10 Mar 2024 11:12:12 +0800 Subject: [PATCH 07/28] fixed attn_mask in flash-attn Former-commit-id: 5e2504f429a342b75bc7a9004cee7b3f75582ee3 --- opensora/models/diffusion/dit/dit.py | 22 ++-- opensora/models/diffusion/latte/latte.py | 143 +++++++++++++---------- 2 files changed, 97 insertions(+), 68 deletions(-) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index e514a213a..62ed31977 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -91,17 +91,21 @@ class Attention(nn.Module): ro_k_t = self.rope(k_t) k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) - if self.attention_mode == 'xformers': # cause loss nan while using with amp + if self.attention_mode == 'xformers': # require pytorch 2.0 with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) - elif self.attention_mode == 'flash': - # cause loss nan while using with amp - # Optionally use the context manager to ensure one of the fused kerenels is run - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + elif self.attention_mode == 'flash': # require pytorch 2.0 + # https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/109 + if attn_mask is None or torch.all(attn_mask.bool()): + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = F.scaled_dot_product_attention(q, k, v, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + else: # turn to xformers + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) elif self.attention_mode == 'math': attn = (q @ k.transpose(-2, -1)) * self.scale @@ -550,4 +554,4 @@ DiT_models = { 'DiT-L/122': DiT_L_122, 'DiT-L/144': DiT_L_144, 'DiT-L/188': DiT_L_188, 'DiT-B/122': DiT_B_122, 'DiT-B/144': DiT_B_144, 'DiT-B/188': DiT_B_188, 'DiT-S/122': DiT_S_122, 'DiT-S/144': DiT_S_144, 'DiT-S/188': DiT_S_188, -} +} \ No newline at end of file diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index e89a87e57..564005dd9 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -30,13 +30,15 @@ try: from ring_attention_pytorch.ring_flash_attention_cuda import ring_flash_attn_cuda except: RING_ATTENTION_IS_AVAILABLE = False - + + # from timm.models.layers.helpers import to_2tuple # from timm.models.layers.trace_utils import _assert def modulate(x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + ################################################################################# # Attention Layers from TIMM # ################################################################################# @@ -86,7 +88,7 @@ class Attention(nn.Module): def forward(self, x, attn_mask): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() - q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c + q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c if attn_mask is not None: attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype) @@ -99,18 +101,22 @@ class Attention(nn.Module): ro_k_t = self.rope(k_t) k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) - if self.attention_mode == 'xformers': # cause loss nan while using with amp + if self.attention_mode == 'xformers': # require pytorch 2.0 with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + + elif self.attention_mode == 'flash': # require pytorch 2.0 + # https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/109 + if attn_mask is None or torch.all(attn_mask.bool()): + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = F.scaled_dot_product_attention(q, k, v, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + else: # turn to xformers + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) - elif self.attention_mode == 'flash': - # cause loss nan while using with amp - # Optionally use the context manager to ensure one of the fused kerenels is run - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 - elif self.attention_mode == 'math': attn = (q @ k.transpose(-2, -1)) * self.scale if attn_mask is not None: @@ -141,6 +147,7 @@ class Attention(nn.Module): attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias + ################################################################################# # Embedding Layers for Timesteps and Class Labels # ################################################################################# @@ -149,6 +156,7 @@ class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. """ + def __init__(self, hidden_size, frequency_embedding_size=256): super().__init__() self.mlp = nn.Sequential( @@ -189,6 +197,7 @@ class LabelEmbedder(nn.Module): """ Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. """ + def __init__(self, num_classes, hidden_size, dropout_prob): super().__init__() use_cfg_embedding = dropout_prob > 0 @@ -223,6 +232,7 @@ class TransformerBlock(nn.Module): """ A Latte tansformer block with adaptive layer norm zero (adaLN-Zero) conditioning. """ + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): super().__init__() self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) @@ -247,6 +257,7 @@ class FinalLayer(nn.Module): """ The final layer of Latte. """ + def __init__(self, hidden_size, patch_size_t, patch_size, out_channels): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) @@ -267,25 +278,26 @@ class Latte(nn.Module): """ Diffusion model with a Transformer backbone. """ + def __init__( - self, - input_size=32, - patch_size=2, - patch_size_t=1, - in_channels=4, - hidden_size=1152, - depth=28, - num_heads=16, - mlp_ratio=4.0, - num_frames=16, - class_dropout_prob=0.1, - num_classes=1000, - learn_sigma=True, - extras=1, - attention_mode='math', - attention_pe_mode=None, - pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w) - intp_vfreq: bool = True, # vision position interpolation + self, + input_size=32, + patch_size=2, + patch_size_t=1, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + num_frames=16, + class_dropout_prob=0.1, + num_classes=1000, + learn_sigma=True, + extras=1, + attention_mode='math', + attention_pe_mode=None, + pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w) + intp_vfreq: bool = True, # vision position interpolation ): super().__init__() self.learn_sigma = learn_sigma @@ -304,11 +316,11 @@ class Latte(nn.Module): if self.extras == 2: self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob) - if self.extras == 78: # timestep + text_embedding + if self.extras == 78: # timestep + text_embedding self.text_embedding_projection = nn.Sequential( - nn.SiLU(), - nn.Linear(77 * 768, hidden_size, bias=True) - ) + nn.SiLU(), + nn.Linear(77 * 768, hidden_size, bias=True) + ) num_patches = self.x_embedder.num_patches # Will use fixed sin-cos embedding: @@ -335,6 +347,7 @@ class Latte(nn.Module): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.constant_(module.bias, 0) + self.apply(_basic_init) # Initialize (and freeze) pos_embed by sin-cos embedding: @@ -387,6 +400,7 @@ class Latte(nn.Module): def ckpt_forward(*inputs): outputs = module(*inputs) return outputs + return ckpt_forward def make_mask(self, attention_mask): @@ -397,10 +411,10 @@ class Latte(nn.Module): # @torch.cuda.amp.autocast() # @torch.compile - def forward(self, - x, - t, - y=None, + def forward(self, + x, + t, + y=None, text_embedding=None, attention_mask=None): """ @@ -425,7 +439,6 @@ class Latte(nn.Module): x = self.x_embedder(x) + self.pos_embed t = self.t_embedder(t) - timestep_spatial = repeat(t, 'n d -> (n c) d', c=self.temp_embed.shape[1]) timestep_temp = repeat(t, 'n d -> (n c) d', c=self.pos_embed.shape[1]) @@ -439,7 +452,7 @@ class Latte(nn.Module): text_embedding_temp = repeat(text_embedding, 'n d -> (n c) d', c=self.pos_embed.shape[1]) for i in range(0, len(self.blocks), 2): - spatial_block, temp_block = self.blocks[i:i+2] + spatial_block, temp_block = self.blocks[i:i + 2] if self.extras == 2: c = timestep_spatial + y_spatial elif self.extras == 78: @@ -473,8 +486,8 @@ class Latte(nn.Module): c = timestep_spatial + y_spatial else: c = timestep_spatial - x = self.final_layer(x, c) - x = self.unpatchify(x) + x = self.final_layer(x, c) + x = self.unpatchify(x) x = rearrange(x, '(b f) c h w -> b f c h w', b=batches) return x @@ -494,7 +507,7 @@ class Latte(nn.Module): # eps, rest = model_out[:, :, :4, ...], model_out[:, :, 4:, ...] 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) + eps = torch.cat([half_eps, half_eps], dim=0) return torch.cat([eps, rest], dim=2) @@ -507,6 +520,7 @@ def get_1d_sincos_temp_embed(embed_dim, length): pos = torch.arange(0, length).unsqueeze(1) return get_1d_sincos_pos_embed_from_grid(embed_dim, pos) + def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): """ grid_size: int of the grid height and width @@ -529,8 +543,8 @@ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): assert embed_dim % 2 == 0 # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) emb = np.concatenate([emb_h, emb_w], axis=1) return emb @@ -545,15 +559,15 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float64) omega /= embed_dim / 2. - omega = 1. / 10000**omega + omega = 1. / 10000 ** omega - pos = pos.reshape(-1) - out = np.einsum('m,d->md', pos, omega) + pos = pos.reshape(-1) + out = np.einsum('m,d->md', pos, omega) - emb_sin = np.sin(out) - emb_cos = np.cos(out) + emb_sin = np.sin(out) + emb_cos = np.cos(out) - emb = np.concatenate([emb_sin, emb_cos], axis=1) + emb = np.concatenate([emb_sin, emb_cos], axis=1) return emb @@ -564,49 +578,59 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): def Latte_XL_122(**kwargs): return Latte(depth=28, hidden_size=1152, patch_size_t=1, patch_size=2, num_heads=16, **kwargs) + def Latte_XL_144(**kwargs): return Latte(depth=28, hidden_size=1152, patch_size_t=1, patch_size=4, num_heads=16, **kwargs) + def Latte_XL_188(**kwargs): return Latte(depth=28, hidden_size=1152, patch_size_t=1, patch_size=8, num_heads=16, **kwargs) + def Latte_L_122(**kwargs): return Latte(depth=24, hidden_size=1024, patch_size_t=1, patch_size=2, num_heads=16, **kwargs) + def Latte_L_144(**kwargs): return Latte(depth=24, hidden_size=1024, patch_size_t=1, patch_size=4, num_heads=16, **kwargs) + def Latte_L_188(**kwargs): return Latte(depth=24, hidden_size=1024, patch_size_t=1, patch_size=8, num_heads=16, **kwargs) + def Latte_B_122(**kwargs): return Latte(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs) + def Latte_B_144(**kwargs): return Latte(depth=12, hidden_size=768, patch_size_t=1, patch_size=4, num_heads=12, **kwargs) + def Latte_B_188(**kwargs): return Latte(depth=12, hidden_size=768, patch_size_t=1, patch_size=8, num_heads=12, **kwargs) + def Latte_S_122(**kwargs): return Latte(depth=12, hidden_size=384, patch_size_t=1, patch_size=2, num_heads=6, **kwargs) + def Latte_S_144(**kwargs): return Latte(depth=12, hidden_size=384, patch_size_t=1, patch_size=4, num_heads=6, **kwargs) + def Latte_S_188(**kwargs): return Latte(depth=12, hidden_size=384, patch_size_t=1, patch_size=8, num_heads=6, **kwargs) Latte_models = { - 'Latte-XL/122': Latte_XL_122, 'Latte-XL/144': Latte_XL_144, 'Latte-XL/188': Latte_XL_188, - 'Latte-L/122': Latte_L_122, 'Latte-L/144': Latte_L_144, 'Latte-L/188': Latte_L_188, - 'Latte-B/122': Latte_B_122, 'Latte-B/144': Latte_B_144, 'Latte-B/188': Latte_B_188, - 'Latte-S/122': Latte_S_122, 'Latte-S/144': Latte_S_144, 'Latte-S/188': Latte_S_188, + 'Latte-XL/122': Latte_XL_122, 'Latte-XL/144': Latte_XL_144, 'Latte-XL/188': Latte_XL_188, + 'Latte-L/122': Latte_L_122, 'Latte-L/144': Latte_L_144, 'Latte-L/188': Latte_L_188, + 'Latte-B/122': Latte_B_122, 'Latte-B/144': Latte_B_144, 'Latte-B/188': Latte_B_188, + 'Latte-S/122': Latte_S_122, 'Latte-S/144': Latte_S_144, 'Latte-S/188': Latte_S_188, } if __name__ == '__main__': - import torch device = "cuda" if torch.cuda.is_available() else "cpu" @@ -615,12 +639,13 @@ if __name__ == '__main__': t = torch.tensor([1, 2, 3]).to(device) y = torch.tensor([1, 2, 3]).to(device) network = Latte_XL_122().to(device) - from thop import profile + from thop import profile + flops, params = profile(network, inputs=(img, t)) - print('FLOPs = ' + str(flops/1000**3) + 'G') - print('Params = ' + str(params/1000**2) + 'M') + print('FLOPs = ' + str(flops / 1000 ** 3) + 'G') + print('Params = ' + str(params / 1000 ** 2) + 'M') # y_embeder = LabelEmbedder(num_classes=101, hidden_size=768, dropout_prob=0.5).to(device) # lora.mark_only_lora_as_trainable(network) # out = y_embeder(y, True) # out = network(img, t, y) - # print(out.shape) + # print(out.shape) \ No newline at end of file From d6d54da131867d5af801ede23cc44a218d8fce3e Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 11:15:14 +0800 Subject: [PATCH 08/28] typo Former-commit-id: 9e87366e906ed19726c8905afb27b5df4d4e1512 --- opensora/models/diffusion/dit/dit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index 62ed31977..3bf6e8942 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -103,7 +103,7 @@ class Attention(nn.Module): x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) else: # turn to xformers - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) @@ -554,4 +554,4 @@ DiT_models = { 'DiT-L/122': DiT_L_122, 'DiT-L/144': DiT_L_144, 'DiT-L/188': DiT_L_188, 'DiT-B/122': DiT_B_122, 'DiT-B/144': DiT_B_144, 'DiT-B/188': DiT_B_188, 'DiT-S/122': DiT_S_122, 'DiT-S/144': DiT_S_144, 'DiT-S/188': DiT_S_188, -} \ No newline at end of file +} From ee3810a585ad7aeab8b64a9520dc78df4ec6edd8 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 11:24:39 +0800 Subject: [PATCH 09/28] Decouple mixed_precision and gradient_accumulation to command args. Former-commit-id: 903d473d1432eb492a87efe184d03439f643b92d --- opensora/train/train.py | 26 ++++++++++++++-------- scripts/accelerate_configs/ddp_config.yaml | 11 +++++++++ scripts/train_256.sh | 20 +++++++++++++++++ 3 files changed, 48 insertions(+), 9 deletions(-) create mode 100644 scripts/accelerate_configs/ddp_config.yaml create mode 100644 scripts/train_256.sh diff --git a/opensora/train/train.py b/opensora/train/train.py index bd4b851e1..9fc88c604 100644 --- a/opensora/train/train.py +++ b/opensora/train/train.py @@ -48,10 +48,8 @@ def main(args): assert torch.cuda.is_available(), "Training currently requires at least one GPU." # Setup accelerator: - accelerator = Accelerator() + accelerator = Accelerator(mixed_precision=args.mixed_precision, gradient_accumulation_steps=args.gradient_accumulation_steps) device = accelerator.device - args.mixed_precision = accelerator.mixed_precision - args.gradient_accumulation_steps = accelerator.gradient_accumulation_steps # Setup an experiment folder: if accelerator.is_main_process: @@ -206,6 +204,14 @@ def main(args): if accelerator.is_main_process: logger.info(f"Training for {num_train_epochs} epochs...") + + if args.mixed_precision == "bf16": + dtype = torch.bfloat16 + elif args.mixed_precision == "fp16": + dtype = torch.float16 + else: + dtype = torch.float32 + for epoch in range(first_epoch, num_train_epochs): logger.info(f"Beginning epoch {epoch}...") for step, (x, y, attn_mask) in enumerate(loader): @@ -218,7 +224,6 @@ def main(args): # Map input images to latent space + normalize latents x = ae.encode(x) # B C T H W -> B T C H W - if args.extras == 78: # text-to-video raise 'T2V training are Not supported at this moment!' elif args.extras == 2: @@ -226,10 +231,11 @@ def main(args): else: model_kwargs = dict(y=None, attention_mask=attn_mask) - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() - opt.zero_grad() + with torch.autocast(device_type='cuda', dtype=dtype): + t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) + loss_dict = diffusion.training_losses(model, x, t, model_kwargs) + loss = loss_dict["loss"].mean() + opt.zero_grad() accelerator.backward(loss) @@ -322,9 +328,11 @@ if __name__ == "__main__": parser.add_argument("--lr-warmup-steps", type=int, default=0) parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math") + parser.add_argument("--gradient-accumulation-steps", type=int, default=1) + parser.add_argument("--mixed-precision", type=str, default=None, choices=[None, "fp16", "bf16"]) parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)") # -------------------------------------- - args = parser.parse_args() + main(args) diff --git a/scripts/accelerate_configs/ddp_config.yaml b/scripts/accelerate_configs/ddp_config.yaml new file mode 100644 index 000000000..fa6cfbdd6 --- /dev/null +++ b/scripts/accelerate_configs/ddp_config.yaml @@ -0,0 +1,11 @@ +compute_environment: LOCAL_MACHINE +distributed_type: MULTI_GPU +fsdp_config: {} +machine_rank: 0 +main_process_ip: null +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +gpu_ids: 0,1,2,3,4,5,6,7 +use_cpu: false \ No newline at end of file diff --git a/scripts/train_256.sh b/scripts/train_256.sh new file mode 100644 index 000000000..9d7912cf5 --- /dev/null +++ b/scripts/train_256.sh @@ -0,0 +1,20 @@ + +accelerate launch \ + --config_file scripts/accelerate_configs/ddp_config.yaml \ + opensora/train/train.py \ + --model Latte-XL/122 \ + --dataset sky \ + --ae stabilityai/sd-vae-ft-mse \ + --data-path datasets/sky_timelapse/sky_train \ + --extras 1 \ + --sample-rate 3 \ + --num-frames 16 \ + --max-image-size 256 \ + --max-train-steps 1000000 \ + --local-batch-size 5 \ + --lr 1e-4 \ + --ckpt-every 500 \ + --log-every 50 \ + --gradient-checkpointing \ + --attention_mode math \ + --mixed-precision fp16 From bfe4f15694058eccfc23d7cbd84743c2806967a1 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 11:25:07 +0800 Subject: [PATCH 10/28] Fix mixed_precision bugs of latte. Former-commit-id: d70a551972e04d95444d46cdbe2bfb16ad099fcb --- opensora/models/diffusion/latte/latte.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index e89a87e57..132dbbea4 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -137,7 +137,7 @@ class Attention(nn.Module): return x def make_attn_bias(self, attn_mask): - attn_bias = torch.where(attn_mask == 0, -1e8, attn_mask) + attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else 1e-4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias From 443bab13640a593fb367ad02490f3c0efa3afbb4 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 11:25:50 +0800 Subject: [PATCH 11/28] Refactor latte. Former-commit-id: 304534269f105578097bc16ab497636ad59fd9b2 --- .../diffusion/latte/configuration_latte.py | 92 +++ .../models/diffusion/latte/modeling_latte.py | 647 ++++++++++++++++++ 2 files changed, 739 insertions(+) create mode 100644 opensora/models/diffusion/latte/configuration_latte.py create mode 100644 opensora/models/diffusion/latte/modeling_latte.py diff --git a/opensora/models/diffusion/latte/configuration_latte.py b/opensora/models/diffusion/latte/configuration_latte.py new file mode 100644 index 000000000..0aaaa9d83 --- /dev/null +++ b/opensora/models/diffusion/latte/configuration_latte.py @@ -0,0 +1,92 @@ +from typing import Union, Tuple + +from ..configuration_base import BaseConfiguration + + +class LatteConfiguration(BaseConfiguration): + + def __init__( + self, + input_size=32, + patch_size=2, + patch_size_t=1, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + num_frames=16, + class_dropout_prob=0.1, + num_classes=1000, + learn_sigma=True, + extras=1, + attention_mode='math', + attention_pe_mode=None, + pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w) + intp_vfreq: bool = True, # vision position interpolation + **kwargs, + ): + super().__init__(**kwargs) + + self.input_size = input_size + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.in_channels = in_channels + self.hidden_size = hidden_size + self.depth = depth + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.num_frames = num_frames + self.class_dropout_prob = class_dropout_prob + self.num_classes = num_classes + self.learn_sigma = learn_sigma + self.extras = learn_sigma + self.attention_mode = attention_mode + self.attention_pe_mode = attention_pe_mode + self.pt_input_size = pt_input_size + self.intp_vfreq = intp_vfreq + + +def Latte_XL_122_Config(**kwargs): + return LatteConfiguration(depth=28, hidden_size=1152, patch_size_t=1, patch_size=2, num_heads=16, **kwargs) + +def Latte_XL_144_Config(**kwargs): + return LatteConfiguration(depth=28, hidden_size=1152, patch_size_t=1, patch_size=4, num_heads=16, **kwargs) + +def Latte_XL_188_Config(**kwargs): + return LatteConfiguration(depth=28, hidden_size=1152, patch_size_t=1, patch_size=8, num_heads=16, **kwargs) + +def Latte_L_122_Config(**kwargs): + return LatteConfiguration(depth=24, hidden_size=1024, patch_size_t=1, patch_size=2, num_heads=16, **kwargs) + +def Latte_L_144_Config(**kwargs): + return LatteConfiguration(depth=24, hidden_size=1024, patch_size_t=1, patch_size=4, num_heads=16, **kwargs) + +def Latte_L_188_Config(**kwargs): + return LatteConfiguration(depth=24, hidden_size=1024, patch_size_t=1, patch_size=8, num_heads=16, **kwargs) + +def Latte_B_122_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=768, patch_size_t=1, patch_size=2, num_heads=12, **kwargs) + +def Latte_B_144_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=768, patch_size_t=1, patch_size=4, num_heads=12, **kwargs) + +def Latte_B_188_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=768, patch_size_t=1, patch_size=8, num_heads=12, **kwargs) + +def Latte_S_122_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=384, patch_size_t=1, patch_size=2, num_heads=6, **kwargs) + +def Latte_S_144_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=384, patch_size_t=1, patch_size=4, num_heads=6, **kwargs) + +def Latte_S_188_Config(**kwargs): + return LatteConfiguration(depth=12, hidden_size=384, patch_size_t=1, patch_size=8, num_heads=6, **kwargs) + + +Latte_configs = { + "Latte-XL/122": Latte_XL_122_Config, "Latte-XL/144": Latte_XL_144_Config, "Latte-XL/188": Latte_XL_188_Config, + "Latte-L/122": Latte_L_122_Config, "Latte-L/144": Latte_L_144_Config, "Latte-L/188": Latte_L_188_Config, + "Latte-B/122": Latte_B_122_Config, "Latte-B/144": Latte_B_144_Config, "Latte-B/188": Latte_B_188_Config, + "Latte-S/122": Latte_S_122_Config, "Latte-S/144": Latte_S_144_Config, "Latte-S/188": Latte_S_188_Config, +} diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py new file mode 100644 index 000000000..1b9018e0b --- /dev/null +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -0,0 +1,647 @@ +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# GLIDE: https://github.com/openai/glide-text2im +# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py +# -------------------------------------------------------- +import math + +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp +import numpy as np +from torch.nn import functional as F +from einops import rearrange, repeat +from typing import Tuple, Union +from timm.layers import to_2tuple +from timm.models.vision_transformer import Mlp, PatchEmbed + +from .VisionRoPE import VisionRotaryEmbeddingFast + +try: + # needs to have https://github.com/corl-team/rebased/ installed + from fla.ops.triton.rebased_fast import parallel_rebased +except: + REBASED_IS_AVAILABLE = False + +try: + # needs to have https://github.com/lucidrains/ring-attention-pytorch installed + from ring_attention_pytorch.ring_flash_attention_cuda import ring_flash_attn_cuda +except: + RING_ATTENTION_IS_AVAILABLE = False + +from .configuration_latte import LatteConfiguration + +# from timm.models.layers.helpers import to_2tuple +# from timm.models.layers.trace_utils import _assert + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + +################################################################################# +# Attention Layers from TIMM # +################################################################################# + +class Attention(nn.Module): + def __init__(self, + dim, + num_heads=8, + qkv_bias=False, + attn_drop=0., + proj_drop=0., + use_lora=False, + attention_mode='math', + eps=1e-12, + causal=True, + ring_bucket_size=1024, + attention_pe_mode=None, + hw: Union[int, Tuple[int, int]] = 16, # (h, w) + pt_hw: Union[int, Tuple[int, int]] = 16, # (h, w) + intp_vfreq: bool = True, # vision position interpolation + ): + super().__init__() + assert dim % num_heads == 0, 'dim should be divisible by num_heads' + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim ** -0.5 + self.attention_mode = attention_mode + self.attention_pe_mode = attention_pe_mode + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + self.eps = eps + self.causal = causal + self.ring_bucket_size = ring_bucket_size + + if self.attention_pe_mode == '2d_rope': + half_head_dim = dim // num_heads // 2 + self.hw = to_2tuple(hw) + self.rope = VisionRotaryEmbeddingFast( + dim=half_head_dim, + pt_hw=to_2tuple(pt_hw), + ft_hw=self.hw if intp_vfreq else None, + ) + + def forward(self, x, attn_mask): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() + q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c + if attn_mask is not None: + attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype) + + if self.attention_pe_mode == '2d_rope': + q_t = q.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_q_t = self.rope(q_t) + q = ro_q_t.view(B, self.num_heads, N, C // self.num_heads) + + k_t = k.view(B, self.num_heads, -1, self.hw[0] * self.hw[1], C // self.num_heads) + ro_k_t = self.rope(k_t) + k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) + + if self.attention_mode == 'xformers': # cause loss nan while using with amp + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + + elif self.attention_mode == 'flash': + # cause loss nan while using with amp + # Optionally use the context manager to ensure one of the fused kerenels is run + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + + elif self.attention_mode == 'math': + attn = (q @ k.transpose(-2, -1)) * self.scale + if attn_mask is not None: + attn_bias = self.make_attn_bias(attn_mask) + attn = attn + attn_bias + attn = attn.softmax(dim=-1) + if torch.any(torch.isnan(attn)): + print('torch.any(torch.isnan(attn))') + attn = attn.masked_fill(torch.isnan(attn), float(0.)) + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + + elif self.attention_mode == 'rebased': + x = parallel_rebased(q, k, v, self.eps, True, True).reshape(B, N, C) + + elif self.attention_mode == 'ring': + x = ring_flash_attn_cuda(q, k, v, causal=self.causal, bucket_size=self.ring_bucket_size).reshape(B, N, C) + + else: + raise NotImplemented + + x = self.proj(x) + x = self.proj_drop(x) + return x + + def make_attn_bias(self, attn_mask): + attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else 1e-4, attn_mask) + attn_bias = torch.where(attn_mask == 1, 0., attn_bias) + return attn_bias + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + +class LabelEmbedder(nn.Module): + """ + Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. + """ + def __init__(self, num_classes, hidden_size, dropout_prob): + super().__init__() + use_cfg_embedding = dropout_prob > 0 + self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) + self.num_classes = num_classes + self.dropout_prob = dropout_prob + + def token_drop(self, labels, force_drop_ids=None): + """ + Drops labels to enable classifier-free guidance. + """ + if force_drop_ids is None: + drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob + else: + drop_ids = force_drop_ids == 1 + labels = torch.where(drop_ids, self.num_classes, labels) + return labels + + def forward(self, labels, train, force_drop_ids=None): + use_dropout = self.dropout_prob > 0 + if (train and use_dropout) or (force_drop_ids is not None): + labels = self.token_drop(labels, force_drop_ids) + embeddings = self.embedding_table(labels) + return embeddings + + +################################################################################# +# Core Latte Model # +################################################################################# + +class TransformerBlock(nn.Module): + """ + A Latte tansformer block with adaptive layer norm zero (adaLN-Zero) conditioning. + """ + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + mlp_hidden_dim = int(hidden_size * mlp_ratio) + approx_gelu = lambda: nn.GELU(approximate="tanh") + self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True) + ) + + def forward(self, x, c, attn_bias): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) + x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), attn_bias) + x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) + return x + + +class FinalLayer(nn.Module): + """ + The final layer of Latte. + """ + def __init__(self, hidden_size, patch_size_t, patch_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size_t * patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class Latte(nn.Module): + """ + Diffusion model with a Transformer backbone. + """ + def __init__(self, config: LatteConfiguration): + super().__init__() + + input_size = config.input_size + patch_size = config.patch_size + patch_size_t = config.patch_size_t + in_channels = config.in_channels + hidden_size = config.hidden_size + depth = config.depth + num_heads = config.num_heads + mlp_ratio = config.mlp_ratio + num_frames = config.num_frames + class_dropout_prob = config.class_dropout_prob + num_classes = config.num_classes + learn_sigma = config.learn_sigma + extras = config.learn_sigma + attention_mode = config.attention_mode + attention_pe_mode = config.attention_pe_mode + pt_input_size = config.pt_input_size + intp_vfreq = config.intp_vfreq + + self.config = config + + self.learn_sigma = learn_sigma + self.in_channels = in_channels + self.out_channels = in_channels * 2 if learn_sigma else in_channels + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.num_heads = num_heads + self.extras = extras + self.num_frames = num_frames + self.hidden_size = hidden_size + self.gradient_checkpointing = False + + self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) + self.t_embedder = TimestepEmbedder(hidden_size) + + if self.extras == 2: + self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob) + if self.extras == 78: # timestep + text_embedding + self.text_embedding_projection = nn.Sequential( + nn.SiLU(), + nn.Linear(77 * 768, hidden_size, bias=True) + ) + + num_patches = self.x_embedder.num_patches + # Will use fixed sin-cos embedding: + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False) + self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False) + self.hidden_size = hidden_size + + if pt_input_size is None: + pt_input_size = input_size + self.blocks = nn.ModuleList([ + TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode, + attention_pe_mode=attention_pe_mode, hw=input_size, pt_hw=pt_input_size, + intp_vfreq=intp_vfreq + ) for _ in range(depth) + ]) + + self.final_layer = FinalLayer(hidden_size, patch_size_t, patch_size, self.out_channels) + self.initialize_weights() + + def initialize_weights(self): + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize (and freeze) pos_embed by sin-cos embedding: + pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5)) + self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) + + temp_embed = get_1d_sincos_temp_embed(self.temp_embed.shape[-1], self.temp_embed.shape[-2]) + self.temp_embed.data.copy_(torch.from_numpy(temp_embed).float().unsqueeze(0)) + + # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): + w = self.x_embedder.proj.weight.data + nn.init.xavier_uniform_(w.view([w.shape[0], -1])) + nn.init.constant_(self.x_embedder.proj.bias, 0) + + if self.extras == 2: + # Initialize label embedding table: + 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) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in Latte blocks: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.final_layer.linear.weight, 0) + nn.init.constant_(self.final_layer.linear.bias, 0) + + def unpatchify(self, x): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.out_channels + p = self.x_embedder.patch_size[0] + h = w = int(x.shape[1] ** 0.5) + assert h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) + x = torch.einsum('nhwpqc->nchpwq', x) + imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p)) + return imgs + + def ckpt_wrapper(self, module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + return ckpt_forward + + def make_mask(self, attention_mask): + 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) + return attention_mask + + # @torch.cuda.amp.autocast() + # @torch.compile + def forward(self, + x, + t, + y=None, + text_embedding=None, + attention_mask=None): + """ + Forward pass of Latte. + x: (N, F, C, H, W) tensor of video inputs + t: (N,) tensor of diffusion timesteps + y: (N,) tensor of class labels + attention_mask: (N, F, H, W) + """ + attention_mask_temproal, attention_mask_spatial = None, None + if attention_mask is not None: + attention_mask_spatial = rearrange(attention_mask, 'b t h w -> (b t) h w') + attention_mask_spatial = self.make_mask(attention_mask_spatial) + + attention_mask_temproal = rearrange(attention_mask, 'b t h w -> (b h w) t') + attention_mask_temproal = self.make_mask(attention_mask_temproal) + + batches, frames, channels, high, weight = x.shape + + x = rearrange(x, 'b f c h w -> (b f) c h w') + + x = self.x_embedder(x) + self.pos_embed + t = self.t_embedder(t) + + + timestep_spatial = repeat(t, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + timestep_temp = repeat(t, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + + if self.extras == 2: + y = self.y_embedder(y, self.training) + y_spatial = repeat(y, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + y_temp = repeat(y, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + elif self.extras == 78: + text_embedding = self.text_embedding_projection(text_embedding.reshape(batches, -1)) + text_embedding_spatial = repeat(text_embedding, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + text_embedding_temp = repeat(text_embedding, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + + for i in range(0, len(self.blocks), 2): + spatial_block, temp_block = self.blocks[i:i+2] + if self.extras == 2: + c = timestep_spatial + y_spatial + elif self.extras == 78: + c = timestep_spatial + text_embedding_spatial + else: + c = timestep_spatial + if self.gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(spatial_block), x, c, attention_mask_spatial) + else: + x = spatial_block(x, c, attention_mask_spatial) + + x = rearrange(x, '(b f) t d -> (b t) f d', b=batches) + # Add Time Embedding + if i == 0: + x = x + self.temp_embed + + if self.extras == 2: + c = timestep_temp + y_temp + elif self.extras == 78: + c = timestep_temp + text_embedding_temp + else: + c = timestep_temp + + if self.gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(temp_block), x, c, attention_mask_temproal) + else: + x = temp_block(x, c, attention_mask_temproal) + x = rearrange(x, '(b t) f d -> (b f) t d', b=batches) + + if self.extras == 2: + c = timestep_spatial + y_spatial + else: + c = timestep_spatial + x = self.final_layer(x, c) + x = self.unpatchify(x) + x = rearrange(x, '(b f) c h w -> b f c h w', b=batches) + return x + + def forward_with_cfg(self, x, t, y=None, cfg_scale=7.0, text_embedding=None, attention_mask=None): + """ + Forward pass of Latte, 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=y, text_embedding=text_embedding, 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:] + # eps, rest = model_out[:, :, :4, ...], model_out[:, :, 4:, ...] + 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=2) + + +################################################################################# +# Sine/Cosine Positional Embedding Functions # +################################################################################# +# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + +def get_1d_sincos_temp_embed(embed_dim, length): + pos = torch.arange(0, length).unsqueeze(1) + return get_1d_sincos_pos_embed_from_grid(embed_dim, pos) + +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): + """ + grid_size: int of the grid height and width + return: + pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_h = np.arange(grid_size, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + assert embed_dim % 2 == 0 + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) + + emb = np.concatenate([emb_h, emb_w], axis=1) + return emb + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2. + omega = 1. / 10000**omega + + pos = pos.reshape(-1) + out = np.einsum('m,d->md', pos, omega) + + emb_sin = np.sin(out) + emb_cos = np.cos(out) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) + return emb + + +################################################################################# +# Latte Configs # +################################################################################# + +from .configuration_latte import ( + Latte_XL_122_Config, Latte_XL_144_Config, Latte_XL_188_Config, + Latte_L_122_Config, Latte_L_144_Config, Latte_L_188_Config, + Latte_B_122_Config, Latte_B_144_Config, Latte_B_188_Config, + Latte_S_122_Config, Latte_S_144_Config, Latte_S_188_Config, +) + +def Latte_XL_122(**kwargs): + return Latte(Latte_XL_122_Config(**kwargs)) + +def Latte_XL_144(**kwargs): + return Latte(Latte_XL_144_Config(**kwargs)) + +def Latte_XL_188(**kwargs): + return Latte(Latte_XL_188_Config(**kwargs)) + +def Latte_L_122(**kwargs): + return Latte(Latte_L_122_Config(**kwargs)) + +def Latte_L_144(**kwargs): + return Latte(Latte_L_144_Config(**kwargs)) + +def Latte_L_188(**kwargs): + return Latte(Latte_L_188_Config(**kwargs)) + +def Latte_B_122(**kwargs): + return Latte(Latte_B_122_Config(**kwargs)) + +def Latte_B_144(**kwargs): + return Latte(Latte_B_144_Config(**kwargs)) + +def Latte_B_188(**kwargs): + return Latte(Latte_B_188_Config(**kwargs)) + +def Latte_S_122(**kwargs): + return Latte(Latte_S_122_Config(**kwargs)) + +def Latte_S_144(**kwargs): + return Latte(Latte_S_144_Config(**kwargs)) + +def Latte_S_188(**kwargs): + return Latte(Latte_S_188_Config(**kwargs)) + + +Latte_models = { + "Latte-XL/122": Latte_XL_122, "Latte-XL/144": Latte_XL_144, "Latte-XL/188": Latte_XL_188, + "Latte-L/122": Latte_L_122, "Latte-L/144": Latte_L_144, "Latte-L/188": Latte_L_188, + "Latte-B/122": Latte_B_122, "Latte-B/144": Latte_B_144, "Latte-B/188": Latte_B_188, + "Latte-S/122": Latte_S_122, "Latte-S/144": Latte_S_144, "Latte-S/188": Latte_S_188, +} + + +Latte_models = { + 'Latte-XL/122': Latte_XL_122, 'Latte-XL/144': Latte_XL_144, 'Latte-XL/188': Latte_XL_188, + 'Latte-L/122': Latte_L_122, 'Latte-L/144': Latte_L_144, 'Latte-L/188': Latte_L_188, + 'Latte-B/122': Latte_B_122, 'Latte-B/144': Latte_B_144, 'Latte-B/188': Latte_B_188, + 'Latte-S/122': Latte_S_122, 'Latte-S/144': Latte_S_144, 'Latte-S/188': Latte_S_188, +} + +if __name__ == '__main__': + + import torch + + device = "cuda" if torch.cuda.is_available() else "cpu" + + img = torch.randn(3, 16, 4, 32, 32).to(device) + t = torch.tensor([1, 2, 3]).to(device) + y = torch.tensor([1, 2, 3]).to(device) + network = Latte_XL_122().to(device) + from thop import profile + flops, params = profile(network, inputs=(img, t)) + print('FLOPs = ' + str(flops/1000**3) + 'G') + print('Params = ' + str(params/1000**2) + 'M') + # y_embeder = LabelEmbedder(num_classes=101, hidden_size=768, dropout_prob=0.5).to(device) + # lora.mark_only_lora_as_trainable(network) + # out = y_embeder(y, True) + # out = network(img, t, y) + # print(out.shape) From 81c455e878868a37b95e1d5fce5ac2f3b5c43dcc Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 11:25:56 +0800 Subject: [PATCH 12/28] Refactor latte. Former-commit-id: 7408fd0ee16247afc3d4ba43b8c8e7dfb1091a97 --- opensora/models/diffusion/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/opensora/models/diffusion/__init__.py b/opensora/models/diffusion/__init__.py index a97eff755..6492cc37f 100644 --- a/opensora/models/diffusion/__init__.py +++ b/opensora/models/diffusion/__init__.py @@ -1,5 +1,6 @@ from .dit.dit import DiT_models -from .latte.latte import Latte_models +from .latte.modeling_latte import Latte_models +# from .latte.latte import Latte_models from .latte_t2v.latte_t2v import LatteT2V Diffusion_models = {} @@ -8,4 +9,3 @@ Diffusion_models.update(Latte_models) def get_t2v_models(args): return LatteT2V.from_pretrained_2d(args.pretrained_model_path, subfolder="transformer", video_length=args.video_length) - \ No newline at end of file From f3274235d8d3dcb704fe6736aaa29bb37167b0f0 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 11:27:33 +0800 Subject: [PATCH 13/28] Support deepspeed zero2 and zero2_offload for latte training. Former-commit-id: c87e95dd357e6fa6d6fb2dcfee781bc49c039e13 --- .../deepspeed_zero2_config.yaml | 13 ++++++++++ .../deepspeed_zero2_offload_config.yaml | 13 ++++++++++ scripts/accelerate_configs/zero2.json | 23 ++++++++++++++++ scripts/accelerate_configs/zero2_offload.json | 26 +++++++++++++++++++ scripts/train_256_dsz2.sh | 20 ++++++++++++++ scripts/train_256_dsz2_off.sh | 20 ++++++++++++++ 6 files changed, 115 insertions(+) create mode 100644 scripts/accelerate_configs/deepspeed_zero2_config.yaml create mode 100644 scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml create mode 100644 scripts/accelerate_configs/zero2.json create mode 100644 scripts/accelerate_configs/zero2_offload.json create mode 100644 scripts/train_256_dsz2.sh create mode 100644 scripts/train_256_dsz2_off.sh diff --git a/scripts/accelerate_configs/deepspeed_zero2_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_config.yaml new file mode 100644 index 000000000..43ec1e80c --- /dev/null +++ b/scripts/accelerate_configs/deepspeed_zero2_config.yaml @@ -0,0 +1,13 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero2.json +fsdp_config: {} +machine_rank: 0 +main_process_ip: null +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +gpu_ids: 0,1,2,3,4,5,6,7 +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml new file mode 100644 index 000000000..9c96bae67 --- /dev/null +++ b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml @@ -0,0 +1,13 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero2_offload.json +fsdp_config: {} +machine_rank: 0 +main_process_ip: null +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +gpu_ids: 0,1,2,3,4,5,6,7 +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/zero2.json b/scripts/accelerate_configs/zero2.json new file mode 100644 index 000000000..c95ebefe0 --- /dev/null +++ b/scripts/accelerate_configs/zero2.json @@ -0,0 +1,23 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "train_micro_batch_size_per_gpu": "auto", + "train_batch_size": "auto", + "gradient_accumulation_steps": "auto", + "zero_optimization": { + "stage": 2, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto" + } +} \ No newline at end of file diff --git a/scripts/accelerate_configs/zero2_offload.json b/scripts/accelerate_configs/zero2_offload.json new file mode 100644 index 000000000..4ea31ddc8 --- /dev/null +++ b/scripts/accelerate_configs/zero2_offload.json @@ -0,0 +1,26 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "train_micro_batch_size_per_gpu": "auto", + "train_batch_size": "auto", + "gradient_accumulation_steps": "auto", + "zero_optimization": { + "stage": 2, + "offload_optimizer": { + "device": "cpu" + }, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto" + } +} \ No newline at end of file diff --git a/scripts/train_256_dsz2.sh b/scripts/train_256_dsz2.sh new file mode 100644 index 000000000..120bf7547 --- /dev/null +++ b/scripts/train_256_dsz2.sh @@ -0,0 +1,20 @@ + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml \ + opensora/train/train.py \ + --model Latte-XL/122 \ + --dataset sky \ + --ae stabilityai/sd-vae-ft-mse \ + --data-path datasets/sky_timelapse/sky_train \ + --extras 1 \ + --sample-rate 3 \ + --num-frames 16 \ + --max-image-size 256 \ + --max-train-steps 1000000 \ + --local-batch-size 5 \ + --lr 1e-4 \ + --ckpt-every 500 \ + --log-every 50 \ + --gradient-checkpointing \ + --attention_mode math \ + --mixed-precision fp16 diff --git a/scripts/train_256_dsz2_off.sh b/scripts/train_256_dsz2_off.sh new file mode 100644 index 000000000..120bf7547 --- /dev/null +++ b/scripts/train_256_dsz2_off.sh @@ -0,0 +1,20 @@ + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml \ + opensora/train/train.py \ + --model Latte-XL/122 \ + --dataset sky \ + --ae stabilityai/sd-vae-ft-mse \ + --data-path datasets/sky_timelapse/sky_train \ + --extras 1 \ + --sample-rate 3 \ + --num-frames 16 \ + --max-image-size 256 \ + --max-train-steps 1000000 \ + --local-batch-size 5 \ + --lr 1e-4 \ + --ckpt-every 500 \ + --log-every 50 \ + --gradient-checkpointing \ + --attention_mode math \ + --mixed-precision fp16 From 61f9c5bced749036f8c3ebec3fea2225d74d3f53 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 12:23:42 +0800 Subject: [PATCH 14/28] Add some comments. Former-commit-id: d5b5bd7743edde6a63dc74457bded0f8c56be1c2 --- opensora/models/diffusion/latte/modeling_latte.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 1b9018e0b..fa3ba86b6 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -141,6 +141,8 @@ class Attention(nn.Module): return x def make_attn_bias(self, attn_mask): + # The numerical range of bfloat16, float16 can't conver -1e8 + # Refer to https://discuss.pytorch.org/t/runtimeerror-value-cannot-be-converted-to-type-at-half-without-overflow-1e-30/109768 attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else 1e-4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias From 28b7fc4bae5f4771db1be33310e3e63d06f44824 Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 13:17:44 +0800 Subject: [PATCH 15/28] Fix typos. Former-commit-id: c10330d373466f3c1ab455bfb391569b9175bc3a --- opensora/models/diffusion/latte/latte.py | 2 +- opensora/models/diffusion/latte/modeling_latte.py | 9 +-------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index 132dbbea4..6b1a3a226 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -137,7 +137,7 @@ class Attention(nn.Module): return x def make_attn_bias(self, attn_mask): - attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else 1e-4, attn_mask) + attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else -1e4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index fa3ba86b6..23b651391 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -143,7 +143,7 @@ class Attention(nn.Module): def make_attn_bias(self, attn_mask): # The numerical range of bfloat16, float16 can't conver -1e8 # Refer to https://discuss.pytorch.org/t/runtimeerror-value-cannot-be-converted-to-type-at-half-without-overflow-1e-30/109768 - attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else 1e-4, attn_mask) + attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else -1e4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias @@ -621,13 +621,6 @@ Latte_models = { } -Latte_models = { - 'Latte-XL/122': Latte_XL_122, 'Latte-XL/144': Latte_XL_144, 'Latte-XL/188': Latte_XL_188, - 'Latte-L/122': Latte_L_122, 'Latte-L/144': Latte_L_144, 'Latte-L/188': Latte_L_188, - 'Latte-B/122': Latte_B_122, 'Latte-B/144': Latte_B_144, 'Latte-B/188': Latte_B_188, - 'Latte-S/122': Latte_S_122, 'Latte-S/144': Latte_S_144, 'Latte-S/188': Latte_S_188, -} - if __name__ == '__main__': import torch From 6dae0b1c0c71d935ddcd985149fcfd135b6d543b Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sun, 10 Mar 2024 15:08:12 +0800 Subject: [PATCH 16/28] deepspeed Former-commit-id: 6179a4b0e7d06caed4201a27679abc64615febe6 --- opensora/models/diffusion/__init__.py | 3 ++- opensora/models/diffusion/dit/dit.py | 6 +++-- opensora/models/diffusion/latte/latte.py | 2 ++ .../models/diffusion/latte/modeling_latte.py | 22 +++++++++++-------- opensora/train/train.py | 7 +++--- scripts/sky/train.sh | 6 ++--- scripts/train_256.sh | 4 ++-- scripts/train_256_dsz2.sh | 4 ++-- scripts/train_256_dsz2_off.sh | 4 ++-- 9 files changed, 33 insertions(+), 25 deletions(-) diff --git a/opensora/models/diffusion/__init__.py b/opensora/models/diffusion/__init__.py index 6492cc37f..2adaeffb5 100644 --- a/opensora/models/diffusion/__init__.py +++ b/opensora/models/diffusion/__init__.py @@ -1,7 +1,7 @@ from .dit.dit import DiT_models from .latte.modeling_latte import Latte_models # from .latte.latte import Latte_models -from .latte_t2v.latte_t2v import LatteT2V +from .latte_t2v.latte_t2v import LatteT2V Diffusion_models = {} Diffusion_models.update(DiT_models) @@ -9,3 +9,4 @@ Diffusion_models.update(Latte_models) def get_t2v_models(args): return LatteT2V.from_pretrained_2d(args.pretrained_model_path, subfolder="transformer", video_length=args.video_length) + \ No newline at end of file diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index 3bf6e8942..7ea2732ec 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -133,7 +133,9 @@ class Attention(nn.Module): return x def make_attn_bias(self, attn_mask): - attn_bias = torch.where(attn_mask == 0, -1e8, attn_mask) + # The numerical range of bfloat16, float16 can't conver -1e8 + # Refer to https://discuss.pytorch.org/t/runtimeerror-value-cannot-be-converted-to-type-at-half-without-overflow-1e-30/109768 + attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else -1e4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias @@ -554,4 +556,4 @@ DiT_models = { 'DiT-L/122': DiT_L_122, 'DiT-L/144': DiT_L_144, 'DiT-L/188': DiT_L_188, 'DiT-B/122': DiT_B_122, 'DiT-B/144': DiT_B_144, 'DiT-B/188': DiT_B_188, 'DiT-S/122': DiT_S_122, 'DiT-S/144': DiT_S_144, 'DiT-S/188': DiT_S_188, -} +} \ No newline at end of file diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index f1cbc860b..8837fc324 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -143,6 +143,8 @@ class Attention(nn.Module): return x def make_attn_bias(self, attn_mask): + # The numerical range of bfloat16, float16 can't conver -1e8 + # Refer to https://discuss.pytorch.org/t/runtimeerror-value-cannot-be-converted-to-type-at-half-without-overflow-1e-30/109768 attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else -1e4, attn_mask) attn_bias = torch.where(attn_mask == 1, 0., attn_bias) return attn_bias diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 23b651391..beefda2f8 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -103,18 +103,22 @@ class Attention(nn.Module): ro_k_t = self.rope(k_t) k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads) - if self.attention_mode == 'xformers': # cause loss nan while using with amp + if self.attention_mode == 'xformers': # require pytorch 2.0 with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 - - elif self.attention_mode == 'flash': - # cause loss nan while using with amp - # Optionally use the context manager to ensure one of the fused kerenels is run - with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): - x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, - dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + elif self.attention_mode == 'flash': # require pytorch 2.0 + # https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/109 + if attn_mask is None or torch.all(attn_mask.bool()): + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = F.scaled_dot_product_attention(q, k, v, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + else: # turn to xformers + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True): + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, + dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) + elif self.attention_mode == 'math': attn = (q @ k.transpose(-2, -1)) * self.scale if attn_mask is not None: diff --git a/opensora/train/train.py b/opensora/train/train.py index 9fc88c604..f3bc24e83 100644 --- a/opensora/train/train.py +++ b/opensora/train/train.py @@ -204,7 +204,7 @@ def main(args): if accelerator.is_main_process: logger.info(f"Training for {num_train_epochs} epochs...") - + if args.mixed_precision == "bf16": dtype = torch.bfloat16 elif args.mixed_precision == "fp16": @@ -224,13 +224,13 @@ def main(args): # Map input images to latent space + normalize latents x = ae.encode(x) # B C T H W -> B T C H W + if args.extras == 78: # text-to-video raise 'T2V training are Not supported at this moment!' elif args.extras == 2: model_kwargs = dict(y=y, attention_mask=attn_mask) else: model_kwargs = dict(y=None, attention_mask=attn_mask) - with torch.autocast(device_type='cuda', dtype=dtype): t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) loss_dict = diffusion.training_losses(model, x, t, model_kwargs) @@ -327,12 +327,11 @@ if __name__ == "__main__": parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--lr-warmup-steps", type=int, default=0) parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math") - parser.add_argument("--gradient-accumulation-steps", type=int, default=1) parser.add_argument("--mixed-precision", type=str, default=None, choices=[None, "fp16", "bf16"]) parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)") # -------------------------------------- - args = parser.parse_args() + args = parser.parse_args() main(args) diff --git a/scripts/sky/train.sh b/scripts/sky/train.sh index 738f6fb91..0363ce3e0 100644 --- a/scripts/sky/train.sh +++ b/scripts/sky/train.sh @@ -1,9 +1,9 @@ -accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train.py \ +CUDA_VISIBLE_DEVICES=5,6 accelerate launch --multi_gpu --num_processes 2 --main_process_port 29503 --mixed_precision bf16 opensora/train/train.py \ --model Latte-XL/122 --dataset sky \ --ae stabilityai/sd-vae-ft-mse \ --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ - --sample-rate 1 --num-frames 128 --max-image-size 256 \ - --max-train-steps 1000000 --local-batch-size 2 --lr 1e-4 --num-workers 10 \ + --sample-rate 1 --num-frames 16 --max-image-size 256 \ + --max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 --num-workers 10 \ --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file diff --git a/scripts/train_256.sh b/scripts/train_256.sh index 9d7912cf5..7016e632a 100644 --- a/scripts/train_256.sh +++ b/scripts/train_256.sh @@ -16,5 +16,5 @@ accelerate launch \ --ckpt-every 500 \ --log-every 50 \ --gradient-checkpointing \ - --attention_mode math \ - --mixed-precision fp16 + --attention_mode flash \ + --mixed-precision bf16 diff --git a/scripts/train_256_dsz2.sh b/scripts/train_256_dsz2.sh index 120bf7547..d64c1dc08 100644 --- a/scripts/train_256_dsz2.sh +++ b/scripts/train_256_dsz2.sh @@ -16,5 +16,5 @@ accelerate launch \ --ckpt-every 500 \ --log-every 50 \ --gradient-checkpointing \ - --attention_mode math \ - --mixed-precision fp16 + --attention_mode flash \ + --mixed-precision bf16 diff --git a/scripts/train_256_dsz2_off.sh b/scripts/train_256_dsz2_off.sh index 120bf7547..d64c1dc08 100644 --- a/scripts/train_256_dsz2_off.sh +++ b/scripts/train_256_dsz2_off.sh @@ -16,5 +16,5 @@ accelerate launch \ --ckpt-every 500 \ --log-every 50 \ --gradient-checkpointing \ - --attention_mode math \ - --mixed-precision fp16 + --attention_mode flash \ + --mixed-precision bf16 From 3d3ea0c8183615abe4bd7be102c6dccda7f8e49f Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 15:17:18 +0800 Subject: [PATCH 17/28] Update README.md Former-commit-id: 582070d4f97c4a4398dcb9a6699f4ab6e980c812 --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 88bb5b06b..baff74a49 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,7 @@ Project stages: - [x] Support mixed precision training. - [x] Add gradient checkpoint. - [x] Support for ReBased and Ring attention. 🤝 Thanks to [@kabachuha](https://github.com/kabachuha) - - [ ] Train using the deepspeed engine. 🙏 **[Need your contribution]** + - [x] Train using the deepspeed engine. 🤝 Thanks to [@sennnnn](https://github.com/sennnnn） - [ ] Integrate with [Colossal-AI](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/59#issue-2170735221) for a cheaper, faster, and more efficient. 🙏 **[Need your contribution]** - [ ] Train with a text condition. Here we could conduct different experiments: - [ ] Train with T5 conditioning. 🚀 **[Require more computation]** From 995454b518d3393c1eb6618d829ffb3acc54967b Mon Sep 17 00:00:00 2001 From: sennnnn Date: Sun, 10 Mar 2024 15:24:40 +0800 Subject: [PATCH 18/28] Add base config. Former-commit-id: d37c551a0421d6b73dee26060b517c46d8217e4f --- .../models/diffusion/configuration_base.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 opensora/models/diffusion/configuration_base.py diff --git a/opensora/models/diffusion/configuration_base.py b/opensora/models/diffusion/configuration_base.py new file mode 100644 index 000000000..31b77aed4 --- /dev/null +++ b/opensora/models/diffusion/configuration_base.py @@ -0,0 +1,24 @@ +import json + +""" +We are currently working on the abstraction of the code structure, +so there is no content here for the moment. +It will be gradually optimized over time. +Contributions from the open-source community are welcome. +""" +class BaseConfiguration: + 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) From b57709108a0b91bc1e8e2f6d278ac6ffce86f7f9 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sun, 10 Mar 2024 18:08:23 +0800 Subject: [PATCH 19/28] train 1080p video Former-commit-id: 74625142ef892f9e7fa459b2f98f2f7ae5b36d6a --- .../{ => latte}/configuration_base.py | 10 +++-- .../diffusion/latte/configuration_latte.py | 2 +- .../models/diffusion/latte/modeling_latte.py | 1 + opensora/train/train.py | 19 +++++---- .../deepspeed_zero3_config.yaml | 13 +++++++ .../deepspeed_zero3_offload_config.yaml | 13 +++++++ scripts/accelerate_configs/zero3.json | 28 +++++++++++++ scripts/accelerate_configs/zero3_offload.json | 39 +++++++++++++++++++ scripts/sky/train.sh | 9 ----- scripts/{ => sky}/train_256.sh | 5 ++- scripts/{ => sky}/train_256_dsz2.sh | 4 +- scripts/{ => sky}/train_256_dsz2_off.sh | 2 +- scripts/sky/train_256_dsz3.sh | 20 ++++++++++ scripts/sky/train_256_dsz3_off.sh | 20 ++++++++++ scripts/sky/train_512.sh | 6 ++- scripts/sky/train_feature.sh | 9 ----- scripts/sky/train_feature_512.sh | 9 ----- 17 files changed, 163 insertions(+), 46 deletions(-) rename opensora/models/diffusion/{ => latte}/configuration_base.py (92%) create mode 100644 scripts/accelerate_configs/deepspeed_zero3_config.yaml create mode 100644 scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml create mode 100644 scripts/accelerate_configs/zero3.json create mode 100644 scripts/accelerate_configs/zero3_offload.json delete mode 100644 scripts/sky/train.sh rename scripts/{ => sky}/train_256.sh (83%) rename scripts/{ => sky}/train_256_dsz2.sh (75%) rename scripts/{ => sky}/train_256_dsz2_off.sh (89%) create mode 100644 scripts/sky/train_256_dsz3.sh create mode 100644 scripts/sky/train_256_dsz3_off.sh delete mode 100644 scripts/sky/train_feature.sh delete mode 100644 scripts/sky/train_feature_512.sh diff --git a/opensora/models/diffusion/configuration_base.py b/opensora/models/diffusion/latte/configuration_base.py similarity index 92% rename from opensora/models/diffusion/configuration_base.py rename to opensora/models/diffusion/latte/configuration_base.py index 31b77aed4..0d76ba50b 100644 --- a/opensora/models/diffusion/configuration_base.py +++ b/opensora/models/diffusion/latte/configuration_base.py @@ -6,19 +6,21 @@ so there is no content here for the moment. It will be gradually optimized over time. Contributions from the open-source community are welcome. """ + + class BaseConfiguration: 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) + return cls(**config_dict) \ No newline at end of file diff --git a/opensora/models/diffusion/latte/configuration_latte.py b/opensora/models/diffusion/latte/configuration_latte.py index 0aaaa9d83..bd0ea10fe 100644 --- a/opensora/models/diffusion/latte/configuration_latte.py +++ b/opensora/models/diffusion/latte/configuration_latte.py @@ -1,6 +1,6 @@ from typing import Union, Tuple -from ..configuration_base import BaseConfiguration +from .configuration_base import BaseConfiguration class LatteConfiguration(BaseConfiguration): diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index beefda2f8..a5c01b525 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -488,6 +488,7 @@ class Latte(nn.Module): x = self.final_layer(x, c) x = self.unpatchify(x) x = rearrange(x, '(b f) c h w -> b f c h w', b=batches) + return x def forward_with_cfg(self, x, t, y=None, cfg_scale=7.0, text_embedding=None, attention_mask=None): diff --git a/opensora/train/train.py b/opensora/train/train.py index f3bc24e83..195e7c168 100644 --- a/opensora/train/train.py +++ b/opensora/train/train.py @@ -65,6 +65,7 @@ def main(args): tb_writer = create_tensorboard(experiment_dir) logger.info(f"Experiment directory created at {experiment_dir}") logger.info(f'{args}') + logger.info(f'{accelerator}') else: logger = create_logger(None) tb_writer = None @@ -124,8 +125,9 @@ def main(args): model = model.to(device) # 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) + if args.use_ema: + ema = deepcopy(model).to(device) # Create an EMA of the model for use after training + requires_grad(ema, False) diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule if args.use_compile: @@ -164,9 +166,10 @@ def main(args): ) # Prepare models for training: - update_ema(ema, model, decay=0) # Ensure EMA is initialized with synced weights + if args.use_ema: + update_ema(ema, model, decay=0) # Ensure EMA is initialized with synced weights + ema.eval() # EMA model should always be in eval mode model.train() # important! This enables embedding dropout for classifier-free guidance - ema.eval() # EMA model should always be in eval mode model, opt, loader, lr_scheduler = accelerator.prepare(model, opt, loader, lr_scheduler) # Variables for monitoring/logging purposes: @@ -224,7 +227,6 @@ def main(args): # Map input images to latent space + normalize latents x = ae.encode(x) # B C T H W -> B T C H W - if args.extras == 78: # text-to-video raise 'T2V training are Not supported at this moment!' elif args.extras == 2: @@ -246,7 +248,8 @@ def main(args): opt.step() - update_ema(ema, model.module) + if args.use_ema: + update_ema(ema, model.module) # Log loss values: running_loss += loss.item() @@ -278,10 +281,11 @@ def main(args): if accelerator.is_main_process: checkpoint = { "model": model.module.state_dict(), - "ema": ema.state_dict(), "opt": opt.state_dict(), "args": args } + if args.use_ema: + checkpoint.update({"ema": ema.state_dict()}) checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" torch.save(checkpoint, checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") @@ -324,6 +328,7 @@ if __name__ == "__main__": parser.add_argument("--resume-from-checkpoint", action="store_true") parser.add_argument("--gradient-checkpointing", action="store_true") parser.add_argument("--use-compile", action="store_true") + parser.add_argument("--use-ema", action="store_true") parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--lr-warmup-steps", type=int, default=0) parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math") diff --git a/scripts/accelerate_configs/deepspeed_zero3_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_config.yaml new file mode 100644 index 000000000..4a4a0935e --- /dev/null +++ b/scripts/accelerate_configs/deepspeed_zero3_config.yaml @@ -0,0 +1,13 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero3.json +fsdp_config: {} +machine_rank: 0 +main_process_ip: null +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +gpu_ids: 0,1,2,3,4,5,6,7 +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml new file mode 100644 index 000000000..b01631caf --- /dev/null +++ b/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml @@ -0,0 +1,13 @@ +compute_environment: LOCAL_MACHINE +distributed_type: DEEPSPEED +deepspeed_config: + deepspeed_config_file: scripts/accelerate_configs/zero3_offload.json +fsdp_config: {} +machine_rank: 0 +main_process_ip: null +main_process_port: 29501 +main_training_function: main +num_machines: 1 +num_processes: 8 +gpu_ids: 0,1,2,3,4,5,6,7 +use_cpu: false \ No newline at end of file diff --git a/scripts/accelerate_configs/zero3.json b/scripts/accelerate_configs/zero3.json new file mode 100644 index 000000000..6917317af --- /dev/null +++ b/scripts/accelerate_configs/zero3.json @@ -0,0 +1,28 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "train_micro_batch_size_per_gpu": "auto", + "train_batch_size": "auto", + "gradient_accumulation_steps": "auto", + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto", + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + } +} \ No newline at end of file diff --git a/scripts/accelerate_configs/zero3_offload.json b/scripts/accelerate_configs/zero3_offload.json new file mode 100644 index 000000000..f4f482a8e --- /dev/null +++ b/scripts/accelerate_configs/zero3_offload.json @@ -0,0 +1,39 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "zero_optimization": { + "stage": 3, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "offload_param": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto", + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "gather_16bit_weights_on_model_save": true + }, + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "steps_per_print": 1e5, + "wall_clock_breakdown": false +} \ No newline at end of file diff --git a/scripts/sky/train.sh b/scripts/sky/train.sh deleted file mode 100644 index 0363ce3e0..000000000 --- a/scripts/sky/train.sh +++ /dev/null @@ -1,9 +0,0 @@ - - -CUDA_VISIBLE_DEVICES=5,6 accelerate launch --multi_gpu --num_processes 2 --main_process_port 29503 --mixed_precision bf16 opensora/train/train.py \ - --model Latte-XL/122 --dataset sky \ - --ae stabilityai/sd-vae-ft-mse \ - --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ - --sample-rate 1 --num-frames 16 --max-image-size 256 \ - --max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 --num-workers 10 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file diff --git a/scripts/train_256.sh b/scripts/sky/train_256.sh similarity index 83% rename from scripts/train_256.sh rename to scripts/sky/train_256.sh index 7016e632a..7cda35802 100644 --- a/scripts/train_256.sh +++ b/scripts/sky/train_256.sh @@ -5,7 +5,7 @@ accelerate launch \ --model Latte-XL/122 \ --dataset sky \ --ae stabilityai/sd-vae-ft-mse \ - --data-path datasets/sky_timelapse/sky_train \ + --data-path /remote-home/yeyang/sky_timelapse/sky_train/ \ --extras 1 \ --sample-rate 3 \ --num-frames 16 \ @@ -17,4 +17,5 @@ accelerate launch \ --log-every 50 \ --gradient-checkpointing \ --attention_mode flash \ - --mixed-precision bf16 + --mixed-precision bf16 \ + diff --git a/scripts/train_256_dsz2.sh b/scripts/sky/train_256_dsz2.sh similarity index 75% rename from scripts/train_256_dsz2.sh rename to scripts/sky/train_256_dsz2.sh index d64c1dc08..006ae900c 100644 --- a/scripts/train_256_dsz2.sh +++ b/scripts/sky/train_256_dsz2.sh @@ -1,11 +1,11 @@ accelerate launch \ - --config_file scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml \ + --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \ opensora/train/train.py \ --model Latte-XL/122 \ --dataset sky \ --ae stabilityai/sd-vae-ft-mse \ - --data-path datasets/sky_timelapse/sky_train \ + --data-path /remote-home/yeyang/sky_timelapse/sky_train/ \ --extras 1 \ --sample-rate 3 \ --num-frames 16 \ diff --git a/scripts/train_256_dsz2_off.sh b/scripts/sky/train_256_dsz2_off.sh similarity index 89% rename from scripts/train_256_dsz2_off.sh rename to scripts/sky/train_256_dsz2_off.sh index d64c1dc08..0b57be566 100644 --- a/scripts/train_256_dsz2_off.sh +++ b/scripts/sky/train_256_dsz2_off.sh @@ -5,7 +5,7 @@ accelerate launch \ --model Latte-XL/122 \ --dataset sky \ --ae stabilityai/sd-vae-ft-mse \ - --data-path datasets/sky_timelapse/sky_train \ + --data-path /remote-home/yeyang/sky_timelapse/sky_train/ \ --extras 1 \ --sample-rate 3 \ --num-frames 16 \ diff --git a/scripts/sky/train_256_dsz3.sh b/scripts/sky/train_256_dsz3.sh new file mode 100644 index 000000000..adfe35ed5 --- /dev/null +++ b/scripts/sky/train_256_dsz3.sh @@ -0,0 +1,20 @@ + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero3_config.yaml \ + opensora/train/train.py \ + --model Latte-XL/122 \ + --dataset sky \ + --ae stabilityai/sd-vae-ft-mse \ + --data-path /remote-home/yeyang/sky_timelapse/sky_train/ \ + --extras 1 \ + --sample-rate 3 \ + --num-frames 16 \ + --max-image-size 256 \ + --max-train-steps 1000000 \ + --local-batch-size 5 \ + --lr 1e-4 \ + --ckpt-every 500 \ + --log-every 50 \ + --gradient-checkpointing \ + --attention_mode flash \ + --mixed-precision bf16 diff --git a/scripts/sky/train_256_dsz3_off.sh b/scripts/sky/train_256_dsz3_off.sh new file mode 100644 index 000000000..71451a55d --- /dev/null +++ b/scripts/sky/train_256_dsz3_off.sh @@ -0,0 +1,20 @@ + +accelerate launch \ + --config_file scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml \ + opensora/train/train.py \ + --model Latte-XL/122 \ + --dataset sky \ + --ae stabilityai/sd-vae-ft-mse \ + --data-path /remote-home/yeyang/sky_timelapse/sky_train/ \ + --extras 1 \ + --sample-rate 3 \ + --num-frames 16 \ + --max-image-size 256 \ + --max-train-steps 1000000 \ + --local-batch-size 5 \ + --lr 1e-4 \ + --ckpt-every 500 \ + --log-every 50 \ + --gradient-checkpointing \ + --attention_mode flash \ + --mixed-precision bf16 diff --git a/scripts/sky/train_512.sh b/scripts/sky/train_512.sh index 4450c8d13..ca3e4947b 100644 --- a/scripts/sky/train_512.sh +++ b/scripts/sky/train_512.sh @@ -1,8 +1,10 @@ -accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train.py \ +accelerate launch \ + --config_file scripts/accelerate_configs/ddp_config.yaml \ + opensora/train/train.py \ --model Latte-XL/122 --dataset sky \ --ae stabilityai/sd-vae-ft-mse \ --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ --sample-rate 3 --num-frames 16 --max-image-size 512 \ --max-train-steps 1000000 --local-batch-size 2 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash + --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash --mixed_precision bf16 diff --git a/scripts/sky/train_feature.sh b/scripts/sky/train_feature.sh deleted file mode 100644 index 1af1d67ec..000000000 --- a/scripts/sky/train_feature.sh +++ /dev/null @@ -1,9 +0,0 @@ - - -accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train_with_feature.py \ - --model Latte-XL/122 --dataset sky_feature \ - --ae stabilityai/sd-vae-ft-mse \ - --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ - --sample-rate 3 --num-frames 128 --max-image-size 256 \ - --max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file diff --git a/scripts/sky/train_feature_512.sh b/scripts/sky/train_feature_512.sh deleted file mode 100644 index 8c3dc09b7..000000000 --- a/scripts/sky/train_feature_512.sh +++ /dev/null @@ -1,9 +0,0 @@ - - -accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train_with_feature.py \ - --model Latte-XL/122 --dataset landscope_feature \ - --ae stabilityai/sd-vae-ft-mse \ - --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ - --sample-rate 3 --num-frames 64 --max-image-size 512 \ - --max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file From 74494966bd4015bf694f1f8bd3141d3c171471bb Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 18:15:21 +0800 Subject: [PATCH 20/28] Update README.md Former-commit-id: 87636c1dc2275c4da77c6368a0f4425bea15dc12 --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index baff74a49..c65e15bb3 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,8 @@ Project stages: ## 📰 News +**[2024.03.10]** Supports training a latent size of 225×90×90(t×h×w), which means we are able to **train 1 minute of 1080P video with 30FPS** (2× interpolated frames and 2× super resolution). + **[2024.03.08]** We support the training code of text condition with 16 frames of 512x512. The code is mainly borrowed from [Latte](https://github.com/Vchitect/Latte). **[2024.03.07]** We support training with 128 frames (when sample rate = 3, which is about 13 seconds) of 256x256, or 64 frames (which is about 6 seconds) of 512x512. From fb62b8ff7e10909993c69cd4f12ab1d333e8db93 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 18:20:28 +0800 Subject: [PATCH 21/28] Update README.md Former-commit-id: 059d9e863958cfd679dbb8d365de6cd933e284fc --- README.md | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index c65e15bb3..3bdca885b 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ Project stages: ## 📰 News -**[2024.03.10]** Supports training a latent size of 225×90×90(t×h×w), which means we are able to **train 1 minute of 1080P video with 30FPS** (2× interpolated frames and 2× super resolution). +**[2024.03.10]** 🚀🚀🚀 This repo supports training a latent size of 225×90×90 (t×h×w), which means we are able to **train 1 minute of 1080P video with 30FPS** (2× interpolated frames and 2× super resolution). **[2024.03.08]** We support the training code of text condition with 16 frames of 512x512. The code is mainly borrowed from [Latte](https://github.com/Vchitect/Latte). @@ -158,22 +158,23 @@ Project stages: ## 🛠️ Requirements and Installation -The requirements are as follows. - -* Python >= 3.8 -* CUDA Version >= 11.7 -* Install required packages - +1. Clone this repository and navigate to Open-Sora-Plan folder ``` git clone https://github.com/PKU-YuanGroup/Open-Sora-Plan cd Open-Sora-Plan +``` +2. Install required packages +``` conda create -n opensora python=3.8 -y conda activate opensora pip install -e . ``` - -* Install optional requirements such as static type checking: - +3. Install additional packages for training cases +``` +pip install -e ".[train]" +pip install flash-attn --no-build-isolation +``` +4. Install optional requirements such as static type checking: ``` pip install -e '.[dev]' ``` From aeb30813417b6d48e927c179cab1b50b5c8f4ca1 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 18:21:50 +0800 Subject: [PATCH 22/28] Update pyproject.toml Former-commit-id: 5c379b6e4515ad75b89cdebbc5cd490c9a821d78 --- pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6b7733f84..9e9ee3f0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,9 +24,8 @@ dependencies = [ ] [project.optional-dependencies] -dev = [ - "mypy==1.8.0" -] +train = ["deepspeed==0.9.5", "pydantic==1.10.13"] +dev = ["mypy==1.8.0"] [project.urls] From 740363460f4267aea4bef3b22c26f303bc619677 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 18:27:50 +0800 Subject: [PATCH 23/28] Update README.md Former-commit-id: 7a62dec060c91b02ca556451266e6f3ef28e03e7 --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 3bdca885b..6b415744d 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ Project stages: ## 📰 News -**[2024.03.10]** 🚀🚀🚀 This repo supports training a latent size of 225×90×90 (t×h×w), which means we are able to **train 1 minute of 1080P video with 30FPS** (2× interpolated frames and 2× super resolution). +**[2024.03.10]** 🚀🚀🚀 This repo supports training a latent size of 225×90×90 (t×h×w), which means we are able to **train 1 minute of 1080P video with 30FPS** (2× interpolated frames and 2× super resolution) under class-condition. **[2024.03.08]** We support the training code of text condition with 16 frames of 512x512. The code is mainly borrowed from [Latte](https://github.com/Vchitect/Latte). From a7165b506e8a6716bb4a5b2012ec09c6cf0d9f30 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sun, 10 Mar 2024 18:34:37 +0800 Subject: [PATCH 24/28] clean script Former-commit-id: 02c61f63a8539a31516b92f5bb7c2d2e428ca7d7 --- scripts/land_scope/train_feature.sh | 9 --------- scripts/land_scope/train_videofeature.sh | 9 --------- scripts/sky/train_512.sh | 10 ---------- 3 files changed, 28 deletions(-) delete mode 100644 scripts/land_scope/train_feature.sh delete mode 100644 scripts/land_scope/train_videofeature.sh delete mode 100644 scripts/sky/train_512.sh diff --git a/scripts/land_scope/train_feature.sh b/scripts/land_scope/train_feature.sh deleted file mode 100644 index 6433b6256..000000000 --- a/scripts/land_scope/train_feature.sh +++ /dev/null @@ -1,9 +0,0 @@ - - -accelerate launch --multi_gpu --num_processes 7 --mixed_precision bf16 opensora/train/train_with_feature.py \ - --model Latte-XL/122 --dataset landscope_feature \ - --ae stabilityai/sd-vae-ft-mse \ - --data-path "/remote-home/yeyang/train_movie/*/*256.npy" --extras 1 \ - --sample-rate 3 --num-frames 128 --max-image-size 256 \ - --max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing \ No newline at end of file diff --git a/scripts/land_scope/train_videofeature.sh b/scripts/land_scope/train_videofeature.sh deleted file mode 100644 index f2cb30d96..000000000 --- a/scripts/land_scope/train_videofeature.sh +++ /dev/null @@ -1,9 +0,0 @@ - - -accelerate launch --multi_gpu --num_processes 7 opensora/train/train_with_feature.py \ - --model Latte-XL/122 --dataset landscope_video_feature \ - --ae ucf101_stride4x4x4 \ - --data-path "/remote-home/yeyang/train_movie/*/*s256-videogpt-f16*.npy" --extras 1 \ - --sample-rate 3 --num-frames 128 --max-image-size 256 \ - --max-train-steps 1000000 --local-batch-size 16 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 \ No newline at end of file diff --git a/scripts/sky/train_512.sh b/scripts/sky/train_512.sh deleted file mode 100644 index ca3e4947b..000000000 --- a/scripts/sky/train_512.sh +++ /dev/null @@ -1,10 +0,0 @@ - -accelerate launch \ - --config_file scripts/accelerate_configs/ddp_config.yaml \ - opensora/train/train.py \ - --model Latte-XL/122 --dataset sky \ - --ae stabilityai/sd-vae-ft-mse \ - --data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \ - --sample-rate 3 --num-frames 16 --max-image-size 512 \ - --max-train-steps 1000000 --local-batch-size 2 --lr 1e-4 \ - --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash --mixed_precision bf16 From f423fc753846985a5c841728efef94b90a9d5f9b Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 20:25:48 +0800 Subject: [PATCH 25/28] Update README.md Former-commit-id: 6d41cfe906f97dbed1cea8c9ed06a1f6dc84308e --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 6b415744d..cd6559ea0 100644 --- a/README.md +++ b/README.md @@ -82,6 +82,7 @@ Project stages: #### Train models that boost resolution and duration - [x] Add [PI](https://arxiv.org/abs/2306.15595) to support out-of-domain size. 🤝 Thanks to [@jpthu17](https://github.com/jpthu17) - [x] Add 2D RoPE to improve generalization ability as [FiT](https://github.com/whlzy/FiT). 🤝 Thanks to [@jpthu17](https://github.com/jpthu17) +- [ ] Compress KV according to [PixArt-sigma](https://pixart-alpha.github.io/PixArt-sigma-project). ⌛ [WIP] - [ ] Train a **low dimension** Video-AE, whether it is VAE or VQVAE. ⌛ [WIP] 🚀 **[Require more computation]** - [x] Extract offline feature. - [x] Train with offline feature. From 93b624e62d0870100be830ed28fcca21d3e0dc9a Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 20:27:46 +0800 Subject: [PATCH 26/28] Update README.md Former-commit-id: 64d147863368aeef479ad25f2e5463fdcb76de46 --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index cd6559ea0..f9a048ff8 100644 --- a/README.md +++ b/README.md @@ -272,9 +272,9 @@ 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. - + From d4aa360c5099276f8d83b3d60cd4b37332f03484 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 21:16:15 +0800 Subject: [PATCH 27/28] Update README.md Former-commit-id: 2808edf11005e04aa730971a786d1a4e3bca88b5 --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index f9a048ff8..6d1a4f08a 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [[Project Page]](https://pku-yuangroup.github.io/Open-Sora-Plan/) [[中文主页]](https://pku-yuangroup.github.io/Open-Sora-Plan/blog_cn.html) [![slack badge](https://img.shields.io/badge/Discord-join-blueviolet?logo=discord&)](https://discord.gg/fqpmStRX) -[![WeChat badge](https://img.shields.io/badge/微信-加入-green?logo=wechat&)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/53#issuecomment-1980312563) +[![WeChat badge](https://img.shields.io/badge/微信-加入-green?logo=wechat&)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/53#issuecomment-1987226516) [![Twitter](https://img.shields.io/badge/-Twitter@LinBin46984-black?logo=twitter&logoColor=1D9BF0)](https://x.com/LinBin46984/status/1763476690385424554?s=20)
[![License](https://img.shields.io/badge/License-MIT-yellow)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/LICENSE) [![GitHub repo contributors](https://img.shields.io/github/contributors-anon/PKU-YuanGroup/Open-Sora-Plan?style=flat&label=Contributors)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/graphs/contributors) From 14e217a6985c51d6da611df40f8a568ba30c2196 Mon Sep 17 00:00:00 2001 From: lb203 <62638829+LinB203@users.noreply.github.com> Date: Sun, 10 Mar 2024 21:37:01 +0800 Subject: [PATCH 28/28] Update README.md Former-commit-id: 1f94f2721f35670aeae8dbfb7c0a5abd06cd2aee --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6d1a4f08a..4c134f825 100644 --- a/README.md +++ b/README.md @@ -116,7 +116,7 @@ Project stages: - [x] Support mixed precision training. - [x] Add gradient checkpoint. - [x] Support for ReBased and Ring attention. 🤝 Thanks to [@kabachuha](https://github.com/kabachuha) - - [x] Train using the deepspeed engine. 🤝 Thanks to [@sennnnn](https://github.com/sennnnn） + - [x] Train using the deepspeed engine. 🤝 Thanks to [@sennnnn](https://github.com/sennnnn) - [ ] Integrate with [Colossal-AI](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/59#issue-2170735221) for a cheaper, faster, and more efficient. 🙏 **[Need your contribution]** - [ ] Train with a text condition. Here we could conduct different experiments: - [ ] Train with T5 conditioning. 🚀 **[Require more computation]**