Files
character-browser/routes/looks.py
Aodhan Collins 29a6723b25 Code review fixes: wardrobe migration, response validation, path traversal guard, deduplication
- 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>
2026-03-22 00:31:27 +00:00

539 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import json
import os
import re
import logging
from utils import clean_html_text
from flask import render_template, request, redirect, url_for, flash, session
from sqlalchemy.orm.attributes import flag_modified
from models import db, Character, Look, 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, _dedup_tags
from services.sync import sync_looks
from services.file_io import get_available_loras, _count_look_assignments
from services.llm import load_prompt, call_llm
from routes.shared import register_common_routes, apply_library_filters
logger = logging.getLogger('gaze')
def _ensure_look_lora_prefix(lora_name):
"""Ensure look LoRA paths have the correct 'Illustrious/Looks/' prefix."""
if not lora_name:
return lora_name
if not lora_name.startswith('Illustrious/Looks/'):
# Add the prefix if missing
if lora_name.startswith('Illustrious/'):
# Has Illustrious but wrong subfolder - replace
parts = lora_name.split('/', 1)
if len(parts) > 1:
lora_name = 'Illustrious/Looks/' + parts[1]
else:
lora_name = 'Illustrious/Looks/' + lora_name
else:
# No prefix at all - add it
lora_name = 'Illustrious/Looks/' + lora_name
return lora_name
def _fix_look_lora_data(lora_data):
"""Fix look LoRA data to ensure correct prefix."""
if not lora_data:
return lora_data
lora_name = lora_data.get('lora_name', '')
if lora_name:
lora_data = lora_data.copy() # Avoid mutating original
lora_data['lora_name'] = _ensure_look_lora_prefix(lora_name)
return lora_data
def register_routes(app):
register_common_routes(app, 'looks')
@app.route('/looks')
def looks_index():
looks, fav, nsfw = apply_library_filters(Look.query, Look)
look_assignments = _count_look_assignments()
return render_template('looks/index.html', looks=looks, look_assignments=look_assignments, favourite_filter=fav, nsfw_filter=nsfw)
@app.route('/looks/rescan', methods=['POST'])
def rescan_looks():
sync_looks()
flash('Database synced with look files.')
return redirect(url_for('looks_index'))
@app.route('/look/<path:slug>')
def look_detail(slug):
look = Look.query.filter_by(slug=slug).first_or_404()
characters = Character.query.order_by(Character.name).all()
# Pre-select the linked characters if set (supports multi-character assignment)
preferences = session.get(f'prefs_look_{slug}')
preview_image = session.get(f'preview_look_{slug}')
# Get linked character IDs (new character_ids JSON field)
linked_character_ids = look.character_ids or []
# Fallback to legacy character_id if character_ids is empty
if not linked_character_ids and look.character_id:
linked_character_ids = [look.character_id]
# Session-selected character for preview (single selection for generation)
selected_character = session.get(f'char_look_{slug}', linked_character_ids[0] if linked_character_ids else '')
# FIX: Add existing_previews scanning (matching other resource routes)
upload_folder = app.config['UPLOAD_FOLDER']
preview_dir = os.path.join(upload_folder, 'looks', slug)
existing_previews = []
if os.path.isdir(preview_dir):
for f in os.listdir(preview_dir):
if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
existing_previews.append(f'looks/{slug}/{f}')
existing_previews.sort()
extra_positive = session.get(f'extra_pos_look_{slug}', '')
extra_negative = session.get(f'extra_neg_look_{slug}', '')
return render_template('looks/detail.html', look=look, characters=characters,
preferences=preferences, preview_image=preview_image,
selected_character=selected_character,
linked_character_ids=linked_character_ids,
existing_previews=existing_previews,
extra_positive=extra_positive, extra_negative=extra_negative)
@app.route('/look/<path:slug>/edit', methods=['GET', 'POST'])
def edit_look(slug):
look = Look.query.filter_by(slug=slug).first_or_404()
characters = Character.query.order_by(Character.name).all()
loras = get_available_loras('characters')
if request.method == 'POST':
look.name = request.form.get('look_name', look.name)
# Handle multiple character IDs from checkboxes
character_ids = request.form.getlist('character_ids')
look.character_ids = character_ids if character_ids else []
# Also update legacy character_id field for backward compatibility
if character_ids:
look.character_id = character_ids[0]
else:
look.character_id = None
new_data = look.data.copy()
new_data['look_name'] = look.name
new_data['character_id'] = look.character_id
new_data['positive'] = request.form.get('positive', '')
new_data['negative'] = request.form.get('negative', '')
lora_name = request.form.get('lora_lora_name', '')
lora_weight = float(request.form.get('lora_lora_weight', 1.0) or 1.0)
lora_triggers = request.form.get('lora_lora_triggers', '')
new_data['lora'] = {'lora_name': lora_name, 'lora_weight': lora_weight, 'lora_triggers': lora_triggers}
for bound in ['lora_weight_min', 'lora_weight_max']:
val_str = request.form.get(f'lora_{bound}', '').strip()
if val_str:
try:
new_data['lora'][bound] = float(val_str)
except ValueError:
pass
new_data['tags'] = {
'origin_series': request.form.get('tag_origin_series', '').strip(),
'origin_type': request.form.get('tag_origin_type', '').strip(),
'nsfw': 'tag_nsfw' in request.form,
}
look.is_nsfw = new_data['tags']['nsfw']
look.data = new_data
flag_modified(look, 'data')
db.session.commit()
if look.filename:
file_path = os.path.join(app.config['LOOKS_DIR'], look.filename)
with open(file_path, 'w') as f:
json.dump(new_data, f, indent=2)
flash(f'Look "{look.name}" updated!')
return redirect(url_for('look_detail', slug=look.slug))
return render_template('looks/edit.html', look=look, characters=characters, loras=loras)
@app.route('/look/<path:slug>/generate', methods=['POST'])
def generate_look_image(slug):
look = Look.query.filter_by(slug=slug).first_or_404()
try:
action = request.form.get('action', 'preview')
selected_fields = request.form.getlist('include_field')
character_slug = request.form.get('character_slug', '')
character = None
# Only load a character when the user explicitly selects one
character = _resolve_character(character_slug)
if character_slug == '__random__' and character:
character_slug = character.slug
elif character_slug and not character:
# fallback: try matching by character_id
character = Character.query.filter_by(character_id=character_slug).first()
# No fallback to look.character_id — looks are self-contained
# Get additional prompts
extra_positive = request.form.get('extra_positive', '').strip()
extra_negative = request.form.get('extra_negative', '').strip()
session[f'prefs_look_{slug}'] = selected_fields
session[f'char_look_{slug}'] = character_slug
session[f'extra_pos_look_{slug}'] = extra_positive
session[f'extra_neg_look_{slug}'] = extra_negative
session.modified = True
lora_triggers = look.data.get('lora', {}).get('lora_triggers', '')
look_positive = look.data.get('positive', '')
with open('comfy_workflow.json', 'r') as f:
workflow = json.load(f)
if character:
# Merge character identity with look LoRA and positive prompt
combined_data = {
'character_id': character.character_id,
'identity': character.data.get('identity', {}),
'defaults': character.data.get('defaults', {}),
'wardrobe': character.data.get('wardrobe', {}).get(character.active_outfit or 'default',
character.data.get('wardrobe', {}).get('default', {})),
'styles': character.data.get('styles', {}),
'lora': _fix_look_lora_data(look.data.get('lora', {})),
'tags': look.data.get('tags', [])
}
_ensure_character_fields(character, selected_fields,
include_wardrobe=False, include_defaults=True)
prompts = build_prompt(combined_data, selected_fields, character.default_fields)
# Append look-specific triggers and positive
extra = ', '.join(filter(None, [lora_triggers, look_positive]))
if extra:
prompts['main'] = _dedup_tags(f"{prompts['main']}, {extra}" if prompts['main'] else extra)
primary_color = character.data.get('styles', {}).get('primary_color', '')
bg = f"{primary_color} simple background" if primary_color else "simple background"
else:
# Look is self-contained: build prompt from its own positive and triggers only
main = _dedup_tags(', '.join(filter(None, ['(solo:1.2)', lora_triggers, look_positive])))
prompts = {'main': main, 'face': '', 'hand': ''}
bg = "simple background"
prompts['main'] = _dedup_tags(f"{prompts['main']}, {bg}" if prompts['main'] else bg)
if extra_positive:
prompts["main"] = f"{prompts['main']}, {extra_positive}"
# Parse optional seed
seed_val = request.form.get('seed', '').strip()
fixed_seed = int(seed_val) if seed_val else None
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, custom_negative=extra_negative or None, checkpoint=ckpt_path,
checkpoint_data=ckpt_data, look=look, fixed_seed=fixed_seed)
char_label = character.name if character else 'no character'
label = f"Look: {look.name} ({char_label}) {action}"
job = _enqueue_job(label, workflow, _make_finalize('looks', slug, Look, action))
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('look_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('look_detail', slug=slug))
@app.route('/look/<path:slug>/generate_character', methods=['POST'])
def generate_character_from_look(slug):
"""Generate a character JSON using a look as the base."""
look = Look.query.filter_by(slug=slug).first_or_404()
# Get or validate inputs
character_name = request.form.get('character_name', look.name)
use_llm = request.form.get('use_llm') == 'on'
# Auto-generate slug
character_slug = re.sub(r'[^a-zA-Z0-9]+', '_', character_name.lower()).strip('_')
character_slug = re.sub(r'[^a-zA-Z0-9_]', '', character_slug)
# Find available filename
base_slug = character_slug
counter = 1
while os.path.exists(os.path.join(app.config['CHARACTERS_DIR'], f"{character_slug}.json")):
character_slug = f"{base_slug}_{counter}"
counter += 1
if use_llm:
# Use LLM to generate character from look context
system_prompt = load_prompt('character_system.txt')
if not system_prompt:
flash('Character system prompt file not found.', 'error')
return redirect(url_for('look_detail', slug=slug))
prompt = f"""Generate a character based on this look description:
Look Name: {look.name}
Positive Prompt: {look.data.get('positive', '')}
Negative Prompt: {look.data.get('negative', '')}
Tags: {', '.join(look.data.get('tags', []))}
LoRA Triggers: {look.data.get('lora', {}).get('lora_triggers', '')}
Create a complete character JSON with identity, styles, and appropriate wardrobe fields.
The character should match the visual style described in the look.
Character Name: {character_name}
Character ID: {character_slug}"""
try:
llm_response = call_llm(prompt, system_prompt)
# Clean response (remove markdown if present)
clean_json = llm_response.replace('```json', '').replace('```', '').strip()
character_data = json.loads(clean_json)
# Enforce IDs
character_data['character_id'] = character_slug
character_data['character_name'] = character_name
# Ensure the character inherits the look's LoRA with correct path
lora_data = _fix_look_lora_data(look.data.get('lora', {}).copy())
character_data['lora'] = lora_data
except Exception as e:
logger.exception("LLM character generation error: %s", e)
flash(f'Failed to generate character with AI: {e}', 'error')
return redirect(url_for('look_detail', slug=slug))
else:
# Create minimal character template
lora_data = _fix_look_lora_data(look.data.get('lora', {}).copy())
character_data = {
"character_id": character_slug,
"character_name": character_name,
"identity": {
"base": lora_data.get('lora_triggers', ''),
"head": "",
"upper_body": "",
"lower_body": "",
"hands": "",
"feet": "",
"additional": ""
},
"defaults": {
"expression": "",
"pose": "",
"scene": ""
},
"wardrobe": {
"base": "",
"head": "",
"upper_body": "",
"lower_body": "",
"hands": "",
"feet": "",
"additional": ""
},
"styles": {
"aesthetic": "",
"primary_color": "",
"secondary_color": "",
"tertiary_color": ""
},
"lora": lora_data,
"tags": look.data.get('tags', [])
}
# Save character JSON
char_path = os.path.join(app.config['CHARACTERS_DIR'], f"{character_slug}.json")
try:
with open(char_path, 'w') as f:
json.dump(character_data, f, indent=2)
except Exception as e:
flash(f'Failed to save character file: {e}', 'error')
return redirect(url_for('look_detail', slug=slug))
# Create DB entry
character = Character(
character_id=character_slug,
slug=character_slug,
name=character_name,
data=character_data
)
db.session.add(character)
db.session.commit()
# Link the look to this character
look.character_id = character_slug
db.session.commit()
flash(f'Character "{character_name}" created from look!', 'success')
return redirect(url_for('detail', slug=character_slug))
@app.route('/look/create', methods=['GET', 'POST'])
def create_look():
characters = Character.query.order_by(Character.name).all()
loras = get_available_loras('characters')
form_data = {}
if request.method == 'POST':
name = request.form.get('name', '').strip()
character_id = request.form.get('character_id', '') or None
lora_name = request.form.get('lora_lora_name', '')
lora_weight = float(request.form.get('lora_lora_weight', 1.0) or 1.0)
lora_triggers = request.form.get('lora_lora_triggers', '')
positive = request.form.get('positive', '')
negative = request.form.get('negative', '')
tags = {
'origin_series': request.form.get('tag_origin_series', '').strip(),
'origin_type': request.form.get('tag_origin_type', '').strip(),
'nsfw': 'tag_nsfw' in request.form,
}
form_data = {
'name': name, 'character_id': character_id,
'lora_lora_name': lora_name, 'lora_lora_weight': lora_weight,
'lora_lora_triggers': lora_triggers, 'positive': positive,
'negative': negative, 'tags': tags,
}
look_id = re.sub(r'[^a-zA-Z0-9_]', '_', name.lower().replace(' ', '_'))
filename = f'{look_id}.json'
file_path = os.path.join(app.config['LOOKS_DIR'], filename)
data = {
'look_id': look_id,
'look_name': name,
'character_id': character_id,
'positive': positive,
'negative': negative,
'lora': {'lora_name': lora_name, 'lora_weight': lora_weight, 'lora_triggers': lora_triggers},
'tags': tags
}
try:
os.makedirs(app.config['LOOKS_DIR'], exist_ok=True)
with open(file_path, 'w') as f:
json.dump(data, f, indent=2)
slug = re.sub(r'[^a-zA-Z0-9_]', '', look_id)
new_look = Look(look_id=look_id, slug=slug, filename=filename, name=name,
character_id=character_id, data=data,
is_nsfw=tags.get('nsfw', False))
db.session.add(new_look)
db.session.commit()
flash(f'Look "{name}" created!')
return redirect(url_for('look_detail', slug=slug))
except Exception as e:
logger.exception("Save error: %s", e)
flash(f"Failed to create look: {e}")
return render_template('looks/create.html', characters=characters, loras=loras, form_data=form_data)
return render_template('looks/create.html', characters=characters, loras=loras, form_data=form_data)
@app.route('/looks/bulk_create', methods=['POST'])
def bulk_create_looks_from_loras():
_s = Settings.query.first()
lora_dir = ((_s.lora_dir_characters if _s else None) or app.config['LORA_DIR']).rstrip('/')
_lora_subfolder = os.path.basename(lora_dir)
if not os.path.exists(lora_dir):
flash('Looks LoRA directory not found.', 'error')
return redirect(url_for('looks_index'))
overwrite = request.form.get('overwrite') == 'true'
skipped = 0
job_ids = []
system_prompt = load_prompt('look_system.txt')
if not system_prompt:
flash('Look system prompt file not found.', 'error')
return redirect(url_for('looks_index'))
looks_dir = app.config['LOOKS_DIR']
for filename in os.listdir(lora_dir):
if not filename.endswith('.safetensors'):
continue
name_base = filename.rsplit('.', 1)[0]
look_id = re.sub(r'[^a-zA-Z0-9_]', '_', name_base.lower())
look_name = re.sub(r'[^a-zA-Z0-9]+', ' ', name_base).title()
json_filename = f"{look_id}.json"
json_path = os.path.join(looks_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(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, look_id, look_name, json_path, looks_dir, _lora_subfolder, html_content, system_prompt, is_existing):
def task_fn(job):
prompt = f"Create a look profile for a character appearance LoRA based on the filename: '{filename}'"
if html_content:
prompt += f"\n\nHere 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()
look_data = json.loads(clean_json)
look_data['look_id'] = look_id
look_data['look_name'] = look_name
if 'lora' not in look_data:
look_data['lora'] = {}
look_data['lora']['lora_name'] = f"Illustrious/{_lora_subfolder}/{filename}"
if not look_data['lora'].get('lora_triggers'):
look_data['lora']['lora_triggers'] = name_base
if look_data['lora'].get('lora_weight') is None:
look_data['lora']['lora_weight'] = 0.8
if look_data['lora'].get('lora_weight_min') is None:
look_data['lora']['lora_weight_min'] = 0.7
if look_data['lora'].get('lora_weight_max') is None:
look_data['lora']['lora_weight_max'] = 1.0
os.makedirs(looks_dir, exist_ok=True)
with open(json_path, 'w') as f:
json.dump(look_data, f, indent=2)
job['result'] = {'name': look_name, 'action': 'overwritten' if is_existing else 'created'}
return task_fn
job = _enqueue_task(f"Create look: {look_name}", make_task(filename, name_base, look_id, look_name, json_path, looks_dir, _lora_subfolder, html_content, system_prompt, is_existing))
job_ids.append(job['id'])
if job_ids:
def sync_task(job):
sync_looks()
job['result'] = {'synced': True}
_enqueue_task("Sync looks 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)} look tasks ({skipped} skipped).')
return redirect(url_for('looks_index'))