mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-27 17:37:39 +08:00
105 lines
5.1 KiB
Python
105 lines
5.1 KiB
Python
import torch
|
|
|
|
from comfy import sd1_clip
|
|
from .lumina2 import Gemma2BTokenizer, LuminaModel
|
|
import comfy.text_encoders.llama
|
|
|
|
|
|
class PixelDiTGemma2_2BModel(sd1_clip.SDClipModel):
|
|
def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
|
|
llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
|
|
if llama_quantization_metadata is not None:
|
|
model_options = model_options.copy()
|
|
model_options["quantization_metadata"] = llama_quantization_metadata
|
|
|
|
super().__init__(
|
|
device=device, layer=layer, layer_idx=layer_idx,
|
|
textmodel_json_config={}, dtype=dtype,
|
|
special_tokens={"start": 2, "pad": 0},
|
|
layer_norm_hidden_state=False,
|
|
model_class=comfy.text_encoders.llama.Gemma2_2B,
|
|
enable_attention_masks=attention_mask,
|
|
return_attention_masks=attention_mask,
|
|
model_options=model_options,
|
|
)
|
|
|
|
|
|
_PIXELDIT_CHI_PROMPT = (
|
|
'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions '
|
|
"suitable for image generation. Evaluate the level of detail in the user prompt:\n"
|
|
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, "
|
|
"and spatial relationships to create vivid and concrete scenes.\n"
|
|
"- If the prompt is already detailed, refine and enhance the existing details slightly without "
|
|
"overcomplicating.\n"
|
|
"Here are examples of how to transform or refine prompts:\n"
|
|
"- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, "
|
|
"sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.\n"
|
|
"- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring "
|
|
"glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus "
|
|
"passing by towering glass skyscrapers.\n"
|
|
"Please generate only the enhanced description for the prompt below and avoid including any "
|
|
"additional commentary or evaluations:\n"
|
|
"User Prompt: "
|
|
)
|
|
|
|
_PIXELDIT_MAX_LENGTH = 300
|
|
_PIXELDIT_CHI_PROMPT_DETECT_PREFIX = 'Given a user prompt, generate an "Enhanced prompt"'
|
|
|
|
|
|
class PixelDiTGemma2Tokenizer(sd1_clip.SD1Tokenizer):
|
|
def __init__(self, embedding_directory=None, tokenizer_data=None):
|
|
if tokenizer_data is None:
|
|
tokenizer_data = {}
|
|
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data,
|
|
name="gemma2_2b", tokenizer=Gemma2BTokenizer)
|
|
|
|
def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
|
|
if not text.strip():
|
|
return super().tokenize_with_weights("", return_word_ids=return_word_ids, disable_weights=True, min_length=_PIXELDIT_MAX_LENGTH)
|
|
|
|
chi_token_count = len(self.gemma2_2b.tokenizer(_PIXELDIT_CHI_PROMPT)["input_ids"])
|
|
combined = text if text.startswith(_PIXELDIT_CHI_PROMPT_DETECT_PREFIX) else _PIXELDIT_CHI_PROMPT + text
|
|
max_length_all = chi_token_count + _PIXELDIT_MAX_LENGTH - 2
|
|
out = super().tokenize_with_weights(combined, return_word_ids=return_word_ids,
|
|
disable_weights=True, min_length=max_length_all)
|
|
out["gemma2_2b"] = [out["gemma2_2b"][0][:max_length_all]]
|
|
return out
|
|
|
|
def untokenize(self, token_weight_pair):
|
|
return self.gemma2_2b.untokenize(token_weight_pair)
|
|
|
|
def state_dict(self):
|
|
return self.gemma2_2b.state_dict()
|
|
|
|
|
|
class PixelDiTGemma2TE(LuminaModel):
|
|
# PixelDiT's select_index: keep BOS + last 299 embeddings of the padded sequence.
|
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
super().__init__(device=device, dtype=dtype, name="gemma2_2b",
|
|
clip_model=PixelDiTGemma2_2BModel, model_options=model_options)
|
|
|
|
def encode_token_weights(self, token_weight_pairs):
|
|
result = super().encode_token_weights(token_weight_pairs)
|
|
cond, pooled = result[0], result[1]
|
|
extra = result[2] if len(result) > 2 else None
|
|
if cond.shape[1] > _PIXELDIT_MAX_LENGTH:
|
|
cond = torch.cat([cond[:, :1], cond[:, -(_PIXELDIT_MAX_LENGTH - 1):]], dim=1)
|
|
if extra is not None and "attention_mask" in extra:
|
|
am = extra["attention_mask"]
|
|
extra["attention_mask"] = torch.cat([am[..., :1], am[..., -(_PIXELDIT_MAX_LENGTH - 1):]], dim=-1)
|
|
if extra is not None:
|
|
return cond, pooled, extra
|
|
return cond, pooled
|
|
|
|
|
|
def pixeldit_te(dtype_llama=None, llama_quantization_metadata=None):
|
|
class PixelDiTTE_(PixelDiTGemma2TE):
|
|
def __init__(self, device="cpu", dtype=None, model_options={}):
|
|
if llama_quantization_metadata is not None:
|
|
model_options = model_options.copy()
|
|
model_options["llama_quantization_metadata"] = llama_quantization_metadata
|
|
if dtype_llama is not None:
|
|
dtype = dtype_llama
|
|
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
|
return PixelDiTTE_
|