mirror of
https://github.com/THUDM/CogVideo.git
synced 2025-04-06 03:57:56 +08:00
91d79fd9a4
Add Pydantic models to handle: - CLI arguments and configuration (Args) - Model components and pipeline (Components) - Training state and parameters (State)
28 lines
492 B
Python
28 lines
492 B
Python
from typing import Any
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class Components(BaseModel):
|
|
# pipeline cls
|
|
pipeline_cls: Any = None
|
|
|
|
# Tokenizers
|
|
tokenizer: Any = None
|
|
tokenizer_2: Any = None
|
|
tokenizer_3: Any = None
|
|
|
|
# Text encoders
|
|
text_encoder: Any = None
|
|
text_encoder_2: Any = None
|
|
text_encoder_3: Any = None
|
|
|
|
# Autoencoder
|
|
vae: Any = None
|
|
|
|
# Denoiser
|
|
transformer: Any = None
|
|
unet: Any = None
|
|
|
|
# Scheduler
|
|
scheduler: Any = None
|