fix: 优化内容

2026-03-25 13:33:48 +08:00
parent f99098ec58
commit 8991f2a2d7
14 changed files with 1417 additions and 277 deletions
--- a/server/index.js
+++ b/server/index.js
@@ -2,14 +2,34 @@ import express from "express";
 import { spawn } from "node:child_process";
 import path from "node:path";
 import { fileURLToPath } from "node:url";
+import fs from "node:fs";
+import crypto from "node:crypto";

 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
 const repoRoot = path.resolve(__dirname, "..");
+const outputsDir = path.join(repoRoot, "outputs");
+fs.mkdirSync(outputsDir, { recursive: true });

 const app = express();
+app.use(express.json({ limit: "2mb" }));
+app.use(
+  "/api/static",
+  express.static(outputsDir, {
+    fallthrough: true,
+    setHeaders: (res) => {
+      // Important: avoid stale video preview.
+      res.setHeader("Cache-Control", "no-cache, no-transform");
+    },
+  })
+);
 app.use(express.static(path.join(__dirname, "public")));

+app.get("/api/health", (_req, res) => {
+  res.setHeader("Cache-Control", "no-cache");
+  res.status(200).json({ ok: true });
+});
+
 function sseHeaders(res) {
  res.setHeader("Content-Type", "text/event-stream; charset=utf-8");
  res.setHeader("Cache-Control", "no-cache, no-transform");
@@ -25,9 +45,46 @@ function sseSend(res, event, data) {
  res.write("\n");
 }

-app.get("/api/run", (req, res) => {
+function newTaskId() {
+  return crypto.randomUUID();
+}
+
+function taskDir(taskId) {
+  return path.join(outputsDir, taskId);
+}
+
+function ensureTaskDir(taskId) {
+  const dir = taskDir(taskId);
+  fs.mkdirSync(dir, { recursive: true });
+  return dir;
+}
+
+function spawnPythonStep({ step, prompt, configPath, mock, globalStyle, character, taskId, sceneIndex }) {
+  const py = process.env.PYTHON_BIN || "python3.10";
+  const args = [
+    "-m",
+    "engine.main",
+    "--prompt",
+    prompt,
+    "--config",
+    configPath,
+    "--step",
+    step,
+    "--task-id",
+    taskId,
+  ];
+  if (sceneIndex) args.push("--scene-index", String(sceneIndex));
+  if (globalStyle) args.push("--global-style", globalStyle);
+  if (character) args.push("--character", character);
+  if (mock) args.push("--mock");
+  return spawn(py, args, { cwd: repoRoot, env: process.env, stdio: ["pipe", "pipe", "pipe"] });
+}
+
+app.get("/api/script", (req, res) => {
  const prompt = String(req.query.prompt || "").trim();
  const mock = String(req.query.mock || "1") === "1";
+  const globalStyle = String(req.query.global_style || "").trim();
+  const character = String(req.query.character || "").trim();
  const configPath = String(req.query.config || "./configs/config.yaml");

  if (!prompt) {
@@ -35,25 +92,21 @@ app.get("/api/run", (req, res) => {
    return;
  }

+  const taskId = newTaskId();
+  ensureTaskDir(taskId);
+
  sseHeaders(res);
+  sseSend(res, "task", JSON.stringify({ task_id: taskId }));
  sseSend(res, "status", "starting");

-  // Unified in-container execution: Node spawns python directly.
-  const py = process.env.PYTHON_BIN || "python";
-  const args = [
-    path.join(repoRoot, "main.py"),
-    "--prompt",
+  const child = spawnPythonStep({
+    step: "script",
    prompt,
-    "--config",
    configPath,
-    "--script-only",
-  ];
-  if (mock) args.push("--mock");
-
-  const child = spawn(py, args, {
-    cwd: repoRoot,
-    env: process.env,
-    stdio: ["ignore", "pipe", "pipe"],
+    mock,
+    globalStyle,
+    character,
+    taskId,
  });

  let buf = "";
@@ -64,14 +117,15 @@ app.get("/api/run", (req, res) => {
    buf = parts.pop() || "";
    for (const line of parts) {
      if (!line) continue;
-      // Forward raw lines. Frontend will parse SCENE_JSON.
-      sseSend(res, "line", line);
+      if (line.startsWith("SCENE_JSON ")) sseSend(res, "scene", line.slice("SCENE_JSON ".length));
+      else if (line.startsWith("PROG ")) sseSend(res, "prog", line.slice("PROG ".length));
+      else sseSend(res, "line", line);
    }
  });

  child.stderr.setEncoding("utf8");
  child.stderr.on("data", (chunk) => {
-    sseSend(res, "stderr", chunk);
+    sseSend(res, "error", chunk);
  });

  req.on("close", () => {
@@ -80,13 +134,177 @@ app.get("/api/run", (req, res) => {

  child.on("exit", (code) => {
    if (buf.trim()) sseSend(res, "line", buf.trim());
-    sseSend(res, "done", String(code ?? 0));
+    sseSend(res, "done", String(code != null ? code : 0));
    res.end();
  });
 });

-const port = Number(process.env.PORT || 3000);
-app.listen(port, () => {
-  console.log(`[server] http://127.0.0.1:${port}`);
+app.post("/api/refine", (req, res) => {
+  const prompt = String((req.body && req.body.prompt) || "").trim();
+  const sceneIndex = Number((req.body && req.body.scene_index) || 1);
+  const scenes = req.body && req.body.scenes;
+  const scene = req.body && req.body.scene;
+  const mock = Boolean((req.body && req.body.mock) != null ? req.body.mock : true);
+  const globalStyle = String((req.body && req.body.global_style) || "").trim();
+  const character = String((req.body && req.body.character) || "").trim();
+  const configPath = String((req.body && req.body.config) || "./configs/config.yaml");
+  const taskId = String((req.body && req.body.task_id) || "").trim() || newTaskId();
+
+  if (!prompt) return res.status(400).json({ error: "missing prompt" });
+  if (!Number.isFinite(sceneIndex) || sceneIndex < 1) return res.status(400).json({ error: "bad scene_index" });
+  if (!Array.isArray(scenes) && (!scene || typeof scene !== "object")) {
+    return res.status(400).json({ error: "missing scene or scenes[]" });
+  }
+  ensureTaskDir(taskId);
+
+  const child = spawnPythonStep({
+    step: "refine",
+    prompt,
+    configPath,
+    mock,
+    globalStyle,
+    character,
+    taskId,
+    sceneIndex,
+  });
+  if (Array.isArray(scenes)) {
+    child.stdin.end(JSON.stringify({ scenes }));
+  } else {
+    child.stdin.end(JSON.stringify({ scene }));
+  }
+
+  let out = "";
+  let err = "";
+  child.stdout.setEncoding("utf8");
+  child.stderr.setEncoding("utf8");
+  child.stdout.on("data", (c) => (out += c));
+  child.stderr.on("data", (c) => (err += c));
+  child.on("exit", (code) => {
+    if (code !== 0) return res.status(500).json({ error: "python failed", stderr: err, stdout: out });
+    const line = out
+      .split(/\r?\n/)
+      .map((s) => s.trim())
+      .find((s) => s.startsWith("SCENE_JSON "));
+    if (!line) return res.status(500).json({ error: "no SCENE_JSON", stderr: err, stdout: out });
+    const payload = JSON.parse(line.slice("SCENE_JSON ".length));
+    return res.json({ task_id: taskId, scene: payload, stderr: err });
+  });
 });

+let isBusy = false;
+
+app.post("/api/render", (req, res) => {
+  const prompt = String((req.body && req.body.prompt) || "").trim();
+  const scenes = req.body && req.body.scenes;
+  const mock = Boolean((req.body && req.body.mock) != null ? req.body.mock : false);
+  const globalStyle = String((req.body && req.body.global_style) || "").trim();
+  const character = String((req.body && req.body.character) || "").trim();
+  const configPath = String((req.body && req.body.config) || "./configs/config.yaml");
+  const taskId = String((req.body && req.body.task_id) || "").trim() || newTaskId();
+
+  if (!prompt) return res.status(400).json({ error: "missing prompt" });
+  if (!Array.isArray(scenes)) return res.status(400).json({ error: "missing scenes[]" });
+  ensureTaskDir(taskId);
+
+  if (isBusy) {
+    return res.status(429).json({ error: "busy", msg: "GPU is busy, try later" });
+  }
+  isBusy = true;
+
+  sseHeaders(res);
+  sseSend(res, "task", JSON.stringify({ task_id: taskId }));
+  sseSend(res, "status", "render_start");
+
+  const child = spawnPythonStep({
+    step: "render",
+    prompt,
+    configPath,
+    mock,
+    globalStyle,
+    character,
+    taskId,
+  });
+  child.stdin.end(JSON.stringify({ scenes }));
+
+  let buf = "";
+  child.stdout.setEncoding("utf8");
+  child.stderr.setEncoding("utf8");
+
+  child.stdout.on("data", (chunk) => {
+    buf += chunk;
+    const parts = buf.split(/\r?\n/);
+    buf = parts.pop() || "";
+    for (const line of parts) {
+      if (!line) continue;
+      if (line.startsWith("PROG ")) sseSend(res, "prog", line.slice("PROG ".length));
+      else if (line.startsWith("RENDER_DONE ")) sseSend(res, "done", line.slice("RENDER_DONE ".length));
+      else sseSend(res, "line", line);
+    }
+  });
+
+  child.stderr.on("data", (chunk) => {
+    sseSend(res, "error", chunk);
+  });
+
+  req.on("close", () => {
+    child.kill("SIGTERM");
+  });
+
+  child.on("exit", (code) => {
+    isBusy = false;
+    if (buf.trim()) sseSend(res, "line", buf.trim());
+    if (code !== 0) sseSend(res, "error", `[ERROR] python exit_code=${code}`);
+    res.end();
+  });
+});
+
+async function runSelfCheck() {
+  const py = process.env.PYTHON_BIN || "python3.10";
+  const checks = [
+    { name: "check_comfy", args: ["scripts/check_comfy.py"] },
+    { name: "inspect_comfy_node", args: ["scripts/inspect_comfy_node.py"] },
+  ];
+  for (const c of checks) {
+    const deadline = Date.now() + 90_000;
+    let lastErr = "";
+    while (Date.now() < deadline) {
+      try {
+        await new Promise((resolve, reject) => {
+          const child = spawn(py, c.args, { cwd: repoRoot, env: process.env, stdio: ["ignore", "pipe", "pipe"] });
+          let out = "";
+          let err = "";
+          child.stdout.setEncoding("utf8");
+          child.stderr.setEncoding("utf8");
+          child.stdout.on("data", (d) => (out += d));
+          child.stderr.on("data", (d) => (err += d));
+          child.on("exit", (code) => {
+            if (code === 0) return resolve(true);
+            reject(new Error(`${c.name} failed (code=${code})\n${err || out}`));
+          });
+        });
+        lastErr = "";
+        break;
+      } catch (e) {
+        lastErr = String(e);
+        await new Promise((r) => setTimeout(r, 2000));
+      }
+    }
+    if (lastErr) {
+      throw new Error(lastErr);
+    }
+  }
+}
+
+const port = Number(process.env.PORT || 3000);
+(async () => {
+  try {
+    await runSelfCheck();
+    app.listen(port, () => {
+      console.log(`[server] http://127.0.0.1:${port}`);
+    });
+  } catch (e) {
+    console.error(String(e));
+    process.exit(1);
+  }
+})();
+