265 lines
11 KiB
Python
265 lines
11 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import base64
|
|
import io
|
|
import numpy as np
|
|
import aiohttp
|
|
from PIL import Image
|
|
from comfy_api.latest import ComfyExtension, io as comfy_io, ui as comfy_ui
|
|
|
|
try:
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
except ImportError:
|
|
PromptServer = None
|
|
|
|
def get_prompt_files():
|
|
"""Scans the 'prompts' directory for .txt files and returns a dictionary."""
|
|
prompt_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'prompts')
|
|
prompt_files = {}
|
|
if os.path.exists(prompt_dir):
|
|
for filename in os.listdir(prompt_dir):
|
|
if filename.endswith(".txt"):
|
|
filepath = os.path.join(prompt_dir, filename)
|
|
try:
|
|
with open(filepath, 'r', encoding='utf-8') as f:
|
|
key = os.path.splitext(filename)[0]
|
|
prompt_files[key] = f.read().strip()
|
|
except Exception as e:
|
|
print(f"BasicOllama: Error reading prompt file {filename}: {e}")
|
|
return prompt_files
|
|
|
|
def rgba_to_rgb(image):
|
|
"""Convert RGBA image to RGB with white background"""
|
|
if image.mode == 'RGBA':
|
|
background = Image.new("RGB", image.size, (255, 255, 255))
|
|
image = Image.alpha_composite(background.convert("RGBA"), image).convert("RGB")
|
|
return image
|
|
|
|
def tensor_to_pil(tensor):
|
|
"""Convert a [H, W, C] tensor to PIL Image"""
|
|
if tensor.dim() == 4:
|
|
tensor = tensor[0]
|
|
|
|
image_np = tensor.cpu().numpy()
|
|
image_np = (image_np * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
if image_np.shape[-1] == 1:
|
|
image_np = np.repeat(image_np, 3, axis=-1)
|
|
|
|
mode = 'RGBA' if image_np.shape[-1] == 4 else 'RGB'
|
|
image = Image.fromarray(image_np, mode=mode)
|
|
return rgba_to_rgb(image)
|
|
|
|
def tensor_to_base64(tensor):
|
|
"""Convert tensor to base64 encoded PNG"""
|
|
image = tensor_to_pil(tensor)
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode()
|
|
|
|
def get_ollama_url():
|
|
config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'config.json')
|
|
ollama_url = "http://localhost:11434"
|
|
if os.path.exists(config_path):
|
|
try:
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
ollama_url = config.get("OLLAMA_URL", ollama_url)
|
|
except Exception as e:
|
|
print(f"BasicOllama: Error reading config.json: {e}")
|
|
return ollama_url
|
|
|
|
async def fetch_ollama_models():
|
|
url = f"{get_ollama_url()}/api/tags"
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(url, timeout=2) as response:
|
|
if response.status == 200:
|
|
data = await response.json()
|
|
models = data.get('models', [])
|
|
return [model['name'] for model in models]
|
|
except Exception:
|
|
pass
|
|
return ["Start Ollama and Refresh"]
|
|
|
|
class BasicOllama(comfy_io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> comfy_io.Schema:
|
|
prompt_templates = get_prompt_files()
|
|
prompt_structures = list(prompt_templates.keys())
|
|
if not prompt_structures:
|
|
prompt_structures = ["None"]
|
|
|
|
return comfy_io.Schema(
|
|
node_id="BasicOllama",
|
|
display_name="Basic Ollama",
|
|
category="Ollama",
|
|
description="Interface with Ollama LLMs, supporting dynamic image inputs.",
|
|
is_output_node=False,
|
|
inputs=[
|
|
comfy_io.String.Input("prompt", multiline=True, tooltip="The main text prompt for the LLM."),
|
|
comfy_io.MultiType.Input(
|
|
comfy_io.Combo.Input("ollama_model",
|
|
options=["Loading..."],
|
|
default="Loading...",
|
|
tooltip="Select the Ollama model to use.",
|
|
socketless=True,
|
|
remote=comfy_io.RemoteOptions(
|
|
route="/basic_ollama/models",
|
|
refresh_button=True
|
|
)),
|
|
types=[comfy_io.String]
|
|
),
|
|
comfy_io.Int.Input(
|
|
"generation_seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xffffffffffffffff,
|
|
step=1,
|
|
control_after_generate=True,
|
|
tooltip="Seed for randomness. Change this to force a new generation."
|
|
),
|
|
comfy_io.Boolean.Input("enable_think", default=True, tooltip="Enable reasoning/thinking for supported models."),
|
|
comfy_io.Combo.Input(
|
|
"saved_sys_prompt",
|
|
options=prompt_structures,
|
|
default=prompt_structures[0],
|
|
tooltip="Select a pre-defined system prompt."
|
|
),
|
|
comfy_io.Boolean.Input("use_sys_prompt_below", default=False, tooltip="If True, use the custom system prompt below instead of the saved one."),
|
|
comfy_io.String.Input("system_prompt", multiline=True, tooltip="Custom system prompt to guide the LLM's behavior."),
|
|
comfy_io.Autogrow.Input("images",
|
|
comfy_io.Autogrow.TemplatePrefix(comfy_io.Image.Input("image", optional=True), prefix="image", min=1, max=10),
|
|
optional=True,
|
|
extra_dict={"autoshrink": True},
|
|
tooltip="Dynamic image inputs for multimodal models.")
|
|
],
|
|
outputs=[
|
|
comfy_io.String.Output(id="text", display_name="text", tooltip="The text generated by the LLM."),
|
|
comfy_io.String.Output(id="thinking", display_name="thinking", tooltip="The reasoning trace (if available).")
|
|
],
|
|
hidden=[comfy_io.Hidden.unique_id]
|
|
)
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, **kwargs) -> bool | str:
|
|
return True
|
|
|
|
@classmethod
|
|
async def execute(cls, prompt, ollama_model, generation_seed, enable_think, use_sys_prompt_below, saved_sys_prompt, system_prompt, images=None, **kwargs) -> comfy_io.NodeOutput:
|
|
ollama_url = get_ollama_url()
|
|
url = f"{ollama_url}/api/generate"
|
|
|
|
system_prompt_content = ""
|
|
if use_sys_prompt_below:
|
|
system_prompt_content = system_prompt
|
|
else:
|
|
prompt_templates = get_prompt_files()
|
|
if saved_sys_prompt in prompt_templates:
|
|
system_prompt_content = prompt_templates[saved_sys_prompt]
|
|
|
|
payload = {
|
|
"model": ollama_model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"keep_alive": "0m",
|
|
"options": {
|
|
"seed": generation_seed
|
|
}
|
|
}
|
|
|
|
# Handle 'think' parameter logic
|
|
payload["think"] = enable_think
|
|
|
|
if system_prompt_content:
|
|
payload["system"] = system_prompt_content
|
|
|
|
# Collect all images from the 'images' dict (Autogrow) and any stray images in kwargs
|
|
provided_images = []
|
|
|
|
# Check the packed 'images' dictionary
|
|
if isinstance(images, dict):
|
|
for key in sorted(images.keys()):
|
|
img_tensor = images[key]
|
|
if isinstance(img_tensor, torch.Tensor):
|
|
for i in range(img_tensor.shape[0]):
|
|
provided_images.append(img_tensor[i])
|
|
|
|
# Check kwargs for any stray image inputs
|
|
for k, v in kwargs.items():
|
|
if k.startswith("image") and isinstance(v, torch.Tensor):
|
|
# Avoid duplicates if already in 'images' dict
|
|
if not isinstance(images, dict) or k not in images:
|
|
for i in range(v.shape[0]):
|
|
provided_images.append(v[i])
|
|
|
|
if provided_images:
|
|
image_data = [tensor_to_base64(img) for img in provided_images]
|
|
payload["images"] = image_data
|
|
|
|
try:
|
|
# Set a timeout (10 minutes) for model inference
|
|
timeout = aiohttp.ClientTimeout(total=600)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.post(url, json=payload) as response:
|
|
if response.status != 200:
|
|
error_text = await response.text()
|
|
# Return empty strings on error to match output signature
|
|
return comfy_io.NodeOutput(f"Error: {response.status} - {error_text}", "")
|
|
|
|
data = await response.json()
|
|
text_output = data.get('response', '')
|
|
thinking_output = data.get('thinking', '')
|
|
|
|
if text_output.strip():
|
|
text_output = cls.clean_text(text_output)
|
|
|
|
return comfy_io.NodeOutput(text_output, thinking_output)
|
|
|
|
except Exception as e:
|
|
return comfy_io.NodeOutput(f"Connection Error: {e}. Ensure Ollama is running at {ollama_url}")
|
|
|
|
@staticmethod
|
|
def clean_text(text):
|
|
clean_text = text.strip()
|
|
if clean_text.startswith("```") and "```" in clean_text[3:]:
|
|
first_block_end = clean_text.find("```", 3)
|
|
if first_block_end > 3:
|
|
language_line_end = clean_text.find("\n", 3)
|
|
if language_line_end > 3 and language_line_end < first_block_end:
|
|
clean_text = clean_text[language_line_end+1:first_block_end].strip()
|
|
else:
|
|
clean_text = clean_text[3:first_block_end].strip()
|
|
|
|
if (clean_text.startswith('"') and clean_text.endswith('"')) or (clean_text.startswith("'") and clean_text.endswith("'")):
|
|
clean_text = clean_text[1:-1].strip()
|
|
|
|
prefixes_to_remove = ["Prompt:", "PROMPT:", "Generated Prompt:", "Final Prompt:"]
|
|
for prefix in prefixes_to_remove:
|
|
if clean_text.startswith(prefix):
|
|
clean_text = clean_text[len(prefix):].strip()
|
|
break
|
|
|
|
return clean_text
|
|
|
|
async def get_ollama_models_endpoint(request):
|
|
models = await fetch_ollama_models()
|
|
return web.json_response(models)
|
|
|
|
if PromptServer:
|
|
try:
|
|
@PromptServer.instance.routes.get("/basic_ollama/models")
|
|
async def get_ollama_models_endpoint_handler(request):
|
|
return await get_ollama_models_endpoint(request)
|
|
except Exception as e:
|
|
print(f"BasicOllama: Could not register API route: {e}")
|
|
|
|
class BasicOllamaExtension(ComfyExtension):
|
|
async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
|
|
return [BasicOllama]
|
|
|
|
async def comfy_entrypoint() -> BasicOllamaExtension:
|
|
return BasicOllamaExtension()
|