47 lines
2.3 KiB
Python
47 lines
2.3 KiB
Python
import os
|
|
import torch
|
|
|
|
from comfy import sd1_clip
|
|
import comfy.model_management
|
|
from .gemma import Gemma2Model
|
|
|
|
from transformers import GemmaTokenizer as TFGemmaTokenizer
|
|
|
|
class GemmaClipModel(sd1_clip.SDClipModel):
|
|
def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs):
|
|
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "config.json")
|
|
special_tokens = {"start": 2, "end": 1, "pad": 0}
|
|
super().__init__(device=device, layer="last", layer_idx=None, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens=special_tokens, model_class=Gemma2Model, enable_attention_masks=True, return_attention_masks=False, model_options=model_options)
|
|
|
|
class SanaClipModel(sd1_clip.SD1ClipModel):
|
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
super().__init__(device=device, dtype=dtype, name="gemma", clip_model=GemmaClipModel, model_options=model_options)
|
|
|
|
class GemmaTokenizer(sd1_clip.SDTokenizer):
|
|
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
tokenizer_path = os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
|
|
"tokenizers", "gemma_tokenizer",
|
|
)
|
|
# TODO: reenable proper logic here - needs comfy version 44db978 or higher
|
|
super().__init__(tokenizer_path, embedding_directory=embedding_directory, pad_with_end=False, embedding_size=2304, embedding_key='gemma', tokenizer_class=TFGemmaTokenizer, has_start_token=False, pad_to_max_length=True, max_length=300, min_length=1)
|
|
self.start_token = 2
|
|
self.end_token = 1
|
|
self.pad_token = 0
|
|
|
|
def tokenize_with_weights(self, text, return_word_ids=False):
|
|
# TODO: see above, this is still just a wrapper for now
|
|
tokens = self.tokenizer(
|
|
text,
|
|
max_length=300,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt"
|
|
)
|
|
batched_tokens = [(x.item(), 1.0) for x in tokens.input_ids[0]]
|
|
return [batched_tokens]
|
|
|
|
class SanaTokenizer(sd1_clip.SD1Tokenizer):
|
|
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
|
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="gemma", tokenizer=GemmaTokenizer)
|