313 lines
12 KiB
Python
313 lines
12 KiB
Python
import subprocess
|
|
import os
|
|
import sys
|
|
import requests
|
|
import json
|
|
import random
|
|
import uuid
|
|
import time
|
|
from .base import tool
|
|
from ..core.config import load_config
|
|
from ..ui.formatter import print_status, console
|
|
|
|
@tool(
|
|
name="get_comfy_status",
|
|
description="Check the status of image generation jobs or the ComfyUI queue",
|
|
parameters={
|
|
"prompt_id": {"type": "string", "description": "Optional: Specific prompt ID to check", "default": ""}
|
|
}
|
|
)
|
|
def get_comfy_status(prompt_id: str = "") -> str:
|
|
"""Queries the ComfyUI API for queue or history status"""
|
|
config = load_config()
|
|
comfy_url = config.get("comfy_url", "http://127.0.0.1:8188").rstrip("/")
|
|
|
|
try:
|
|
if prompt_id:
|
|
# Check history for specific prompt
|
|
resp = requests.get(f"{comfy_url}/history/{prompt_id}", timeout=5)
|
|
if resp.status_code == 200:
|
|
history = resp.json()
|
|
if prompt_id in history:
|
|
return f"Status for {prompt_id}: Completed. Result saved."
|
|
|
|
# Check queue if not in history
|
|
resp = requests.get(f"{comfy_url}/queue", timeout=5)
|
|
if resp.status_code == 200:
|
|
queue_data = resp.json()
|
|
for job in queue_data.get("queue_running", []):
|
|
if job[1] == prompt_id: return f"Status for {prompt_id}: Currently Running."
|
|
for job in queue_data.get("queue_pending", []):
|
|
if job[1] == prompt_id: return f"Status for {prompt_id}: Pending in Queue."
|
|
|
|
return f"Status for {prompt_id}: ID not found in current queue or history."
|
|
else:
|
|
# General queue status
|
|
resp = requests.get(f"{comfy_url}/queue", timeout=5)
|
|
if resp.status_code == 200:
|
|
queue_data = resp.json()
|
|
running = len(queue_data.get("queue_running", []))
|
|
pending = len(queue_data.get("queue_pending", []))
|
|
return f"ComfyUI Queue Status: {running} running, {pending} pending."
|
|
return f"Error: Could not retrieve queue from {comfy_url}"
|
|
except Exception as e:
|
|
return f"Error connecting to ComfyUI: {str(e)}"
|
|
|
|
@tool(
|
|
name="generate_image",
|
|
description="Generate an image using the local Stable Diffusion workflow (ComfyUI)",
|
|
parameters={
|
|
"prompt": {"type": "string", "description": "The image description"}
|
|
}
|
|
)
|
|
def generate_image(prompt: str) -> str:
|
|
"""Invokes the ComfyUI API with WebSocket for realtime progress tracking"""
|
|
config = load_config()
|
|
comfy_url = config.get("comfy_url", "http://127.0.0.1:8188").rstrip("/")
|
|
image_model = config.get("image_model", "dreamshaper_8_pruned.safetensors")
|
|
|
|
# Determine dimensions (SDXL prefers 1024x1024)
|
|
width, height = (512, 512)
|
|
if "xl" in image_model.lower():
|
|
width, height = (1024, 1024)
|
|
|
|
client_id = str(uuid.uuid4())
|
|
|
|
try:
|
|
import websocket
|
|
except ImportError:
|
|
return "Error: 'websocket-client' not installed. Please run 'pip install websocket-client'."
|
|
|
|
try:
|
|
# Define the workflow directly
|
|
workflow = {
|
|
"1": {
|
|
"inputs": {"ckpt_name": image_model},
|
|
"class_type": "CheckpointLoaderSimple"
|
|
},
|
|
"2": {
|
|
"inputs": {
|
|
"text": f"masterpiece, best quality, highly detailed, {prompt}",
|
|
"clip": ["1", 1]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"3": {
|
|
"inputs": {
|
|
"text": "blurry, low quality, distorted, watermark, nsfw",
|
|
"clip": ["1", 1]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"4": {
|
|
"inputs": {"width": width, "height": height, "batch_size": 1},
|
|
"class_type": "EmptyLatentImage"
|
|
},
|
|
"5": {
|
|
"inputs": {
|
|
"seed": random.randint(1, 1000000000), "steps": 20, "cfg": 7.0,
|
|
"sampler_name": "dpmpp_2m", "scheduler": "karras",
|
|
"denoise": 1.0, "model": ["1", 0],
|
|
"positive": ["2", 0], "negative": ["3", 0],
|
|
"latent_image": ["4", 0]
|
|
},
|
|
"class_type": "KSampler"
|
|
},
|
|
"6": {
|
|
"inputs": {"samples": ["5", 0], "vae": ["1", 2]},
|
|
"class_type": "VAEDecode"
|
|
},
|
|
"7": {
|
|
"inputs": {
|
|
"filename_prefix": "ollama_cli_",
|
|
"images": ["6", 0]
|
|
},
|
|
"class_type": "SaveImage"
|
|
}
|
|
}
|
|
|
|
# Connect to WebSocket
|
|
ws_url = comfy_url.replace("http://", "ws://").replace("https://", "wss://") + f"/ws?clientId={client_id}"
|
|
ws = websocket.WebSocket()
|
|
ws.connect(ws_url)
|
|
|
|
# Queue the prompt
|
|
payload = {"prompt": workflow, "client_id": client_id}
|
|
response = requests.post(f"{comfy_url}/prompt", json=payload, timeout=10)
|
|
|
|
if response.status_code != 200:
|
|
return f"Error from ComfyUI: {response.status_code} - {response.text}"
|
|
|
|
data = response.json()
|
|
prompt_id = data.get('prompt_id')
|
|
|
|
print_status(f"Realtime monitoring active (Prompt ID: {prompt_id})")
|
|
|
|
found_file = None
|
|
from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn
|
|
|
|
with Progress(
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
|
TimeElapsedColumn(),
|
|
console=console
|
|
) as progress:
|
|
task = progress.add_task("[cyan]Queued...", total=20) # 20 steps default
|
|
|
|
while True:
|
|
out = ws.recv()
|
|
if not out: break
|
|
|
|
if isinstance(out, str):
|
|
msg = json.loads(out)
|
|
if msg['type'] == 'progress':
|
|
p_data = msg['data']
|
|
if p_data['prompt_id'] == prompt_id:
|
|
progress.update(task, description="[green]Sampling...", completed=p_data['value'], total=p_data['max'])
|
|
|
|
if msg['type'] == 'executing':
|
|
e_data = msg['data']
|
|
if e_data['prompt_id'] == prompt_id:
|
|
node = e_data['node']
|
|
if node is None: # Execution finished
|
|
break
|
|
# Update description based on node type if possible (or just generic)
|
|
progress.update(task, description=f"[yellow]Node {node}...")
|
|
|
|
if msg['type'] == 'executed':
|
|
ex_data = msg['data']
|
|
if ex_data['prompt_id'] == prompt_id:
|
|
# We can extract the filename here!
|
|
if 'output' in ex_data and 'images' in ex_data['output']:
|
|
images = ex_data['output']['images']
|
|
if images:
|
|
img_name = images[0]['filename']
|
|
output_path = os.path.expanduser(config.get("comfy_output_path", "~/ComfyUI/output"))
|
|
found_file = os.path.join(output_path, img_name)
|
|
else:
|
|
continue # Binary data (previews) ignored for now
|
|
|
|
ws.close()
|
|
|
|
if found_file and os.path.exists(found_file):
|
|
if sys.platform == "darwin":
|
|
subprocess.run(["open", found_file])
|
|
else:
|
|
subprocess.run(["xdg-open", found_file])
|
|
return f"Image generation complete!\nSaved to: {found_file}\nOpening image now..."
|
|
|
|
return f"Image generation finished, but could not locate file automatically. Prompt ID: {prompt_id}"
|
|
|
|
except Exception as e:
|
|
return f"Error generating image: {str(e)}"
|
|
|
|
@tool(
|
|
name="analyze_image",
|
|
description="Analyze or describe a local image file using AI vision",
|
|
parameters={
|
|
"image_path": {"type": "string", "description": "Path to the image file"},
|
|
"prompt": {"type": "string", "description": "What to ask about the image", "default": "Describe this image in detail."}
|
|
}
|
|
)
|
|
def analyze_image(image_path: str, prompt: str = "Describe this image in detail.") -> str:
|
|
"""Invokes the Ollama vision API"""
|
|
config = load_config()
|
|
vision_model = config.get("vision_model", "llama3.2-vision")
|
|
|
|
from ..core.ollama import OllamaClient
|
|
client = OllamaClient(config.get("ollama_url"))
|
|
|
|
print_status(f"Analyzing image with {vision_model}...")
|
|
return client.describe_image(image_path, prompt, model=vision_model)
|
|
|
|
@tool(
|
|
name="speak_text",
|
|
description="Convert text to speech using Piper TTS (high quality)",
|
|
parameters={
|
|
"text": {"type": "string", "description": "Text to speak"},
|
|
"output_file": {"type": "string", "description": "Output filename", "default": "output.wav"}
|
|
}
|
|
)
|
|
def speak_text(text: str, output_file: str = "output.wav") -> str:
|
|
"""Uses the best available TTS engine (Kokoro > Piper > System)"""
|
|
config = load_config()
|
|
|
|
# 1. Try Kokoro (Ultra-Realistic)
|
|
kokoro_python = os.path.expanduser(config.get("kokoro_path", "~/.ollama-cli/kokoro/venv/bin/python"))
|
|
if os.path.exists(kokoro_python):
|
|
try:
|
|
# Inline script to run in Kokoro venv
|
|
kokoro_script = f"""
|
|
import sys
|
|
import soundfile as sf
|
|
import os
|
|
try:
|
|
from kokoro import KPipeline
|
|
import torch
|
|
except ImportError:
|
|
print("Dependencies missing")
|
|
sys.exit(1)
|
|
|
|
text = sys.argv[1]
|
|
output_file = sys.argv[2]
|
|
|
|
# Initialize pipeline
|
|
pipeline = KPipeline(lang_code='a')
|
|
generator = pipeline(text, voice='af_bella', speed=1)
|
|
|
|
all_audio = []
|
|
for gs, ps, audio in generator:
|
|
if audio is not None:
|
|
all_audio.append(audio)
|
|
|
|
if all_audio:
|
|
import numpy as np
|
|
full_audio = np.concatenate(all_audio)
|
|
sf.write(output_file, full_audio, 24000)
|
|
print("Success")
|
|
else:
|
|
print("No audio generated")
|
|
sys.exit(1)
|
|
"""
|
|
# Run in venv
|
|
subprocess.run([kokoro_python, "-c", kokoro_script, text, output_file],
|
|
check=True, capture_output=True)
|
|
|
|
if sys.platform == "darwin":
|
|
subprocess.run(["afplay", output_file])
|
|
return f"Audio saved to {output_file} and played (using Ultra-Realistic Kokoro TTS)"
|
|
except Exception as e:
|
|
# Continue to fallback
|
|
pass
|
|
|
|
# 2. Try Piper (High-Quality)
|
|
piper_config = config.get("piper", {})
|
|
piper_bin = os.path.expanduser(piper_config.get("path", "~/.ollama-cli/piper/piper"))
|
|
model_path = os.path.expanduser(piper_config.get("model", "~/.ollama-cli/piper/voice.onnx"))
|
|
|
|
if os.path.exists(piper_bin) and os.path.exists(model_path):
|
|
try:
|
|
# Piper reads from stdin and writes to file
|
|
cmd = f'echo "{text}" | {piper_bin} --model {model_path} --output_file {output_file}'
|
|
subprocess.run(cmd, shell=True, check=True, capture_output=True)
|
|
|
|
# Auto-play on macOS
|
|
if sys.platform == "darwin":
|
|
subprocess.run(["afplay", output_file])
|
|
|
|
return f"Audio saved to {output_file} and played (using high-quality Piper TTS)"
|
|
except Exception as e:
|
|
pass
|
|
|
|
# 3. Fallback to system 'say' command on macOS
|
|
if sys.platform == "darwin":
|
|
try:
|
|
subprocess.run(["say", "-o", output_file, text], check=True)
|
|
subprocess.run(["afplay", output_file])
|
|
return f"Audio saved to {output_file} and played (using system 'say' fallback)"
|
|
except Exception as e:
|
|
return f"Error using 'say': {e}"
|
|
|
|
return "TTS engine not found. Please run the installer to set up Piper or Kokoro."
|