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.5-0.85)
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()
    return r.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()
        
        json_match = re.search(r'\{[\s\S]*\}', raw)
        if json_match:
            raw_json = json_match.group(0)
        else:
            raw_json = raw
            
        data = json.loads(raw_json)
        
        return ExpandedPrompt(
            style_name=data.get("style_name", "custom_local"),
            positive_prompt=data["positive_prompt"],
            negative_prompt=data["negative_prompt"],
            cfg=float(data.get("cfg", 7.5)),
            denoise=float(data.get("denoise", 0.72)),
            steps=int(data.get("steps", 30)),
            reasoning=data.get("reasoning", ""),
            source="ollama_local"
        )
    except Exception as e:
        logger.warning(f"Ollama failed, using sync fallback: {e}")
        return expand_prompt_simple(keywords, room_type)


def expand_prompt_simple(keywords: list[str], room_type: str = "living_room") -> ExpandedPrompt:
    ctx = ROOM_CONTEXTS.get(room_type.replace(" ", "_"), ROOM_CONTEXTS["living_room"])
    kw_str = ", ".join(keywords)
    positive = f"{kw_str} interior design, {', '.join(ctx['key_elements'][:4])}, photorealistic {room_type.replace('_', ' ')} interior, architectural photography, 8k resolution, photorealistic"
    negative = "(worst quality, low quality, illustration, 3d render, 2d, painting, cartoon, sketch), blurry, distorted, extra windows, unrealistic lighting, structural changes"
    return ExpandedPrompt(
        style_name="fallback", positive_prompt=positive, negative_prompt=negative, 
        cfg=7.5, denoise=0.72, steps=30, reasoning="No LLM", source="fallback"
    )

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