merge main branch.
Former-commit-id: 5f78da5eae2b18305bb4a0a6dc7ab771e9b96f60
This commit is contained in:
sennnnn
@@ -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.
|
||||
|
||||
<p align="center">
|
||||
<img src="assets/loss.jpg" width=60%>
|
||||
<img src="assets/loss.png" width=90%>
|
||||
</p>
|
||||
|
||||
#### 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** |
|
||||
| --- | --- |
|
||||
| <img src="assets/demo_0006.gif"> | <img src="assets/demo_0011.gif"> |
|
||||
|
||||
|
||||
## 🤝 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.
|
||||
|
||||
|
||||
+12
-7
@@ -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='')
|
||||
|
||||
@@ -29,3 +29,5 @@ class VQVAEConfiguration(VideoBaseConfiguration):
|
||||
self.downsample = downsample
|
||||
|
||||
self.no_pos_embd = no_pos_embd
|
||||
|
||||
self.hidden_size = n_hiddens
|
||||
|
||||
@@ -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]
|
||||
@@ -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 <PIL.GifImagePlugin.GifImageFile> to <PIL.JpegImagePlugin.JpegImageFile>
|
||||
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 # [<PIL.JpegImagePlugin.JpegImageFile>, ...]
|
||||
|
||||
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
|
||||
@@ -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."}
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
+41
-50
@@ -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)
|
||||
|
||||
+12
-6
@@ -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)")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
server1 slots=8 # your server name and GPU in total
|
||||
server2 slots=8
|
||||
@@ -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
|
||||
+3
-2
@@ -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
|
||||
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
--output_dir results/videogpt_444_128 \
|
||||
--bf16 True \
|
||||
--fp16 False \
|
||||
--report_to tensorboard
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user