77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import folder_paths
|
|
|
|
from .loader import load_t5
|
|
from ..utils.dtype import string_to_dtype, dtype_list
|
|
|
|
# initialize custom folder path
|
|
# TODO: integrate with `extra_model_paths.yaml`
|
|
os.makedirs(
|
|
os.path.join(folder_paths.models_dir,"t5"),
|
|
exist_ok = True,
|
|
)
|
|
folder_paths.folder_names_and_paths["t5"] = (
|
|
[os.path.join(folder_paths.models_dir,"t5")],
|
|
folder_paths.supported_pt_extensions
|
|
)
|
|
|
|
class T5v11Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"t5v11_name": (folder_paths.get_filename_list("t5"),),
|
|
"t5v11_ver": (["xxl"],),
|
|
"path_type": (["folder", "file"],),
|
|
"device": (["auto", "cpu", "gpu"],{"default":"cpu"}),
|
|
"dtype": (dtype_list,),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("T5",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "ExtraModels/T5"
|
|
TITLE = "T5v1.1 Loader"
|
|
|
|
def load_model(self, t5v11_name, t5v11_ver, path_type, device, dtype):
|
|
if "bnb" in dtype:
|
|
assert device == "gpu", "BitsAndBytes only works on CUDA! Set device to 'gpu'."
|
|
dtype = string_to_dtype(dtype, "text_encoder")
|
|
if device == "cpu":
|
|
assert dtype in [None, torch.float32], f"Can't use dtype '{dtype}' with CPU! Set dtype to 'default'."
|
|
|
|
return (load_t5(
|
|
model_type = "t5v11",
|
|
model_ver = t5v11_ver,
|
|
model_path = folder_paths.get_full_path("t5", t5v11_name),
|
|
path_type = path_type,
|
|
device = device,
|
|
dtype = dtype,
|
|
),)
|
|
|
|
class T5TextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True}),
|
|
"T5": ("T5",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "ExtraModels/T5"
|
|
TITLE = "T5 Text Encode"
|
|
|
|
def encode(self, text, T5=None):
|
|
tokens = T5.tokenize(text)
|
|
cond = T5.encode_from_tokens(tokens)
|
|
return ([[cond, {}]], )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"T5v11Loader" : T5v11Loader,
|
|
"T5TextEncode" : T5TextEncode,
|
|
}
|