mirror of
https://github.com/THUDM/CogVideo.git
synced 2025-04-06 03:57:56 +08:00
93b906b3fb
Add docstring to train_frames field in State schema to explicitly indicate that it includes one image padding frame
27 lines
768 B
Python
27 lines
768 B
Python
import torch
|
|
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
from pydantic import BaseModel, field_validator
|
|
|
|
class State(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
|
|
train_frames: int # user-defined training frames, **containing one image padding frame**
|
|
train_height: int
|
|
train_width: int
|
|
|
|
transformer_config: Dict[str, Any] = None
|
|
|
|
weight_dtype: torch.dtype = torch.float32
|
|
num_trainable_parameters: int = 0
|
|
overwrote_max_train_steps: bool = False
|
|
num_update_steps_per_epoch: int = 0
|
|
total_batch_size_count: int = 0
|
|
|
|
generator: torch.Generator | None = None
|
|
|
|
validation_prompts: List[str] = []
|
|
validation_images: List[Path | None] = []
|
|
validation_videos: List[Path | None] = []
|