MediaCrawler/mac-agent-os-main/05_tools/10_dashboard/routes/ave_scene.py

"""
routes/ave_scene.py — 场景生成器（分层字段 → 提示词 → 图像）

与角色生成器对称的完整能力：
  · GET  /scene-schema              字段定义（10 层）
  · GET  /scene-views               视角定义（5 个）
  · POST /scene/preview             提示词预览（**零费用**）
  · POST /scene/generate            生成（单张 / 套装；**需 confirm_cost**）
  · GET  /scene/jobs/{job_id}       进度
"""
from __future__ import annotations

import sys
import threading
import uuid
from pathlib import Path
from typing import Optional

from fastapi import APIRouter, HTTPException

from routes.ave_assets import _AVE_SCRIPTS, _store, logger

router = APIRouter()

_SCENE_JOBS: dict = {}


@router.get("/schema")
def api_scene_schema():
    """场景分层字段定义（10 层，供前端渲染表单）"""
    try:
        if str(_AVE_SCRIPTS) not in sys.path:
            sys.path.insert(0, str(_AVE_SCRIPTS))
        from scene_generator.prompt_assembler import (
            SCENE_LAYERS, SCENE_VIEWS, scene_layer_summary, parse_prompt_schema)
    except Exception as e:
        raise HTTPException(500, detail=f"场景模块加载失败: {e}")
    return {"status": "ok", "layers": SCENE_LAYERS, "views": SCENE_VIEWS,
            "summary": scene_layer_summary(), "parse_schema": parse_prompt_schema()}


@router.get("/presets")
def api_scene_presets():
    """场景预设（成套提示词，**零费用**）

    返回 `config/scene_presets.yaml` —— 前端可一键载入各机位提示词。
    """
    from pathlib import Path as _P
    import yaml as _yaml
    try:
        from routes.ave_assets import _AVE_CONFIG
        cands = [_P(_AVE_CONFIG) / "scene_presets.yaml"]
    except Exception:
        cands = []
    cands.append(_P(__file__).resolve().parents[2] / "09_ave" / "config" / "scene_presets.yaml")
    p = next((c for c in cands if c.exists()), None)
    if not p:
        return {"status": "ok", "presets": {}, "note": "预设文件不存在"}
    try:
        d = _yaml.safe_load(p.read_text(encoding="utf-8")) or {}
    except Exception as e:
        raise HTTPException(500, detail=f"预设解析失败: {e}")
    return {"status": "ok", "presets": d.get("presets") or {}, "source": str(p)}


@router.post("/preview")
def api_scene_preview(data: dict):
    """组装场景提示词（**纯计算，零费用**）

    Body: { layers: {...}, view?: "wide"|"front"|..., extra?: str, set?: bool }
    · set=true  → 返回 5 视角完整套装提示词
    """
    try:
        if str(_AVE_SCRIPTS) not in sys.path:
            sys.path.insert(0, str(_AVE_SCRIPTS))
        from scene_generator.prompt_assembler import (
            build_scene_prompt, build_scene_view_prompt, build_scene_set)
    except Exception as e:
        raise HTTPException(500, detail=f"场景模块加载失败: {e}")
    data = data or {}
    layers = data.get("layers") or {}
    prompts = data.get("prompts") or {}
    preset_id = (data.get("preset_id") or "").strip()

    # ★ 预设模式：从 config/scene_presets.yaml 取该预设的各机位成品提示词
    if preset_id and not prompts:
        try:
            from pathlib import Path as _P
            import yaml as _yaml
            from routes.ave_assets import _AVE_CONFIG
            pp = _P(_AVE_CONFIG) / "scene_presets.yaml"
            if pp.exists():
                _d = _yaml.safe_load(pp.read_text(encoding="utf-8")) or {}
                prompts = ((_d.get("presets") or {}).get(preset_id) or {}).get("slots") or {}
        except Exception as e:
            raise HTTPException(400, detail=f"预设读取失败: {e}")
    if not layers and not prompts:
        raise HTTPException(400, detail="需要 layers（场景字段）、prompts（各机位成品提示词）或 preset_id（预设）之一")
    extra = (data.get("extra") or "").strip()
    if data.get("set"):
        return {"status": "ok", "set": build_scene_set(layers, prompts=prompts, extra=extra)}
    view = (data.get("view") or "").strip()
    # ★ 优先用成品提示词（用户自带 / 预设）
    if prompts and str(prompts.get(view) or "").strip():
        prompt = str(prompts[view]).strip()
    elif view:
        prompt = build_scene_view_prompt(view, layers, extra=extra)
    else:
        prompt = build_scene_prompt(layers, extra=extra)
    return {"status": "ok", "view": view or "(主图)", "prompt": prompt}


@router.post("/generate")
def api_scene_generate(data: dict):
    """生成场景图（**真实扣费**，约 ¥0.2/张；套装 5 张 ≈ ¥1）

    Body: {
      layers: {...}, scene_name?: str, view?: str,      # 单张模式
      set?: bool,                                       # 套装模式（5 视角）
      anchor_image?: str,                               # 复用已有主图
      confirm_cost: true
    }
    """
    if not data.get("confirm_cost"):
        raise HTTPException(400, detail="⛔ 生成场景图会真实扣费（约 ¥0.2/张），需显式确认（confirm_cost=true）")
    layers = data.get("layers") or {}
    prompts = data.get("prompts") or {}
    preset_id = (data.get("preset_id") or "").strip()

    # ★ 预设模式：从 config/scene_presets.yaml 取该预设的各机位成品提示词
    if preset_id and not prompts:
        try:
            from pathlib import Path as _P
            import yaml as _yaml
            from routes.ave_assets import _AVE_CONFIG
            pp = _P(_AVE_CONFIG) / "scene_presets.yaml"
            if pp.exists():
                _d = _yaml.safe_load(pp.read_text(encoding="utf-8")) or {}
                prompts = ((_d.get("presets") or {}).get(preset_id) or {}).get("slots") or {}
        except Exception as e:
            raise HTTPException(400, detail=f"预设读取失败: {e}")
    if not layers and not prompts:
        raise HTTPException(400, detail="需要 layers（场景字段）、prompts（各机位成品提示词）或 preset_id（预设）之一")

    job_id = uuid.uuid4().hex[:12]
    _SCENE_JOBS[job_id] = {"job_id": job_id, "status": "running", "progress": 0,
                           "note": "准备中", "result": None, "error": None}

    def _run():
        def _cb(p, n):
            j = _SCENE_JOBS.get(job_id)
            if j:
                j["progress"] = p
                j["note"] = n
        try:
            if str(_AVE_SCRIPTS) not in sys.path:
                sys.path.insert(0, str(_AVE_SCRIPTS))
            from scene_generator.generator import build_scene_set_images, generate_scene_image
            from scene_generator.prompt_assembler import build_scene_prompt, build_scene_view_prompt

            name = (data.get("scene_name") or "").strip()
            extra = (data.get("extra") or "").strip()
            if data.get("set"):
                out = build_scene_set_images(
                    layers, scene_name=name,
                    anchor_image=(data.get("anchor_image") or "").strip(),
                    extra=extra, prompts=prompts, progress_cb=_cb)
            else:
                # 单张模式
                from character_sheet import CACHE_DIR
                view = (data.get("view") or "wide").strip()
                # ★ 优先用成品提示词（用户自带 / 预设）
                if prompts and str(prompts.get(view) or "").strip():
                    prompt = str(prompts[view]).strip()
                elif view != "wide":
                    prompt = build_scene_view_prompt(view, layers, extra=extra)
                else:
                    prompt = build_scene_prompt(layers, extra=extra)
                base = CACHE_DIR.parent / "scene_sheet" / (name or "scene")
                base.mkdir(parents=True, exist_ok=True)
                dst = base / f"{view}.png"
                _cb(30, "生成中")
                r = generate_scene_image(prompt, str(dst),
                                         ref_image=(data.get("anchor_image") or "").strip())
                out = {"views": {view: {"path": r.get("path", ""), "ok": r.get("ok"),
                                        "error": r.get("error", "")}},
                       "done": 1 if r.get("ok") else 0, "cost_estimate": 0.2,
                       "note": "单张场景图完成" if r.get("ok") else f"失败: {r.get('error','')}"}
                if r.get("ok") and name:
                    try:
                        from scene_generator.generator import _save_scene_to_assets
                        out["saved_to"] = _save_scene_to_assets(name, out, base)
                    except Exception as e:
                        out["save_error"] = str(e)[:200]
            _SCENE_JOBS[job_id].update({"status": "done", "progress": 100,
                                        "note": out.get("note", "完成"), "result": out})
        except Exception as e:
            logger.warning("场景生成失败 %s: %s", job_id, e)
            _SCENE_JOBS[job_id].update({"status": "failed", "error": str(e)[:400]})

    threading.Thread(target=_run, daemon=True).start()
    return {"status": "ok", "job_id": job_id,
            "message": "已开始生成。用 GET /api/ave/scene/jobs/{job_id} 查进度。"}


@router.get("/jobs/{job_id}")
def api_scene_job(job_id: str):
    j = _SCENE_JOBS.get(job_id)
    if not j:
        raise HTTPException(404, detail="任务不存在")
    return j