472 lines
24 KiB
Python
472 lines
24 KiB
Python
import json
|
||
import os
|
||
import re
|
||
import logging
|
||
import random
|
||
from flask import render_template, request, redirect, url_for, flash, session, current_app
|
||
from werkzeug.utils import secure_filename
|
||
from models import db, Character, Preset, Outfit, Action, Style, Scene, Detailer, Checkpoint, Look, Settings
|
||
from sqlalchemy.orm.attributes import flag_modified
|
||
from services.prompts import build_prompt, _dedup_tags, _resolve_character, _ensure_character_fields, _append_background
|
||
from services.workflow import _prepare_workflow, _get_default_checkpoint
|
||
from services.job_queue import _enqueue_job, _make_finalize
|
||
from services.sync import sync_presets, _resolve_preset_entity, _resolve_preset_fields, _PRESET_ENTITY_MAP
|
||
from services.llm import load_prompt, call_llm
|
||
from utils import allowed_file
|
||
|
||
logger = logging.getLogger('gaze')
|
||
|
||
|
||
def register_routes(app):
|
||
|
||
@app.route('/presets')
|
||
def presets_index():
|
||
presets = Preset.query.order_by(Preset.filename).all()
|
||
return render_template('presets/index.html', presets=presets)
|
||
|
||
@app.route('/preset/<path:slug>')
|
||
def preset_detail(slug):
|
||
preset = Preset.query.filter_by(slug=slug).first_or_404()
|
||
preview_path = session.get(f'preview_preset_{slug}')
|
||
extra_positive = session.get(f'extra_pos_preset_{slug}', '')
|
||
extra_negative = session.get(f'extra_neg_preset_{slug}', '')
|
||
return render_template('presets/detail.html', preset=preset, preview_path=preview_path,
|
||
extra_positive=extra_positive, extra_negative=extra_negative)
|
||
|
||
@app.route('/preset/<path:slug>/generate', methods=['POST'])
|
||
def generate_preset_image(slug):
|
||
preset = Preset.query.filter_by(slug=slug).first_or_404()
|
||
|
||
try:
|
||
action = request.form.get('action', 'preview')
|
||
|
||
# Get additional prompts
|
||
extra_positive = request.form.get('extra_positive', '').strip()
|
||
extra_negative = request.form.get('extra_negative', '').strip()
|
||
session[f'extra_pos_preset_{slug}'] = extra_positive
|
||
session[f'extra_neg_preset_{slug}'] = extra_negative
|
||
session.modified = True
|
||
|
||
data = preset.data
|
||
|
||
# Resolve entities
|
||
char_cfg = data.get('character', {})
|
||
character = _resolve_preset_entity('character', char_cfg.get('character_id'))
|
||
if not character:
|
||
character = Character.query.order_by(db.func.random()).first()
|
||
|
||
outfit_cfg = data.get('outfit', {})
|
||
action_cfg = data.get('action', {})
|
||
style_cfg = data.get('style', {})
|
||
scene_cfg = data.get('scene', {})
|
||
detailer_cfg = data.get('detailer', {})
|
||
look_cfg = data.get('look', {})
|
||
ckpt_cfg = data.get('checkpoint', {})
|
||
|
||
outfit = _resolve_preset_entity('outfit', outfit_cfg.get('outfit_id'))
|
||
action_obj = _resolve_preset_entity('action', action_cfg.get('action_id'))
|
||
style_obj = _resolve_preset_entity('style', style_cfg.get('style_id'))
|
||
scene_obj = _resolve_preset_entity('scene', scene_cfg.get('scene_id'))
|
||
detailer_obj = _resolve_preset_entity('detailer', detailer_cfg.get('detailer_id'))
|
||
look_obj = _resolve_preset_entity('look', look_cfg.get('look_id'))
|
||
|
||
# Checkpoint: form override > preset config > session default
|
||
checkpoint_override = request.form.get('checkpoint_override', '').strip()
|
||
if checkpoint_override:
|
||
ckpt_obj = Checkpoint.query.filter_by(checkpoint_path=checkpoint_override).first()
|
||
ckpt_path = checkpoint_override
|
||
ckpt_data = ckpt_obj.data if ckpt_obj else None
|
||
else:
|
||
preset_ckpt = ckpt_cfg.get('checkpoint_path')
|
||
if preset_ckpt == 'random':
|
||
ckpt_obj = Checkpoint.query.order_by(db.func.random()).first()
|
||
ckpt_path = ckpt_obj.checkpoint_path if ckpt_obj else None
|
||
ckpt_data = ckpt_obj.data if ckpt_obj else None
|
||
elif preset_ckpt:
|
||
ckpt_obj = Checkpoint.query.filter_by(checkpoint_path=preset_ckpt).first()
|
||
ckpt_path = preset_ckpt
|
||
ckpt_data = ckpt_obj.data if ckpt_obj else None
|
||
else:
|
||
ckpt_path, ckpt_data = _get_default_checkpoint()
|
||
|
||
# Resolve selected fields from preset toggles
|
||
selected_fields = _resolve_preset_fields(data)
|
||
|
||
# Build combined data for prompt building
|
||
active_wardrobe = char_cfg.get('fields', {}).get('wardrobe', {}).get('outfit', 'default')
|
||
wardrobe_source = outfit.data.get('wardrobe', {}) if outfit else None
|
||
if wardrobe_source is None:
|
||
wardrobe_source = character.get_active_wardrobe() if character else {}
|
||
|
||
combined_data = {
|
||
'character_id': character.character_id if character else 'unknown',
|
||
'identity': character.data.get('identity', {}) if character else {},
|
||
'defaults': character.data.get('defaults', {}) if character else {},
|
||
'wardrobe': wardrobe_source,
|
||
'styles': character.data.get('styles', {}) if character else {},
|
||
'lora': (look_obj.data.get('lora', {}) if look_obj
|
||
else (character.data.get('lora', {}) if character else {})),
|
||
'tags': (character.data.get('tags', []) if character else []) + data.get('tags', []),
|
||
}
|
||
|
||
# Build extras prompt from secondary resources
|
||
extras_parts = []
|
||
if action_obj:
|
||
action_fields = action_cfg.get('fields', {})
|
||
from utils import _BODY_GROUP_KEYS
|
||
for key in _BODY_GROUP_KEYS:
|
||
val_cfg = action_fields.get(key, True)
|
||
if val_cfg == 'random':
|
||
val_cfg = random.choice([True, False])
|
||
if val_cfg:
|
||
val = action_obj.data.get('action', {}).get(key, '')
|
||
if val:
|
||
extras_parts.append(val)
|
||
if action_cfg.get('use_lora', True):
|
||
trg = action_obj.data.get('lora', {}).get('lora_triggers', '')
|
||
if trg:
|
||
extras_parts.append(trg)
|
||
extras_parts.extend(action_obj.data.get('tags', []))
|
||
if style_obj:
|
||
s = style_obj.data.get('style', {})
|
||
if s.get('artist_name'):
|
||
extras_parts.append(f"by {s['artist_name']}")
|
||
if s.get('artistic_style'):
|
||
extras_parts.append(s['artistic_style'])
|
||
if style_cfg.get('use_lora', True):
|
||
trg = style_obj.data.get('lora', {}).get('lora_triggers', '')
|
||
if trg:
|
||
extras_parts.append(trg)
|
||
if scene_obj:
|
||
scene_fields = scene_cfg.get('fields', {})
|
||
for key in ['background', 'foreground', 'furniture', 'colors', 'lighting', 'theme']:
|
||
val_cfg = scene_fields.get(key, True)
|
||
if val_cfg == 'random':
|
||
val_cfg = random.choice([True, False])
|
||
if val_cfg:
|
||
val = scene_obj.data.get('scene', {}).get(key, '')
|
||
if val:
|
||
extras_parts.append(val)
|
||
if scene_cfg.get('use_lora', True):
|
||
trg = scene_obj.data.get('lora', {}).get('lora_triggers', '')
|
||
if trg:
|
||
extras_parts.append(trg)
|
||
extras_parts.extend(scene_obj.data.get('tags', []))
|
||
if detailer_obj:
|
||
prompt_val = detailer_obj.data.get('prompt', '')
|
||
if isinstance(prompt_val, list):
|
||
extras_parts.extend(p for p in prompt_val if p)
|
||
elif prompt_val:
|
||
extras_parts.append(prompt_val)
|
||
if detailer_cfg.get('use_lora', True):
|
||
trg = detailer_obj.data.get('lora', {}).get('lora_triggers', '')
|
||
if trg:
|
||
extras_parts.append(trg)
|
||
|
||
with open('comfy_workflow.json', 'r') as f:
|
||
workflow = json.load(f)
|
||
|
||
prompts = build_prompt(combined_data, selected_fields, default_fields=None,
|
||
active_outfit=active_wardrobe)
|
||
if extras_parts:
|
||
extra_str = ', '.join(filter(None, extras_parts))
|
||
prompts['main'] = _dedup_tags(f"{prompts['main']}, {extra_str}" if prompts['main'] else extra_str)
|
||
|
||
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
|
||
|
||
# Resolution: form override > preset config > workflow default
|
||
res_cfg = data.get('resolution', {})
|
||
form_width = request.form.get('width', '').strip()
|
||
form_height = request.form.get('height', '').strip()
|
||
if form_width and form_height:
|
||
gen_width = int(form_width)
|
||
gen_height = int(form_height)
|
||
elif res_cfg.get('random', False):
|
||
_RES_OPTIONS = [
|
||
(1024, 1024), (1152, 896), (896, 1152), (1344, 768),
|
||
(768, 1344), (1280, 800), (800, 1280),
|
||
]
|
||
gen_width, gen_height = random.choice(_RES_OPTIONS)
|
||
else:
|
||
gen_width = res_cfg.get('width') or None
|
||
gen_height = res_cfg.get('height') or None
|
||
|
||
workflow = _prepare_workflow(
|
||
workflow, character, prompts,
|
||
checkpoint=ckpt_path, checkpoint_data=ckpt_data,
|
||
custom_negative=extra_negative or None,
|
||
outfit=outfit if outfit_cfg.get('use_lora', True) else None,
|
||
action=action_obj if action_cfg.get('use_lora', True) else None,
|
||
style=style_obj if style_cfg.get('use_lora', True) else None,
|
||
scene=scene_obj if scene_cfg.get('use_lora', True) else None,
|
||
detailer=detailer_obj if detailer_cfg.get('use_lora', True) else None,
|
||
look=look_obj,
|
||
fixed_seed=fixed_seed,
|
||
width=gen_width,
|
||
height=gen_height,
|
||
)
|
||
|
||
label = f"Preset: {preset.name} – {action}"
|
||
job = _enqueue_job(label, workflow, _make_finalize('presets', slug, Preset, action))
|
||
|
||
session[f'preview_preset_{slug}'] = None
|
||
session.modified = True
|
||
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return {'status': 'queued', 'job_id': job['id']}
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
|
||
except Exception as e:
|
||
logger.exception("Generation error (preset %s): %s", slug, 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('preset_detail', slug=slug))
|
||
|
||
@app.route('/preset/<path:slug>/replace_cover_from_preview', methods=['POST'])
|
||
def replace_preset_cover_from_preview(slug):
|
||
preset = Preset.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(current_app.config['UPLOAD_FOLDER'], preview_path)):
|
||
preset.image_path = preview_path
|
||
db.session.commit()
|
||
flash('Cover image updated!')
|
||
else:
|
||
flash('No valid preview image selected.', 'error')
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
|
||
@app.route('/preset/<path:slug>/upload', methods=['POST'])
|
||
def upload_preset_image(slug):
|
||
preset = Preset.query.filter_by(slug=slug).first_or_404()
|
||
if 'image' not in request.files:
|
||
flash('No file uploaded.')
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
file = request.files['image']
|
||
if file.filename == '':
|
||
flash('No file selected.')
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
filename = secure_filename(file.filename)
|
||
folder = os.path.join(current_app.config['UPLOAD_FOLDER'], f'presets/{slug}')
|
||
os.makedirs(folder, exist_ok=True)
|
||
file.save(os.path.join(folder, filename))
|
||
preset.image_path = f'presets/{slug}/{filename}'
|
||
db.session.commit()
|
||
flash('Image uploaded!')
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
|
||
@app.route('/preset/<path:slug>/edit', methods=['GET', 'POST'])
|
||
def edit_preset(slug):
|
||
preset = Preset.query.filter_by(slug=slug).first_or_404()
|
||
if request.method == 'POST':
|
||
name = request.form.get('preset_name', preset.name)
|
||
preset.name = name
|
||
|
||
def _tog(val):
|
||
"""Convert form value ('true'/'false'/'random') to JSON toggle value."""
|
||
if val == 'random':
|
||
return 'random'
|
||
return val == 'true'
|
||
|
||
def _entity_id(val):
|
||
return val if val else None
|
||
|
||
char_id = _entity_id(request.form.get('char_character_id'))
|
||
new_data = {
|
||
'preset_id': preset.preset_id,
|
||
'preset_name': name,
|
||
'character': {
|
||
'character_id': char_id,
|
||
'use_lora': request.form.get('char_use_lora') == 'on',
|
||
'fields': {
|
||
'identity': {k: _tog(request.form.get(f'id_{k}', 'true'))
|
||
for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']},
|
||
'defaults': {k: _tog(request.form.get(f'def_{k}', 'false'))
|
||
for k in ['expression', 'pose', 'scene']},
|
||
'wardrobe': {
|
||
'outfit': request.form.get('wardrobe_outfit', 'default') or 'default',
|
||
'fields': {k: _tog(request.form.get(f'wd_{k}', 'true'))
|
||
for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']},
|
||
},
|
||
},
|
||
},
|
||
'outfit': {'outfit_id': _entity_id(request.form.get('outfit_id')),
|
||
'use_lora': request.form.get('outfit_use_lora') == 'on'},
|
||
'action': {'action_id': _entity_id(request.form.get('action_id')),
|
||
'use_lora': request.form.get('action_use_lora') == 'on',
|
||
'fields': {k: _tog(request.form.get(f'act_{k}', 'true'))
|
||
for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']}},
|
||
'style': {'style_id': _entity_id(request.form.get('style_id')),
|
||
'use_lora': request.form.get('style_use_lora') == 'on'},
|
||
'scene': {'scene_id': _entity_id(request.form.get('scene_id')),
|
||
'use_lora': request.form.get('scene_use_lora') == 'on',
|
||
'fields': {k: _tog(request.form.get(f'scn_{k}', 'true'))
|
||
for k in ['background', 'foreground', 'furniture', 'colors', 'lighting', 'theme']}},
|
||
'detailer': {'detailer_id': _entity_id(request.form.get('detailer_id')),
|
||
'use_lora': request.form.get('detailer_use_lora') == 'on'},
|
||
'look': {'look_id': _entity_id(request.form.get('look_id'))},
|
||
'checkpoint': {'checkpoint_path': _entity_id(request.form.get('checkpoint_path'))},
|
||
'resolution': {
|
||
'width': int(request.form.get('res_width', 1024)),
|
||
'height': int(request.form.get('res_height', 1024)),
|
||
'random': request.form.get('res_random') == 'on',
|
||
},
|
||
'tags': [t.strip() for t in request.form.get('tags', '').split(',') if t.strip()],
|
||
}
|
||
|
||
preset.data = new_data
|
||
flag_modified(preset, "data")
|
||
db.session.commit()
|
||
|
||
if preset.filename:
|
||
file_path = os.path.join(current_app.config['PRESETS_DIR'], preset.filename)
|
||
with open(file_path, 'w') as f:
|
||
json.dump(new_data, f, indent=2)
|
||
|
||
flash('Preset saved!')
|
||
return redirect(url_for('preset_detail', slug=slug))
|
||
|
||
characters = Character.query.order_by(Character.name).all()
|
||
outfits = Outfit.query.order_by(Outfit.name).all()
|
||
actions = Action.query.order_by(Action.name).all()
|
||
styles = Style.query.order_by(Style.name).all()
|
||
scenes = Scene.query.order_by(Scene.name).all()
|
||
detailers = Detailer.query.order_by(Detailer.name).all()
|
||
looks = Look.query.order_by(Look.name).all()
|
||
checkpoints = Checkpoint.query.order_by(Checkpoint.name).all()
|
||
return render_template('presets/edit.html', preset=preset,
|
||
characters=characters, outfits=outfits, actions=actions,
|
||
styles=styles, scenes=scenes, detailers=detailers,
|
||
looks=looks, checkpoints=checkpoints)
|
||
|
||
@app.route('/preset/<path:slug>/save_json', methods=['POST'])
|
||
def save_preset_json(slug):
|
||
preset = Preset.query.filter_by(slug=slug).first_or_404()
|
||
try:
|
||
new_data = json.loads(request.form.get('json_data', ''))
|
||
preset.data = new_data
|
||
preset.name = new_data.get('preset_name', preset.name)
|
||
flag_modified(preset, "data")
|
||
db.session.commit()
|
||
if preset.filename:
|
||
file_path = os.path.join(current_app.config['PRESETS_DIR'], preset.filename)
|
||
with open(file_path, 'w') as f:
|
||
json.dump(new_data, f, indent=2)
|
||
return {'success': True}
|
||
except Exception as e:
|
||
return {'success': False, 'error': str(e)}, 400
|
||
|
||
@app.route('/preset/<path:slug>/clone', methods=['POST'])
|
||
def clone_preset(slug):
|
||
original = Preset.query.filter_by(slug=slug).first_or_404()
|
||
new_data = dict(original.data)
|
||
|
||
base_id = f"{original.preset_id}_copy"
|
||
new_id = base_id
|
||
counter = 1
|
||
while Preset.query.filter_by(preset_id=new_id).first():
|
||
new_id = f"{base_id}_{counter}"
|
||
counter += 1
|
||
|
||
new_slug = re.sub(r'[^a-zA-Z0-9_]', '', new_id)
|
||
new_data['preset_id'] = new_id
|
||
new_data['preset_name'] = f"{original.name} (Copy)"
|
||
new_filename = f"{new_id}.json"
|
||
|
||
os.makedirs(current_app.config['PRESETS_DIR'], exist_ok=True)
|
||
with open(os.path.join(current_app.config['PRESETS_DIR'], new_filename), 'w') as f:
|
||
json.dump(new_data, f, indent=2)
|
||
|
||
new_preset = Preset(preset_id=new_id, slug=new_slug, filename=new_filename,
|
||
name=new_data['preset_name'], data=new_data)
|
||
db.session.add(new_preset)
|
||
db.session.commit()
|
||
flash(f"Cloned as '{new_data['preset_name']}'")
|
||
return redirect(url_for('preset_detail', slug=new_slug))
|
||
|
||
@app.route('/presets/rescan', methods=['POST'])
|
||
def rescan_presets():
|
||
sync_presets()
|
||
flash('Preset library synced.')
|
||
return redirect(url_for('presets_index'))
|
||
|
||
@app.route('/preset/create', methods=['GET', 'POST'])
|
||
def create_preset():
|
||
if request.method == 'POST':
|
||
name = request.form.get('name', '').strip()
|
||
description = request.form.get('description', '').strip()
|
||
use_llm = request.form.get('use_llm') == 'on'
|
||
|
||
safe_id = re.sub(r'[^a-zA-Z0-9]+', '_', name.lower()).strip('_') or 'preset'
|
||
safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', safe_id)
|
||
base_id = safe_id
|
||
counter = 1
|
||
while os.path.exists(os.path.join(current_app.config['PRESETS_DIR'], f"{safe_id}.json")):
|
||
safe_id = f"{base_id}_{counter}"
|
||
safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', safe_id)
|
||
counter += 1
|
||
|
||
if use_llm and description:
|
||
system_prompt = load_prompt('preset_system.txt')
|
||
if not system_prompt:
|
||
flash('Preset system prompt file not found.', 'error')
|
||
return redirect(request.url)
|
||
try:
|
||
llm_response = call_llm(
|
||
f"Create a preset profile named '{name}' based on this description: {description}",
|
||
system_prompt
|
||
)
|
||
clean_json = llm_response.replace('```json', '').replace('```', '').strip()
|
||
preset_data = json.loads(clean_json)
|
||
except Exception as e:
|
||
logger.exception("LLM error creating preset: %s", e)
|
||
flash(f"AI generation failed: {e}", 'error')
|
||
return redirect(request.url)
|
||
else:
|
||
preset_data = {
|
||
'character': {'character_id': 'random', 'use_lora': True,
|
||
'fields': {
|
||
'identity': {k: True for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']},
|
||
'defaults': {k: False for k in ['expression', 'pose', 'scene']},
|
||
'wardrobe': {'outfit': 'default',
|
||
'fields': {k: True for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']}},
|
||
}},
|
||
'outfit': {'outfit_id': None, 'use_lora': True},
|
||
'action': {'action_id': None, 'use_lora': True,
|
||
'fields': {k: True for k in ['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']}},
|
||
'style': {'style_id': None, 'use_lora': True},
|
||
'scene': {'scene_id': None, 'use_lora': True,
|
||
'fields': {k: True for k in ['background', 'foreground', 'furniture', 'colors', 'lighting', 'theme']}},
|
||
'detailer': {'detailer_id': None, 'use_lora': True},
|
||
'look': {'look_id': None},
|
||
'checkpoint': {'checkpoint_path': None},
|
||
'resolution': {'width': 1024, 'height': 1024, 'random': False},
|
||
'tags': [],
|
||
}
|
||
|
||
preset_data['preset_id'] = safe_id
|
||
preset_data['preset_name'] = name
|
||
|
||
os.makedirs(current_app.config['PRESETS_DIR'], exist_ok=True)
|
||
file_path = os.path.join(current_app.config['PRESETS_DIR'], f"{safe_id}.json")
|
||
with open(file_path, 'w') as f:
|
||
json.dump(preset_data, f, indent=2)
|
||
|
||
new_preset = Preset(preset_id=safe_id, slug=safe_slug,
|
||
filename=f"{safe_id}.json", name=name, data=preset_data)
|
||
db.session.add(new_preset)
|
||
db.session.commit()
|
||
flash(f"Preset '{name}' created!")
|
||
return redirect(url_for('edit_preset', slug=safe_slug))
|
||
|
||
return render_template('presets/create.html')
|
||
|
||
@app.route('/get_missing_presets')
|
||
def get_missing_presets():
|
||
missing = Preset.query.filter((Preset.image_path == None) | (Preset.image_path == '')).order_by(Preset.filename).all()
|
||
return {'missing': [{'slug': p.slug, 'name': p.name} for p in missing]}
|