Auto Image

"""
title: Auto Image Generator
author: J3tze
version: 5.7.0
description: Auto-detect image requests and generate via RunPod ComfyUI. SDXL-compatible prompts, async HTTP, leak scrubber. Enhanced spontaneous detection.
type: filter
"""

import aiohttp
import json
import base64
import random
import re
import asyncio
import os
import time

import uuid

import logging
from typing import Callable, Awaitable, Any, Optional
from pydantic import BaseModel, Field

from open_webui.models.users import Users
from open_webui.utils.chat import generate_chat_completion

# Custom logger with unique prefix for easy filtering
class PrefixedLogger:
    def __init__(self, name):
        self._logger = logging.getLogger(name)
    def info(self, msg): self._logger.info(f"[AUTO_IMG] {msg}")
    def warning(self, msg): self._logger.warning(f"[AUTO_IMG] {msg}")
    def error(self, msg, **kwargs): self._logger.error(f"[AUTO_IMG] {msg}", **kwargs)

logger = PrefixedLogger(__name__)


def _parse_bool(content: str) -> bool:
    if not content:
        return False
    first = content.strip().split()[0].strip().lower()
    return first in {"true", "yes", "1", "sure", "ok"}


class Filter:
    class Valves(BaseModel):
        # --- IDENTITY & STYLE ---
        physical_appearance: str = Field(
            default="",
            description="Character's immutable physical traits (hair, eyes, body, tattoos, bionic arm). Always included in positive prompt.",
        )
        default_outfit: str = Field(
            default="",
            description="Character's default clothing. Gets added to NEGATIVE prompt when user requests different outfit.",
        )
        positive_style: str = Field(
            default="masterpiece, best quality, 8k, detailed, dramatic lighting, cinematic",
            description="Global positive style tags.",
        )
        negative_style: str = Field(
            default="text, watermark, low quality, bad anatomy, blur, ugly, error, jpeg artifacts",
            description="Negative prompt.",
        )

        # --- BRAIN ---
        nanogpt_api_key: str = Field(default="", description="NanoGPT API Key.")
        nanogpt_model: str = Field(
            default="Qwen/Qwen2.5-72B-Instruct",
            description="NanoGPT model for prompt generation.",
        )
        nanogpt_system_prompt: str = Field(
            default=(
                "You are an expert Stable Diffusion prompter.\\n"
                "Merge the User Request with the Base Character. Return a JSON object.\\n"
                "FORMAT: { \"positive\": \"tags for image\", \"negative_additions\": \"tags to AVOID\" }\\n"
                "RULES:\\n"
                "1. TAG COUNT: Generate at least 50 comma-separated tags in 'positive'. Be detailed and descriptive.\\n"
                "2. OUTFIT CHANGES: If user wants different clothing:\\n"
                "   - Put NEW outfit tags in 'positive'\\n"
                "   - Put ALL of Base Character's default clothing in 'negative_additions'\\n"
                "3. LOCATION: If user specifies location, put base environment in 'negative_additions'\\n"
                "4. IDENTITY: Always keep physical traits (hair, eyes, body, tattoos, bionic arm)\\n"
                "5. POV: Use 'POV, first person' for first-person perspective shots.\\n"
                "7. OUTPUT: Return ONLY the JSON object."
            ),
            description="System instruction for the Prompt Enhancer LLM.",
        )
        context_depth: int = Field(default=5, description="Context depth.")

        # --- AUTO-DETECTION ---
        enable_auto_detection: bool = Field(default=True, description="Enable auto image generation.")
        logic_model: str = Field(default="glm-4-air", description="Model for detecting image requests.")
        force_triggers: str = Field(
            default="send me a picture,send me a photo,send me an image,send a selfie,generate an image,draw me,take a photo,show me a picture,show me an image,what you look like,see you,see what you,another image,another picture,another photo,one more image",
            description="Comma-separated trigger phrases.",
        )
        
        # --- SPONTANEOUS IMAGES ---
        enable_spontaneous: bool = Field(
            default=False,
            description="Enable spontaneous selfies when contextually appropriate.",
        )
        spontaneous_chance: int = Field(
            default=15,
            description="Max probability (0-100%) for spontaneous image. Actual chance depends on context.",
        )
        spontaneous_cooldown: int = Field(
            default=300,
            description="Minimum seconds between spontaneous images (default: 5 minutes).",
        )
        spontaneous_boost_keywords: str = Field(
            default="wish you could see,look at me,what are you wearing,dressed,cute,beautiful,show me,my outfit,looking good,how do i look",
            description="Comma-separated keywords that boost spontaneous trigger chance.",
        )

        # --- INFRASTRUCTURE ---
        runpod_api_key: str = Field(default="", description="RunPod API Key.")
        runpod_full_url: str = Field(
            default="https://api.runpod.ai/v2/YOUR_ENDPOINT_ID/runsync",
            description="RunPod URL.",
        )
        runpod_timeout: int = Field(default=300, description="Timeout in seconds.")
        cold_boot_timeout: int = Field(
            default=120,
            description="Extra timeout (seconds) for RunPod cold boots.",
        )

        # --- COMFYUI ---
        workflow_json: str = Field(default="", description="ComfyUI workflow JSON.")
        auto_detect_nodes: bool = Field(default=True, description="Auto-detect nodes.")
        positive_node_id: str = Field(default="", description="Positive Prompt Node ID.")
        negative_node_id: str = Field(default="", description="Negative Prompt Node ID.")
        seed_node_ids: str = Field(default="", description="Comma-separated Seed Node IDs.")

        # --- IMAGE SERVING ---
        image_base_url: str = Field(
            default="https://your-domain.com",
            description="Base URL for serving generated images (e.g. https://example.com). Images will be at {base_url}/images/{filename}.",
        )

        # --- DEBUG ---
        debug_mode: bool = Field(default=True, description="Enable debug logging.")

    def __init__(self):
        self.valves = self.Valves()
        self.last_spontaneous = 0  # Timestamp for cooldown tracking

    async def inlet(
        self, body: dict, __event_emitter__: Callable[[Any], Awaitable[None]], **kwargs
    ) -> dict:
        if not body or "messages" not in body:
            return body
        messages = body["messages"]
        if len(messages) < 1:
            return body

        if messages[-1]["role"] == "user":
            messages[-1]["content"] += (
                " [SYSTEM: You have a built-in image generator. "
                "If the user asks for an image/photo/selfie, DO NOT use web search or any other tool. "
                "Just reply naturally with a description or confirmation - the image will be generated automatically. "
                "IMPORTANT: Never output visible tags like [IMG_GENERATION], [IMAGE_GENERATED], or any bracketed system tags in your response. "
                "These are internal triggers and must never appear in your visible output.]"
            )
        return body

    async def outlet(
        self,
        body: dict,
        __event_emitter__: Callable[[Any], Awaitable[None]],
        __request__: Any,
        __user__: Optional[dict] = None,
        **kwargs,
    ) -> dict:
        logger.info("Outlet called")  # Entry point log
        
        if not body or "messages" not in body or not __user__:
            logger.info("Early return: missing body/messages/user")
            return body
        messages = body["messages"]
        if len(messages) < 2:
            logger.info(f"Early return: only {len(messages)} messages")
            return body

        # Clean leaks
        # Capture raw content for detection
        raw_assistant_msg = messages[-1]["content"]
        assistant_msg_lower = raw_assistant_msg.lower()
        
        # Clean leaks (ALWAYS clean the output for the user)
        # We clean AFTER capturing raw, so we can detect invisible triggers
        cleaned = self._clean_leaks(raw_assistant_msg)
        if cleaned != raw_assistant_msg:
            body["messages"][-1]["content"] = cleaned
            messages[-1]["content"] = cleaned

        if not self.valves.enable_auto_detection:
            logger.info("Auto-detection disabled")
            return body

        user_msg = messages[-2]["content"].lower()
        user_obj = Users.get_user_by_id(__user__["id"])

        # Skip if this is a dashboard response
        if "credits" in assistant_msg_lower and ("nanogpt" in assistant_msg_lower or "runpod" in assistant_msg_lower):
            logger.info("Skipping: Dashboard response detected")
            return body
        if "usage monitor" in assistant_msg_lower or "end transmission" in assistant_msg_lower:
            logger.info("Skipping: Dashboard response detected")
            return body

        triggers = [t.strip().lower() for t in self.valves.force_triggers.split(",") if t.strip()]
        should_gen = any(t in user_msg for t in triggers)
        is_selfie = any(t in user_msg for t in ["me", "selfie", "you", "yourself"])
        
        logger.info(f"User msg: {user_msg[:100]}... | Force trigger match: {should_gen}")

        # Check for explicit image request via LLM
        if not should_gen:
            should_gen = await self._llm_check(
                user_msg, "Return 'true' only if user wants a photo/image. Otherwise 'false'.",
                __request__, user_obj
            )
            logger.info(f"LLM check result: {should_gen}")
        
        # Check if assistant's response indicates it wants to send an image (using RAW content)
        if not should_gen:
            response_triggers = [
                '"generate_image"', '"action":', "strike a pose", "strikes a pose",
                "snap a pic", "take a picture", "let me show you", "here's a visual",
                "you want a visual", "pushing the resolution", "check this out"
            ]
            if any(t in assistant_msg_lower for t in response_triggers):
                should_gen = True
                is_selfie = True
                logger.info("Response-based trigger: assistant wants to send image")
        
        # Check for spontaneous image opportunity
        spontaneous_triggered = False
        if not should_gen and self.valves.enable_spontaneous and self.valves.nanogpt_api_key:
            spontaneous_triggered = await self._check_spontaneous(user_msg, assistant_msg)
            if spontaneous_triggered:
                should_gen = True
                is_selfie = True  # Spontaneous images are selfies
                logger.info("Spontaneous selfie triggered!")

        if should_gen:
            async def emit_status(desc: str, done: bool = False):
                if __event_emitter__:
                    await __event_emitter__({"type": "status", "data": {"description": desc, "done": done}})
                    if not done:
                        await asyncio.sleep(0.5)

            async def emit_message(content: str):
                """Emit message to append content to chat - same as action button uses."""
                if __event_emitter__:
                    await __event_emitter__({"type": "message", "data": {"content": content}})

            logger.info("Triggered - starting image generation")
            await emit_status("Analyzing scene context...")

            history = messages[-self.valves.context_depth:]
            chat_log = "\n".join([f"{m['role']}: {m['content']}" for m in history])
            pos, neg = await self._get_smart_prompt(chat_log, is_selfie)
            
            logger.info(f"Prompt: {pos[:150]}")

            await emit_status("Connecting to RunPod...")
            
            result = await self._generate_image(pos, neg, emit_status)
            
            logger.info(f"Gen result: {result}")

            if result and isinstance(result, dict) and "url" in result:
                url = result["url"]
                
                # Add cache-busting timestamp
                generation_timestamp = int(time.time())
                
                # Clean up prompts for display (replace all newlines/carriage returns with spaces)
                clean_positive = re.sub(r'[\r\n]+', ' ', result.get("positive", ""))
                clean_positive = re.sub(r'\s{2,}', ' ', clean_positive).strip()
                clean_negative = re.sub(r'[\r\n]+', ' ', result.get("negative", ""))
                clean_negative = re.sub(r'\s{2,}', ' ', clean_negative).strip()
                
                # HTML Dropdown for details with JSON block
                details_json = {
                    "prompt": clean_positive,
                    "negative": clean_negative,
                    "seed": result.get("seed", "N/A"),
                    "file_size": result.get("file_size", "N/A"),
                    "generated_at": generation_timestamp,  # Cache buster
                }
                
                details_html = (
                    f'<details>\n'
                    f'<summary>Generation Details</summary>\n\n'
                    f'```json\n{json.dumps(details_json, indent=4)}\n```\n\n'
                    f'</details>'
                )

                # Append image and details to the assistant's response content
                # Add HTML comment with timestamp to force unique content (cache buster)
                image_markdown = f"\n\n![Generated Image]({url})\n\n{details_html}\n\n<!-- gen_{generation_timestamp} -->"
                
                # messages is a reference to body['messages'], so we only need to update one
                messages[-1]["content"] += image_markdown
                logger.info("Appended image to response")
                await emit_status("Image attached!", done=True)
            else:
                logger.warning(f"Invalid result: {result}")
                await emit_status("Image generation failed", done=True)
        else:
            logger.info("No trigger detected, skipping image gen")

        # Final cleanup pass - catch any remaining leaks before returning
        final_content = messages[-1]["content"]
        final_cleaned = self._clean_leaks(final_content)
        if final_cleaned != final_content:
            messages[-1]["content"] = final_cleaned
            logger.info("Final cleanup removed additional leaks")

        return body

    def _clean_leaks(self, text: str) -> str:
        """Remove leaked system prompts, tool indicators, and JSON triggers from AI responses."""
        # Preserve dashboard output
        if "API Usage Dashboard" in text or "💰" in text:
            return text
        
        # Extract and preserve our Generation Details block if present
        details_block = ""
        if "<details>" in text and "<summary>Generation Details</summary>" in text:
            start = text.find("<details>")
            end = text.find("</details>") + len("</details>")
            if start != -1 and end > start:
                details_block = text[start:end]
                text = text[:start] + text[end:]
        
        cleaned = text
        original_len = len(text)
        
        # 1. Remove HTML comment wrappers (safety net for any leaked triggers)
        while "<!--" in cleaned:
            start = cleaned.find("<!--")
            end = cleaned.find("-->", start)
            if start != -1 and end != -1:
                cleaned = cleaned[:start] + cleaned[end + 3:]
            elif start != -1:
                # Truncated comment - remove to end
                cleaned = cleaned[:start]
                break
            else:
                break
        
        # 2. Remove escaped HTML comments
        cleaned = re.sub(r"&lt;!--[\s\S]*?--&gt;", "", cleaned)
        
        # 3. Remove bracketed system indicators (but NOT markdown image alt text like ![Generated Image])
        cleaned = re.sub(r"\[SYSTEM:.*?\]", "", cleaned, flags=re.IGNORECASE | re.DOTALL)
        cleaned = re.sub(r"\[Tool.*?\]", "", cleaned, flags=re.IGNORECASE | re.DOTALL)
        # Explicit patterns for common leaks
        cleaned = re.sub(r"\[IMG\]", "", cleaned, flags=re.IGNORECASE)  # Just [IMG]
        cleaned = re.sub(r"\[IMG_?GENERATION\]", "", cleaned, flags=re.IGNORECASE)
        cleaned = re.sub(r"\[IMAGE_?GENERAT\w*\]", "", cleaned, flags=re.IGNORECASE)
        cleaned = re.sub(r"\[Image[^\]]*\]", "", cleaned, flags=re.IGNORECASE)  # [Image anything]
        cleaned = re.sub(r"\[.*?trigger.*?\]", "", cleaned, flags=re.IGNORECASE)
        cleaned = re.sub(r"\[.*?automatic.*?\]", "", cleaned, flags=re.IGNORECASE)
        cleaned = re.sub(r"\[.*?silent.*?\]", "", cleaned, flags=re.IGNORECASE)  # [silently], [triggers silently]
        # Remove photo/image action tags
        cleaned = re.sub(r"\[.*?(?:takes?|snaps?|sends?).*?(?:photo|pic|selfie).*?\]", "", cleaned, flags=re.IGNORECASE)
        # Generic uppercase-only tags (not markdown images)
        cleaned = re.sub(r"(?<!\!)\[[A-Z][A-Z\s_]+\]", "", cleaned)
        
        # 4. Remove JSON trigger blocks (but NOT the details dropdown JSON)
        cleaned = re.sub(r'\{\s*"action"\s*:\s*"[^"]*"\s*\}', "", cleaned, flags=re.DOTALL)
        
        # 5. Remove explanatory notes about image generation
        cleaned = re.sub(r"\(Note:.*?image.*?\)", "", cleaned, flags=re.IGNORECASE | re.DOTALL)
        
        # 6. Clean up empty code blocks and excessive newlines
        cleaned = re.sub(r"```\w*\s*```", "", cleaned)
        cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
        
        # Log if we cleaned anything
        if len(cleaned) != original_len:
            logger.info(f"Cleaned {original_len - len(cleaned)} chars from response")
        
        # Re-append preserved details block
        result = cleaned.strip()
        if details_block:
            result = result + "\n\n" + details_block
        
        return result

    async def _get_smart_prompt(self, chat_log: str, is_selfie: bool = False) -> tuple:
        try:
            selfie_mod = ", selfie, looking at camera, close up" if is_selfie else ""
            if not self.valves.nanogpt_api_key:
                return (
                    f"{self.valves.positive_style}, {self.valves.physical_appearance}, {self.valves.default_outfit}{selfie_mod}",
                    self.valves.negative_style,
                )

            system_instruction = self.valves.nanogpt_system_prompt

            payload = {
                "model": self.valves.nanogpt_model,
                "messages": [
                    {"role": "system", "content": system_instruction},
                    {"role": "user", "content": f"PHYSICAL APPEARANCE (always keep):\n{self.valves.physical_appearance}\n\nDEFAULT OUTFIT (put in negative if changed):\n{self.valves.default_outfit}\n\nHISTORY:\n{chat_log}"},
                ],
                "temperature": 0.3,
            }

            # Retry logic for NanoGPT (handles occasional timeouts)
            raw_content = None
            max_retries = 2
            for attempt in range(max_retries):
                try:
                    async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=45)) as session:
                        async with session.post(
                            "https://nano-gpt.com/api/v1/chat/completions",
                            json=payload,
                            headers={"Authorization": f"Bearer {self.valves.nanogpt_api_key}", "Content-Type": "application/json"},
                        ) as resp:
                            data = await resp.json()
                            logger.info(f"NanoGPT response status: {resp.status}")
                            
                            if "error" in data:
                                logger.warning(f"NanoGPT API error: {data['error']}")
                                raise Exception(f"API error: {data['error']}")
                            
                            if "choices" not in data:
                                logger.warning(f"NanoGPT unexpected response: {json.dumps(data)[:500]}")
                                raise Exception(f"No 'choices' in response")
                            
                            raw_content = data["choices"][0]["message"]["content"].strip()
                            logger.info(f"NanoGPT raw response: {raw_content[:200]}")
                            break  # Success, exit retry loop
                except asyncio.TimeoutError:
                    logger.warning(f"NanoGPT timeout (attempt {attempt + 1}/{max_retries})")
                    if attempt == max_retries - 1:
                        raise  # Re-raise on final attempt
                    await asyncio.sleep(1)
            
            if not raw_content:
                raise Exception("NanoGPT returned empty content")

            # Parse JSON or fallback to string
            ai_tags = raw_content
            neg_additions = ""

            try:
                # Remove Markdown code blocks if present
                clean_json = re.sub(r"```json\s*|\s*```", "", raw_content, flags=re.IGNORECASE).strip()
                parsed = json.loads(clean_json)
                if isinstance(parsed, dict):
                    ai_tags = parsed.get("positive", "")
                    neg_additions = parsed.get("negative_additions", "")
                    logger.info(f"Parsed: positive={ai_tags[:100]}..., neg_additions={neg_additions[:100] if neg_additions else 'NONE'}")
            except json.JSONDecodeError:
                logger.warning("NanoGPT returned non-JSON text. Using raw output as positive tags.")
                # Strip fences to prevent nested-markdown issues in the dropdown
                ai_tags = re.sub(r"```json\s*|\s*```", "", raw_content, flags=re.IGNORECASE).strip()

            final_negative = self.valves.negative_style
            if neg_additions:
                final_negative = f"{neg_additions}, {final_negative}"
                logger.info(f"NanoGPT provided neg_additions: {neg_additions[:100]}")
            else:
                # Fallback: detect if user asked for outfit change but NanoGPT didn't apply it
                outfit_keywords = ["dress", "bikini", "swimsuit", "lingerie", "underwear", "pajamas", "nightgown", "skirt", "gown", "uniform"]
                chat_lower = chat_log.lower()
                detected_keywords = [kw for kw in outfit_keywords if kw in chat_lower]
                logger.info(f"Outfit fallback check: keywords={detected_keywords}, default_outfit={'SET' if self.valves.default_outfit else 'EMPTY'}")
                if detected_keywords and self.valves.default_outfit:
                    final_negative = f"{self.valves.default_outfit}, {final_negative}"
                    logger.info(f"Fallback ACTIVATED: added default_outfit to negative")

            logger.info(f"Final positive preview: {ai_tags[:150]}...")
            logger.info(f"Final negative preview: {final_negative[:150]}...")
            
            return (
                f"{ai_tags}, {self.valves.physical_appearance}, {self.valves.positive_style}{selfie_mod}",
                final_negative,
            )
        except Exception as e:
            error_msg = str(e) or type(e).__name__
            logger.warning(f"Smart prompt failed: {error_msg}")
            selfie_mod = ", selfie, looking at camera, close up" if is_selfie else ""
            
            # Even in fallback, check for outfit keywords
            final_neg = self.valves.negative_style
            outfit_keywords = ["dress", "bikini", "swimsuit", "lingerie", "underwear", "skirt", "gown"]
            chat_lower = chat_log.lower()
            if any(kw in chat_lower for kw in outfit_keywords) and self.valves.default_outfit:
                final_neg = f"{self.valves.default_outfit}, {final_neg}"
                logger.info(f"Fallback: detected outfit keyword, added default_outfit to negative")
            
            return (
                f"{self.valves.positive_style}, {self.valves.physical_appearance}{selfie_mod}",
                final_neg,
            )

    async def _generate_image(self, pos: str, neg: str, emit_status) -> Optional[dict]:
        try:
            if not self.valves.workflow_json or not self.valves.runpod_api_key:
                await emit_status("Error: Missing config")
                return None

            workflow = json.loads(self.valves.workflow_json)
            
            # Generate truly random seed using system entropy
            random.seed(int.from_bytes(os.urandom(8), 'big') ^ int(time.time() * 1000000))
            seed_val = random.randint(1, 2**32 - 1)
            logger.info(f"Generated seed: {seed_val}")
            
            resolution = self._get_resolution(workflow)

            if self.valves.auto_detect_nodes:
                detected = self._auto_detect_nodes(workflow)
                p_id, n_id = detected.get("positive", ""), detected.get("negative", "")
                seed_ids = detected.get("seeds", [])
                logger.info(f"Auto-detected: pos={p_id}, neg={n_id}, seeds={seed_ids}")
            else:
                p_id = self.valves.positive_node_id.strip()
                n_id = self.valves.negative_node_id.strip()
                seed_ids = [s.strip() for s in self.valves.seed_node_ids.split(",") if s.strip()]

            if p_id and p_id in workflow:
                workflow[p_id]["inputs"]["text"] = pos
                logger.info(f"Injected positive prompt into node {p_id}")
            else:
                logger.warning(f"Positive node '{p_id}' not found!")
            
            if n_id and n_id in workflow:
                workflow[n_id]["inputs"]["text"] = neg
            else:
                logger.warning(f"Negative node '{n_id}' not found!")
            
            # Inject seeds
            for s_id in seed_ids:
                if s_id in workflow and "inputs" in workflow[s_id]:
                    inputs = workflow[s_id]["inputs"]
                    if "seed" in inputs:
                        inputs["seed"] = seed_val
                    if "noise_seed" in inputs:
                        inputs["noise_seed"] = seed_val

            # Replace ALL seeds in workflow
            self._replace_all_seeds(workflow, seed_val)
            self._recursive_fix(workflow, pos, neg)

            # --- ROBUST RUNPOD HANDLER ---
            headers = {
                "Authorization": f"Bearer {self.valves.runpod_api_key}",
                "Content-Type": "application/json",
            }
            
            start_time = time.time()
            total_timeout = self.valves.runpod_timeout + self.valves.cold_boot_timeout
            
            # Initial submit
            await emit_status("Sending to RunPod...")
            job_id = None
            
            async with aiohttp.ClientSession() as session:
                try:
                    async with session.post(
                        self.valves.runpod_full_url,
                        json={"input": {"workflow": workflow}},
                        headers=headers,
                        timeout=aiohttp.ClientTimeout(total=30) # Short timeout for initial ack
                    ) as resp:
                        data = await resp.json()
                        if self.valves.debug_mode:
                            logger.info(f"RunPod submit response: {json.dumps(data)}")
                            
                        job_id = data.get("id")
                        job_status = data.get("status", "").upper()
                except Exception as e:
                    logger.error(f"Failed to submit job: {e}")
                    await emit_status("Failed to connect to RunPod")
                    return None

            if not job_id:
                if "error" in data:
                    await emit_status(f"RunPod Error: {data.get('error')}")
                else:
                    await emit_status("RunPod returned no Job ID")
                return None

            # Polling Loop
            poll_url = self.valves.runpod_full_url.replace("/runsync", f"/status/{job_id}")
            logger.info(f"Job {job_id} submitted. Status: {job_status}. Starting poll loop.")
            
            consecutive_network_errors = 0
            max_network_errors = 10  # Prevent infinite loop on persistent network issues
            
            while (time.time() - start_time) < total_timeout:
                # 1. Check Status
                if job_status == "COMPLETED":
                    await emit_status("Generation complete! Downloading...")
                    break
                elif job_status == "FAILED":
                    await emit_status("RunPod job failed.")
                    return None
                
                # 2. Smart Polling Delay
                elapsed = time.time() - start_time
                if job_status == "IN_QUEUE":
                    if elapsed > 15:
                        await emit_status(f"GPU Warming up... ({int(elapsed)}s)")
                        sleep_time = 5  # Slow poll for queue
                    else:
                        await emit_status("Queued...")
                        sleep_time = 2
                elif job_status == "IN_PROGRESS":
                    await emit_status(f"Rendering... ({int(elapsed)}s)")
                    sleep_time = 2  # Fast poll for progress
                else:
                    await emit_status(f"Status: {job_status}...")
                    sleep_time = 3
                
                await asyncio.sleep(sleep_time)
                
                # 3. Fetch Update (Short-lived session)
                try:
                    async with aiohttp.ClientSession() as session:
                        async with session.get(
                            poll_url, 
                            headers=headers,
                            timeout=aiohttp.ClientTimeout(total=15)  # Increased from 10s to 15s
                        ) as resp:
                            data = await resp.json()
                            new_status = data.get("status", job_status).upper()
                            
                            # Reset error counter on successful poll
                            if new_status != job_status:
                                consecutive_network_errors = 0
                            
                            job_status = new_status
                            if self.valves.debug_mode and job_status == "COMPLETED":
                                logger.info(f"Job completed data: {str(data)[:200]}...")
                except Exception as e:
                    consecutive_network_errors += 1
                    logger.warning(f"Poll network error #{consecutive_network_errors} (retrying): {e}")
                    
                    # If we hit too many consecutive errors, give up
                    if consecutive_network_errors >= max_network_errors:
                        logger.error(f"Too many consecutive network errors ({consecutive_network_errors}). Aborting.")
                        await emit_status("Network errors - unable to check status")
                        return None
                    
                    # Don't count network errors against total timeout - just retry
                    await asyncio.sleep(3)  # Longer backoff on network error
                    continue
            else:
                # Final status check before giving up - RunPod might have completed while we were timing out
                logger.info("Timeout reached - attempting final status check...")
                try:
                    async with aiohttp.ClientSession() as session:
                        async with session.get(
                            poll_url, 
                            headers=headers,
                            timeout=aiohttp.ClientTimeout(total=20)
                        ) as resp:
                            data = await resp.json()
                            job_status = data.get("status", job_status).upper()
                            logger.info(f"Final check: status is {job_status}")
                            
                            # If it completed, continue to extraction
                            if job_status == "COMPLETED":
                                logger.info("Job completed during final check!")
                                await emit_status("Generation complete! Downloading...")
                            else:
                                await emit_status(f"Timeout waiting for RunPod (last status: {job_status})")
                                return None
                except Exception as e:
                    logger.error(f"Final status check failed: {e}")
                    await emit_status("Timeout waiting for RunPod")
                    return None

            # --- IMAGE EXTRACTION ---
            if "error" in data:
                await emit_status(f"RunPod error")
                return None

            img_result = self._deep_find_img(data)
            if not img_result:
                if isinstance(data, dict):
                    logger.warning(f"No image found. keys: {list(data.keys())}")
                    if "status" in data:
                        logger.warning(f"Status: {data['status']}")
                await emit_status("No image in response")
                return None

            if img_result["type"] == "url":
                return {"url": img_result["data"], "seed": seed_val, "positive": pos, "negative": neg, "file_size": "External URL"}

            await emit_status("Saving locally...")
            
            img_data = img_result["data"]
            if "base64," in img_data:
                img_data = img_data.split("base64,")[1]

            decoded = base64.b64decode(img_data)
            
            # Generate unique filename
            filename = f"img_{uuid.uuid4().hex[:12]}.png"
            
            # Save to shared volume path
            # We use /app/backend/data/images which is persistent and shared
            output_dir = "/app/backend/data/images"
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
            
            output_path = os.path.join(output_dir, filename)
            
            with open(output_path, "wb") as f:
                f.write(decoded)
            
            # Calculate file size in human-readable format
            file_size_bytes = len(decoded)
            if file_size_bytes >= 1024 * 1024:
                file_size_str = f"{file_size_bytes / (1024 * 1024):.2f} MB"
            elif file_size_bytes >= 1024:
                file_size_str = f"{file_size_bytes / 1024:.1f} KB"
            else:
                file_size_str = f"{file_size_bytes} bytes"

            if self.valves.debug_mode:
                logger.info(f"Saved image to {output_path} ({file_size_str})")

            # Return public URL matching Nginx config and metadata
            return {
                "url": f"{self.valves.image_base_url}/images/{filename}",
                "seed": seed_val,
                "positive": pos,
                "negative": neg,
                "file_size": file_size_str,
            }

        except Exception as e:
            logger.error(f"Image gen failed: {e}", exc_info=True)
            await emit_status(f"Error: {str(e)[:50]}")
            return None

    def _replace_all_seeds(self, workflow: dict, seed_val: int):
        """Replace ALL seed values in the entire workflow."""
        for node_id, node in workflow.items():
            if not isinstance(node, dict) or "inputs" not in node:
                continue
            inputs = node["inputs"]
            if "seed" in inputs and isinstance(inputs["seed"], (int, float, str)):
                inputs["seed"] = seed_val
            if "noise_seed" in inputs and isinstance(inputs["noise_seed"], (int, float, str)):
                inputs["noise_seed"] = seed_val

    def _auto_detect_nodes(self, workflow: dict) -> dict:
        result = {"positive": "", "negative": "", "seeds": []}
        text_types = {"CLIPTextEncode", "CLIPTextEncodeSDXL"}
        sampler_types = {"KSampler", "KSamplerAdvanced", "SamplerCustom", "FaceDetailer", "DetailerForEach"}
        
        sampler_conns = {}
        for nid, node in workflow.items():
            if not isinstance(node, dict):
                continue
            ctype = node.get("class_type", "")
            inputs = node.get("inputs", {})
            if ctype in sampler_types or "sampler" in ctype.lower() or "Sampler" in ctype:
                sampler_conns[nid] = {"positive": [], "negative": []}
                for key in ["positive", "negative"]:
                    if key in inputs and isinstance(inputs[key], list):
                        sampler_conns[nid][key].append(str(inputs[key][0]))
                if "seed" in inputs or "noise_seed" in inputs:
                    result["seeds"].append(nid)
        
        # RandomNoise nodes
        for nid, node in workflow.items():
            if not isinstance(node, dict):
                continue
            ctype = node.get("class_type", "")
            inputs = node.get("inputs", {})
            if "Noise" in ctype and "noise_seed" in inputs:
                if nid not in result["seeds"]:
                    result["seeds"].append(nid)

        pos_cand, neg_cand = [], []
        for nid, node in workflow.items():
            if not isinstance(node, dict):
                continue
            ctype = node.get("class_type", "")
            inputs = node.get("inputs", {})
            if (ctype in text_types or "TextEncode" in ctype) and "text" in inputs:
                for sid, conns in sampler_conns.items():
                    if nid in conns["positive"]:
                        pos_cand.append(nid)
                        break
                    elif nid in conns["negative"]:
                        neg_cand.append(nid)
                        break
                else:
                    title = node.get("_meta", {}).get("title", "").lower()
                    if "negative" in title:
                        neg_cand.append(nid)
                    elif not pos_cand:
                        pos_cand.append(nid)
                    elif not neg_cand:
                        neg_cand.append(nid)

        if pos_cand:
            result["positive"] = pos_cand[0]
        if neg_cand:
            result["negative"] = neg_cand[0]
        return result

    def _recursive_fix(self, obj, pos, neg):
        """Replace text placeholders."""
        if isinstance(obj, dict):
            for k, v in obj.items():
                if isinstance(v, (dict, list)):
                    self._recursive_fix(v, pos, neg)
                elif k == "text" and isinstance(v, str):
                    if "%prompt%" in v:
                        obj[k] = v.replace("%prompt%", pos)
                    if "%negative_prompt%" in v:
                        obj[k] = v.replace("%negative_prompt%", neg)
        elif isinstance(obj, list):
            for item in obj:
                self._recursive_fix(item, pos, neg)

    def _deep_find_img(self, obj) -> Optional[dict]:
        if isinstance(obj, str):
            if len(obj) > 1000 and ("iVBOR" in obj or "/9j/" in obj):
                return {"type": "base64", "data": obj}
            if obj.startswith("data:image/"):
                return {"type": "base64", "data": obj}
            if obj.startswith(("http://", "https://")) and any(e in obj.lower() for e in [".png", ".jpg", ".jpeg", ".webp"]):
                return {"type": "url", "data": obj}
        if isinstance(obj, list):
            for item in obj:
                r = self._deep_find_img(item)
                if r:
                    return r
        if isinstance(obj, dict):
            for key in ["output", "images", "image", "data", "message", "result"]:
                if key in obj:
                    r = self._deep_find_img(obj[key])
                    if r:
                        return r
            for v in obj.values():
                r = self._deep_find_img(v)
                if r:
                    return r
        return None

    async def _llm_check(self, prompt: str, system: str, req, user) -> bool:
        try:
            r = await generate_chat_completion(
                request=req,
                form_data={"model": self.valves.logic_model, "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": prompt},
                ], "stream": False},
                user=user,
            )
            content = r["choices"][0]["message"]["content"]
            result = _parse_bool(content)
            logger.info(f"LLM check response: '{content[:50]}' -> {result}")
            return result
        except Exception as e:
            logger.warning(f"LLM check failed: {e}")
            return False

    def _get_resolution(self, workflow):
        try:
            for node in workflow.values():
                if not isinstance(node, dict): continue
                if node.get("class_type") == "EmptyLatentImage":
                     inputs = node.get("inputs", {})
                     w = inputs.get("width")
                     h = inputs.get("height")
                     return f"{w}x{h}"
        except:
            pass
        return "Unknown"

    async def _check_spontaneous(self, user_msg: str, assistant_msg: str) -> bool:
        """Use NanoGPT to decide if a spontaneous selfie would be appropriate."""
        try:
            # 1. Cooldown check - avoid spamming spontaneous images
            if time.time() - self.last_spontaneous < self.valves.spontaneous_cooldown:
                logger.info(f"Spontaneous cooldown active ({int(self.valves.spontaneous_cooldown - (time.time() - self.last_spontaneous))}s remaining)")
                return False
            
            combined_text = (user_msg + assistant_msg).lower()
            
            # 2. Keyword boost pre-filter - saves API calls if no relevant keywords
            boost_keywords = [kw.strip().lower() for kw in self.valves.spontaneous_boost_keywords.split(",") if kw.strip()]
            keyword_matches = [kw for kw in boost_keywords if kw in combined_text]
            keyword_boost = len(keyword_matches)
            
            # If no keywords match and random check fails, skip the API call entirely
            if keyword_boost == 0 and random.randint(1, 100) > 30:  # 70% chance to skip if no keywords
                logger.info("Spontaneous: No keyword matches, skipping API check")
                return False
            
            logger.info(f"Spontaneous keyword matches: {keyword_matches[:5]} (boost: +{keyword_boost})")
            
            # 3. Richer context window (500 chars each instead of 300)
            context = f"User: {user_msg[:500]}\nAssistant: {assistant_msg[:500]}"
            
            # 4. More specific prompt for roleplay scenarios
            system_prompt = (
                "You analyze roleplay conversations to decide if sending a spontaneous selfie feels natural. "
                "HIGH appropriateness (7-10): Showing off outfit changes, celebrating achievements, "
                "flirty/playful moments, 'wish you could see' statements, intimate scenarios, physical descriptions. "
                "MEDIUM (4-6): General positivity, casual chat, mild emotional connection. "
                "LOW (1-3): Technical discussions, questions needing answers, tense moments, serious topics. "
                "Reply with ONLY a single digit 1-10."
            )
            
            payload = {
                "model": self.valves.nanogpt_model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": context},
                ],
                "temperature": 0.2,
            }
            
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15)) as session:
                async with session.post(
                    "https://nano-gpt.com/api/v1/chat/completions",
                    json=payload,
                    headers={"Authorization": f"Bearer {self.valves.nanogpt_api_key}", "Content-Type": "application/json"},
                ) as resp:
                    data = await resp.json()
                    
                    if "error" in data or "choices" not in data:
                        logger.warning(f"Spontaneous check API error: {data.get('error', 'No choices')}")
                        return False
                    
                    score_text = data["choices"][0]["message"]["content"].strip()
                    # Extract first number from response
                    score = 5  # default
                    for char in score_text:
                        if char.isdigit():
                            score = int(char)
                            break
                    
                    # Apply keyword boost to score (max +2)
                    boosted_score = min(10, score + min(keyword_boost, 2))
                    logger.info(f"Spontaneous score: {score}/10 + {min(keyword_boost, 2)} boost = {boosted_score}/10 (response: {score_text[:30]})")
                    
                    # Combine AI score with random chance
                    if boosted_score < 5:
                        return False
                    
                    # Scale probability: score 5=low chance, score 10=max chance
                    scaled_chance = (boosted_score - 4) / 6 * self.valves.spontaneous_chance  # 0-max%
                    roll = random.randint(1, 100)
                    triggered = roll <= scaled_chance
                    
                    logger.info(f"Spontaneous roll: {roll} vs {scaled_chance:.1f}% -> {'TRIGGERED' if triggered else 'skip'}")
                    
                    # Update cooldown timestamp if triggered
                    if triggered:
                        self.last_spontaneous = time.time()
                    
                    return triggered
                    
        except Exception as e:
            logger.warning(f"Spontaneous check failed: {e}")
            return False