[refactor] rename

Former-commit-id: ec317605d24223527b8c2d4b4d2ee876f86abbaf

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)
 

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 = {}

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,
 )
 
 

opensora/models/ae/videobase/videogpt/__init__.py

@@ -1,4 +0,0 @@
-from .configuration_videogpt import VideoGPTConfiguration
-from .modeling_videogpt import VideoGPTVQVAE
-from .trainer_videogpt import VideoGPTTrainer
-from .dataset_videogpt import VideoGPTDataset

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

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

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']

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)
         

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

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)