sayan/Project_Velocity
1
1
Fork 0
forked from sagnik/Project_Velocity
Code Issues Pull Requests Activity

feat: Overlay the mathematical Sun Path over the live camera feed or 3D model view (#8)

Browse Source 
#7 Task completed.

Co-authored-by: Sayan Datta <sayan@Sayans-MacBook-Air.local>
Reviewed-on: sagnik/Project_Velocity#8
This commit is contained in:
sayan
2026-03-21 17:01:06 +05:30
parent 6c98affe53
commit 023ba48da2
16 changed files with 3344 additions and 2939 deletions
Download Patch File Download Diff File Expand all files Collapse all files

342
comfy_engine/scripts/README.md
View File

@@ -1,171 +1,171 @@
# Dream Weaver Automation Scripts
This directory contains Python automation scripts for the Dream Weaver interior restyling workflow.
## Scripts Overview
### 1. dreamweaver_batch_processor.py
Main batch processing controller for automated image restyling.
**Features:**
- Directory monitoring for automatic job queueing
- Automatic mask caching for improved performance
- Queue management with status tracking
- Support for all three processing phases
- WebSocket integration with ComfyUI for real-time status
**Usage:**
```bash
# Process single image
python dreamweaver_batch_processor.py --input image.jpg --style scandinavian_minimalist --phase 1
# Process all images in directory
python dreamweaver_batch_processor.py --batch --style art_deco_luxe --phase 2
# Start directory monitoring mode
python dreamweaver_batch_processor.py --monitor
```
### 2. mask_preprocessor.py
Utility for preprocessing and caching segmentation masks.
**Features:**
- Offline mask generation and caching
- Mask refinement (grow, feather, invert)
- Multi-region mask support (walls, floor, ceiling)
- Batch preprocessing for entire directories
- Cache management and statistics
**Usage:**
```bash
# Preprocess single image
python mask_preprocessor.py --image image.jpg
# Preprocess entire directory
python mask_preprocessor.py --directory ../test_inputs/
# Show cache statistics
python mask_preprocessor.py --stats
# Clear all cached masks
python mask_preprocessor.py --clear-cache
# Custom mask parameters
python mask_preprocessor.py --image image.jpg --grow 5 --feather 8
```
## Configuration
Scripts use configuration from `CONFIG` dictionary in each file. Key settings:
- `comfyui_server`: ComfyUI HTTP endpoint (default: http://localhost:8188)
- `comfyui_ws`: ComfyUI WebSocket endpoint (default: ws://localhost:8188/ws)
- `input_directory`: Default input images directory
- `output_directory`: Generated images output directory
- `cache_directory`: Mask cache storage location
- `batch_size`: Number of images to process in batch (Phase 3)
## Integration with ComfyUI
These scripts require ComfyUI to be running with the Dream Weaver workflows loaded.
**Starting ComfyUI:**
```bash
cd Project_Velocity/comfy_engine
python main.py --fp16 --lowvram
```
**For Production (Dual RTX PRO 6000):**
```bash
python main.py --bf16 --highvram --xformers --gpu-batch-size 8
```
## Workflow Files
Scripts reference these workflow JSON files:
- `workflows/dreamweaver_phase1_depth.json` - Single ControlNet (RTX 3080Ti)
- `workflows/dreamweaver_phase2_multicontrol.json` - Multi-ControlNet (RTX 3080Ti)
- `workflows/dreamweaver_phase3_batch.json` - Batch processing (Dual RTX PRO 6000)
## Style Templates
Available style templates (located in `../prompts/`):
- `scandinavian_minimalist` - Light, airy Nordic design
- `art_deco_luxe` - Glamorous 1920s aesthetic
- `cyberpunk_neon` - High-tech futuristic
- `biophilic_organic` - Nature-connected sustainable
- `japandi_fusion` - Japanese-Scandinavian blend
## Dependencies
Install required packages:
```bash
pip install -r ../requirements.txt
```
## Logging
Scripts output logs to:
- Console (real-time)
- `dreamweaver_batch.log` (file)
Log level can be adjusted in script `logging.basicConfig()` calls.
## Hardware Requirements
**Phase 1 & 2 (Development):**
- NVIDIA RTX 3080Ti (12GB VRAM)
- 32GB System RAM
- SSD Storage
**Phase 3 (Production):**
- Dual NVIDIA RTX PRO 6000 Blackwell (96GB VRAM each)
- 128GB System RAM
- NVMe SSD Storage
- NVLink enabled for GPU memory pooling
## API Reference
### ComfyUI Endpoints Used
- `POST /prompt` - Submit workflow to queue
- `GET /queue` - Get queue status
- `WS /ws` - WebSocket for real-time updates
### Job Status Values
- `pending` - Waiting in queue
- `processing` - Currently generating
- `completed` - Successfully finished
- `failed` - Error occurred
## Troubleshooting
**Connection Refused Error:**
- Ensure ComfyUI is running
- Check server URL in configuration
- Verify firewall settings
**Out of Memory:**
- Reduce batch size
- Lower resolution
- Enable tiled VAE decoding
**Mask Cache Issues:**
- Clear cache: `python mask_preprocessor.py --clear-cache`
- Check cache directory permissions
- Verify available disk space
## Development
To extend functionality:
1. Modify `BatchProcessor` class for new processing logic
2. Add new style templates in `../prompts/`
3. Update workflow JSON files for new ControlNet configurations
## Support
For issues related to:
- **Scripts**: Check logs in `dreamweaver_batch.log`
- **ComfyUI**: Refer to ComfyUI documentation
- **Workflows**: See technical specification in `../docs/DREAMWEAVER_TECHNICAL_SPEC.md`
# Dream Weaver Automation Scripts
This directory contains Python automation scripts for the Dream Weaver interior restyling workflow.
## Scripts Overview
### 1. dreamweaver_batch_processor.py
Main batch processing controller for automated image restyling.
**Features:**
- Directory monitoring for automatic job queueing
- Automatic mask caching for improved performance
- Queue management with status tracking
- Support for all three processing phases
- WebSocket integration with ComfyUI for real-time status
**Usage:**
```bash
# Process single image
python dreamweaver_batch_processor.py --input image.jpg --style scandinavian_minimalist --phase 1
# Process all images in directory
python dreamweaver_batch_processor.py --batch --style art_deco_luxe --phase 2
# Start directory monitoring mode
python dreamweaver_batch_processor.py --monitor
```
### 2. mask_preprocessor.py
Utility for preprocessing and caching segmentation masks.
**Features:**
- Offline mask generation and caching
- Mask refinement (grow, feather, invert)
- Multi-region mask support (walls, floor, ceiling)
- Batch preprocessing for entire directories
- Cache management and statistics
**Usage:**
```bash
# Preprocess single image
python mask_preprocessor.py --image image.jpg
# Preprocess entire directory
python mask_preprocessor.py --directory ../test_inputs/
# Show cache statistics
python mask_preprocessor.py --stats
# Clear all cached masks
python mask_preprocessor.py --clear-cache
# Custom mask parameters
python mask_preprocessor.py --image image.jpg --grow 5 --feather 8
```
## Configuration
Scripts use configuration from `CONFIG` dictionary in each file. Key settings:
- `comfyui_server`: ComfyUI HTTP endpoint (default: http://localhost:8188)
- `comfyui_ws`: ComfyUI WebSocket endpoint (default: ws://localhost:8188/ws)
- `input_directory`: Default input images directory
- `output_directory`: Generated images output directory
- `cache_directory`: Mask cache storage location
- `batch_size`: Number of images to process in batch (Phase 3)
## Integration with ComfyUI
These scripts require ComfyUI to be running with the Dream Weaver workflows loaded.
**Starting ComfyUI:**
```bash
cd Project_Velocity/comfy_engine
python main.py --fp16 --lowvram
```
**For Production (Dual RTX PRO 6000):**
```bash
python main.py --bf16 --highvram --xformers --gpu-batch-size 8
```
## Workflow Files
Scripts reference these workflow JSON files:
- `workflows/dreamweaver_phase1_depth.json` - Single ControlNet (RTX 3080Ti)
- `workflows/dreamweaver_phase2_multicontrol.json` - Multi-ControlNet (RTX 3080Ti)
- `workflows/dreamweaver_phase3_batch.json` - Batch processing (Dual RTX PRO 6000)
## Style Templates
Available style templates (located in `../prompts/`):
- `scandinavian_minimalist` - Light, airy Nordic design
- `art_deco_luxe` - Glamorous 1920s aesthetic
- `cyberpunk_neon` - High-tech futuristic
- `biophilic_organic` - Nature-connected sustainable
- `japandi_fusion` - Japanese-Scandinavian blend
## Dependencies
Install required packages:
```bash
pip install -r ../requirements.txt
```
## Logging
Scripts output logs to:
- Console (real-time)
- `dreamweaver_batch.log` (file)
Log level can be adjusted in script `logging.basicConfig()` calls.
## Hardware Requirements
**Phase 1 & 2 (Development):**
- NVIDIA RTX 3080Ti (12GB VRAM)
- 32GB System RAM
- SSD Storage
**Phase 3 (Production):**
- Dual NVIDIA RTX PRO 6000 Blackwell (96GB VRAM each)
- 128GB System RAM
- NVMe SSD Storage
- NVLink enabled for GPU memory pooling
## API Reference
### ComfyUI Endpoints Used
- `POST /prompt` - Submit workflow to queue
- `GET /queue` - Get queue status
- `WS /ws` - WebSocket for real-time updates
### Job Status Values
- `pending` - Waiting in queue
- `processing` - Currently generating
- `completed` - Successfully finished
- `failed` - Error occurred
## Troubleshooting
**Connection Refused Error:**
- Ensure ComfyUI is running
- Check server URL in configuration
- Verify firewall settings
**Out of Memory:**
- Reduce batch size
- Lower resolution
- Enable tiled VAE decoding
**Mask Cache Issues:**
- Clear cache: `python mask_preprocessor.py --clear-cache`
- Check cache directory permissions
- Verify available disk space
## Development
To extend functionality:
1. Modify `BatchProcessor` class for new processing logic
2. Add new style templates in `../prompts/`
3. Update workflow JSON files for new ControlNet configurations
## Support
For issues related to:
- **Scripts**: Check logs in `dreamweaver_batch.log`
- **ComfyUI**: Refer to ComfyUI documentation
- **Workflows**: See technical specification in `../docs/DREAMWEAVER_TECHNICAL_SPEC.md`

1244
comfy_engine/scripts/a100_deployment_executor.py
View File

File diff suppressed because it is too large Load Diff

996
comfy_engine/scripts/dreamweaver_batch_processor.py
View File

@@ -1,498 +1,498 @@
#!/usr/bin/env python3
"""
Dream Weaver Batch Processor
============================
Automated batch processing script for Dream Weaver interior restyling workflow.
Handles directory monitoring, automatic mask caching, and queue management.
Target Hardware: Dual NVIDIA RTX PRO 6000 Blackwell (96GB GDDR7 each)
Author: Project Velocity Team
Version: 1.0.0
"""
import os
import sys
import json
import time
import hashlib
import asyncio
import argparse
import logging
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
import requests
import websockets
import aiofiles
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
# Configuration
CONFIG = {
"comfyui_server": "http://localhost:8188",
"comfyui_ws": "ws://localhost:8188/ws",
"input_directory": "Project_Velocity/comfy_engine/test_inputs/",
"output_directory": "Project_Velocity/comfy_engine/test_outputs/",
"cache_directory": "Project_Velocity/comfy_engine/cache/masks/",
"workflow_phase1": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase1_depth.json",
"workflow_phase2": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase2_multicontrol.json",
"workflow_phase3": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase3_batch.json",
"batch_size": 8,
"target_resolution": (1024, 1024),
"enable_mask_cache": True,
"gpu_sharding": True,
"dual_gpu": True,
}
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('dreamweaver_batch.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger('DreamWeaver')
@dataclass
class ProcessingJob:
"""Represents a single image processing job."""
job_id: str
input_path: str
output_path: str
style_template: str
phase: int
status: str = "pending"
created_at: datetime = None
started_at: datetime = None
completed_at: datetime = None
error_message: str = None
mask_cached: bool = False
def __post_init__(self):
if self.created_at is None:
self.created_at = datetime.now()
def to_dict(self) -> Dict:
return {
"job_id": self.job_id,
"input_path": self.input_path,
"output_path": self.output_path,
"style_template": self.style_template,
"phase": self.phase,
"status": self.status,
"created_at": self.created_at.isoformat() if self.created_at else None,
"started_at": self.started_at.isoformat() if self.started_at else None,
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
"error_message": self.error_message,
"mask_cached": self.mask_cached
}
class MaskCacheManager:
"""Manages caching of segmentation masks for improved performance."""
def __init__(self, cache_dir: str):
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Mask cache initialized at: {self.cache_dir}")
def _get_cache_key(self, image_path: str) -> str:
"""Generate cache key from image content hash."""
hasher = hashlib.md5()
with open(image_path, 'rb') as f:
hasher.update(f.read())
return hasher.hexdigest()
def get_cached_mask(self, image_path: str) -> Optional[str]:
"""Retrieve cached mask path if it exists."""
cache_key = self._get_cache_key(image_path)
cached_path = self.cache_dir / f"{cache_key}.png"
if cached_path.exists():
logger.info(f"Cache hit for {image_path}")
return str(cached_path)
return None
def cache_mask(self, image_path: str, mask_path: str) -> str:
"""Cache a mask file for future use."""
cache_key = self._get_cache_key(image_path)
cached_path = self.cache_dir / f"{cache_key}.png"
import shutil
shutil.copy2(mask_path, cached_path)
logger.info(f"Cached mask for {image_path} at {cached_path}")
return str(cached_path)
class ComfyUIClient:
"""Client for communicating with ComfyUI server."""
def __init__(self, server_url: str, ws_url: str):
self.server_url = server_url
self.ws_url = ws_url
self.client_id = self._generate_client_id()
logger.info(f"ComfyUI client initialized with ID: {self.client_id}")
def _generate_client_id(self) -> str:
"""Generate unique client ID."""
return f"dreamweaver_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{os.urandom(4).hex()}"
async def submit_workflow(self, workflow: Dict, input_image: str) -> str:
"""Submit a workflow to ComfyUI queue."""
# Update workflow with input image
for node_id, node in workflow.items():
if node.get("class_type") == "LoadImage":
node["inputs"]["image"] = input_image
if node.get("class_type") == "LoadImageBatch":
node["inputs"]["directory"] = os.path.dirname(input_image)
payload = {
"prompt": workflow,
"client_id": self.client_id
}
response = requests.post(
f"{self.server_url}/prompt",
json=payload
)
response.raise_for_status()
result = response.json()
prompt_id = result.get("prompt_id")
logger.info(f"Submitted workflow with prompt_id: {prompt_id}")
return prompt_id
async def get_queue_status(self) -> Dict:
"""Get current queue status."""
response = requests.get(f"{self.server_url}/queue")
return response.json()
async def wait_for_completion(self, prompt_id: str, timeout: int = 300) -> bool:
"""Wait for workflow completion via WebSocket."""
start_time = time.time()
async with websockets.connect(
f"{self.ws_url}?clientId={self.client_id}"
) as websocket:
while time.time() - start_time < timeout:
try:
message = await asyncio.wait_for(
websocket.recv(),
timeout=5.0
)
data = json.loads(message)
if data.get("type") == "executing":
if data["data"].get("prompt_id") == prompt_id:
node_id = data["data"].get("node")
logger.debug(f"Executing node: {node_id}")
elif data.get("type") == "completed":
if data["data"].get("prompt_id") == prompt_id:
logger.info(f"Workflow {prompt_id} completed")
return True
elif data.get("type") == "error":
logger.error(f"Workflow error: {data}")
return False
except asyncio.TimeoutError:
continue
logger.warning(f"Workflow {prompt_id} timed out")
return False
class BatchProcessor:
"""Main batch processing controller."""
def __init__(self, config: Dict):
self.config = config
self.queue: List[ProcessingJob] = []
self.processing = False
self.cache_manager = MaskCacheManager(config["cache_directory"])
self.comfy_client = ComfyUIClient(
config["comfyui_server"],
config["comfyui_ws"]
)
# Load workflow templates
self.workflows = self._load_workflows()
# Ensure output directory exists
Path(config["output_directory"]).mkdir(parents=True, exist_ok=True)
def _load_workflows(self) -> Dict[int, Dict]:
"""Load workflow JSON files."""
workflows = {}
workflow_paths = {
1: self.config["workflow_phase1"],
2: self.config["workflow_phase2"],
3: self.config["workflow_phase3"]
}
for phase, path in workflow_paths.items():
try:
with open(path, 'r') as f:
workflows[phase] = json.load(f)
logger.info(f"Loaded Phase {phase} workflow")
except Exception as e:
logger.error(f"Failed to load Phase {phase} workflow: {e}")
return workflows
def add_job(self, input_path: str, style_template: str = "scandinavian_minimalist", phase: int = 1) -> str:
"""Add a new processing job to the queue."""
job_id = hashlib.md5(f"{input_path}_{time.time()}".encode()).hexdigest()[:12]
output_filename = f"{Path(input_path).stem}_restyled_{job_id}.png"
output_path = os.path.join(self.config["output_directory"], output_filename)
job = ProcessingJob(
job_id=job_id,
input_path=input_path,
output_path=output_path,
style_template=style_template,
phase=phase
)
# Check if mask is cached
if self.config["enable_mask_cache"]:
cached_mask = self.cache_manager.get_cached_mask(input_path)
job.mask_cached = cached_mask is not None
self.queue.append(job)
logger.info(f"Added job {job_id} to queue. Queue size: {len(self.queue)}")
return job_id
async def process_single(self, job: ProcessingJob) -> bool:
"""Process a single job."""
job.status = "processing"
job.started_at = datetime.now()
try:
logger.info(f"Processing job {job.job_id}: {job.input_path}")
# Get workflow for phase
workflow = self.workflows.get(job.phase)
if not workflow:
raise ValueError(f"Workflow for phase {job.phase} not found")
# Submit to ComfyUI
prompt_id = await self.comfy_client.submit_workflow(
workflow,
job.input_path
)
# Wait for completion
success = await self.comfy_client.wait_for_completion(prompt_id)
if success:
job.status = "completed"
job.completed_at = datetime.now()
logger.info(f"Job {job.job_id} completed successfully")
return True
else:
job.status = "failed"
job.error_message = "Workflow execution failed or timed out"
logger.error(f"Job {job.job_id} failed")
return False
except Exception as e:
job.status = "failed"
job.error_message = str(e)
logger.error(f"Error processing job {job.job_id}: {e}")
return False
async def process_batch(self, jobs: List[ProcessingJob]) -> List[bool]:
"""Process multiple jobs in batch (Phase 3)."""
if not jobs:
return []
logger.info(f"Processing batch of {len(jobs)} jobs")
results = []
# For batch processing, use Phase 3 workflow
workflow = self.workflows.get(3)
if not workflow:
logger.warning("Phase 3 workflow not available, processing sequentially")
for job in jobs:
result = await self.process_single(job)
results.append(result)
return results
# TODO: Implement true batch processing with Phase 3 workflow
# This would require grouping images and processing together
for job in jobs:
result = await self.process_single(job)
results.append(result)
return results
async def run(self):
"""Main processing loop."""
logger.info("Starting batch processor")
self.processing = True
while self.processing:
# Get pending jobs
pending_jobs = [j for j in self.queue if j.status == "pending"]
if not pending_jobs:
await asyncio.sleep(1)
continue
# Check if batch processing is appropriate
if len(pending_jobs) >= self.config["batch_size"] and self.config.get("dual_gpu"):
# Process in batches for Phase 3
batch = pending_jobs[:self.config["batch_size"]]
await self.process_batch(batch)
else:
# Process single job with appropriate phase
job = pending_jobs[0]
await self.process_single(job)
def stop(self):
"""Stop the processing loop."""
logger.info("Stopping batch processor")
self.processing = False
def get_status(self) -> Dict:
"""Get current processing status."""
total = len(self.queue)
pending = len([j for j in self.queue if j.status == "pending"])
processing = len([j for j in self.queue if j.status == "processing"])
completed = len([j for j in self.queue if j.status == "completed"])
failed = len([j for j in self.queue if j.status == "failed"])
return {
"total_jobs": total,
"pending": pending,
"processing": processing,
"completed": completed,
"failed": failed,
"is_running": self.processing
}
class InputDirectoryHandler(FileSystemEventHandler):
"""Handles new file events in input directory."""
def __init__(self, processor: BatchProcessor):
self.processor = processor
def on_created(self, event):
if not event.is_directory:
file_path = event.src_path
if file_path.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
logger.info(f"New image detected: {file_path}")
self.processor.add_job(file_path)
def load_style_template(template_name: str) -> str:
"""Load a style template from prompts directory."""
template_path = Path("Project_Velocity/comfy_engine/prompts/") / f"{template_name}.txt"
if template_path.exists():
with open(template_path, 'r') as f:
content = f.read()
# Extract positive prompt
lines = content.split('\n')
positive_lines = []
in_positive = False
for line in lines:
if 'POSITIVE PROMPT:' in line:
in_positive = True
continue
if in_positive and line.startswith('Style Weight:'):
break
if in_positive and line.strip() and not line.startswith('-'):
positive_lines.append(line.strip())
return ' '.join(positive_lines)
return ""
async def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="Dream Weaver Batch Processor"
)
parser.add_argument(
"--monitor",
action="store_true",
help="Enable directory monitoring mode"
)
parser.add_argument(
"--input",
type=str,
help="Single input image to process"
)
parser.add_argument(
"--style",
type=str,
default="scandinavian_minimalist",
choices=["scandinavian_minimalist", "art_deco_luxe", "cyberpunk_neon", "biophilic_organic", "japandi_fusion"],
help="Style template to apply"
)
parser.add_argument(
"--phase",
type=int,
default=1,
choices=[1, 2, 3],
help="Processing phase to use"
)
parser.add_argument(
"--batch",
action="store_true",
help="Process all images in input directory"
)
args = parser.parse_args()
# Initialize processor
processor = BatchProcessor(CONFIG)
if args.input:
# Process single image
job_id = processor.add_job(args.input, args.style, args.phase)
await processor.process_single(processor.queue[-1])
print(f"Processed image: {args.input}")
print(f"Job ID: {job_id}")
elif args.batch:
# Process all images in directory
input_dir = Path(CONFIG["input_directory"])
image_files = list(input_dir.glob("*.jpg")) + list(input_dir.glob("*.png"))
for img_file in image_files:
processor.add_job(str(img_file), args.style, args.phase)
await processor.run()
elif args.monitor:
# Start directory monitoring
event_handler = InputDirectoryHandler(processor)
observer = Observer()
observer.schedule(
event_handler,
CONFIG["input_directory"],
recursive=False
)
observer.start()
logger.info(f"Started monitoring: {CONFIG['input_directory']}")
try:
# Run processor
await processor.run()
except KeyboardInterrupt:
processor.stop()
observer.stop()
observer.join()
else:
print("No action specified. Use --help for usage information.")
if __name__ == "__main__":
asyncio.run(main())
#!/usr/bin/env python3
"""
Dream Weaver Batch Processor
============================
Automated batch processing script for Dream Weaver interior restyling workflow.
Handles directory monitoring, automatic mask caching, and queue management.
Target Hardware: Dual NVIDIA RTX PRO 6000 Blackwell (96GB GDDR7 each)
Author: Project Velocity Team
Version: 1.0.0
"""
import os
import sys
import json
import time
import hashlib
import asyncio
import argparse
import logging
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
import requests
import websockets
import aiofiles
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
# Configuration
CONFIG = {
"comfyui_server": "http://localhost:8188",
"comfyui_ws": "ws://localhost:8188/ws",
"input_directory": "Project_Velocity/comfy_engine/test_inputs/",
"output_directory": "Project_Velocity/comfy_engine/test_outputs/",
"cache_directory": "Project_Velocity/comfy_engine/cache/masks/",
"workflow_phase1": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase1_depth.json",
"workflow_phase2": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase2_multicontrol.json",
"workflow_phase3": "Project_Velocity/comfy_engine/workflows/dreamweaver_phase3_batch.json",
"batch_size": 8,
"target_resolution": (1024, 1024),
"enable_mask_cache": True,
"gpu_sharding": True,
"dual_gpu": True,
}
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('dreamweaver_batch.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger('DreamWeaver')
@dataclass
class ProcessingJob:
"""Represents a single image processing job."""
job_id: str
input_path: str
output_path: str
style_template: str
phase: int
status: str = "pending"
created_at: datetime = None
started_at: datetime = None
completed_at: datetime = None
error_message: str = None
mask_cached: bool = False
def __post_init__(self):
if self.created_at is None:
self.created_at = datetime.now()
def to_dict(self) -> Dict:
return {
"job_id": self.job_id,
"input_path": self.input_path,
"output_path": self.output_path,
"style_template": self.style_template,
"phase": self.phase,
"status": self.status,
"created_at": self.created_at.isoformat() if self.created_at else None,
"started_at": self.started_at.isoformat() if self.started_at else None,
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
"error_message": self.error_message,
"mask_cached": self.mask_cached
}
class MaskCacheManager:
"""Manages caching of segmentation masks for improved performance."""
def __init__(self, cache_dir: str):
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Mask cache initialized at: {self.cache_dir}")
def _get_cache_key(self, image_path: str) -> str:
"""Generate cache key from image content hash."""
hasher = hashlib.md5()
with open(image_path, 'rb') as f:
hasher.update(f.read())
return hasher.hexdigest()
def get_cached_mask(self, image_path: str) -> Optional[str]:
"""Retrieve cached mask path if it exists."""
cache_key = self._get_cache_key(image_path)
cached_path = self.cache_dir / f"{cache_key}.png"
if cached_path.exists():
logger.info(f"Cache hit for {image_path}")
return str(cached_path)
return None
def cache_mask(self, image_path: str, mask_path: str) -> str:
"""Cache a mask file for future use."""
cache_key = self._get_cache_key(image_path)
cached_path = self.cache_dir / f"{cache_key}.png"
import shutil
shutil.copy2(mask_path, cached_path)
logger.info(f"Cached mask for {image_path} at {cached_path}")
return str(cached_path)
class ComfyUIClient:
"""Client for communicating with ComfyUI server."""
def __init__(self, server_url: str, ws_url: str):
self.server_url = server_url
self.ws_url = ws_url
self.client_id = self._generate_client_id()
logger.info(f"ComfyUI client initialized with ID: {self.client_id}")
def _generate_client_id(self) -> str:
"""Generate unique client ID."""
return f"dreamweaver_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{os.urandom(4).hex()}"
async def submit_workflow(self, workflow: Dict, input_image: str) -> str:
"""Submit a workflow to ComfyUI queue."""
# Update workflow with input image
for node_id, node in workflow.items():
if node.get("class_type") == "LoadImage":
node["inputs"]["image"] = input_image
if node.get("class_type") == "LoadImageBatch":
node["inputs"]["directory"] = os.path.dirname(input_image)
payload = {
"prompt": workflow,
"client_id": self.client_id
}
response = requests.post(
f"{self.server_url}/prompt",
json=payload
)
response.raise_for_status()
result = response.json()
prompt_id = result.get("prompt_id")
logger.info(f"Submitted workflow with prompt_id: {prompt_id}")
return prompt_id
async def get_queue_status(self) -> Dict:
"""Get current queue status."""
response = requests.get(f"{self.server_url}/queue")
return response.json()
async def wait_for_completion(self, prompt_id: str, timeout: int = 300) -> bool:
"""Wait for workflow completion via WebSocket."""
start_time = time.time()
async with websockets.connect(
f"{self.ws_url}?clientId={self.client_id}"
) as websocket:
while time.time() - start_time < timeout:
try:
message = await asyncio.wait_for(
websocket.recv(),
timeout=5.0
)
data = json.loads(message)
if data.get("type") == "executing":
if data["data"].get("prompt_id") == prompt_id:
node_id = data["data"].get("node")
logger.debug(f"Executing node: {node_id}")
elif data.get("type") == "completed":
if data["data"].get("prompt_id") == prompt_id:
logger.info(f"Workflow {prompt_id} completed")
return True
elif data.get("type") == "error":
logger.error(f"Workflow error: {data}")
return False
except asyncio.TimeoutError:
continue
logger.warning(f"Workflow {prompt_id} timed out")
return False
class BatchProcessor:
"""Main batch processing controller."""
def __init__(self, config: Dict):
self.config = config
self.queue: List[ProcessingJob] = []
self.processing = False
self.cache_manager = MaskCacheManager(config["cache_directory"])
self.comfy_client = ComfyUIClient(
config["comfyui_server"],
config["comfyui_ws"]
)
# Load workflow templates
self.workflows = self._load_workflows()
# Ensure output directory exists
Path(config["output_directory"]).mkdir(parents=True, exist_ok=True)
def _load_workflows(self) -> Dict[int, Dict]:
"""Load workflow JSON files."""
workflows = {}
workflow_paths = {
1: self.config["workflow_phase1"],
2: self.config["workflow_phase2"],
3: self.config["workflow_phase3"]
}
for phase, path in workflow_paths.items():
try:
with open(path, 'r') as f:
workflows[phase] = json.load(f)
logger.info(f"Loaded Phase {phase} workflow")
except Exception as e:
logger.error(f"Failed to load Phase {phase} workflow: {e}")
return workflows
def add_job(self, input_path: str, style_template: str = "scandinavian_minimalist", phase: int = 1) -> str:
"""Add a new processing job to the queue."""
job_id = hashlib.md5(f"{input_path}_{time.time()}".encode()).hexdigest()[:12]
output_filename = f"{Path(input_path).stem}_restyled_{job_id}.png"
output_path = os.path.join(self.config["output_directory"], output_filename)
job = ProcessingJob(
job_id=job_id,
input_path=input_path,
output_path=output_path,
style_template=style_template,
phase=phase
)
# Check if mask is cached
if self.config["enable_mask_cache"]:
cached_mask = self.cache_manager.get_cached_mask(input_path)
job.mask_cached = cached_mask is not None
self.queue.append(job)
logger.info(f"Added job {job_id} to queue. Queue size: {len(self.queue)}")
return job_id
async def process_single(self, job: ProcessingJob) -> bool:
"""Process a single job."""
job.status = "processing"
job.started_at = datetime.now()
try:
logger.info(f"Processing job {job.job_id}: {job.input_path}")
# Get workflow for phase
workflow = self.workflows.get(job.phase)
if not workflow:
raise ValueError(f"Workflow for phase {job.phase} not found")
# Submit to ComfyUI
prompt_id = await self.comfy_client.submit_workflow(
workflow,
job.input_path
)
# Wait for completion
success = await self.comfy_client.wait_for_completion(prompt_id)
if success:
job.status = "completed"
job.completed_at = datetime.now()
logger.info(f"Job {job.job_id} completed successfully")
return True
else:
job.status = "failed"
job.error_message = "Workflow execution failed or timed out"
logger.error(f"Job {job.job_id} failed")
return False
except Exception as e:
job.status = "failed"
job.error_message = str(e)
logger.error(f"Error processing job {job.job_id}: {e}")
return False
async def process_batch(self, jobs: List[ProcessingJob]) -> List[bool]:
"""Process multiple jobs in batch (Phase 3)."""
if not jobs:
return []
logger.info(f"Processing batch of {len(jobs)} jobs")
results = []
# For batch processing, use Phase 3 workflow
workflow = self.workflows.get(3)
if not workflow:
logger.warning("Phase 3 workflow not available, processing sequentially")
for job in jobs:
result = await self.process_single(job)
results.append(result)
return results
# TODO: Implement true batch processing with Phase 3 workflow
# This would require grouping images and processing together
for job in jobs:
result = await self.process_single(job)
results.append(result)
return results
async def run(self):
"""Main processing loop."""
logger.info("Starting batch processor")
self.processing = True
while self.processing:
# Get pending jobs
pending_jobs = [j for j in self.queue if j.status == "pending"]
if not pending_jobs:
await asyncio.sleep(1)
continue
# Check if batch processing is appropriate
if len(pending_jobs) >= self.config["batch_size"] and self.config.get("dual_gpu"):
# Process in batches for Phase 3
batch = pending_jobs[:self.config["batch_size"]]
await self.process_batch(batch)
else:
# Process single job with appropriate phase
job = pending_jobs[0]
await self.process_single(job)
def stop(self):
"""Stop the processing loop."""
logger.info("Stopping batch processor")
self.processing = False
def get_status(self) -> Dict:
"""Get current processing status."""
total = len(self.queue)
pending = len([j for j in self.queue if j.status == "pending"])
processing = len([j for j in self.queue if j.status == "processing"])
completed = len([j for j in self.queue if j.status == "completed"])
failed = len([j for j in self.queue if j.status == "failed"])
return {
"total_jobs": total,
"pending": pending,
"processing": processing,
"completed": completed,
"failed": failed,
"is_running": self.processing
}
class InputDirectoryHandler(FileSystemEventHandler):
"""Handles new file events in input directory."""
def __init__(self, processor: BatchProcessor):
self.processor = processor
def on_created(self, event):
if not event.is_directory:
file_path = event.src_path
if file_path.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
logger.info(f"New image detected: {file_path}")
self.processor.add_job(file_path)
def load_style_template(template_name: str) -> str:
"""Load a style template from prompts directory."""
template_path = Path("Project_Velocity/comfy_engine/prompts/") / f"{template_name}.txt"
if template_path.exists():
with open(template_path, 'r') as f:
content = f.read()
# Extract positive prompt
lines = content.split('\n')
positive_lines = []
in_positive = False
for line in lines:
if 'POSITIVE PROMPT:' in line:
in_positive = True
continue
if in_positive and line.startswith('Style Weight:'):
break
if in_positive and line.strip() and not line.startswith('-'):
positive_lines.append(line.strip())
return ' '.join(positive_lines)
return ""
async def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="Dream Weaver Batch Processor"
)
parser.add_argument(
"--monitor",
action="store_true",
help="Enable directory monitoring mode"
)
parser.add_argument(
"--input",
type=str,
help="Single input image to process"
)
parser.add_argument(
"--style",
type=str,
default="scandinavian_minimalist",
choices=["scandinavian_minimalist", "art_deco_luxe", "cyberpunk_neon", "biophilic_organic", "japandi_fusion"],
help="Style template to apply"
)
parser.add_argument(
"--phase",
type=int,
default=1,
choices=[1, 2, 3],
help="Processing phase to use"
)
parser.add_argument(
"--batch",
action="store_true",
help="Process all images in input directory"
)
args = parser.parse_args()
# Initialize processor
processor = BatchProcessor(CONFIG)
if args.input:
# Process single image
job_id = processor.add_job(args.input, args.style, args.phase)
await processor.process_single(processor.queue[-1])
print(f"Processed image: {args.input}")
print(f"Job ID: {job_id}")
elif args.batch:
# Process all images in directory
input_dir = Path(CONFIG["input_directory"])
image_files = list(input_dir.glob("*.jpg")) + list(input_dir.glob("*.png"))
for img_file in image_files:
processor.add_job(str(img_file), args.style, args.phase)
await processor.run()
elif args.monitor:
# Start directory monitoring
event_handler = InputDirectoryHandler(processor)
observer = Observer()
observer.schedule(
event_handler,
CONFIG["input_directory"],
recursive=False
)
observer.start()
logger.info(f"Started monitoring: {CONFIG['input_directory']}")
try:
# Run processor
await processor.run()
except KeyboardInterrupt:
processor.stop()
observer.stop()
observer.join()
else:
print("No action specified. Use --help for usage information.")
if __name__ == "__main__":
asyncio.run(main())

840
comfy_engine/scripts/dw_gateway_v2.py
View File

@@ -1,420 +1,420 @@
#!/usr/bin/env python3
"""
Dream Weaver API Gateway v2 — Dynamic Keyword → Local LLM → ComfyUI Pipeline
========================================================================
Port: 8080 (public-facing)
ComfyUI: localhost:8188 (internal)
NEW IN v2:
- POST /dream-weaver now accepts keywords[] + room_type for LLM-based prompt generation
- POST /dream-weaver/expand — expand keywords to prompt WITHOUT generating (preview)
- GET /room-types — list available room types
- Uses local Ollama model (qwen3.5:27b) for prompt expansion (no cloud API dependencies)
Environment variables:
OLLAMA_URL — Ollama server (default: http://localhost:11434)
OLLAMA_MODEL — Model name (default: qwen3.5:27b)
"""
import asyncio, json, time, uuid, io, sys, os, logging
from pathlib import Path
from typing import Optional, List
import httpx
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, BackgroundTasks
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
# Add scripts dir to path so we can import prompt_expander
SCRIPTS_DIR = Path(__file__).parent / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))
try:
from prompt_expander import expand_prompt, expand_prompt_simple, ROOM_CONTEXTS, ExpandedPrompt
LLM_AVAILABLE = True
except ImportError:
LLM_AVAILABLE = False
logging.warning("prompt_expander not found — LLM expansion disabled")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("DreamWeaverGateway")
COMFY = "http://127.0.0.1:8188"
COMFY_ROOT = "/opt/dlami/nvme/ComfyUI"
app = FastAPI(
title="Dream Weaver API v2",
version="2.0.0",
description="Dynamic keyword-to-interior-design generation powered by LLM + ComfyUI"
)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
# In-memory job store (swap for Redis in production)
jobs: dict = {}
# ─── Models ──────────────────────────────────────────────────────────────────
class ExpandRequest(BaseModel):
keywords: List[str]
room_type: str = "living_room"
additional_notes: str = ""
class ExpandResponse(BaseModel):
style_name: str
positive_prompt: str
negative_prompt: str
cfg: float
denoise: float
steps: int
reasoning: str
source: str
# ─── ComfyUI helpers ──────────────────────────────────────────────────────────
async def upload_to_comfy(data: bytes, filename: str) -> str:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(f"{COMFY}/upload/image",
files={"image": (filename, data, "image/jpeg")},
data={"overwrite": "true"})
r.raise_for_status()
return r.json()["name"]
def build_workflow(img_name: str, expanded: "ExpandedPrompt") -> dict:
"""Build ComfyUI API workflow from an ExpandedPrompt result."""
return {
"1": {"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "realvisxlV50_v50LightningBakedvae.safetensors"}},
"2": {"class_type": "LoadImage",
"inputs": {"image": img_name, "upload": "image"}},
"3": {"class_type": "CLIPTextEncode", # Positive prompt
"inputs": {"text": expanded.positive_prompt, "clip": ["1", 1]}},
"4": {"class_type": "CLIPTextEncode", # Negative prompt
"inputs": {"text": expanded.negative_prompt, "clip": ["1", 1]}},
"5": {"class_type": "VAEEncode",
"inputs": {"pixels": ["2", 0], "vae": ["1", 2]}},
"6": {"class_type": "KSampler",
"inputs": {"model": ["1", 0],
"positive": ["3", 0],
"negative": ["4", 0],
"latent_image": ["5", 0],
"seed": int(time.time()) % 999983,
"steps": expanded.steps,
"cfg": expanded.cfg,
"sampler_name": "dpmpp_2m",
"scheduler": "karras",
"denoise": expanded.denoise}},
"7": {"class_type": "VAEDecode",
"inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
"8": {"class_type": "SaveImage",
"inputs": {"images": ["7", 0],
"filename_prefix": f"dw_{expanded.style_name.replace(' ', '_')[:30]}"}},
}
async def queue_prompt(workflow: dict) -> str:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(f"{COMFY}/prompt",
json={"prompt": workflow, "client_id": str(uuid.uuid4())})
r.raise_for_status()
return r.json()["prompt_id"]
async def poll_result(prompt_id: str, timeout: int = 300):
start = time.time()
async with httpx.AsyncClient(timeout=10) as client:
while time.time() - start < timeout:
r = await client.get(f"{COMFY}/history/{prompt_id}")
if r.status_code == 200:
h = r.json().get(prompt_id, {})
if h.get("status", {}).get("status_str") == "error":
return None, h.get("status", {}).get("messages", ["unknown"])
imgs = [img for nd in h.get("outputs", {}).values()
for img in nd.get("images", [])]
if imgs:
return imgs[0], None
await asyncio.sleep(2)
return None, "timeout"
async def background_poll(job_id: str, prompt_id: str):
img, err = await poll_result(prompt_id)
if img:
jobs[job_id].update({"status": "done", "output": img, "completed": time.time()})
else:
jobs[job_id].update({"status": "error", "error": str(err)})
# ─── Endpoints ───────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
comfy_ok = False
try:
async with httpx.AsyncClient(timeout=5) as c:
r = await c.get(f"{COMFY}/system_stats")
comfy_ok = r.status_code == 200
except Exception:
pass
return {
"status": "ok",
"comfyui": comfy_ok,
"gpu": "4x NVIDIA L4 (96GB VRAM)",
"model": "RealVisXL V5.0 Lightning",
"llm_expansion": LLM_AVAILABLE,
"version": "2.0.0"
}
@app.get("/room-types")
async def room_types():
"""List all supported room types with their context."""
if not LLM_AVAILABLE:
return {"room_types": ["bedroom", "living_room", "bathroom", "kitchen",
"dining_room", "home_office", "hallway", "balcony"]}
return {
"room_types": {
k: {
"description": v["description"],
"key_elements": v["key_elements"]
}
for k, v in ROOM_CONTEXTS.items()
}
}
@app.post("/dream-weaver/expand", response_model=ExpandResponse)
async def expand_endpoint(req: ExpandRequest):
"""
Preview the LLM-generated prompt WITHOUT submitting to ComfyUI.
Use this to let the user review/edit the prompt before generating.
Request body:
{
"keywords": ["blue marble", "gold veins", "renaissance", "sharp contours"],
"room_type": "bedroom",
"additional_notes": "luxury hotel feel"
}
"""
if not req.keywords:
raise HTTPException(status_code=400, detail="keywords list cannot be empty")
try:
if LLM_AVAILABLE:
result = await asyncio.to_thread(
expand_prompt,
keywords=req.keywords,
room_type=req.room_type,
additional_notes=req.additional_notes
)
else:
result = expand_prompt_simple(req.keywords, req.room_type)
except Exception as e:
logger.error(f"Prompt expansion failed: {e}")
raise HTTPException(status_code=500, detail=f"LLM expansion failed: {str(e)}")
return ExpandResponse(
style_name=result.style_name,
positive_prompt=result.positive_prompt,
negative_prompt=result.negative_prompt,
cfg=result.cfg,
denoise=result.denoise,
steps=result.steps,
reasoning=result.reasoning,
source=result.source
)
@app.post("/dream-weaver")
async def dream_weaver(
image: UploadFile = File(...),
# ── Dynamic keyword mode (new) ──
keywords: str = Form(default=""), # comma-separated: "blue marble, gold, renaissance"
room_type: str = Form(default="living_room"),
additional_notes: str = Form(default=""),
# ── Optional overrides ──
custom_positive: str = Form(default=""), # skip LLM, use this prompt directly
custom_negative: str = Form(default=""),
denoise: float = Form(default=0.0), # 0.0 = use LLM recommendation
cfg_scale: float = Form(default=0.0), # 0.0 = use LLM recommendation
):
"""
Submit a room photo for AI redesign using dynamic keyword → LLM → ComfyUI pipeline.
Two modes:
1. KEYWORD MODE (recommended): Provide keywords + room_type, LLM generates prompt
2. DIRECT MODE: Provide custom_positive + custom_negative to bypass LLM
Returns job_id for async polling.
"""
job_id = str(uuid.uuid4())
jobs[job_id] = {"status": "uploading", "created": time.time()}
try:
# Upload image to ComfyUI
data = await image.read()
filename = f"dw_{job_id[:8]}_{image.filename or 'room.jpg'}"
comfy_name = await upload_to_comfy(data, filename)
jobs[job_id]["status"] = "expanding_prompt"
# ── Determine prompt ──────────────────────────────────────────────
if custom_positive:
# Direct mode — user provided prompts explicitly
from dataclasses import dataclass
@dataclass
class DirectPrompt:
style_name: str = "custom"
positive_prompt: str = custom_positive
negative_prompt: str = custom_negative or (
"(worst quality, low quality, illustration, 3d render, painting, cartoon, sketch), "
"blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes"
)
cfg: float = cfg_scale or 7.5
denoise: float = denoise or 0.72
steps: int = 30
reasoning: str = "Direct user input"
source: str = "direct"
expanded = DirectPrompt()
elif keywords:
# Keyword mode — expand via LLM
kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
if LLM_AVAILABLE:
expanded = await asyncio.to_thread(
expand_prompt,
keywords=kw_list,
room_type=room_type,
additional_notes=additional_notes
)
else:
expanded = expand_prompt_simple(kw_list, room_type)
# Apply manual overrides if provided
if denoise > 0:
expanded.denoise = denoise
if cfg_scale > 0:
expanded.cfg = cfg_scale
else:
raise HTTPException(status_code=400,
detail="Provide either 'keywords' or 'custom_positive'")
jobs[job_id].update({
"status": "queued",
"style": expanded.style_name,
"prompt_source": expanded.source,
"positive_prompt": expanded.positive_prompt,
"negative_prompt": expanded.negative_prompt,
"room_type": room_type,
})
# Submit workflow
wf = build_workflow(comfy_name, expanded)
prompt_id = await queue_prompt(wf)
jobs[job_id].update({"status": "processing", "prompt_id": prompt_id})
# Start background polling
asyncio.create_task(background_poll(job_id, prompt_id))
return {
"job_id": job_id,
"status": "processing",
"style": expanded.style_name,
"prompt_preview": expanded.positive_prompt[:120] + "...",
"reasoning": expanded.reasoning,
"poll_url": f"/dream-weaver/status/{job_id}",
"result_url": f"/dream-weaver/result/{job_id}"
}
except HTTPException:
raise
except Exception as e:
jobs[job_id] = {"status": "error", "error": str(e)}
logger.error(f"Generation failed: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/dream-weaver/status/{job_id}")
async def status(job_id: str):
job = jobs.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found")
result = {k: v for k, v in job.items() if k != "output"}
result["ready"] = job.get("status") == "done"
if result["ready"]:
result["result_url"] = f"/dream-weaver/result/{job_id}"
return result
@app.get("/dream-weaver/result/{job_id}")
async def result(job_id: str):
job = jobs.get(job_id)
if not job or job.get("status") != "done":
raise HTTPException(status_code=404, detail="Result not ready")
img = job["output"]
url = (f"{COMFY}/view?filename={img['filename']}"
f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
async with httpx.AsyncClient(timeout=30) as c:
r = await c.get(url)
return StreamingResponse(
io.BytesIO(r.content),
media_type="image/png",
headers={"Content-Disposition": f"attachment; filename=dreamweaver_{job_id[:8]}.png"}
)
@app.post("/dream-weaver/sync")
async def dream_weaver_sync(
image: UploadFile = File(...),
keywords: str = Form(default=""),
room_type: str = Form(default="living_room"),
additional_notes: str = Form(default=""),
custom_positive: str = Form(default=""),
custom_negative: str = Form(default=""),
):
"""
Blocking version — waits up to 120s and returns image bytes directly.
Use for testing. Prefer async /dream-weaver for production.
"""
data = await image.read()
filename = f"sync_{uuid.uuid4().hex[:8]}_{image.filename or 'room.jpg'}"
comfy_name = await upload_to_comfy(data, filename)
if custom_positive:
from dataclasses import dataclass
@dataclass
class _P:
style_name = "custom"
positive_prompt = custom_positive
negative_prompt = custom_negative or "(worst quality, low quality), blurry, structural changes"
cfg = 7.5; denoise = 0.72; steps = 30
reasoning = ""; source = "direct"
expanded = _P()
elif keywords:
kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
expanded = (expand_prompt(kw_list, room_type, additional_notes)
if LLM_AVAILABLE else expand_prompt_simple(kw_list, room_type))
else:
raise HTTPException(status_code=400, detail="Provide keywords or custom_positive")
wf = build_workflow(comfy_name, expanded)
prompt_id = await queue_prompt(wf)
img, err = await poll_result(prompt_id, timeout=120)
if err:
raise HTTPException(status_code=500, detail=str(err))
url = (f"{COMFY}/view?filename={img['filename']}"
f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
async with httpx.AsyncClient(timeout=30) as c:
r = await c.get(url)
return StreamingResponse(io.BytesIO(r.content), media_type="image/png",
headers={"X-Style": expanded.style_name,
"X-Prompt-Source": expanded.source})
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8082")), log_level="info")
#!/usr/bin/env python3
"""
Dream Weaver API Gateway v2 — Dynamic Keyword → Local LLM → ComfyUI Pipeline
========================================================================
Port: 8080 (public-facing)
ComfyUI: localhost:8188 (internal)
NEW IN v2:
- POST /dream-weaver now accepts keywords[] + room_type for LLM-based prompt generation
- POST /dream-weaver/expand — expand keywords to prompt WITHOUT generating (preview)
- GET /room-types — list available room types
- Uses local Ollama model (qwen3.5:27b) for prompt expansion (no cloud API dependencies)
Environment variables:
OLLAMA_URL — Ollama server (default: http://localhost:11434)
OLLAMA_MODEL — Model name (default: qwen3.5:27b)
"""
import asyncio, json, time, uuid, io, sys, os, logging
from pathlib import Path
from typing import Optional, List
import httpx
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, BackgroundTasks
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
# Add scripts dir to path so we can import prompt_expander
SCRIPTS_DIR = Path(__file__).parent / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))
try:
from prompt_expander import expand_prompt, expand_prompt_simple, ROOM_CONTEXTS, ExpandedPrompt
LLM_AVAILABLE = True
except ImportError:
LLM_AVAILABLE = False
logging.warning("prompt_expander not found — LLM expansion disabled")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("DreamWeaverGateway")
COMFY = "http://127.0.0.1:8188"
COMFY_ROOT = "/opt/dlami/nvme/ComfyUI"
app = FastAPI(
title="Dream Weaver API v2",
version="2.0.0",
description="Dynamic keyword-to-interior-design generation powered by LLM + ComfyUI"
)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
# In-memory job store (swap for Redis in production)
jobs: dict = {}
# ─── Models ──────────────────────────────────────────────────────────────────
class ExpandRequest(BaseModel):
keywords: List[str]
room_type: str = "living_room"
additional_notes: str = ""
class ExpandResponse(BaseModel):
style_name: str
positive_prompt: str
negative_prompt: str
cfg: float
denoise: float
steps: int
reasoning: str
source: str
# ─── ComfyUI helpers ──────────────────────────────────────────────────────────
async def upload_to_comfy(data: bytes, filename: str) -> str:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(f"{COMFY}/upload/image",
files={"image": (filename, data, "image/jpeg")},
data={"overwrite": "true"})
r.raise_for_status()
return r.json()["name"]
def build_workflow(img_name: str, expanded: "ExpandedPrompt") -> dict:
"""Build ComfyUI API workflow from an ExpandedPrompt result."""
return {
"1": {"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "realvisxlV50_v50LightningBakedvae.safetensors"}},
"2": {"class_type": "LoadImage",
"inputs": {"image": img_name, "upload": "image"}},
"3": {"class_type": "CLIPTextEncode", # Positive prompt
"inputs": {"text": expanded.positive_prompt, "clip": ["1", 1]}},
"4": {"class_type": "CLIPTextEncode", # Negative prompt
"inputs": {"text": expanded.negative_prompt, "clip": ["1", 1]}},
"5": {"class_type": "VAEEncode",
"inputs": {"pixels": ["2", 0], "vae": ["1", 2]}},
"6": {"class_type": "KSampler",
"inputs": {"model": ["1", 0],
"positive": ["3", 0],
"negative": ["4", 0],
"latent_image": ["5", 0],
"seed": int(time.time()) % 999983,
"steps": expanded.steps,
"cfg": expanded.cfg,
"sampler_name": "dpmpp_2m",
"scheduler": "karras",
"denoise": expanded.denoise}},
"7": {"class_type": "VAEDecode",
"inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
"8": {"class_type": "SaveImage",
"inputs": {"images": ["7", 0],
"filename_prefix": f"dw_{expanded.style_name.replace(' ', '_')[:30]}"}},
}
async def queue_prompt(workflow: dict) -> str:
async with httpx.AsyncClient(timeout=30) as client:
r = await client.post(f"{COMFY}/prompt",
json={"prompt": workflow, "client_id": str(uuid.uuid4())})
r.raise_for_status()
return r.json()["prompt_id"]
async def poll_result(prompt_id: str, timeout: int = 300):
start = time.time()
async with httpx.AsyncClient(timeout=10) as client:
while time.time() - start < timeout:
r = await client.get(f"{COMFY}/history/{prompt_id}")
if r.status_code == 200:
h = r.json().get(prompt_id, {})
if h.get("status", {}).get("status_str") == "error":
return None, h.get("status", {}).get("messages", ["unknown"])
imgs = [img for nd in h.get("outputs", {}).values()
for img in nd.get("images", [])]
if imgs:
return imgs[0], None
await asyncio.sleep(2)
return None, "timeout"
async def background_poll(job_id: str, prompt_id: str):
img, err = await poll_result(prompt_id)
if img:
jobs[job_id].update({"status": "done", "output": img, "completed": time.time()})
else:
jobs[job_id].update({"status": "error", "error": str(err)})
# ─── Endpoints ───────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
comfy_ok = False
try:
async with httpx.AsyncClient(timeout=5) as c:
r = await c.get(f"{COMFY}/system_stats")
comfy_ok = r.status_code == 200
except Exception:
pass
return {
"status": "ok",
"comfyui": comfy_ok,
"gpu": "4x NVIDIA L4 (96GB VRAM)",
"model": "RealVisXL V5.0 Lightning",
"llm_expansion": LLM_AVAILABLE,
"version": "2.0.0"
}
@app.get("/room-types")
async def room_types():
"""List all supported room types with their context."""
if not LLM_AVAILABLE:
return {"room_types": ["bedroom", "living_room", "bathroom", "kitchen",
"dining_room", "home_office", "hallway", "balcony"]}
return {
"room_types": {
k: {
"description": v["description"],
"key_elements": v["key_elements"]
}
for k, v in ROOM_CONTEXTS.items()
}
}
@app.post("/dream-weaver/expand", response_model=ExpandResponse)
async def expand_endpoint(req: ExpandRequest):
"""
Preview the LLM-generated prompt WITHOUT submitting to ComfyUI.
Use this to let the user review/edit the prompt before generating.
Request body:
{
"keywords": ["blue marble", "gold veins", "renaissance", "sharp contours"],
"room_type": "bedroom",
"additional_notes": "luxury hotel feel"
}
"""
if not req.keywords:
raise HTTPException(status_code=400, detail="keywords list cannot be empty")
try:
if LLM_AVAILABLE:
result = await asyncio.to_thread(
expand_prompt,
keywords=req.keywords,
room_type=req.room_type,
additional_notes=req.additional_notes
)
else:
result = expand_prompt_simple(req.keywords, req.room_type)
except Exception as e:
logger.error(f"Prompt expansion failed: {e}")
raise HTTPException(status_code=500, detail=f"LLM expansion failed: {str(e)}")
return ExpandResponse(
style_name=result.style_name,
positive_prompt=result.positive_prompt,
negative_prompt=result.negative_prompt,
cfg=result.cfg,
denoise=result.denoise,
steps=result.steps,
reasoning=result.reasoning,
source=result.source
)
@app.post("/dream-weaver")
async def dream_weaver(
image: UploadFile = File(...),
# ── Dynamic keyword mode (new) ──
keywords: str = Form(default=""), # comma-separated: "blue marble, gold, renaissance"
room_type: str = Form(default="living_room"),
additional_notes: str = Form(default=""),
# ── Optional overrides ──
custom_positive: str = Form(default=""), # skip LLM, use this prompt directly
custom_negative: str = Form(default=""),
denoise: float = Form(default=0.0), # 0.0 = use LLM recommendation
cfg_scale: float = Form(default=0.0), # 0.0 = use LLM recommendation
):
"""
Submit a room photo for AI redesign using dynamic keyword → LLM → ComfyUI pipeline.
Two modes:
1. KEYWORD MODE (recommended): Provide keywords + room_type, LLM generates prompt
2. DIRECT MODE: Provide custom_positive + custom_negative to bypass LLM
Returns job_id for async polling.
"""
job_id = str(uuid.uuid4())
jobs[job_id] = {"status": "uploading", "created": time.time()}
try:
# Upload image to ComfyUI
data = await image.read()
filename = f"dw_{job_id[:8]}_{image.filename or 'room.jpg'}"
comfy_name = await upload_to_comfy(data, filename)
jobs[job_id]["status"] = "expanding_prompt"
# ── Determine prompt ──────────────────────────────────────────────
if custom_positive:
# Direct mode — user provided prompts explicitly
from dataclasses import dataclass
@dataclass
class DirectPrompt:
style_name: str = "custom"
positive_prompt: str = custom_positive
negative_prompt: str = custom_negative or (
"(worst quality, low quality, illustration, 3d render, painting, cartoon, sketch), "
"blurry, distorted, deformed, extra windows, unrealistic lighting, structural changes"
)
cfg: float = cfg_scale or 7.5
denoise: float = denoise or 0.72
steps: int = 30
reasoning: str = "Direct user input"
source: str = "direct"
expanded = DirectPrompt()
elif keywords:
# Keyword mode — expand via LLM
kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
if LLM_AVAILABLE:
expanded = await asyncio.to_thread(
expand_prompt,
keywords=kw_list,
room_type=room_type,
additional_notes=additional_notes
)
else:
expanded = expand_prompt_simple(kw_list, room_type)
# Apply manual overrides if provided
if denoise > 0:
expanded.denoise = denoise
if cfg_scale > 0:
expanded.cfg = cfg_scale
else:
raise HTTPException(status_code=400,
detail="Provide either 'keywords' or 'custom_positive'")
jobs[job_id].update({
"status": "queued",
"style": expanded.style_name,
"prompt_source": expanded.source,
"positive_prompt": expanded.positive_prompt,
"negative_prompt": expanded.negative_prompt,
"room_type": room_type,
})
# Submit workflow
wf = build_workflow(comfy_name, expanded)
prompt_id = await queue_prompt(wf)
jobs[job_id].update({"status": "processing", "prompt_id": prompt_id})
# Start background polling
asyncio.create_task(background_poll(job_id, prompt_id))
return {
"job_id": job_id,
"status": "processing",
"style": expanded.style_name,
"prompt_preview": expanded.positive_prompt[:120] + "...",
"reasoning": expanded.reasoning,
"poll_url": f"/dream-weaver/status/{job_id}",
"result_url": f"/dream-weaver/result/{job_id}"
}
except HTTPException:
raise
except Exception as e:
jobs[job_id] = {"status": "error", "error": str(e)}
logger.error(f"Generation failed: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/dream-weaver/status/{job_id}")
async def status(job_id: str):
job = jobs.get(job_id)
if not job:
raise HTTPException(status_code=404, detail="Job not found")
result = {k: v for k, v in job.items() if k != "output"}
result["ready"] = job.get("status") == "done"
if result["ready"]:
result["result_url"] = f"/dream-weaver/result/{job_id}"
return result
@app.get("/dream-weaver/result/{job_id}")
async def result(job_id: str):
job = jobs.get(job_id)
if not job or job.get("status") != "done":
raise HTTPException(status_code=404, detail="Result not ready")
img = job["output"]
url = (f"{COMFY}/view?filename={img['filename']}"
f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
async with httpx.AsyncClient(timeout=30) as c:
r = await c.get(url)
return StreamingResponse(
io.BytesIO(r.content),
media_type="image/png",
headers={"Content-Disposition": f"attachment; filename=dreamweaver_{job_id[:8]}.png"}
)
@app.post("/dream-weaver/sync")
async def dream_weaver_sync(
image: UploadFile = File(...),
keywords: str = Form(default=""),
room_type: str = Form(default="living_room"),
additional_notes: str = Form(default=""),
custom_positive: str = Form(default=""),
custom_negative: str = Form(default=""),
):
"""
Blocking version — waits up to 120s and returns image bytes directly.
Use for testing. Prefer async /dream-weaver for production.
"""
data = await image.read()
filename = f"sync_{uuid.uuid4().hex[:8]}_{image.filename or 'room.jpg'}"
comfy_name = await upload_to_comfy(data, filename)
if custom_positive:
from dataclasses import dataclass
@dataclass
class _P:
style_name = "custom"
positive_prompt = custom_positive
negative_prompt = custom_negative or "(worst quality, low quality), blurry, structural changes"
cfg = 7.5; denoise = 0.72; steps = 30
reasoning = ""; source = "direct"
expanded = _P()
elif keywords:
kw_list = [k.strip() for k in keywords.split(",") if k.strip()]
expanded = (expand_prompt(kw_list, room_type, additional_notes)
if LLM_AVAILABLE else expand_prompt_simple(kw_list, room_type))
else:
raise HTTPException(status_code=400, detail="Provide keywords or custom_positive")
wf = build_workflow(comfy_name, expanded)
prompt_id = await queue_prompt(wf)
img, err = await poll_result(prompt_id, timeout=120)
if err:
raise HTTPException(status_code=500, detail=str(err))
url = (f"{COMFY}/view?filename={img['filename']}"
f"&subfolder={img.get('subfolder','')}&type={img.get('type','output')}")
async with httpx.AsyncClient(timeout=30) as c:
r = await c.get(url)
return StreamingResponse(io.BytesIO(r.content), media_type="image/png",
headers={"X-Style": expanded.style_name,
"X-Prompt-Source": expanded.source})
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8082")), log_level="info")

776
comfy_engine/scripts/mask_preprocessor.py
View File

@@ -1,388 +1,388 @@
#!/usr/bin/env python3
"""
Dream Weaver Mask Preprocessor
==============================
Utility script for preprocessing and caching segmentation masks.
Enables offline mask generation to speed up production workflows.
Target Hardware: Dual NVIDIA RTX PRO 6000 Blackwell
Author: Project Velocity Team
Version: 1.0.0
"""
import os
import sys
import json
import hashlib
import argparse
import logging
from pathlib import Path
from typing import List, Optional, Tuple, Dict
from dataclasses import dataclass
import numpy as np
from PIL import Image
import cv2
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('MaskPreprocessor')
@dataclass
class MaskConfig:
"""Configuration for mask generation."""
grow_pixels: int = 3
feather_pixels: int = 5
threshold: float = 0.3
target_classes: List[str] = None
def __post_init__(self):
if self.target_classes is None:
self.target_classes = ["wall", "floor", "ceiling"]
class MaskPreprocessor:
"""Preprocesses and caches segmentation masks for Dream Weaver."""
def __init__(self, cache_dir: str = "Project_Velocity/comfy_engine/cache/masks/"):
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
self.config = MaskConfig()
logger.info(f"MaskPreprocessor initialized. Cache directory: {self.cache_dir}")
def _get_image_hash(self, image_path: str) -> str:
"""Generate MD5 hash of image for caching."""
hasher = hashlib.md5()
with open(image_path, 'rb') as f:
hasher.update(f.read())
return hasher.hexdigest()
def _get_cache_path(self, image_path: str, suffix: str = "") -> Path:
"""Generate cache file path for an image."""
image_hash = self._get_image_hash(image_path)
filename = f"{image_hash}{suffix}.png"
return self.cache_dir / filename
def is_cached(self, image_path: str, suffix: str = "") -> bool:
"""Check if a mask is already cached."""
cache_path = self._get_cache_path(image_path, suffix)
return cache_path.exists()
def load_from_cache(self, image_path: str, suffix: str = "") -> Optional[np.ndarray]:
"""Load mask from cache if available."""
cache_path = self._get_cache_path(image_path, suffix)
if cache_path.exists():
logger.info(f"Loading cached mask from {cache_path}")
mask = cv2.imread(str(cache_path), cv2.IMREAD_GRAYSCALE)
return mask
return None
def save_to_cache(self, image_path: str, mask: np.ndarray, suffix: str = "") -> str:
"""Save mask to cache."""
cache_path = self._get_cache_path(image_path, suffix)
cv2.imwrite(str(cache_path), mask)
logger.info(f"Saved mask to cache: {cache_path}")
return str(cache_path)
def create_structural_mask(self, image_path: str, mask_data: np.ndarray) -> np.ndarray:
"""
Create a structural preservation mask from segmentation data.
This mask identifies walls, floors, ceilings that must be preserved.
"""
# Ensure binary mask
if len(mask_data.shape) == 3:
mask_data = cv2.cvtColor(mask_data, cv2.COLOR_BGR2GRAY)
_, binary_mask = cv2.threshold(
mask_data,
int(255 * self.config.threshold),
255,
cv2.THRESH_BINARY
)
return binary_mask.astype(np.uint8)
def grow_mask(self, mask: np.ndarray, pixels: int = None) -> np.ndarray:
"""
Grow (dilate) the mask by specified pixels.
This prevents edge bleeding by expanding the mask slightly.
"""
if pixels is None:
pixels = self.config.grow_pixels
kernel = np.ones((pixels * 2 + 1, pixels * 2 + 1), np.uint8)
grown_mask = cv2.dilate(mask, kernel, iterations=1)
return grown_mask
def feather_mask(self, mask: np.ndarray, pixels: int = None) -> np.ndarray:
"""
Apply Gaussian blur to feather mask edges.
Creates smooth transitions at boundaries.
"""
if pixels is None:
pixels = self.config.feather_pixels
# Ensure odd kernel size
kernel_size = pixels * 2 + 1
feathered = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
return feathered
def invert_mask(self, mask: np.ndarray) -> np.ndarray:
"""Invert mask (structural -> stylable or vice versa)."""
return cv2.bitwise_not(mask)
def combine_masks(self, masks: List[np.ndarray], operation: str = "union") -> np.ndarray:
"""
Combine multiple masks.
operation: 'union' (OR), 'intersection' (AND), 'difference'
"""
if not masks:
return None
result = masks[0].copy()
for mask in masks[1:]:
if operation == "union":
result = cv2.bitwise_or(result, mask)
elif operation == "intersection":
result = cv2.bitwise_and(result, mask)
elif operation == "difference":
result = cv2.bitwise_and(result, cv2.bitwise_not(mask))
return result
def create_multi_region_mask(
self,
image_path: str,
regions: Dict[str, np.ndarray]
) -> Dict[str, np.ndarray]:
"""
Create masks for multiple regions (walls, floor, ceiling, etc.)
Returns dictionary of processed masks.
"""
processed_masks = {}
for region_name, mask_data in regions.items():
logger.info(f"Processing mask for region: {region_name}")
# Create base mask
base_mask = self.create_structural_mask(image_path, mask_data)
# Grow mask to prevent edge bleeding
grown_mask = self.grow_mask(base_mask)
# Feather edges
feathered_mask = self.feather_mask(grown_mask)
# Cache the processed mask
cache_path = self.save_to_cache(
image_path,
feathered_mask,
suffix=f"_{region_name}"
)
processed_masks[region_name] = {
"mask": feathered_mask,
"cache_path": cache_path
}
# Create combined structural mask
all_structural = [m["mask"] for m in processed_masks.values()]
combined_structural = self.combine_masks(all_structural, operation="union")
# Create stylable mask (inverse of structural)
stylable_mask = self.invert_mask(combined_structural)
# Save combined masks
structural_cache = self.save_to_cache(
image_path,
combined_structural,
suffix="_structural"
)
stylable_cache = self.save_to_cache(
image_path,
stylable_mask,
suffix="_stylable"
)
processed_masks["combined_structural"] = {
"mask": combined_structural,
"cache_path": structural_cache
}
processed_masks["stylable"] = {
"mask": stylable_mask,
"cache_path": stylable_cache
}
return processed_masks
def preprocess_image(self, image_path: str) -> Dict:
"""
Complete preprocessing pipeline for a single image.
Returns metadata about generated masks.
"""
logger.info(f"Preprocessing image: {image_path}")
# Check if already cached
if self.is_cached(image_path, "_structural"):
logger.info(f"Image already preprocessed: {image_path}")
return {
"image_path": image_path,
"cached": True,
"masks": {
"structural": str(self._get_cache_path(image_path, "_structural")),
"stylable": str(self._get_cache_path(image_path, "_stylable"))
}
}
# Load image for reference
img = cv2.imread(image_path)
if img is None:
raise ValueError(f"Could not load image: {image_path}")
height, width = img.shape[:2]
# Create placeholder masks (in production, these would come from SAM)
# This simulates wall, floor, ceiling segmentation
regions = {}
# Wall mask (upper portion)
wall_mask = np.zeros((height, width), dtype=np.uint8)
wall_mask[0:int(height*0.6), :] = 255
regions["wall"] = wall_mask
# Floor mask (lower portion)
floor_mask = np.zeros((height, width), dtype=np.uint8)
floor_mask[int(height*0.6):, :] = 255
regions["floor"] = floor_mask
# Ceiling mask (top portion)
ceiling_mask = np.zeros((height, width), dtype=np.uint8)
ceiling_mask[0:int(height*0.15), :] = 255
regions["ceiling"] = ceiling_mask
# Process all regions
processed = self.create_multi_region_mask(image_path, regions)
return {
"image_path": image_path,
"cached": False,
"dimensions": (width, height),
"masks": {
name: data["cache_path"]
for name, data in processed.items()
}
}
def batch_preprocess(self, directory: str, pattern: str = "*.jpg") -> List[Dict]:
"""Preprocess all images in a directory."""
input_dir = Path(directory)
image_files = list(input_dir.glob(pattern))
image_files.extend(list(input_dir.glob("*.png")))
results = []
for img_file in image_files:
try:
result = self.preprocess_image(str(img_file))
results.append(result)
except Exception as e:
logger.error(f"Failed to preprocess {img_file}: {e}")
return results
def clear_cache(self):
"""Clear all cached masks."""
for cache_file in self.cache_dir.glob("*.png"):
cache_file.unlink()
logger.info("Cache cleared")
def get_cache_stats(self) -> Dict:
"""Get cache statistics."""
cache_files = list(self.cache_dir.glob("*.png"))
total_size = sum(f.stat().st_size for f in cache_files)
return {
"cached_files": len(cache_files),
"total_size_mb": total_size / (1024 * 1024),
"cache_directory": str(self.cache_dir)
}
def main():
"""Main entry point for command-line usage."""
parser = argparse.ArgumentParser(
description="Dream Weaver Mask Preprocessor"
)
parser.add_argument(
"--image",
type=str,
help="Single image to preprocess"
)
parser.add_argument(
"--directory",
type=str,
help="Directory of images to preprocess"
)
parser.add_argument(
"--cache-dir",
type=str,
default="Project_Velocity/comfy_engine/cache/masks/",
help="Cache directory for masks"
)
parser.add_argument(
"--grow",
type=int,
default=3,
help="Pixels to grow mask (dilation)"
)
parser.add_argument(
"--feather",
type=int,
default=5,
help="Pixels to feather mask edges"
)
parser.add_argument(
"--clear-cache",
action="store_true",
help="Clear all cached masks"
)
parser.add_argument(
"--stats",
action="store_true",
help="Show cache statistics"
)
args = parser.parse_args()
# Initialize preprocessor
preprocessor = MaskPreprocessor(cache_dir=args.cache_dir)
preprocessor.config.grow_pixels = args.grow
preprocessor.config.feather_pixels = args.feather
if args.clear_cache:
preprocessor.clear_cache()
return
if args.stats:
stats = preprocessor.get_cache_stats()
print(json.dumps(stats, indent=2))
return
if args.image:
result = preprocessor.preprocess_image(args.image)
print(json.dumps(result, indent=2))
elif args.directory:
results = preprocessor.batch_preprocess(args.directory)
print(json.dumps(results, indent=2))
print(f"\nProcessed {len(results)} images")
else:
print("No action specified. Use --help for usage information.")
if __name__ == "__main__":
main()
#!/usr/bin/env python3
"""
Dream Weaver Mask Preprocessor
==============================
Utility script for preprocessing and caching segmentation masks.
Enables offline mask generation to speed up production workflows.
Target Hardware: Dual NVIDIA RTX PRO 6000 Blackwell
Author: Project Velocity Team
Version: 1.0.0
"""
import os
import sys
import json
import hashlib
import argparse
import logging
from pathlib import Path
from typing import List, Optional, Tuple, Dict
from dataclasses import dataclass
import numpy as np
from PIL import Image
import cv2
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('MaskPreprocessor')
@dataclass
class MaskConfig:
"""Configuration for mask generation."""
grow_pixels: int = 3
feather_pixels: int = 5
threshold: float = 0.3
target_classes: List[str] = None
def __post_init__(self):
if self.target_classes is None:
self.target_classes = ["wall", "floor", "ceiling"]
class MaskPreprocessor:
"""Preprocesses and caches segmentation masks for Dream Weaver."""
def __init__(self, cache_dir: str = "Project_Velocity/comfy_engine/cache/masks/"):
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
self.config = MaskConfig()
logger.info(f"MaskPreprocessor initialized. Cache directory: {self.cache_dir}")
def _get_image_hash(self, image_path: str) -> str:
"""Generate MD5 hash of image for caching."""
hasher = hashlib.md5()
with open(image_path, 'rb') as f:
hasher.update(f.read())
return hasher.hexdigest()
def _get_cache_path(self, image_path: str, suffix: str = "") -> Path:
"""Generate cache file path for an image."""
image_hash = self._get_image_hash(image_path)
filename = f"{image_hash}{suffix}.png"
return self.cache_dir / filename
def is_cached(self, image_path: str, suffix: str = "") -> bool:
"""Check if a mask is already cached."""
cache_path = self._get_cache_path(image_path, suffix)
return cache_path.exists()
def load_from_cache(self, image_path: str, suffix: str = "") -> Optional[np.ndarray]:
"""Load mask from cache if available."""
cache_path = self._get_cache_path(image_path, suffix)
if cache_path.exists():
logger.info(f"Loading cached mask from {cache_path}")
mask = cv2.imread(str(cache_path), cv2.IMREAD_GRAYSCALE)
return mask
return None
def save_to_cache(self, image_path: str, mask: np.ndarray, suffix: str = "") -> str:
"""Save mask to cache."""
cache_path = self._get_cache_path(image_path, suffix)
cv2.imwrite(str(cache_path), mask)
logger.info(f"Saved mask to cache: {cache_path}")
return str(cache_path)
def create_structural_mask(self, image_path: str, mask_data: np.ndarray) -> np.ndarray:
"""
Create a structural preservation mask from segmentation data.
This mask identifies walls, floors, ceilings that must be preserved.
"""
# Ensure binary mask
if len(mask_data.shape) == 3:
mask_data = cv2.cvtColor(mask_data, cv2.COLOR_BGR2GRAY)
_, binary_mask = cv2.threshold(
mask_data,
int(255 * self.config.threshold),
255,
cv2.THRESH_BINARY
)
return binary_mask.astype(np.uint8)
def grow_mask(self, mask: np.ndarray, pixels: int = None) -> np.ndarray:
"""
Grow (dilate) the mask by specified pixels.
This prevents edge bleeding by expanding the mask slightly.
"""
if pixels is None:
pixels = self.config.grow_pixels
kernel = np.ones((pixels * 2 + 1, pixels * 2 + 1), np.uint8)
grown_mask = cv2.dilate(mask, kernel, iterations=1)
return grown_mask
def feather_mask(self, mask: np.ndarray, pixels: int = None) -> np.ndarray:
"""
Apply Gaussian blur to feather mask edges.
Creates smooth transitions at boundaries.
"""
if pixels is None:
pixels = self.config.feather_pixels
# Ensure odd kernel size
kernel_size = pixels * 2 + 1
feathered = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
return feathered
def invert_mask(self, mask: np.ndarray) -> np.ndarray:
"""Invert mask (structural -> stylable or vice versa)."""
return cv2.bitwise_not(mask)
def combine_masks(self, masks: List[np.ndarray], operation: str = "union") -> np.ndarray:
"""
Combine multiple masks.
operation: 'union' (OR), 'intersection' (AND), 'difference'
"""
if not masks:
return None
result = masks[0].copy()
for mask in masks[1:]:
if operation == "union":
result = cv2.bitwise_or(result, mask)
elif operation == "intersection":
result = cv2.bitwise_and(result, mask)
elif operation == "difference":
result = cv2.bitwise_and(result, cv2.bitwise_not(mask))
return result
def create_multi_region_mask(
self,
image_path: str,
regions: Dict[str, np.ndarray]
) -> Dict[str, np.ndarray]:
"""
Create masks for multiple regions (walls, floor, ceiling, etc.)
Returns dictionary of processed masks.
"""
processed_masks = {}
for region_name, mask_data in regions.items():
logger.info(f"Processing mask for region: {region_name}")
# Create base mask
base_mask = self.create_structural_mask(image_path, mask_data)
# Grow mask to prevent edge bleeding
grown_mask = self.grow_mask(base_mask)
# Feather edges
feathered_mask = self.feather_mask(grown_mask)
# Cache the processed mask
cache_path = self.save_to_cache(
image_path,
feathered_mask,
suffix=f"_{region_name}"
)
processed_masks[region_name] = {
"mask": feathered_mask,
"cache_path": cache_path
}
# Create combined structural mask
all_structural = [m["mask"] for m in processed_masks.values()]
combined_structural = self.combine_masks(all_structural, operation="union")
# Create stylable mask (inverse of structural)
stylable_mask = self.invert_mask(combined_structural)
# Save combined masks
structural_cache = self.save_to_cache(
image_path,
combined_structural,
suffix="_structural"
)
stylable_cache = self.save_to_cache(
image_path,
stylable_mask,
suffix="_stylable"
)
processed_masks["combined_structural"] = {
"mask": combined_structural,
"cache_path": structural_cache
}
processed_masks["stylable"] = {
"mask": stylable_mask,
"cache_path": stylable_cache
}
return processed_masks
def preprocess_image(self, image_path: str) -> Dict:
"""
Complete preprocessing pipeline for a single image.
Returns metadata about generated masks.
"""
logger.info(f"Preprocessing image: {image_path}")
# Check if already cached
if self.is_cached(image_path, "_structural"):
logger.info(f"Image already preprocessed: {image_path}")
return {
"image_path": image_path,
"cached": True,
"masks": {
"structural": str(self._get_cache_path(image_path, "_structural")),
"stylable": str(self._get_cache_path(image_path, "_stylable"))
}
}
# Load image for reference
img = cv2.imread(image_path)
if img is None:
raise ValueError(f"Could not load image: {image_path}")
height, width = img.shape[:2]
# Create placeholder masks (in production, these would come from SAM)
# This simulates wall, floor, ceiling segmentation
regions = {}
# Wall mask (upper portion)
wall_mask = np.zeros((height, width), dtype=np.uint8)
wall_mask[0:int(height*0.6), :] = 255
regions["wall"] = wall_mask
# Floor mask (lower portion)
floor_mask = np.zeros((height, width), dtype=np.uint8)
floor_mask[int(height*0.6):, :] = 255
regions["floor"] = floor_mask
# Ceiling mask (top portion)
ceiling_mask = np.zeros((height, width), dtype=np.uint8)
ceiling_mask[0:int(height*0.15), :] = 255
regions["ceiling"] = ceiling_mask
# Process all regions
processed = self.create_multi_region_mask(image_path, regions)
return {
"image_path": image_path,
"cached": False,
"dimensions": (width, height),
"masks": {
name: data["cache_path"]
for name, data in processed.items()
}
}
def batch_preprocess(self, directory: str, pattern: str = "*.jpg") -> List[Dict]:
"""Preprocess all images in a directory."""
input_dir = Path(directory)
image_files = list(input_dir.glob(pattern))
image_files.extend(list(input_dir.glob("*.png")))
results = []
for img_file in image_files:
try:
result = self.preprocess_image(str(img_file))
results.append(result)
except Exception as e:
logger.error(f"Failed to preprocess {img_file}: {e}")
return results
def clear_cache(self):
"""Clear all cached masks."""
for cache_file in self.cache_dir.glob("*.png"):
cache_file.unlink()
logger.info("Cache cleared")
def get_cache_stats(self) -> Dict:
"""Get cache statistics."""
cache_files = list(self.cache_dir.glob("*.png"))
total_size = sum(f.stat().st_size for f in cache_files)
return {
"cached_files": len(cache_files),
"total_size_mb": total_size / (1024 * 1024),
"cache_directory": str(self.cache_dir)
}
def main():
"""Main entry point for command-line usage."""
parser = argparse.ArgumentParser(
description="Dream Weaver Mask Preprocessor"
)
parser.add_argument(
"--image",
type=str,
help="Single image to preprocess"
)
parser.add_argument(
"--directory",
type=str,
help="Directory of images to preprocess"
)
parser.add_argument(
"--cache-dir",
type=str,
default="Project_Velocity/comfy_engine/cache/masks/",
help="Cache directory for masks"
)
parser.add_argument(
"--grow",
type=int,
default=3,
help="Pixels to grow mask (dilation)"
)
parser.add_argument(
"--feather",
type=int,
default=5,
help="Pixels to feather mask edges"
)
parser.add_argument(
"--clear-cache",
action="store_true",
help="Clear all cached masks"
)
parser.add_argument(
"--stats",
action="store_true",
help="Show cache statistics"
)
args = parser.parse_args()
# Initialize preprocessor
preprocessor = MaskPreprocessor(cache_dir=args.cache_dir)
preprocessor.config.grow_pixels = args.grow
preprocessor.config.feather_pixels = args.feather
if args.clear_cache:
preprocessor.clear_cache()
return
if args.stats:
stats = preprocessor.get_cache_stats()
print(json.dumps(stats, indent=2))
return
if args.image:
result = preprocessor.preprocess_image(args.image)
print(json.dumps(result, indent=2))
elif args.directory:
results = preprocessor.batch_preprocess(args.directory)
print(json.dumps(results, indent=2))
print(f"\nProcessed {len(results)} images")
else:
print("No action specified. Use --help for usage information.")
if __name__ == "__main__":
main()

412
comfy_engine/scripts/prompt_expander.py
View File

@@ -1,206 +1,206 @@
#!/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)
#!/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)