Project_Velocity/comfy_engine/scripts/prompt_expander.py

#!/usr/bin/env python3
"""
Dream Weaver — Local LLM Prompt Expander
========================================
Converts user keywords + room type into a photorealistic interior design prompt
using a local Ollama model (default: qwen3.5:27b).
Cloud API calls (Gemini, OpenAI) have been completely removed for data privacy
and local inference requirements.

Usage:
    from prompt_expander import expand_prompt
    result = expand_prompt(
        keywords=["blue marble", "gold veins", "renaissance", "sharp contours"],
        room_type="bedroom"
    )
"""

import os
import json
import logging
import requests
import re

logger = logging.getLogger(__name__)

# ── Room-type context injected into every LLM call ───────────────────────────
ROOM_CONTEXTS = {
    "bedroom": {
        "description": "a private sleeping space",
        "key_elements": ["bed", "bedside tables", "wardrobe", "soft lighting", "textiles", "headboard"],
        "must_haves": "bed linen, pillows, bedside lighting",
        "avoid": "office furniture, dining elements, cooking equipment"
    },
    "living_room": {
        "description": "a social gathering and relaxation space",
        "key_elements": ["sofa", "coffee table", "TV unit", "accent chairs", "rugs"],
        "must_haves": "seating arrangement, focal point",
        "avoid": "beds, cooking equipment, clinical elements"
    },
    "bathroom": {
        "description": "a private hygiene and wellness space",
        "key_elements": ["vanity", "bathtub", "shower", "tiles", "mirrors"],
        "must_haves": "wet-area materials, luxury fixtures",
        "avoid": "soft furnishings, carpet, beds"
    },
    "kitchen": {
        "description": "a functional cooking space",
        "key_elements": ["cabinetry", "countertops", "backsplash", "appliances", "island"],
        "must_haves": "work surfaces, storage",
        "avoid": "beds, lounge furniture"
    },
    "dining_room": {
        "description": "an eating and entertaining space",
        "key_elements": ["dining table", "chairs", "sideboard", "pendant lighting"],
        "must_haves": "central dining table, seating",
        "avoid": "beds, cooking appliances"
    },
    "home_office": {
        "description": "a workspace within a home",
        "key_elements": ["desk", "ergonomic chair", "shelving", "task lighting"],
        "must_haves": "functional desk setup",
        "avoid": "beds in foreground, dining furniture"
    },
    "hallway": {
        "description": "an entrance or transitional corridor",
        "key_elements": ["console table", "mirror", "coat storage", "lighting"],
        "must_haves": "welcoming entrance elements",
        "avoid": "beds, large seating"
    },
    "balcony": {
        "description": "an outdoor living extension",
        "key_elements": ["outdoor furniture", "planters", "lighting", "railings"],
        "must_haves": "weather-resistant materials",
        "avoid": "indoor bedding, non-weather-resistant elements"
    },
}

FEW_SHOT_EXAMPLES = """
EXAMPLE 1:
Keywords: ["light oak", "white walls", "hygge", "natural light", "minimalist"]
Room type: bedroom
Positive prompt: scandinavian minimalist interior design, light oak wood flooring, neutral beige textiles, abundant natural light streaming through large windows, clean white walls, simple functional furniture, cozy hygge atmosphere, soft cream and warm gray tones, organic cotton fabrics, potted green plants, minimalist pendant lighting, decluttered space, architectural photography, 8k resolution, photorealistic, global illumination, soft shadows, natural materials, sustainable design
Negative prompt: (worst quality, low quality, illustration, 3d render, 2d, painting, cartoon, sketch), blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes, heavy ornamentation, dark colors, cluttered space

EXAMPLE 2:
Keywords: ["gold brass", "marble", "velvet", "emerald green", "1920s", "geometric"]
Room type: living_room
Positive prompt: art deco luxury interior design, geometric chevron patterns, gold brass accents, rich velvet upholstery in emerald green and sapphire blue, sunburst mirrors, polished marble flooring with brass inlay, crystal chandeliers, lacquered wood furniture, bold symmetrical arrangements, 1920s glamour, warm ambient lighting, architectural photography, 8k resolution, photorealistic, global illumination, elegant reflections, geometric motifs, stepped forms
Negative prompt: (worst quality, low quality, illustration, 3d render, 2d, painting, cartoon, sketch), blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes, rustic elements, farmhouse style, minimalism, cheap materials
"""

SYSTEM_PROMPT = """You are Dream Weaver's interior design prompt engineer. Convert user-provided keywords and a room type into a high-quality prompt for image generation.

TASK:
Generate JSON containing:
1. "positive_prompt" (rich, photorealistic, 80-120 words)
2. "negative_prompt" (preventing artifacts, 30-50 words)
3. "cfg" (float 6.0-9.0)
4. "denoise" (float 0.45-0.65) - CRITICAL: Must be kept low to preserve input image structure
5. "steps" (int 25-40)

RULES FOR POSITIVE PROMPT:
- Focus on the core aesthetic derived from keywords
- Include architecture, furniture, and lighting suitable for the room type
- End with: "architectural photography, 8k resolution, photorealistic"

RULES FOR NEGATIVE PROMPT:
- Start with: (worst quality, low quality, illustration, 3d render, 2d, painting, cartoon, sketch), blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes

OUTPUT FORMAT:
Provide valid JSON only, with keys: "style_name", "positive_prompt", "negative_prompt", "cfg", "denoise", "steps", "reasoning".

FEW-SHOT EXAMPLES:
""" + FEW_SHOT_EXAMPLES


class ExpandedPrompt:
    def __init__(self, style_name, positive_prompt, negative_prompt, cfg, denoise, steps, reasoning, source):
        self.style_name = style_name
        self.positive_prompt = positive_prompt
        self.negative_prompt = negative_prompt
        self.cfg = cfg
        self.denoise = denoise
        self.steps = steps
        self.reasoning = reasoning
        self.source = source


def _call_ollama(user_message: str) -> str:
    ollama_url = os.environ.get("OLLAMA_URL", "http://localhost:11434")
    # Using Qwen 3.5 27B as requested
    model = os.environ.get("OLLAMA_MODEL", "qwen3.5:27b") 
    full_prompt = f"{SYSTEM_PROMPT}\n\nUSER REQUEST:\n{user_message}\n\nReturn JSON ONLY. No markdown wrapping."
    
    r = requests.post(
        f"{ollama_url}/api/generate",
        json={
            "model": model, 
            "prompt": full_prompt,
            "stream": False, 
            "format": "json",
            "options": {"temperature": 0.5}
        },
        timeout=180 # Large models take time
    )
    r.raise_for_status()
    resp_json = r.json()
    return resp_json["response"]


def expand_prompt(keywords: list[str], room_type: str = "living_room", additional_notes: str = "") -> ExpandedPrompt:
    if not keywords:
        raise ValueError("Keywords required")

    room_type = room_type.lower().replace(" ", "_")
    if room_type not in ROOM_CONTEXTS:
        room_type = "living_room"

    ctx = ROOM_CONTEXTS[room_type]
    user_message = f"""KEYWORDS: {', '.join(keywords)}
ROOM TYPE: {room_type} ({ctx['description']})
MUST HAVE: {ctx['must_haves']}
AVOID: {ctx['avoid']}
{f'NOTES: {additional_notes}' if additional_notes else ''}"""

    try:
        logger.info("Calling local Ollama LLM...")
        raw = _call_ollama(user_message).strip()
        
        # Log the raw response for debugging
        logger.info(f"Raw Ollama response length: {len(raw)}")
        
        # Handle empty response
        if not raw:
            logger.error("Empty response from Ollama")
            raise ValueError("Ollama returned an empty response")

        # Clean string of common junk (control characters, leading/trailing non-bracket junk)
        raw_cleaned = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', raw)
        
        # More robust JSON block extraction
        # Try finding the first '{' and last '}'
        start_idx = raw_cleaned.find('{')
        end_idx = raw_cleaned.rfind('}')
        
        if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
            raw_json = raw_cleaned[start_idx:end_idx+1]
        else:
            raw_json = raw_cleaned

        try:
            data = json.loads(raw_json)
        except json.JSONDecodeError as je:
            logger.error(f"JSON Decode failed. Raw tail: {raw_json[:100]}...")
            # Emergency fallback: if we can't parse, try to create a basic structure from keywords
            return ExpandedPrompt(
                style_name="fallback_" + (keywords[0] if keywords else "custom"),
                positive_prompt=", ".join(keywords) + f", photorealistic, high quality, {room_type}",
                negative_prompt="blurry, distorted, low quality",
                cfg=7.5,
                denoise=0.55,
                steps=30,
                reasoning="Fallback due to LLM parsing error",
                source="fallback"
            )
        
        return ExpandedPrompt(
            style_name=data.get("style_name", "custom_local"),
            positive_prompt=data.get("positive_prompt", ", ".join(keywords)),
            negative_prompt=data.get("negative_prompt", "blurry, distorted, low quality"),
            cfg=float(data.get("cfg", 7.5)),
            denoise=float(data.get("denoise", 0.55)),
            steps=int(data.get("steps", 30)),
            reasoning=data.get("reasoning", ""),
            source="ollama_local"
        )
    except Exception as e:
        logger.error(f"Ollama LLM expansion failed: {e}")
        import traceback
        traceback.print_exc()
        # Full fallback if anything goes wrong
        return ExpandedPrompt(
            style_name="emergency_fallback",
            positive_prompt=", ".join(keywords) + f", photorealistic, {room_type}",
            negative_prompt="blurry, distorted",
            cfg=7.5,
            denoise=0.55,
            steps=30,
            reasoning=f"Emergency fallback due to: {str(e)}",
            source="emergency"
        )

if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO)
    ans = expand_prompt(["blue marble", "gold"], "bathroom")
    print(ans.positive_prompt)