|
|
#!/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)
|
|
|
|
