Vito
@@ -2,11 +2,13 @@
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import server
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
# Import database operations
|
||||
try:
|
||||
@@ -243,6 +245,11 @@ class PromptManagerAPI:
|
||||
async def test_image_link_route(request):
|
||||
return await self.test_image_link(request)
|
||||
|
||||
# Scan endpoint
|
||||
@routes.post("/prompt_manager/scan")
|
||||
async def scan_images_route(request):
|
||||
return await self.scan_images(request)
|
||||
|
||||
print("[PromptManager] All routes registered with decorator pattern")
|
||||
|
||||
async def search_prompts(self, request):
|
||||
@@ -1617,3 +1624,265 @@ class PromptManagerAPI:
|
||||
'success': False,
|
||||
'error': f'Failed to restore database: {str(e)}'
|
||||
}, status=500)
|
||||
|
||||
async def scan_images(self, request):
|
||||
"""
|
||||
Scan ComfyUI output images for prompt metadata and add them to the database.
|
||||
Streams progress updates to the client.
|
||||
"""
|
||||
import json
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from aiohttp import web
|
||||
|
||||
async def stream_response():
|
||||
try:
|
||||
# Find ComfyUI output directory
|
||||
output_dir = self._find_comfyui_output_dir()
|
||||
if not output_dir:
|
||||
yield f"data: {json.dumps({'type': 'error', 'message': 'ComfyUI output directory not found'})}\n\n"
|
||||
return
|
||||
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Scanning for PNG files...', 'processed': 0, 'found': 0})}\n\n"
|
||||
|
||||
# Find all PNG files
|
||||
png_files = list(Path(output_dir).rglob("*.png"))
|
||||
total_files = len(png_files)
|
||||
|
||||
if total_files == 0:
|
||||
yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0})}\n\n"
|
||||
return
|
||||
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} PNG files to process...', 'processed': 0, 'found': 0})}\n\n"
|
||||
|
||||
processed_count = 0
|
||||
found_count = 0
|
||||
added_count = 0
|
||||
|
||||
for i, png_file in enumerate(png_files):
|
||||
try:
|
||||
# Extract metadata from PNG
|
||||
metadata = self._extract_comfyui_metadata(str(png_file))
|
||||
processed_count += 1
|
||||
|
||||
if metadata:
|
||||
print(f"[PromptManager] Found metadata in {os.path.basename(png_file)}: {list(metadata.keys())}")
|
||||
|
||||
# Parse ComfyUI prompt data
|
||||
parsed_data = self._parse_comfyui_prompt(metadata)
|
||||
print(f"[PromptManager] Parsed data keys: {list(parsed_data.keys())}, has prompt: {bool(parsed_data.get('prompt'))}, has parameters: {bool(parsed_data.get('parameters'))}")
|
||||
|
||||
# Check if we found any meaningful prompt data
|
||||
if parsed_data.get('prompt') or parsed_data.get('parameters'):
|
||||
found_count += 1
|
||||
|
||||
# Extract readable prompt text
|
||||
prompt_text = self._extract_readable_prompt(parsed_data)
|
||||
|
||||
# Debug: print what we found and its type
|
||||
if prompt_text:
|
||||
print(f"[PromptManager] Found prompt in {os.path.basename(png_file)} (type: {type(prompt_text)}): {str(prompt_text)[:100]}...")
|
||||
else:
|
||||
print(f"[PromptManager] No readable prompt found in {os.path.basename(png_file)}, parsed_data keys: {list(parsed_data.keys())}")
|
||||
|
||||
# Ensure prompt_text is a string
|
||||
if prompt_text and not isinstance(prompt_text, str):
|
||||
print(f"[PromptManager] Converting prompt_text from {type(prompt_text)} to string")
|
||||
prompt_text = str(prompt_text)
|
||||
|
||||
if prompt_text and prompt_text.strip():
|
||||
# Try to save to database (will skip duplicates)
|
||||
try:
|
||||
# Generate hash for duplicate detection
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
except ImportError:
|
||||
import sys
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
prompt_hash = generate_prompt_hash(prompt_text.strip())
|
||||
|
||||
# Check if prompt already exists
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
# Link image to existing prompt
|
||||
try:
|
||||
self.db.link_image_to_prompt(existing['id'], str(png_file))
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Failed to link image {png_file}: {e}")
|
||||
else:
|
||||
# Save new prompt
|
||||
prompt_id = self.db.save_prompt(
|
||||
text=prompt_text.strip(),
|
||||
category='scanned',
|
||||
tags=['auto-scanned'],
|
||||
notes=f'Auto-scanned from {os.path.basename(png_file)}',
|
||||
prompt_hash=prompt_hash
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
added_count += 1
|
||||
|
||||
# Link image to new prompt
|
||||
try:
|
||||
self.db.link_image_to_prompt(prompt_id, str(png_file))
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Failed to link image {png_file}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Failed to save prompt from {png_file}: {e}")
|
||||
|
||||
# Update progress every 10 files or so
|
||||
if i % 10 == 0 or i == total_files - 1:
|
||||
progress = int((i + 1) / total_files * 100)
|
||||
status = f"Processing file {i + 1}/{total_files}..."
|
||||
|
||||
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': status, 'processed': processed_count, 'found': found_count})}\n\n"
|
||||
|
||||
# Small delay to allow UI updates
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Error processing {png_file}: {e}")
|
||||
continue
|
||||
|
||||
# Send completion message
|
||||
yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count})}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Scan error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
|
||||
|
||||
# Return streaming response
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason='OK',
|
||||
headers={
|
||||
'Content-Type': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
'Connection': 'keep-alive'
|
||||
}
|
||||
)
|
||||
|
||||
await response.prepare(request)
|
||||
|
||||
async for chunk in stream_response():
|
||||
await response.write(chunk.encode('utf-8'))
|
||||
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
def _find_comfyui_output_dir(self):
|
||||
"""Find the ComfyUI output directory."""
|
||||
|
||||
# Common ComfyUI installation paths
|
||||
possible_paths = [
|
||||
os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "..", "..", "output"),
|
||||
os.path.join(os.path.expanduser("~"), "ComfyUI", "output"),
|
||||
os.path.join(os.getcwd(), "output"),
|
||||
os.path.join(os.getcwd(), "..", "output"),
|
||||
os.path.join(os.getcwd(), "..", "..", "output"),
|
||||
]
|
||||
|
||||
for path in possible_paths:
|
||||
abs_path = os.path.abspath(path)
|
||||
if os.path.exists(abs_path) and os.path.isdir(abs_path):
|
||||
print(f"[PromptManager] Found ComfyUI output directory: {abs_path}")
|
||||
return abs_path
|
||||
|
||||
print("[PromptManager] ComfyUI output directory not found")
|
||||
return None
|
||||
|
||||
def _extract_comfyui_metadata(self, image_path):
|
||||
"""Extract ComfyUI metadata from a PNG file."""
|
||||
try:
|
||||
with Image.open(image_path) as img:
|
||||
metadata = {}
|
||||
if hasattr(img, 'text'):
|
||||
for key, value in img.text.items():
|
||||
metadata[key] = value
|
||||
return metadata
|
||||
except Exception as e:
|
||||
print(f"[PromptManager] Error reading {image_path}: {e}")
|
||||
return {}
|
||||
|
||||
def _parse_comfyui_prompt(self, metadata):
|
||||
"""Parse ComfyUI prompt data from metadata."""
|
||||
result = {
|
||||
'prompt': None,
|
||||
'workflow': None,
|
||||
'parameters': {}
|
||||
}
|
||||
|
||||
# Check for direct prompt field
|
||||
if 'prompt' in metadata:
|
||||
try:
|
||||
prompt_data = json.loads(metadata['prompt'])
|
||||
result['prompt'] = prompt_data
|
||||
except json.JSONDecodeError:
|
||||
result['prompt'] = metadata['prompt']
|
||||
|
||||
# Check for workflow
|
||||
if 'workflow' in metadata:
|
||||
try:
|
||||
workflow_data = json.loads(metadata['workflow'])
|
||||
result['workflow'] = workflow_data
|
||||
except json.JSONDecodeError:
|
||||
result['workflow'] = metadata['workflow']
|
||||
|
||||
# Check for other common ComfyUI fields
|
||||
common_fields = ['parameters', 'positive', 'negative', 'steps', 'cfg', 'sampler', 'scheduler', 'seed']
|
||||
for field in common_fields:
|
||||
if field in metadata:
|
||||
try:
|
||||
result['parameters'][field] = json.loads(metadata[field])
|
||||
except json.JSONDecodeError:
|
||||
result['parameters'][field] = metadata[field]
|
||||
|
||||
return result
|
||||
|
||||
def _extract_readable_prompt(self, parsed_data):
|
||||
"""Extract human-readable prompt text from ComfyUI JSON structure."""
|
||||
# Helper function to convert any value to string safely
|
||||
def safe_to_string(value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
elif isinstance(value, list):
|
||||
# Join list elements with spaces
|
||||
return ' '.join(str(item) for item in value if item)
|
||||
elif value is not None:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
# Check if prompt is already a string
|
||||
if isinstance(parsed_data.get('prompt'), str):
|
||||
return parsed_data['prompt']
|
||||
|
||||
# Check if prompt is a simple value that can be converted
|
||||
if parsed_data.get('prompt') and not isinstance(parsed_data.get('prompt'), dict):
|
||||
return safe_to_string(parsed_data['prompt'])
|
||||
|
||||
prompt_data = parsed_data.get('prompt')
|
||||
if isinstance(prompt_data, dict):
|
||||
# ComfyUI stores prompts in various node types
|
||||
for node_id, node_data in prompt_data.items():
|
||||
if isinstance(node_data, dict):
|
||||
class_type = node_data.get('class_type', '')
|
||||
inputs = node_data.get('inputs', {})
|
||||
|
||||
# Common prompt node types in ComfyUI
|
||||
if class_type in ['CLIPTextEncode', 'CLIPTextEncodeSDXL', 'PromptText']:
|
||||
if 'text' in inputs:
|
||||
return safe_to_string(inputs['text'])
|
||||
elif 'text' in inputs:
|
||||
return safe_to_string(inputs['text'])
|
||||
|
||||
# Check parameters for positive prompt
|
||||
if parsed_data.get('parameters', {}).get('positive'):
|
||||
return safe_to_string(parsed_data['parameters']['positive'])
|
||||
|
||||
return None
|
||||
|
||||
+185
@@ -188,6 +188,10 @@
|
||||
class="px-4 py-2 bg-amber-600 hover:bg-amber-700 text-white text-sm font-medium rounded-lg transition-colors">
|
||||
📁 Restore DB
|
||||
</button>
|
||||
<button id="scanBtn"
|
||||
class="px-4 py-2 bg-pink-600 hover:bg-pink-700 text-white text-sm font-medium rounded-lg transition-colors">
|
||||
📸 Scan Images
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -536,6 +540,68 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Scan Modal -->
|
||||
<div id="scanModal" class="fixed inset-0 bg-black bg-opacity-50 hidden items-center justify-center z-50">
|
||||
<div class="bg-gray-800 rounded-xl p-6 max-w-lg w-full mx-4 border border-gray-700">
|
||||
<h3 class="text-xl font-semibold text-gray-100 mb-6">📸 Scan ComfyUI Output Images</h3>
|
||||
|
||||
<div class="space-y-4">
|
||||
<div class="bg-orange-900/30 border border-orange-600 rounded-lg p-4">
|
||||
<div class="flex items-start space-x-3">
|
||||
<span class="text-orange-400 text-xl">⚠️</span>
|
||||
<div>
|
||||
<h4 class="text-orange-300 font-medium mb-2">Important Information</h4>
|
||||
<ul class="text-orange-200 text-sm space-y-1">
|
||||
<li>• This will scan all PNG files in your ComfyUI output directory</li>
|
||||
<li>• Depending on the size of your output folder, this could take several minutes</li>
|
||||
<li>• This is an experimental feature - we recommend backing up your database first</li>
|
||||
<li>• Only images with ComfyUI metadata will be processed</li>
|
||||
<li>• Duplicate prompts will be automatically detected and skipped</li>
|
||||
</ul>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="bg-blue-900/30 border border-blue-600 rounded-lg p-4">
|
||||
<div class="flex items-center justify-between">
|
||||
<span class="text-blue-300 font-medium">Backup Database First (Recommended)</span>
|
||||
<button id="quickBackupBtn"
|
||||
class="px-4 py-2 bg-blue-600 hover:bg-blue-700 text-white text-sm font-medium rounded-lg transition-colors">
|
||||
💾 Backup Now
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="scanProgress" class="hidden space-y-3">
|
||||
<div class="bg-gray-700 rounded-lg p-4">
|
||||
<div class="flex items-center justify-between mb-2">
|
||||
<span class="text-gray-300 text-sm font-medium">Scanning Progress</span>
|
||||
<span id="scanStatusText" class="text-gray-400 text-sm">Preparing...</span>
|
||||
</div>
|
||||
<div class="w-full bg-gray-600 rounded-full h-2">
|
||||
<div id="scanProgressBar" class="bg-pink-600 h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
|
||||
</div>
|
||||
<div class="flex justify-between text-xs text-gray-400 mt-2">
|
||||
<span id="scanCount">0 files processed</span>
|
||||
<span id="scanFound">0 prompts found</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="flex justify-end space-x-3 mt-6">
|
||||
<button id="cancelScan"
|
||||
class="px-6 py-2 bg-gray-600 hover:bg-gray-700 text-white font-medium rounded-lg transition-colors">
|
||||
Cancel
|
||||
</button>
|
||||
<button id="startScan"
|
||||
class="px-6 py-2 bg-pink-600 hover:bg-pink-700 text-white font-medium rounded-lg transition-colors">
|
||||
Start Scan
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
class PromptAdmin {
|
||||
constructor() {
|
||||
@@ -579,6 +645,7 @@
|
||||
document.getElementById("maintenanceBtn").addEventListener("click", () => this.showMaintenanceModal());
|
||||
document.getElementById("backupBtn").addEventListener("click", () => this.backupDatabase());
|
||||
document.getElementById("restoreBtn").addEventListener("click", () => this.showRestoreModal());
|
||||
document.getElementById("scanBtn").addEventListener("click", () => this.showScanModal());
|
||||
|
||||
// Modals
|
||||
this.bindModalEvents();
|
||||
@@ -615,6 +682,11 @@
|
||||
document.getElementById("selectAllMaintenanceBtn").addEventListener("click", () => this.selectAllMaintenance());
|
||||
document.getElementById("clearAllMaintenanceBtn").addEventListener("click", () => this.clearAllMaintenance());
|
||||
|
||||
// Scan modal
|
||||
document.getElementById("startScan").addEventListener("click", () => this.startScan());
|
||||
document.getElementById("cancelScan").addEventListener("click", () => this.hideModal("scanModal"));
|
||||
document.getElementById("quickBackupBtn").addEventListener("click", () => this.quickBackup());
|
||||
|
||||
// Close modals on backdrop click
|
||||
document.querySelectorAll("[id$='Modal']").forEach((modal) => {
|
||||
modal.addEventListener("click", (e) => {
|
||||
@@ -1989,6 +2061,119 @@
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Scan functionality
|
||||
showScanModal() {
|
||||
// Reset progress display
|
||||
document.getElementById("scanProgress").classList.add("hidden");
|
||||
document.getElementById("startScan").disabled = false;
|
||||
document.getElementById("startScan").textContent = "Start Scan";
|
||||
this.showModal("scanModal");
|
||||
}
|
||||
|
||||
async quickBackup() {
|
||||
try {
|
||||
const response = await fetch("/prompt_manager/backup");
|
||||
if (response.ok) {
|
||||
const blob = await response.blob();
|
||||
const url = window.URL.createObjectURL(blob);
|
||||
const a = document.createElement("a");
|
||||
a.href = url;
|
||||
a.download = `prompts-backup-${new Date().toISOString().split("T")[0]}.db`;
|
||||
document.body.appendChild(a);
|
||||
a.click();
|
||||
window.URL.revokeObjectURL(url);
|
||||
document.body.removeChild(a);
|
||||
this.showNotification("Database backup created successfully!", "success");
|
||||
} else {
|
||||
throw new Error("Backup failed");
|
||||
}
|
||||
} catch (error) {
|
||||
this.showNotification("Failed to create backup", "error");
|
||||
}
|
||||
}
|
||||
|
||||
async startScan() {
|
||||
// Show progress section
|
||||
document.getElementById("scanProgress").classList.remove("hidden");
|
||||
document.getElementById("startScan").disabled = true;
|
||||
document.getElementById("startScan").textContent = "Scanning...";
|
||||
|
||||
// Reset progress
|
||||
this.updateScanProgress(0, "Initializing scan...", 0, 0);
|
||||
|
||||
try {
|
||||
const response = await fetch("/prompt_manager/scan", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({})
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}: ${response.statusText}`);
|
||||
}
|
||||
|
||||
// Handle streaming response
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
const chunk = decoder.decode(value);
|
||||
const lines = chunk.split('\n');
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.trim().startsWith('data: ')) {
|
||||
try {
|
||||
const data = JSON.parse(line.substring(6));
|
||||
if (data.type === 'progress') {
|
||||
this.updateScanProgress(
|
||||
data.progress,
|
||||
data.status,
|
||||
data.processed,
|
||||
data.found
|
||||
);
|
||||
} else if (data.type === 'complete') {
|
||||
this.completeScan(data.processed, data.found, data.added);
|
||||
return;
|
||||
} else if (data.type === 'error') {
|
||||
throw new Error(data.message);
|
||||
}
|
||||
} catch (e) {
|
||||
console.log('Non-JSON line:', line);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
this.showNotification(`Scan failed: ${error.message}`, "error");
|
||||
document.getElementById("startScan").disabled = false;
|
||||
document.getElementById("startScan").textContent = "Start Scan";
|
||||
}
|
||||
}
|
||||
|
||||
updateScanProgress(progress, status, processed, found) {
|
||||
document.getElementById("scanProgressBar").style.width = `${progress}%`;
|
||||
document.getElementById("scanStatusText").textContent = status;
|
||||
document.getElementById("scanCount").textContent = `${processed} files processed`;
|
||||
document.getElementById("scanFound").textContent = `${found} prompts found`;
|
||||
}
|
||||
|
||||
completeScan(processed, found, added) {
|
||||
this.updateScanProgress(100, "Scan completed!", processed, found);
|
||||
document.getElementById("startScan").disabled = false;
|
||||
document.getElementById("startScan").textContent = "Start New Scan";
|
||||
|
||||
this.showNotification(
|
||||
`Scan completed! Processed ${processed} files, found ${found} prompts, added ${added} new prompts to database.`,
|
||||
"success"
|
||||
);
|
||||
|
||||
// Refresh the data
|
||||
this.loadInitialData();
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize the admin interface
|
||||
|
||||
Reference in New Issue
Block a user