fix workflow

server.py

@@ -3,8 +3,10 @@
 DaSiWa I2V/FLF2V API Server для ComfyUI.
 Работает рядом с ComfyUI на той же машине.
 
-Принимает HTTP запросы с HMAC авторизацией,
-отправляет workflow в ComfyUI, возвращает видео.
+Асинхронный API (как RunPod):
+  POST /run        → {"id": "job_id", "status": "IN_QUEUE"}
+  GET  /status/ID  → {"id": ..., "status": "IN_QUEUE|IN_PROGRESS|COMPLETED|FAILED", "output": ...}
+  GET  /health     → {"status": "ok", "comfyui": "ok", "queue": 0}
 """
 
 import os
@@ -18,6 +20,8 @@ import random
 import logging
 import binascii
 import subprocess
+import threading
+import queue
 import urllib.request
 import urllib.parse
 import websocket as ws_client
@@ -69,6 +73,14 @@ used_nonces = set()
 # WebSocket client ID
 ws_client_id = str(uuid.uuid4())
 
+# ============================================================================
+# Job Queue (асинхронная очередь как в RunPod)
+# ============================================================================
+
+job_queue = queue.Queue()
+jobs = {}  # job_id -> {status, input, output, error, created_at, started_at, completed_at}
+jobs_lock = threading.Lock()
+
 
 # ============================================================================
 # Утилиты
@@ -185,6 +197,162 @@ def generate_video(prompt):
     return None
 
 
+# ============================================================================
+# Background Worker (обработка задач из очереди)
+# ============================================================================
+
+def build_prompt(job_input, image_path, last_image_path, use_flf2v):
+    """Загружает workflow и патчит параметрами задачи."""
+    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)
+
+    # Node 5: Positive prompt
+    prompt["5"]["inputs"]["text"] = job_input.get("prompt", "")
+
+    # Node 6: Negative prompt (use default or custom)
+    negative_prompt = job_input.get("negative_prompt", prompt["6"]["inputs"]["text"])
+    prompt["6"]["inputs"]["text"] = negative_prompt
+
+    # Node 7: Load first frame image
+    prompt["7"]["inputs"]["image"] = image_path
+
+    # Node 15: Load last frame image (for FLF2V 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:
+        # I2V mode: switch to WanImageToVideo, remove end_image
+        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")
+
+    # Node 8: WanFirstLastFrameToVideo / WanImageToVideo
+    prompt["8"]["inputs"]["width"] = width
+    prompt["8"]["inputs"]["height"] = height
+    prompt["8"]["inputs"]["length"] = length
+
+    # Node 11: 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
+
+    # Node 12: 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
+
+    # Node 14: Video output
+    prompt["14"]["inputs"]["frame_rate"] = fps
+
+    return prompt, seed, width, height
+
+
+def cleanup_comfy_output():
+    """Очистка 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
+
+
+def worker_loop():
+    """Фоновый воркер — берёт задачи из очереди и выполняет по одной."""
+    logger.info("⚙️ Worker thread started")
+    while True:
+        job_id = job_queue.get()  # блокируется пока нет задач
+        with jobs_lock:
+            job = jobs.get(job_id)
+        if not job:
+            continue
+
+        logger.info("=" * 60)
+        logger.info(f"🎬 Job {job_id}: Начинаем генерацию")
+        logger.info("=" * 60)
+
+        with jobs_lock:
+            job["status"] = "IN_PROGRESS"
+            job["started_at"] = time.time()
+
+        job_input = job["input"]
+        temp_dir = job["temp_dir"]
+
+        try:
+            # Обработка изображений
+            image_path, has_image = process_image_input(job_input, "image", temp_dir)
+            if not has_image:
+                raise ValueError("No input image provided")
+
+            last_image_path, use_flf2v = process_image_input(job_input, "last_image", temp_dir)
+            mode = "FLF2V" if use_flf2v else "I2V"
+            logger.info(f"🎬 Job {job_id}: Режим {mode}")
+
+            # Сборка промпта
+            prompt, seed, width, height = build_prompt(job_input, image_path, last_image_path, use_flf2v)
+            logger.info(f"📐 Job {job_id}: {width}x{height}, seed {seed}")
+
+            # Генерация
+            video_b64 = generate_video(prompt)
+
+            if not video_b64:
+                raise RuntimeError("Video generation failed — no output from ComfyUI")
+
+            elapsed = time.time() - job["started_at"]
+            logger.info(f"✅ Job {job_id}: Видео готово за {elapsed:.1f}s")
+
+            with jobs_lock:
+                job["status"] = "COMPLETED"
+                job["completed_at"] = time.time()
+                job["output"] = {
+                    "video": video_b64,
+                    "seed": seed,
+                    "mode": mode,
+                    "elapsed": round(elapsed, 1)
+                }
+
+        except Exception as e:
+            logger.error(f"❌ Job {job_id}: {e}", exc_info=True)
+            with jobs_lock:
+                job["status"] = "FAILED"
+                job["completed_at"] = time.time()
+                job["error"] = str(e)
+
+        finally:
+            # Очистка
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir, ignore_errors=True)
+            cleanup_comfy_output()
+
+        job_queue.task_done()
+
+
 # ============================================================================
 # API Endpoints
 # ============================================================================
@@ -222,146 +390,85 @@ def health():
     return jsonify({
         "status": "ok",
         "comfyui": comfy_status,
+        "queue": job_queue.qsize(),
         "timestamp": int(time.time())
     })
 
 
-@app.route("/generate", methods=["POST"])
-def generate():
-    """Основной endpoint для генерации видео."""
-    start_time = time.time()
-
+@app.route("/run", methods=["POST"])
+def run_job():
+    """Отправляет задачу в очередь. Возвращает job_id сразу."""
     job_input = request.json or {}
 
-    logger.info("=" * 60)
-    logger.info("🎬 Новый запрос на генерацию")
-    logger.info("=" * 60)
+    # Валидация: должно быть хотя бы одно изображение
+    has_image = any(k in job_input and job_input[k]
+                    for k in ("image_base64", "image_url", "image_path"))
+    if not has_image:
+        return jsonify({"error": "No input image. Use image_base64, image_url, or image_path"}), 400
 
-    # Логирование (без 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)}")
+    job_id = str(uuid.uuid4())
+    temp_dir = os.path.join("/tmp", f"job_{job_id[:8]}")
 
-    task_id = f"task_{uuid.uuid4().hex[:8]}"
-    temp_dir = os.path.join("/tmp", task_id)
+    # Логирование (без base64)
+    log_input = {k: (f"[{len(v)}chars]" if k.endswith("_base64") else v)
+                 for k, v in job_input.items()}
+    logger.info(f"📥 Job {job_id}: поставлен в очередь")
+    logger.info(f"   Параметры: {json.dumps(log_input, ensure_ascii=False)}")
 
-    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
+    with jobs_lock:
+        jobs[job_id] = {
+            "status": "IN_QUEUE",
+            "input": job_input,
+            "temp_dir": temp_dir,
+            "output": None,
+            "error": None,
+            "created_at": time.time(),
+            "started_at": None,
+            "completed_at": None,
+        }
 
-        last_image_path, use_flf2v = process_image_input(job_input, "last_image", temp_dir)
+    job_queue.put(job_id)
 
-        mode = "FLF2V" if use_flf2v else "I2V"
-        logger.info(f"🎬 Режим: {mode}")
+    return jsonify({
+        "id": job_id,
+        "status": "IN_QUEUE"
+    })
 
-        # === Загрузка 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)
+@app.route("/status/<job_id>", methods=["GET"])
+def job_status(job_id):
+    """Получить статус задачи. Когда COMPLETED — возвращает результат."""
+    with jobs_lock:
+        job = jobs.get(job_id)
 
-        # === Параметры генерации ===
-        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 not job:
+        return jsonify({"error": "Job not found"}), 404
 
-        if seed == -1:
-            seed = random.randint(0, 2**63 - 1)
+    response = {
+        "id": job_id,
+        "status": job["status"],
+    }
 
-        logger.info(f"📐 {width}x{height}, {length} frames, {steps} steps, CFG {cfg}, seed {seed}")
+    if job["status"] == "COMPLETED":
+        response["output"] = job["output"]
+    elif job["status"] == "FAILED":
+        response["error"] = job["error"]
 
-        # === Заполнение workflow ===
+    return jsonify(response)
 
-        # 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
+@app.route("/purge/<job_id>", methods=["POST"])
+def purge_job(job_id):
+    """Удалить завершённую задачу из памяти (освободить RAM от base64 видео)."""
+    with jobs_lock:
+        job = jobs.get(job_id)
+        if not job:
+            return jsonify({"error": "Job not found"}), 404
+        if job["status"] in ("IN_QUEUE", "IN_PROGRESS"):
+            return jsonify({"error": "Cannot purge active job"}), 400
+        del jobs[job_id]
 
-        # 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
+    return jsonify({"id": job_id, "purged": True})
 
 
 # ============================================================================
@@ -370,10 +477,15 @@ def generate():
 
 if __name__ == "__main__":
     logger.info("=" * 60)
-    logger.info("🚀 DaSiWa API Server")
+    logger.info("🚀 DaSiWa API Server (async worker mode)")
     logger.info(f"   ComfyUI:  http://{COMFY_HOST}:{COMFY_PORT}")
     logger.info(f"   API Port: {API_PORT}")
     logger.info(f"   Workflow: {WORKFLOW_FILE}")
+    logger.info("   Endpoints:")
+    logger.info("     POST /run         → поставить задачу")
+    logger.info("     GET  /status/<id> → статус / результат")
+    logger.info("     POST /purge/<id>  → удалить задачу из памяти")
+    logger.info("     GET  /health      → здоровье")
     logger.info("=" * 60)
 
     # Проверяем подключение к ComfyUI
@@ -381,6 +493,11 @@ if __name__ == "__main__":
         urllib.request.urlopen(f"http://{COMFY_HOST}:{COMFY_PORT}/", timeout=5)
         logger.info("✅ ComfyUI доступен")
     except Exception:
-        logger.warning("⚠️  ComfyUI недоступен — запросы будут ждать")
+        logger.warning("⚠️  ComfyUI недоступен — запросы будут ждать"
+                        " пока ComfyUI запустится")
+
+    # Запуск фонового воркера
+    worker_thread = threading.Thread(target=worker_loop, daemon=True)
+    worker_thread.start()
 
     app.run(host="0.0.0.0", port=API_PORT, debug=False)