import json import logging import os import random import secrets import subprocess import requests from flask import flash, jsonify, redirect, render_template, request, session, url_for from models import Checkpoint, Settings, db logger = logging.getLogger('gaze') def register_routes(app): @app.template_global() def random_gen_image(category, slug): """Return a random generated image path for a resource, or None.""" folder = os.path.join(app.config['UPLOAD_FOLDER'], category, slug) try: images = [f for f in os.listdir(folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))] except FileNotFoundError: return None if not images: return None return f"{category}/{slug}/{random.choice(images)}" @app.context_processor def inject_comfyui_ws_url(): url = app.config.get('COMFYUI_URL', 'http://127.0.0.1:8188') # If the URL is localhost/127.0.0.1, replace it with the current request's host # so that remote clients connect to the correct machine for WebSockets. if '127.0.0.1' in url or 'localhost' in url: host = request.host.split(':')[0] url = url.replace('127.0.0.1', host).replace('localhost', host) # Convert http/https to ws/wss ws_url = url.replace('http://', 'ws://').replace('https://', 'wss://') return dict(COMFYUI_WS_URL=f"{ws_url}/ws") @app.context_processor def inject_default_checkpoint(): checkpoints = Checkpoint.query.order_by(Checkpoint.name).all() return dict(all_checkpoints=checkpoints, default_checkpoint_path=session.get('default_checkpoint', '')) @app.route('/set_default_checkpoint', methods=['POST']) def set_default_checkpoint(): checkpoint_path = request.form.get('checkpoint_path', '') session['default_checkpoint'] = checkpoint_path session.modified = True # Persist to database Settings so it survives across server restarts try: settings = Settings.query.first() if not settings: settings = Settings() db.session.add(settings) settings.default_checkpoint = checkpoint_path db.session.commit() logger.info("Default checkpoint saved to database: %s", checkpoint_path) except Exception as e: logger.error(f"Failed to persist checkpoint to database: {e}") db.session.rollback() # Also persist to comfy_workflow.json for backwards compatibility try: workflow_path = 'comfy_workflow.json' with open(workflow_path, 'r') as f: workflow = json.load(f) # Update node 4 (CheckpointLoaderSimple) with the new checkpoint if '4' in workflow and 'inputs' in workflow['4']: workflow['4']['inputs']['ckpt_name'] = checkpoint_path with open(workflow_path, 'w') as f: json.dump(workflow, f, indent=2) except Exception as e: logger.error(f"Failed to persist checkpoint to workflow file: {e}") return {'status': 'ok'} @app.route('/api/status/comfyui') def api_status_comfyui(): """Return whether ComfyUI is reachable.""" url = app.config.get('COMFYUI_URL', 'http://127.0.0.1:8188') try: resp = requests.get(f'{url}/system_stats', timeout=3) if resp.ok: return {'status': 'ok'} except Exception: pass return {'status': 'error'} @app.route('/api/comfyui/loaded_checkpoint') def api_comfyui_loaded_checkpoint(): """Return the checkpoint name from the most recently completed ComfyUI job.""" url = app.config.get('COMFYUI_URL', 'http://127.0.0.1:8188') try: resp = requests.get(f'{url}/history', timeout=3) if not resp.ok: return {'checkpoint': None} history = resp.json() if not history: return {'checkpoint': None} # Sort by timestamp descending, take the most recent job latest = max(history.values(), key=lambda j: j.get('status', {}).get('status_str', '')) # Node "4" is the checkpoint loader in the workflow nodes = latest.get('prompt', [None, None, {}])[2] ckpt_name = nodes.get('4', {}).get('inputs', {}).get('ckpt_name') return {'checkpoint': ckpt_name} except Exception: return {'checkpoint': None} @app.route('/api/status/mcp') def api_status_mcp(): """Return whether the danbooru-mcp Docker container is running.""" try: result = subprocess.run( ['docker', 'ps', '--filter', 'name=danbooru-mcp', '--format', '{{.Names}}'], capture_output=True, text=True, timeout=5, ) if 'danbooru-mcp' in result.stdout: return {'status': 'ok'} except Exception: pass return {'status': 'error'} @app.route('/api/status/llm') def api_status_llm(): """Return whether the configured LLM provider is reachable.""" try: settings = Settings.query.first() if not settings: return {'status': 'error', 'message': 'Settings not configured'} is_local = settings.llm_provider != 'openrouter' if not is_local: # Check OpenRouter if not settings.openrouter_api_key: return {'status': 'error', 'message': 'API key not configured'} # Try to fetch models list as a lightweight check headers = { "Authorization": f"Bearer {settings.openrouter_api_key}", } resp = requests.get("https://openrouter.ai/api/v1/models", headers=headers, timeout=5) if resp.ok: return {'status': 'ok', 'provider': 'OpenRouter'} else: # Check local provider (Ollama or LMStudio) if not settings.local_base_url: return {'status': 'error', 'message': 'Base URL not configured'} # Try to reach the models endpoint url = f"{settings.local_base_url.rstrip('/')}/models" resp = requests.get(url, timeout=5) if resp.ok: return {'status': 'ok', 'provider': settings.llm_provider.title()} except Exception as e: return {'status': 'error', 'message': str(e)} return {'status': 'error'} @app.route('/api/status/character-mcp') def api_status_character_mcp(): """Return whether the character-mcp Docker container is running.""" try: result = subprocess.run( ['docker', 'ps', '--format', '{{.Names}}'], capture_output=True, text=True, timeout=5, ) # Check if any container name contains 'character-mcp' if any('character-mcp' in line for line in result.stdout.splitlines()): return {'status': 'ok'} except Exception: pass return {'status': 'error'} @app.route('/get_openrouter_models', methods=['POST']) def get_openrouter_models(): api_key = request.form.get('api_key') if not api_key: return {'error': 'API key is required'}, 400 headers = {"Authorization": f"Bearer {api_key}"} try: response = requests.get("https://openrouter.ai/api/v1/models", headers=headers) response.raise_for_status() models = response.json().get('data', []) # Return simplified list of models return {'models': [{'id': m['id'], 'name': m.get('name', m['id'])} for m in models]} except Exception as e: return {'error': str(e)}, 500 @app.route('/get_local_models', methods=['POST']) def get_local_models(): base_url = request.form.get('base_url') if not base_url: return {'error': 'Base URL is required'}, 400 try: response = requests.get(f"{base_url.rstrip('/')}/models") response.raise_for_status() models = response.json().get('data', []) # Ollama/LMStudio often follow the same structure as OpenAI return {'models': [{'id': m['id'], 'name': m.get('name', m['id'])} for m in models]} except Exception as e: return {'error': str(e)}, 500 @app.route('/settings', methods=['GET', 'POST']) def settings(): settings = Settings.query.first() if not settings: settings = Settings() db.session.add(settings) db.session.commit() if request.method == 'POST': settings.llm_provider = request.form.get('llm_provider', 'openrouter') settings.openrouter_api_key = request.form.get('api_key') settings.openrouter_model = request.form.get('model') settings.local_base_url = request.form.get('local_base_url') settings.local_model = request.form.get('local_model') settings.lora_dir_characters = request.form.get('lora_dir_characters') or settings.lora_dir_characters settings.lora_dir_outfits = request.form.get('lora_dir_outfits') or settings.lora_dir_outfits settings.lora_dir_actions = request.form.get('lora_dir_actions') or settings.lora_dir_actions settings.lora_dir_styles = request.form.get('lora_dir_styles') or settings.lora_dir_styles settings.lora_dir_scenes = request.form.get('lora_dir_scenes') or settings.lora_dir_scenes settings.lora_dir_detailers = request.form.get('lora_dir_detailers') or settings.lora_dir_detailers settings.checkpoint_dirs = request.form.get('checkpoint_dirs') or settings.checkpoint_dirs db.session.commit() flash('Settings updated successfully!') return redirect(url_for('settings')) return render_template('settings.html', settings=settings) @app.route('/api/key/regenerate', methods=['POST']) def regenerate_api_key(): settings = Settings.query.first() if not settings: settings = Settings() db.session.add(settings) settings.api_key = secrets.token_hex(32) db.session.commit() logger.info("API key regenerated") return jsonify({'api_key': settings.api_key})