From 7a78ef2a97c0c2d7a2a368d5adf1338d2e1d7dfc Mon Sep 17 00:00:00 2001
From: qqingzheng <2533221180@qq.com>
Date: Wed, 6 Mar 2024 15:59:13 +0000
Subject: [PATCH] [refactor] configuration default values

Former-commit-id: d7ca7685d6a8c6d015082dda241bc98ab156770d
---
 .../ae/videobase/vqvae/configuration_vqvae.py | 20 +++++++++----------
 train_scripts/train_videogpt.py               | 19 +++++-------------
 2 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
index a21f52dfd..b472101d4 100644
--- a/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
+++ b/opensora/models/ae/videobase/vqvae/configuration_vqvae.py
@@ -1,17 +1,17 @@
 from ..configuration_videobase import VideoBaseConfiguration
-from dataclasses import dataclass
-from typing import Tuple
+from dataclasses import dataclass, field
+from typing import Tuple, Union
 
 @dataclass
 class VQVAEConfiguration:
-    embedding_dim: int
-    n_codes: int
-    n_hiddens: int
-    n_res_layers: int
-    resolution: int
-    sequence_length: int
-    downsample: str
-    no_pos_embd: bool
+    embedding_dim: int = field(default=256)
+    n_codes: int = field(default=2048)
+    n_hiddens: int = field(default=240)
+    n_res_layers: int = field(default=4)
+    resolution: int = field(default=128)
+    sequence_length: int = field(default=16)
+    downsample: Union[str, Tuple[int, int, int]] = field(default=(4,4,4))
+    no_pos_embd: bool = field(default=True)
     def __post_init__(self):
         if isinstance(self.downsample, str):
             self.downsample = tuple(map(int, self.downsample.split(',')))
\ No newline at end of file
diff --git a/train_scripts/train_videogpt.py b/train_scripts/train_videogpt.py
index 7fcf247ff..532c94995 100644
--- a/train_scripts/train_videogpt.py
+++ b/train_scripts/train_videogpt.py
@@ -23,18 +23,9 @@ class VQVAETrainingArgument(TrainingArguments):
         default=False, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."}
     )
 
-def train(args, VQVAE_args, training_args):
+def train(args, vqvae_args, training_args):
     # Load Model
-    model_config = VQVAEConfiguration(
-        embedding_dim=args.embedding_dim,
-        n_codes=args.n_codes,
-        n_hiddens=args.n_hiddens,
-        n_res_layers=args.n_res_layers,
-        sequence_length=args.sequence_length,
-        downsample=args.downsample,
-        resolution=args.resolution
-    )
-    model = VQVAEModel(model_config)
+    model = VQVAEModel(vqvae_args)
     # Load Dataset
     dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length)
 
@@ -45,8 +36,8 @@ def train(args, VQVAE_args, training_args):
 
 if __name__ == "__main__":
     parser = HfArgumentParser((VQVAEArgument, VQVAETrainingArgument))
-    VQVAE_args, training_args = parser.parse_args_into_dataclasses()
-    args = argparse.Namespace(**vars(VQVAE_args), **vars(training_args))
+    vqvae_args, training_args = parser.parse_args_into_dataclasses()
+    args = argparse.Namespace(**vars(vqvae_args), **vars(training_args))
     set_seed(args.seed)
 
-    train(args, VQVAE_args, training_args)
+    train(args, vqvae_args, training_args)