mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-12-20 03:23:00 +08:00
e1e50e017c
Since model saving is not implemented for the mixed quant system this breaks model saving for every scaled fp8 model which needs to be fixed before this gets merged.
39 lines
1.8 KiB
Python
39 lines
1.8 KiB
Python
from comfy import sd1_clip
|
|
import comfy.text_encoders.sd3_clip
|
|
import os
|
|
from transformers import T5TokenizerFast
|
|
|
|
|
|
class T5XXLModel(comfy.text_encoders.sd3_clip.T5XXLModel):
|
|
def __init__(self, **kwargs):
|
|
kwargs["attention_mask"] = True
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
class MochiT5XXL(sd1_clip.SD1ClipModel):
|
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
super().__init__(device=device, dtype=dtype, name="t5xxl", clip_model=T5XXLModel, model_options=model_options)
|
|
|
|
|
|
class T5XXLTokenizer(sd1_clip.SDTokenizer):
|
|
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
|
|
super().__init__(tokenizer_path, embedding_directory=embedding_directory, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256, tokenizer_data=tokenizer_data)
|
|
|
|
|
|
class MochiT5Tokenizer(sd1_clip.SD1Tokenizer):
|
|
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="t5xxl", tokenizer=T5XXLTokenizer)
|
|
|
|
|
|
def mochi_te(dtype_t5=None, t5_quantization_metadata=None):
|
|
class MochiTEModel_(MochiT5XXL):
|
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
if t5_quantization_metadata is not None:
|
|
model_options = model_options.copy()
|
|
model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
|
|
if dtype is None:
|
|
dtype = dtype_t5
|
|
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
|
return MochiTEModel_
|