134 lines
5.4 KiB
Python
134 lines
5.4 KiB
Python
import json
|
|
import logging
|
|
import requests
|
|
from flask import current_app
|
|
|
|
logger = logging.getLogger('gaze')
|
|
|
|
|
|
def get_loaded_checkpoint():
|
|
"""Return the checkpoint path currently loaded in ComfyUI, or None."""
|
|
try:
|
|
url = current_app.config.get('COMFYUI_URL', 'http://127.0.0.1:8188')
|
|
resp = requests.get(f'{url}/history', timeout=3)
|
|
if resp.ok:
|
|
history = resp.json()
|
|
if history:
|
|
latest = max(history.values(), key=lambda j: j.get('status', {}).get('status_str', ''))
|
|
nodes = latest.get('prompt', [None, None, {}])[2]
|
|
return nodes.get('4', {}).get('inputs', {}).get('ckpt_name')
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _ensure_checkpoint_loaded(checkpoint_path):
|
|
"""Check if the desired checkpoint is loaded in ComfyUI, and force reload if not."""
|
|
if not checkpoint_path:
|
|
return
|
|
|
|
try:
|
|
# Get currently loaded checkpoint from ComfyUI history
|
|
url = current_app.config.get('COMFYUI_URL', 'http://127.0.0.1:8188')
|
|
resp = requests.get(f'{url}/history', timeout=3)
|
|
if resp.ok:
|
|
history = resp.json()
|
|
if history:
|
|
latest = max(history.values(), key=lambda j: j.get('status', {}).get('status_str', ''))
|
|
nodes = latest.get('prompt', [None, None, {}])[2]
|
|
loaded_ckpt = nodes.get('4', {}).get('inputs', {}).get('ckpt_name')
|
|
|
|
# If the loaded checkpoint matches what we want, no action needed
|
|
if loaded_ckpt == checkpoint_path:
|
|
logger.info(f"Checkpoint {checkpoint_path} already loaded in ComfyUI")
|
|
return
|
|
|
|
# Checkpoint doesn't match or couldn't determine - force unload all models
|
|
logger.info(f"Forcing ComfyUI to unload models to ensure {checkpoint_path} loads")
|
|
requests.post(f'{url}/free', json={'unload_models': True}, timeout=5)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to check/force checkpoint reload: {e}")
|
|
|
|
|
|
def queue_prompt(prompt_workflow, client_id=None):
|
|
"""POST a workflow to ComfyUI's /prompt endpoint."""
|
|
# Ensure the checkpoint in the workflow is loaded in ComfyUI
|
|
checkpoint_path = prompt_workflow.get('4', {}).get('inputs', {}).get('ckpt_name')
|
|
_ensure_checkpoint_loaded(checkpoint_path)
|
|
|
|
p = {"prompt": prompt_workflow}
|
|
if client_id:
|
|
p["client_id"] = client_id
|
|
|
|
url = current_app.config['COMFYUI_URL']
|
|
|
|
# Log the full request being sent to ComfyUI
|
|
logger.debug("=" * 80)
|
|
logger.debug("COMFYUI REQUEST - Sending prompt to %s/prompt", url)
|
|
logger.debug("Checkpoint: %s", checkpoint_path)
|
|
logger.debug("Client ID: %s", client_id if client_id else "(none)")
|
|
logger.debug("Full workflow JSON:")
|
|
logger.debug(json.dumps(prompt_workflow, indent=2))
|
|
logger.debug("=" * 80)
|
|
|
|
data = json.dumps(p).encode('utf-8')
|
|
response = requests.post(f"{url}/prompt", data=data)
|
|
response_json = response.json()
|
|
|
|
# Log the response from ComfyUI
|
|
logger.debug("COMFYUI RESPONSE - Status: %s", response.status_code)
|
|
logger.debug("Response JSON: %s", json.dumps(response_json, indent=2))
|
|
if 'prompt_id' in response_json:
|
|
logger.info("ComfyUI accepted prompt with ID: %s", response_json['prompt_id'])
|
|
else:
|
|
logger.error("ComfyUI rejected prompt: %s", response_json)
|
|
logger.debug("=" * 80)
|
|
|
|
return response_json
|
|
|
|
|
|
def get_history(prompt_id):
|
|
"""Poll ComfyUI /history for results of a given prompt_id."""
|
|
url = current_app.config['COMFYUI_URL']
|
|
response = requests.get(f"{url}/history/{prompt_id}")
|
|
history_json = response.json()
|
|
|
|
# Log detailed history response for debugging
|
|
if prompt_id in history_json:
|
|
logger.debug("=" * 80)
|
|
logger.debug("COMFYUI HISTORY - Prompt ID: %s", prompt_id)
|
|
logger.debug("Status: %s", response.status_code)
|
|
|
|
# Extract key information from the history
|
|
prompt_data = history_json[prompt_id]
|
|
if 'status' in prompt_data:
|
|
logger.debug("Generation status: %s", prompt_data['status'])
|
|
|
|
if 'outputs' in prompt_data:
|
|
logger.debug("Outputs available: %s", list(prompt_data['outputs'].keys()))
|
|
for node_id, output in prompt_data['outputs'].items():
|
|
if 'images' in output:
|
|
logger.debug(" Node %s produced %d image(s)", node_id, len(output['images']))
|
|
for img in output['images']:
|
|
logger.debug(" - %s (subfolder: %s, type: %s)",
|
|
img.get('filename'), img.get('subfolder'), img.get('type'))
|
|
|
|
logger.debug("Full history response:")
|
|
logger.debug(json.dumps(history_json, indent=2))
|
|
logger.debug("=" * 80)
|
|
else:
|
|
logger.debug("History not yet available for prompt ID: %s", prompt_id)
|
|
|
|
return history_json
|
|
|
|
|
|
def get_image(filename, subfolder, folder_type):
|
|
"""Retrieve a generated image from ComfyUI's /view endpoint."""
|
|
url = current_app.config['COMFYUI_URL']
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
logger.debug("Fetching image from ComfyUI: filename=%s, subfolder=%s, type=%s",
|
|
filename, subfolder, folder_type)
|
|
response = requests.get(f"{url}/view", params=data)
|
|
logger.debug("Image retrieved: %d bytes (status: %s)", len(response.content), response.status_code)
|
|
return response.content
|