413 lines
16 KiB
Python
413 lines
16 KiB
Python
import time
|
|
|
|
from aiohttp import web
|
|
import os
|
|
import json
|
|
import folder_paths
|
|
import server # Import server for node_info
|
|
import uuid # For generating unique filenames
|
|
import aiohttp
|
|
|
|
HIDDEN_FOLDERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".hidden_folders")
|
|
|
|
def read_hidden_folders():
|
|
if not os.path.exists(HIDDEN_FOLDERS_FILE):
|
|
return []
|
|
with open(HIDDEN_FOLDERS_FILE, 'r') as f:
|
|
return [line.strip() for line in f.readlines()]
|
|
|
|
def write_hidden_folders(hidden_folders):
|
|
with open(HIDDEN_FOLDERS_FILE, 'w') as f:
|
|
f.write('\n'.join(hidden_folders))
|
|
|
|
|
|
async def get_hello(request: web.Request) -> web.Response:
|
|
return web.json_response({"status": "success", "message": "Hello from the CozyGen API!"})
|
|
|
|
async def toggle_folder_visibility(request: web.Request) -> web.Response:
|
|
data = await request.json()
|
|
folder_to_toggle = data.get('folder')
|
|
|
|
if not folder_to_toggle:
|
|
return web.json_response({"error": "Missing 'folder' in request body"}, status=400)
|
|
|
|
hidden_folders = read_hidden_folders()
|
|
|
|
if folder_to_toggle in hidden_folders:
|
|
hidden_folders.remove(folder_to_toggle)
|
|
else:
|
|
hidden_folders.append(folder_to_toggle)
|
|
|
|
write_hidden_folders(hidden_folders)
|
|
|
|
return web.json_response({"status": "success", "hidden_folders": hidden_folders})
|
|
|
|
async def get_gallery_files(request: web.Request) -> web.Response:
|
|
subfolder = request.rel_url.query.get('subfolder', '')
|
|
show_hidden_str = request.rel_url.query.get('show_hidden', 'false')
|
|
show_hidden = show_hidden_str.lower() == 'true'
|
|
|
|
try:
|
|
page = int(request.rel_url.query.get('page', '1'))
|
|
per_page = int(request.rel_url.query.get('per_page', '20'))
|
|
except ValueError:
|
|
return web.json_response({"error": "Invalid page or per_page parameter"}, status=400)
|
|
|
|
output_directory = folder_paths.get_output_directory()
|
|
hidden_folders = read_hidden_folders()
|
|
|
|
# Security: Prevent directory traversal
|
|
gallery_path = os.path.normpath(os.path.join(output_directory, subfolder))
|
|
if not gallery_path.startswith(output_directory):
|
|
return web.json_response({"error": "Unauthorized path"}, status=403)
|
|
|
|
if not os.path.exists(gallery_path) or not os.path.isdir(gallery_path):
|
|
return web.json_response({"error": "Gallery directory not found"}, status=404)
|
|
|
|
items = os.listdir(gallery_path)
|
|
gallery_items = []
|
|
|
|
for item_name in items:
|
|
item_path = os.path.join(gallery_path, item_name)
|
|
relative_item_path = os.path.join(subfolder, item_name)
|
|
|
|
if not show_hidden and relative_item_path in hidden_folders:
|
|
continue
|
|
|
|
if os.path.isdir(item_path):
|
|
mod_time = os.path.getmtime(item_path)
|
|
gallery_items.append({
|
|
"filename": item_name,
|
|
"type": "directory",
|
|
"subfolder": os.path.join(subfolder, item_name),
|
|
"mod_time": mod_time,
|
|
"hidden": relative_item_path in hidden_folders
|
|
})
|
|
elif item_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp', '.mp4', '.webm', '.mp3', '.wav', '.flac')):
|
|
mod_time = os.path.getmtime(item_path)
|
|
gallery_items.append({
|
|
"filename": item_name,
|
|
"type": "output",
|
|
"subfolder": subfolder,
|
|
"mod_time": mod_time
|
|
})
|
|
|
|
# Sort items: directories first, then by modification time
|
|
gallery_items.sort(key=lambda x: (x['type'] == 'directory', x.get('mod_time', 0)), reverse=True)
|
|
|
|
# Pagination
|
|
total_items = len(gallery_items)
|
|
total_pages = (total_items + per_page - 1) // per_page
|
|
start_index = (page - 1) * per_page
|
|
end_index = start_index + per_page
|
|
paginated_items = gallery_items[start_index:end_index]
|
|
|
|
# Remove mod_time before sending
|
|
for item in paginated_items:
|
|
if 'mod_time' in item:
|
|
del item['mod_time']
|
|
|
|
return web.json_response({
|
|
"items": paginated_items,
|
|
"page": page,
|
|
"per_page": per_page,
|
|
"total_pages": total_pages,
|
|
"total_items": total_items
|
|
})
|
|
|
|
async def upload_image(request: web.Request) -> web.Response:
|
|
reader = await request.multipart()
|
|
field = await reader.next()
|
|
|
|
if field.name != 'image':
|
|
return web.json_response({"error": "Expected field 'image'"}, status=400)
|
|
|
|
filename = field.filename
|
|
if not filename:
|
|
return web.json_response({"error": "No filename provided"}, status=400)
|
|
|
|
# Generate a unique filename to prevent collisions
|
|
unique_filename = f"{uuid.uuid4()}_{filename}"
|
|
|
|
# Save to the input directory of ComfyUI
|
|
input_dir = folder_paths.get_input_directory()
|
|
save_path = os.path.join(input_dir, unique_filename)
|
|
|
|
size = 0
|
|
with open(save_path, 'wb') as f:
|
|
while True:
|
|
chunk = await field.read_chunk()
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
size += len(chunk)
|
|
|
|
return web.json_response({"filename": unique_filename, "size": size})
|
|
|
|
async def get_workflow_list(request: web.Request) -> web.Response:
|
|
workflows_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflows")
|
|
if not os.path.exists(workflows_dir):
|
|
return web.json_response({"error": "Workflows directory not found"}, status=404)
|
|
|
|
workflow_files = [f for f in os.listdir(workflows_dir) if f.endswith('.json')]
|
|
return web.json_response({"workflows": workflow_files})
|
|
|
|
async def get_workflow_file(request: web.Request) -> web.Response:
|
|
filename = request.match_info.get('filename', '')
|
|
workflows_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflows")
|
|
workflow_path = os.path.join(workflows_dir, filename)
|
|
|
|
if not os.path.exists(workflow_path):
|
|
return web.json_response({"error": f"Workflow file '{filename}' not found"}, status=404)
|
|
|
|
try:
|
|
with open(workflow_path, 'r', encoding='utf-8') as f:
|
|
workflow_content = json.load(f)
|
|
return web.json_response(workflow_content)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"error": f"Invalid JSON in workflow file '{filename}'"}, status=400)
|
|
except Exception:
|
|
return web.json_response({"error": f"Error reading workflow file: {filename}"}, status=500)
|
|
|
|
import comfy.samplers
|
|
|
|
# Get all valid model folder types from ComfyUI itself
|
|
valid_model_types = folder_paths.folder_names_and_paths.keys()
|
|
|
|
# A map for aliases to official folder_paths names
|
|
alias_map = {
|
|
"unet": "unet_gguf"
|
|
}
|
|
|
|
async def get_choices(request: web.Request) -> web.Response:
|
|
choice_type = request.rel_url.query.get('type', '')
|
|
|
|
if not choice_type:
|
|
return web.json_response({"error": "Missing 'type' query parameter"}, status=400)
|
|
|
|
# Alias map for backward compatibility
|
|
alias_map = {
|
|
"samplers_list": "sampler",
|
|
"schedulers_list": "scheduler",
|
|
"unet": "unet_gguf" # Example of another potential alias
|
|
}
|
|
resolved_choice_type = alias_map.get(choice_type, choice_type)
|
|
|
|
choices = []
|
|
if resolved_choice_type == "scheduler":
|
|
choices = comfy.samplers.KSampler.SCHEDULERS
|
|
elif resolved_choice_type == "sampler":
|
|
choices = comfy.samplers.KSampler.SAMPLERS
|
|
elif resolved_choice_type in valid_model_types:
|
|
choices = folder_paths.get_filename_list(resolved_choice_type)
|
|
else:
|
|
return web.json_response({"error": f"Invalid choice type: {choice_type}"}, status=400)
|
|
|
|
return web.json_response({"choices": choices})
|
|
|
|
import asyncio
|
|
from collections import deque
|
|
|
|
# Global queue to hold generation jobs
|
|
# Each job is a dict: {'id': str, 'status': str, 'data': dict, 'workflowName': str}
|
|
server_queue = deque()
|
|
queue_lock = asyncio.Lock()
|
|
|
|
async def submit_to_queue(request: web.Request) -> web.Response:
|
|
"""Adds a new workflow generation job to the server-side queue."""
|
|
try:
|
|
payload = await request.json()
|
|
prompt = payload.get('prompt')
|
|
workflow_name = payload.get('workflowName', 'Unknown')
|
|
|
|
if not prompt:
|
|
return web.json_response({"error": "Missing 'prompt' in request body"}, status=400)
|
|
|
|
job_id = str(uuid.uuid4())
|
|
job = {
|
|
"id": job_id,
|
|
"status": "queued",
|
|
"data": prompt,
|
|
"workflowName": workflow_name,
|
|
"timestamp": time.time(),
|
|
"thumbnailUrl": None, # Will be populated on completion
|
|
"comfy_prompt_id": None, # Will be populated by the worker
|
|
}
|
|
|
|
async with queue_lock:
|
|
server_queue.append(job)
|
|
|
|
return web.json_response({"status": "success", "job_id": job_id})
|
|
except json.JSONDecodeError:
|
|
return web.json_response({"error": "Invalid JSON in request body"}, status=400)
|
|
except Exception:
|
|
return web.json_response({"error": "An internal server error occurred"}, status=500)
|
|
|
|
async def get_queue_status(request: web.Request) -> web.Response:
|
|
"""Returns the current state of the generation queue."""
|
|
async with queue_lock:
|
|
# Simple pruning: if queue is long, trim oldest completed
|
|
while len(server_queue) > 50:
|
|
try:
|
|
oldest_completed = next(j for j in server_queue if j['status'] in ['completed', 'error'])
|
|
server_queue.remove(oldest_completed)
|
|
except StopIteration:
|
|
break # No completed jobs to prune
|
|
|
|
queue_list = sorted(list(server_queue), key=lambda j: j.get('timestamp', 0), reverse=True)
|
|
return web.json_response({"queue": queue_list})
|
|
|
|
async def queue_worker():
|
|
"""The background worker that processes the server-side queue."""
|
|
prompt_url = "http://127.0.0.1:8188/prompt"
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
while True:
|
|
job_to_run = None
|
|
async with queue_lock:
|
|
queued_job = next((job for job in server_queue if job['status'] == 'queued'), None)
|
|
if queued_job:
|
|
queued_job['status'] = 'running'
|
|
job_to_run = queued_job
|
|
|
|
if job_to_run:
|
|
try:
|
|
job_id = job_to_run['id']
|
|
prompt = job_to_run['data']
|
|
|
|
# Dynamically inject the job_id into the workflow
|
|
output_node_id = None
|
|
for node_id, node_data in prompt.items():
|
|
if node_data["class_type"] in ["CozyGenOutput", "CozyGenVideoOutput"]:
|
|
output_node_id = node_id
|
|
break
|
|
|
|
if output_node_id:
|
|
# Find a unique ID for our new node
|
|
job_id_node_id = str(max([int(k) for k in prompt.keys() if k.isdigit()] + [0]) + 1)
|
|
|
|
# Add the new CozyGenJobID node
|
|
prompt[job_id_node_id] = {
|
|
"inputs": {"job_id": job_id},
|
|
"class_type": "CozyGenJobID",
|
|
}
|
|
|
|
# Link the output node to our new job_id node
|
|
prompt[output_node_id]["inputs"]["job_id"] = [job_id_node_id, 0]
|
|
|
|
# Prepare the payload for the /prompt endpoint (no longer needs extra_data)
|
|
payload = {"prompt": prompt}
|
|
|
|
async with session.post(prompt_url, json=payload) as response:
|
|
if response.status == 200:
|
|
response_data = await response.json()
|
|
comfy_prompt_id = response_data.get('prompt_id')
|
|
async with queue_lock:
|
|
job_to_run['comfy_prompt_id'] = comfy_prompt_id
|
|
else:
|
|
async with queue_lock:
|
|
job_to_run['status'] = 'queued'
|
|
|
|
except Exception:
|
|
async with queue_lock:
|
|
job_to_run['status'] = 'queued'
|
|
|
|
await asyncio.sleep(1)
|
|
|
|
async def cozygen_startup(app):
|
|
"""Starts the CozyGen background worker on server startup."""
|
|
asyncio.create_task(queue_worker())
|
|
|
|
|
|
async def get_presets(request: web.Request) -> web.Response:
|
|
workflow_name = request.rel_url.query.get('workflow', '')
|
|
if not workflow_name:
|
|
return web.json_response({"error": "Missing 'workflow' query parameter"}, status=400)
|
|
|
|
presets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets")
|
|
workflow_presets_dir = os.path.join(presets_dir, os.path.splitext(workflow_name)[0])
|
|
|
|
if not os.path.exists(workflow_presets_dir):
|
|
return web.json_response({"presets": []})
|
|
|
|
preset_files = [f for f in os.listdir(workflow_presets_dir) if f.endswith('.json')]
|
|
return web.json_response({"presets": preset_files})
|
|
|
|
async def get_preset_file(request: web.Request) -> web.Response:
|
|
workflow_name = request.match_info.get('workflow_name', '')
|
|
preset_name = request.match_info.get('preset_name', '')
|
|
|
|
presets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets")
|
|
preset_path = os.path.join(presets_dir, os.path.splitext(workflow_name)[0], preset_name)
|
|
|
|
if not os.path.exists(preset_path):
|
|
return web.json_response({"error": f"Preset file '{preset_name}' not found"}, status=404)
|
|
|
|
try:
|
|
with open(preset_path, 'r', encoding='utf-8') as f:
|
|
preset_content = json.load(f)
|
|
return web.json_response(preset_content)
|
|
except Exception as e:
|
|
return web.json_response({"error": f"Error reading preset file: {e}"}, status=500)
|
|
|
|
async def save_preset(request: web.Request) -> web.Response:
|
|
data = await request.json()
|
|
workflow_name = data.get('workflow')
|
|
preset_name = data.get('preset_name')
|
|
form_data = data.get('form_data')
|
|
|
|
if not all([workflow_name, preset_name, form_data]):
|
|
return web.json_response({"error": "Missing required fields"}, status=400)
|
|
|
|
presets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets")
|
|
workflow_presets_dir = os.path.join(presets_dir, os.path.splitext(workflow_name)[0])
|
|
os.makedirs(workflow_presets_dir, exist_ok=True)
|
|
|
|
preset_path = os.path.join(workflow_presets_dir, f"{preset_name}.json")
|
|
|
|
try:
|
|
with open(preset_path, 'w', encoding='utf-8') as f:
|
|
json.dump(form_data, f, indent=4)
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return web.json_response({"error": f"Error saving preset file: {e}"}, status=500)
|
|
|
|
async def delete_preset(request: web.Request) -> web.Response:
|
|
data = await request.json()
|
|
workflow_name = data.get('workflow')
|
|
preset_name = data.get('preset_name')
|
|
|
|
if not all([workflow_name, preset_name]):
|
|
return web.json_response({"error": "Missing required fields"}, status=400)
|
|
|
|
presets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "presets")
|
|
preset_path = os.path.join(presets_dir, os.path.splitext(workflow_name)[0], preset_name)
|
|
|
|
if not os.path.exists(preset_path):
|
|
return web.json_response({"error": "Preset not found"}, status=404)
|
|
|
|
try:
|
|
os.remove(preset_path)
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return web.json_response({"error": f"Error deleting preset: {e}"}, status=500)
|
|
|
|
routes = [
|
|
web.get('/cozygen/hello', get_hello),
|
|
web.get('/cozygen/gallery', get_gallery_files),
|
|
web.post('/cozygen/toggle_folder_visibility', toggle_folder_visibility),
|
|
web.post('/cozygen/upload_image', upload_image),
|
|
web.get('/cozygen/workflows', get_workflow_list),
|
|
web.get('/cozygen/workflows/{filename}', get_workflow_file),
|
|
web.get('/cozygen/get_choices', get_choices),
|
|
web.post('/cozygen/submit', submit_to_queue),
|
|
web.get('/cozygen/queue', get_queue_status),
|
|
web.get('/cozygen/presets', get_presets),
|
|
web.get('/cozygen/presets/{workflow_name}/{preset_name}', get_preset_file),
|
|
web.post('/cozygen/presets', save_preset),
|
|
web.delete('/cozygen/presets', delete_preset),
|
|
]
|
|
|
|
# Register the startup function correctly
|
|
server.PromptServer.instance.app.on_startup.append(cozygen_startup)
|