Flow-Factory Adapter 인터페이스 조사 — Cosmos-Predict2 매핑 표 확정

TL;DR

FF BaseAdapter 실제 소스(flow_factory.models.abc) 기준 abstract 메서드 4개 확인 — load_pipeline, decode_latents, forward, inference. encode_prompt/encode_image는 optional override.
레퍼런스 구현 Wan2_I2V_Adapter 5개 메서드 전부 추적 — Cosmos 매핑 표 6항목 완성.
계획 문서(plan-vggrpo_phase2_rl.md)의 잘못된 import 경로 3개 발견 및 정정 → Task 4–16 전체 적용 필요.

abstract 메서드

Cosmos 매핑 항목

잘못된 경로 정정

1배경 / 목적

Phase 2 RL Sub-plan #1의 Task 3은 src/cosmos_ff/adapter.py 작성이다. 그 전에 Flow-Factory(FF)의 BaseAdapter 추상화를 실제 소스 코드 기준으로 정확히 파악해야 한다. 계획 문서에 기재된 FF 경로가 실제와 다를 경우 Task 4–16 전체에 걸쳐 잘못된 import가 누적될 수 있다.

조사 동기: 계획 문서는 Wan2.1 시나리오로 초안이 작성됐고, Cosmos-Predict2로 base가 교체된 뒤 adapter 경로/시그니처 검증이 이루어지지 않은 상태였다.

조사 범위: ① BaseAdapter + BaseSample 추상화, ② PointwiseRewardModel / GroupwiseRewardModel, ③ 레퍼런스 구현 Wan2_I2V_Adapter 전 메서드, ④ Cosmos-Predict2 매핑 가능성 검토.

2작업 내용

① BaseAdapter (`flow_factory.models.abc`, line 95)

필수 abstract 메서드 4개:

@abstractmethod def load_pipeline(self) -> DiffusionPipeline: ... # line 186 @abstractmethod def decode_latents(self, latents, **kwargs): ... # line 2124 @abstractmethod def forward(self, *args, **kwargs) -> SDESchedulerOutput: ... # line 2137 @abstractmethod def inference(self, *args, **kwargs) -> List[BaseSample]: ... # line 2148

encode_prompt / encode_image는 abstract가 아닌 non-abstract 메서드(기본값 return None). 서브클래스가 override 가능하나 필수 아님. → Cosmos adapter는 최소한 4개만 구현해도 컴파일 통과.

load_scheduler는 concrete 구현 제공(line 190). __init__에서 self.pipeline = self.load_pipeline() → self.pipeline.scheduler = self.load_scheduler() 순서 호출.

② BaseSample (`flow_factory.samples.BaseSample`)

주요 필드 구조:

필드	Shape / Type	설명
`timesteps`	(T+1,)	denoising trajectory
`all_latents`	(num_steps, Seq_len, C)	전체 latent 이력
`log_probs`	(num_steps,)	step별 log probability
`latent_index_map`	(T+1,) LongTensor	latent 인덱스 매핑
`image / video / audio`	C×H×W / T×C×H×W	생성 결과 (자동 canonicalize)
`prompt_embeds`	tensor	텍스트 임베딩 (CFG 포함)

_shared_fields = frozenset({'height', 'width', 'latent_index_map', 'log_prob_index_map'}) — 배치 내 공유 필드(stack 시 첫 번째 값만 사용). DDP/FSDP 호환을 위해 __init_subclass__에서 PyTorch pytree node로 자동 등록.

③ Reward base classes (`flow_factory.rewards.abc`)

@dataclass class RewardModelOutput(BaseOutput): # line 37 rewards: Union[torch.Tensor, np.ndarray, List[float]] extra_info: Optional[Dict[str, Any]] = None class PointwiseRewardModel(BaseRewardModel): # line 78 # rewards shape: (batch_size,) — 배치당 독립 reward class GroupwiseRewardModel(BaseRewardModel): # line 125 # 동일 unique_id 그룹 전체를 한 번에 받음 — pairwise/ranking reward

use_tensor_inputs: bool = False이면 PIL Images 입력, True이면 torch.Tensor 입력. VGGRPO reward는 PointwiseRewardModel로 구현하고 use_tensor_inputs=True 설정 예정.

④ 레퍼런스 구현 — Wan2_I2V_Adapter 전 메서드 추적

flow_factory.models.wan.wan2_i2v.Wan2_I2V_Adapter(BaseAdapter)를 Cosmos analog로 분석:

메서드	Wan2_I2V 구현 요약
`load_pipeline`	`WanImageToVideoPipeline.from_pretrained(model_path)`
`encode_prompt`	T5 tokenizer → `text_encoder().last_hidden_state` → seq masking → zero-pad to max_seq_len. CFG: guidance_scale > 1.0
`encode_image`	`pipeline.image_processor` → `image_encoder(hidden_states[-2])`. Wan2.2는 `first_frame_mask` per-token expand (Cosmos에는 없는 패턴)
`inference`	set_timesteps → encode → prepare_latents → denoising loop (forward 호출) → decode_latents → `List[WanI2VSample]`
`forward`	단일 denoising step: transformer → CFG → scheduler.step. `compute_log_prob=True`일 때 log_probs 채움

3결과 — Cosmos-Predict2 매핑 표

FF 메서드	Cosmos analog	주요 고려사항
`load_pipeline`	`Predict2Video2WorldModel.__init__` + `Video2WorldPipeline.from_config`	pseudo-pipeline wrap 필요. FSDP off(single-GPU dev). pipeline.scheduler를 FF SDE scheduler로 교체
`default_target_modules`	`add_lora_to_model`의 q/k/v/output/mlp 타겟	cosmos 코드에 LoRA target 목록 이미 존재 → 해당 문자열 그대로 반환
`encode_prompt`	`CosmosT5TextEncoder` (`imaginaire.auxiliary.text_encoder`)	T5-XXL, max_len=512, dim=1024. `return_dict=False` → last_hidden_state 직접 접근
`encode_image`	`pipe.vae.encode(first_frame_tensor)`	`pipe.tokenizer`는 text tokenizer. 영상 인코딩은 VAE: resize→normalize→encode 순서
`inference`	`RectifiedFlowAB2Scheduler.set_timesteps` + step loop	SDE step 교체 필요(`sde.py`). `trajectory_indices='all'`로 전체 수집 후 log_probs 채움
`forward`	`DiT.forward` + x0→v 변환 + `flow_sde_step`	x0-pred → paper v-pred algebra. `return_dict=False`로 [0]만 사용. Wan2.2의 first_frame_mask expand 패턴은 Cosmos에 없음

계획 문서 잘못된 경로 정정 (Task 4–16 전체 적용 필요)

계획 문서 경로 (틀림)	실제 경로 (정확)
`flow_factory.models.base_adapter`	`flow_factory.models.abc`
`flow_factory.samples.base.BaseSample`	`flow_factory.samples.BaseSample`
`flow_factory.models.wan21_t2v`	`flow_factory.models.wan.wan2_i2v` (I2V) / `wan.wan2_t2v` (T2V)

4Takeaway

최소 구현 계약 확정: Cosmos adapter는 load_pipeline, decode_latents, forward, inference 4개만 구현하면 FF 컴파일 통과. encode_prompt/encode_image는 필요 시 추가.

경로 오류 조기 발견: 계획 문서의 잘못된 import 3개를 코드 작성 전에 정정. Task 4–16 구현 시 flow_factory.models.abc, flow_factory.samples, flow_factory.models.wan.wan2_i2v를 기준으로 사용해야 한다.

구현 전략 결론

Wan2_I2V_Adapter를 직접 레퍼런스로 사용하되, Cosmos 전용 차이점(pseudo-pipeline, x0→v 변환, first_frame_mask 없음, T5-XXL 텍스트 인코더)만 대체하는 전략이 최소 리스크. encode_image는 VAE encode 경로로 단순화 가능 — Wan2.2의 image_encoder 복잡도 불필요.

5Next Steps

Torch 업그레이드 필수 (Task 4 이전): B200 sm_100은 torch 2.6.0+cu124에 없음. nightly / 2.7+ 설치 확인 후 adapter 구현 착수.
pipe.tokenizer 속성명 검증: Cosmos Video2WorldPipeline에서 .tokenizer가 text tokenizer인지 VAE tokenizer인지 실제 소스 확인 필요. 현재 매핑 표에서 VAE는 pipe.vae로 추정했으나 검증 전.
DiT.forward 시그니처 직접 확인: positional vs kwargs, fps/padding_mask 처리 방식 검증 필요 (Task 5 진입 전).
FF submodule pyproject 패치: transformers >= 4.51.3 미적용 상태, uncommitted. fork-branch 커밋 또는 patches/ff-transformers.patch로 관리 방식 결정.
CosmosPredict2Sample 설계: _shared_fields에 first_frame_latent 포함 여부 결정 필요 — Wan2_I2V_Adapter의 _shared_fields 패턴 참고.