Create comfy ui loader
This commit is contained in:
@@ -0,0 +1,386 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
DaSiWa I2V/FLF2V API Server для ComfyUI.
|
||||
Работает рядом с ComfyUI на той же машине.
|
||||
|
||||
Принимает HTTP запросы с HMAC авторизацией,
|
||||
отправляет workflow в ComfyUI, возвращает видео.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import uuid
|
||||
import time
|
||||
import shutil
|
||||
import base64
|
||||
import random
|
||||
import logging
|
||||
import binascii
|
||||
import subprocess
|
||||
import urllib.request
|
||||
import urllib.parse
|
||||
import websocket as ws_client
|
||||
|
||||
from flask import Flask, request, jsonify
|
||||
from hmac_auth import verify_request
|
||||
|
||||
# ============================================================================
|
||||
# Конфигурация
|
||||
# ============================================================================
|
||||
|
||||
COMFY_HOST = os.getenv("COMFY_HOST", "127.0.0.1")
|
||||
COMFY_PORT = os.getenv("COMFY_PORT", "8188")
|
||||
API_PORT = int(os.getenv("API_PORT", "8080"))
|
||||
|
||||
WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow_api.json")
|
||||
KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "keys.json")
|
||||
|
||||
COMFY_OUTPUT_DIR = os.getenv("COMFY_OUTPUT_DIR", "/ComfyUI/output")
|
||||
|
||||
# ============================================================================
|
||||
# Инициализация
|
||||
# ============================================================================
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s"
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# Загрузка ключей
|
||||
if not os.path.exists(KEYS_FILE):
|
||||
logger.error(f"❌ Файл ключей не найден: {KEYS_FILE}")
|
||||
logger.error(" Запусти: python generate_keys.py")
|
||||
sys.exit(1)
|
||||
|
||||
with open(KEYS_FILE, "r") as f:
|
||||
keys = json.load(f)
|
||||
CLIENT_ID = keys["client_id"]
|
||||
SECRET_KEY = keys["secret_key"]
|
||||
|
||||
logger.info(f"🔐 Ключи загружены. Client ID: {CLIENT_ID}")
|
||||
|
||||
# Множество использованных nonce (защита от replay-атак)
|
||||
used_nonces = set()
|
||||
|
||||
# WebSocket client ID
|
||||
ws_client_id = str(uuid.uuid4())
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Утилиты
|
||||
# ============================================================================
|
||||
|
||||
def to_nearest_multiple_of_16(value):
|
||||
"""Округляет значение до ближайшего кратного 16."""
|
||||
try:
|
||||
numeric_value = float(value)
|
||||
except Exception:
|
||||
raise ValueError(f"width/height value is not a number: {value}")
|
||||
adjusted = int(round(numeric_value / 16.0) * 16)
|
||||
return max(adjusted, 16)
|
||||
|
||||
|
||||
def save_base64_to_file(base64_data, temp_dir, filename):
|
||||
"""Сохраняет base64 данные в файл."""
|
||||
decoded = base64.b64decode(base64_data)
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
file_path = os.path.abspath(os.path.join(temp_dir, filename))
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(decoded)
|
||||
return file_path
|
||||
|
||||
|
||||
def download_file(url, temp_dir, filename):
|
||||
"""Скачивает файл по URL."""
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
file_path = os.path.abspath(os.path.join(temp_dir, filename))
|
||||
result = subprocess.run(
|
||||
["wget", "-O", file_path, "--no-verbose", url],
|
||||
capture_output=True, text=True, timeout=120
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Download failed: {result.stderr}")
|
||||
return file_path
|
||||
|
||||
|
||||
def process_image_input(job_input, prefix, temp_dir):
|
||||
"""
|
||||
Обрабатывает входные данные изображения.
|
||||
prefix: "image" или "last_image"
|
||||
Возвращает (file_path, True) или (None, False)
|
||||
"""
|
||||
path_key = f"{prefix}_path"
|
||||
url_key = f"{prefix}_url"
|
||||
b64_key = f"{prefix}_base64"
|
||||
|
||||
if path_key in job_input and job_input[path_key]:
|
||||
return job_input[path_key], True
|
||||
elif url_key in job_input and job_input[url_key]:
|
||||
return download_file(job_input[url_key], temp_dir, f"{prefix}.png"), True
|
||||
elif b64_key in job_input and job_input[b64_key]:
|
||||
return save_base64_to_file(job_input[b64_key], temp_dir, f"{prefix}.png"), True
|
||||
|
||||
return None, False
|
||||
|
||||
|
||||
def queue_prompt(prompt):
|
||||
"""Отправляет prompt в ComfyUI."""
|
||||
url = f"http://{COMFY_HOST}:{COMFY_PORT}/prompt"
|
||||
data = json.dumps({"prompt": prompt, "client_id": ws_client_id}).encode("utf-8")
|
||||
req = urllib.request.Request(url, data=data)
|
||||
return json.loads(urllib.request.urlopen(req).read())
|
||||
|
||||
|
||||
def get_history(prompt_id):
|
||||
"""Получает историю выполнения prompt."""
|
||||
url = f"http://{COMFY_HOST}:{COMFY_PORT}/history/{prompt_id}"
|
||||
with urllib.request.urlopen(url) as response:
|
||||
return json.loads(response.read())
|
||||
|
||||
|
||||
def generate_video(prompt):
|
||||
"""Подключается к ComfyUI по WebSocket, запускает генерацию, ждёт результат."""
|
||||
ws_url = f"ws://{COMFY_HOST}:{COMFY_PORT}/ws?clientId={ws_client_id}"
|
||||
|
||||
ws = ws_client.WebSocket()
|
||||
ws.connect(ws_url)
|
||||
logger.info("🔌 WebSocket подключён к ComfyUI")
|
||||
|
||||
prompt_id = queue_prompt(prompt)["prompt_id"]
|
||||
logger.info(f"📤 Prompt отправлен: {prompt_id}")
|
||||
|
||||
# Ждём завершения
|
||||
while True:
|
||||
out = ws.recv()
|
||||
if isinstance(out, str):
|
||||
message = json.loads(out)
|
||||
if message["type"] == "executing":
|
||||
data = message["data"]
|
||||
if data["node"] is None and data["prompt_id"] == prompt_id:
|
||||
break
|
||||
|
||||
ws.close()
|
||||
logger.info("✅ Генерация завершена")
|
||||
|
||||
# Извлекаем видео
|
||||
history = get_history(prompt_id)[prompt_id]
|
||||
for node_id in history["outputs"]:
|
||||
node_output = history["outputs"][node_id]
|
||||
if "gifs" in node_output:
|
||||
for video in node_output["gifs"]:
|
||||
video_path = video["fullpath"]
|
||||
with open(video_path, "rb") as f:
|
||||
video_b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
# Очистка
|
||||
try:
|
||||
os.remove(video_path)
|
||||
except OSError:
|
||||
pass
|
||||
return video_b64
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# API Endpoints
|
||||
# ============================================================================
|
||||
|
||||
@app.before_request
|
||||
def check_hmac_auth():
|
||||
"""Проверяет HMAC подпись для всех запросов кроме health check."""
|
||||
if request.path == "/health":
|
||||
return None
|
||||
|
||||
body = request.get_data()
|
||||
headers = {
|
||||
"X-Client-Id": request.headers.get("X-Client-Id", ""),
|
||||
"X-Timestamp": request.headers.get("X-Timestamp", ""),
|
||||
"X-Nonce": request.headers.get("X-Nonce", ""),
|
||||
"X-Signature": request.headers.get("X-Signature", ""),
|
||||
}
|
||||
|
||||
is_valid, error = verify_request(body, headers, SECRET_KEY, CLIENT_ID, used_nonces)
|
||||
if not is_valid:
|
||||
logger.warning(f"🚫 Auth failed: {error} from {request.remote_addr}")
|
||||
return jsonify({"error": "Unauthorized", "detail": error}), 401
|
||||
|
||||
|
||||
@app.route("/health", methods=["GET"])
|
||||
def health():
|
||||
"""Health check — без авторизации."""
|
||||
try:
|
||||
url = f"http://{COMFY_HOST}:{COMFY_PORT}/"
|
||||
urllib.request.urlopen(url, timeout=5)
|
||||
comfy_status = "ok"
|
||||
except Exception:
|
||||
comfy_status = "unavailable"
|
||||
|
||||
return jsonify({
|
||||
"status": "ok",
|
||||
"comfyui": comfy_status,
|
||||
"timestamp": int(time.time())
|
||||
})
|
||||
|
||||
|
||||
@app.route("/generate", methods=["POST"])
|
||||
def generate():
|
||||
"""Основной endpoint для генерации видео."""
|
||||
start_time = time.time()
|
||||
|
||||
job_input = request.json or {}
|
||||
|
||||
logger.info("=" * 60)
|
||||
logger.info("🎬 Новый запрос на генерацию")
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Логирование (без base64 данных)
|
||||
log_input = {k: v for k, v in job_input.items()
|
||||
if not k.endswith("_base64")}
|
||||
logger.info(f"Параметры: {json.dumps(log_input, ensure_ascii=False)}")
|
||||
|
||||
task_id = f"task_{uuid.uuid4().hex[:8]}"
|
||||
temp_dir = os.path.join("/tmp", task_id)
|
||||
|
||||
try:
|
||||
# === Обработка изображений ===
|
||||
image_path, has_image = process_image_input(job_input, "image", temp_dir)
|
||||
if not has_image:
|
||||
return jsonify({"error": "No input image provided. Use image_base64, image_url, or image_path"}), 400
|
||||
|
||||
last_image_path, use_flf2v = process_image_input(job_input, "last_image", temp_dir)
|
||||
|
||||
mode = "FLF2V" if use_flf2v else "I2V"
|
||||
logger.info(f"🎬 Режим: {mode}")
|
||||
|
||||
# === Загрузка workflow ===
|
||||
if not os.path.exists(WORKFLOW_FILE):
|
||||
return jsonify({"error": f"Workflow file not found: {WORKFLOW_FILE}"}), 500
|
||||
|
||||
with open(WORKFLOW_FILE, "r") as f:
|
||||
prompt = json.load(f)
|
||||
|
||||
# === Параметры генерации ===
|
||||
width = to_nearest_multiple_of_16(job_input.get("width", 528))
|
||||
height = to_nearest_multiple_of_16(job_input.get("height", 768))
|
||||
length = job_input.get("length", 81)
|
||||
steps = job_input.get("steps", 4)
|
||||
cfg = job_input.get("cfg", 1.0)
|
||||
seed = job_input.get("seed", -1)
|
||||
fps = job_input.get("fps", 16)
|
||||
sampler_name = job_input.get("sampler_name", "euler")
|
||||
scheduler = job_input.get("scheduler", "linear_quadratic")
|
||||
|
||||
if seed == -1:
|
||||
seed = random.randint(0, 2**63 - 1)
|
||||
|
||||
logger.info(f"📐 {width}x{height}, {length} frames, {steps} steps, CFG {cfg}, seed {seed}")
|
||||
|
||||
# === Заполнение workflow ===
|
||||
|
||||
# Positive prompt
|
||||
prompt["5"]["inputs"]["text"] = job_input.get("prompt", "")
|
||||
|
||||
# Negative prompt
|
||||
negative_prompt = job_input.get("negative_prompt", prompt["6"]["inputs"]["text"])
|
||||
prompt["6"]["inputs"]["text"] = negative_prompt
|
||||
|
||||
# First frame image
|
||||
prompt["7"]["inputs"]["image"] = image_path
|
||||
|
||||
# FLF2V / I2V mode
|
||||
if use_flf2v and last_image_path:
|
||||
prompt["15"]["inputs"]["image"] = last_image_path
|
||||
logger.info(f"🎬 FLF2V: last frame = {last_image_path}")
|
||||
else:
|
||||
prompt["8"]["class_type"] = "WanImageToVideo"
|
||||
if "end_image" in prompt["8"]["inputs"]:
|
||||
del prompt["8"]["inputs"]["end_image"]
|
||||
if "15" in prompt:
|
||||
del prompt["15"]
|
||||
logger.info("🎬 I2V: single image mode")
|
||||
|
||||
# Video dimensions
|
||||
prompt["8"]["inputs"]["width"] = width
|
||||
prompt["8"]["inputs"]["height"] = height
|
||||
prompt["8"]["inputs"]["length"] = length
|
||||
|
||||
# KSampler High
|
||||
prompt["11"]["inputs"]["noise_seed"] = seed
|
||||
prompt["11"]["inputs"]["steps"] = steps
|
||||
prompt["11"]["inputs"]["cfg"] = cfg
|
||||
prompt["11"]["inputs"]["sampler_name"] = sampler_name
|
||||
prompt["11"]["inputs"]["scheduler"] = scheduler
|
||||
prompt["11"]["inputs"]["end_at_step"] = steps // 2
|
||||
|
||||
# KSampler Low
|
||||
prompt["12"]["inputs"]["noise_seed"] = seed
|
||||
prompt["12"]["inputs"]["steps"] = steps
|
||||
prompt["12"]["inputs"]["cfg"] = cfg
|
||||
prompt["12"]["inputs"]["sampler_name"] = sampler_name
|
||||
prompt["12"]["inputs"]["scheduler"] = scheduler
|
||||
prompt["12"]["inputs"]["start_at_step"] = steps // 2
|
||||
|
||||
# Video output
|
||||
prompt["14"]["inputs"]["frame_rate"] = fps
|
||||
|
||||
# === Генерация ===
|
||||
video_b64 = generate_video(prompt)
|
||||
|
||||
if not video_b64:
|
||||
return jsonify({"error": "Video generation failed — no output"}), 500
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
logger.info(f"✅ Видео сгенерировано за {elapsed:.1f}s")
|
||||
|
||||
return jsonify({
|
||||
"video": video_b64,
|
||||
"seed": seed,
|
||||
"mode": mode,
|
||||
"elapsed": round(elapsed, 1)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Ошибка: {e}", exc_info=True)
|
||||
return jsonify({"error": str(e)}), 500
|
||||
|
||||
finally:
|
||||
# Очистка temp файлов
|
||||
if os.path.exists(temp_dir):
|
||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||
# Очистка output ComfyUI
|
||||
try:
|
||||
if os.path.exists(COMFY_OUTPUT_DIR):
|
||||
for fname in os.listdir(COMFY_OUTPUT_DIR):
|
||||
fpath = os.path.join(COMFY_OUTPUT_DIR, fname)
|
||||
if os.path.isfile(fpath):
|
||||
os.unlink(fpath)
|
||||
elif os.path.isdir(fpath):
|
||||
shutil.rmtree(fpath)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Запуск
|
||||
# ============================================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("=" * 60)
|
||||
logger.info("🚀 DaSiWa API Server")
|
||||
logger.info(f" ComfyUI: http://{COMFY_HOST}:{COMFY_PORT}")
|
||||
logger.info(f" API Port: {API_PORT}")
|
||||
logger.info(f" Workflow: {WORKFLOW_FILE}")
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Проверяем подключение к ComfyUI
|
||||
try:
|
||||
urllib.request.urlopen(f"http://{COMFY_HOST}:{COMFY_PORT}/", timeout=5)
|
||||
logger.info("✅ ComfyUI доступен")
|
||||
except Exception:
|
||||
logger.warning("⚠️ ComfyUI недоступен — запросы будут ждать")
|
||||
|
||||
app.run(host="0.0.0.0", port=API_PORT, debug=False)
|
||||
Reference in New Issue
Block a user