[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:
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user