29a6723b25
- Migrate 11 character JSONs from old wardrobe keys to _BODY_GROUP_KEYS format - Add is_favourite/is_nsfw columns to Preset model - Add HTTP response validation and timeouts to ComfyUI client - Add path traversal protection on replace cover route - Deduplicate services/mcp.py (4 functions → 2 generic + 2 wrappers) - Extract apply_library_filters() and clean_html_text() shared helpers - Add named constants for 17 ComfyUI workflow node IDs - Fix bare except clauses in services/llm.py - Fix tags schema in ensure_default_outfit() (list → dict) - Convert f-string logging to lazy % formatting - Add 5-minute polling timeout to frontend waitForJob() - Improve migration error handling (non-duplicate errors log at WARNING) - Update CLAUDE.md to reflect all changes Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
394 lines
18 KiB
Python
394 lines
18 KiB
Python
import json
|
||
import os
|
||
import re
|
||
import logging
|
||
|
||
from flask import render_template, request, redirect, url_for, flash, session
|
||
from sqlalchemy.orm.attributes import flag_modified
|
||
|
||
from models import db, Character, Detailer, Action, Settings
|
||
from services.workflow import _prepare_workflow, _get_default_checkpoint
|
||
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_detailers
|
||
from services.file_io import get_available_loras
|
||
from services.llm import load_prompt, call_llm
|
||
from utils import _BODY_GROUP_KEYS, 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, 'detailers')
|
||
|
||
def _queue_detailer_generation(detailer_obj, character=None, selected_fields=None, client_id=None, action=None, extra_positive=None, extra_negative=None, fixed_seed=None):
|
||
if character:
|
||
combined_data = character.data.copy()
|
||
combined_data['character_id'] = character.character_id
|
||
|
||
# Capture detailer prompt for injection into main prompt later
|
||
detailer_prompt = detailer_obj.data.get('prompt', '')
|
||
|
||
# Merge detailer lora triggers if present
|
||
detailer_lora = detailer_obj.data.get('lora', {})
|
||
if detailer_lora.get('lora_triggers'):
|
||
if 'lora' not in combined_data: combined_data['lora'] = {}
|
||
combined_data['lora']['lora_triggers'] = f"{combined_data['lora'].get('lora_triggers', '')}, {detailer_lora['lora_triggers']}"
|
||
|
||
# Merge character identity and wardrobe fields into selected_fields
|
||
if selected_fields:
|
||
_ensure_character_fields(character, selected_fields)
|
||
else:
|
||
# Auto-include essential character fields (minimal set for batch/default generation)
|
||
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 _BODY_GROUP_KEYS:
|
||
if wardrobe.get(key):
|
||
selected_fields.append(f'wardrobe::{key}')
|
||
selected_fields.extend(['lora::lora_triggers'])
|
||
|
||
default_fields = detailer_obj.default_fields
|
||
active_outfit = character.active_outfit
|
||
else:
|
||
# Detailer only - no character
|
||
detailer_prompt = detailer_obj.data.get('prompt', '')
|
||
combined_data = {
|
||
'character_id': detailer_obj.detailer_id,
|
||
'lora': detailer_obj.data.get('lora', {}),
|
||
}
|
||
if not selected_fields:
|
||
selected_fields = ['lora::lora_triggers']
|
||
default_fields = detailer_obj.default_fields
|
||
active_outfit = 'default'
|
||
|
||
with open('comfy_workflow.json', 'r') as f:
|
||
workflow = json.load(f)
|
||
|
||
prompts = build_prompt(combined_data, selected_fields, default_fields, active_outfit=active_outfit)
|
||
|
||
# Inject detailer prompt directly into main prompt
|
||
if detailer_prompt:
|
||
prompt_str = detailer_prompt if isinstance(detailer_prompt, str) else ', '.join(detailer_prompt)
|
||
prompts['main'] = f"{prompts['main']}, {prompt_str}" if prompts['main'] else prompt_str
|
||
|
||
_append_background(prompts, character)
|
||
|
||
if extra_positive:
|
||
prompts["main"] = f"{prompts['main']}, {extra_positive}"
|
||
|
||
ckpt_path, ckpt_data = _get_default_checkpoint()
|
||
workflow = _prepare_workflow(workflow, character, prompts, detailer=detailer_obj, action=action, custom_negative=extra_negative or None, checkpoint=ckpt_path, checkpoint_data=ckpt_data, fixed_seed=fixed_seed)
|
||
return workflow
|
||
|
||
@app.route('/detailers')
|
||
def detailers_index():
|
||
detailers, fav, nsfw = apply_library_filters(Detailer.query, Detailer)
|
||
return render_template('detailers/index.html', detailers=detailers, favourite_filter=fav, nsfw_filter=nsfw)
|
||
|
||
@app.route('/detailers/rescan', methods=['POST'])
|
||
def rescan_detailers():
|
||
sync_detailers()
|
||
flash('Database synced with detailer files.')
|
||
return redirect(url_for('detailers_index'))
|
||
|
||
@app.route('/detailer/<path:slug>')
|
||
def detailer_detail(slug):
|
||
detailer = Detailer.query.filter_by(slug=slug).first_or_404()
|
||
characters = Character.query.order_by(Character.name).all()
|
||
actions = Action.query.order_by(Action.name).all()
|
||
|
||
# Load state from session
|
||
preferences = session.get(f'prefs_detailer_{slug}')
|
||
preview_image = session.get(f'preview_detailer_{slug}')
|
||
selected_character = session.get(f'char_detailer_{slug}')
|
||
selected_action = session.get(f'action_detailer_{slug}')
|
||
extra_positive = session.get(f'extra_pos_detailer_{slug}', '')
|
||
extra_negative = session.get(f'extra_neg_detailer_{slug}', '')
|
||
|
||
# List existing preview images
|
||
upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], f"detailers/{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"detailers/{slug}/{f}" for f in files]
|
||
|
||
return render_template('detailers/detail.html', detailer=detailer, characters=characters,
|
||
actions=actions, preferences=preferences, preview_image=preview_image,
|
||
selected_character=selected_character, selected_action=selected_action,
|
||
extra_positive=extra_positive, extra_negative=extra_negative,
|
||
existing_previews=existing_previews)
|
||
|
||
@app.route('/detailer/<path:slug>/edit', methods=['GET', 'POST'])
|
||
def edit_detailer(slug):
|
||
detailer = Detailer.query.filter_by(slug=slug).first_or_404()
|
||
loras = get_available_loras('detailers')
|
||
|
||
if request.method == 'POST':
|
||
try:
|
||
# 1. Update basic fields
|
||
detailer.name = request.form.get('detailer_name')
|
||
|
||
# 2. Rebuild the data dictionary
|
||
new_data = detailer.data.copy()
|
||
new_data['detailer_name'] = detailer.name
|
||
|
||
# Update prompt (stored as a plain string)
|
||
new_data['prompt'] = request.form.get('detailer_prompt', '')
|
||
|
||
# Update lora section
|
||
if 'lora' in new_data:
|
||
for key in new_data['lora'].keys():
|
||
form_key = f"lora_{key}"
|
||
if form_key in request.form:
|
||
val = request.form.get(form_key)
|
||
if key == 'lora_weight':
|
||
try: val = float(val)
|
||
except: val = 1.0
|
||
new_data['lora'][key] = val
|
||
|
||
# LoRA weight randomization bounds
|
||
for bound in ['lora_weight_min', 'lora_weight_max']:
|
||
val_str = request.form.get(f'lora_{bound}', '').strip()
|
||
if val_str:
|
||
try:
|
||
new_data.setdefault('lora', {})[bound] = float(val_str)
|
||
except ValueError:
|
||
pass
|
||
else:
|
||
new_data.setdefault('lora', {}).pop(bound, None)
|
||
|
||
# Update structured tags
|
||
new_data['tags'] = {
|
||
'associated_resource': request.form.get('tag_associated_resource', '').strip(),
|
||
'adetailer_targets': request.form.getlist('tag_adetailer_targets'),
|
||
'nsfw': 'tag_nsfw' in request.form,
|
||
}
|
||
detailer.is_nsfw = new_data['tags']['nsfw']
|
||
|
||
detailer.data = new_data
|
||
flag_modified(detailer, "data")
|
||
|
||
# 3. Write back to JSON file
|
||
detailer_file = detailer.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', detailer.detailer_id)}.json"
|
||
file_path = os.path.join(app.config['DETAILERS_DIR'], detailer_file)
|
||
|
||
with open(file_path, 'w') as f:
|
||
json.dump(new_data, f, indent=2)
|
||
|
||
db.session.commit()
|
||
flash('Detailer updated successfully!')
|
||
return redirect(url_for('detailer_detail', slug=slug))
|
||
|
||
except Exception as e:
|
||
logger.exception("Edit error: %s", e)
|
||
flash(f"Error saving changes: {str(e)}")
|
||
|
||
return render_template('detailers/edit.html', detailer=detailer, loras=loras)
|
||
|
||
@app.route('/detailer/<path:slug>/generate', methods=['POST'])
|
||
def generate_detailer_image(slug):
|
||
detailer_obj = Detailer.query.filter_by(slug=slug).first_or_404()
|
||
|
||
try:
|
||
# Get action type
|
||
action = request.form.get('action', 'preview')
|
||
|
||
# Get selected fields
|
||
selected_fields = request.form.getlist('include_field')
|
||
|
||
# Get selected character (if any)
|
||
character_slug = request.form.get('character_slug', '')
|
||
character = _resolve_character(character_slug)
|
||
if character_slug == '__random__' and character:
|
||
character_slug = character.slug
|
||
|
||
# Get selected action (if any)
|
||
action_slug = request.form.get('action_slug', '')
|
||
action_obj = Action.query.filter_by(slug=action_slug).first() if action_slug else None
|
||
|
||
# Get additional prompts
|
||
extra_positive = request.form.get('extra_positive', '').strip()
|
||
extra_negative = request.form.get('extra_negative', '').strip()
|
||
|
||
# Save preferences
|
||
session[f'char_detailer_{slug}'] = character_slug
|
||
session[f'action_detailer_{slug}'] = action_slug
|
||
session[f'extra_pos_detailer_{slug}'] = extra_positive
|
||
session[f'extra_neg_detailer_{slug}'] = extra_negative
|
||
session[f'prefs_detailer_{slug}'] = selected_fields
|
||
session.modified = True
|
||
|
||
# Parse optional seed
|
||
seed_val = request.form.get('seed', '').strip()
|
||
fixed_seed = int(seed_val) if seed_val else None
|
||
|
||
# Build workflow using helper
|
||
workflow = _queue_detailer_generation(detailer_obj, character, selected_fields, action=action_obj, extra_positive=extra_positive, extra_negative=extra_negative, fixed_seed=fixed_seed)
|
||
|
||
char_label = character.name if character else 'no character'
|
||
label = f"Detailer: {detailer_obj.name} ({char_label}) – {action}"
|
||
job = _enqueue_job(label, workflow, _make_finalize('detailers', slug, Detailer, action))
|
||
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return {'status': 'queued', 'job_id': job['id']}
|
||
|
||
return redirect(url_for('detailer_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('detailer_detail', slug=slug))
|
||
|
||
@app.route('/detailers/bulk_create', methods=['POST'])
|
||
def bulk_create_detailers_from_loras():
|
||
_s = Settings.query.first()
|
||
detailers_lora_dir = ((_s.lora_dir_detailers if _s else None) or '/ImageModels/lora/Illustrious/Detailers').rstrip('/')
|
||
_lora_subfolder = os.path.basename(detailers_lora_dir)
|
||
if not os.path.exists(detailers_lora_dir):
|
||
flash('Detailers LoRA directory not found.', 'error')
|
||
return redirect(url_for('detailers_index'))
|
||
|
||
overwrite = request.form.get('overwrite') == 'true'
|
||
skipped = 0
|
||
job_ids = []
|
||
|
||
system_prompt = load_prompt('detailer_system.txt')
|
||
if not system_prompt:
|
||
flash('Detailer system prompt file not found.', 'error')
|
||
return redirect(url_for('detailers_index'))
|
||
|
||
detailers_dir = app.config['DETAILERS_DIR']
|
||
|
||
for filename in os.listdir(detailers_lora_dir):
|
||
if filename.endswith('.safetensors'):
|
||
name_base = filename.rsplit('.', 1)[0]
|
||
detailer_id = re.sub(r'[^a-zA-Z0-9_]', '_', name_base.lower())
|
||
detailer_name = re.sub(r'[^a-zA-Z0-9]+', ' ', name_base).title()
|
||
|
||
json_filename = f"{detailer_id}.json"
|
||
json_path = os.path.join(detailers_dir, json_filename)
|
||
|
||
is_existing = os.path.exists(json_path)
|
||
if is_existing and not overwrite:
|
||
skipped += 1
|
||
continue
|
||
|
||
html_filename = f"{name_base}.html"
|
||
html_path = os.path.join(detailers_lora_dir, html_filename)
|
||
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 %s: %s", html_filename, e)
|
||
|
||
def make_task(filename, name_base, detailer_id, detailer_name, json_path, _lora_subfolder, html_content, system_prompt, is_existing):
|
||
def task_fn(job):
|
||
prompt = f"Describe a detailer LoRA for AI image generation based on the filename: '{filename}'"
|
||
if html_content:
|
||
prompt += f"\n\nHere is descriptive text and metadata extracted from an associated HTML file for this LoRA:\n###\n{html_content[:3000]}\n###"
|
||
|
||
llm_response = call_llm(prompt, system_prompt)
|
||
clean_json = llm_response.replace('```json', '').replace('```', '').strip()
|
||
detailer_data = json.loads(clean_json)
|
||
|
||
detailer_data['detailer_id'] = detailer_id
|
||
detailer_data['detailer_name'] = detailer_name
|
||
|
||
if 'lora' not in detailer_data: detailer_data['lora'] = {}
|
||
detailer_data['lora']['lora_name'] = f"Illustrious/{_lora_subfolder}/{filename}"
|
||
|
||
if not detailer_data['lora'].get('lora_triggers'):
|
||
detailer_data['lora']['lora_triggers'] = name_base
|
||
if detailer_data['lora'].get('lora_weight') is None:
|
||
detailer_data['lora']['lora_weight'] = 1.0
|
||
if detailer_data['lora'].get('lora_weight_min') is None:
|
||
detailer_data['lora']['lora_weight_min'] = 0.7
|
||
if detailer_data['lora'].get('lora_weight_max') is None:
|
||
detailer_data['lora']['lora_weight_max'] = 1.0
|
||
|
||
with open(json_path, 'w') as f:
|
||
json.dump(detailer_data, f, indent=2)
|
||
|
||
job['result'] = {'name': detailer_name, 'action': 'overwritten' if is_existing else 'created'}
|
||
return task_fn
|
||
|
||
job = _enqueue_task(f"Create detailer: {detailer_name}", make_task(filename, name_base, detailer_id, detailer_name, json_path, _lora_subfolder, html_content, system_prompt, is_existing))
|
||
job_ids.append(job['id'])
|
||
|
||
if job_ids:
|
||
def sync_task(job):
|
||
sync_detailers()
|
||
job['result'] = {'synced': True}
|
||
_enqueue_task("Sync detailers DB", sync_task)
|
||
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return {'success': True, 'queued': len(job_ids), 'skipped': skipped}
|
||
flash(f'Queued {len(job_ids)} detailer tasks ({skipped} skipped).')
|
||
return redirect(url_for('detailers_index'))
|
||
|
||
@app.route('/detailer/create', methods=['GET', 'POST'])
|
||
def create_detailer():
|
||
form_data = {}
|
||
|
||
if request.method == 'POST':
|
||
name = request.form.get('name')
|
||
slug = request.form.get('filename', '').strip()
|
||
|
||
form_data = {'name': name, 'filename': slug}
|
||
|
||
if not slug:
|
||
slug = re.sub(r'[^a-zA-Z0-9]+', '_', name.lower()).strip('_')
|
||
|
||
safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', slug)
|
||
if not safe_slug:
|
||
safe_slug = 'detailer'
|
||
|
||
base_slug = safe_slug
|
||
counter = 1
|
||
while os.path.exists(os.path.join(app.config['DETAILERS_DIR'], f"{safe_slug}.json")):
|
||
safe_slug = f"{base_slug}_{counter}"
|
||
counter += 1
|
||
|
||
detailer_data = {
|
||
"detailer_id": safe_slug,
|
||
"detailer_name": name,
|
||
"prompt": "",
|
||
"lora": {
|
||
"lora_name": "",
|
||
"lora_weight": 1.0,
|
||
"lora_triggers": ""
|
||
}
|
||
}
|
||
|
||
try:
|
||
file_path = os.path.join(app.config['DETAILERS_DIR'], f"{safe_slug}.json")
|
||
with open(file_path, 'w') as f:
|
||
json.dump(detailer_data, f, indent=2)
|
||
|
||
new_detailer = Detailer(
|
||
detailer_id=safe_slug, slug=safe_slug, filename=f"{safe_slug}.json",
|
||
name=name, data=detailer_data
|
||
)
|
||
db.session.add(new_detailer)
|
||
db.session.commit()
|
||
|
||
flash('Detailer created successfully!')
|
||
return redirect(url_for('detailer_detail', slug=safe_slug))
|
||
except Exception as e:
|
||
logger.exception("Save error: %s", e)
|
||
flash(f"Failed to create detailer: {e}")
|
||
return render_template('detailers/create.html', form_data=form_data)
|
||
|
||
return render_template('detailers/create.html', form_data=form_data)
|
||
|