fix: 优化架构

engine/comfy_client.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import json
+import time
 import uuid
 from dataclasses import dataclass
 from pathlib import Path
@@ -186,3 +187,215 @@ class ComfyClient:
 
             # unreachable
             # return ComfyResult(prompt_id=prompt_id, output_files=last_files)
+
+
+# ---------------------------------------------------------------------------
+# Minimal "text->image" helpers (used by shot rendering)
+# ---------------------------------------------------------------------------
+
+
+def _build_simple_workflow(
+    prompt_text: str,
+    *,
+    seed: int,
+    ckpt_name: str,
+    width: int,
+    height: int,
+    steps: int = 20,
+    cfg: float = 8.0,
+    sampler_name: str = "euler",
+    scheduler: str = "normal",
+    denoise: float = 1.0,
+    filename_prefix: str = "shot",
+    negative_text: str = "low quality, blurry",
+) -> dict[str, Any]:
+    # Best-effort workflow. If your ComfyUI nodes/models differ, generation must fallback.
+    return {
+        "3": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": int(seed),
+                "steps": int(steps),
+                "cfg": float(cfg),
+                "sampler_name": sampler_name,
+                "scheduler": scheduler,
+                "denoise": float(denoise),
+                "model": ["4", 0],
+                "positive": ["6", 0],
+                "negative": ["7", 0],
+                "latent_image": ["5", 0],
+            },
+        },
+        "4": {
+            "class_type": "CheckpointLoaderSimple",
+            "inputs": {
+                "ckpt_name": ckpt_name,
+            },
+        },
+        "5": {
+            "class_type": "EmptyLatentImage",
+            "inputs": {
+                "width": int(width),
+                "height": int(height),
+                "batch_size": 1,
+            },
+        },
+        "6": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {
+                "text": prompt_text,
+                "clip": ["4", 1],
+            },
+        },
+        "7": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {
+                "text": negative_text,
+                "clip": ["4", 1],
+            },
+        },
+        "8": {
+            "class_type": "VAEDecode",
+            "inputs": {
+                "samples": ["3", 0],
+                "vae": ["4", 2],
+            },
+        },
+        "9": {
+            "class_type": "SaveImage",
+            "inputs": {
+                "images": ["8", 0],
+                "filename_prefix": filename_prefix,
+            },
+        },
+    }
+
+
+def _queue_prompt(base_url: str, workflow: dict[str, Any], client_id: str) -> str:
+    r = httpx.post(
+        base_url.rstrip("/") + "/prompt",
+        json={"prompt": workflow, "client_id": client_id},
+        timeout=30.0,
+    )
+    r.raise_for_status()
+    data = r.json()
+    pid = data.get("prompt_id")
+    if not isinstance(pid, str) or not pid:
+        raise RuntimeError(f"Unexpected /prompt response: {data}")
+    return pid
+
+
+def _get_history_item(base_url: str, prompt_id: str) -> dict[str, Any] | None:
+    for url in (f"{base_url.rstrip('/')}/history/{prompt_id}", f"{base_url.rstrip('/')}/history"):
+        try:
+            r = httpx.get(url, timeout=30.0)
+            if r.status_code == 404:
+                continue
+            r.raise_for_status()
+            data = r.json()
+            if isinstance(data, dict):
+                if prompt_id in data and isinstance(data[prompt_id], dict):
+                    return data[prompt_id]
+                if url.endswith(f"/{prompt_id}") and isinstance(data, dict):
+                    return data
+            return None
+        except Exception:
+            continue
+    return None
+
+
+def _extract_first_image_view_target(history_item: dict[str, Any]) -> tuple[str, str] | None:
+    outputs = history_item.get("outputs")
+    if not isinstance(outputs, dict):
+        return None
+
+    def walk(v: Any) -> list[dict[str, Any]]:
+        found: list[dict[str, Any]] = []
+        if isinstance(v, dict):
+            if isinstance(v.get("filename"), str) and v.get("filename").strip():
+                found.append(v)
+            for vv in v.values():
+                found.extend(walk(vv))
+        elif isinstance(v, list):
+            for vv in v:
+                found.extend(walk(vv))
+        return found
+
+    candidates = walk(outputs)
+    for c in candidates:
+        fn = str(c.get("filename", "")).strip()
+        sf = str(c.get("subfolder", "") or "").strip()
+        if fn:
+            return fn, sf
+    return None
+
+
+def generate_image(
+    prompt_text: str,
+    output_dir: str | Path,
+    *,
+    cfg: AppConfig | None = None,
+    timeout_s: int = 60,
+    retry: int = 2,
+    width: int | None = None,
+    height: int | None = None,
+    filename_prefix: str = "shot",
+    ckpt_candidates: list[str] | None = None,
+    negative_text: str | None = None,
+) -> Path:
+    cfg2 = cfg or AppConfig.load("./configs/config.yaml")
+    base_url = str(cfg2.get("app.comfy_base_url", "http://comfyui:8188")).rstrip("/")
+
+    out_dir = Path(output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    if width is None or height is None:
+        mock_size = cfg2.get("video.mock_size", [1024, 576])
+        width = int(width or mock_size[0])
+        height = int(height or mock_size[1])
+
+    if negative_text is None:
+        negative_text = "low quality, blurry"
+
+    if ckpt_candidates is None:
+        ckpt_candidates = [
+            "v1-5-pruned-emaonly.ckpt",
+            "v1-5-pruned-emaonly.safetensors",
+            "sd-v1-5-tiny.safetensors",
+        ]
+
+    last_err: Exception | None = None
+    for _attempt in range(max(1, retry)):
+        for ckpt_name in ckpt_candidates:
+            client_id = str(uuid.uuid4())
+            seed = int(uuid.uuid4().int % 2_147_483_647)
+            workflow = _build_simple_workflow(
+                prompt_text,
+                seed=seed,
+                ckpt_name=ckpt_name,
+                width=width,
+                height=height,
+                filename_prefix=filename_prefix,
+                    negative_text=negative_text,
+            )
+            try:
+                prompt_id = _queue_prompt(base_url, workflow, client_id)
+                start = time.time()
+                while time.time() - start < timeout_s:
+                    item = _get_history_item(base_url, prompt_id)
+                    if isinstance(item, dict):
+                        img_target = _extract_first_image_view_target(item)
+                        if img_target:
+                            filename, subfolder = img_target
+                            view_url = f"{base_url}/view?filename={filename}&subfolder={subfolder}"
+                            img_resp = httpx.get(view_url, timeout=60.0)
+                            img_resp.raise_for_status()
+                            image_path = out_dir / filename
+                            image_path.write_bytes(img_resp.content)
+                            return image_path
+                    time.sleep(1.0)
+            except Exception as e:
+                last_err = e
+                continue
+
+    raise RuntimeError(f"ComfyUI image generation failed after retries: {last_err}")