From da08e767dbe8fbf9fc3ec962063ada85a0108a61 Mon Sep 17 00:00:00 2001 From: qqingzheng <2533221180@qq.com> Date: Wed, 6 Mar 2024 15:48:19 +0000 Subject: [PATCH] [refactor] rename Former-commit-id: ec317605d24223527b8c2d4b4d2ee876f86abbaf --- examples/rec_video.py | 4 +-- opensora/models/ae/__init__.py | 4 +-- opensora/models/ae/videobase/__init__.py | 10 +++--- .../models/ae/videobase/videogpt/__init__.py | 4 --- .../models/ae/videobase/vqvae/__init__.py | 4 +++ .../configuration_vqvae.py} | 2 +- .../dataset_vqvae.py} | 2 +- .../modeling_vqvae.py} | 13 ++++---- .../trainer_vqvae.py} | 2 +- train_scripts/train_videogpt.py | 31 +++++++++---------- 10 files changed, 37 insertions(+), 39 deletions(-) delete mode 100644 opensora/models/ae/videobase/videogpt/__init__.py create mode 100644 opensora/models/ae/videobase/vqvae/__init__.py rename opensora/models/ae/videobase/{videogpt/configuration_videogpt.py => vqvae/configuration_vqvae.py} (93%) rename opensora/models/ae/videobase/{videogpt/dataset_videogpt.py => vqvae/dataset_vqvae.py} (98%) rename opensora/models/ae/videobase/{videogpt/modeling_videogpt.py => vqvae/modeling_vqvae.py} (98%) rename opensora/models/ae/videobase/{videogpt/trainer_videogpt.py => vqvae/trainer_vqvae.py} (93%) diff --git a/examples/rec_video.py b/examples/rec_video.py index 7c9e343db..160c8e337 100644 --- a/examples/rec_video.py +++ b/examples/rec_video.py @@ -8,7 +8,7 @@ from torchvision.io import read_video from torchvision.transforms import Lambda, Compose from torchvision.transforms._transforms_video import RandomCropVideo from torch.nn import functional as F -from opensora.models.ae import VideoGPTVQVAE +from opensora.models.ae import VQVAEModel import argparse def array_to_video(image_array, fps=30.0, output_file='output_video.mp4'): @@ -82,7 +82,7 @@ def main(args): sample_rate = args.sample_rate device = torch.device('cuda') - vqvae = VideoGPTVQVAE.download_and_load_model(args.ckpt) + vqvae = VQVAEModel.download_and_load_model(args.ckpt) vqvae.eval() vqvae = vqvae.to(device) diff --git a/opensora/models/ae/__init__.py b/opensora/models/ae/__init__.py index 20f9f7ba1..d5f25d2a8 100644 --- a/opensora/models/ae/__init__.py +++ b/opensora/models/ae/__init__.py @@ -1,8 +1,8 @@ from .imagebase import imagebase_ae, vae, vqvae, imagebase_ae_stride, imagebase_ae_channel from .videobase import videobase_ae, videovae, videovqvae, videobase_ae_stride, videobase_ae_channel from .videobase import ( - VideoGPTConfiguration, - VideoGPTVQVAE + VQVAEConfiguration, + VQVAEModel ) ae_stride_config = {} diff --git a/opensora/models/ae/videobase/__init__.py b/opensora/models/ae/videobase/__init__.py index 0747182ea..623c1d077 100644 --- a/opensora/models/ae/videobase/__init__.py +++ b/opensora/models/ae/videobase/__init__.py @@ -1,11 +1,11 @@ """ new code """ -from .videogpt import ( - VideoGPTConfiguration, - VideoGPTVQVAE, - VideoGPTTrainer, - VideoGPTDataset, +from .vqvae import ( + VQVAEConfiguration, + VQVAEModel, + VQVAETrainer, + VQVAEDataset, ) diff --git a/opensora/models/ae/videobase/videogpt/__init__.py b/opensora/models/ae/videobase/videogpt/__init__.py deleted file mode 100644 index e90e4ef48..000000000 --- a/opensora/models/ae/videobase/videogpt/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .configuration_videogpt import VideoGPTConfiguration -from .modeling_videogpt import VideoGPTVQVAE -from .trainer_videogpt import VideoGPTTrainer -from .dataset_videogpt import VideoGPTDataset diff --git a/opensora/models/ae/videobase/vqvae/__init__.py b/opensora/models/ae/videobase/vqvae/__init__.py new file mode 100644 index 000000000..c16062399 --- /dev/null +++ b/opensora/models/ae/videobase/vqvae/__init__.py @@ -0,0 +1,4 @@ +from .configuration_vqvae import VQVAEConfiguration +from .modeling_vqvae import VQVAEModel +from .trainer_vqvae import VQVAETrainer +from .dataset_vqvae import VQVAEDataset \ No newline at end of file diff --git a/opensora/models/ae/videobase/videogpt/configuration_videogpt.py b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py similarity index 93% rename from opensora/models/ae/videobase/videogpt/configuration_videogpt.py rename to opensora/models/ae/videobase/vqvae/configuration_vqvae.py index 93e0f4faa..a21f52dfd 100644 --- a/opensora/models/ae/videobase/videogpt/configuration_videogpt.py +++ b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from typing import Tuple @dataclass -class VideoGPTConfiguration: +class VQVAEConfiguration: embedding_dim: int n_codes: int n_hiddens: int diff --git a/opensora/models/ae/videobase/videogpt/dataset_videogpt.py b/opensora/models/ae/videobase/vqvae/dataset_vqvae.py similarity index 98% rename from opensora/models/ae/videobase/videogpt/dataset_videogpt.py rename to opensora/models/ae/videobase/vqvae/dataset_vqvae.py index df467cc6f..c7f28c06a 100644 --- a/opensora/models/ae/videobase/videogpt/dataset_videogpt.py +++ b/opensora/models/ae/videobase/vqvae/dataset_vqvae.py @@ -10,7 +10,7 @@ import torch.nn.functional as F from torchvision.datasets.video_utils import VideoClips # Copied from https://github.com/wilson1yan/VideoGPT -class VideoGPTDataset(data.Dataset): +class VQVAEDataset(data.Dataset): """ Generic dataset for videos files stored in folders Returns BCTHW videos in the range [-0.5, 0.5] """ exts = ['avi', 'mp4', 'webm'] diff --git a/opensora/models/ae/videobase/videogpt/modeling_videogpt.py b/opensora/models/ae/videobase/vqvae/modeling_vqvae.py similarity index 98% rename from opensora/models/ae/videobase/videogpt/modeling_videogpt.py rename to opensora/models/ae/videobase/vqvae/modeling_vqvae.py index f5637495b..68f8d0a9f 100644 --- a/opensora/models/ae/videobase/videogpt/modeling_videogpt.py +++ b/opensora/models/ae/videobase/vqvae/modeling_vqvae.py @@ -5,8 +5,7 @@ import numpy as np import torch.distributed as dist import torch.nn.functional as F import math -from transformers import LlamaForCausalLM -from .configuration_videogpt import VideoGPTConfiguration +from .configuration_vqvae import VQVAEConfiguration # Copied from https://github.com/wilson1yan/VideoGPT def view_range(x, i, j, shape): @@ -639,7 +638,7 @@ class SparseAttention(nn.Module): return view_range(out, 2, 3, old_shape) # Modified from https://github.com/wilson1yan/VideoGPT -class VideoGPTVQVAE(VideoBaseAE): +class VQVAEModel(VideoBaseAE): DOWNLOADED_VQVAE = { 'bair_stride4x2x2': '1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L', @@ -648,7 +647,7 @@ class VideoGPTVQVAE(VideoBaseAE): 'kinetics_stride2x4x4': '1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB' } - def __init__(self, config: VideoGPTConfiguration): + def __init__(self, config: VQVAEConfiguration): super().__init__() self.embedding_dim = config.embedding_dim self.n_codes = config.n_codes @@ -687,7 +686,7 @@ class VideoGPTVQVAE(VideoBaseAE): hyper_parameters = vars(model_cpkt.get("hyper_parameters").get("args")) print(hyper_parameters) state_dict = model_cpkt.get("state_dict") - model = VideoGPTVQVAE(config=VideoGPTConfiguration(**hyper_parameters)) + model = VQVAEModel(config=VQVAEConfiguration(**hyper_parameters)) model.load_state_dict(state_dict) return model else: @@ -696,6 +695,6 @@ class VideoGPTVQVAE(VideoBaseAE): @staticmethod def download_and_load_model(model_name, cache_dir=None): from .....utils.downloader import gdown_download - path = gdown_download(VideoGPTVQVAE.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir) - return VideoGPTVQVAE.load_from_checkpoint(path) + path = gdown_download(VQVAEModel.DOWNLOADED_VQVAE[model_name], model_name, cache_dir=cache_dir) + return VQVAEModel.load_from_checkpoint(path) \ No newline at end of file diff --git a/opensora/models/ae/videobase/videogpt/trainer_videogpt.py b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py similarity index 93% rename from opensora/models/ae/videobase/videogpt/trainer_videogpt.py rename to opensora/models/ae/videobase/vqvae/trainer_vqvae.py index fc5f4ada4..c2ede2ed0 100644 --- a/opensora/models/ae/videobase/videogpt/trainer_videogpt.py +++ b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py @@ -2,7 +2,7 @@ from torch import nn from transformers import Trainer import torch.nn.functional as F -class VideoGPTTrainer(Trainer): +class VQVAETrainer(Trainer): def compute_loss(self, model, inputs, return_outputs=False): model = model.module diff --git a/train_scripts/train_videogpt.py b/train_scripts/train_videogpt.py index 30d5c7aa2..7fcf247ff 100644 --- a/train_scripts/train_videogpt.py +++ b/train_scripts/train_videogpt.py @@ -1,13 +1,12 @@ import sys sys.path.append(".") from opensora.models.ae.videobase import ( - VideoGPTVQVAE, - VideoGPTConfiguration, - VideoGPTDataset, - VideoGPTTrainer, + VQVAEModel, + VQVAEConfiguration, + VQVAEDataset, + VQVAETrainer, ) import argparse -import accelerate from typing import Optional from accelerate.utils import set_seed from transformers import HfArgumentParser, TrainingArguments @@ -15,18 +14,18 @@ from dataclasses import dataclass, field @dataclass -class VideoGPTArgument(VideoGPTConfiguration): +class VQVAEArgument(VQVAEConfiguration): data_path: str = field(default=None, metadata={"help": "data path"}) @dataclass -class VideoGPTTrainingArgument(TrainingArguments): +class VQVAETrainingArgument(TrainingArguments): remove_unused_columns: Optional[bool] = field( default=False, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."} ) -def train(args, videogpt_args, training_args): +def train(args, VQVAE_args, training_args): # Load Model - model_config = VideoGPTConfiguration( + model_config = VQVAEConfiguration( embedding_dim=args.embedding_dim, n_codes=args.n_codes, n_hiddens=args.n_hiddens, @@ -35,19 +34,19 @@ def train(args, videogpt_args, training_args): downsample=args.downsample, resolution=args.resolution ) - model = VideoGPTVQVAE(model_config) + model = VQVAEModel(model_config) # Load Dataset - dataset = VideoGPTDataset(args.data_path, sequence_length=args.sequence_length) + dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length) # Load Trainer - trainer = VideoGPTTrainer(model, training_args, train_dataset=dataset) + trainer = VQVAETrainer(model, training_args, train_dataset=dataset) trainer.train() if __name__ == "__main__": - parser = HfArgumentParser((VideoGPTArgument, VideoGPTTrainingArgument)) - videogpt_args, training_args = parser.parse_args_into_dataclasses() - args = argparse.Namespace(**vars(videogpt_args), **vars(training_args)) + parser = HfArgumentParser((VQVAEArgument, VQVAETrainingArgument)) + VQVAE_args, training_args = parser.parse_args_into_dataclasses() + args = argparse.Namespace(**vars(VQVAE_args), **vars(training_args)) set_seed(args.seed) - train(args, videogpt_args, training_args) + train(args, VQVAE_args, training_args)