Files
hao-ai-lab-FastVideo/opensora/models/ae/__init__.py
T
LinB203 ff36d291d8 train with image
Former-commit-id: 209b9a8ad2f19a50e46f9b7ec68dd781898e491a
2024-03-19 23:21:57 +08:00

28 lines
784 B
Python

from .imagebase import imagebase_ae, imagebase_ae_stride, imagebase_ae_channel
from .videobase import videobase_ae, videobase_ae_stride, videobase_ae_channel
from .videobase import (
VQVAEConfiguration,
VQVAEModel,
VQVAEDataset,
VQVAETrainer,
CausalVQVAEModel,
CausalVQVAEConfiguration,
CausalVQVAEDataset,
CausalVQVAETrainer
)
ae_stride_config = {}
ae_stride_config.update(imagebase_ae_stride)
ae_stride_config.update(videobase_ae_stride)
ae_channel_config = {}
ae_channel_config.update(imagebase_ae_channel)
ae_channel_config.update(videobase_ae_channel)
def getae(args):
"""deprecation"""
ae = imagebase_ae.get(args.ae, None)
if ae is None:
ae = videobase_ae.get(args.ae, None)
assert ae is not None
return ae(args.ae)