From 6a47d615bb8f28894a380ebdb994216440bca143 Mon Sep 17 00:00:00 2001 From: danieljanata Date: Thu, 15 Jan 2026 02:10:27 +0100 Subject: [PATCH] Delete database directory --- database/__init__.py | 6 - database/db_manager.py | 561 -------------------------------------- database/image_handler.py | 239 ---------------- 3 files changed, 806 deletions(-) delete mode 100644 database/__init__.py delete mode 100644 database/db_manager.py delete mode 100644 database/image_handler.py diff --git a/database/__init__.py b/database/__init__.py deleted file mode 100644 index 4003844..0000000 --- a/database/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Database management for ComfyUI Prompting System""" - -from .db_manager import PromptDatabase -from .image_handler import ImageHandler - -__all__ = ['PromptDatabase', 'ImageHandler'] diff --git a/database/db_manager.py b/database/db_manager.py deleted file mode 100644 index d67b236..0000000 --- a/database/db_manager.py +++ /dev/null @@ -1,561 +0,0 @@ -"""Main database manager for prompt storage and retrieval""" - -import json -import hashlib -from pathlib import Path -from typing import Optional, List, Dict, Any -from datetime import datetime, timedelta -from .image_handler import ImageHandler -import sys -import os - -# Add parent directory to path for imports -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from utils.path_utils import get_database_path, normalize_path - - -class PromptDatabase: - """Manages prompt database with JSON storage and base64 thumbnails""" - - VERSION = "1.0" - - def __init__(self, db_path: Optional[Path] = None): - """ - Initialize PromptDatabase. - - Args: - db_path: Optional custom path to database file - """ - self.db_path = db_path or get_database_path() - self.image_handler = ImageHandler() - self.data = self._load_or_create() - - def _load_or_create(self) -> Dict[str, Any]: - """Load existing database or create new one""" - if self.db_path.exists(): - try: - with open(self.db_path, 'r', encoding='utf-8') as f: - data = json.load(f) - print(f"[PromptSystem] Loaded database with {len(data.get('prompts', []))} prompts") - return data - except Exception as e: - print(f"[PromptSystem] Error loading database: {e}") - return self._create_empty_database() - else: - return self._create_empty_database() - - def _create_empty_database(self) -> Dict[str, Any]: - """Create empty database structure""" - return { - "version": self.VERSION, - "created": datetime.now().isoformat(), - "last_updated": datetime.now().isoformat(), - "settings": { - "auto_cleanup_enabled": True, - "auto_cleanup_days": 30, - "max_thumbnails": 3 - }, - "categories": [], - "tags": [], - "prompts": [] - } - - def save(self) -> bool: - """Save database to disk""" - try: - self.data['last_updated'] = datetime.now().isoformat() - - # Ensure directory exists - self.db_path.parent.mkdir(parents=True, exist_ok=True) - - # Write to temporary file first, then rename (atomic operation) - temp_path = self.db_path.with_suffix('.tmp') - with open(temp_path, 'w', encoding='utf-8') as f: - json.dump(self.data, f, indent=2, ensure_ascii=False) - - # Rename temp to actual file - temp_path.replace(self.db_path) - - print(f"[PromptSystem] Database saved to {self.db_path}") - return True - - except Exception as e: - print(f"[PromptSystem] Error saving database: {e}") - return False - - def _generate_hash(self, text: str) -> str: - """Generate SHA256 hash of text""" - return hashlib.sha256(text.encode('utf-8')).hexdigest() - - def _get_next_id(self) -> int: - """Get next available prompt ID""" - if not self.data['prompts']: - return 1 - return max(p['id'] for p in self.data['prompts']) + 1 - - def add_prompt( - self, - text: str, - category: str, - tags: Optional[List[str]] = None, - rating: int = 0, - notes: str = "", - image_path: Optional[str] = None - ) -> int: - """ - Add new prompt to database. - - Args: - text: Prompt text - category: Category name - tags: List of tags - rating: Rating 0-5 - notes: Optional notes - image_path: Optional path to image for thumbnail - - Returns: - Prompt ID - """ - tags = tags or [] - prompt_hash = self._generate_hash(text) - - # Check for duplicate - existing = self.find_by_hash(prompt_hash) - if existing: - print(f"[PromptSystem] Prompt already exists with ID {existing['id']}") - return existing['id'] - - # Create thumbnails if image provided - thumbnails = [] - if image_path: - thumb_data = self.image_handler.compress_and_encode(image_path) - if thumb_data: - thumbnails.append({ - 'data': thumb_data, - 'locked': False, - 'timestamp': datetime.now().isoformat(), - 'source_image': str(Path(image_path).name) - }) - - # Create new prompt - prompt_id = self._get_next_id() - new_prompt = { - "id": prompt_id, - "text": text, - "category": category, - "tags": tags, - "rating": rating, - "notes": notes, - "created_at": datetime.now().isoformat(), - "updated_at": datetime.now().isoformat(), - "used_count": 1, - "hash": prompt_hash, - "thumbnails": thumbnails, - "generation_history": [] - } - - self.data['prompts'].append(new_prompt) - - # Update categories and tags - if category and category not in self.data['categories']: - self.data['categories'].append(category) - - for tag in tags: - if tag and tag not in self.data['tags']: - self.data['tags'].append(tag) - - self.save() - print(f"[PromptSystem] Added new prompt with ID {prompt_id}") - return prompt_id - - def update_prompt( - self, - prompt_id: int, - text: Optional[str] = None, - category: Optional[str] = None, - tags: Optional[List[str]] = None, - rating: Optional[int] = None, - notes: Optional[str] = None - ) -> bool: - """ - Update existing prompt. - - Args: - prompt_id: Prompt ID to update - text: New text (optional) - category: New category (optional) - tags: New tags (optional) - rating: New rating (optional) - notes: New notes (optional) - - Returns: - True if successful - """ - prompt = self.get_prompt(prompt_id) - if not prompt: - print(f"[PromptSystem] Prompt {prompt_id} not found") - return False - - # Update fields - if text is not None: - prompt['text'] = text - prompt['hash'] = self._generate_hash(text) - - if category is not None: - prompt['category'] = category - if category not in self.data['categories']: - self.data['categories'].append(category) - - if tags is not None: - prompt['tags'] = tags - for tag in tags: - if tag not in self.data['tags']: - self.data['tags'].append(tag) - - if rating is not None: - prompt['rating'] = max(0, min(5, rating)) - - if notes is not None: - prompt['notes'] = notes - - prompt['updated_at'] = datetime.now().isoformat() - - self.save() - print(f"[PromptSystem] Updated prompt {prompt_id}") - return True - - def get_prompt(self, prompt_id: int) -> Optional[Dict[str, Any]]: - """Get prompt by ID""" - for prompt in self.data['prompts']: - if prompt['id'] == prompt_id: - return prompt - return None - - def find_by_hash(self, prompt_hash: str) -> Optional[Dict[str, Any]]: - """Find prompt by hash""" - for prompt in self.data['prompts']: - if prompt['hash'] == prompt_hash: - return prompt - return None - - def delete_prompt(self, prompt_id: int) -> bool: - """Delete prompt by ID""" - initial_count = len(self.data['prompts']) - self.data['prompts'] = [p for p in self.data['prompts'] if p['id'] != prompt_id] - - if len(self.data['prompts']) < initial_count: - self.save() - print(f"[PromptSystem] Deleted prompt {prompt_id}") - return True - - return False - - def add_thumbnail(self, prompt_id: int, image_path: str) -> bool: - """ - Add thumbnail to prompt with lock/unlock mechanism. - - Args: - prompt_id: Prompt ID - image_path: Path to image file - - Returns: - True if successful - """ - prompt = self.get_prompt(prompt_id) - if not prompt: - return False - - max_thumbnails = self.data['settings'].get('max_thumbnails', 3) - current_thumbnails = prompt.get('thumbnails', []) - - # Use image handler to manage thumbnails - updated_thumbnails = self.image_handler.manage_thumbnails( - current_thumbnails, - image_path, - max_thumbnails - ) - - prompt['thumbnails'] = updated_thumbnails - prompt['updated_at'] = datetime.now().isoformat() - - self.save() - return True - - def lock_thumbnail(self, prompt_id: int, thumbnail_index: int) -> bool: - """Lock a thumbnail""" - prompt = self.get_prompt(prompt_id) - if not prompt: - return False - - prompt['thumbnails'] = self.image_handler.lock_thumbnail( - prompt.get('thumbnails', []), - thumbnail_index - ) - - self.save() - return True - - def unlock_thumbnail(self, prompt_id: int, thumbnail_index: int) -> bool: - """Unlock a thumbnail""" - prompt = self.get_prompt(prompt_id) - if not prompt: - return False - - prompt['thumbnails'] = self.image_handler.unlock_thumbnail( - prompt.get('thumbnails', []), - thumbnail_index - ) - - self.save() - return True - - def add_generation_history( - self, - prompt_id: int, - full_prompt: str, - output_image: str, - workflow_snapshot: Optional[Dict] = None - ) -> bool: - """ - Add generation history entry to prompt. - - Args: - prompt_id: Prompt ID - full_prompt: Full concatenated prompt used - output_image: Path to output image - workflow_snapshot: Optional workflow data - - Returns: - True if successful - """ - prompt = self.get_prompt(prompt_id) - if not prompt: - return False - - history_entry = { - 'timestamp': datetime.now().isoformat(), - 'full_prompt': full_prompt, - 'output_image': normalize_path(output_image), - 'workflow_snapshot': workflow_snapshot - } - - if 'generation_history' not in prompt: - prompt['generation_history'] = [] - - prompt['generation_history'].append(history_entry) - prompt['used_count'] = prompt.get('used_count', 0) + 1 - prompt['updated_at'] = datetime.now().isoformat() - - self.save() - return True - - def search_prompts( - self, - search_text: Optional[str] = None, - category: Optional[str] = None, - tags: Optional[List[str]] = None, - min_rating: int = 0, - max_results: int = 100 - ) -> List[Dict[str, Any]]: - """ - Search prompts with filters. - - Args: - search_text: Text to search in prompt content - category: Filter by category - tags: Filter by tags (any match) - min_rating: Minimum rating filter - max_results: Maximum results to return - - Returns: - List of matching prompts - """ - results = [] - - for prompt in self.data['prompts']: - # Rating filter - if prompt.get('rating', 0) < min_rating: - continue - - # Category filter - if category and prompt.get('category') != category: - continue - - # Tags filter - if tags: - prompt_tags = set(prompt.get('tags', [])) - if not any(tag in prompt_tags for tag in tags): - continue - - # Text search - if search_text: - search_lower = search_text.lower() - if search_lower not in prompt.get('text', '').lower(): - if search_lower not in prompt.get('notes', '').lower(): - continue - - results.append(prompt) - - if len(results) >= max_results: - break - - # Sort by used_count and rating - results.sort(key=lambda p: (p.get('rating', 0), p.get('used_count', 0)), reverse=True) - - return results - - def get_latest_prompt_by_category(self, category: str) -> Optional[Dict[str, Any]]: - """Get most recent prompt in category""" - category_prompts = [p for p in self.data['prompts'] if p.get('category') == category] - - if not category_prompts: - return None - - # Sort by updated_at - category_prompts.sort(key=lambda p: p.get('updated_at', ''), reverse=True) - return category_prompts[0] - - def cleanup_old_unrated(self) -> int: - """ - Remove old unrated prompts based on settings. - - Returns: - Number of prompts removed - """ - if not self.data['settings'].get('auto_cleanup_enabled', True): - return 0 - - cleanup_days = self.data['settings'].get('auto_cleanup_days', 30) - cutoff_date = datetime.now() - timedelta(days=cleanup_days) - cutoff_iso = cutoff_date.isoformat() - - initial_count = len(self.data['prompts']) - - # Remove prompts that are: - # - Unrated (rating = 0) - # - Older than cutoff date - self.data['prompts'] = [ - p for p in self.data['prompts'] - if p.get('rating', 0) > 0 or p.get('created_at', '') > cutoff_iso - ] - - removed = initial_count - len(self.data['prompts']) - - if removed > 0: - self.save() - print(f"[PromptSystem] Cleaned up {removed} old unrated prompts") - - return removed - - def merge_database(self, other_db_path: Path) -> Dict[str, int]: - """ - Merge another database into this one. - - Logic: - - Deduplicate by hash - - Higher rating wins - - Sum used_count - - Keep older created_at - - Merge thumbnails (keep locked + newest) - - Args: - other_db_path: Path to other database JSON file - - Returns: - Dictionary with merge statistics - """ - try: - with open(other_db_path, 'r', encoding='utf-8') as f: - other_data = json.load(f) - - stats = { - 'added': 0, - 'merged': 0, - 'skipped': 0 - } - - for other_prompt in other_data.get('prompts', []): - other_hash = other_prompt.get('hash') - - # Find existing by hash - existing = self.find_by_hash(other_hash) - - if not existing: - # New prompt - add it - # Reassign ID to avoid conflicts - new_id = self._get_next_id() - other_prompt['id'] = new_id - self.data['prompts'].append(other_prompt) - stats['added'] += 1 - else: - # Merge with existing - # Higher rating wins - if other_prompt.get('rating', 0) > existing.get('rating', 0): - existing['rating'] = other_prompt['rating'] - - # Sum used_count - existing['used_count'] = ( - existing.get('used_count', 0) + - other_prompt.get('used_count', 0) - ) - - # Keep older created_at - if other_prompt.get('created_at', '') < existing.get('created_at', ''): - existing['created_at'] = other_prompt['created_at'] - - # Merge thumbnails - existing['thumbnails'] = self.image_handler.merge_thumbnails( - existing.get('thumbnails', []), - other_prompt.get('thumbnails', []), - self.data['settings'].get('max_thumbnails', 3) - ) - - # Merge generation history - existing_history = existing.get('generation_history', []) - other_history = other_prompt.get('generation_history', []) - existing['generation_history'] = existing_history + other_history - - # Update timestamp - existing['updated_at'] = datetime.now().isoformat() - - stats['merged'] += 1 - - # Merge categories and tags - for cat in other_data.get('categories', []): - if cat not in self.data['categories']: - self.data['categories'].append(cat) - - for tag in other_data.get('tags', []): - if tag not in self.data['tags']: - self.data['tags'].append(tag) - - self.save() - print(f"[PromptSystem] Merge complete: {stats}") - return stats - - except Exception as e: - print(f"[PromptSystem] Error merging database: {e}") - return {'error': str(e)} - - def export_to_file(self, export_path: Path) -> bool: - """Export database to file""" - try: - with open(export_path, 'w', encoding='utf-8') as f: - json.dump(self.data, f, indent=2, ensure_ascii=False) - print(f"[PromptSystem] Exported database to {export_path}") - return True - except Exception as e: - print(f"[PromptSystem] Error exporting database: {e}") - return False - - def get_statistics(self) -> Dict[str, Any]: - """Get database statistics""" - return { - 'total_prompts': len(self.data['prompts']), - 'total_categories': len(self.data['categories']), - 'total_tags': len(self.data['tags']), - 'rated_prompts': len([p for p in self.data['prompts'] if p.get('rating', 0) > 0]), - 'unrated_prompts': len([p for p in self.data['prompts'] if p.get('rating', 0) == 0]), - 'prompts_with_thumbnails': len([p for p in self.data['prompts'] if p.get('thumbnails')]), - 'total_generations': sum(p.get('used_count', 0) for p in self.data['prompts']) - } diff --git a/database/image_handler.py b/database/image_handler.py deleted file mode 100644 index 74bd866..0000000 --- a/database/image_handler.py +++ /dev/null @@ -1,239 +0,0 @@ -"""Handle image compression, base64 encoding, and thumbnail management""" - -import io -import base64 -from pathlib import Path -from typing import Optional, List, Dict -from PIL import Image -from datetime import datetime - - -class ImageHandler: - """Manages image thumbnails with lock/unlock mechanism""" - - def __init__(self, thumbnail_size: tuple = (150, 150), quality: int = 70): - """ - Initialize ImageHandler. - - Args: - thumbnail_size: Tuple of (width, height) for thumbnails - quality: JPEG quality (0-100) - """ - self.thumbnail_size = thumbnail_size - self.quality = quality - - def compress_and_encode(self, image_path: str) -> Optional[str]: - """ - Compress image and encode to base64 data URI. - - Args: - image_path: Path to source image - - Returns: - Base64 data URI string or None on error - """ - try: - # Open and convert to RGB (handles RGBA, grayscale, etc.) - img = Image.open(image_path) - - # Convert RGBA to RGB if needed - if img.mode in ('RGBA', 'LA', 'P'): - background = Image.new('RGB', img.size, (255, 255, 255)) - if img.mode == 'P': - img = img.convert('RGBA') - background.paste(img, mask=img.split()[-1] if img.mode in ('RGBA', 'LA') else None) - img = background - elif img.mode != 'RGB': - img = img.convert('RGB') - - # Create thumbnail (maintains aspect ratio) - img.thumbnail(self.thumbnail_size, Image.Resampling.LANCZOS) - - # Encode to JPEG in memory - buffer = io.BytesIO() - img.save(buffer, format='JPEG', quality=self.quality, optimize=True) - buffer.seek(0) - - # Encode to base64 - img_base64 = base64.b64encode(buffer.read()).decode('utf-8') - - # Return as data URI - return f"data:image/jpeg;base64,{img_base64}" - - except Exception as e: - print(f"[PromptSystem] Error compressing image {image_path}: {e}") - return None - - def decode_and_save(self, base64_data: str, output_path: str) -> bool: - """ - Decode base64 data URI and save to file. - - Args: - base64_data: Base64 data URI string - output_path: Path to save the image - - Returns: - True if successful, False otherwise - """ - try: - # Remove data URI prefix if present - if base64_data.startswith('data:image'): - base64_data = base64_data.split(',', 1)[1] - - # Decode base64 - img_data = base64.b64decode(base64_data) - - # Save to file - with open(output_path, 'wb') as f: - f.write(img_data) - - return True - - except Exception as e: - print(f"[PromptSystem] Error decoding and saving image: {e}") - return False - - def manage_thumbnails( - self, - current_thumbnails: List[Dict], - new_image_path: str, - max_thumbnails: int = 3 - ) -> List[Dict]: - """ - Manage thumbnail slots with lock/unlock mechanism. - - Logic: - - Unlocked slots are always replaced with the newest image - - Locked slots are preserved - - Maximum of max_thumbnails slots - - When all locked and adding new, remove oldest locked - - Args: - current_thumbnails: List of thumbnail dicts with 'data', 'locked', 'timestamp' - new_image_path: Path to new image to add - max_thumbnails: Maximum number of thumbnails to keep - - Returns: - Updated list of thumbnails - """ - # Encode new image - new_thumbnail_data = self.compress_and_encode(new_image_path) - if not new_thumbnail_data: - return current_thumbnails - - new_thumbnail = { - 'data': new_thumbnail_data, - 'locked': False, - 'timestamp': datetime.now().isoformat(), - 'source_image': str(Path(new_image_path).name) - } - - # If no thumbnails yet, just add it - if not current_thumbnails: - return [new_thumbnail] - - # Find unlocked slot - unlocked_index = None - for i, thumb in enumerate(current_thumbnails): - if not thumb.get('locked', False): - unlocked_index = i - break - - # Replace unlocked slot - if unlocked_index is not None: - current_thumbnails[unlocked_index] = new_thumbnail - return current_thumbnails - - # All locked - add new slot if under max - if len(current_thumbnails) < max_thumbnails: - current_thumbnails.append(new_thumbnail) - return current_thumbnails - - # All locked and at max - remove oldest locked - # Sort by timestamp, keep newest (max_thumbnails - 1), then add new - current_thumbnails.sort(key=lambda x: x.get('timestamp', ''), reverse=True) - return current_thumbnails[:max_thumbnails - 1] + [new_thumbnail] - - def lock_thumbnail(self, thumbnails: List[Dict], index: int) -> List[Dict]: - """ - Lock a specific thumbnail by index. - - Args: - thumbnails: List of thumbnail dicts - index: Index of thumbnail to lock - - Returns: - Updated thumbnails list - """ - if 0 <= index < len(thumbnails): - thumbnails[index]['locked'] = True - return thumbnails - - def unlock_thumbnail(self, thumbnails: List[Dict], index: int) -> List[Dict]: - """ - Unlock a specific thumbnail by index. - - Args: - thumbnails: List of thumbnail dicts - index: Index of thumbnail to unlock - - Returns: - Updated thumbnails list - """ - if 0 <= index < len(thumbnails): - thumbnails[index]['locked'] = False - return thumbnails - - def merge_thumbnails( - self, - thumbnails_a: List[Dict], - thumbnails_b: List[Dict], - max_thumbnails: int = 3 - ) -> List[Dict]: - """ - Merge two thumbnail lists, keeping locked ones and newest. - - Args: - thumbnails_a: First thumbnail list - thumbnails_b: Second thumbnail list - max_thumbnails: Maximum thumbnails to keep - - Returns: - Merged thumbnail list - """ - # Collect all locked thumbnails - locked = [t for t in thumbnails_a if t.get('locked', False)] - locked.extend([t for t in thumbnails_b if t.get('locked', False)]) - - # Remove duplicates based on data (same image) - seen_data = set() - unique_locked = [] - for thumb in locked: - data_hash = hash(thumb.get('data', '')) - if data_hash not in seen_data: - seen_data.add(data_hash) - unique_locked.append(thumb) - - # If more than max, keep newest locked - if len(unique_locked) > max_thumbnails: - unique_locked.sort(key=lambda x: x.get('timestamp', ''), reverse=True) - return unique_locked[:max_thumbnails] - - # If under max, add newest unlocked to fill - if len(unique_locked) < max_thumbnails: - unlocked = [t for t in thumbnails_a if not t.get('locked', False)] - unlocked.extend([t for t in thumbnails_b if not t.get('locked', False)]) - - # Sort by timestamp, newest first - unlocked.sort(key=lambda x: x.get('timestamp', ''), reverse=True) - - # Add unlocked until we reach max - for thumb in unlocked: - if len(unique_locked) >= max_thumbnails: - break - data_hash = hash(thumb.get('data', '')) - if data_hash not in seen_data: - seen_data.add(data_hash) - unique_locked.append(thumb) - - return unique_locked