96 lines
2.4 KiB
Python
96 lines
2.4 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import folder_paths
|
|
|
|
from .loader import load_t5
|
|
from ..utils.dtype import string_to_dtype
|
|
|
|
# initialize custom folder path
|
|
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.folder_names_and_paths.get("t5", [[],set()])[0]
|
|
],
|
|
folder_paths.supported_pt_extensions
|
|
)
|
|
|
|
dtypes = [
|
|
"default",
|
|
"auto (comfy)",
|
|
"FP32",
|
|
"FP16",
|
|
# Note: remove these at some point
|
|
"bnb8bit",
|
|
"bnb4bit",
|
|
]
|
|
try: torch.float8_e5m2
|
|
except AttributeError: print("Torch version too old for FP8")
|
|
else: dtypes += ["FP8 E4M3", "FP8 E5M2"]
|
|
|
|
class T5v11Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
devices = ["auto", "cpu", "gpu"]
|
|
# hack for using second GPU as offload
|
|
for k in range(1, torch.cuda.device_count()):
|
|
devices.append(f"cuda:{k}")
|
|
return {
|
|
"required": {
|
|
"t5v11_name": (folder_paths.get_filename_list("t5"),),
|
|
"t5v11_ver": (["xxl"],),
|
|
"path_type": (["folder", "file"],),
|
|
"device": (devices, {"default":"cpu"}),
|
|
"dtype": (dtypes,),
|
|
}
|
|
}
|
|
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" or device.startswith("cuda"), "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,
|
|
}
|