EgoX를 LongLive(causal AR)로 이식

TL;DR

LongLive upstream import 후 3개 파일 수정(+56/-6) + egox 모듈 신규로 exo→ego causal 학습 경로 구현
20-step smoke (2 GPU): loss 0.348→0.204, 3.2 s/step, 62.8GB

+56/-6

core 수정 라인

3.2s

s/step

smoke jobs

1 배경 / 목적

EgoX(exo→ego 비디오 생성)를 Wan2.2 양방향 DiT 대신 LongLive 2.0 (causal block-AR, KV cache) 위에서 학습해 긴 영상(609f, 20s)을 생성하려는 시도. LongLive는 T2V 전용이라 exo 조건·ego prior 입력 경로가 없다.

기존 한계: 양방향 Wan2.2 EgoX는 81f 단위 생성이라 긴 영상은 chunk 이어붙이기 필요. LongLive upstream(NVlabs/LongLive @4ce82ab, 238 files)을 import한 직후라 EgoX 조건 경로가 전무.

2 작업 내용

[Import] b89d7cb LongLive2.0 @ 4ce82ab (238 files, +41,866 lines) [수정] wan_5b/modules/causal_model.py (+44/-6) - CausalWanAttentionBlock / CausalHead: modulation e가 [B,F,S,6,C] token-level이어도 처리 (token_level 분기, Triton adaln은 token_level 시 비활성) - _forward_train / _forward_inference: exo_cols 인자 → 프레임당 [exo|ego] 중 exo 토큰(앞 exo_cols 컬럼)에 t=0 modulation 적용 - teacher-forcing clean_x 쪽에도 per-frame 조건 채널 y concat [수정] utils/wan_5b_wrapper.py: forward에 y, exo_cols 전달 [신규] model/egox_diffusion.py EgoXCausalDiffusion.egox_loss (exo는 clean·loss 제외, ego만 noise+flow loss, error recycling은 ego 컬럼에만) [신규] utils/egox_dataset.py, scripts/egox/{train_egox.py,infer_egox.py,eval_ckpt.sh,eval_wild.sh,smoke_t2v.*} [신규] configs/egox_ar.yaml, egox_infer.yaml, egox_smoke_t2v.yaml

레이아웃: latent 프레임당 [exo 448x768 | ego 448x448] 가로 concat → 28 x (48+28) latent, 14x38=532 tokens/frame. ego prior는 extra input channel y (prior_channels=48, exo 영역은 0)로 patch_embedding 확장(widen_patch_embedding). 81f→21 latent→7 block (num_frame_per_block=3).

추론: LongLive pipeline은 그대로 두고 EgoXGeneratorAdapter가 generator 앞에서 clean exo latent 삽입 + prior 채널 전달 (num_frames=81/201/609, latent 수가 block 크기 3의 배수).

3 결과 (검증 수치)

Job	내용	GPU	결과
190684~190703	smoke_t2v (LongLive T2V 동작 확인, container 옵션 시행착오)	1	COMPLETED 2m34s / 2m46s / 3m43s
190721	EgoX train smoke (max_steps 20, no_save)	2	COMPLETED 5m56s, 3.2 s/step, mem 62.8GB, loss 0.348→0.204, grad_norm ~0.29
190735 / 190773	본 학습 첫 제출 (8 GPU, 4 GPU extra)	8/4	CANCELLED (시작 전, 0s) → own 4 GPU로 재제출

smoke 통과: token-level modulation + exo t=0 경로 + y 채널 확장 모두 forward/backward NaN 없이 20 step 수행. 4 GPU 본 학습은 mem 36.1GB/GPU로 FSDP full shard 여유 충분.

4 Takeaway

causal LongLive에 exo 조건을 "프레임 내 token-level timestep modulation"으로 주입하는 방식이 코드 변경 3개 파일(+56/-6)로 가능. LongLive pipeline/KV cache 로직을 수정하지 않고 adapter만으로 추론 가능해 upstream 추종 비용이 낮다.

5 Next Steps

10K step 학습 후 ego 품질 정량 평가 (→ 별도 카드 egox_ar_10k_train_eval)
token_level 시 Triton adaln 비활성 → inference 속도 저하 여부 측정, 필요 시 token-level 커널화
변경분 미커밋 상태(M 3 files, ?? 신규 다수) — 커밋 필요