Admin/ollama-cli
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhancement Suite: Anti-loop protection, Kokoro TTS integration, ComfyUI status monitoring, and auto-open image functionality

Browse Source
This commit is contained in:
Paul Bailey
2026-02-23 00:59:20 +00:00
parent d7efe008b8
commit 90fcba28bc
41 changed files with 1982 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

175
src/ollama_cli/tools/media.py
View File

@@ -7,6 +7,56 @@ import random
from .base import tool
from ..core.config import load_config
@tool(
name="generate_image",
description="Generate an image using the local Stable Diffusion workflow (ComfyUI)",
parameters={
"prompt": {"type": "string", "description": "The image description"}
}
)
@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)",
@@ -18,7 +68,13 @@ def generate_image(prompt: str) -> str:
"""Invokes the ComfyUI API directly with a predefined workflow"""
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)
try:
# Check if ComfyUI is running
try:
@@ -27,10 +83,9 @@ def generate_image(prompt: str) -> str:
return f"Error: ComfyUI server ({comfy_url}) appears to be down. Please start it first or configure the URL with '/config comfy <url>'."
# Define the workflow directly
# Based on the user's previous generate_image.py, but using the configured URL
workflow = {
"1": {
"inputs": {"ckpt_name": "dreamshaper_8_pruned.safetensors"},
"inputs": {"ckpt_name": image_model},
"class_type": "CheckpointLoaderSimple"
},
"2": {
@@ -48,7 +103,7 @@ def generate_image(prompt: str) -> str:
"class_type": "CLIPTextEncode"
},
"4": {
"inputs": {"width": 512, "height": 512, "batch_size": 1},
"inputs": {"width": width, "height": height, "batch_size": 1},
"class_type": "EmptyLatentImage"
},
"5": {
@@ -80,7 +135,44 @@ def generate_image(prompt: str) -> str:
if response.status_code == 200:
data = response.json()
prompt_id = data.get('prompt_id', 'unknown')
return f"Image generation queued successfully on {comfy_url}.\nPrompt ID: {prompt_id}\nCheck your ComfyUI output folder for 'ollama_cli_' prefixed files."
# Start monitoring for the new image
output_path = os.path.expanduser(config.get("comfy_output_path", "~/ComfyUI/output"))
# Helper to find the newest file with our prefix
def find_newest_image():
import glob
files = glob.glob(os.path.join(output_path, "ollama_cli_*"))
if not files: return None
return max(files, key=os.path.getctime)
initial_newest = find_newest_image()
# Polling loop (max 60 seconds)
import time
found_file = None
for _ in range(30):
# Check status via API
status_check = get_comfy_status(prompt_id)
if "Completed" in status_check:
# Give filesystem a moment to sync
time.sleep(1)
current_newest = find_newest_image()
if current_newest and current_newest != initial_newest:
found_file = current_newest
break
time.sleep(2)
if found_file:
# Open the image automatically
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 queued successfully on {comfy_url}.\nPrompt ID: {prompt_id}\nMonitoring timed out, but the image should appear in {output_path} soon."
else:
return f"Error form ComfyUI: {response.status_code} - {response.text}"
@@ -96,26 +188,83 @@ def generate_image(prompt: str) -> str:
}
)
def speak_text(text: str, output_file: str = "output.wav") -> str:
"""Uses the high-quality Piper TTS installed by the installer"""
piper_dir = os.path.expanduser("~/.ollama-cli/piper")
piper_bin = os.path.join(piper_dir, "piper")
model_path = os.path.join(piper_dir, "voice.onnx")
"""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)
return f"Audio saved to {output_file} (using high-quality Piper TTS)"
# 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:
return f"Error using Piper: {e}. Falling back to system TTS..."
pass
# Fallback to simple 'say' command on macOS
# 3. Fallback to system 'say' command on macOS
if sys.platform == "darwin":
try:
subprocess.run(["say", "-o", output_file, text], check=True)
return f"Audio saved to {output_file} (using system 'say' fallback)"
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."
return "No compatible TTS engine found. Please run the installer."