From abec2c27c5adc4db3f726988c61ca3dc24c45470 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 00:43:04 +0000 Subject: [PATCH 1/3] feat(diffusion): add encoder_hub for family-specific frozen-encoder logic Review feedback on PR #90: TrainPipelineConfig should only drive the training backend, so frozen-encoder loading/encoding lives in its own rollout-side hub, dispatched by the resolved diffusion model family. Wan2.2 provides UMT5+VAE loading (from the explicit --sft-encoder-checkpoint introduced by the consumer PR), sample encoding, and its 4k+1 frame-count constraint. No callers in this PR; the SFT PR (#90) is stacked on top and wires argument validation and the encode actor pool to this hub. Co-Authored-By: Claude Fable 5 --- miles/rollout/encoder_hub/__init__.py | 17 ++++++++ miles/rollout/encoder_hub/wan2_2.py | 58 +++++++++++++++++++++++++++ 2 files changed, 75 insertions(+) create mode 100644 miles/rollout/encoder_hub/__init__.py create mode 100644 miles/rollout/encoder_hub/wan2_2.py diff --git a/miles/rollout/encoder_hub/__init__.py b/miles/rollout/encoder_hub/__init__.py new file mode 100644 index 00000000..4e04b245 --- /dev/null +++ b/miles/rollout/encoder_hub/__init__.py @@ -0,0 +1,17 @@ +"""Frozen-encoder logic per model family, decoupled from the training-side TrainPipelineConfig. + +Each family module provides: +- ``load_encoder(args, device)``: load the frozen encode components (tokenizer/text + encoder/VAE) from the ``--sft-encoder-checkpoint`` HF name or path; +- ``encode_sample(encoder, pixels, prompt, generator)``: encode one media/prompt pair + into a cached train sample (clean latent + cond kwargs); +- ``validate_args(args)``: family-specific encode constraints. +""" + + +def get_encoder(family: str | None): + if family == "wan2_2": + from miles.rollout.encoder_hub import wan2_2 + + return wan2_2 + raise ValueError(f"no encoder_hub entry for model family {family!r}") diff --git a/miles/rollout/encoder_hub/wan2_2.py b/miles/rollout/encoder_hub/wan2_2.py new file mode 100644 index 00000000..c8b487ac --- /dev/null +++ b/miles/rollout/encoder_hub/wan2_2.py @@ -0,0 +1,58 @@ +"""Wan2.2 frozen encoders (UMT5 text encoder + VAE) for offline SFT encoding.""" + +from __future__ import annotations + +import torch + + +def validate_args(args) -> None: + if (args.diffusion_output_num_frames - 1) % 4 != 0: + raise ValueError("--diffusion-output-num-frames must be 4k+1 for the Wan VAE temporal stride") + + +def load_encoder(args, device: torch.device) -> dict: + from diffusers import AutoencoderKLWan + from transformers import AutoTokenizer, UMT5EncoderModel + + ckpt = args.sft_encoder_checkpoint + tokenizer = AutoTokenizer.from_pretrained(ckpt, subfolder="tokenizer") + text_encoder = UMT5EncoderModel.from_pretrained(ckpt, subfolder="text_encoder", torch_dtype=torch.bfloat16).to( + device + ) + vae = AutoencoderKLWan.from_pretrained(ckpt, subfolder="vae", torch_dtype=torch.float32).to(device) + view = (1, vae.config.z_dim, 1, 1, 1) + return { + "device": device, + "tokenizer": tokenizer, + "text_encoder": text_encoder, + "vae": vae, + "latents_mean": torch.tensor(vae.config.latents_mean).view(view).to(device), + "latents_std": torch.tensor(vae.config.latents_std).view(view).to(device), + } + + +@torch.no_grad() +def encode_sample(encoder: dict, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict: + from diffusers.pipelines.wan.pipeline_wan import prompt_clean + + device = encoder["device"] + latent = encoder["vae"].encode(pixels.unsqueeze(0).to(device, torch.float32)).latent_dist.sample(generator) + latent = (latent - encoder["latents_mean"]) / encoder["latents_std"] + + inputs = encoder["tokenizer"]( + [prompt_clean(prompt)], + padding="max_length", + max_length=512, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + embeds = encoder["text_encoder"](inputs.input_ids.to(device), inputs.attention_mask.to(device)).last_hidden_state + embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 + + return { + "latent": latent[0].to(torch.float16).cpu(), + "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, + "prompt": prompt, + } From 2a3c9ced672bc36fddf1be2e8e4a2b54eb1808fd Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 06:03:58 +0000 Subject: [PATCH 2/3] docs(encoder_hub): record measured rationale for cache storage dtypes Co-Authored-By: Claude Fable 5 --- miles/rollout/encoder_hub/wan2_2.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/miles/rollout/encoder_hub/wan2_2.py b/miles/rollout/encoder_hub/wan2_2.py index c8b487ac..096de2cf 100644 --- a/miles/rollout/encoder_hub/wan2_2.py +++ b/miles/rollout/encoder_hub/wan2_2.py @@ -51,6 +51,10 @@ def encode_sample(encoder: dict, pixels: torch.Tensor, prompt: str, generator: t embeds = encoder["text_encoder"](inputs.input_ids.to(device), inputs.attention_mask.to(device)).last_hidden_state embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 + # Storage dtypes: the normalized latent is unit-scale (std ~0.8, max|x| ~4 on real + # clips), where fp16's 10-bit mantissa stores ~8x tighter than bf16 (measured rms + # rel err 2e-4 vs 1.7e-3) with no range risk; the embeds are a bf16 model's output, + # so bf16 keeps them bit-exact. The trainer upcasts both before use. return { "latent": latent[0].to(torch.float16).cpu(), "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, From 28a93be3e42ed6ab0acad161a596224115a08ffc Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 6 Aug 2026 06:11:24 +0000 Subject: [PATCH 3/3] fix(encoder_hub): align Wan encode dtypes with sglang-d MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Match the rollout engine's Wan pipeline config (vae_precision fp32, text_encoder_precisions fp32, DiT bf16): UMT5 now computes in fp32, and both cached tensors store bf16 — the precision the DiT boundary sees on the rollout path. Co-Authored-By: Claude Fable 5 --- miles/rollout/encoder_hub/wan2_2.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/miles/rollout/encoder_hub/wan2_2.py b/miles/rollout/encoder_hub/wan2_2.py index 096de2cf..6a587c89 100644 --- a/miles/rollout/encoder_hub/wan2_2.py +++ b/miles/rollout/encoder_hub/wan2_2.py @@ -16,7 +16,7 @@ def load_encoder(args, device: torch.device) -> dict: ckpt = args.sft_encoder_checkpoint tokenizer = AutoTokenizer.from_pretrained(ckpt, subfolder="tokenizer") - text_encoder = UMT5EncoderModel.from_pretrained(ckpt, subfolder="text_encoder", torch_dtype=torch.bfloat16).to( + text_encoder = UMT5EncoderModel.from_pretrained(ckpt, subfolder="text_encoder", torch_dtype=torch.float32).to( device ) vae = AutoencoderKLWan.from_pretrained(ckpt, subfolder="vae", torch_dtype=torch.float32).to(device) @@ -51,12 +51,8 @@ def encode_sample(encoder: dict, pixels: torch.Tensor, prompt: str, generator: t embeds = encoder["text_encoder"](inputs.input_ids.to(device), inputs.attention_mask.to(device)).last_hidden_state embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 - # Storage dtypes: the normalized latent is unit-scale (std ~0.8, max|x| ~4 on real - # clips), where fp16's 10-bit mantissa stores ~8x tighter than bf16 (measured rms - # rel err 2e-4 vs 1.7e-3) with no range risk; the embeds are a bf16 model's output, - # so bf16 keeps them bit-exact. The trainer upcasts both before use. return { - "latent": latent[0].to(torch.float16).cpu(), + "latent": latent[0].to(torch.bfloat16).cpu(), "cond_kwargs": {"encoder_hidden_states": embeds.to(torch.bfloat16).cpu()}, "prompt": prompt, }