character-browser/routes/generator.py

import logging

from flask import render_template, request, redirect, url_for, flash
from models import Preset
from services.generation import generate_from_preset
from services.file_io import get_available_checkpoints
from services.comfyui import get_loaded_checkpoint
from services.workflow import _get_default_checkpoint
from services.sync import _resolve_preset_entity

logger = logging.getLogger('gaze')


def register_routes(app):

    @app.route('/generator', methods=['GET'])
    def generator():
        presets = Preset.query.order_by(Preset.name).all()
        checkpoints = get_available_checkpoints()

        selected_ckpt = get_loaded_checkpoint()
        if not selected_ckpt:
            default_path, _ = _get_default_checkpoint()
            selected_ckpt = default_path

        # Pre-select preset from query param
        preset_slug = request.args.get('preset', '')

        return render_template('generator.html',
                               presets=presets,
                               checkpoints=checkpoints,
                               selected_ckpt=selected_ckpt,
                               preset_slug=preset_slug)

    @app.route('/generator/generate', methods=['POST'])
    def generator_generate():
        preset_slug = request.form.get('preset_slug', '').strip()
        if not preset_slug:
            return {'error': 'No preset selected'}, 400

        preset = Preset.query.filter_by(slug=preset_slug).first()
        if not preset:
            return {'error': 'Preset not found'}, 404

        try:
            overrides = {
                'checkpoint': request.form.get('checkpoint', '').strip() or None,
                'extra_positive': request.form.get('extra_positive', '').strip(),
                'extra_negative': request.form.get('extra_negative', '').strip(),
                'action': 'preview',
            }

            seed_val = request.form.get('seed', '').strip()
            if seed_val:
                overrides['seed'] = int(seed_val)

            width = request.form.get('width', '').strip()
            height = request.form.get('height', '').strip()
            if width and height:
                overrides['width'] = int(width)
                overrides['height'] = int(height)

            job = generate_from_preset(preset, overrides, save_category='generator')

            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'status': 'queued', 'job_id': job['id']}

            flash("Generation queued.")
            return redirect(url_for('generator', preset=preset_slug))

        except Exception as e:
            logger.exception("Generator error: %s", e)
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'error': str(e)}, 500
            flash(f"Error: {str(e)}")
            return redirect(url_for('generator', preset=preset_slug))

    @app.route('/generator/preset_info', methods=['GET'])
    def generator_preset_info():
        """Return resolved entity names for a preset (for the summary panel)."""
        slug = request.args.get('slug', '')
        if not slug:
            return {'error': 'slug required'}, 400

        preset = Preset.query.filter_by(slug=slug).first()
        if not preset:
            return {'error': 'not found'}, 404

        data = preset.data
        info = {}

        # Character
        char_cfg = data.get('character', {})
        char_id = char_cfg.get('character_id')
        if char_id == 'random':
            info['character'] = 'Random'
        elif char_id:
            obj = _resolve_preset_entity('character', char_id)
            info['character'] = obj.name if obj else char_id
        else:
            info['character'] = None

        # Secondary entities
        for key, label in [('outfit', 'outfit'), ('action', 'action'), ('style', 'style'),
                           ('scene', 'scene'), ('detailer', 'detailer'), ('look', 'look')]:
            cfg = data.get(key, {})
            eid = cfg.get(f'{key}_id')
            if eid == 'random':
                info[label] = 'Random'
            elif eid:
                obj = _resolve_preset_entity(key, eid)
                info[label] = obj.name if obj else eid
            else:
                info[label] = None

        # Checkpoint
        ckpt_cfg = data.get('checkpoint', {})
        ckpt_path = ckpt_cfg.get('checkpoint_path')
        if ckpt_path == 'random':
            info['checkpoint'] = 'Random'
        elif ckpt_path:
            info['checkpoint'] = ckpt_path.split('/')[-1].replace('.safetensors', '')
        else:
            info['checkpoint'] = 'Default'

        # Resolution
        res_cfg = data.get('resolution', {})
        if res_cfg.get('random'):
            info['resolution'] = 'Random'
        elif res_cfg.get('width') and res_cfg.get('height'):
            info['resolution'] = f"{res_cfg['width']}x{res_cfg['height']}"
        else:
            info['resolution'] = 'Default'

        return info