character-browser/routes/checkpoints.py

import json
import os
import re
import logging

from flask import render_template, request, redirect, url_for, flash, session, current_app
from werkzeug.utils import secure_filename
from sqlalchemy.orm.attributes import flag_modified

from models import db, Checkpoint, Character, Settings
from services.workflow import _prepare_workflow, _get_default_checkpoint, _apply_checkpoint_settings, _log_workflow_prompts
from services.job_queue import _enqueue_job, _make_finalize, _enqueue_task
from services.prompts import build_prompt, _resolve_character, _ensure_character_fields, _append_background
from services.sync import sync_checkpoints, _default_checkpoint_data
from services.file_io import get_available_checkpoints
from services.llm import load_prompt, call_llm
from utils import allowed_file, clean_html_text
from routes.shared import register_common_routes, apply_library_filters

logger = logging.getLogger('gaze')


def register_routes(app):
    register_common_routes(app, 'checkpoints')

    def _build_checkpoint_workflow(ckpt_obj, character=None, fixed_seed=None, extra_positive=None, extra_negative=None):
        """Build and return a prepared ComfyUI workflow dict for a checkpoint generation."""
        with open('comfy_workflow.json', 'r') as f:
            workflow = json.load(f)

        if character:
            combined_data = character.data.copy()
            combined_data['character_id'] = character.character_id
            selected_fields = []
            for key in ['base', 'head']:
                if character.data.get('identity', {}).get(key):
                    selected_fields.append(f'identity::{key}')
            selected_fields.append('special::name')
            wardrobe = character.get_active_wardrobe()
            for key in ['base', 'upper_body', 'lower_body']:
                if wardrobe.get(key):
                    selected_fields.append(f'wardrobe::{key}')
            prompts = build_prompt(combined_data, selected_fields, None, active_outfit=character.active_outfit)
            _append_background(prompts, character)
        else:
            prompts = {
                "main": "masterpiece, best quality, 1girl, solo, simple background, looking at viewer",
                "face": "masterpiece, best quality",
                "hand": "masterpiece, best quality",
            }

        if extra_positive:
            prompts["main"] = f"{prompts['main']}, {extra_positive}"

        workflow = _prepare_workflow(workflow, character, prompts, checkpoint=ckpt_obj.checkpoint_path,
                                     checkpoint_data=ckpt_obj.data or {}, custom_negative=extra_negative or None, fixed_seed=fixed_seed)
        return workflow

    @app.route('/checkpoints')
    def checkpoints_index():
        checkpoints, fav, nsfw = apply_library_filters(Checkpoint.query, Checkpoint)
        return render_template('checkpoints/index.html', checkpoints=checkpoints, favourite_filter=fav, nsfw_filter=nsfw)

    @app.route('/checkpoints/rescan', methods=['POST'])
    def rescan_checkpoints():
        sync_checkpoints()
        flash('Checkpoint list synced from disk.')
        return redirect(url_for('checkpoints_index'))

    @app.route('/checkpoint/<path:slug>')
    def checkpoint_detail(slug):
        ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404()
        characters = Character.query.order_by(Character.name).all()
        preview_image = session.get(f'preview_checkpoint_{slug}')
        selected_character = session.get(f'char_checkpoint_{slug}')
        extra_positive = session.get(f'extra_pos_checkpoint_{slug}', '')
        extra_negative = session.get(f'extra_neg_checkpoint_{slug}', '')

        # List existing preview images
        upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], f"checkpoints/{slug}")
        existing_previews = []
        if os.path.isdir(upload_dir):
            files = sorted([f for f in os.listdir(upload_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))], reverse=True)
            existing_previews = [f"checkpoints/{slug}/{f}" for f in files]

        return render_template('checkpoints/detail.html', ckpt=ckpt, characters=characters,
                               preview_image=preview_image, selected_character=selected_character,
                               existing_previews=existing_previews,
                               extra_positive=extra_positive, extra_negative=extra_negative)

    @app.route('/checkpoint/<path:slug>/generate', methods=['POST'])
    def generate_checkpoint_image(slug):
        ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404()
        try:
            character_slug = request.form.get('character_slug', '')
            character = _resolve_character(character_slug)
            if character_slug == '__random__' and character:
                character_slug = character.slug

            # Get additional prompts
            extra_positive = request.form.get('extra_positive', '').strip()
            extra_negative = request.form.get('extra_negative', '').strip()

            session[f'char_checkpoint_{slug}'] = character_slug
            session[f'extra_pos_checkpoint_{slug}'] = extra_positive
            session[f'extra_neg_checkpoint_{slug}'] = extra_negative
            session.modified = True

            seed_val = request.form.get('seed', '').strip()
            fixed_seed = int(seed_val) if seed_val else None
            workflow = _build_checkpoint_workflow(ckpt, character, fixed_seed=fixed_seed, extra_positive=extra_positive, extra_negative=extra_negative)

            char_label = character.name if character else 'random'
            label = f"Checkpoint: {ckpt.name} ({char_label})"
            job = _enqueue_job(label, workflow, _make_finalize('checkpoints', slug, Checkpoint))

            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'status': 'queued', 'job_id': job['id']}
            return redirect(url_for('checkpoint_detail', slug=slug))
        except Exception as e:
            logger.exception("Generation error: %s", e)
            if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
                return {'error': str(e)}, 500
            flash(f"Error during generation: {str(e)}")
            return redirect(url_for('checkpoint_detail', slug=slug))

    @app.route('/checkpoints/bulk_create', methods=['POST'])
    def bulk_create_checkpoints():
        checkpoints_dir = app.config.get('CHECKPOINTS_DIR', 'data/checkpoints')
        os.makedirs(checkpoints_dir, exist_ok=True)

        overwrite = request.form.get('overwrite') == 'true'
        skipped = 0
        written_directly = 0
        job_ids = []

        system_prompt = load_prompt('checkpoint_system.txt')
        if not system_prompt:
            flash('Checkpoint system prompt file not found.', 'error')
            return redirect(url_for('checkpoints_index'))

        dirs = [
            (app.config.get('ILLUSTRIOUS_MODELS_DIR', ''), 'Illustrious'),
            (app.config.get('NOOB_MODELS_DIR', ''), 'Noob'),
        ]

        for dirpath, family in dirs:
            if not dirpath or not os.path.exists(dirpath):
                continue

            for filename in sorted(os.listdir(dirpath)):
                if not (filename.endswith('.safetensors') or filename.endswith('.ckpt')):
                    continue

                checkpoint_path = f"{family}/{filename}"
                name_base = filename.rsplit('.', 1)[0]
                safe_id = re.sub(r'[^a-zA-Z0-9_]', '_', checkpoint_path.rsplit('.', 1)[0]).lower().strip('_')
                json_filename = f"{safe_id}.json"
                json_path = os.path.join(checkpoints_dir, json_filename)

                is_existing = os.path.exists(json_path)
                if is_existing and not overwrite:
                    skipped += 1
                    continue

                # Look for a matching HTML file alongside the model file
                html_path = os.path.join(dirpath, f"{name_base}.html")
                html_content = ""
                if os.path.exists(html_path):
                    try:
                        with open(html_path, 'r', encoding='utf-8', errors='ignore') as hf:
                            html_raw = hf.read()
                            html_content = clean_html_text(html_raw)
                    except Exception as e:
                        logger.error("Error reading HTML for %s: %s", filename, e)

                defaults = _default_checkpoint_data(checkpoint_path, filename)

                if html_content:
                    # Has HTML companion — enqueue LLM task
                    def make_task(filename, checkpoint_path, json_path, html_content, system_prompt, defaults, is_existing):
                        def task_fn(job):
                            prompt = (
                                f"Generate checkpoint metadata JSON for the model file: '{filename}' "
                                f"(checkpoint_path: '{checkpoint_path}').\n\n"
                                f"Here is descriptive text extracted from an associated HTML file:\n###\n{html_content[:3000]}\n###"
                            )
                            try:
                                llm_response = call_llm(prompt, system_prompt)
                                clean_json = llm_response.replace('```json', '').replace('```', '').strip()
                                ckpt_data = json.loads(clean_json)
                                ckpt_data['checkpoint_path'] = checkpoint_path
                                ckpt_data['checkpoint_name'] = filename
                                for key, val in defaults.items():
                                    if key not in ckpt_data or ckpt_data[key] is None:
                                        ckpt_data[key] = val
                            except Exception as e:
                                logger.error("LLM error for %s: %s. Using defaults.", filename, e)
                                ckpt_data = defaults

                            with open(json_path, 'w') as f:
                                json.dump(ckpt_data, f, indent=2)

                            job['result'] = {'name': filename, 'action': 'overwritten' if is_existing else 'created'}
                        return task_fn

                    job = _enqueue_task(f"Create checkpoint: {filename}", make_task(filename, checkpoint_path, json_path, html_content, system_prompt, defaults, is_existing))
                    job_ids.append(job['id'])
                else:
                    # No HTML — write defaults directly (no LLM needed)
                    try:
                        with open(json_path, 'w') as f:
                            json.dump(defaults, f, indent=2)
                        written_directly += 1
                    except Exception as e:
                        logger.error("Error saving JSON for %s: %s", filename, e)

        needs_sync = len(job_ids) > 0 or written_directly > 0

        if needs_sync:
            if job_ids:
                # Sync after all LLM tasks complete
                def sync_task(job):
                    sync_checkpoints()
                    job['result'] = {'synced': True}
                _enqueue_task("Sync checkpoints DB", sync_task)
            else:
                # No LLM tasks — sync immediately
                sync_checkpoints()

        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return {'success': True, 'queued': len(job_ids), 'written_directly': written_directly, 'skipped': skipped}
        flash(f'Queued {len(job_ids)} checkpoint tasks, {written_directly} written directly ({skipped} skipped).')
        return redirect(url_for('checkpoints_index'))