bemoregtandClaude Sonnet 4.6
9a0c3dc1a0
- LoadFFTNetModel / FFTNetGenerate ComfyUI nodes - FFTNet architecture (SpectralFilter, LocalWindowMixing, FFTNetMixer) - upgrade/: replace LLM self-attention with FFTNet + MPS-safe fine-tuning script - Auto-detects self_attn return count (1/2/3) for transformers version compat - MPS FFT CPU offload workaround for Apple Silicon Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
413 lines
17 KiB
Python
413 lines
17 KiB
Python
"""
|
||
ComfyUI custom nodes for FFTNet inference.
|
||
|
||
Nodes:
|
||
• LoadFFTNetModel – load a trained FFTNet checkpoint + tokenizer
|
||
(standalone FFTNetForCausalLM 또는 Upgrade된 LLaMA 모두 지원)
|
||
• FFTNetGenerate – generate text from a prompt
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import threading
|
||
import torch
|
||
from .fftnet_arch import FFTNetForCausalLM
|
||
from .upgrade.replace_attention import FFTNET_MARKER
|
||
|
||
|
||
# ── Tokenizer helpers ─────────────────────────────────────────────────────────
|
||
|
||
def _load_tokenizer(tokenizer_path: str | None):
|
||
"""
|
||
Try to load a tokenizer in order of preference:
|
||
1. transformers AutoTokenizer from model directory
|
||
2. tiktoken (gpt-2 encoding)
|
||
3. Minimal character-level fallback
|
||
|
||
Returns an object with .encode(str)->list[int] and .decode(list[int])->str
|
||
"""
|
||
# 1) transformers
|
||
if tokenizer_path and os.path.isdir(tokenizer_path):
|
||
try:
|
||
from transformers import AutoTokenizer
|
||
tok = AutoTokenizer.from_pretrained(tokenizer_path)
|
||
return tok
|
||
except Exception as e:
|
||
print(f"[FFTNet] transformers tokenizer failed: {e}")
|
||
|
||
# 2) tiktoken (gpt-2)
|
||
try:
|
||
import tiktoken
|
||
enc = tiktoken.get_encoding("gpt2")
|
||
|
||
class TiktokenWrapper:
|
||
def encode(self, text: str):
|
||
return enc.encode(text)
|
||
def decode(self, ids):
|
||
return enc.decode(ids)
|
||
|
||
return TiktokenWrapper()
|
||
except Exception:
|
||
pass
|
||
|
||
# 3) Character-level fallback
|
||
class CharTokenizer:
|
||
def encode(self, text: str):
|
||
return [ord(c) for c in text]
|
||
def decode(self, ids):
|
||
return "".join(chr(i) for i in ids)
|
||
|
||
print("[FFTNet] Warning: using character-level tokenizer (install tiktoken for better results)")
|
||
return CharTokenizer()
|
||
|
||
|
||
# ── Thread helpers (escape inference_mode) ────────────────────────────────────
|
||
|
||
def _run_in_thread(fn, *args, **kwargs):
|
||
"""ComfyUI wraps node execution in torch.inference_mode(); run in fresh thread to escape."""
|
||
result = [None]
|
||
error = [None]
|
||
|
||
def target():
|
||
try:
|
||
result[0] = fn(*args, **kwargs)
|
||
except Exception as e:
|
||
error[0] = e
|
||
|
||
t = threading.Thread(target=target)
|
||
t.start()
|
||
t.join()
|
||
if error[0] is not None:
|
||
raise error[0]
|
||
return result[0]
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════════════════════════
|
||
# Node 1 – LoadFFTNetModel
|
||
# ══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class LoadFFTNetModel:
|
||
"""
|
||
Load a trained FFTNet checkpoint.
|
||
|
||
The checkpoint (.pt / .pth) should be a dict saved via:
|
||
torch.save({
|
||
"model_state_dict": model.state_dict(),
|
||
"config": { ... }, # optional – model hyper-parameters
|
||
"tokenizer": "path/or/name" # optional
|
||
}, path)
|
||
|
||
If the checkpoint contains only the state_dict (no wrapper dict),
|
||
model hyper-parameters must be supplied through this node's inputs.
|
||
A config.json next to the checkpoint is also auto-detected.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"checkpoint_path": ("STRING", {
|
||
"default": "/path/to/fftnet_model.pt",
|
||
"multiline": False,
|
||
}),
|
||
"device": (["auto", "cpu", "cuda", "mps"],),
|
||
},
|
||
"optional": {
|
||
# Override / supply model config manually
|
||
"vocab_size": ("INT", {"default": 50257, "min": 256, "max": 256000}),
|
||
"d_model": ("INT", {"default": 512, "min": 64, "max": 4096}),
|
||
"n_layers": ("INT", {"default": 6, "min": 1, "max": 64}),
|
||
"d_ff": ("INT", {"default": 2048, "min": 64, "max": 16384}),
|
||
"max_seq_len": ("INT", {"default": 512, "min": 64, "max": 8192}),
|
||
"window_size": ("INT", {"default": 4, "min": 1, "max": 64}),
|
||
# Optional: directory or path where the tokenizer lives
|
||
"tokenizer_path": ("STRING", {"default": "", "multiline": False}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("FFTNET_MODEL",)
|
||
RETURN_NAMES = ("fftnet_model",)
|
||
FUNCTION = "load_model"
|
||
CATEGORY = "FFTNet"
|
||
DESCRIPTION = (
|
||
"학습된 FFTNet 모델 체크포인트를 로드합니다. "
|
||
"체크포인트 파일(.pt/.pth)과 선택적으로 토크나이저 경로를 입력하세요."
|
||
)
|
||
|
||
def load_model(
|
||
self,
|
||
checkpoint_path: str,
|
||
device: str,
|
||
vocab_size: int = 50257,
|
||
d_model: int = 512,
|
||
n_layers: int = 6,
|
||
d_ff: int = 2048,
|
||
max_seq_len: int = 512,
|
||
window_size: int = 4,
|
||
tokenizer_path: str = "",
|
||
):
|
||
return _run_in_thread(
|
||
self._load_inner,
|
||
checkpoint_path, device,
|
||
vocab_size, d_model, n_layers, d_ff, max_seq_len, window_size,
|
||
tokenizer_path,
|
||
)
|
||
|
||
def _load_inner(
|
||
self,
|
||
checkpoint_path, device_str,
|
||
vocab_size, d_model, n_layers, d_ff, max_seq_len, window_size,
|
||
tokenizer_path,
|
||
):
|
||
# ── resolve device ──────────────────────────────────────────────────
|
||
if device_str == "auto":
|
||
if torch.cuda.is_available():
|
||
device = torch.device("cuda")
|
||
elif torch.backends.mps.is_available():
|
||
device = torch.device("mps")
|
||
else:
|
||
device = torch.device("cpu")
|
||
else:
|
||
device = torch.device(device_str)
|
||
|
||
print(f"[FFTNet] Loading checkpoint: {checkpoint_path} (device={device})")
|
||
|
||
# ── load checkpoint ─────────────────────────────────────────────────
|
||
if not os.path.isfile(checkpoint_path):
|
||
raise FileNotFoundError(f"[FFTNet] Checkpoint not found: {checkpoint_path}")
|
||
|
||
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||
|
||
# ── Upgraded LLaMA 체크포인트 분기 ──────────────────────────────────
|
||
if isinstance(ckpt, dict) and ckpt.get("type") == FFTNET_MARKER:
|
||
return self._load_upgraded_inner(
|
||
ckpt, checkpoint_path, device, tokenizer_path, window_size
|
||
)
|
||
|
||
# Support both bare state_dict and wrapped dict
|
||
if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
|
||
state_dict = ckpt["model_state_dict"]
|
||
ckpt_config = ckpt.get("config", {})
|
||
ckpt_tok_path = ckpt.get("tokenizer", "")
|
||
elif isinstance(ckpt, dict) and all(isinstance(v, torch.Tensor) for v in ckpt.values()):
|
||
state_dict = ckpt
|
||
ckpt_config = {}
|
||
ckpt_tok_path = ""
|
||
else:
|
||
# e.g. full model object
|
||
state_dict = ckpt.state_dict() if hasattr(ckpt, "state_dict") else ckpt
|
||
ckpt_config = getattr(ckpt, "config", {}) if hasattr(ckpt, "config") else {}
|
||
ckpt_tok_path = ""
|
||
|
||
# ── try auto-detect config.json next to checkpoint ──────────────────
|
||
cfg_json = os.path.join(os.path.dirname(checkpoint_path), "config.json")
|
||
if os.path.isfile(cfg_json):
|
||
with open(cfg_json) as f:
|
||
file_config = json.load(f)
|
||
print(f"[FFTNet] Loaded config from {cfg_json}")
|
||
else:
|
||
file_config = {}
|
||
|
||
# Priority: node inputs > checkpoint config > file config > defaults
|
||
def _pick(key, node_val, default):
|
||
if key in ckpt_config:
|
||
return ckpt_config[key]
|
||
if key in file_config:
|
||
return file_config[key]
|
||
return node_val if node_val != default else default
|
||
|
||
config = {
|
||
"vocab_size": _pick("vocab_size", vocab_size, 50257),
|
||
"d_model": _pick("d_model", d_model, 512),
|
||
"n_layers": _pick("n_layers", n_layers, 6),
|
||
"d_ff": _pick("d_ff", d_ff, 2048),
|
||
"max_seq_len": _pick("max_seq_len", max_seq_len, 512),
|
||
"window_size": _pick("window_size", window_size, 4),
|
||
"dropout": ckpt_config.get("dropout", file_config.get("dropout", 0.0)),
|
||
"eos_token_id": ckpt_config.get("eos_token_id", file_config.get("eos_token_id", None)),
|
||
}
|
||
print(f"[FFTNet] Model config: {config}")
|
||
|
||
# ── build model & load weights ──────────────────────────────────────
|
||
model = FFTNetForCausalLM(config)
|
||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||
if missing:
|
||
print(f"[FFTNet] Warning – missing keys: {missing[:5]} ...")
|
||
if unexpected:
|
||
print(f"[FFTNet] Warning – unexpected keys: {unexpected[:5]} ...")
|
||
model.to(device)
|
||
model.eval()
|
||
print(f"[FFTNet] Model loaded ({sum(p.numel() for p in model.parameters()):,} params)")
|
||
|
||
# ── tokenizer ───────────────────────────────────────────────────────
|
||
tok_search = tokenizer_path or ckpt_tok_path or os.path.dirname(checkpoint_path)
|
||
tokenizer = _load_tokenizer(tok_search if tok_search else None)
|
||
print(f"[FFTNet] Tokenizer: {type(tokenizer).__name__}")
|
||
|
||
bundle = {
|
||
"model": model,
|
||
"tokenizer": tokenizer,
|
||
"device": device,
|
||
"config": config,
|
||
"model_type": "standalone",
|
||
}
|
||
return (bundle,)
|
||
|
||
def _load_upgraded_inner(self, ckpt, checkpoint_path, device, tokenizer_path, window_size):
|
||
"""Upgrade된 LLaMA 체크포인트 로드 (transformers 필요)."""
|
||
from .upgrade.replace_attention import load_upgraded_model
|
||
|
||
ws = ckpt.get("fftnet_config", {}).get("window_size", window_size)
|
||
model, tokenizer, fftnet_cfg = load_upgraded_model(checkpoint_path, device, window_size=ws)
|
||
|
||
# 노드에서 별도로 tokenizer_path 를 지정한 경우 우선 적용
|
||
if tokenizer_path:
|
||
alt = _load_tokenizer(tokenizer_path)
|
||
if alt is not None:
|
||
tokenizer = alt
|
||
|
||
if tokenizer is None:
|
||
tokenizer = _load_tokenizer(None)
|
||
|
||
cfg = ckpt.get("fftnet_config", {})
|
||
bundle = {
|
||
"model": model,
|
||
"tokenizer": tokenizer,
|
||
"device": device,
|
||
"config": cfg,
|
||
"model_type": "upgraded",
|
||
}
|
||
print(f"[FFTNet] Upgraded LLaMA 로드 완료 (window_size={ws})")
|
||
return (bundle,)
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════════════════════════
|
||
# Node 2 – FFTNetGenerate
|
||
# ══════════════════════════════════════════════════════════════════════════════
|
||
|
||
class FFTNetGenerate:
|
||
"""
|
||
Run autoregressive text generation with a loaded FFTNet model.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"fftnet_model": ("FFTNET_MODEL",),
|
||
"prompt": ("STRING", {
|
||
"default": "Once upon a time",
|
||
"multiline": True,
|
||
}),
|
||
"max_new_tokens": ("INT", {"default": 200, "min": 1, "max": 4096}),
|
||
"temperature": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 2.0, "step": 0.01}),
|
||
"top_k": ("INT", {"default": 50, "min": 0, "max": 1000}),
|
||
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"repetition_penalty": ("FLOAT", {"default": 1.1, "min": 1.0, "max": 5.0, "step": 0.05}),
|
||
"seed": ("INT", {"default": -1, "min": -1, "max": 2**31 - 1}),
|
||
},
|
||
"optional": {
|
||
"system_prompt": ("STRING", {
|
||
"default": "",
|
||
"multiline": True,
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("generated_text",)
|
||
OUTPUT_NODE = True
|
||
FUNCTION = "generate"
|
||
CATEGORY = "FFTNet"
|
||
DESCRIPTION = (
|
||
"FFTNet 모델로 텍스트를 생성합니다. "
|
||
"프롬프트를 입력하면 모델이 이어지는 텍스트를 자기회귀적으로 생성합니다."
|
||
)
|
||
|
||
def generate(
|
||
self,
|
||
fftnet_model: dict,
|
||
prompt: str,
|
||
max_new_tokens: int,
|
||
temperature: float,
|
||
top_k: int,
|
||
top_p: float,
|
||
repetition_penalty: float,
|
||
seed: int,
|
||
system_prompt: str = "",
|
||
):
|
||
return _run_in_thread(
|
||
self._generate_inner,
|
||
fftnet_model, prompt, max_new_tokens,
|
||
temperature, top_k, top_p, repetition_penalty,
|
||
seed, system_prompt,
|
||
)
|
||
|
||
def _generate_inner(
|
||
self,
|
||
bundle: dict,
|
||
prompt: str,
|
||
max_new_tokens: int,
|
||
temperature: float,
|
||
top_k: int,
|
||
top_p: float,
|
||
repetition_penalty: float,
|
||
seed: int,
|
||
system_prompt: str,
|
||
):
|
||
model = bundle["model"]
|
||
tokenizer = bundle["tokenizer"]
|
||
device = bundle["device"]
|
||
model_type = bundle.get("model_type", "standalone")
|
||
|
||
if seed >= 0:
|
||
torch.manual_seed(seed)
|
||
|
||
full_text = f"{system_prompt.strip()}\n\n{prompt}" if system_prompt.strip() else prompt
|
||
|
||
# ── tokenize ─────────────────────────────────────────────────────────
|
||
token_ids = tokenizer.encode(full_text)
|
||
if hasattr(token_ids, "input_ids"):
|
||
ids = token_ids.input_ids
|
||
token_ids = ids[0].tolist() if isinstance(ids, torch.Tensor) else list(ids)
|
||
elif isinstance(token_ids, torch.Tensor):
|
||
token_ids = token_ids.squeeze().tolist()
|
||
|
||
input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
|
||
prompt_len = input_ids.shape[1]
|
||
print(f"[FFTNet] [{model_type}] Prompt tokens: {prompt_len} | max_new: {max_new_tokens}")
|
||
|
||
# ── generate ─────────────────────────────────────────────────────────
|
||
if model_type == "upgraded":
|
||
# HuggingFace model.generate() 사용
|
||
output_ids = model.generate(
|
||
input_ids,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=True,
|
||
temperature=max(temperature, 1e-4),
|
||
top_k=top_k if top_k > 0 else None,
|
||
top_p=top_p,
|
||
repetition_penalty=repetition_penalty,
|
||
pad_token_id=tokenizer.eos_token_id,
|
||
)
|
||
else:
|
||
# standalone FFTNetForCausalLM.generate()
|
||
output_ids = model.generate(
|
||
input_ids,
|
||
max_new_tokens=max_new_tokens,
|
||
temperature=temperature,
|
||
top_k=top_k,
|
||
top_p=top_p,
|
||
repetition_penalty=repetition_penalty,
|
||
)
|
||
|
||
# ── decode ───────────────────────────────────────────────────────────
|
||
new_ids = output_ids[0, prompt_len:].tolist()
|
||
try:
|
||
decoded = tokenizer.decode(new_ids, skip_special_tokens=True)
|
||
except TypeError:
|
||
decoded = tokenizer.decode(new_ids)
|
||
|
||
print(f"[FFTNet] Generated {len(new_ids)} tokens.\n{decoded}")
|
||
return (decoded,)
|