[feat] v2 torch backend: CLIP image-encoder adapter + image_encoder component kind (i2v groundwork)

Adds CLIPImageEncoder (encode_image -> the DiT's encoder_hidden_states_image) + the generic builder's
image_encoder maker (ImageEncoderLoader + ImageProcessorLoader), registered as the cuda 'image_encoder'
kind. Mirrors fastvideo's ImageEncodingStage. CPU-verified (component-kinds + lazy invariant, 211/2);
GPU path marked BRINGUP/written-not-run (processor subfolder + dtype to confirm on a real i2v checkpoint),
matching how the rest of the torch backend was originally landed. Reusable by the Wan i2v cluster + many
bucket-C models (Hunyuan/Cosmos i2v). Next i2v increments: the mask+cond latent construction + the
concat-into-DiT-input loop, then the card + registry + GPU-verify with Wan2.1-Fun-1.3B-InP.
This commit is contained in:
Will Lin
2026-06-18 11:35:18 +00:00
parent 060bdd85b6
commit 60ca94796c
3 changed files with 32 additions and 4 deletions
+28 -1
View File
@@ -367,6 +367,25 @@ class Gemma(TorchComponent):
return _to_numpy(video.squeeze(0)), _to_numpy(audio.squeeze(0))
class CLIPImageEncoder(TorchComponent):
"""CLIP-vision image encoder for Wan i2v: ``encode_image(image) -> image_embeds`` (the DiT's
``encoder_hidden_states_image``). Mirrors fastvideo's ImageEncodingStage — the HF image processor
preprocesses, the encoder returns ``last_hidden_state``. BRINGUP: written-not-run; GPU-verify the
processor/encoder subfolders + dtype against a real i2v checkpoint (see GPU_BRINGUP.md)."""
def __init__(self, module, processor, *, device, dtype):
super().__init__(module, device=device, dtype=dtype)
self.processor = processor
@torch.no_grad()
def encode_image(self, image):
inputs = self.processor(images=image, return_tensors="pt").to(self.device)
with self._ctx():
out = self.module(**inputs)
embeds = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
return _to_numpy(embeds.squeeze(0))
# --------------------------------------------------------------------------- #
# Upsampler + audio (LTX-2) #
# --------------------------------------------------------------------------- #
@@ -445,6 +464,13 @@ def _make_text_encoder(spec, instance, platform, args):
return cls(module, tokenizer, device=device, dtype=dtype)
def _make_image_encoder(spec, instance, platform, args):
module = load_component("ImageEncoderLoader", spec.checkpoint, args) # CLIP vision
# the HF image processor is a sibling subfolder (BRINGUP: confirm path on a real i2v checkpoint)
processor = load_component("ImageProcessorLoader", os.path.join(_model_root(spec), "image_processor"), args)
return CLIPImageEncoder(module, processor, device=_device(platform), dtype=_native_dtype(module))
def _make_upsampler(spec, instance, platform, args):
module = load_component("UpsamplerLoader", spec.checkpoint, args) # LTX2LatentUpsampler
vae = instance.component("vae")
@@ -468,7 +494,8 @@ def _make_vocoder(spec, instance, platform, args):
_MAKERS = {
"dit": _make_dit, "vae": _make_vae, "text_encoder": _make_text_encoder,
"upsampler": _make_upsampler, "audio_vae": _make_audio_vae, "vocoder": _make_vocoder,
"image_encoder": _make_image_encoder, "upsampler": _make_upsampler,
"audio_vae": _make_audio_vae, "vocoder": _make_vocoder,
}
+2 -1
View File
@@ -51,7 +51,8 @@ def _flow_sde_cuda(*args, **kwargs):
_B = "v2.platform.backends.torch_backend"
_KIND_SOURCE = {
"dit": "WanDiT/LTX2DiT", "vae": "WanVAE/LTX2VAE", "text_encoder": "T5Encoder/Gemma",
"upsampler": "LTX2Upsampler", "audio_vae": "LTX2AudioVAE", "vocoder": "LTX2Vocoder",
"image_encoder": "CLIPImageEncoder", "upsampler": "LTX2Upsampler",
"audio_vae": "LTX2AudioVAE", "vocoder": "LTX2Vocoder",
}
for _kind, _cls in _KIND_SOURCE.items():
register_component(_kind, _build_component, device="cuda", available=_cuda_available,
+2 -2
View File
@@ -42,8 +42,8 @@ def _make_cuda_available(monkeypatch):
# --------------------------------------------------------------------------- #
def test_cuda_has_all_components_and_two_solver_ops():
comps = sorted(r["kind"] for r in component_matrix() if r["device"] == "cuda")
# + upsampler (LTX-2 spatial), audio_vae + vocoder (LTX-2.3 T2VS audio branch)
assert comps == ["audio_vae", "dit", "text_encoder", "upsampler", "vae", "vocoder"]
# + image_encoder (i2v CLIP), upsampler, audio_vae + vocoder (LTX-2.3 T2VS audio branch)
assert comps == ["audio_vae", "dit", "image_encoder", "text_encoder", "upsampler", "vae", "vocoder"]
ops = sorted(r["op"] for r in kernel_matrix() if r["device"] == "cuda")
assert ops == ["flow_match_step", "flow_sde_step"]