30 lines
892 B
Python
30 lines
892 B
Python
from typing import List
|
|
|
|
import torch
|
|
|
|
from toolkit.models.v2.text_encoders.umt5 import (
|
|
PatchedT5Tokenizer,
|
|
UMT5TextEncoder,
|
|
)
|
|
|
|
|
|
def get_umt5_encoder(
|
|
model_path: str,
|
|
tokenizer_subfolder: str = None,
|
|
encoder_subfolder: str = None,
|
|
torch_dtype: str = torch.bfloat16,
|
|
# reserved for the comfy-weights flip (Phase 2); accepted for
|
|
# signature compatibility, not consulted yet
|
|
comfy_files: List[str] = None,
|
|
):
|
|
"""Load the UMT5 tokenizer + encoder. Thin compatibility wrapper around
|
|
toolkit/models/v2/text_encoders/umt5.py."""
|
|
tokenizer = UMT5TextEncoder.load_tokenizer(
|
|
model_path, subfolder=tokenizer_subfolder or ""
|
|
)
|
|
print(f"Using {model_path} for UMT5 encoder.")
|
|
text_encoder = UMT5TextEncoder.load_model(
|
|
model_path, dtype=torch_dtype, subfolder=encoder_subfolder or ""
|
|
)
|
|
return tokenizer, text_encoder
|