miku-discord/bot/utils/image_generation.py

"""
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
from utils.logger import get_logger

logger = get_logger('media')

# 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:
                    logger.debug(f"Found recent image: {file_path}")
                    return file_path
                    
        except Exception as e:
            logger.error(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):
            logger.error(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
                            logger.debug(f"ComfyUI found at: {url}")
                            break
            except:
                continue
        
        if not comfyui_url:
            logger.error(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:
                    logger.error(f"ComfyUI request failed: {response.status}")
                    return None
                
                result = await response.json()
                prompt_id = result.get("prompt_id")
                
                if not prompt_id:
                    logger.error("No prompt_id received from ComfyUI")
                    return None
                
                logger.info(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):
                                                logger.info(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):
                                                    logger.info(f"Image generated successfully: {alt_path}")
                                                    return alt_path
                                                else:
                                                    logger.warning(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
                                logger.debug("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)
                
                logger.error("ComfyUI generation timed out")
                
                # Final fallback: look for the most recent image
                logger.debug("Trying final fallback: most recent image...")
                fallback_image = find_latest_generated_image(prompt_id)
                if fallback_image:
                    logger.info(f"Found image via fallback method: {fallback_image}")
                    return fallback_image
                
                return None
                
    except Exception as e:
        logger.error(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
            logger.info(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)
                
                logger.info(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)
                
                logger.error(f"Image generation failed for prompt: {prompt}")
                return False
                
    except Exception as e:
        logger.error(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)
        }