character-browser/routes/generator.py

import json
import logging
from flask import render_template, request, redirect, url_for, flash, session, current_app
from models import db, Character, Outfit, Action, Style, Scene, Detailer, Checkpoint
from services.prompts import build_prompt, build_extras_prompt
from services.workflow import _prepare_workflow, _get_default_checkpoint
from services.job_queue import _enqueue_job, _make_finalize
from services.file_io import get_available_checkpoints

logger = logging.getLogger('gaze')


def register_routes(app):

    @app.route('/generator', methods=['GET', 'POST'])
    def generator():
        characters = Character.query.order_by(Character.name).all()
        checkpoints = get_available_checkpoints()
        actions = Action.query.order_by(Action.name).all()
        outfits = Outfit.query.order_by(Outfit.name).all()
        scenes = Scene.query.order_by(Scene.name).all()
        styles = Style.query.order_by(Style.name).all()
        detailers = Detailer.query.order_by(Detailer.name).all()

        if not checkpoints:
            checkpoints = ["Noob/oneObsession_v19Atypical.safetensors"]

        if request.method == 'POST':
            char_slug = request.form.get('character')
            checkpoint = request.form.get('checkpoint')
            custom_positive = request.form.get('positive_prompt', '')
            custom_negative = request.form.get('negative_prompt', '')

            action_slugs = request.form.getlist('action_slugs')
            outfit_slugs = request.form.getlist('outfit_slugs')
            scene_slugs = request.form.getlist('scene_slugs')
            style_slugs = request.form.getlist('style_slugs')
            detailer_slugs = request.form.getlist('detailer_slugs')
            override_prompt = request.form.get('override_prompt', '').strip()
            width = request.form.get('width') or 1024
            height = request.form.get('height') or 1024

            character = Character.query.filter_by(slug=char_slug).first_or_404()

            sel_actions = Action.query.filter(Action.slug.in_(action_slugs)).all() if action_slugs else []
            sel_outfits = Outfit.query.filter(Outfit.slug.in_(outfit_slugs)).all() if outfit_slugs else []
            sel_scenes = Scene.query.filter(Scene.slug.in_(scene_slugs)).all() if scene_slugs else []
            sel_styles = Style.query.filter(Style.slug.in_(style_slugs)).all() if style_slugs else []
            sel_detailers = Detailer.query.filter(Detailer.slug.in_(detailer_slugs)).all() if detailer_slugs else []

            try:
                with open('comfy_workflow.json', 'r') as f:
                    workflow = json.load(f)

                # Build base prompts from character defaults
                prompts = build_prompt(character.data, default_fields=character.default_fields)

                if override_prompt:
                    prompts["main"] = override_prompt
                else:
                    extras = build_extras_prompt(sel_actions, sel_outfits, sel_scenes, sel_styles, sel_detailers)
                    combined = prompts["main"]
                    if extras:
                        combined = f"{combined}, {extras}"
                    if custom_positive:
                        combined = f"{combined}, {custom_positive}"
                    prompts["main"] = combined

                # Parse optional seed
                seed_val = request.form.get('seed', '').strip()
                fixed_seed = int(seed_val) if seed_val else None

                # Prepare workflow - first selected item per category supplies its LoRA slot
                ckpt_obj = Checkpoint.query.filter_by(checkpoint_path=checkpoint).first() if checkpoint else None
                workflow = _prepare_workflow(
                    workflow, character, prompts, checkpoint, custom_negative,
                    outfit=sel_outfits[0] if sel_outfits else None,
                    action=sel_actions[0] if sel_actions else None,
                    style=sel_styles[0] if sel_styles else None,
                    detailer=sel_detailers[0] if sel_detailers else None,
                    scene=sel_scenes[0] if sel_scenes else None,
                    width=width,
                    height=height,
                    checkpoint_data=ckpt_obj.data if ckpt_obj else None,
                    fixed_seed=fixed_seed,
                )

                print(f"Queueing generator prompt for {character.character_id}")

                _finalize = _make_finalize('characters', character.slug)
                label = f"Generator: {character.name}"
                job = _enqueue_job(label, workflow, _finalize)

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

                flash("Generation queued.")
            except Exception as e:
                print(f"Generator error: {e}")
                if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                    return {'error': str(e)}, 500
                flash(f"Error: {str(e)}")

        return render_template('generator.html', characters=characters, checkpoints=checkpoints,
                               actions=actions, outfits=outfits, scenes=scenes,
                               styles=styles, detailers=detailers)

    @app.route('/generator/preview_prompt', methods=['POST'])
    def generator_preview_prompt():
        char_slug = request.form.get('character')
        if not char_slug:
            return {'error': 'No character selected'}, 400

        character = Character.query.filter_by(slug=char_slug).first()
        if not character:
            return {'error': 'Character not found'}, 404

        action_slugs = request.form.getlist('action_slugs')
        outfit_slugs = request.form.getlist('outfit_slugs')
        scene_slugs = request.form.getlist('scene_slugs')
        style_slugs = request.form.getlist('style_slugs')
        detailer_slugs = request.form.getlist('detailer_slugs')
        custom_positive = request.form.get('positive_prompt', '')

        sel_actions = Action.query.filter(Action.slug.in_(action_slugs)).all() if action_slugs else []
        sel_outfits = Outfit.query.filter(Outfit.slug.in_(outfit_slugs)).all() if outfit_slugs else []
        sel_scenes = Scene.query.filter(Scene.slug.in_(scene_slugs)).all() if scene_slugs else []
        sel_styles = Style.query.filter(Style.slug.in_(style_slugs)).all() if style_slugs else []
        sel_detailers = Detailer.query.filter(Detailer.slug.in_(detailer_slugs)).all() if detailer_slugs else []

        prompts = build_prompt(character.data, default_fields=character.default_fields)
        extras = build_extras_prompt(sel_actions, sel_outfits, sel_scenes, sel_styles, sel_detailers)
        combined = prompts["main"]
        if extras:
            combined = f"{combined}, {extras}"
        if custom_positive:
            combined = f"{combined}, {custom_positive}"

        return {'prompt': combined}