58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
"""
|
|
-----------------------------------------------------------------------------
|
|
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
and proprietary rights in and to this software, related documentation
|
|
and any modifications thereto. Any use, reproduction, disclosure or
|
|
distribution of this software and related documentation without an express
|
|
license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
-----------------------------------------------------------------------------
|
|
"""
|
|
|
|
from typing import Literal, Optional
|
|
|
|
import attrs
|
|
|
|
|
|
@attrs.define(slots=False)
|
|
class ModelConfig:
|
|
# vae
|
|
vae_conf: str = "vae.configs.part_woenc"
|
|
vae_ckpt_path: Optional[str] = None
|
|
|
|
# learn & generate parts
|
|
use_parts: bool = False
|
|
part_embed_mode: Literal["element", "part", "part2_only"] = "part2_only"
|
|
shuffle_parts: bool = False
|
|
use_num_parts_cond: bool = False
|
|
|
|
# flow matching hyper-params
|
|
flow_shift: float = 1.0
|
|
logitnorm_mean: float = 0.0
|
|
logitnorm_std: float = 1.0
|
|
|
|
# image encoder
|
|
dino_model: Literal["dinov2_vitl14_reg", "dinov2_vitg14"] = "dinov2_vitg14"
|
|
|
|
# backbone DiT
|
|
hidden_dim: int = 1536
|
|
num_heads: int = 16
|
|
num_layers: int = 24
|
|
qknorm: bool = True
|
|
qknorm_type: Literal["LayerNorm", "RMSNorm"] = "RMSNorm"
|
|
use_pos_embed: bool = False
|
|
|
|
# latent code
|
|
latent_size: Optional[int] = None # if None, will load from vae
|
|
latent_dim: Optional[int] = None
|
|
|
|
# preload vae weights
|
|
preload_vae: bool = True
|
|
|
|
# preload dinov2 weights
|
|
preload_dinov2: bool = True
|
|
|
|
# init weights from a pretrained checkpoint
|
|
pretrain_path: Optional[str] = None
|