diff --git a/v2/platform/backends/torch_backend.py b/v2/platform/backends/torch_backend.py index a3dd9153c..105cf6427 100644 --- a/v2/platform/backends/torch_backend.py +++ b/v2/platform/backends/torch_backend.py @@ -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, } diff --git a/v2/platform/backends/torch_cuda.py b/v2/platform/backends/torch_cuda.py index 7681cfacb..5492fc33c 100644 --- a/v2/platform/backends/torch_cuda.py +++ b/v2/platform/backends/torch_cuda.py @@ -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, diff --git a/v2/tests/test_torch_backend.py b/v2/tests/test_torch_backend.py index 2be1fec94..f8b724daa 100644 --- a/v2/tests/test_torch_backend.py +++ b/v2/tests/test_torch_backend.py @@ -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"]