110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
import torch
|
|
import os
|
|
import folder_paths
|
|
import comfy.model_management
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
|
from peft import PeftModel
|
|
from .resampler import SDXLResampler
|
|
|
|
|
|
class LLMClipLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
file_list = folder_paths.get_filename_list("clip")
|
|
return {
|
|
"required": {
|
|
"base_repo_id": ("STRING", {"default": "unsloth/Qwen3-4B-unsloth-bnb-4bit"}),
|
|
"resampler_name": (file_list,),
|
|
},
|
|
"optional": {
|
|
"lora_repo_id": ("STRING", {"default": "molbal/qwen-clip-resampler-adapter"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLM_CLIP_BUNDLE",)
|
|
FUNCTION = "load_llm_clip"
|
|
CATEGORY = "LLM_CLIP"
|
|
TITLE = "LLM CLIP Loader (Full HF)"
|
|
|
|
def load_llm_clip(self, base_repo_id, resampler_name, lora_repo_id=None):
|
|
device = comfy.model_management.text_encoder_device()
|
|
resampler_path = folder_paths.get_full_path("clip", resampler_name)
|
|
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=torch.float16,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_use_double_quant=True,
|
|
llm_int8_enable_fp32_cpu_offload=True
|
|
)
|
|
|
|
print(f"LLM CLIP: Loading base {base_repo_id}...")
|
|
tokenizer = AutoTokenizer.from_pretrained(base_repo_id, trust_remote_code=True)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
base_repo_id,
|
|
quantization_config=bnb_config,
|
|
trust_remote_code=True,
|
|
device_map="auto"
|
|
)
|
|
|
|
if lora_repo_id and lora_repo_id.strip() != "":
|
|
print(f"LLM CLIP: Applying LoRA from {lora_repo_id}...")
|
|
model = PeftModel.from_pretrained(model, lora_repo_id)
|
|
|
|
# 1. Initialize Resampler
|
|
hidden_size = model.config.hidden_size
|
|
# 2. CRITICAL: Move to device AND cast to HALF immediately
|
|
resampler = SDXLResampler(llm_dim=hidden_size, sdxl_dim=2048).to(device).half()
|
|
|
|
resampler.load_state_dict(torch.load(resampler_path, map_location=device))
|
|
resampler.eval()
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
return ({"model": model, "tokenizer": tokenizer, "resampler": resampler},)
|
|
|
|
|
|
class LLMClipTextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip_bundle": ("LLM_CLIP_BUNDLE",),
|
|
"text": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "LLM_CLIP"
|
|
|
|
def encode(self, clip_bundle, text):
|
|
model = clip_bundle["model"]
|
|
tokenizer = clip_bundle["tokenizer"]
|
|
resampler = clip_bundle["resampler"]
|
|
|
|
model_device = next(model.parameters()).device
|
|
resampler_device = next(resampler.parameters()).device
|
|
|
|
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256)
|
|
inputs = {k: v.to(model_device) for k, v in inputs.items()}
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs, output_hidden_states=True)
|
|
|
|
# 3. CRITICAL: Ensure hidden states are cast to match resampler (Half)
|
|
last_hidden = outputs.hidden_states[-1].to(resampler_device).half()
|
|
|
|
# Run resampler
|
|
cond, pooled = resampler(last_hidden)
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
# Final output tensors for ComfyUI need to be Float32 (Comfy's standard)
|
|
return ([[cond.float(), {"pooled_output": pooled.float()}]],)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLMClipLoader": LLMClipLoader,
|
|
"LLMClipTextEncode": LLMClipTextEncode
|
|
} |