import json import os import re import time 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 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 logger = logging.getLogger('gaze') def register_routes(app): 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_specs', 'hair', 'eyes']: 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 ['full_body', 'top', 'bottom']: 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 = Checkpoint.query.order_by(Checkpoint.name).all() return render_template('checkpoints/index.html', checkpoints=checkpoints) @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/') 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//upload', methods=['POST']) def upload_checkpoint_image(slug): ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404() if 'image' not in request.files: flash('No file part') return redirect(url_for('checkpoint_detail', slug=slug)) file = request.files['image'] if file.filename == '': flash('No selected file') return redirect(url_for('checkpoint_detail', slug=slug)) if file and allowed_file(file.filename): folder = os.path.join(app.config['UPLOAD_FOLDER'], f"checkpoints/{slug}") os.makedirs(folder, exist_ok=True) filename = secure_filename(file.filename) file.save(os.path.join(folder, filename)) ckpt.image_path = f"checkpoints/{slug}/{filename}" db.session.commit() flash('Image uploaded successfully!') return redirect(url_for('checkpoint_detail', slug=slug)) @app.route('/checkpoint//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: print(f"Generation error: {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('/checkpoint//replace_cover_from_preview', methods=['POST']) def replace_checkpoint_cover_from_preview(slug): ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404() preview_path = request.form.get('preview_path') if preview_path and os.path.exists(os.path.join(app.config['UPLOAD_FOLDER'], preview_path)): ckpt.image_path = preview_path db.session.commit() flash('Cover image updated!') else: flash('No valid preview image selected.', 'error') return redirect(url_for('checkpoint_detail', slug=slug)) @app.route('/checkpoint//save_json', methods=['POST']) def save_checkpoint_json(slug): ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404() try: new_data = json.loads(request.form.get('json_data', '')) except (ValueError, TypeError) as e: return {'success': False, 'error': f'Invalid JSON: {e}'}, 400 ckpt.data = new_data flag_modified(ckpt, 'data') db.session.commit() checkpoints_dir = app.config.get('CHECKPOINTS_DIR', 'data/checkpoints') file_path = os.path.join(checkpoints_dir, f'{ckpt.slug}.json') with open(file_path, 'w') as f: json.dump(new_data, f, indent=2) return {'success': True} @app.route('/get_missing_checkpoints') def get_missing_checkpoints(): missing = Checkpoint.query.filter((Checkpoint.image_path == None) | (Checkpoint.image_path == '')).order_by(Checkpoint.name).all() return {'missing': [{'slug': c.slug, 'name': c.name} for c in missing]} @app.route('/clear_all_checkpoint_covers', methods=['POST']) def clear_all_checkpoint_covers(): for ckpt in Checkpoint.query.all(): ckpt.image_path = None db.session.commit() return {'success': True} @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' created_count = 0 skipped_count = 0 overwritten_count = 0 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_count += 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() clean_html = re.sub(r']*>.*?', '', html_raw, flags=re.DOTALL) clean_html = re.sub(r']*>.*?', '', clean_html, flags=re.DOTALL) clean_html = re.sub(r']*>', '', clean_html) clean_html = re.sub(r'<[^>]+>', ' ', clean_html) html_content = ' '.join(clean_html.split()) except Exception as e: print(f"Error reading HTML for {filename}: {e}") defaults = _default_checkpoint_data(checkpoint_path, filename) if html_content: try: print(f"Asking LLM to describe checkpoint: {filename}") 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###" ) llm_response = call_llm(prompt, system_prompt) clean_json = llm_response.replace('```json', '').replace('```', '').strip() ckpt_data = json.loads(clean_json) # Enforce fixed fields ckpt_data['checkpoint_path'] = checkpoint_path ckpt_data['checkpoint_name'] = filename # Fill missing fields with defaults for key, val in defaults.items(): if key not in ckpt_data or ckpt_data[key] is None: ckpt_data[key] = val time.sleep(0.5) except Exception as e: print(f"LLM error for {filename}: {e}. Using defaults.") ckpt_data = defaults else: ckpt_data = defaults try: with open(json_path, 'w') as f: json.dump(ckpt_data, f, indent=2) if is_existing: overwritten_count += 1 else: created_count += 1 except Exception as e: print(f"Error saving JSON for {filename}: {e}") if created_count > 0 or overwritten_count > 0: sync_checkpoints() msg = f'Successfully processed checkpoints: {created_count} created, {overwritten_count} overwritten.' if skipped_count > 0: msg += f' (Skipped {skipped_count} existing)' flash(msg) else: flash(f'No checkpoints created or overwritten. {skipped_count} existing entries found.') return redirect(url_for('checkpoints_index'))