KyleBae1017/stldm-sevir-vae

STLDM SEVIR VAE

STLDM의 latent diffusion 컴포넌트 중 VAE 부분만 추출한 체크포인트. JEPACast Phase 1처럼 VAE를 frozen tokenizer로 재사용하는 setup에서 그대로 로드해서 사용.

출처

• 원본 STLDM end-to-end 체크포인트: 20260406181528_1gpu_reimp/stldm_..._final.pt
• 그 중 backbone.vae.* 키만 추출 (38 keys)

아키텍처

• VAE encoder + decoder (4-stage ConvSC, hid_S=32, N_S=4)
• 입력: (B, 1, 128, 128) SEVIR VIL frame, [0, 1]
• Latent: mean/log_var 각각 (B, 32, 32, 32) (공간 1/4 다운샘플)

로드 예시 (이 repo가 가정하는 STLDM 코드베이스 기준)

from stldm.modules import VAE
import torch
from huggingface_hub import hf_hub_download

vae_path = hf_hub_download(
    repo_id="KyleBae1017/stldm-sevir-vae",
    filename="vae_only.pt",
)
vae = VAE(C_in=1, hid_S=32, N_S=4)
state = torch.load(vae_path, map_location="cpu")
missing, unexpected = vae.load_state_dict(state, strict=True)
assert not missing and not unexpected
vae.eval()
for p in vae.parameters():
    p.requires_grad_(False)

train.py에서 바로 쓰기

먼저 위 hf_hub_download로 받아서 로컬에 두고:

python train.py ... --ae_ckpt <local_path> --ae_eval --freeze_vae ...