import json import os import logging from flask import render_template, request, current_app from models import ( db, Character, Look, Outfit, Action, Style, Scene, Detailer, Checkpoint, ) logger = logging.getLogger('gaze') # Category config: (model_class, id_field, name_field, url_prefix, detail_route) _SEARCH_CATEGORIES = { 'characters': (Character, 'character_id', 'character_name', '/character', 'detail'), 'looks': (Look, 'look_id', 'look_name', '/look', 'look_detail'), 'outfits': (Outfit, 'outfit_id', 'outfit_name', '/outfit', 'outfit_detail'), 'actions': (Action, 'action_id', 'action_name', '/action', 'action_detail'), 'styles': (Style, 'style_id', 'style_name', '/style', 'style_detail'), 'scenes': (Scene, 'scene_id', 'scene_name', '/scene', 'scene_detail'), 'detailers': (Detailer, 'detailer_id', 'detailer_name', '/detailer', 'detailer_detail'), 'checkpoints': (Checkpoint, 'checkpoint_id', 'checkpoint_name', '/checkpoint', 'checkpoint_detail'), } def register_routes(app): def _search_resources(query_str, category_filter='all', nsfw_filter='all'): """Search resources by name, tag values, and prompt field content.""" results = [] q = query_str.lower() categories = _SEARCH_CATEGORIES if category_filter == 'all' else {category_filter: _SEARCH_CATEGORIES.get(category_filter)} categories = {k: v for k, v in categories.items() if v} for cat_name, (model_class, id_field, name_field, url_prefix, detail_route) in categories.items(): entities = model_class.query.all() for entity in entities: # Apply NSFW filter if nsfw_filter == 'sfw' and getattr(entity, 'is_nsfw', False): continue if nsfw_filter == 'nsfw' and not getattr(entity, 'is_nsfw', False): continue match_context = None score = 0 # 1. Name match (highest priority) if q in entity.name.lower(): score = 100 if entity.name.lower() == q else 90 match_context = f"Name: {entity.name}" # 2. Tag match if not match_context: data = entity.data or {} tags = data.get('tags', {}) if isinstance(tags, dict): for key, val in tags.items(): if key == 'nsfw': continue if isinstance(val, str) and q in val.lower(): score = 70 match_context = f"Tag {key}: {val}" break elif isinstance(val, list): for item in val: if isinstance(item, str) and q in item.lower(): score = 70 match_context = f"Tag {key}: {item}" break # 3. Prompt field match (search JSON data values) if not match_context: data_str = json.dumps(entity.data or {}).lower() if q in data_str: score = 30 # Find which field matched for context data = entity.data or {} for section_key, section_val in data.items(): if section_key in ('lora', 'tags'): continue if isinstance(section_val, dict): for field_key, field_val in section_val.items(): if isinstance(field_val, str) and q in field_val.lower(): match_context = f"{section_key}.{field_key}: ...{field_val[:80]}..." break elif isinstance(section_val, str) and q in section_val.lower(): match_context = f"{section_key}: ...{section_val[:80]}..." if match_context: break if not match_context: match_context = "Matched in data" if match_context: results.append({ 'type': cat_name, 'slug': entity.slug, 'name': entity.name, 'match_context': match_context, 'score': score, 'is_favourite': getattr(entity, 'is_favourite', False), 'is_nsfw': getattr(entity, 'is_nsfw', False), 'image_path': entity.image_path, 'detail_route': detail_route, }) results.sort(key=lambda x: (-x['score'], x['name'].lower())) return results def _search_images(query_str, nsfw_filter='all'): """Search gallery images by prompt metadata in sidecar JSON files.""" results = [] q = query_str.lower() upload_folder = app.config['UPLOAD_FOLDER'] gallery_cats = ['characters', 'actions', 'outfits', 'scenes', 'styles', 'detailers', 'checkpoints', 'looks', 'presets', 'generator'] for cat in gallery_cats: cat_folder = os.path.join(upload_folder, cat) if not os.path.isdir(cat_folder): continue try: slugs = os.listdir(cat_folder) except OSError: continue for item_slug in slugs: item_folder = os.path.join(cat_folder, item_slug) if not os.path.isdir(item_folder): continue try: files = os.listdir(item_folder) except OSError: continue for filename in files: if not filename.lower().endswith('.json'): continue # This is a sidecar JSON sidecar_path = os.path.join(item_folder, filename) try: with open(sidecar_path) as f: sidecar = json.load(f) except (OSError, json.JSONDecodeError): continue # NSFW filter if nsfw_filter == 'sfw' and sidecar.get('is_nsfw'): continue if nsfw_filter == 'nsfw' and not sidecar.get('is_nsfw'): continue # Search positive/negative prompts stored in sidecar sidecar_str = json.dumps(sidecar).lower() if q in sidecar_str: img_filename = filename.rsplit('.', 1)[0] # Find the actual image file img_path = None for ext in ('.png', '.jpg', '.jpeg', '.webp'): candidate = f"{cat}/{item_slug}/{img_filename}{ext}" if os.path.isfile(os.path.join(upload_folder, candidate)): img_path = candidate break if img_path: results.append({ 'path': img_path, 'category': cat, 'slug': item_slug, 'is_favourite': sidecar.get('is_favourite', False), 'is_nsfw': sidecar.get('is_nsfw', False), 'match_context': f"Found in sidecar metadata", }) if len(results) >= 50: break return results[:50] @app.route('/search') def search(): q = request.args.get('q', '').strip() category = request.args.get('category', 'all') nsfw = request.args.get('nsfw', 'all') search_type = request.args.get('type', 'all') resources = [] images = [] if q: if search_type in ('all', 'resources'): resources = _search_resources(q, category, nsfw) if search_type in ('all', 'images'): images = _search_images(q, nsfw) # Group resources by type grouped = {} for r in resources: grouped.setdefault(r['type'], []).append(r) return render_template( 'search.html', query=q, category=category, nsfw_filter=nsfw, search_type=search_type, grouped_resources=grouped, images=images, total_resources=len(resources), total_images=len(images), )