improve text encoding

2025-08-30 22:12:56 -07:00
parent 055e52e5ea
commit 642b89edca
2 changed files with 323 additions and 51 deletions
@@ -59,58 +59,37 @@ class TextEncodingStage(PipelineStage):
        assert len(self.text_encoders) == len(
            fastvideo_args.pipeline_config.text_encoder_configs)

-        for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
-                self.tokenizers,
-                self.text_encoders,
-                fastvideo_args.pipeline_config.text_encoder_configs,
-                fastvideo_args.pipeline_config.preprocess_text_funcs,
-                fastvideo_args.pipeline_config.postprocess_text_funcs,
-                strict=True):
+        # Encode positive prompt with all available encoders
+        assert batch.prompt is not None
+        prompt_text: str | list[str] = batch.prompt
+        all_indices: list[int] = list(range(len(self.text_encoders)))
+        prompt_embeds_list, prompt_masks_list = self.encode_text(
+            prompt_text,
+            fastvideo_args,
+            encoder_index=all_indices,
+            return_attention_mask=True,
+        )
+        for pe in prompt_embeds_list:
+            batch.prompt_embeds.append(pe)
+        if batch.prompt_attention_mask is not None:
+            for am in prompt_masks_list:
+                batch.prompt_attention_mask.append(am)

-            assert isinstance(batch.prompt, str | list)
-            if isinstance(batch.prompt, str):
-                batch.prompt = [batch.prompt]
-            texts = []
-            for prompt_str in batch.prompt:
-                texts.append(preprocess_func(prompt_str))
-            text_inputs = tokenizer(texts,
-                                    **encoder_config.tokenizer_kwargs).to(
-                                        get_local_torch_device())
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs = text_encoder(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-            prompt_embeds = postprocess_func(outputs)
-            batch.prompt_embeds.append(prompt_embeds)
-            if batch.prompt_attention_mask is not None:
-                batch.prompt_attention_mask.append(attention_mask)
-
-            if batch.do_classifier_free_guidance:
-                assert isinstance(batch.negative_prompt, str)
-                negative_text = preprocess_func(batch.negative_prompt)
-                negative_text_inputs = tokenizer(
-                    negative_text, **encoder_config.tokenizer_kwargs).to(
-                        get_local_torch_device())
-                negative_input_ids = negative_text_inputs["input_ids"]
-                negative_attention_mask = negative_text_inputs["attention_mask"]
-                with set_forward_context(current_timestep=0,
-                                         attn_metadata=None):
-                    negative_outputs = text_encoder(
-                        input_ids=negative_input_ids,
-                        attention_mask=negative_attention_mask,
-                        output_hidden_states=True,
-                    )
-                negative_prompt_embeds = postprocess_func(negative_outputs)
-
-                assert batch.negative_prompt_embeds is not None
-                batch.negative_prompt_embeds.append(negative_prompt_embeds)
-                if batch.negative_attention_mask is not None:
-                    batch.negative_attention_mask.append(
-                        negative_attention_mask)
+        # Encode negative prompt if CFG is enabled
+        if batch.do_classifier_free_guidance:
+            assert isinstance(batch.negative_prompt, str)
+            neg_embeds_list, neg_masks_list = self.encode_text(
+                batch.negative_prompt,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
+            assert batch.negative_prompt_embeds is not None
+            for ne in neg_embeds_list:
+                batch.negative_prompt_embeds.append(ne)
+            if batch.negative_attention_mask is not None:
+                for nm in neg_masks_list:
+                    batch.negative_attention_mask.append(nm)

        return batch

@@ -129,6 +108,171 @@ class TextEncodingStage(PipelineStage):
                         V.none_or_list)
        return result

+    @torch.no_grad()
+    def encode_text(
+        self,
+        text: str | list[str],
+        fastvideo_args: FastVideoArgs,
+        encoder_index: int | list[int] | None = None,
+        return_attention_mask: bool = False,
+        return_type: str = "list",  # one of: "list", "dict", "stack"
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+        max_length: int | None = None,
+        truncation: bool | None = None,
+        padding: bool | str | None = None,
+    ):
+        """
+        Encode plain text using selected text encoder(s) and return embeddings.
+
+        Args:
+            text: A single string or a list of strings to encode.
+            fastvideo_args: The inference arguments providing pipeline config,
+                including tokenizer and encoder settings, preprocess and postprocess
+                functions.
+            encoder_index: Encoder selector by index. Accepts an int or list of ints.
+            return_attention_mask: If True, also return attention masks for each
+                selected encoder.
+            return_type: "list" (default) returns a list aligned with selection;
+                "dict" returns a dict keyed by encoder index as a string; "stack" stacks along a
+                new first dimension (requires matching shapes).
+            device: Optional device override for inputs; defaults to local torch device.
+            dtype: Optional dtype to cast returned embeddings to.
+            max_length: Optional per-call tokenizer override.
+            truncation: Optional per-call tokenizer override.
+            padding: Optional per-call tokenizer override.
+
+        Returns:
+            Depending on return_type and return_attention_mask:
+            - list: List[Tensor] or (List[Tensor], List[Tensor])
+            - dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])
+            - stack: Tensor of shape [num_encoders, ...] or a tuple with stacked
+              attention masks
+        """
+
+        assert len(self.tokenizers) == len(self.text_encoders)
+        assert len(self.text_encoders) == len(
+            fastvideo_args.pipeline_config.text_encoder_configs)
+
+        # Resolve selection into indices
+        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
+        if encoder_index is None:
+            indices: list[int] = [0]
+        elif isinstance(encoder_index, int):
+            indices = [encoder_index]
+        else:
+            indices = list(encoder_index)
+        # validate range
+        num_encoders = len(self.text_encoders)
+        for idx in indices:
+            if idx < 0 or idx >= num_encoders:
+                raise IndexError(
+                    f"encoder index {idx} out of range [0, {num_encoders-1}]")
+
+        # Validate indices are within range
+        num_encoders = len(self.text_encoders)
+
+        # Normalize input to list[str]
+        assert isinstance(text, str | list)
+        if isinstance(text, str):
+            texts: list[str] = [text]
+        else:
+            texts = text
+
+        embeds_list: list[torch.Tensor] = []
+        attn_masks_list: list[torch.Tensor] = []
+
+        preprocess_funcs = fastvideo_args.pipeline_config.preprocess_text_funcs
+        postprocess_funcs = fastvideo_args.pipeline_config.postprocess_text_funcs
+        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
+
+        if return_type not in ("list", "dict", "stack"):
+            raise ValueError(
+                f"Invalid return_type '{return_type}'. Expected one of: 'list', 'dict', 'stack'"
+            )
+
+        target_device = device if device is not None else get_local_torch_device(
+        )
+
+        for i in indices:
+            tokenizer = self.tokenizers[i]
+            text_encoder = self.text_encoders[i]
+            encoder_config = encoder_cfgs[i]
+            preprocess_func = preprocess_funcs[i]
+            postprocess_func = postprocess_funcs[i]
+
+            processed_texts: list[str] = []
+            for prompt_str in texts:
+                processed_texts.append(preprocess_func(prompt_str))
+
+            tok_kwargs = dict(encoder_config.tokenizer_kwargs)
+            if max_length is not None:
+                tok_kwargs["max_length"] = max_length
+            if truncation is not None:
+                tok_kwargs["truncation"] = truncation
+            if padding is not None:
+                tok_kwargs["padding"] = padding
+
+            text_inputs = tokenizer(processed_texts,
+                                    **tok_kwargs).to(target_device)
+
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs = text_encoder(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+
+            prompt_embeds = postprocess_func(outputs)
+            if dtype is not None:
+                prompt_embeds = prompt_embeds.to(dtype=dtype)
+            embeds_list.append(prompt_embeds)
+            if return_attention_mask:
+                attn_masks_list.append(attention_mask)
+
+        # Shape results according to return_type
+        if return_type == "list":
+            if return_attention_mask:
+                return embeds_list, attn_masks_list
+            return embeds_list
+
+        if return_type == "dict":
+            key_strs = [str(i) for i in indices]
+            embeds_dict = {
+                k: v
+                for k, v in zip(key_strs, embeds_list, strict=False)
+            }
+            if return_attention_mask:
+                attn_dict = {
+                    k: v
+                    for k, v in zip(key_strs, attn_masks_list, strict=False)
+                }
+                return embeds_dict, attn_dict
+            return embeds_dict
+
+        # return_type == "stack"
+        # Validate shapes are compatible
+        base_shape = list(embeds_list[0].shape)
+        for t in embeds_list[1:]:
+            if list(t.shape) != base_shape:
+                raise ValueError(
+                    f"Cannot stack embeddings with differing shapes: {[list(t.shape) for t in embeds_list]}"
+                )
+        stacked_embeds = torch.stack(embeds_list, dim=0)
+        if return_attention_mask:
+            base_mask_shape = list(attn_masks_list[0].shape)
+            for m in attn_masks_list[1:]:
+                if list(m.shape) != base_mask_shape:
+                    raise ValueError(
+                        f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}"
+                    )
+            stacked_masks = torch.stack(attn_masks_list, dim=0)
+            return stacked_embeds, stacked_masks
+        return stacked_embeds
+
    def verify_output(self, batch: ForwardBatch,
                      fastvideo_args: FastVideoArgs) -> VerificationResult:
        """Verify text encoding stage outputs."""
@@ -0,0 +1,128 @@
+import torch
+import types
+import pytest
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, BaseEncoderOutput
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+class TensorDict(dict):
+    def to(self, device):
+        return TensorDict({k: v.to(device) for k, v in self.items()})
+
+class FakeTokenizer:
+    def __call__(self, texts, **kwargs):
+        B = len(texts)
+        seq_len = int(kwargs.get("max_length", 4))
+        return TensorDict({
+            "input_ids": torch.arange(B * seq_len).view(B, seq_len),
+            "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
+        })
+
+class FakeTextEncoder(torch.nn.Module):
+    def __init__(self, hidden_size=8):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, input_ids, attention_mask, output_hidden_states=True):
+        B, T = input_ids.shape
+        last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
+        return types.SimpleNamespace(last_hidden_state=last_hidden_state)
+
+def id_preprocess(x: str) -> str:
+    return x
+
+def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    # [B, T, H] -> [B, H]
+    return outputs.last_hidden_state.mean(dim=1)
+
+def make_args(num_encoders=2, text_len=4, hidden_size=8):
+    enc_cfgs = []
+    preprocess_fns = []
+    postprocess_fns = []
+    for _ in range(num_encoders):
+        arch = TextEncoderArchConfig(text_len=text_len)
+        enc_cfgs.append(TextEncoderConfig(arch_config=arch))
+        preprocess_fns.append(id_preprocess)
+        postprocess_fns.append(take_mean_postprocess)
+    pipe_cfg = PipelineConfig(
+        text_encoder_configs=tuple(enc_cfgs),
+        text_encoder_precisions=tuple(["fp32"] * num_encoders),
+        preprocess_text_funcs=tuple(preprocess_fns),
+        postprocess_text_funcs=tuple(postprocess_fns),
+    )
+    return FastVideoArgs(model_path="", pipeline_config=pipe_cfg), hidden_size
+
+def make_stage(num_encoders=2, hidden_size=8):
+    tokenizers = [FakeTokenizer() for _ in range(num_encoders)]
+    encoders = [FakeTextEncoder(hidden_size=hidden_size) for _ in range(num_encoders)]
+    return TextEncodingStage(text_encoders=encoders, tokenizers=tokenizers)
+
+def test_encode_text_selection_and_shapes():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # list return, two encoders
+    embeds = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1])
+    assert isinstance(embeds, list) and len(embeds) == 2
+    for e in embeds:
+        assert e.shape == (2, hidden)
+
+    # with masks
+    embeds2, masks2 = stage.encode_text("a", fastvideo_args, encoder_index=[1], return_attention_mask=True)
+    assert len(embeds2) == 1 and len(masks2) == 1
+    assert embeds2[0].shape == (1, hidden)
+    assert masks2[0].shape == (1, 4)
+
+    # dict return
+    d = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="dict")
+    assert set(d.keys()) == {"0", "1"}
+    assert d["0"].shape == (2, hidden)
+
+    # stack return
+    s = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="stack")
+    assert s.shape == (2, 2, hidden)  # [encoders, batch, hidden]
+
+    # overrides: dtype + max_length
+    e3, m3 = stage.encode_text(["a"], fastvideo_args, encoder_index=[0], dtype=torch.float16, return_attention_mask=True, max_length=3)
+    assert e3[0].dtype == torch.float16
+    assert m3[0].shape[1] == 3
+
+def test_forward_integration_cfg_off_and_on():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # CFG off
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="a cat",
+        negative_prompt="",
+        do_classifier_free_guidance=False,
+        prompt_embeds=[],
+        negative_prompt_embeds=None,
+        prompt_attention_mask=[],
+        negative_attention_mask=None,
+    )
+    out = stage.forward(batch, fastvideo_args)
+    assert len(out.prompt_embeds) == 2
+    for e in out.prompt_embeds:
+        assert e.shape[1] == hidden
+
+    # CFG on
+    batch2 = ForwardBatch(
+        data_type="video",
+        prompt=["a cat", "a dog"],
+        negative_prompt="bad picture",
+        do_classifier_free_guidance=True,
+        prompt_embeds=[],
+        negative_prompt_embeds=[],
+        prompt_attention_mask=[],
+        negative_attention_mask=[],
+    )
+    out2 = stage.forward(batch2, fastvideo_args)
+    assert len(out2.prompt_embeds) == 2
+    assert len(out2.negative_prompt_embeds) == 2
+    assert len(out2.prompt_attention_mask) == 2
+    assert len(out2.negative_attention_mask) == 2