From 642b89edca386cea86c832e631de0936727aeead Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 30 Aug 2025 22:12:56 -0700 Subject: [PATCH] improve text encoding --- fastvideo/pipelines/stages/text_encoding.py | 246 +++++++++++++++---- fastvideo/tests/stages/test_text_encoding.py | 128 ++++++++++ 2 files changed, 323 insertions(+), 51 deletions(-) create mode 100644 fastvideo/tests/stages/test_text_encoding.py diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 5a25236da..4353f1b26 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -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.""" diff --git a/fastvideo/tests/stages/test_text_encoding.py b/fastvideo/tests/stages/test_text_encoding.py new file mode 100644 index 000000000..51bd1ca77 --- /dev/null +++ b/fastvideo/tests/stages/test_text_encoding.py @@ -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