aodhan/character-browser
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add background job queue system for generation

Browse Source 
- Implements sequential job queue with background worker thread (_enqueue_job, _queue_worker)
- All generate routes now return job_id instead of prompt_id; frontend polls /api/queue/<id>/status
- Queue management UI in navbar with live badge, job list, pause/resume/remove controls
- Fix: replaced url_for() calls inside finalize callbacks with direct string paths (url_for raises RuntimeError without request context in background threads)
- Batch cover generation now uses two-phase pattern: queue all jobs upfront, then poll concurrently via Promise.all so page navigation doesn't interrupt the process
- Strengths gallery sweep migrated to same two-phase pattern; sgStop() cancels queued jobs server-side
- LoRA weight randomisation via lora_weight_min/lora_weight_max already present in _resolve_lora_weight

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Aodhan Collins
2026-03-03 02:32:50 +00:00
parent ae7ba961c1
commit 3c828a170f
21 changed files with 1451 additions and 2146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

651
app.py
View File

@@ -6,6 +6,9 @@ import requests
import random
import asyncio
import subprocess
import threading
import uuid
from collections import deque
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from flask import Flask, render_template, request, redirect, url_for, flash, session
@@ -39,6 +42,201 @@ app.config['SESSION_PERMANENT'] = False
db.init_app(app)
Session(app)
# ---------------------------------------------------------------------------
# Generation Job Queue
# ---------------------------------------------------------------------------
# Each job is a dict:
# id — unique UUID string
# label — human-readable description (e.g. "Tifa Lockhart – preview")
# status — 'pending' | 'processing' | 'done' | 'failed' | 'paused' | 'removed'
# workflow — the fully-prepared ComfyUI workflow dict
# finalize_fn — callable(comfy_prompt_id, job) that saves the image; called after ComfyUI finishes
# error — error message string (when status == 'failed')
# result — dict with image_url etc. (set by finalize_fn on success)
# created_at — unix timestamp
# comfy_prompt_id — the prompt_id returned by ComfyUI (set when processing starts)
_job_queue_lock = threading.Lock()
_job_queue = deque() # ordered list of job dicts (pending + paused + processing)
_job_history = {} # job_id -> job dict (all jobs ever added, for status lookup)
_queue_worker_event = threading.Event() # signals worker that a new job is available
def _enqueue_job(label, workflow, finalize_fn):
"""Add a generation job to the queue. Returns the job dict."""
job = {
'id': str(uuid.uuid4()),
'label': label,
'status': 'pending',
'workflow': workflow,
'finalize_fn': finalize_fn,
'error': None,
'result': None,
'created_at': time.time(),
'comfy_prompt_id': None,
}
with _job_queue_lock:
_job_queue.append(job)
_job_history[job['id']] = job
_queue_worker_event.set()
return job
def _queue_worker():
"""Background thread: processes jobs from _job_queue sequentially."""
while True:
_queue_worker_event.wait()
_queue_worker_event.clear()
while True:
job = None
with _job_queue_lock:
# Find the first pending job
for j in _job_queue:
if j['status'] == 'pending':
job = j
break
if job is None:
break # No pending jobs — go back to waiting
# Mark as processing
with _job_queue_lock:
job['status'] = 'processing'
try:
with app.app_context():
# Send workflow to ComfyUI
prompt_response = queue_prompt(job['workflow'])
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI rejected job: {prompt_response.get('error', 'unknown error')}")
comfy_id = prompt_response['prompt_id']
with _job_queue_lock:
job['comfy_prompt_id'] = comfy_id
# Poll until done (max ~10 minutes)
max_retries = 300
finished = False
while max_retries > 0:
history = get_history(comfy_id)
if comfy_id in history:
finished = True
break
time.sleep(2)
max_retries -= 1
if not finished:
raise Exception("ComfyUI generation timed out")
# Run the finalize callback (saves image to disk / DB)
# finalize_fn(comfy_prompt_id, job) — job is passed so callback can store result
job['finalize_fn'](comfy_id, job)
with _job_queue_lock:
job['status'] = 'done'
except Exception as e:
print(f"[Queue] Job {job['id']} failed: {e}")
with _job_queue_lock:
job['status'] = 'failed'
job['error'] = str(e)
# Remove completed/failed jobs from the active queue (keep in history)
with _job_queue_lock:
try:
_job_queue.remove(job)
except ValueError:
pass # Already removed (e.g. by user)
# Start the background worker thread
_worker_thread = threading.Thread(target=_queue_worker, daemon=True, name='queue-worker')
_worker_thread.start()
# ---------------------------------------------------------------------------
# Queue API routes
# ---------------------------------------------------------------------------
@app.route('/api/queue')
def api_queue_list():
"""Return the current queue as JSON."""
with _job_queue_lock:
jobs = [
{
'id': j['id'],
'label': j['label'],
'status': j['status'],
'error': j['error'],
'created_at': j['created_at'],
}
for j in _job_queue
]
return {'jobs': jobs, 'count': len(jobs)}
@app.route('/api/queue/count')
def api_queue_count():
"""Return just the count of active (non-done, non-failed) jobs."""
with _job_queue_lock:
count = sum(1 for j in _job_queue if j['status'] in ('pending', 'processing', 'paused'))
return {'count': count}
@app.route('/api/queue/<job_id>/remove', methods=['POST'])
def api_queue_remove(job_id):
"""Remove a pending or paused job from the queue."""
with _job_queue_lock:
job = _job_history.get(job_id)
if not job:
return {'error': 'Job not found'}, 404
if job['status'] == 'processing':
return {'error': 'Cannot remove a job that is currently processing'}, 400
try:
_job_queue.remove(job)
except ValueError:
pass # Already not in queue
job['status'] = 'removed'
return {'status': 'ok'}
@app.route('/api/queue/<job_id>/pause', methods=['POST'])
def api_queue_pause(job_id):
"""Toggle pause/resume on a pending job."""
with _job_queue_lock:
job = _job_history.get(job_id)
if not job:
return {'error': 'Job not found'}, 404
if job['status'] == 'pending':
job['status'] = 'paused'
elif job['status'] == 'paused':
job['status'] = 'pending'
_queue_worker_event.set()
else:
return {'error': f'Cannot pause/resume job with status {job["status"]}'}, 400
return {'status': 'ok', 'new_status': job['status']}
@app.route('/api/queue/<job_id>/status')
def api_queue_job_status(job_id):
"""Return the status of a specific job."""
with _job_queue_lock:
job = _job_history.get(job_id)
if not job:
return {'error': 'Job not found'}, 404
resp = {
'id': job['id'],
'label': job['label'],
'status': job['status'],
'error': job['error'],
'comfy_prompt_id': job['comfy_prompt_id'],
}
if job.get('result'):
resp['result'] = job['result']
return resp
# Path to the danbooru-mcp docker-compose project, relative to this file.
MCP_TOOLS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tools')
MCP_COMPOSE_DIR = os.path.join(MCP_TOOLS_DIR, 'danbooru-mcp')
@@ -1284,7 +1482,6 @@ def generator():
checkpoint = request.form.get('checkpoint')
custom_positive = request.form.get('positive_prompt', '')
custom_negative = request.form.get('negative_prompt', '')
client_id = request.form.get('client_id')
action_slugs = request.form.getlist('action_slugs')
outfit_slugs = request.form.getlist('outfit_slugs')
@@ -1336,44 +1533,33 @@ def generator():
)
print(f"Queueing generator prompt for {character.character_id}")
prompt_response = queue_prompt(workflow, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
_char_slug = character.slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
char_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"characters/{_char_slug}")
os.makedirs(char_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(char_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"characters/{_char_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
return
label = f"Generator: {character.name}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
flash("Generation started...")
max_retries = 120
while max_retries > 0:
history = get_history(prompt_id)
if prompt_id in history:
outputs = history[prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
char_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"characters/{character.slug}")
os.makedirs(char_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(char_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"characters/{character.slug}/{filename}"
return render_template('generator.html',
characters=characters, checkpoints=checkpoints,
actions=actions, outfits=outfits, scenes=scenes,
styles=styles, detailers=detailers,
generated_image=relative_path, selected_char=char_slug, selected_ckpt=checkpoint)
time.sleep(2)
max_retries -= 1
flash("Generation timed out.")
return {'status': 'queued', 'job_id': job['id']}
flash("Generation queued.")
except Exception as e:
print(f"Generator error: {e}")
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
@@ -2194,7 +2380,6 @@ def generate_image(slug):
try:
# Get action type
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
# Get selected fields
selected_fields = request.form.getlist('include_field')
@@ -2202,17 +2387,45 @@ def generate_image(slug):
# Save preferences
session[f'prefs_{slug}'] = selected_fields
# Queue generation using helper
prompt_response = _queue_generation(character, action, selected_fields, client_id=client_id)
# Build workflow
with open('comfy_workflow.json', 'r') as f:
workflow = json.load(f)
prompts = build_prompt(character.data, selected_fields, character.default_fields, character.active_outfit)
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, checkpoint=ckpt_path, checkpoint_data=ckpt_data)
# Finalize callback — runs in background thread after ComfyUI finishes
_action = action
_slug = slug
_char_name = character.name
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
char_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"characters/{_slug}")
os.makedirs(char_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(char_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"characters/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action == 'replace':
char_obj = Character.query.filter_by(slug=_slug).first()
if char_obj:
char_obj.image_path = relative_path
db.session.commit()
return
label = f"{character.name} – {action}"
job = _enqueue_job(label, workflow, _finalize)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
# Return JSON if AJAX request
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('detail', slug=slug))
@@ -2510,11 +2723,10 @@ def generate_outfit_image(slug):
try:
# Get action type
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
# Get selected fields
selected_fields = request.form.getlist('include_field')
# Get selected character (if any)
character_slug = request.form.get('character_slug', '')
character = None
@@ -2610,17 +2822,38 @@ def generate_outfit_image(slug):
# Prepare workflow - pass both character and outfit for dual LoRA support
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, outfit=outfit, checkpoint=ckpt_path, checkpoint_data=ckpt_data)
prompt_response = queue_prompt(workflow, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
# Return JSON if AJAX request
_action = action
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
outfit_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"outfits/{_slug}")
os.makedirs(outfit_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(outfit_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"outfits/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action == 'replace':
outfit_obj = Outfit.query.filter_by(slug=_slug).first()
if outfit_obj:
outfit_obj.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Outfit: {outfit.name} ({char_label}) – {action}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('outfit_detail', slug=slug))
@@ -3027,8 +3260,7 @@ def generate_action_image(slug):
try:
# Get action type
action_type = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
# Get selected fields
selected_fields = request.form.getlist('include_field')
@@ -3209,16 +3441,38 @@ def generate_action_image(slug):
# Prepare workflow
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, action=action_obj, checkpoint=ckpt_path, checkpoint_data=ckpt_data)
prompt_response = queue_prompt(workflow, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
_action_type = action_type
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
action_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"actions/{_slug}")
os.makedirs(action_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(action_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"actions/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action_type == 'replace':
action_db = Action.query.filter_by(slug=_slug).first()
if action_db:
action_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Action: {action_obj.name} ({char_label}) – {action_type}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('action_detail', slug=slug))
@@ -3654,7 +3908,8 @@ def upload_style_image(slug):
return redirect(url_for('style_detail', slug=slug))
def _queue_style_generation(style_obj, character=None, selected_fields=None, client_id=None):
def _build_style_workflow(style_obj, character=None, selected_fields=None):
"""Build and return a prepared ComfyUI workflow dict for a style generation."""
if character:
combined_data = character.data.copy()
combined_data['character_id'] = character.character_id
@@ -3733,7 +3988,7 @@ def _queue_style_generation(style_obj, character=None, selected_fields=None, cli
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, style=style_obj, checkpoint=ckpt_path, checkpoint_data=ckpt_data)
return queue_prompt(workflow, client_id=client_id)
return workflow
@app.route('/style/<path:slug>/generate', methods=['POST'])
def generate_style_image(slug):
@@ -3742,7 +3997,6 @@ def generate_style_image(slug):
try:
# Get action type
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
# Get selected fields
selected_fields = request.form.getlist('include_field')
@@ -3763,16 +4017,40 @@ def generate_style_image(slug):
session[f'char_style_{slug}'] = character_slug
session[f'prefs_style_{slug}'] = selected_fields
# Queue generation using helper
prompt_response = _queue_style_generation(style_obj, character, selected_fields, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
# Build workflow using helper (returns workflow dict, not prompt_response)
workflow = _build_style_workflow(style_obj, character, selected_fields)
_action = action
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
style_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"styles/{_slug}")
os.makedirs(style_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(style_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"styles/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action == 'replace':
style_db = Style.query.filter_by(slug=_slug).first()
if style_db:
style_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Style: {style_obj.name} ({char_label}) – {action}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('style_detail', slug=slug))
@@ -3899,7 +4177,8 @@ def generate_missing_styles():
style_slug = style_obj.slug
try:
print(f"Batch generating style: {style_obj.name} with character {character.name}")
prompt_response = _queue_style_generation(style_obj, character=character)
workflow = _build_style_workflow(style_obj, character=character)
prompt_response = queue_prompt(workflow)
prompt_id = prompt_response['prompt_id']
max_retries = 120
@@ -4370,7 +4649,7 @@ def _queue_scene_generation(scene_obj, character=None, selected_fields=None, cli
# For scene generation, we want to ensure Node 20 is handled in _prepare_workflow
ckpt_path, ckpt_data = _get_default_checkpoint()
workflow = _prepare_workflow(workflow, character, prompts, scene=scene_obj, checkpoint=ckpt_path, checkpoint_data=ckpt_data)
return queue_prompt(workflow, client_id=client_id)
return workflow
@app.route('/scene/<path:slug>/generate', methods=['POST'])
def generate_scene_image(slug):
@@ -4379,7 +4658,6 @@ def generate_scene_image(slug):
try:
# Get action type
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
# Get selected fields
selected_fields = request.form.getlist('include_field')
@@ -4400,16 +4678,40 @@ def generate_scene_image(slug):
session[f'char_scene_{slug}'] = character_slug
session[f'prefs_scene_{slug}'] = selected_fields
# Queue generation using helper
prompt_response = _queue_scene_generation(scene_obj, character, selected_fields, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
# Build workflow using helper
workflow = _queue_scene_generation(scene_obj, character, selected_fields)
_action = action
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
scene_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"scenes/{_slug}")
os.makedirs(scene_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(scene_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"scenes/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action == 'replace':
scene_db = Scene.query.filter_by(slug=_slug).first()
if scene_db:
scene_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Scene: {scene_obj.name} ({char_label}) – {action}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('scene_detail', slug=slug))
@@ -4920,7 +5222,7 @@ def _queue_detailer_generation(detailer_obj, character=None, selected_fields=Non
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)
return queue_prompt(workflow, client_id=client_id)
return workflow
@app.route('/detailer/<path:slug>/generate', methods=['POST'])
def generate_detailer_image(slug):
@@ -4928,8 +5230,7 @@ def generate_detailer_image(slug):
try:
# Get action type
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
action_type = request.form.get('action', 'preview')
# Get selected fields
selected_fields = request.form.getlist('include_field')
@@ -4948,7 +5249,7 @@ def generate_detailer_image(slug):
# Get selected action (if any)
action_slug = request.form.get('action_slug', '')
action = Action.query.filter_by(slug=action_slug).first() if action_slug else None
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()
@@ -4961,16 +5262,40 @@ def generate_detailer_image(slug):
session[f'extra_neg_detailer_{slug}'] = extra_negative
session[f'prefs_detailer_{slug}'] = selected_fields
# Queue generation using helper
prompt_response = _queue_detailer_generation(detailer_obj, character, selected_fields, client_id=client_id, action=action, extra_positive=extra_positive, extra_negative=extra_negative)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
prompt_id = prompt_response['prompt_id']
# Build workflow using helper
workflow = _queue_detailer_generation(detailer_obj, character, selected_fields, action=action_obj, extra_positive=extra_positive, extra_negative=extra_negative)
_action_type = action_type
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
detailer_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"detailers/{_slug}")
os.makedirs(detailer_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(detailer_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"detailers/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action_type == 'replace':
detailer_db = Detailer.query.filter_by(slug=_slug).first()
if detailer_db:
detailer_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Detailer: {detailer_obj.name} ({char_label}) – {action_type}"
job = _enqueue_job(label, workflow, _finalize)
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('detailer_detail', slug=slug))
@@ -5316,7 +5641,8 @@ def _apply_checkpoint_settings(workflow, ckpt_data):
return workflow
def _queue_checkpoint_generation(ckpt_obj, character=None, client_id=None):
def _build_checkpoint_workflow(ckpt_obj, character=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)
@@ -5344,14 +5670,12 @@ def _queue_checkpoint_generation(ckpt_obj, character=None, client_id=None):
workflow = _prepare_workflow(workflow, character, prompts, checkpoint=ckpt_obj.checkpoint_path,
checkpoint_data=ckpt_obj.data or {})
return queue_prompt(workflow, client_id=client_id)
return workflow
@app.route('/checkpoint/<path:slug>/generate', methods=['POST'])
def generate_checkpoint_image(slug):
ckpt = Checkpoint.query.filter_by(slug=slug).first_or_404()
try:
client_id = request.form.get('client_id')
character_slug = request.form.get('character_slug', '')
character = None
if character_slug == '__random__':
@@ -5363,14 +5687,37 @@ def generate_checkpoint_image(slug):
character = Character.query.filter_by(slug=character_slug).first()
session[f'char_checkpoint_{slug}'] = character_slug
prompt_response = _queue_checkpoint_generation(ckpt, character, client_id=client_id)
workflow = _build_checkpoint_workflow(ckpt, character)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
ckpt_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"checkpoints/{_slug}")
os.makedirs(ckpt_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(ckpt_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"checkpoints/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
ckpt_db = Checkpoint.query.filter_by(slug=_slug).first()
if ckpt_db:
ckpt_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'random'
label = f"Checkpoint: {ckpt.name} ({char_label})"
job = _enqueue_job(label, workflow, _finalize)
prompt_id = prompt_response['prompt_id']
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('checkpoint_detail', slug=slug))
except Exception as e:
print(f"Checkpoint generation error: {e}")
@@ -5648,7 +5995,6 @@ def generate_look_image(slug):
try:
action = request.form.get('action', 'preview')
client_id = request.form.get('client_id')
selected_fields = request.form.getlist('include_field')
character_slug = request.form.get('character_slug', '')
@@ -5718,13 +6064,37 @@ def generate_look_image(slug):
workflow = _prepare_workflow(workflow, character, prompts, checkpoint=ckpt_path,
checkpoint_data=ckpt_data, look=look)
prompt_response = queue_prompt(workflow, client_id=client_id)
if 'prompt_id' not in prompt_response:
raise Exception(f"ComfyUI failed: {prompt_response.get('error', 'Unknown error')}")
_action = action
_slug = slug
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id]['outputs']
for node_id in outputs:
if 'images' in outputs[node_id]:
image_info = outputs[node_id]['images'][0]
image_data = get_image(image_info['filename'], image_info['subfolder'], image_info['type'])
look_folder = os.path.join(app.config['UPLOAD_FOLDER'], f"looks/{_slug}")
os.makedirs(look_folder, exist_ok=True)
filename = f"gen_{int(time.time())}.png"
file_path = os.path.join(look_folder, filename)
with open(file_path, 'wb') as f:
f.write(image_data)
relative_path = f"looks/{_slug}/{filename}"
image_url = f'/static/uploads/{relative_path}'
job['result'] = {'image_url': image_url, 'relative_path': relative_path}
if _action == 'replace':
look_db = Look.query.filter_by(slug=_slug).first()
if look_db:
look_db.image_path = relative_path
db.session.commit()
return
char_label = character.name if character else 'no character'
label = f"Look: {look.name} ({char_label}) – {action}"
job = _enqueue_job(label, workflow, _finalize)
prompt_id = prompt_response['prompt_id']
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
return {'status': 'queued', 'prompt_id': prompt_id}
return {'status': 'queued', 'job_id': job['id']}
return redirect(url_for('look_detail', slug=slug))
except Exception as e:
@@ -6533,7 +6903,6 @@ def strengths_generate(category, slug):
try:
strength_value = float(request.form.get('strength_value', 1.0))
fixed_seed = int(request.form.get('seed', random.randint(1, 10**15)))
client_id = request.form.get('client_id', '')
# Resolve character: prefer POST body value (reflects current page dropdown),
# then fall back to session.
@@ -6584,9 +6953,35 @@ def strengths_generate(category, slug):
custom_negative=extra_negative
)
result = queue_prompt(workflow, client_id)
prompt_id = result.get('prompt_id', '')
return {'status': 'queued', 'prompt_id': prompt_id}
_category = category
_slug = slug
_strength_value = strength_value
_fixed_seed = fixed_seed
def _finalize(comfy_prompt_id, job):
history = get_history(comfy_prompt_id)
outputs = history[comfy_prompt_id].get('outputs', {})
img_data = None
for node_output in outputs.values():
for img in node_output.get('images', []):
img_data = get_image(img['filename'], img.get('subfolder', ''), img.get('type', 'output'))
break
if img_data:
break
if not img_data:
raise Exception('no image in output')
strength_str = f"{_strength_value:.2f}".replace('.', '_')
upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], _category, _slug, 'strengths')
os.makedirs(upload_dir, exist_ok=True)
out_filename = f"strength_{strength_str}_seed_{_fixed_seed}.png"
out_path = os.path.join(upload_dir, out_filename)
with open(out_path, 'wb') as f:
f.write(img_data)
relative = f"{_category}/{_slug}/strengths/{out_filename}"
job['result'] = {'image_url': f"/static/uploads/{relative}", 'strength_value': _strength_value}
label = f"Strengths: {entity.name} @ {strength_value:.2f}"
job = _enqueue_job(label, workflow, _finalize)
return {'status': 'queued', 'job_id': job['id']}
except Exception as e:
print(f"[Strengths] generate error: {e}")