372 lines
15 KiB
Python
372 lines
15 KiB
Python
"""
|
|
ComfyUI BuimenLabo - Unified Package
|
|
A comprehensive collection of ComfyUI custom nodes including:
|
|
- Multi-language locale toggle with 20 language support
|
|
- AI-powered pose analysis with Gemini integration
|
|
- Advanced LLM text generation with llama-cpp-python
|
|
- Smart ControlNet loader for SD/SDXL models
|
|
- AI-powered prompt translation
|
|
|
|
Author: BuimenLabo
|
|
Blog: https://note.com/hirodream44
|
|
"""
|
|
|
|
# Import version information
|
|
try:
|
|
from ._version import __version__, PACKAGE_NAME, AUTHOR, BLOG_URL, print_version_info
|
|
print(f"🚀 {PACKAGE_NAME} v{__version__} - {AUTHOR}")
|
|
print(f"📝 Blog: {BLOG_URL}")
|
|
except ImportError:
|
|
__version__ = "1.0.0"
|
|
print(f"🚀 ComfyUI BuimenLabo v{__version__}")
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
import random
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
print("🚀 Loading ComfyUI BuimenLabo unified package...")
|
|
|
|
# Import all packages as-is (simple copy approach)
|
|
try:
|
|
# 1. Gemini Pose Analyzer
|
|
from .gemini_pose_analyzer import NODE_CLASS_MAPPINGS as GEMINI_MAPPINGS
|
|
from .gemini_pose_analyzer import NODE_DISPLAY_NAME_MAPPINGS as GEMINI_DISPLAY_MAPPINGS
|
|
print("✅ GeminiPoseAnalyzer loaded")
|
|
except Exception as e:
|
|
print(f"⚠️ GeminiPoseAnalyzer failed: {e}")
|
|
GEMINI_MAPPINGS = {}
|
|
GEMINI_DISPLAY_MAPPINGS = {}
|
|
|
|
try:
|
|
# 2. LlamaCpp Text (working version)
|
|
from .llama_cpp_node import (
|
|
LlamaCppLoader,
|
|
LlamaCppGenerate,
|
|
LlamaCppUnload,
|
|
LlamaCppSafeUnload,
|
|
LlamaCppMemoryInfo,
|
|
LlamaCppCompleteUnload,
|
|
LlamaCppAIO,
|
|
)
|
|
LLAMA_MAPPINGS = {
|
|
"LlamaCppLoader": LlamaCppLoader,
|
|
"LlamaCppGenerate": LlamaCppGenerate,
|
|
"LlamaCppUnload": LlamaCppUnload,
|
|
"LlamaCppSafeUnload": LlamaCppSafeUnload,
|
|
"LlamaCppMemoryInfo": LlamaCppMemoryInfo,
|
|
"LlamaCppCompleteUnload": LlamaCppCompleteUnload,
|
|
"LlamaCppAIO": LlamaCppAIO,
|
|
}
|
|
LLAMA_DISPLAY_MAPPINGS = {
|
|
"LlamaCppLoader": "Llama-CPP Loader",
|
|
"LlamaCppGenerate": "Llama-CPP Generate",
|
|
"LlamaCppUnload": "Llama-CPP Unload VRAM",
|
|
"LlamaCppSafeUnload": "Llama-CPP Safe Unload",
|
|
"LlamaCppMemoryInfo": "Llama-CPP Memory Info",
|
|
"LlamaCppCompleteUnload": "Llama-CPP Complete Unload",
|
|
"LlamaCppAIO": "🚀 Llama-CPP All-In-One",
|
|
}
|
|
print("✅ LlamaCppText loaded")
|
|
except Exception as e:
|
|
print(f"⚠️ LlamaCppText failed: {e}")
|
|
LLAMA_MAPPINGS = {}
|
|
LLAMA_DISPLAY_MAPPINGS = {}
|
|
|
|
try:
|
|
# 3. Multi ControlNet Loader
|
|
from .multi_controlnet_loader import NODE_CLASS_MAPPINGS as CONTROLNET_MAPPINGS
|
|
from .multi_controlnet_loader import NODE_DISPLAY_NAME_MAPPINGS as CONTROLNET_DISPLAY_MAPPINGS
|
|
print("✅ MultiControlNetLoader loaded")
|
|
except Exception as e:
|
|
print(f"⚠️ MultiControlNetLoader failed: {e}")
|
|
CONTROLNET_MAPPINGS = {}
|
|
CONTROLNET_DISPLAY_MAPPINGS = {}
|
|
|
|
try:
|
|
# 4. Prompt Translator
|
|
from .prompt_translator import NODE_CLASS_MAPPINGS as TRANSLATOR_MAPPINGS
|
|
from .prompt_translator import NODE_DISPLAY_NAME_MAPPINGS as TRANSLATOR_DISPLAY_MAPPINGS
|
|
print("✅ PromptTranslator loaded")
|
|
except Exception as e:
|
|
print(f"⚠️ PromptTranslator failed: {e}")
|
|
TRANSLATOR_MAPPINGS = {}
|
|
TRANSLATOR_DISPLAY_MAPPINGS = {}
|
|
|
|
try:
|
|
# 5. Locale Toggle
|
|
from .locale_toggle import NODE_CLASS_MAPPINGS as LOCALE_MAPPINGS
|
|
from .locale_toggle import NODE_DISPLAY_NAME_MAPPINGS as LOCALE_DISPLAY_MAPPINGS
|
|
print("✅ LocaleToggle loaded")
|
|
except Exception as e:
|
|
print(f"⚠️ LocaleToggle failed: {e}")
|
|
LOCALE_MAPPINGS = {}
|
|
LOCALE_DISPLAY_MAPPINGS = {}
|
|
|
|
# Setup LlamaCpp routes (from working version)
|
|
def setup_llama_routes():
|
|
"""Setup llama generation routes using ComfyUI's standard method."""
|
|
try:
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
import json
|
|
|
|
routes = PromptServer.instance.routes
|
|
|
|
@routes.post('/llama_cpp/generate')
|
|
async def handle_llama_generation(request):
|
|
try:
|
|
print(f"🔄 Llama generation request received")
|
|
|
|
# Get form data
|
|
data = await request.post()
|
|
|
|
instruction_prompt = data.get('instruction_prompt', '')
|
|
input_text = data.get('input_text', '')
|
|
model_path = data.get('model_path', '')
|
|
max_tokens = int(data.get('max_tokens', '128'))
|
|
temperature = float(data.get('temperature', '0.8'))
|
|
top_k = int(data.get('top_k', '40'))
|
|
top_p = float(data.get('top_p', '0.95'))
|
|
repeat_penalty = float(data.get('repeat_penalty', '1.1'))
|
|
n_ctx = int(data.get('n_ctx', '2048'))
|
|
n_gpu_layers = int(data.get('n_gpu_layers', '40'))
|
|
nsfw_filter = data.get('nsfw_filter', 'true').lower() == 'true'
|
|
auto_unload = data.get('auto_unload', 'true').lower() == 'true'
|
|
|
|
print(f"📋 Request data: model={model_path}, text='{input_text[:50]}...', max_tokens={max_tokens}")
|
|
|
|
if not input_text.strip() and not instruction_prompt.strip():
|
|
return web.json_response({'error': 'No input text provided'}, status=400)
|
|
|
|
# Import llama functions
|
|
from .llama_cpp_node import _get_or_create_model, _cleanup_memory
|
|
import random
|
|
|
|
# Load model
|
|
try:
|
|
llm = _get_or_create_model(model_path, n_ctx, n_gpu_layers, False)
|
|
print(f"✅ Model loaded for API: {model_path}")
|
|
except Exception as e:
|
|
print(f"❌ Model loading failed: {e}")
|
|
return web.json_response({'error': f'Model loading failed: {str(e)}'}, status=500)
|
|
|
|
# Generate text
|
|
try:
|
|
# Build prompt（指示と内容を明確に分離）
|
|
instruction = instruction_prompt if instruction_prompt else ""
|
|
content = input_text if input_text else "A beautiful scene"
|
|
|
|
# デフォルト指示を設定（ユーザー指示がない場合）
|
|
if not instruction.strip():
|
|
instruction = """You are a prompt generator for Stable Diffusion.
|
|
|
|
Task:
|
|
Convert the input description into ONE optimized English prompt for Stable Diffusion.
|
|
Rules:
|
|
- Output format: single line of comma-separated phrases only
|
|
- Stay faithful to the original description's content and intent
|
|
- Include: subject, setting/environment, lighting, mood/atmosphere, artistic style, quality tags
|
|
- Use descriptive English terms that accurately represent the input
|
|
- Do not censor or sanitize the content - convert faithfully
|
|
- End cleanly after final quality tag
|
|
|
|
Output:
|
|
<prompt>"""
|
|
|
|
# ランダムなバリエーション用のリスト（適度に調整）
|
|
variation_prompts = [
|
|
"Focus on artistic composition and visual details.",
|
|
"Emphasize mood and atmospheric elements.",
|
|
"Include creative styling and artistic elements.",
|
|
"Add detailed lighting and visual effects."
|
|
]
|
|
|
|
quality_tags = [
|
|
"masterpiece, best quality, highly detailed",
|
|
"high quality, detailed artwork",
|
|
"ultra detailed, premium quality",
|
|
"best quality, intricate details"
|
|
]
|
|
|
|
# 内部固定システムプロンプト
|
|
system_prompt = "You are a creative writing assistant. Respond in English."
|
|
|
|
# Apply NSFW filter if requested
|
|
if nsfw_filter:
|
|
safety_note = " IMPORTANT: If inappropriate content, respond with 'I cannot generate prompts for inappropriate content.' Otherwise, respond ONLY in English."
|
|
wrapped_prompt = f"""{system_prompt}
|
|
|
|
{instruction}{safety_note}
|
|
|
|
Input to convert: {content}
|
|
|
|
Response (English only):"""
|
|
else:
|
|
# 最小限のランダム性のみ追加
|
|
if random.choice([True, False]):
|
|
quality_hint = random.choice(quality_tags)
|
|
wrapped_prompt = f"""{system_prompt}
|
|
|
|
{instruction}
|
|
|
|
Input to convert: {content}
|
|
|
|
Hint: Consider quality tags like "{quality_hint}"
|
|
|
|
Response (English only):"""
|
|
else:
|
|
wrapped_prompt = f"""{system_prompt}
|
|
|
|
{instruction}
|
|
|
|
Input to convert: {content}
|
|
|
|
Response (English only):"""
|
|
|
|
# Generate
|
|
seed = random.randint(1, 2**31 - 1)
|
|
response = llm.create_completion(
|
|
prompt=wrapped_prompt,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
repeat_penalty=repeat_penalty,
|
|
seed=seed,
|
|
stream=False,
|
|
echo=False,
|
|
stop=["</s>"]
|
|
)
|
|
|
|
if isinstance(response, dict) and "choices" in response:
|
|
result_text = response["choices"][0]["text"]
|
|
else:
|
|
result_text = str(response)
|
|
|
|
# Clean up result
|
|
result_text = result_text.strip()
|
|
for stop_token in ["</s>", "<eos>"]:
|
|
if stop_token in result_text:
|
|
result_text = result_text.split(stop_token)[0]
|
|
|
|
if not result_text or result_text.strip() == "":
|
|
result_text = "レスポンスが生成されませんでした。"
|
|
|
|
print(f"✅ Generation completed: {len(result_text)} characters")
|
|
|
|
# Auto unload if requested
|
|
if auto_unload:
|
|
print("🧹 Auto unloading VRAM...")
|
|
try:
|
|
_cleanup_memory()
|
|
print("✅ Auto unload completed")
|
|
except Exception as cleanup_e:
|
|
print(f"⚠️ Auto unload warning: {cleanup_e}")
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'generated_text': result_text,
|
|
'input_text': f"{instruction}\n\n{content}",
|
|
'token_count': len(result_text.split()),
|
|
'model_path': model_path,
|
|
'auto_unload': auto_unload
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"❌ Generation error: {e}")
|
|
return web.json_response({'error': f'Generation failed: {str(e)}'}, status=500)
|
|
|
|
except Exception as e:
|
|
print(f"❌ API error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
@routes.post('/llama_cpp/unload')
|
|
async def handle_llama_unload(request):
|
|
try:
|
|
print(f"🗑️ Llama unload request received")
|
|
|
|
# Import cleanup functions
|
|
from .llama_cpp_node import _cleanup_memory, _force_cleanup_memory
|
|
|
|
# Execute force memory cleanup for better VRAM release
|
|
_force_cleanup_memory()
|
|
|
|
print(f"✅ VRAM cleanup completed")
|
|
|
|
return web.json_response({
|
|
'success': True,
|
|
'message': 'VRAM has been successfully freed. All models have been unloaded from memory.'
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"❌ Unload API error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({
|
|
'success': False,
|
|
'error': str(e)
|
|
}, status=500)
|
|
|
|
print("✅ Llama-CPP routes registered")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Could not register llama routes: {e}")
|
|
|
|
# Setup routes when server is available
|
|
try:
|
|
setup_llama_routes()
|
|
print("🚀 BuimenLabo: Routes setup completed successfully")
|
|
except Exception as e:
|
|
print(f"⚠️ BuimenLabo: Routes setup failed: {e}")
|
|
|
|
# Combine all node mappings with BuimenLabo prefix
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# Add with prefixes to avoid conflicts
|
|
for name, cls in GEMINI_MAPPINGS.items():
|
|
NODE_CLASS_MAPPINGS[f"BuimenLabo_{name}"] = cls
|
|
for name, display in GEMINI_DISPLAY_MAPPINGS.items():
|
|
NODE_DISPLAY_NAME_MAPPINGS[f"BuimenLabo_{name}"] = f"🔍 BuimenLabo {display}"
|
|
|
|
for name, cls in LLAMA_MAPPINGS.items():
|
|
NODE_CLASS_MAPPINGS[f"BuimenLabo_{name}"] = cls
|
|
for name, display in LLAMA_DISPLAY_MAPPINGS.items():
|
|
NODE_DISPLAY_NAME_MAPPINGS[f"BuimenLabo_{name}"] = f"🦙 BuimenLabo {display}"
|
|
|
|
for name, cls in CONTROLNET_MAPPINGS.items():
|
|
NODE_CLASS_MAPPINGS[f"BuimenLabo_{name}"] = cls
|
|
for name, display in CONTROLNET_DISPLAY_MAPPINGS.items():
|
|
NODE_DISPLAY_NAME_MAPPINGS[f"BuimenLabo_{name}"] = f"🎛️ BuimenLabo {display}"
|
|
|
|
for name, cls in TRANSLATOR_MAPPINGS.items():
|
|
NODE_CLASS_MAPPINGS[f"BuimenLabo_{name}"] = cls
|
|
for name, display in TRANSLATOR_DISPLAY_MAPPINGS.items():
|
|
NODE_DISPLAY_NAME_MAPPINGS[f"BuimenLabo_{name}"] = f"🌐 BuimenLabo {display}"
|
|
|
|
for name, cls in LOCALE_MAPPINGS.items():
|
|
NODE_CLASS_MAPPINGS[f"BuimenLabo_{name}"] = cls
|
|
for name, display in LOCALE_DISPLAY_MAPPINGS.items():
|
|
NODE_DISPLAY_NAME_MAPPINGS[f"BuimenLabo_{name}"] = f"🌍 BuimenLabo {display}"
|
|
|
|
# Set WEB_DIRECTORY for JavaScript files
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
# Package information
|
|
print("📦 ComfyUI BuimenLabo package loaded successfully!")
|
|
print(f"🎯 Total nodes: {len(NODE_CLASS_MAPPINGS)}")
|
|
print("📝 Blog: https://note.com/hirodream44")
|
|
print("🔗 Features: Multi-language, AI pose analysis, LLM generation, Smart ControlNet, Prompt translation")
|
|
|
|
# Export for ComfyUI
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] |