diff --git a/README.md b/README.md
index 4c134f825..bc498407d 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,7 @@ Project stages:
- [x] Setup environment. ๐ค Thanks to [@nameless1117](https://github.com/nameless1117)
- [ ] Add docker file. โ [WIP] ๐ค Thanks to [@Mon-ius](https://github.com/Mon-ius), [@SimonLeeGit](https://github.com/SimonLeeGit)
- [ ] Enable type hints for functions. [@RuslanPeresy](https://github.com/RuslanPeresy), ๐ **[Need your contribution]**
+- [x] Resume from checkpoint.
- [x] Add Video-VQGAN model, which is borrowed from [VideoGPT](https://github.com/wilson1yan/VideoGPT).
- [x] Support variable aspect ratios, resolutions, durations training on [DiT](https://github.com/facebookresearch/DiT).
- [x] Support Dynamic mask input inspired by [FiT](https://github.com/whlzy/FiT).
@@ -76,6 +77,7 @@ Project stages:
- [x] Refactor VideoGPT. ๐ค Thanks to [@qqingzheng](https://github.com/qqingzheng), [@luo3300612](https://github.com/luo3300612), [@sennnnn](https://github.com/sennnnn)
- [x] Add sampling script.
- [ ] Add DDP sampling script. โ [WIP]
+- [x] Use accelerate on multi-node. ๐ค Thanks to [@sysuyy](https://github.com/sysuyy)
- [x] Incorporate [SiT](https://github.com/willisma/SiT). ๐ค Thanks to [@khan-yin](https://github.com/khan-yin)
- [x] Add evaluation scripts (FVD, CLIP score). ๐ค Thanks to [@rain305f](https://github.com/rain305f)
@@ -83,6 +85,7 @@ Project stages:
- [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]
+- [x] Support deepspeed for videogpt training. ๐ค Thanks to [@sennnnn](https://github.com/sennnnn)
- [ ] 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.
@@ -101,9 +104,9 @@ Project stages:
- [ ] Add prompt captioner.
- [ ] Collect training data.
- [ ] Need video-text pairs with poor caption. ๐ **[Need your contribution]**
- - [ ] Extract multi-frame descriptions by large image-language models. ๐ **[Need your contribution]**
+ - [ ] Extract multi-frame descriptions by large image-language models. ๐ค Thanks to [@HowardLi1984](https://github.com/HowardLi1984)
- [ ] Extract video description by large video-language models. ๐ **[Need your contribution]**
- - [ ] Integrate captions to get a dense caption by using a large language model, such as GPT-4. ๐ **[Need your contribution]**
+ - [ ] Integrate captions to get a dense caption by using a large language model, such as GPT-4. ๐ค Thanks to [@HowardLi1984](https://github.com/HowardLi1984)
- [ ] Train a captioner to refine captions. ๐ **[Require more computation]**
#### Train the 1080p model on video2text dataset
@@ -232,15 +235,17 @@ We present four reconstructed videos in this demonstration, arranged from left t
Please refer to the document [VQVAE](docs/VQVAE.md).
-### VideoDiT (DiT)
+### Video Diffusion Transformer
#### Training
```
sh scripts/train.sh
```
+The current resources are only enough for us to do primary experiments on the Sky dataset.
+
-
+
#### Sampling
@@ -248,6 +253,11 @@ sh scripts/train.sh
sh scripts/sample.sh
```
+Below is a visualization of the sampling results.
+
+| **12s 256x256** | **25s 256x256** |
+| --- | --- |
+| 
| 
|
## ๐ค How to Contribute to the Open-Sora Plan Community
@@ -263,8 +273,8 @@ For more details, please refer to the [Contribution Guidelines](docs/Contributio
## ๐ Acknowledgement
* [Latte](https://github.com/Vchitect/Latte): The **main codebase** we built upon and it is an wonderful video gererated model.
-* [DiT](https://github.com/facebookresearch/DiT): Scalable Diffusion Models with Transformers.
* [VideoGPT](https://github.com/wilson1yan/VideoGPT): Video Generation using VQ-VAE and Transformers.
+* [DiT](https://github.com/facebookresearch/DiT): Scalable Diffusion Models with Transformers.
* [FiT](https://github.com/whlzy/FiT): Flexible Vision Transformer for Diffusion Model.
* [Positional Interpolation](https://arxiv.org/abs/2306.15595): Extending Context Window of Large Language Models via Positional Interpolation.
diff --git a/examples/rec_video.py b/examples/rec_video.py
index 51eab1cbe..5acbeb6fb 100644
--- a/examples/rec_video.py
+++ b/examples/rec_video.py
@@ -1,19 +1,22 @@
-import sys
-sys.path.append(".")
-
import random
+import argparse
+from typing import Optional
+
import cv2
+import imageio
import numpy as np
import numpy.typing as npt
import torch
from decord import VideoReader, cpu
+from torch.nn import functional as F
from pytorchvideo.transforms import ShortSideScale
from torchvision.transforms import Lambda, Compose
from torchvision.transforms._transforms_video import RandomCropVideo
-from torch.nn import functional as F
-from typing import Optional
+
+import sys
+sys.path.append(".")
from opensora.models.ae import VQVAEModel
-import argparse
+
def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None:
height, width, channels = image_array[0].shape
@@ -32,7 +35,8 @@ def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'outpu
x = (x + 0.5)
x = x.permute(1, 2, 3, 0).numpy() # (C, T, H, W) -> (T, H, W, C)
x = (255*x).astype(np.uint8)
- array_to_video(x, fps=fps, output_file=output_file)
+ # array_to_video(x, fps=fps, output_file=output_file)
+ imageio.mimwrite(output_file, x, fps=fps, quality=9)
return
def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
@@ -101,6 +105,7 @@ def main(args: argparse.Namespace):
# custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4')
custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)
+
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--video-path', type=str, default='')
diff --git a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
index 450735aeb..90ac29cfa 100644
--- a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
+++ b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
@@ -29,3 +29,5 @@ class VQVAEConfiguration(VideoBaseConfiguration):
self.downsample = downsample
self.no_pos_embd = no_pos_embd
+
+ self.hidden_size = n_hiddens
diff --git a/opensora/models/captioner/caption_refiner/README.md b/opensora/models/captioner/caption_refiner/README.md
new file mode 100644
index 000000000..cf1ae266a
--- /dev/null
+++ b/opensora/models/captioner/caption_refiner/README.md
@@ -0,0 +1,38 @@
+# Refiner for Video Caption
+
+Transform the short caption annotations from video datasets into the long and detailed caption annotations.
+
+* Add detailed description for background scene.
+* Add detailed description for object attributes, including color, material, pose.
+* Add detailed description for object-level spatial relationship.
+
+## ๐ ๏ธ Extra Requirements and Installation
+
+* openai == 0.28.0
+* jsonlines == 4.0.0
+* nltk == 3.8.1
+* Install the LLaMA-Accessory:
+
+you also need to download the weight of SPHINX to ./ckpt/ folder
+
+## ๐๏ธ Refining
+
+The refining instruction is in [demo_for_refiner.py](demo_for_refiner.py).
+
+```bash
+python demo_for_refiner.py --root_path $path_to_repo$ --api_key $openai_api_key$
+```
+
+### Refining Demos
+
+```bash
+[original caption]: A red mustang parked in a showroom with american flags hanging from the ceiling.
+```
+
+```bash
+[refine caption]: This scene depicts a red Mustang parked in a showroom with American flags hanging from the ceiling. The showroom likely serves as a space for showcasing and purchasing cars, and the Mustang is displayed prominently near the flags and ceiling. The scene also features a large window and other objects. Overall, it seems to take place in a car show or dealership.
+```
+
+- [ ] Add GPT-3.5-Turbo for caption summarization. โ [WIP]
+- [ ] Add LLAVA-1.6. โ [WIP]
+- [ ] More descriptions. โ [WIP]
\ No newline at end of file
diff --git a/opensora/models/captioner/caption_refiner/caption_refiner.py b/opensora/models/captioner/caption_refiner/caption_refiner.py
new file mode 100644
index 000000000..23952f6d9
--- /dev/null
+++ b/opensora/models/captioner/caption_refiner/caption_refiner.py
@@ -0,0 +1,122 @@
+import itertools
+import numpy as np
+from PIL import Image
+from PIL import ImageSequence
+from nltk import pos_tag, word_tokenize
+
+from LLaMA2_Accessory.SPHINX import SPHINXModel
+from gpt_combinator import caption_summary
+
+class CaptionRefiner():
+ def __init__(self, sample_num, add_detect=True, add_pos=True, add_attr=True,
+ openai_api_key=None, openai_api_base=None,
+ ):
+ self.sample_num = sample_num
+ self.ADD_DETECTION_OBJ = add_detect
+ self.ADD_POS = add_pos
+ self.ADD_ATTR = add_attr
+ self.openai_api_key = openai_api_key
+ self.openai_api_base =openai_api_base
+
+ def video_load_split(self, video_path=None):
+ frame_img_list, sampled_img_list = [], []
+
+ if ".gif" in video_path:
+ img = Image.open(video_path)
+ # process every frame in GIF from to
+ for frame in ImageSequence.Iterator(img):
+ frame_np = np.array(frame.copy().convert('RGB').getdata(),dtype=np.uint8).reshape(frame.size[1],frame.size[0],3)
+ frame_img = Image.fromarray(np.uint8(frame_np))
+ frame_img_list.append(frame_img)
+ elif ".mp4" in video_path:
+ pass
+
+ # sample frames from the mp4/gif
+ for i in range(0, len(frame_img_list), int(len(frame_img_list)/self.sample_num)):
+ sampled_img_list.append(frame_img_list[i])
+
+ return sampled_img_list # [, ...]
+
+ def caption_refine(self, video_path, org_caption, model_path):
+ sampled_imgs = self.video_load_split(video_path)
+
+ model = SPHINXModel.from_pretrained(
+ pretrained_path=model_path,
+ with_visual=True
+ )
+
+ existing_objects, scene_description = [], []
+ text = word_tokenize(org_caption)
+ existing_objects = [word for word,tag in pos_tag(text) if tag in ["NN", "NNS", "NNP"]]
+ if self.ADD_DETECTION_OBJ:
+ # Detect the objects and scene in the sampled images
+
+ qas = [["Where is this scene in the picture most likely to take place?", None]]
+ sc_response = model.generate_response(qas, sampled_imgs[0], max_gen_len=1024, temperature=0.9, top_p=0.5, seed=0)
+ scene_description.append(sc_response)
+
+ # # Lacking accuracy
+ # for img in sampled_imgs:
+ # qas = [["Please detect the objects in the image.", None]]
+ # response = model.generate_response(qas, img, max_gen_len=1024, temperature=0.9, top_p=0.5, seed=0)
+ # print(response)
+
+ object_attrs = []
+ if self.ADD_ATTR:
+ # Detailed Description for all the objects in the sampled images
+ for obj in existing_objects:
+ obj_attr = []
+ for img in sampled_imgs:
+ qas = [["Please describe the attribute of the {}, including color, position, etc".format(obj), None]]
+ response = model.generate_response(qas, img, max_gen_len=1024, temperature=0.9, top_p=0.5, seed=0)
+ obj_attr.append(response)
+ object_attrs.append({obj : obj_attr})
+
+ space_relations = []
+ if self.ADD_POS:
+ obj_pairs = list(itertools.combinations(existing_objects, 2))
+ # Description for the relationship between each object in the sample images
+ for obj_pair in obj_pairs:
+ qas = [["What is the spatial relationship between the {} and the {}? Please describe in lease than twenty words".format(obj_pair[0], obj_pair[1]), None]]
+ response = model.generate_response(qas, img, max_gen_len=1024, temperature=0.9, top_p=0.5, seed=0)
+ space_relations.append(response)
+
+ return dict(
+ org_caption = org_caption,
+ scene_description = scene_description,
+ existing_objects = existing_objects,
+ object_attrs = object_attrs,
+ space_relations = space_relations,
+ )
+
+ def gpt_summary(self, total_captions):
+ # combine all captions into a detailed long caption
+ detailed_caption = ""
+
+ if "org_caption" in total_captions.keys():
+ detailed_caption += "In summary, "+ total_captions['org_caption']
+
+ if "scene_description" in total_captions.keys():
+ detailed_caption += "We first describe the whole scene. "+total_captions['scene_description'][-1]
+
+ if "existing_objects" in total_captions.keys():
+ tmp_sentence = "There are multiple objects in the video, including "
+ for obj in total_captions['existing_objects']:
+ tmp_sentence += obj+", "
+ detailed_caption += tmp_sentence
+
+ # if "object_attrs" in total_captions.keys():
+ # caption_summary(
+ # caption_list="",
+ # api_key=self.openai_api_key,
+ # api_base=self.openai_api_base,
+ # )
+
+ if "space_relations" in total_captions.keys():
+ tmp_sentence = "As for the spatial relationship. "
+ for sentence in total_captions['space_relations']: tmp_sentence += sentence
+ detailed_caption += tmp_sentence
+
+ detailed_caption = caption_summary(detailed_caption, self.open_api_key, self.open_api_base)
+
+ return detailed_caption
\ No newline at end of file
diff --git a/opensora/models/captioner/caption_refiner/dataset/test_videos/captions.json b/opensora/models/captioner/caption_refiner/dataset/test_videos/captions.json
new file mode 100644
index 000000000..098a352f2
--- /dev/null
+++ b/opensora/models/captioner/caption_refiner/dataset/test_videos/captions.json
@@ -0,0 +1 @@
+{"video1.gif": "A red mustang parked in a showroom with american flags hanging from the ceiling.", "video2.gif": "An aerial view of a city with a river running through it."}
\ No newline at end of file
diff --git a/opensora/models/captioner/caption_refiner/demo_for_refiner.py b/opensora/models/captioner/caption_refiner/demo_for_refiner.py
new file mode 100644
index 000000000..c7c0bfc5e
--- /dev/null
+++ b/opensora/models/captioner/caption_refiner/demo_for_refiner.py
@@ -0,0 +1,28 @@
+import argparse
+from caption_refiner import CaptionRefiner
+from gpt_combinator import caption_summary, caption_qa
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="question-answer-generation-using-gpt-3")
+ parser.add_argument("--root_path", required=True, help="The path to repo.")
+ parser.add_argument("--api_key", required=True, help="OpenAI API key.")
+ args = parser.parse_args()
+ return args
+
+if __name__ == "__main__":
+ args = parse_args()
+ myrefiner = CaptionRefiner(
+ sample_num=6, add_detect=True, add_pos=True, add_attr=True,
+ openai_api_key = args.api_key,
+ openai_api_base = "https://one-api.bltcy.top/v1",
+ )
+
+ results = myrefiner.caption_refine(
+ video_path="./dataset/test_videos/video1.gif",
+ org_caption="A red mustang parked in a showroom with american flags hanging from the ceiling.",
+ model_path = args.root_path + "/ckpts/SPHINX-Tiny",
+ )
+
+ final_caption = myrefiner.gpt_summary(results)
+
+ print(final_caption)
diff --git a/opensora/models/captioner/caption_refiner/gpt_combinator.py b/opensora/models/captioner/caption_refiner/gpt_combinator.py
new file mode 100644
index 000000000..c0a6f0dff
--- /dev/null
+++ b/opensora/models/captioner/caption_refiner/gpt_combinator.py
@@ -0,0 +1,93 @@
+import openai
+import ast
+
+def caption_qa(caption_list, api_key, api_base):
+ openai.api_key = api_key
+ openai.api_base = api_base
+
+ question = "What is the color of a red apple"
+ answer = "red"
+ pred = "green"
+ try:
+ # Compute the correctness score
+ completion = openai.ChatCompletion.create(
+ model="gpt-3.5-turbo",
+ # model="gpt-4",
+ # model="gpt-4-vision-compatible",
+ messages=[
+ {
+ "role": "system",
+ "content":
+ "You are an intelligent chatbot designed for evaluating the correctness of generative outputs for question-answer pairs. "
+ "Your task is to compare the predicted answer with the correct answer and determine if they match meaningfully. Here's how you can accomplish the task:"
+ "------"
+ "##INSTRUCTIONS: "
+ "- Focus on the meaningful match between the predicted answer and the correct answer.\n"
+ "- Consider synonyms or paraphrases as valid matches.\n"
+ "- Evaluate the correctness of the prediction compared to the answer."
+ },
+ {
+ "role": "user",
+ "content":
+ "Please evaluate the following video-based question-answer pair:\n\n"
+ f"Question: {question}\n"
+ f"Correct Answer: {answer}\n"
+ f"Predicted Answer: {pred}\n\n"
+ "Provide your evaluation only as a yes/no and score where the score is an integer value between 0 and 5, with 5 indicating the highest meaningful match. "
+ "Please generate the response in the form of a Python dictionary string with keys 'pred' and 'score', where value of 'pred' is a string of 'yes' or 'no' and value of 'score' is in INTEGER, not STRING."
+ "DO NOT PROVIDE ANY OTHER OUTPUT TEXT OR EXPLANATION. Only provide the Python dictionary string. "
+ "For example, your response should look like this: {'pred': 'yes', 'score': 4.8}."
+ }
+ ]
+ )
+ # Convert response to a Python dictionary.
+ response_message = completion["choices"][0]["message"]["content"]
+ response_dict = ast.literal_eval(response_message)
+ print(response_dict)
+
+ except Exception as e:
+ print(f"Error processing file : {e}")
+
+
+def caption_summary(long_caption, api_key, api_base):
+ """
+ apply GPT3-Turbo as the combination for original caption and the prompted captions for a video
+ """
+ openai.api_key = api_key
+ openai.api_base = api_base
+
+ try:
+ # Compute the correctness score
+ completion = openai.ChatCompletion.create(
+ model="gpt-3.5-turbo",
+ messages=[
+ {
+ "role": "system",
+ "content":
+ "You are an intelligent chatbot designed for summarizing from a long sentence. "
+ },
+ {
+ "role": "user",
+ "content":
+ "Please summarize the following sentences. Make it shorter than 70 words."
+ f"the long sentence: {long_caption}\n"
+ "Provide your summarization with less than 70 words. "
+ "DO NOT PROVIDE ANY OTHER TEXT OR EXPLANATION. Only provide the summary sentence. "
+ }
+ ]
+ )
+ # "Please generate the response in the form of a Python dictionary string with keys 'pred' and 'score', where value of 'pred' is a string of 'yes' or 'no' and value of 'score' is in INTEGER, not STRING."
+ # "DO NOT PROVIDE ANY OTHER OUTPUT TEXT OR EXPLANATION. Only provide the Python dictionary string. "
+ # "For example, your response should look like this: {'summary': 'your summary sentence'}."
+
+ # Convert response to a Python dictionary.
+ response_message = completion["choices"][0]["message"]["content"]
+ response_dict = ast.literal_eval(response_message)
+
+ except Exception as e:
+ print(f"Error processing file : {e}")
+
+ return response_dict
+
+if __name__ == "__main__":
+ caption_summary()
\ No newline at end of file
diff --git a/opensora/models/diffusion/__init__.py b/opensora/models/diffusion/__init__.py
index ee6f9de2e..387f33af6 100644
--- a/opensora/models/diffusion/__init__.py
+++ b/opensora/models/diffusion/__init__.py
@@ -6,7 +6,6 @@ from .latte.modeling_latte import Latte_models
from .latte_t2v.latte_t2v import LatteT2V
Diffusion_models = {}
-Diffusion_models.update(DiT_models)
Diffusion_models.update(Latte_models)
def get_t2v_models(args):
diff --git a/opensora/models/diffusion/dit/VisionRoPE.py b/opensora/models/diffusion/dit/VisionRoPE.py
deleted file mode 100644
index c6b01f572..000000000
--- a/opensora/models/diffusion/dit/VisionRoPE.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# 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
deleted file mode 100644
index 7ea2732ec..000000000
--- a/opensora/models/diffusion/dit/dit.py
+++ /dev/null
@@ -1,559 +0,0 @@
-# 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
-# --------------------------------------------------------
-from typing import Final, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import numpy as np
-import math
-import torch.utils.checkpoint as cp
-from einops import rearrange, repeat
-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
-from torch.nn import functional as F
-
-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
-
-
-class Attention(nn.Module):
- fused_attn: Final[bool]
-
- def __init__(
- self,
- dim: int,
- num_heads: int = 8,
- qkv_bias: bool = False,
- qk_norm: bool = False,
- attn_drop: float = 0.,
- proj_drop: float = 0.,
- norm_layer: nn.Module = nn.LayerNorm,
- attention_mode='math',
- 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
- ) -> None:
- super().__init__()
- assert dim % num_heads == 0, 'dim should be divisible by num_heads'
- self.num_heads = num_heads
- 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()
- self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
- self.attn_drop = nn.Dropout(attn_drop)
- 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, 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)
- 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': # 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)
-
- 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:
- 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):
- # 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
-
-def modulate(x, shift, scale):
- return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-#################################################################################
-# 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, C) 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 DiT Model #
-#################################################################################
-
-class DiTBlock(nn.Module):
- """
- A DiT 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, attention_mask):
- 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_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
- return x
-
-
-class FinalLayer(nn.Module):
- """
- The final layer of DiT.
- """
- 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 DiT(nn.Module):
- """
- Diffusion model with a Transformer backbone.
- """
- def __init__(
- self,
- input_size=32,
- patch_size=2,
- patch_size_t=1,
- in_channels=256,
- 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.gradient_checkpointing = False
-
- 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.hidden_size = hidden_size
- # import ipdb;ipdb.set_trace()
- 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)
- 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)
-
- 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,
- 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)
-
- # Initialize label embedding table:
- if self.extras == 2:
- nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
-
- # Initialize timestep embedding MLP:
- nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
- nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-
- # Zero-out adaLN modulation layers in DiT 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
-
- 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)
- t: (B,) tensor of diffusion timesteps
- y: (B,) tensor of class labels
- """
- if x.ndim == 4:
- raise NotImplementedError
-
- 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)
-
-
-
- # print(num_patches_height, num_patches_width, x.shape, pos_embed.shape)
- x = rearrange(x, 'b t c h w -> (b t) c h w')
- x = self.x_embedder(x) + self.pos_embed # (BT, N, C), where N = H * W / patch_size ** 2
- x = rearrange(x, '(b t) n c -> (b n) t c', t=self.t)
- x = x + self.temp_embed
- x = rearrange(x, '(b n) t c -> b (t n) c', b=B)
-
- t = self.t_embedder(t) # (B, C)
- if self.extras == 2:
- y = self.y_embedder(y, self.training) # (B, C)
- c = t + y # (B, C)
- else:
- c = t
- for block in self.blocks: # (B, N, C)
- if self.gradient_checkpointing and self.training:
- x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, c, attention_mask) # (B, N, C)
- else:
- x = block(x, c, attention_mask)
-
- x = rearrange(x, 'b (t n) c -> (b t) n c', t=self.t)
- c = repeat(c, 'b c -> (b t) c', t=self.t)
- x = self.final_layer(x, c) # (B, N, patch_size ** 2 * out_channels)
- x = self.unpatchify(x) # (B, out_channels, H, W)
- x = rearrange(x, '(b t) c h w -> b t c h w', t=self.t)
- return x
-
- def forward_with_cfg(self, x, t, y, cfg_scale, attention_mask):
- """
- Forward pass of DiT, but also batches the unconDiTional forward pass for classifier-free guidance.
- """
- # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
- half = x[: len(x) // 2]
- combined = torch.cat([half, half], dim=0)
- model_out = self.forward(combined, t, y, 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:]
- cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
- half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
- eps = torch.cat([half_eps, half_eps], dim=0)
- return torch.cat([eps, rest], dim=1)
-
-
-#################################################################################
-# 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
-
-
-
-#################################################################################
-# DiT Configs #
-#################################################################################
-
-def DiT_XL_122(**kwargs):
- return DiT(depth=28, hidden_size=1152, patch_size_t=1, patch_size=2, num_heads=16, **kwargs)
-
-def DiT_XL_144(**kwargs):
- return DiT(depth=28, hidden_size=1152, patch_size_t=1, patch_size=4, num_heads=16, **kwargs)
-
-def DiT_XL_188(**kwargs):
- return DiT(depth=28, hidden_size=1152, patch_size_t=1, patch_size=8, num_heads=16, **kwargs)
-
-def DiT_L_122(**kwargs):
- return DiT(depth=24, hidden_size=1024, patch_size_t=1, patch_size=2, num_heads=16, **kwargs)
-
-def DiT_L_144(**kwargs):
- return DiT(depth=24, hidden_size=1024, patch_size_t=1, patch_size=4, num_heads=16, **kwargs)
-
-def DiT_L_188(**kwargs):
- return DiT(depth=24, hidden_size=1024, patch_size_t=1, patch_size=8, num_heads=16, **kwargs)
-
-def DiT_B_122(**kwargs):
- return DiT(depth=12, hidden_size=768, patch_size_t=1, patch_size=2, num_heads=12, **kwargs)
-
-def DiT_B_144(**kwargs):
- return DiT(depth=12, hidden_size=768, patch_size_t=1, patch_size=4, num_heads=12, **kwargs)
-
-def DiT_B_188(**kwargs):
- return DiT(depth=12, hidden_size=768, patch_size_t=1, patch_size=8, num_heads=12, **kwargs)
-
-def DiT_S_122(**kwargs):
- return DiT(depth=12, hidden_size=384, patch_size_t=1, patch_size=2, num_heads=6, **kwargs)
-
-def DiT_S_144(**kwargs):
- return DiT(depth=12, hidden_size=384, patch_size_t=1, patch_size=4, num_heads=6, **kwargs)
-
-def DiT_S_188(**kwargs):
- return DiT(depth=12, hidden_size=384, patch_size_t=1, patch_size=8, num_heads=6, **kwargs)
-
-
-DiT_models = {
- 'DiT-XL/122': DiT_XL_122, 'DiT-XL/144': DiT_XL_144, 'DiT-XL/188': DiT_XL_188,
- '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/configuration_latte.py b/opensora/models/diffusion/latte/configuration_latte.py
index 1f61fcd65..f21ef6ecb 100644
--- a/opensora/models/diffusion/latte/configuration_latte.py
+++ b/opensora/models/diffusion/latte/configuration_latte.py
@@ -6,8 +6,8 @@ class LatteConfiguration:
def __init__(
self,
- input_size=32,
- patch_size=2,
+ input_size=(32, 32),
+ patch_size=(2, 2),
patch_size_t=1,
in_channels=4,
hidden_size=1152,
@@ -20,8 +20,10 @@ class LatteConfiguration:
learn_sigma=True,
extras=1,
attention_mode='math',
+ compress_kv=False,
attention_pe_mode=None,
pt_input_size: Union[int, Tuple[int, int]] = None, # (h, w)
+ pt_num_frames: Union[int, Tuple[int, int]] = None, # (num_frames, 1)
intp_vfreq: bool = True, # vision position interpolation
**kwargs,
):
@@ -41,8 +43,10 @@ class LatteConfiguration:
self.learn_sigma = learn_sigma
self.extras = learn_sigma
self.attention_mode = attention_mode
+ self.compress_kv = compress_kv
self.attention_pe_mode = attention_pe_mode
self.pt_input_size = pt_input_size
+ self.pt_num_frames = pt_num_frames
self.intp_vfreq = intp_vfreq
def to_json_string(self):
diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py
index a5c01b525..f5fcd805d 100644
--- a/opensora/models/diffusion/latte/modeling_latte.py
+++ b/opensora/models/diffusion/latte/modeling_latte.py
@@ -44,7 +44,7 @@ def modulate(x, shift, scale):
#################################################################################
# Attention Layers from TIMM #
#################################################################################
-
+#
class Attention(nn.Module):
def __init__(self,
dim,
@@ -61,6 +61,7 @@ class Attention(nn.Module):
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
+ compress_kv: bool = False
):
super().__init__()
assert dim % num_heads == 0, 'dim should be divisible by num_heads'
@@ -118,7 +119,7 @@ class Attention(nn.Module):
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:
@@ -151,6 +152,149 @@ class Attention(nn.Module):
attn_bias = torch.where(attn_mask == 1, 0., attn_bias)
return attn_bias
+
+# 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, 16), # (h, w)
+# pt_hw: Union[int, Tuple[int, int]] = (16, 16), # (h, w)
+# intp_vfreq: bool = True, # vision position interpolation
+# compress_kv: bool = False
+# ):
+# 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.q = nn.Linear(dim, dim, bias=qkv_bias)
+# self.kv = nn.Linear(dim, dim * 2, 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
+#
+# self.hw = hw
+# if self.attention_pe_mode == '2d_rope':
+# half_head_dim = dim // num_heads // 2
+# self.rope_q = VisionRotaryEmbeddingFast(
+# dim=half_head_dim,
+# pt_hw=(pt_hw[0], pt_hw[1]),
+# ft_hw=(self.hw[0], self.hw[1]) if intp_vfreq else None,
+# )
+# if self.compress_kv:
+# self.hw_compress = (math.ceil(self.hw[0] // 2), math.ceil(self.hw[1] // 2))
+# self.rope_k = VisionRotaryEmbeddingFast(
+# dim=half_head_dim,
+# pt_hw=(math.ceil(pt_hw[0] // 2), math.ceil(pt_hw[1] // 2)),
+# ft_hw=self.hw_compress if intp_vfreq else None,
+# )
+# else:
+# self.hw_compress = self.hw
+# self.rope_k = VisionRotaryEmbeddingFast(
+# dim=half_head_dim,
+# pt_hw=(pt_hw[0], pt_hw[1]),
+# ft_hw=(self.hw[0], self.hw[1]) if intp_vfreq else None,
+# )
+#
+#
+# self.compress_kv = compress_kv
+# if self.compress_kv:
+# self.sr = nn.Conv2d(dim, dim, kernel_size=(2, 1 if 1 in self.hw else 2), stride=(2, 1 if 1 in self.hw else 2))
+# self.norm = nn.LayerNorm(dim)
+#
+# 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
+# # import ipdb;ipdb.set_trace()
+# q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+#
+# if self.sr_ratio > 1:
+# x_ = x.permute(0, 2, 1).reshape(B, C, self.hw[0], self.hw[1])
+# x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
+# x_ = self.norm(x_)
+# kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+# else:
+# kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+# k, v = kv[0], kv[1]
+#
+#
+# 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(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_compress[0] * self.hw_compress[1], C // self.num_heads)
+# ro_k_t = self.rope_k(k_t)
+# k = ro_k_t.view(B, self.num_heads, N, C // self.num_heads)
+#
+# 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)
+#
+# 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:
+# 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):
+# # 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
+
#################################################################################
# Embedding Layers for Timesteps and Class Labels #
#################################################################################
@@ -294,8 +438,10 @@ class Latte(nn.Module):
learn_sigma = config.learn_sigma
extras = config.learn_sigma
attention_mode = config.attention_mode
+ compress_kv = config.compress_kv
attention_pe_mode = config.attention_pe_mode
pt_input_size = config.pt_input_size
+ pt_num_frames = config.pt_num_frames
intp_vfreq = config.intp_vfreq
self.config = config
@@ -309,6 +455,7 @@ class Latte(nn.Module):
self.extras = extras
self.num_frames = num_frames
self.hidden_size = hidden_size
+ self.compress_kv = compress_kv
self.gradient_checkpointing = False
self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
@@ -330,12 +477,28 @@ class Latte(nn.Module):
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)
- ])
+ if pt_num_frames is None:
+ pt_num_frames = num_frames
+ self.blocks = []
+ for i in range(depth):
+ if i % 2 == 0:
+ m = TransformerBlock(
+ hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
+ attention_pe_mode=attention_pe_mode,
+ hw=(input_size[0] // patch_size, input_size[1] // patch_size),
+ pt_hw=(pt_input_size[0] // patch_size, pt_input_size[1] // patch_size),
+ intp_vfreq=intp_vfreq, compress_kv=compress_kv
+ )
+ else:
+ m = TransformerBlock(
+ hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
+ attention_pe_mode=attention_pe_mode,
+ hw=(num_frames, 1),
+ pt_hw=(pt_num_frames, 1),
+ intp_vfreq=intp_vfreq, compress_kv=compress_kv
+ )
+ self.blocks.append(m)
+ self.blocks = nn.ModuleList(self.blocks)
self.final_layer = FinalLayer(hidden_size, patch_size_t, patch_size, self.out_channels)
self.initialize_weights()
diff --git a/opensora/models/frame_interpolation/interpolation.py b/opensora/models/frame_interpolation/interpolation.py
index e4ccd5ba7..0c60bed8a 100644
--- a/opensora/models/frame_interpolation/interpolation.py
+++ b/opensora/models/frame_interpolation/interpolation.py
@@ -21,6 +21,17 @@ from utils.build_utils import build_from_cfg
from utils.utils import InputPadder
+AMT_G = {
+ 'name': 'networks.AMT-G.Model',
+ 'params':{
+ 'corr_radius': 3,
+ 'corr_lvls': 4,
+ 'num_flows': 5,
+ }
+}
+
+
+
def init(device="cuda"):
'''
@@ -89,14 +100,13 @@ def get_input_video_from_path(input_path, device="cuda"):
return inputs, scale, padder
-def load_model(cfg_path, ckpt_path, device="cuda"):
+def load_model(ckpt_path, device="cuda"):
'''
load the frame interpolation model.
'''
-
- network_cfg = OmegaConf.load(cfg_path).network
- network_name = network_cfg.name
+ network_cfg = AMT_G
+ network_name = network_cfg['name']
print(f'Loading [{network_name}] from [{ckpt_path}]...')
model = build_from_cfg(network_cfg)
ckpt = torch.load(ckpt_path)
@@ -137,7 +147,7 @@ def interpolater(model, inputs, scale, padder, iters=1):
return outputs
-def write(outputs, output_path, frame_rate=30):
+def write(outputs, input_path, output_path, frame_rate=30):
'''
write results to the output_path.
'''
@@ -148,9 +158,10 @@ def write(outputs, output_path, frame_rate=30):
size = outputs[0].shape[2:][::-1]
- video_count = len(glob.glob(os.path.join(output_path, "*.mp4")))
+ _, file_name_with_extension = os.path.split(input_path)
+ file_name, _ = os.path.splitext(file_name_with_extension)
- save_video_path = f'{output_path}/demo_{video_count:04d}.mp4'
+ save_video_path = f'{output_path}/output_{file_name}.mp4'
writer = cv2.VideoWriter(save_video_path, cv2.VideoWriter_fourcc(*"mp4v"),
frame_rate, size)
@@ -165,17 +176,15 @@ def write(outputs, output_path, frame_rate=30):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
- parser.add_argument('-c', '--config', type=str, default='cfgs/AMT-G.yaml')
- parser.add_argument('-p', '--ckpt', type=str, default='amt-g.pth')
- parser.add_argument('-n', '--niters', type=int, default=6)
- parser.add_argument('-i', '--input', default="test.mp4")
- parser.add_argument('-o', '--output_path', type=str, default='results')
- parser.add_argument('-r', '--frame_rate', type=int, default=30)
+ parser.add_argument('--ckpt', type=str, default='amt-g.pth', help="The pretrained model.")
+ parser.add_argument('--niters', type=int, default=1, help="Iter of Interpolation. The number of frames will be double after per iter.")
+ parser.add_argument('--input', default="test.mp4", help="Input video.")
+ parser.add_argument('--output_path', type=str, default='results', help="Output path.")
+ parser.add_argument('--frame_rate', type=int, default=30, help="Frames rate of the output video.")
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- cfg_path = args.config
ckpt_path = args.ckpt
input_path = args.input
output_path = args.output_path
@@ -183,6 +192,6 @@ if __name__ == "__main__":
frame_rate = int(args.frame_rate)
inputs, scale, padder = get_input_video_from_path(input_path, device)
- model = load_model(cfg_path, ckpt_path, device)
+ model = load_model(ckpt_path, device)
outputs = interpolater(model, inputs, scale, padder, iters)
- write(outputs, output_path, frame_rate)
+ write(outputs, input_path, output_path, frame_rate)
diff --git a/opensora/models/frame_interpolation/readme.md b/opensora/models/frame_interpolation/readme.md
new file mode 100644
index 000000000..54cbfa2ee
--- /dev/null
+++ b/opensora/models/frame_interpolation/readme.md
@@ -0,0 +1,17 @@
+#### Frame Interpolation
+
+We use AMT as our frame interpolation model. (Thanks [AMT](https://github.com/MCG-NKU/AMT)) After sampling, you can use frame interpolation model to interpolate your video smoothly.
+
+1. Download the pretrained weights from [AMT](https://github.com/MCG-NKU/AMT), we recommend using the largest model AMT-G to achieve the best performance.
+2. Run the script of frame interpolation.
+```
+python opensora/models/frame_interpolation/interpolation.py --ckpt /path/to/ckpt --niters 1 --input /path/to/input/video.mp4 --output_path /path/to/output/floder --frame_rate 30
+```
+3. The output video will be stored at output_path and its duration time is equal `the total number of frames after frame interpolation / the frame rate`
+##### Frame Interpolation Specific Settings
+
+* `--ckpt`: Pretrained model of [AMT](https://github.com/MCG-NKU/AMT). We use AMT-G as our frame interpolation model.
+* `--niter`: Iterations of interpolation. With $m$ input frames, `[N_ITER]` $=n$ corresponds to $2^n\times (m-1)+1$ output frames.
+* `--input`: Path of the input video.
+* `--output_path`: Folder Path of the output video.
+* `--frame_rate"`: Frame rate of the output video.
diff --git a/opensora/sample/sample.py b/opensora/sample/sample.py
index 34d485334..9f127b085 100644
--- a/opensora/sample/sample.py
+++ b/opensora/sample/sample.py
@@ -10,6 +10,7 @@ import os
import sys
from accelerate import Accelerator
+from tqdm import tqdm
from opensora.dataset import ae_denorm
from opensora.models.ae import ae_channel_config, getae, ae_stride_config
@@ -23,8 +24,8 @@ import argparse
from einops import rearrange
import imageio
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
+# torch.backends.cuda.matmul.allow_tf32 = True
+# torch.backends.cudnn.allow_tf32 = True
def main(args):
# Setup PyTorch:
@@ -39,7 +40,7 @@ def main(args):
using_cfg = args.cfg_scale > 1.0
# Load model:
- latent_size = args.image_size // ae_stride_config[args.ae][1]
+ latent_size = (args.image_size // ae_stride_config[args.ae][1], args.image_size // ae_stride_config[args.ae][2])
args.latent_size = latent_size
model = Diffusion_models[args.model](
input_size=latent_size,
@@ -65,56 +66,45 @@ def main(args):
diffusion = create_diffusion(str(args.num_sampling_steps))
ae = getae(args).to(device)
- # if args.attention_mode == 'flash':
- # ae.to(dtype=torch.bfloat16)
- # model.to(dtype=torch.bfloat16)
+ bar = tqdm(range(args.num_sample))
+ for i in bar:
+ # Create sampling noise:
+ z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.module.in_channels, latent_size[0], latent_size[1], device=device)
- # Labels to condition the model with (feel free to change):
+ # Setup classifier-free guidance:
+ if using_cfg:
+ z = torch.cat([z, z], 0)
+ y = torch.randint(0, args.num_classes, (1,), device=device)
+ y_null = torch.tensor([args.num_classes] * 1, device=device)
+ y = torch.cat([y, y_null], dim=0)
+ model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
+ sample_fn = model.forward_with_cfg
+ else:
+ sample_fn = model.forward
+ model_kwargs = dict(y=None)
- # Create sampling noise:
- # if args.attention_mode == 'flash':
- # z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, dtype=torch.bfloat16, device=device)
- # else:
- z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.module.in_channels, latent_size, latent_size, device=device)
+ # Sample images:
+ if args.sample_method == 'ddim':
+ samples = diffusion.ddim_sample_loop(
+ sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+ )
+ elif args.sample_method == 'ddpm':
+ samples = diffusion.p_sample_loop(
+ sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+ )
- # Setup classifier-free guidance:
- if using_cfg:
- z = torch.cat([z, z], 0)
- y = torch.randint(0, args.num_classes, (1,), device=device)
- y_null = torch.tensor([args.num_classes] * 1, device=device)
- y = torch.cat([y, y_null], dim=0)
- model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
- sample_fn = model.forward_with_cfg
- else:
- sample_fn = model.forward
- model_kwargs = dict(y=None)
+ with torch.no_grad():
+ samples = ae.decode(samples)
+ # Save and display images:
- # Sample images:
- if args.sample_method == 'ddim':
- samples = diffusion.ddim_sample_loop(
- sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
- )
- elif args.sample_method == 'ddpm':
- samples = diffusion.p_sample_loop(
- sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
- )
+ if not os.path.exists(args.save_video_path):
+ os.makedirs(args.save_video_path)
- # if args.attention_mode == 'flash':
- # samples = samples.to(dtype=torch.bfloat16)
-
- with torch.no_grad():
- samples = ae.decode(samples)
- # Save and display images:
-
- if not os.path.exists(args.save_video_path):
- os.makedirs(args.save_video_path)
-
-
- video_ = (ae_denorm[args.ae](samples[0]) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
- video_save_path = os.path.join(args.save_video_path, 'sample' + '.mp4')
- print(video_save_path)
- imageio.mimwrite(video_save_path, video_, fps=args.fps, quality=9)
- print('save path {}'.format(args.save_video_path))
+ video_ = (ae_denorm[args.ae](samples[0]) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
+ video_save_path = os.path.join(args.save_video_path, f'sample_{i:03d}' + '.mp4')
+ print(video_save_path)
+ imageio.mimwrite(video_save_path, video_, fps=args.fps, quality=9)
+ print('save path {}'.format(args.save_video_path))
if __name__ == "__main__":
@@ -126,12 +116,13 @@ if __name__ == "__main__":
parser.add_argument("--fps", type=int, default=10)
parser.add_argument("--num-classes", type=int, default=101)
parser.add_argument("--num-frames", type=int, default=16)
- parser.add_argument("--image-size", type=int, default=256, choices=[256, 512])
+ parser.add_argument("--image-size", type=int, default=256)
parser.add_argument("--extras", type=int, default=1)
parser.add_argument("--num-sampling-steps", type=int, default=250)
+ parser.add_argument("--num-sample", type=int, default=1)
parser.add_argument("--cfg-scale", type=float, default=1.0)
parser.add_argument("--use-compile", action="store_true")
parser.add_argument("--sample-method", type=str, default='ddpm')
- parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
+ parser.add_argument("--attention-mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
args = parser.parse_args()
main(args)
diff --git a/opensora/train/train.py b/opensora/train/train.py
index 195e7c168..0ca36ea4f 100644
--- a/opensora/train/train.py
+++ b/opensora/train/train.py
@@ -52,14 +52,17 @@ def main(args):
device = accelerator.device
# Setup an experiment folder:
- if accelerator.is_main_process:
- os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders)
+ if args.resume_from_checkpoint:
+ experiment_dir = f"{args.results_dir}/{args.resume_from_checkpoint}"
+ else:
experiment_index = len(glob(f"{args.results_dir}/*"))
model_string_name = args.model.replace("/", "-") # e.g., Latte-XL/2 --> Latte-XL-2 (for naming folders)
num_frame_string = 'F' + str(args.num_frames) + 'S' + str(args.sample_rate)
experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}-{num_frame_string}-{args.dataset}" # Create an experiment folder
experiment_dir = get_experiment_dir(experiment_dir, args)
- checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
+ checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
+ if accelerator.is_main_process:
+ os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders)
os.makedirs(checkpoint_dir, exist_ok=True)
logger = create_logger(experiment_dir)
tb_writer = create_tensorboard(experiment_dir)
@@ -192,8 +195,11 @@ def main(args):
dirs = [d for d in dirs if d.endswith("pt")]
dirs = sorted(dirs, key=lambda x: int(x.split(".")[0]))
path = dirs[-1]
+ state_dict = torch.load(os.path.join(os.path.join(experiment_dir, 'checkpoints'), path), map_location='cpu')
+ print(state_dict['args'])
+ state_dict = {'module.'+k: v for k, v in state_dict['model'].items()}
logger.info(f"Resuming from checkpoint {path}")
- model.load_state(os.path.join(dirs, path))
+ model.load_state_dict(state_dict)
train_steps = int(path.split(".")[0])
first_epoch = train_steps // num_update_steps_per_epoch
@@ -325,13 +331,13 @@ if __name__ == "__main__":
parser.add_argument("--num-frames", type=int, default=16)
parser.add_argument("--max-image-size", type=int, default=128)
parser.add_argument("--dynamic-frames", action="store_true")
- parser.add_argument("--resume-from-checkpoint", action="store_true")
+ parser.add_argument("--resume-from-checkpoint", type=str, default=None)
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")
+ 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)")
diff --git a/opensora/train/train_videogpt.py b/opensora/train/train_videogpt.py
index a9c102b43..e323e67f6 100644
--- a/opensora/train/train_videogpt.py
+++ b/opensora/train/train_videogpt.py
@@ -38,8 +38,7 @@ def train(args, vqvae_args, training_args):
# Load Model
model = VQVAEModel(config)
# Load Dataset
- dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length)
-
+ dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length, resolution=config.resolution)
# Load Trainer
trainer = VQVAETrainer(model, training_args, train_dataset=dataset)
trainer.train()
diff --git a/scripts/accelerate_configs/hostfile b/scripts/accelerate_configs/hostfile
new file mode 100644
index 000000000..5e384ed0f
--- /dev/null
+++ b/scripts/accelerate_configs/hostfile
@@ -0,0 +1,2 @@
+server1 slots=8 # your server name and GPU in total
+server2 slots=8
diff --git a/scripts/accelerate_configs/multi_node_example.yaml b/scripts/accelerate_configs/multi_node_example.yaml
new file mode 100644
index 000000000..c462d62cc
--- /dev/null
+++ b/scripts/accelerate_configs/multi_node_example.yaml
@@ -0,0 +1,18 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero3_offload.json
+ deepspeed_hostfile: # FullPath here
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: # main process ip
+main_process_port: 29501
+main_training_function: main
+num_machines: 2
+num_processes: 16
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/scripts/sample.sh b/scripts/sample.sh
index c6ddd8064..bf0e7f49d 100644
--- a/scripts/sample.sh
+++ b/scripts/sample.sh
@@ -1,5 +1,6 @@
accelerate launch --num_processes 1 --main_process_port 29501 --mixed_precision bf16 opensora/sample/sample.py \
- --model Latte-XL/122 --ae stabilityai/sd-vae-ft-mse \
+ --model Latte-XL/122 \
+ --ae stabilityai/sd-vae-ft-mse \
--ckpt results/011-Latte-XL-122-F128S3-landscope_feature-Gc-Amp/checkpoints/0005000.pt --extras 1 \
- --fps 10 --num-frames 128 --image-size 256 --num-sampling-steps 250 --attention_mode flash
+ --fps 10 --num-frames 128 --image-size 256 --num-sampling-steps 250 --attention-mode flash
diff --git a/scripts/sky/train_256.sh b/scripts/sky/train_256.sh
index 7cda35802..1ae97a1b3 100644
--- a/scripts/sky/train_256.sh
+++ b/scripts/sky/train_256.sh
@@ -8,7 +8,7 @@ accelerate launch \
--data-path /remote-home/yeyang/sky_timelapse/sky_train/ \
--extras 1 \
--sample-rate 3 \
- --num-frames 16 \
+ --num-frames 128 \
--max-image-size 256 \
--max-train-steps 1000000 \
--local-batch-size 5 \
diff --git a/scripts/ucf101/train_256.sh b/scripts/ucf101/train_256.sh
new file mode 100644
index 000000000..7bb45639f
--- /dev/null
+++ b/scripts/ucf101/train_256.sh
@@ -0,0 +1,23 @@
+
+accelerate launch \
+ --config_file scripts/accelerate_configs/ddp_config.yaml \
+ opensora/train/train.py \
+ --model Latte-XL/122 \
+ --dataset ucf101 \
+ --num-classes 101 \
+ --ae stabilityai/sd-vae-ft-mse \
+ --data-path /remote-home/yeyang/UCF-101 \
+ --extras 2 \
+ --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/videogpt/train_videogpt.sh b/scripts/videogpt/train_videogpt.sh
index 45a9558fc..be7f74fd3 100644
--- a/scripts/videogpt/train_videogpt.sh
+++ b/scripts/videogpt/train_videogpt.sh
@@ -1,26 +1,30 @@
+export ACCELERATE_GRADIENT_ACCUMULATION_STEPS=1
+
accelerate launch \
- --config_file scripts/accelerate_configs/default_config.yaml \
+ --config_file scripts/accelerate_configs/ddp_config.yaml \
opensora/train/train_videogpt.py \
--do_train \
--seed 1234 \
- --data_path UCF-101 \
- --per_device_train_batch_size 8 \
- --gradient_accumulation_steps 8 \
- --learning_rate 1e-5 \
+ --data_path "datasets/UCF-101/" \
+ --per_device_train_batch_size 32 \
+ --gradient_accumulation_steps $ACCELERATE_GRADIENT_ACCUMULATION_STEPS \
+ --learning_rate 7e-4 \
--weight_decay 0. \
--num_train_epochs 2 \
- --embedding_dim 512 \
--lr_scheduler_type cosine \
--max_grad_norm 1.0 \
--save_strategy steps \
--save_total_limit 5 \
--logging_steps 5 \
- --save_steps 5 \
- --n_codes 2048 \
+ --save_steps 1000 \
+ --n_codes 1024 \
--n_hiddens 240 \
+ --embedding_dim 4 \
--n_res_layers 4 \
--downsample "4,4,4" \
--resolution 128 \
--sequence_length 16 \
- --output_dir results/videogpt \
- --report_to tensorboard
\ No newline at end of file
+ --output_dir results/videogpt_444_128 \
+ --bf16 True \
+ --fp16 False \
+ --report_to tensorboard
diff --git a/scripts/videogpt/train_videogpt_dsz2.sh b/scripts/videogpt/train_videogpt_dsz2.sh
new file mode 100644
index 000000000..7875333eb
--- /dev/null
+++ b/scripts/videogpt/train_videogpt_dsz2.sh
@@ -0,0 +1,30 @@
+export ACCELERATE_GRADIENT_ACCUMULATION_STEPS=1
+
+accelerate launch \
+ --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
+ opensora/train/train_videogpt.py \
+ --do_train \
+ --seed 1234 \
+ --data_path "datasets/UCF-101/" \
+ --per_device_train_batch_size 32 \
+ --gradient_accumulation_steps $ACCELERATE_GRADIENT_ACCUMULATION_STEPS \
+ --learning_rate 7e-4 \
+ --weight_decay 0. \
+ --num_train_epochs 2 \
+ --lr_scheduler_type cosine \
+ --max_grad_norm 1.0 \
+ --save_strategy steps \
+ --save_total_limit 5 \
+ --logging_steps 5 \
+ --save_steps 10000 \
+ --n_codes 1024 \
+ --n_hiddens 240 \
+ --embedding_dim 4 \
+ --n_res_layers 4 \
+ --downsample "4,4,4" \
+ --resolution 128 \
+ --sequence_length 16 \
+ --output_dir results/videogpt_444_128 \
+ --bf16 True \
+ --fp16 False \
+ --report_to tensorboard
diff --git a/scripts/videogpt/train_videogpt_dsz3.sh b/scripts/videogpt/train_videogpt_dsz3.sh
new file mode 100644
index 000000000..484084236
--- /dev/null
+++ b/scripts/videogpt/train_videogpt_dsz3.sh
@@ -0,0 +1,30 @@
+export ACCELERATE_GRADIENT_ACCUMULATION_STEPS=1
+
+accelerate launch \
+ --config_file scripts/accelerate_configs/deepspeed_zero3_config.yaml \
+ opensora/train/train_videogpt.py \
+ --do_train \
+ --seed 1234 \
+ --data_path "datasets/UCF-101/" \
+ --per_device_train_batch_size 32 \
+ --gradient_accumulation_steps $ACCELERATE_GRADIENT_ACCUMULATION_STEPS \
+ --learning_rate 7e-4 \
+ --weight_decay 0. \
+ --num_train_epochs 2 \
+ --lr_scheduler_type cosine \
+ --max_grad_norm 1.0 \
+ --save_strategy steps \
+ --save_total_limit 5 \
+ --logging_steps 5 \
+ --save_steps 10000 \
+ --n_codes 1024 \
+ --n_hiddens 240 \
+ --embedding_dim 4 \
+ --n_res_layers 4 \
+ --downsample "4,4,4" \
+ --resolution 128 \
+ --sequence_length 16 \
+ --output_dir results/videogpt_444_128 \
+ --bf16 True \
+ --fp16 False \
+ --report_to tensorboard