496 lines
20 KiB
Python
496 lines
20 KiB
Python
# transformers_api.py
|
||
"""Utilities for loading & running local HuggingFace Transformer models in ComfyUI LLM-Toolkit.
|
||
|
||
The helper focuses on *offline* (local) inference and therefore does not require any Internet
|
||
connectivity once the model weights have been downloaded. It provides a thin wrapper that mirrors
|
||
our other provider_* modules (openai_api, gemini_api, ollama_api …) so that *send_request* can call
|
||
it with a familiar signature.
|
||
|
||
Key design goals
|
||
----------------
|
||
• *Precision* – Load models in half-precision (float16 / bfloat16) by default to minimise VRAM
|
||
footprint. Users may override via the *precision* kwarg.
|
||
• *AWQ* – If the model repo or directory name contains the substring "AWQ" we will attempt to load
|
||
it with the *autoawq* loader (if available). Otherwise we fallback to standard HF loading.
|
||
• *Multi-modal* – Basic support for Alibaba Qwen2.5-Omni (text-image-audio) and Qwen-VL/VL-like
|
||
models is included. The caller can pass *base64_images* (List[str]) and *base64_audio* (List[str])
|
||
which will be converted to PIL images / raw audio arrays via *qwen_omni_utils* when available.
|
||
• *Model cache* – Keep models & tokenisers in a global dict so subsequent calls are instant.
|
||
• *Thread safety* – The cache is guarded by a simple asyncio lock.
|
||
|
||
The public coroutine *send_transformers_request* follows the same signature as other providers so
|
||
that integration inside *send_request.py* is trivial.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
import asyncio
|
||
import base64
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
import textwrap
|
||
|
||
import torch
|
||
|
||
# Transformers is an optional dependency for the toolkit – guard the import
|
||
try:
|
||
from transformers import (
|
||
AutoModelForCausalLM,
|
||
AutoTokenizer,
|
||
TextIteratorStreamer,
|
||
)
|
||
|
||
# Optional – might not exist in every transformers release. We import them
|
||
# in a nested try so failure does *not* disable the whole provider.
|
||
try:
|
||
from transformers import (
|
||
Qwen2_5OmniForConditionalGeneration, # type: ignore
|
||
Qwen2_5OmniProcessor, # type: ignore
|
||
)
|
||
except (ImportError, AttributeError): # pragma: no cover
|
||
Qwen2_5OmniForConditionalGeneration = None # type: ignore
|
||
Qwen2_5OmniProcessor = None # type: ignore
|
||
|
||
TRANSFORMERS_AVAILABLE = True
|
||
except Exception as _e: # pragma: no cover – keep ComfyUI boot-time lightweight
|
||
TRANSFORMERS_AVAILABLE = False # Will raise later when used
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Model cache (global across requests)
|
||
# -----------------------------------------------------------------------------
|
||
_MODEL_CACHE: Dict[str, Dict[str, Any]] = {}
|
||
_CACHE_LOCK = asyncio.Lock()
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Helper – determine torch dtype from precision string
|
||
# -----------------------------------------------------------------------------
|
||
|
||
_PRECISION_MAP = {
|
||
"fp16": torch.float16,
|
||
"float16": torch.float16,
|
||
"bf16": torch.bfloat16,
|
||
"bfloat16": torch.bfloat16,
|
||
"fp32": torch.float32,
|
||
"float32": torch.float32,
|
||
}
|
||
|
||
def _to_dtype(precision: str | None) -> torch.dtype:
|
||
if not precision:
|
||
return torch.float16
|
||
precision = precision.lower()
|
||
return _PRECISION_MAP.get(precision, torch.float16)
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Helper – decide which device to place the model on
|
||
# -----------------------------------------------------------------------------
|
||
|
||
def _best_device() -> str:
|
||
if torch.cuda.is_available():
|
||
return "cuda"
|
||
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): # macOS Metal
|
||
return "mps"
|
||
return "cpu"
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Loader – main entry used by *send_transformers_request*
|
||
# -----------------------------------------------------------------------------
|
||
|
||
def _is_awq_model(model_name: str) -> bool:
|
||
return "awq" in model_name.lower()
|
||
|
||
async def _ensure_model_loaded(
|
||
model_name_or_path: str,
|
||
*,
|
||
precision: str = "fp16",
|
||
device: str | None = None,
|
||
trust_remote_code: bool = True,
|
||
) -> Dict[str, Any]:
|
||
"""Load *model_name_or_path* into memory once and cache it.
|
||
|
||
Returns a dict with keys *model*, *tokenizer*, and optionally *processor* (for Qwen-Omni).
|
||
Calls are thread-safe (asyncio).
|
||
"""
|
||
if not TRANSFORMERS_AVAILABLE:
|
||
raise RuntimeError(
|
||
"transformers_api – HuggingFace transformers not installed. Please `pip install transformers`"
|
||
)
|
||
|
||
async with _CACHE_LOCK:
|
||
if model_name_or_path in _MODEL_CACHE:
|
||
return _MODEL_CACHE[model_name_or_path]
|
||
|
||
device = device or _best_device()
|
||
dtype = _to_dtype(precision)
|
||
|
||
logger.info(
|
||
"[Transformers-API] Loading model %s | device=%s | dtype=%s",
|
||
model_name_or_path,
|
||
device,
|
||
dtype,
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Special-case Qwen2.5-Omni multimodal
|
||
# ------------------------------------------------------------------
|
||
if (
|
||
"qwen2.5-omni" in model_name_or_path.lower()
|
||
and Qwen2_5OmniForConditionalGeneration is not None
|
||
and Qwen2_5OmniProcessor is not None
|
||
):
|
||
try:
|
||
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
|
||
model_name_or_path,
|
||
torch_dtype=dtype,
|
||
device_map="auto" if device == "cuda" else None, # CPU / MPS handled by HF
|
||
trust_remote_code=trust_remote_code,
|
||
)
|
||
processor = Qwen2_5OmniProcessor.from_pretrained(model_name_or_path)
|
||
tokenizer = None # Processor handles tokenisation
|
||
_MODEL_CACHE[model_name_or_path] = {
|
||
"model": model,
|
||
"tokenizer": tokenizer,
|
||
"processor": processor,
|
||
"modal": "omni",
|
||
}
|
||
return _MODEL_CACHE[model_name_or_path]
|
||
except Exception as exc:
|
||
logger.error("Failed to load Qwen2.5-Omni model %s – %s", model_name_or_path, exc)
|
||
# Fall through to generic loader below instead of raising
|
||
|
||
# If the user requested an Omni model but the specialised classes are unavailable, log a hint.
|
||
if "qwen2.5-omni" in model_name_or_path.lower() and Qwen2_5OmniForConditionalGeneration is None:
|
||
logger.warning(
|
||
"Qwen2.5-Omni support not available in your transformers build; "
|
||
"loading as a plain text model. Install a newer transformers pre-release "
|
||
"to enable multimodal features."
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# AWQ (AutoAWQ) support – try specialised loader first
|
||
# ------------------------------------------------------------------
|
||
if _is_awq_model(model_name_or_path):
|
||
try:
|
||
from awq import AutoAWQForCausalLM # type: ignore
|
||
|
||
model = AutoAWQForCausalLM.from_pretrained( # noqa: F401 – optional import
|
||
model_name_or_path,
|
||
torch_dtype=dtype,
|
||
device_map="auto" if device == "cuda" else None,
|
||
trust_remote_code=trust_remote_code,
|
||
)
|
||
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
|
||
_MODEL_CACHE[model_name_or_path] = {
|
||
"model": model,
|
||
"tokenizer": tokenizer,
|
||
"modal": "text",
|
||
}
|
||
return _MODEL_CACHE[model_name_or_path]
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"AWQ specific loader failed for %s (continuing with normal HF loader): %s",
|
||
model_name_or_path,
|
||
exc,
|
||
)
|
||
# Fall through to normal loader
|
||
|
||
# ------------------------------------------------------------------
|
||
# Generic text models (AutoModel)
|
||
# ------------------------------------------------------------------
|
||
try:
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_name_or_path,
|
||
torch_dtype=dtype,
|
||
device_map="auto" if device == "cuda" else None,
|
||
trust_remote_code=trust_remote_code,
|
||
)
|
||
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
|
||
_MODEL_CACHE[model_name_or_path] = {
|
||
"model": model,
|
||
"tokenizer": tokenizer,
|
||
"modal": "text",
|
||
}
|
||
return _MODEL_CACHE[model_name_or_path]
|
||
except Exception as exc:
|
||
logger.error("Could not load model %s via transformers – %s", model_name_or_path, exc)
|
||
raise
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Utility – multimodal helpers (images / audio) for Qwen-Omni
|
||
# -----------------------------------------------------------------------------
|
||
|
||
try:
|
||
from qwen_omni_utils import process_mm_info # provided by qwen-omni-utils package
|
||
except Exception: # pragma: no cover – optional dependency
|
||
def process_mm_info(*_args, **_kwargs): # type: ignore
|
||
raise NotImplementedError(
|
||
"qwen_omni_utils not installed – required for image/audio inputs with Qwen-Omni"
|
||
)
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Public coroutine – mirror *send_openai_request* API
|
||
# -----------------------------------------------------------------------------
|
||
async def send_transformers_request(
|
||
*,
|
||
base64_images: List[str],
|
||
base64_audio: List[str],
|
||
model: str,
|
||
system_message: str,
|
||
user_message: str,
|
||
messages: List[Dict[str, Any]] | None,
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 1024,
|
||
top_p: float = 0.9,
|
||
repeat_penalty: float = 1.1,
|
||
precision: str = "fp16",
|
||
**kwargs,
|
||
) -> Dict[str, Any]:
|
||
"""Run local inference with a HF transformer checkpoint.
|
||
|
||
Returns an *OpenAI-style* dict so the rest of the toolkit can remain provider-agnostic.
|
||
"""
|
||
if not TRANSFORMERS_AVAILABLE:
|
||
return {
|
||
"choices": [
|
||
{"message": {"content": "Transformers not installed. Please `pip install transformers`."}}]
|
||
}
|
||
|
||
try:
|
||
artefacts = await _ensure_model_loaded(model, precision=precision)
|
||
mdl = artefacts["model"]
|
||
tokenizer = artefacts.get("tokenizer")
|
||
processor = artefacts.get("processor")
|
||
modal_type = artefacts.get("modal", "text")
|
||
|
||
device = mdl.device if hasattr(mdl, "device") else torch.device("cpu")
|
||
mdl.eval()
|
||
|
||
# Build the conversation list for chat-style models
|
||
if messages is None:
|
||
messages = []
|
||
if system_message:
|
||
messages = [{"role": "system", "content": system_message}] + messages
|
||
if user_message:
|
||
messages = messages + [{"role": "user", "content": user_message}]
|
||
|
||
# ------------------------------------------------------------------
|
||
# Qwen2.5-Omni – multimodal path
|
||
# ------------------------------------------------------------------
|
||
if modal_type == "omni":
|
||
if processor is None:
|
||
raise RuntimeError("Qwen-Omni processor not initialised")
|
||
|
||
# Convert base64 images / audio to the expected format via helper
|
||
audios, images, videos = process_mm_info(
|
||
messages, # uses the conversation to pull URLs but we feed empty lists next
|
||
use_audio_in_video=False,
|
||
)
|
||
# Replace the image/audio placeholders explicitly – images list expects raw bytes or PIL
|
||
# For simplicity we treat base64_images as url-style – processor can handle raw bytes too.
|
||
for b64 in base64_images or []:
|
||
images.append({"type": "image", "image": f"data:image/png;base64,{b64}"})
|
||
for b64 in base64_audio or []:
|
||
audios.append({"type": "audio", "audio": base64.b64decode(b64)})
|
||
|
||
text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
|
||
inputs = processor(text=text_prompt, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True)
|
||
inputs = inputs.to(device)
|
||
|
||
out_ids, out_audio = mdl.generate(
|
||
**inputs,
|
||
do_sample=True,
|
||
temperature=temperature,
|
||
max_new_tokens=max_tokens,
|
||
top_p=top_p,
|
||
repetition_penalty=repeat_penalty,
|
||
return_audio=False, # We'll stick to text for now
|
||
)
|
||
output_texts = processor.batch_decode(out_ids, skip_special_tokens=True)
|
||
content = output_texts[0]
|
||
else:
|
||
# ------------------------------------------------------------------
|
||
# Generic text-only generation
|
||
# ------------------------------------------------------------------
|
||
if tokenizer is None:
|
||
raise RuntimeError("Tokenizer missing for text model")
|
||
|
||
# Compose single prompt from messages for simple models
|
||
prompt_parts: List[str] = []
|
||
for msg in messages:
|
||
if msg["role"] == "system":
|
||
prompt_parts.append(f"[INST] {msg['content']} [/INST]")
|
||
elif msg["role"] == "user":
|
||
prompt_parts.append(f"<|user|>: {msg['content']}")
|
||
elif msg["role"] == "assistant":
|
||
prompt_parts.append(f"<|assistant|>: {msg['content']}")
|
||
prompt = "\n".join(prompt_parts) + "\n<|assistant|>:"
|
||
|
||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
||
with torch.no_grad():
|
||
gen_ids = mdl.generate(
|
||
input_ids,
|
||
do_sample=True,
|
||
temperature=temperature,
|
||
max_new_tokens=max_tokens,
|
||
top_p=top_p,
|
||
repetition_penalty=repeat_penalty,
|
||
)
|
||
content = tokenizer.decode(gen_ids[0][input_ids.shape[-1] :], skip_special_tokens=True)
|
||
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"content": content,
|
||
}
|
||
}
|
||
]
|
||
}
|
||
except Exception as exc:
|
||
logger.error("transformers_api – inference error: %s", exc, exc_info=True)
|
||
return {"choices": [{"message": {"content": f"Error: {exc}"}}]}
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Public coroutine – streaming version (yields chunks)
|
||
# -----------------------------------------------------------------------------
|
||
async def send_transformers_request_stream(
|
||
*,
|
||
base64_images: List[str],
|
||
base64_audio: List[str],
|
||
model: str,
|
||
system_message: str,
|
||
user_message: str,
|
||
messages: List[Dict[str, Any]] | None,
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 1024,
|
||
top_p: float = 0.9,
|
||
repeat_penalty: float = 1.1,
|
||
precision: str = "fp16",
|
||
**kwargs,
|
||
):
|
||
"""Async generator that yields text chunks as they are generated."""
|
||
if not TRANSFORMERS_AVAILABLE:
|
||
yield "[Transformers not installed]"
|
||
return
|
||
|
||
try:
|
||
artefacts = await _ensure_model_loaded(model, precision=precision)
|
||
mdl = artefacts["model"]
|
||
tokenizer = artefacts.get("tokenizer")
|
||
modal_type = artefacts.get("modal", "text")
|
||
|
||
if modal_type != "text":
|
||
# For now fall back to full generation for multimodal models.
|
||
full_resp = await send_transformers_request(
|
||
base64_images=base64_images,
|
||
base64_audio=base64_audio,
|
||
model=model,
|
||
system_message=system_message,
|
||
user_message=user_message,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
top_p=top_p,
|
||
repeat_penalty=repeat_penalty,
|
||
precision=precision,
|
||
)
|
||
content = full_resp.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||
if content:
|
||
for _chunk in textwrap.wrap(content, 120):
|
||
yield _chunk
|
||
return
|
||
|
||
if tokenizer is None:
|
||
raise RuntimeError("Tokenizer missing for text model (streaming)")
|
||
|
||
device = mdl.device if hasattr(mdl, "device") else torch.device("cpu")
|
||
mdl.eval()
|
||
|
||
# Build prompt
|
||
if messages is None:
|
||
messages = []
|
||
if system_message:
|
||
messages = [{"role": "system", "content": system_message}] + messages
|
||
if user_message:
|
||
messages = messages + [{"role": "user", "content": user_message}]
|
||
|
||
prompt_parts: List[str] = []
|
||
for msg in messages:
|
||
if msg["role"] == "system":
|
||
prompt_parts.append(f"[INST] {msg['content']} [/INST]")
|
||
elif msg["role"] == "user":
|
||
prompt_parts.append(f"<|user|>: {msg['content']}")
|
||
elif msg["role"] == "assistant":
|
||
prompt_parts.append(f"<|assistant|>: {msg['content']}")
|
||
prompt = "\n".join(prompt_parts) + "\n<|assistant|>:"
|
||
|
||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
||
|
||
from threading import Thread
|
||
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||
|
||
generation_kwargs = dict(
|
||
input_ids=input_ids,
|
||
streamer=streamer,
|
||
do_sample=True,
|
||
temperature=temperature,
|
||
max_new_tokens=max_tokens,
|
||
top_p=top_p,
|
||
repetition_penalty=repeat_penalty,
|
||
)
|
||
|
||
# Run generate in background thread so we can iterate streamer concurrently
|
||
gen_thread = Thread(target=mdl.generate, kwargs=generation_kwargs, daemon=True)
|
||
gen_thread.start()
|
||
|
||
loop = asyncio.get_running_loop()
|
||
|
||
stop = False
|
||
while not stop:
|
||
def _next_chunk():
|
||
try:
|
||
return next(streamer)
|
||
except StopIteration:
|
||
return None
|
||
|
||
chunk = await loop.run_in_executor(None, _next_chunk)
|
||
if chunk is None:
|
||
stop = True
|
||
elif chunk:
|
||
yield chunk
|
||
|
||
gen_thread.join(timeout=0)
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"transformers_api – streaming failed (%s). Falling back to blocking generation.",
|
||
exc,
|
||
)
|
||
try:
|
||
full_resp = await send_transformers_request(
|
||
base64_images=base64_images,
|
||
base64_audio=base64_audio,
|
||
model=model,
|
||
system_message=system_message,
|
||
user_message=user_message,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
top_p=top_p,
|
||
repeat_penalty=repeat_penalty,
|
||
precision=precision,
|
||
)
|
||
content = full_resp.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||
if content:
|
||
for _chunk in textwrap.wrap(content, 120):
|
||
yield _chunk
|
||
except Exception as e2:
|
||
logger.error("transformers_api – fallback blocking generation also failed: %s", e2, exc_info=True)
|
||
yield f"[Error: {e2}]" |