125 lines
5.0 KiB
Python
125 lines
5.0 KiB
Python
#
|
|
# Text pipeline ported from Apple ml-starflow utils/common.py::preprocess_text
|
|
# and utils/common.py::encode_text (commit 20984834c1a68ee679faa8ed9cda2842f02b042d),
|
|
# restricted to the single text encoder used by STARFlow-V (google/flan-t5-xl,
|
|
# a plain HF encoder-only model, no chat-template / Gemma / label-embedding path).
|
|
# For licensing see accompanying LICENSE file.
|
|
# Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
#
|
|
"""
|
|
Prompt wrapping and FLAN-T5-XL text encoding for STARFlow-V.
|
|
|
|
Exact upstream wrapping order (verified against utils/common.py):
|
|
1. fps wrap: "A video with {fps} fps:\n{prompt}\n" (only if fps is not None)
|
|
2. noise prepend: "Noise Level {int(noise_std*1000)}:\n{wrapped}"
|
|
3. aspect append: "{wrapped}\n in a {aspect_ratio} aspect ratio.\n"
|
|
|
|
The unconditional ("") caption used for classifier-free guidance goes through
|
|
the exact same wrapping (it receives the same fps/noise_std/aspect_ratio
|
|
metadata as the conditional caption), matching upstream sample.py which
|
|
appends "" to the caption list *before* calling encode_text.
|
|
"""
|
|
import os
|
|
from typing import List, Optional, Union
|
|
|
|
import torch
|
|
|
|
|
|
def preprocess_text(
|
|
text: List[str],
|
|
aspect_ratio: Optional[str] = None,
|
|
fps: Optional[Union[int, torch.Tensor]] = None,
|
|
noise_std: Optional[Union[float, torch.Tensor]] = None,
|
|
) -> List[str]:
|
|
if fps is not None:
|
|
if isinstance(fps, torch.Tensor):
|
|
fps = [int(f) for f in fps.tolist()]
|
|
elif isinstance(fps, int):
|
|
fps = [fps] * len(text)
|
|
text = [f"A video with {f} fps:\n{txt}\n" if f > 0 else f"An image:\n{txt}\n"
|
|
for txt, f in zip(text, fps)]
|
|
|
|
if noise_std is not None:
|
|
if isinstance(noise_std, torch.Tensor):
|
|
noise_std = [int(n * 1000) for n in noise_std.view(-1).tolist()]
|
|
elif isinstance(noise_std, float):
|
|
noise_std = [int(noise_std * 1000)] * len(text)
|
|
text = [f'Noise Level {n}:\n{txt}' for n, txt in zip(noise_std, text)]
|
|
|
|
if aspect_ratio is not None:
|
|
text = [f"{txt}\n in a {aspect_ratio} aspect ratio.\n" for txt in text]
|
|
|
|
return text
|
|
|
|
|
|
@torch.no_grad()
|
|
def encode_text(text_encoder, tokenizer, text: List[str], max_length: int, device,
|
|
return_tokens: bool = False, **kwargs):
|
|
text = preprocess_text(text, **kwargs)
|
|
tokenized_outputs = tokenizer(
|
|
text, padding="max_length", truncation=True, return_tensors="pt", max_length=max_length)
|
|
tokenized_outputs = {key: val.to(device) for key, val in tokenized_outputs.items()}
|
|
y = text_encoder(**tokenized_outputs).last_hidden_state
|
|
y = y * tokenized_outputs['attention_mask'].unsqueeze(-1) # mask out padding
|
|
if return_tokens:
|
|
return y, tokenized_outputs
|
|
return y
|
|
|
|
|
|
class STARFlowVTextEncoder:
|
|
"""
|
|
Loads a local FLAN-T5-XL directory (tokenizer + encoder-only T5 weights)
|
|
with no network access. Kept dependency-light: only transformers is
|
|
required (no sentencepiece network fetch, no accelerate device_map
|
|
auto-download).
|
|
"""
|
|
|
|
hidden_size = 2048 # google/flan-t5-xl d_model
|
|
|
|
def __init__(self, model_dir: str, dtype: torch.dtype = torch.bfloat16):
|
|
if not os.path.isdir(model_dir):
|
|
raise FileNotFoundError(
|
|
f"STARFlow-V text encoder directory not found: {model_dir}. "
|
|
f"Provide a local FLAN-T5-XL checkout (config.json, tokenizer files, "
|
|
f"and model weights); this loader never fetches from the network."
|
|
)
|
|
from transformers import AutoTokenizer, T5EncoderModel, T5Config
|
|
|
|
self.model_dir = model_dir
|
|
self.dtype = dtype
|
|
|
|
config = T5Config.from_pretrained(model_dir, local_files_only=True)
|
|
if config.d_model != self.hidden_size:
|
|
raise ValueError(
|
|
f"Expected FLAN-T5-XL (d_model={self.hidden_size}), got d_model={config.d_model} "
|
|
f"from {model_dir}. STARFlow-V's checkpoint proj_txt layers require this exact size."
|
|
)
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
|
|
self.model = T5EncoderModel.from_pretrained(
|
|
model_dir, config=config, torch_dtype=dtype, local_files_only=True)
|
|
self.model.requires_grad_(False)
|
|
self.model.eval()
|
|
|
|
self.device = torch.device("cpu")
|
|
self.load_device = torch.device("cpu")
|
|
|
|
def to(self, device):
|
|
self.model.to(device)
|
|
self.device = torch.device(device)
|
|
return self
|
|
|
|
def model_size(self) -> int:
|
|
return sum(p.numel() * p.element_size() for p in self.model.parameters())
|
|
|
|
@torch.no_grad()
|
|
def encode(self, prompts: List[str], txt_size: int, device,
|
|
aspect_ratio: Optional[str] = None,
|
|
fps: Optional[int] = None,
|
|
noise_std: Optional[float] = None) -> torch.Tensor:
|
|
self.model.to(device)
|
|
return encode_text(
|
|
self.model, self.tokenizer, prompts, txt_size, device,
|
|
aspect_ratio=aspect_ratio, fps=fps, noise_std=noise_std,
|
|
)
|