koko210Serve
d58be3b33e
- Remove Ollama-specific files (Dockerfile.ollama, entrypoint.sh) - Replace all query_ollama imports and calls with query_llama - Remove langchain-ollama dependency from requirements.txt - Update all utility files (autonomous, kindness, image_generation, etc.) - Update README.md documentation references - Maintain backward compatibility alias in llm.py
403 lines
18 KiB
Python
403 lines
18 KiB
Python
"""
|
|
Image Generation System for Miku Bot
|
|
Natural language detection and ComfyUI integration
|
|
"""
|
|
|
|
import aiohttp
|
|
import asyncio
|
|
import glob
|
|
import json
|
|
import os
|
|
import re
|
|
import tempfile
|
|
import time
|
|
from typing import Optional, Tuple
|
|
import globals
|
|
from utils.llm import query_llama
|
|
|
|
# Image generation detection patterns
|
|
IMAGE_REQUEST_PATTERNS = [
|
|
# Direct requests
|
|
r'\b(?:draw|generate|create|make|show me|paint|sketch|illustrate)\b.*\b(?:image|picture|art|artwork|drawing|painting|illustration)\b',
|
|
r'\b(?:i\s+(?:want|would like|need)\s+(?:to see|an?\s+)?(?:image|picture|art|artwork|drawing|painting|illustration))\b',
|
|
r'\b(?:can you|could you|please)\s+(?:draw|generate|create|make|show me|paint|sketch|illustrate)\b',
|
|
r'\b(?:image|picture|art|artwork|drawing|painting|illustration)\s+of\b',
|
|
|
|
# Visual requests about Miku
|
|
r'\b(?:show me|let me see)\s+(?:you|miku|yourself)\b',
|
|
r'\b(?:what do you look like|how do you look)\b',
|
|
r'\b(?:i\s+(?:want|would like)\s+to see)\s+(?:you|miku|yourself)\b',
|
|
r'\bsee\s+(?:you|miku|yourself)(?:\s+(?:in|with|doing|wearing))?\b',
|
|
|
|
# Activity-based visual requests
|
|
r'\b(?:you|miku|yourself)\s+(?:swimming|dancing|singing|playing|wearing|in|with|doing)\b.*\b(?:pool|water|stage|outfit|clothes|dress)\b',
|
|
r'\b(?:visualize|envision|imagine)\s+(?:you|miku|yourself)\b',
|
|
|
|
# Artistic requests
|
|
r'\b(?:artistic|art|visual)\s+(?:representation|depiction|version)\s+of\s+(?:you|miku|yourself)\b',
|
|
]
|
|
|
|
# Compile patterns for efficiency
|
|
COMPILED_PATTERNS = [re.compile(pattern, re.IGNORECASE) for pattern in IMAGE_REQUEST_PATTERNS]
|
|
|
|
async def detect_image_request(message_content: str) -> Tuple[bool, Optional[str]]:
|
|
"""
|
|
Detect if a message is requesting image generation using natural language.
|
|
|
|
Returns:
|
|
Tuple[bool, Optional[str]]: (is_image_request, extracted_prompt)
|
|
"""
|
|
content = message_content.lower().strip()
|
|
|
|
# Quick rejection for very short messages
|
|
if len(content) < 5:
|
|
return False, None
|
|
|
|
# Check against patterns
|
|
for pattern in COMPILED_PATTERNS:
|
|
if pattern.search(content):
|
|
# Extract the prompt by cleaning up the message
|
|
prompt = extract_image_prompt(message_content)
|
|
return True, prompt
|
|
|
|
return False, None
|
|
|
|
def extract_image_prompt(message_content: str) -> str:
|
|
"""
|
|
Extract and clean the image prompt from the user's message.
|
|
Convert natural language to a proper image generation prompt.
|
|
"""
|
|
content = message_content.strip()
|
|
|
|
# Remove common prefixes that don't help with image generation
|
|
prefixes_to_remove = [
|
|
r'^(?:hey\s+)?miku,?\s*',
|
|
r'^(?:can you|could you|please|would you)\s*',
|
|
r'^(?:i\s+(?:want|would like|need)\s+(?:to see|you to|an?)?)\s*',
|
|
r'^(?:show me|let me see)\s*',
|
|
r'^(?:draw|generate|create|make|paint|sketch|illustrate)\s*(?:me\s*)?(?:an?\s*)?(?:image|picture|art|artwork|drawing|painting|illustration)?\s*(?:of\s*)?',
|
|
]
|
|
|
|
cleaned = content
|
|
for prefix in prefixes_to_remove:
|
|
cleaned = re.sub(prefix, '', cleaned, flags=re.IGNORECASE).strip()
|
|
|
|
# If the cleaned prompt is too short or generic, enhance it
|
|
if len(cleaned) < 10 or cleaned.lower() in ['you', 'yourself', 'miku']:
|
|
cleaned = "Hatsune Miku"
|
|
|
|
# Ensure Miku is mentioned if the user said "you" or "yourself"
|
|
if re.search(r'\b(?:you|yourself)\b', content, re.IGNORECASE) and not re.search(r'\bmiku\b', cleaned, re.IGNORECASE):
|
|
# Replace "you" with "Hatsune Miku" instead of just prepending
|
|
cleaned = re.sub(r'\byou\b', 'Hatsune Miku', cleaned, flags=re.IGNORECASE)
|
|
cleaned = re.sub(r'\byourself\b', 'Hatsune Miku', cleaned, flags=re.IGNORECASE)
|
|
|
|
return cleaned
|
|
|
|
def find_latest_generated_image(prompt_id: str, expected_filename: str = None) -> Optional[str]:
|
|
"""
|
|
Find the most recently generated image in the ComfyUI output directory.
|
|
This handles cases where the exact filename from API doesn't match the file system.
|
|
"""
|
|
output_dirs = [
|
|
"ComfyUI/output",
|
|
"/app/ComfyUI/output"
|
|
]
|
|
|
|
for output_dir in output_dirs:
|
|
if not os.path.exists(output_dir):
|
|
continue
|
|
|
|
try:
|
|
# Get all image files in the directory
|
|
image_extensions = ['.png', '.jpg', '.jpeg', '.webp']
|
|
all_files = []
|
|
|
|
for ext in image_extensions:
|
|
pattern = os.path.join(output_dir, f"*{ext}")
|
|
all_files.extend(glob.glob(pattern))
|
|
|
|
if not all_files:
|
|
continue
|
|
|
|
# Sort by modification time (most recent first)
|
|
all_files.sort(key=os.path.getmtime, reverse=True)
|
|
|
|
# If we have an expected filename, try to find it first
|
|
if expected_filename:
|
|
for file_path in all_files:
|
|
if os.path.basename(file_path) == expected_filename:
|
|
return file_path
|
|
|
|
# Otherwise, return the most recent image (within last 10 minutes)
|
|
recent_threshold = time.time() - 600 # 10 minutes
|
|
for file_path in all_files:
|
|
if os.path.getmtime(file_path) > recent_threshold:
|
|
print(f"🎨 Found recent image: {file_path}")
|
|
return file_path
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Error searching in {output_dir}: {e}")
|
|
continue
|
|
|
|
return None
|
|
|
|
async def generate_image_with_comfyui(prompt: str) -> Optional[str]:
|
|
"""
|
|
Generate an image using ComfyUI with the provided prompt.
|
|
|
|
Args:
|
|
prompt: The image generation prompt
|
|
|
|
Returns:
|
|
Optional[str]: Path to the generated image file, or None if failed
|
|
"""
|
|
try:
|
|
# Load the workflow template
|
|
workflow_path = "Miku_BasicWorkflow.json"
|
|
if not os.path.exists(workflow_path):
|
|
print(f"❌ Workflow template not found: {workflow_path}")
|
|
return None
|
|
|
|
with open(workflow_path, 'r') as f:
|
|
workflow_data = json.load(f)
|
|
|
|
# Replace the prompt placeholder
|
|
workflow_json = json.dumps(workflow_data)
|
|
workflow_json = workflow_json.replace("_POSITIVEPROMPT_", prompt)
|
|
workflow_data = json.loads(workflow_json)
|
|
|
|
# Prepare the request payload
|
|
payload = {"prompt": workflow_data}
|
|
|
|
# Send request to ComfyUI (try different Docker networking options)
|
|
comfyui_urls = [
|
|
"http://host.docker.internal:8188", # Docker Desktop
|
|
"http://172.17.0.1:8188", # Default Docker bridge gateway
|
|
"http://localhost:8188" # Fallback (if network_mode: host)
|
|
]
|
|
|
|
# Try each URL until one works
|
|
comfyui_url = None
|
|
for url in comfyui_urls:
|
|
try:
|
|
async with aiohttp.ClientSession() as test_session:
|
|
timeout = aiohttp.ClientTimeout(total=2)
|
|
async with test_session.get(f"{url}/system_stats", timeout=timeout) as test_response:
|
|
if test_response.status == 200:
|
|
comfyui_url = url
|
|
print(f"✅ ComfyUI found at: {url}")
|
|
break
|
|
except:
|
|
continue
|
|
|
|
if not comfyui_url:
|
|
print(f"❌ ComfyUI not reachable at any of: {comfyui_urls}")
|
|
return None
|
|
async with aiohttp.ClientSession() as session:
|
|
# Submit the generation request
|
|
async with session.post(f"{comfyui_url}/prompt", json=payload) as response:
|
|
if response.status != 200:
|
|
print(f"❌ ComfyUI request failed: {response.status}")
|
|
return None
|
|
|
|
result = await response.json()
|
|
prompt_id = result.get("prompt_id")
|
|
|
|
if not prompt_id:
|
|
print("❌ No prompt_id received from ComfyUI")
|
|
return None
|
|
|
|
print(f"🎨 ComfyUI generation started with prompt_id: {prompt_id}")
|
|
|
|
# Poll for completion (timeout after 5 minutes)
|
|
timeout = 300 # 5 minutes
|
|
start_time = time.time()
|
|
|
|
while time.time() - start_time < timeout:
|
|
# Check if generation is complete
|
|
async with session.get(f"{comfyui_url}/history/{prompt_id}") as hist_response:
|
|
if hist_response.status == 200:
|
|
history = await hist_response.json()
|
|
|
|
if prompt_id in history:
|
|
# Generation complete, find the output image
|
|
outputs = history[prompt_id].get("outputs", {})
|
|
|
|
# Look for image outputs (usually in nodes with "images" key)
|
|
for node_id, node_output in outputs.items():
|
|
if "images" in node_output:
|
|
images = node_output["images"]
|
|
if images:
|
|
# Get the first image
|
|
image_info = images[0]
|
|
filename = image_info["filename"]
|
|
subfolder = image_info.get("subfolder", "")
|
|
|
|
# Construct the full path (adjust for Docker mount)
|
|
if subfolder:
|
|
image_path = os.path.join("ComfyUI", "output", subfolder, filename)
|
|
else:
|
|
image_path = os.path.join("ComfyUI", "output", filename)
|
|
|
|
# Verify the file exists before returning
|
|
if os.path.exists(image_path):
|
|
print(f"✅ Image generated successfully: {image_path}")
|
|
return image_path
|
|
else:
|
|
# Try alternative paths in case of different mounting
|
|
alt_path = os.path.join("/app/ComfyUI/output", filename)
|
|
if os.path.exists(alt_path):
|
|
print(f"✅ Image generated successfully: {alt_path}")
|
|
return alt_path
|
|
else:
|
|
print(f"⚠️ Generated image not found at expected paths: {image_path} or {alt_path}")
|
|
continue
|
|
|
|
# If we couldn't find the image via API, try the fallback method
|
|
print("🔍 Image not found via API, trying fallback method...")
|
|
fallback_image = find_latest_generated_image(prompt_id)
|
|
if fallback_image:
|
|
return fallback_image
|
|
|
|
# Wait before polling again
|
|
await asyncio.sleep(2)
|
|
|
|
print("❌ ComfyUI generation timed out")
|
|
|
|
# Final fallback: look for the most recent image
|
|
print("🔍 Trying final fallback: most recent image...")
|
|
fallback_image = find_latest_generated_image(prompt_id)
|
|
if fallback_image:
|
|
print(f"✅ Found image via fallback method: {fallback_image}")
|
|
return fallback_image
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error in generate_image_with_comfyui: {e}")
|
|
return None
|
|
|
|
async def handle_image_generation_request(message, prompt: str) -> bool:
|
|
"""
|
|
Handle the complete image generation workflow for a user request.
|
|
|
|
Args:
|
|
message: Discord message object
|
|
prompt: Extracted image prompt
|
|
|
|
Returns:
|
|
bool: True if image was successfully generated and sent
|
|
"""
|
|
try:
|
|
# Generate a contextual response about what we're creating
|
|
is_dm = message.guild is None
|
|
guild_id = message.guild.id if message.guild else None
|
|
user_id = str(message.author.id)
|
|
|
|
# Create a response about starting image generation
|
|
response_prompt = f"A user asked you to create an image with this description: '{prompt}'. Respond enthusiastically that you're creating this image for them. Keep it short and excited!"
|
|
|
|
response_type = "dm_response" if is_dm else "server_response"
|
|
initial_response = await query_llama(response_prompt, user_id=user_id, guild_id=guild_id, response_type=response_type)
|
|
|
|
# Send initial response
|
|
initial_msg = await message.channel.send(initial_response)
|
|
|
|
# Start typing to show we're working
|
|
async with message.channel.typing():
|
|
# Generate the image
|
|
print(f"🎨 Starting image generation for prompt: {prompt}")
|
|
image_path = await generate_image_with_comfyui(prompt)
|
|
|
|
if image_path and os.path.exists(image_path):
|
|
# Send the image
|
|
import discord
|
|
with open(image_path, 'rb') as f:
|
|
file = discord.File(f, filename=f"miku_generated_{int(time.time())}.png")
|
|
|
|
# Create a follow-up message about the completed image
|
|
completion_prompt = f"You just finished creating an image based on '{prompt}'. Make a short, excited comment about the completed artwork!"
|
|
completion_response = await query_llama(completion_prompt, user_id=user_id, guild_id=guild_id, response_type=response_type)
|
|
|
|
await message.channel.send(completion_response, file=file)
|
|
|
|
print(f"✅ Image sent successfully to {message.author.display_name}")
|
|
|
|
# Log to DM history if it's a DM
|
|
if is_dm:
|
|
from utils.dm_logger import dm_logger
|
|
dm_logger.log_conversation(user_id, message.content, f"{initial_response}\n[Generated image: {prompt}]", attachments=["generated_image.png"])
|
|
|
|
return True
|
|
else:
|
|
# Image generation failed
|
|
error_prompt = "You tried to create an image but something went wrong with the generation process. Apologize briefly and suggest they try again later."
|
|
error_response = await query_llama(error_prompt, user_id=user_id, guild_id=guild_id, response_type=response_type)
|
|
await message.channel.send(error_response)
|
|
|
|
print(f"❌ Image generation failed for prompt: {prompt}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error in handle_image_generation_request: {e}")
|
|
|
|
# Send error message
|
|
try:
|
|
await message.channel.send("Sorry, I had trouble creating that image. Please try again later!")
|
|
except:
|
|
pass
|
|
|
|
return False
|
|
|
|
async def check_comfyui_status() -> dict:
|
|
"""
|
|
Check the status of ComfyUI and the workflow template.
|
|
|
|
Returns:
|
|
dict: Status information
|
|
"""
|
|
try:
|
|
import aiohttp
|
|
|
|
# Check if ComfyUI workflow template exists
|
|
workflow_exists = os.path.exists("Miku_BasicWorkflow.json")
|
|
|
|
# Check if ComfyUI is running (try different Docker networking options)
|
|
comfyui_running = False
|
|
comfyui_url = "http://host.docker.internal:8188" # Default
|
|
|
|
comfyui_urls = [
|
|
"http://host.docker.internal:8188", # Docker Desktop
|
|
"http://172.17.0.1:8188", # Default Docker bridge gateway
|
|
"http://localhost:8188" # Fallback (if network_mode: host)
|
|
]
|
|
|
|
for url in comfyui_urls:
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
timeout = aiohttp.ClientTimeout(total=3)
|
|
async with session.get(f"{url}/system_stats", timeout=timeout) as response:
|
|
if response.status == 200:
|
|
comfyui_running = True
|
|
comfyui_url = url
|
|
break
|
|
except:
|
|
continue
|
|
|
|
return {
|
|
"workflow_template_exists": workflow_exists,
|
|
"comfyui_running": comfyui_running,
|
|
"comfyui_url": comfyui_url,
|
|
"ready": workflow_exists and comfyui_running
|
|
}
|
|
|
|
except Exception as e:
|
|
return {
|
|
"workflow_template_exists": False,
|
|
"comfyui_running": False,
|
|
"comfyui_url": "http://localhost:8188",
|
|
"ready": False,
|
|
"error": str(e)
|
|
}
|