from __future__ import annotations import asyncio import json import os import shutil import subprocess import uuid from contextlib import asynccontextmanager from pathlib import Path from typing import Literal from dotenv import load_dotenv from fastapi import BackgroundTasks, FastAPI, File, HTTPException, UploadFile from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse from pydantic import BaseModel, Field load_dotenv() JOBS_DIR = Path(os.getenv("JOBS_DIR", "./jobs")).resolve() JOBS_DIR.mkdir(parents=True, exist_ok=True) CORS_ORIGINS = [o.strip() for o in os.getenv("CORS_ORIGINS", "http://localhost:4290").split(",") if o.strip()] LLM_BASE_URL = os.getenv("LLM_BASE_URL", "").strip() LLM_API_KEY = os.getenv("LLM_API_KEY", "").strip() ASR_MODEL = os.getenv("ASR_MODEL", "whisper-1") TRANSLATE_MODEL = os.getenv("TRANSLATE_MODEL", "gemini-2.5-flash") REWRITE_MODEL = os.getenv("REWRITE_MODEL", "gemini-2.5-pro") VISION_MODEL = os.getenv("VISION_MODEL", "gemini-2.5-flash") IMAGE_MODEL = os.getenv("IMAGE_MODEL", "gemini-3-pro-image-preview") # OpenAI 客户端(OpenAI 兼容网关,含 SKG ezlink) from openai import OpenAI _llm_client: OpenAI | None = None def llm() -> OpenAI: global _llm_client if _llm_client is None: if not LLM_API_KEY: raise RuntimeError("LLM_API_KEY 未配置") _llm_client = OpenAI(base_url=LLM_BASE_URL or None, api_key=LLM_API_KEY) return _llm_client # Pipeline 状态: # created → downloading → downloaded(停,等用户点解析)→ splitting → frames_extracted # → transcribing → transcribed | failed JobStatus = Literal[ "created", "downloading", "downloaded", "splitting", "frames_extracted", "transcribing", "transcribed", "failed", ] KEYFRAME_COUNT = int(os.getenv("KEYFRAME_COUNT", "5")) class GeneratedImage(BaseModel): id: str # uuid hex 12 prompt: str model: str mode: str = "edit" # "edit"(带参考图) | "text"(纯文字) url: str # /jobs/{job_id}/frames/{idx}/gen/{id}.jpg selected: bool = False created_at: float = 0.0 class KeyElement(BaseModel): """关键帧里识别 / 用户提取的元素,可单独抠图给下游做"二创素材层" """ id: str # uuid hex 8 name_zh: str name_en: str = "" position: str = "" source: Literal["auto", "manual", "region"] = "manual" region: dict | None = None # 抠图(注:Gemini 输出 JPEG 不支持真透明,所以让模型在纯白 / 纯黑底上输出) cutout_id: str | None = None # 已抠图 → /jobs/{id}/frames/{idx}/elements/{element_id}/cutout.jpg cutout_background: Literal["white", "black"] = "white" created_at: float = 0.0 class KeyFrame(BaseModel): index: int timestamp: float url: str description: dict | None = None # vision 模型识别结果 {scene, objects, style, suggested_prompt} cleaned_url: str | None = None # 清洗后干净版(待应用)→ /jobs/{id}/frames/{idx}/cleaned.jpg cleaned_applied: bool = False # 是否已用清洗版替换原图(替换后 cleaned_url=null) elements: list[KeyElement] = [] # 提取的元素清单(持久化) generated_images: list[GeneratedImage] = [] class TranscriptSegment(BaseModel): index: int start: float end: float en: str zh: str = "" class Job(BaseModel): id: str url: str status: JobStatus = "created" progress: int = 0 message: str = "" video_url: str = "" duration: float = 0.0 width: int = 0 height: int = 0 frames: list[KeyFrame] = Field(default_factory=list) transcript: list[TranscriptSegment] = Field(default_factory=list) error: str = "" JOBS: dict[str, Job] = {} def job_dir(job_id: str) -> Path: d = JOBS_DIR / job_id d.mkdir(parents=True, exist_ok=True) return d def save_state(job: Job) -> None: (job_dir(job.id) / "state.json").write_text(job.model_dump_json(indent=2)) def update(job: Job, **kw) -> None: for k, v in kw.items(): setattr(job, k, v) save_state(job) @asynccontextmanager async def lifespan(_: FastAPI): # 启动时从磁盘恢复 jobs(简化版:只列目录) for p in JOBS_DIR.iterdir(): if p.is_dir() and (p / "state.json").exists(): try: JOBS[p.name] = Job.model_validate_json((p / "state.json").read_text()) except Exception: pass yield app = FastAPI(title="SKG TK 二创 API", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=CORS_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ---------- Pipeline 实现 ---------- def run(cmd: list[str], cwd: Path | None = None) -> str: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) if res.returncode != 0: # ffmpeg 把 banner 写 stderr,挑最后几行(真错误一般在末尾) tail = "\n".join(res.stderr.splitlines()[-12:]) or res.stderr[-800:] raise RuntimeError(f"cmd failed: {' '.join(cmd[:3])}... · {tail}") return res.stdout # ---- 启发式选帧工具 ---- import imagehash import numpy as np from PIL import Image def _sharpness(img_path: Path) -> float: """Laplacian variance:值越大越清晰,模糊/转场帧值低。""" g = np.asarray(Image.open(img_path).convert("L").resize((320, 180)), dtype=np.float32) lap = (-4 * g[1:-1, 1:-1] + g[:-2, 1:-1] + g[2:, 1:-1] + g[1:-1, :-2] + g[1:-1, 2:]) return float(lap.var()) def _select_keyframes(candidates: list[Path], n: int, dup_threshold: int = 8) -> list[Path]: """ candidates: 按时间排序的候选帧路径 n: 目标帧数 dup_threshold: pHash 汉明距离 < 此值视为相似(默认 8,64bit hash 大致 ~12.5% 像素差) """ if len(candidates) <= n: return candidates # 算 pHash + sharpness items = [] for i, p in enumerate(candidates): try: img = Image.open(p) h = imagehash.phash(img) s = _sharpness(p) items.append({"path": p, "idx": i, "hash": h, "sharp": s}) except Exception: continue # 去重:相似帧保留 sharpness 高的 deduped: list[dict] = [] for it in items: dup = None for kept in deduped: if (it["hash"] - kept["hash"]) < dup_threshold: dup = kept break if dup is None: deduped.append(it) elif it["sharp"] > dup["sharp"]: deduped[deduped.index(dup)] = it # 时序分桶:把候选时间轴等分 n 段,每段取去重后 sharpness 最高的 total = len(candidates) buckets: list[list[dict]] = [[] for _ in range(n)] for it in deduped: b = min(int(it["idx"] * n / total), n - 1) buckets[b].append(it) selected: list[dict] = [] for b in buckets: if b: selected.append(max(b, key=lambda x: x["sharp"])) # 空桶补足:从未选的 deduped 里按 sharpness 排序补 chosen_paths = {it["path"] for it in selected} remaining = sorted([it for it in deduped if it["path"] not in chosen_paths], key=lambda x: -x["sharp"]) while len(selected) < n and remaining: selected.append(remaining.pop(0)) # 按时间排序输出 selected.sort(key=lambda x: x["idx"]) return [it["path"] for it in selected] def ffprobe_meta(mp4: Path) -> dict: out = run([ "ffprobe", "-v", "error", "-print_format", "json", "-show_streams", "-show_format", str(mp4), ]) return json.loads(out) async def pipeline_download(job_id: str) -> None: """阶段 1:仅下载(或上传跳过),落 source.mp4,停在 downloaded 等用户点解析。""" job = JOBS[job_id] d = job_dir(job_id) try: mp4 = d / "source.mp4" if mp4.exists(): update(job, status="downloading", message="本地上传 · 跳过下载", progress=15) else: update(job, status="downloading", message="yt-dlp 下载中…", progress=5) run([ "yt-dlp", "-f", "best[ext=mp4]/best", "-o", str(mp4), "--no-warnings", "--no-playlist", "--retries", "3", job.url, ]) if not mp4.exists(): raise RuntimeError("下载完成但找不到 source.mp4") meta = ffprobe_meta(mp4) v_stream = next((s for s in meta["streams"] if s["codec_type"] == "video"), None) duration = float(meta["format"]["duration"]) update( job, status="downloaded", video_url=f"/jobs/{job_id}/video.mp4", duration=duration, width=int(v_stream["width"]) if v_stream else 0, height=int(v_stream["height"]) if v_stream else 0, progress=25, message=f"视频就绪 · {duration:.1f}s · 等待解析", ) except Exception as e: update(job, status="failed", error=str(e), message="下载失败") async def pipeline_analyze(job_id: str, frame_count: int = KEYFRAME_COUNT) -> None: """阶段 2:拆音轨 + 抽关键帧 + ASR + 翻译。需要 source.mp4 已存在。""" job = JOBS[job_id] d = job_dir(job_id) try: mp4 = d / "source.mp4" if not mp4.exists(): raise RuntimeError("source.mp4 不存在,先完成下载") update(job, status="splitting", message="ffmpeg 拆分音轨…", progress=35) wav = d / "audio.wav" run([ "ffmpeg", "-y", "-i", str(mp4), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(wav), ]) n = max(1, min(int(frame_count), 20)) # 候选数:n 的 6 倍或至少 24,封顶 60 candidate_count = max(24, min(60, n * 6)) update(job, message=f"抽取候选 {candidate_count} 张…", progress=45) frames_dir = d / "frames" if frames_dir.exists(): shutil.rmtree(frames_dir) frames_dir.mkdir(parents=True) cand_dir = d / "candidates" if cand_dir.exists(): shutil.rmtree(cand_dir) cand_dir.mkdir(parents=True) # 1) 均匀采样大批候选(fast seek,每张 < 0.5s) duration = max(float(job.duration or 1.0), 0.1) step = duration / (candidate_count + 1) candidate_meta: list[tuple[Path, float]] = [] # (path, timestamp) for i in range(candidate_count): t = step * (i + 1) out = cand_dir / f"c_{i:03d}.jpg" run([ "ffmpeg", "-y", "-ss", str(t), "-i", str(mp4), "-frames:v", "1", "-pix_fmt", "yuvj420p", "-q:v", "3", str(out), ]) if out.exists(): candidate_meta.append((out, t)) # 2) D 启发式选 n 张:pHash 去重 + Laplacian 清晰度 + 时序分桶 update(job, message=f"启发式筛选 {n} / {len(candidate_meta)} 张…", progress=60) cand_paths = [m[0] for m in candidate_meta] ts_by_path = {m[0]: m[1] for m in candidate_meta} chosen = _select_keyframes(cand_paths, n) # 3) 落盘到 frames/.jpg renamed: list[KeyFrame] = [] chosen_sorted = sorted(chosen, key=lambda p: ts_by_path[p]) for i, src in enumerate(chosen_sorted): dst = frames_dir / f"{i:03d}.jpg" shutil.copyfile(src, dst) renamed.append(KeyFrame( index=i, timestamp=round(ts_by_path[src], 2), url=f"/jobs/{job_id}/frames/{i}.jpg", )) # 4) 清理候选目录 shutil.rmtree(cand_dir, ignore_errors=True) update( job, status="frames_extracted", frames=renamed, progress=70, message=f"已抽取 {len(renamed)} 张关键帧", ) # 自动接 ASR + 翻译 await pipeline_transcribe(job_id) except Exception as e: update(job, status="failed", error=str(e), message="解析失败") # ---------- Gemini ASR + 翻译 ---------- def _transcribe_sync(wav: Path) -> list[dict]: """whisper-1 verbose_json → segments[{start, end, text}]""" with wav.open("rb") as f: resp = llm().audio.transcriptions.create( file=(wav.name, f, "audio/wav"), model=ASR_MODEL, response_format="verbose_json", timestamp_granularities=["segment"], ) raw = resp.model_dump() if hasattr(resp, "model_dump") else resp segments = raw.get("segments") or [] # 兜底:网关如果不返回 segments,把全文当一段 if not segments and raw.get("text"): segments = [{"start": 0.0, "end": float(raw.get("duration", 0) or 0), "text": raw["text"]}] return segments def _translate_sync(segments: list[dict]) -> list[str]: """gemini-2.5-flash 批量翻译为中文,按段返回""" payload = [{"i": i, "en": s.get("text", "").strip()} for i, s in enumerate(segments)] prompt = ( "你是字幕翻译。把下列英文字幕段翻译为简体中文,保持原意、口语化、自然流畅。" "严格返回 JSON 数组,不要任何 markdown 或多余文字,schema: " '[{"i": 0, "zh": "..."}, ...]\n\n输入:\n' + json.dumps(payload, ensure_ascii=False) ) resp = llm().chat.completions.create( model=TRANSLATE_MODEL, messages=[{"role": "user", "content": prompt}], response_format={"type": "json_object"}, temperature=0.2, ) content = resp.choices[0].message.content or "[]" try: data = json.loads(content) if isinstance(data, dict): for k in ("data", "items", "result", "translations"): if k in data and isinstance(data[k], list): data = data[k] break if not isinstance(data, list): data = [] except json.JSONDecodeError: data = [] zh_by_idx: dict[int, str] = {} for it in data: if isinstance(it, dict) and "i" in it: zh_by_idx[int(it["i"])] = str(it.get("zh", "")) return [zh_by_idx.get(i, "") for i in range(len(segments))] async def pipeline_transcribe(job_id: str) -> None: job = JOBS[job_id] d = job_dir(job_id) wav = d / "audio.wav" try: if not wav.exists(): raise RuntimeError("audio.wav 不存在") if not LLM_API_KEY: # 无 key 模式:mock 数据 update(job, status="transcribing", message="ASR (mock) …", progress=75) await asyncio.sleep(1.0) mock = [ TranscriptSegment(index=0, start=0.0, end=3.5, en="Welcome back, today we're testing something new.", zh="欢迎回来,今天我们要测试一些新东西。"), TranscriptSegment(index=1, start=3.5, end=7.2, en="This device looks really sleek and minimal.", zh="这个设备看起来非常时尚和简约。"), ] update(job, transcript=mock, status="transcribed", progress=100, message="转录完成(MOCK · 未设 LLM_API_KEY)") return # 1) whisper ASR update(job, status="transcribing", message=f"{ASR_MODEL} 转录中…", progress=78) segments = await asyncio.to_thread(_transcribe_sync, wav) if not segments: raise RuntimeError("ASR 返回 0 段(可能无人声 / 格式问题)") # 先把英文段落落到 job 上(让 UI 提前看到,翻译再补 zh) en_only = [ TranscriptSegment( index=i, start=float(s.get("start", 0)), end=float(s.get("end", 0)), en=str(s.get("text", "")).strip(), zh="", ) for i, s in enumerate(segments) ] update(job, transcript=en_only, message=f"ASR 完成 · {len(en_only)} 段,开始翻译…", progress=88) # 2) Gemini 翻译 zh_list = await asyncio.to_thread(_translate_sync, segments) full = [ TranscriptSegment( index=seg.index, start=seg.start, end=seg.end, en=seg.en, zh=zh_list[i] if i < len(zh_list) else "", ) for i, seg in enumerate(en_only) ] update(job, transcript=full, status="transcribed", progress=100, message=f"转录完成 · {len(full)} 段({ASR_MODEL} + {TRANSLATE_MODEL})") except Exception as e: update(job, status="failed", error=str(e), message="转录失败") def _image_edit_call( image_path: Path, prompt: str, model: str | None = None, models: list[str] | None = None, fallback_text: bool = False, max_attempts: int = 3, max_side: int = 1024, ) -> tuple[bytes, str]: """通用 image edit 调用 · 失败重试 + 可选 text fallback。 返回 (image_bytes, effective_mode) where effective_mode in {"edit","text"}。 失败 raise RuntimeError。 输入图自动 resize 到 max_side(默认 1024)边长后再 base64,避免大图把 Gemini function call 输入挤超阈值导致 incomplete_generation。 models: 多模型轮换列表,重试时换 model;不传则单一 model 重试。""" import base64 as b64lib import io as _io import time as _time import httpx from PIL import Image as _PILImage if not LLM_API_KEY: raise RuntimeError("LLM_API_KEY 未配置") # model 优先级:models 列表 > 单个 model 参数 > IMAGE_MODEL if models and len(models) > 0: models_cycle = list(models) else: models_cycle = [model or IMAGE_MODEL] model = models_cycle[0] # 缩到 max_side 内 try: im = _PILImage.open(image_path) if max(im.size) > max_side: im.thumbnail((max_side, max_side), _PILImage.LANCZOS) buf = _io.BytesIO() im.convert("RGB").save(buf, format="JPEG", quality=88) img_bytes_in = buf.getvalue() except Exception: # PIL 失败兜底走原文件 img_bytes_in = image_path.read_bytes() img_b64 = b64lib.b64encode(img_bytes_in).decode("ascii") data_uri = f"data:image/jpeg;base64,{img_b64}" plan: list[str] = ["edit"] * max_attempts if fallback_text: plan.append("text") last_err = "" resp_data: dict = {} effective_mode = "edit" for attempt, current_mode in enumerate(plan): # 多模型轮换:第 N 次重试用第 N 个 model(不够时用最后一个) current_model = models_cycle[min(attempt, len(models_cycle) - 1)] try: if current_mode == "edit": with httpx.Client(timeout=120) as client: r = client.post( f"{LLM_BASE_URL}/images/generations", headers={ "Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json", }, json={"model": current_model, "prompt": prompt, "image": data_uri, "n": 1}, ) r.raise_for_status() resp_data = r.json() else: resp = llm().images.generate(model=current_model, prompt=prompt, n=1) resp_data = resp.model_dump() if hasattr(resp, "model_dump") else {"data": [{"b64_json": resp.data[0].b64_json}]} if resp_data.get("data"): effective_mode = current_mode model = current_model # 记录实际成功的 model break err_obj = resp_data.get("error") or {} last_err = f"empty data · {err_obj.get('code', '')} · {str(err_obj.get('message', ''))[:200]} · model={current_model}" except httpx.HTTPStatusError as e: body = e.response.text transient = ( e.response.status_code >= 500 or "incomplete_generation" in body or "rate_limit" in body or "timeout" in body.lower() ) last_err = f"HTTP {e.response.status_code}: {body[:200]} · model={current_model}" if not transient: raise RuntimeError(f"image edit HTTP {e.response.status_code}: {body[:300]}") except Exception as e: last_err = f"{type(e).__name__}: {e} · model={current_model}" if attempt < len(plan) - 1: next_model = models_cycle[min(attempt + 1, len(models_cycle) - 1)] tag = f"retry {attempt + 1}/{len(plan)} → {next_model}" print(f"[image edit {tag}] {last_err}", flush=True) _time.sleep(1.0) data_arr = resp_data.get("data", []) if not data_arr: raise RuntimeError(f"image edit failed after {len(plan)} attempts: {last_err}") b64 = data_arr[0].get("b64_json") if not b64: raise RuntimeError("image edit returned no b64_json") return b64lib.b64decode(b64), effective_mode # ---------- API 路由 ---------- class CreateJobReq(BaseModel): url: str class TranslateReq(BaseModel): text: str target: Literal["en", "zh"] = "en" @app.post("/translate") def translate_text(req: TranslateReq) -> dict: """单条文本翻译(给生图自定义提取元素 zh→en 用)""" import re as _re text = req.text.strip() if not text: return {"text": ""} if not LLM_API_KEY: raise HTTPException(503, "LLM_API_KEY 未配置") target_label = "English" if req.target == "en" else "Simplified Chinese" prompt = ( f"Translate the following text into concise {target_label}, suitable as an element " "label in an image-generation prompt. Output only the translation itself — no quotes, " "no punctuation, no explanation, no markdown.\n\n" f"Input: {text}" ) try: resp = llm().chat.completions.create( model=TRANSLATE_MODEL, messages=[{"role": "user", "content": prompt}], temperature=0.2, max_tokens=200, ) out = (resp.choices[0].message.content or "").strip() if not out: rc = getattr(resp.choices[0].message, "reasoning_content", "") or "" if rc: out = rc.strip().splitlines()[-1].strip() out = _re.sub(r'^[\'"「『]+|[\'"」』]+$', "", out).strip() return {"text": out} except Exception as e: raise HTTPException(500, f"translate failed: {e}") @app.get("/health") def health() -> dict: return { "ok": True, "llm_configured": bool(LLM_API_KEY), "base_url": LLM_BASE_URL or "openai-default", "models": { "asr": ASR_MODEL, "translate": TRANSLATE_MODEL, "rewrite": REWRITE_MODEL, }, } @app.post("/jobs", response_model=Job) async def create_job(req: CreateJobReq, bg: BackgroundTasks) -> Job: if not req.url.strip(): raise HTTPException(400, "url required") job_id = uuid.uuid4().hex[:12] job = Job(id=job_id, url=req.url.strip()) JOBS[job_id] = job save_state(job) bg.add_task(pipeline_download, job_id) return job @app.post("/jobs/upload", response_model=Job) async def create_job_from_upload(bg: BackgroundTasks, file: UploadFile = File(...)) -> Job: if not file.filename: raise HTTPException(400, "file required") ext = Path(file.filename).suffix.lower() if ext not in {".mp4", ".mov", ".webm", ".mkv", ".m4v"}: raise HTTPException(400, f"unsupported video format: {ext}") job_id = uuid.uuid4().hex[:12] d = job_dir(job_id) mp4 = d / "source.mp4" with mp4.open("wb") as f: while chunk := await file.read(1024 * 1024): f.write(chunk) if not mp4.exists() or mp4.stat().st_size == 0: raise HTTPException(500, "upload failed") job = Job(id=job_id, url=f"upload://{file.filename}") JOBS[job_id] = job save_state(job) bg.add_task(pipeline_download, job_id) return job @app.post("/jobs/{job_id}/analyze", response_model=Job) async def trigger_analyze(job_id: str, bg: BackgroundTasks, frames: int = KEYFRAME_COUNT) -> Job: job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") if job.status not in {"downloaded", "frames_extracted", "transcribed", "failed"}: raise HTTPException(409, f"status must be downloaded/failed, got {job.status}") bg.add_task(pipeline_analyze, job_id, frames) return job @app.post("/jobs/{job_id}/frames", response_model=Job) def add_manual_frame(job_id: str, t: float) -> Job: """从指定时间戳手动抽 1 帧追加到 job.frames""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") if not job.video_url: raise HTTPException(400, "video not ready") d = job_dir(job_id) mp4 = d / "source.mp4" if not mp4.exists(): raise HTTPException(400, "source.mp4 missing") frames_dir = d / "frames" frames_dir.mkdir(parents=True, exist_ok=True) # 新 index:max(existing)+1(即使列表已按 ts 排序,文件名用 index 保持稳定) next_idx = max((f.index for f in job.frames), default=-1) + 1 out = frames_dir / f"{next_idx:03d}.jpg" try: run([ "ffmpeg", "-y", "-ss", str(t), "-i", str(mp4), "-frames:v", "1", "-pix_fmt", "yuvj420p", "-q:v", "3", str(out), ]) except RuntimeError as e: raise HTTPException(500, f"ffmpeg failed: {e}") new_frame = KeyFrame( index=next_idx, timestamp=round(float(t), 2), url=f"/jobs/{job_id}/frames/{next_idx}.jpg", ) merged = sorted(list(job.frames) + [new_frame], key=lambda f: f.timestamp) update(job, frames=merged, message=f"已手动加帧({t:.1f}s),共 {len(merged)} 张") return job @app.get("/jobs/{job_id}", response_model=Job) def get_job(job_id: str) -> Job: job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") return job @app.post("/jobs/{job_id}/transcribe", response_model=Job) async def trigger_transcribe(job_id: str, bg: BackgroundTasks) -> Job: job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") if job.status != "frames_extracted": raise HTTPException(409, f"status must be frames_extracted, got {job.status}") bg.add_task(pipeline_transcribe, job_id) return job @app.get("/jobs/{job_id}/video.mp4") def get_video(job_id: str): p = job_dir(job_id) / "source.mp4" if not p.exists(): raise HTTPException(404, "video not found") return FileResponse(p, media_type="video/mp4") @app.get("/jobs/{job_id}/frames/{idx}.jpg") def get_frame(job_id: str, idx: int): p = job_dir(job_id) / "frames" / f"{idx:03d}.jpg" if not p.exists(): raise HTTPException(404, "frame not found") return FileResponse(p, media_type="image/jpeg") class GenerateReq(BaseModel): prompt: str extra_prompt: str = "" # ✓ 需要的元素(正向) negative_prompt: str = "" # ✗ 不需要的元素(负向) model: str = "" # 留空用 IMAGE_MODEL 默认 mode: str = "edit" # "edit" 带参考图,"text" 纯文字 from_selected: bool = False # True 时优先用 frame.selected 的生成图作 reference(迭代),否则原关键帧 @app.post("/jobs/{job_id}/frames/{idx}/generate", response_model=Job) def generate_image(job_id: str, idx: int, req: GenerateReq) -> Job: """根据关键帧 + prompt 生成新图(image-to-image 或 text-to-image)""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") frame_path = job_dir(job_id) / "frames" / f"{idx:03d}.jpg" if not frame_path.exists(): raise HTTPException(404, "frame file missing") # 决定 i2i 参考图:from_selected=True 且存在 selected 生成图 → 用它(迭代);否则原关键帧 reference_path = frame_path reference_source = "keyframe" if req.from_selected: sel = next((g for g in frame.generated_images if g.selected), None) if sel: sel_path = job_dir(job_id) / "gen" / f"{idx:03d}_{sel.id}.jpg" if sel_path.exists(): reference_path = sel_path reference_source = f"gen:{sel.id[:6]}" full_prompt = req.prompt.strip() if req.extra_prompt.strip(): full_prompt = f"{full_prompt}. Include: {req.extra_prompt.strip()}" if req.negative_prompt.strip(): full_prompt = f"{full_prompt}. Avoid: {req.negative_prompt.strip()}" if not full_prompt: raise HTTPException(400, "prompt required") model = req.model or IMAGE_MODEL gen_id = uuid.uuid4().hex[:12] import base64 as b64lib import time as _time import httpx img_b64: str | None = None if req.mode == "edit": img_b64 = b64lib.b64encode(reference_path.read_bytes()).decode("ascii") # 尝试 i2i 最多 3 次,全失败时降级 text-only 再试 1 次 plan: list[str] = ([req.mode] * 3) if req.mode == "edit" else [req.mode] if req.mode == "edit": plan.append("text") # i2i 都失败时自动降级 resp_data: dict = {} last_err = "" effective_mode = req.mode for attempt, current_mode in enumerate(plan): try: if current_mode == "edit": data_uri = f"data:image/jpeg;base64,{img_b64}" # OpenAI SDK 不直接支持 image 参数,用底层 httpx with httpx.Client(timeout=120) as client: r = client.post( f"{LLM_BASE_URL}/images/generations", headers={ "Authorization": f"Bearer {LLM_API_KEY}", "Content-Type": "application/json", }, json={ "model": model, "prompt": full_prompt, "image": data_uri, "n": 1, }, ) r.raise_for_status() resp_data = r.json() else: # text-only resp = llm().images.generate(model=model, prompt=full_prompt, n=1) resp_data = resp.model_dump() if hasattr(resp, "model_dump") else {"data": [{"b64_json": resp.data[0].b64_json}]} if resp_data.get("data"): effective_mode = current_mode break err_obj = resp_data.get("error") or {} last_err = f"empty data · {err_obj.get('code', '')} · {str(err_obj.get('message', ''))[:200]}" except httpx.HTTPStatusError as e: body = e.response.text transient = ( e.response.status_code >= 500 or "incomplete_generation" in body or "rate_limit" in body or "timeout" in body.lower() ) last_err = f"HTTP {e.response.status_code}: {body[:200]}" if not transient: raise HTTPException(500, f"image gen HTTP {e.response.status_code}: {body[:300]}") except Exception as e: last_err = f"{type(e).__name__}: {e}" if attempt < len(plan) - 1: next_mode = plan[attempt + 1] tag = f"fallback → {next_mode}" if next_mode != current_mode else f"retry {attempt + 1}/{len(plan)}" print(f"[image gen {tag}] {last_err}", flush=True) _time.sleep(1.5 * (attempt + 1)) data_arr = resp_data.get("data", []) if not data_arr: raise HTTPException(500, f"image gen failed after {len(plan)} attempts: {last_err}") item = data_arr[0] b64 = item.get("b64_json") if not b64: raise HTTPException(500, "image gen returned no b64_json") # 保存到本地 jobs//gen/_.jpg gen_dir = job_dir(job_id) / "gen" gen_dir.mkdir(parents=True, exist_ok=True) out_path = gen_dir / f"{idx:03d}_{gen_id}.jpg" out_path.write_bytes(b64lib.b64decode(b64)) new_gen = GeneratedImage( id=gen_id, prompt=full_prompt, model=model, mode=effective_mode, url=f"/jobs/{job_id}/frames/{idx}/gen/{gen_id}.jpg", selected=False, created_at=_time.time(), ) # 写回 job.frames for f in job.frames: if f.index == idx: f.generated_images = f.generated_images + [new_gen] update(job, frames=job.frames, message=f"生图完成 · 分镜 {idx + 1}") return job @app.get("/jobs/{job_id}/frames/{idx}/gen/{gen_id}.jpg") def get_generated_image(job_id: str, idx: int, gen_id: str): p = job_dir(job_id) / "gen" / f"{idx:03d}_{gen_id}.jpg" if not p.exists(): raise HTTPException(404, "generated image not found") return FileResponse(p, media_type="image/jpeg") class SelectGenReq(BaseModel): selected: bool @app.post("/jobs/{job_id}/frames/{idx}/gen/{gen_id}/select", response_model=Job) def select_generated(job_id: str, idx: int, gen_id: str, req: SelectGenReq) -> Job: job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") for f in job.frames: if f.index != idx: continue for g in f.generated_images: # 单选:该帧只能选一张 if g.id == gen_id: g.selected = req.selected else: g.selected = False break update(job, frames=job.frames) return job @app.post("/jobs/{job_id}/frames/{idx}/describe", response_model=Job) def describe_frame(job_id: str, idx: int) -> Job: """调 vision 模型识别该关键帧,返回结构化描述。""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") p = job_dir(job_id) / "frames" / f"{idx:03d}.jpg" if not p.exists(): raise HTTPException(404, "frame file not found") import base64 as b64lib import re as _re img_b64 = b64lib.b64encode(p.read_bytes()).decode("ascii") prompt = ( "请识别这张图,输出严格 JSON(不要 markdown 不要解释,不要思考):\n" '{\n' ' "scene": "一句话描述场景",\n' ' "objects": [{"name": "物体名(中文)", "position": "在画面哪里", "color": "颜色", "extract_prompt": "用于提取该元素的英文 prompt"}],\n' ' "style": "整体风格 / 打光 / 色调(一句话)",\n' ' "suggested_prompt": "适合用作下游生图的完整英文 prompt"\n' '}\n' "要求:objects 列出 3-8 个画面里**可独立提取**的主要元素,extract_prompt 用于后续 image edit 模型。" ) last_err = "" data = None for attempt in range(3): try: resp = llm().chat.completions.create( model=VISION_MODEL, messages=[{"role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}}, ]}], response_format={"type": "json_object"}, temperature=0.3, max_tokens=3000, ) content = (resp.choices[0].message.content or "").strip() if not content: # thinking 模型可能 content 空;尝试取 reasoning_content 里挖 JSON rc = getattr(resp.choices[0].message, "reasoning_content", "") or "" m = _re.search(r"\{[\s\S]*\}", rc) content = m.group(0) if m else "" # 剥掉 ```json ... ``` 包装 content = _re.sub(r"^```(?:json)?\s*|\s*```$", "", content).strip() if not content: last_err = f"empty content (attempt {attempt + 1})" continue data = json.loads(content) break except json.JSONDecodeError as e: last_err = f"json decode (attempt {attempt + 1}): {e} · raw[:200]={content[:200]}" print(f"[vision retry] {last_err}", flush=True) continue except Exception as e: last_err = f"vision call (attempt {attempt + 1}): {e}" print(f"[vision retry] {last_err}", flush=True) continue if data is None: raise HTTPException(500, last_err or "vision failed after 3 retries") # 写回 job new_frames = [] for f in job.frames: if f.index == idx: f.description = data new_frames.append(f) update(job, frames=new_frames, message=f"识别完成 · 分镜 {idx + 1}") return job # ---------- 清洗水印 / 元素提取(关键帧二阶段加工) ---------- class CleanupReq(BaseModel): # 多个相对坐标矩形 0-1,限制清洗范围;空 / None = 全图清洗 regions: list[dict] | None = None # [{"x","y","w","h"}, ...] def _region_to_phrase(r: dict) -> str: """把相对坐标矩形转成简短方位描述给 prompt 用(避免百分号 / 括号触发模型异常)""" x = max(0.0, min(1.0, float(r.get("x", 0)))) y = max(0.0, min(1.0, float(r.get("y", 0)))) w = max(0.0, min(1.0 - x, float(r.get("w", 0)))) h = max(0.0, min(1.0 - y, float(r.get("h", 0)))) if w <= 0 or h <= 0: return "" cx, cy = x + w / 2, y + h / 2 hpos = "left" if cx < 0.4 else "right" if cx > 0.6 else "middle" vpos = "top" if cy < 0.4 else "bottom" if cy > 0.6 else "middle" if hpos == "middle" and vpos == "middle": return "center" if hpos == "middle": return vpos if vpos == "middle": return hpos return f"{vpos} {hpos}" @app.post("/jobs/{job_id}/frames/{idx}/cleanup", response_model=Job) def cleanup_frame(job_id: str, idx: int, req: CleanupReq | None = None) -> Job: """调 nano-banana image edit 清洗关键帧:去水印 / @用户名 / 字幕 / 平台 logo。 输出干净版到 jobs//cleaned/.jpg,写回 frame.cleaned_url。 可选 region: 限定只清洗框内区域。""" import time as _time job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") frame_path = job_dir(job_id) / "frames" / f"{idx:03d}.jpg" if not frame_path.exists(): raise HTTPException(404, "frame file missing") region_phrases: list[str] = [] if req and req.regions: for r in req.regions: p = _region_to_phrase(r) if p: region_phrases.append(p) region_phrases = list(dict.fromkeys(region_phrases)) # prompt 用"重画一张副本"语义而非"erase / remove only X" — 避免 Gemini 走 mask/inpainting # function call 路径(实测该路径在 SKG 网关上 100% 触发 incomplete_generation) if region_phrases: if len(region_phrases) == 1: zones = f"the {region_phrases[0]} area" else: zones = ", ".join(region_phrases) + " areas" prompt = ( f"Recreate this image as a clean version: remove the text and graphics in {zones}, " "keep the rest of the scene identical." ) else: prompt = ( "Recreate this image as a clean version without watermarks, captions, " "hashtags, usernames, or platform logos. Keep the composition and style." ) # 模型轮换:nano-banana-pro 失败时换 flash 系列 models = [ IMAGE_MODEL, # gemini-3-pro-image-preview (nano-banana-pro) "gemini-3.1-flash-image-preview", "gemini-2.5-flash-image", ] try: img_bytes, _mode = _image_edit_call( frame_path, prompt, models=models, fallback_text=False, max_attempts=3, ) except RuntimeError as e: raise HTTPException(500, f"cleanup failed: {e}") out_dir = job_dir(job_id) / "cleaned" out_dir.mkdir(parents=True, exist_ok=True) out_path = out_dir / f"{idx:03d}.jpg" out_path.write_bytes(img_bytes) new_frames = [] for f in job.frames: if f.index == idx: f.cleaned_url = f"/jobs/{job_id}/frames/{idx}/cleaned.jpg?t={int(_time.time())}" f.cleaned_applied = False # 重新清洗:重置"已应用"状态 new_frames.append(f) update(job, frames=new_frames, message=f"清洗完成 · 分镜 {idx + 1}") return job @app.get("/jobs/{job_id}/frames/{idx}/cleaned.jpg") def get_cleaned_frame(job_id: str, idx: int): p = job_dir(job_id) / "cleaned" / f"{idx:03d}.jpg" if not p.exists(): raise HTTPException(404, "cleaned frame not found") return FileResponse(p, media_type="image/jpeg") @app.delete("/jobs/{job_id}/frames/{idx}/cleanup", response_model=Job) def discard_cleaned(job_id: str, idx: int) -> Job: """丢弃待应用的清洗版(不影响已应用的)""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") p = job_dir(job_id) / "cleaned" / f"{idx:03d}.jpg" if p.exists(): try: p.unlink() except OSError: pass new_frames = [] for f in job.frames: if f.index == idx: f.cleaned_url = None new_frames.append(f) update(job, frames=new_frames, message=f"丢弃清洗版 · 分镜 {idx + 1}") return job @app.post("/jobs/{job_id}/frames/{idx}/cleanup/apply", response_model=Job) def apply_cleaned(job_id: str, idx: int) -> Job: """用清洗版替换原关键帧:物理覆盖 frames/{idx}.jpg ← cleaned/{idx}.jpg。 原图作备份 → orig/{idx}.jpg(首次替换时备份,后续替换跳过)。 替换后 frame.cleaned_url 清空(不再有"待应用"清洗版)""" import shutil as _shutil job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") cleaned_path = job_dir(job_id) / "cleaned" / f"{idx:03d}.jpg" if not cleaned_path.exists(): raise HTTPException(404, "no cleaned version to apply") frame_path = job_dir(job_id) / "frames" / f"{idx:03d}.jpg" # 首次替换:把原图备份到 orig/{idx}.jpg orig_dir = job_dir(job_id) / "orig" orig_dir.mkdir(parents=True, exist_ok=True) orig_backup = orig_dir / f"{idx:03d}.jpg" if not orig_backup.exists() and frame_path.exists(): _shutil.copy2(frame_path, orig_backup) # 用 cleaned 覆盖 frames/ _shutil.copy2(cleaned_path, frame_path) # 删 cleaned 文件(已经"应用",不再是单独的待选版本) try: cleaned_path.unlink() except OSError: pass new_frames = [] for f in job.frames: if f.index == idx: f.cleaned_url = None f.cleaned_applied = True new_frames.append(f) update(job, frames=new_frames, message=f"已替换分镜 {idx + 1} 为清洗版") return job class AddElementReq(BaseModel): name_zh: str name_en: str = "" position: str = "" source: Literal["auto", "manual", "region"] = "manual" region: dict | None = None @app.post("/jobs/{job_id}/frames/{idx}/elements", response_model=Job) def add_element(job_id: str, idx: int, req: AddElementReq) -> Job: """加一条元素 · 若 name_en 缺则自动 zh→en 翻译""" import time as _time import re as _re job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") name_zh = req.name_zh.strip() if not name_zh: raise HTTPException(400, "name_zh required") name_en = req.name_en.strip() if not name_en and LLM_API_KEY: try: prompt = ( "Translate the following text into concise English, suitable as an element label " "in an image-generation prompt. Output only the translation — no quotes, no punctuation, " f"no explanation.\n\nInput: {name_zh}" ) resp = llm().chat.completions.create( model=TRANSLATE_MODEL, messages=[{"role": "user", "content": prompt}], temperature=0.2, max_tokens=200, ) out = (resp.choices[0].message.content or "").strip() if not out: rc = getattr(resp.choices[0].message, "reasoning_content", "") or "" if rc: out = rc.strip().splitlines()[-1].strip() name_en = _re.sub(r'^[\'"「『]+|[\'"」』]+$', "", out).strip() except Exception as e: print(f"[add_element translate failed] {e}", flush=True) name_en = "" el = KeyElement( id=uuid.uuid4().hex[:8], name_zh=name_zh, name_en=name_en, position=req.position.strip(), source=req.source, region=req.region, created_at=_time.time(), ) new_frames = [] for f in job.frames: if f.index == idx: f.elements = f.elements + [el] new_frames.append(f) update(job, frames=new_frames, message=f"加入元素 · 分镜 {idx + 1} · {name_zh}") return job @app.delete("/jobs/{job_id}/frames/{idx}/elements/{element_id}", response_model=Job) def delete_element(job_id: str, idx: int, element_id: str) -> Job: job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") new_frames = [] removed = False for f in job.frames: if f.index == idx: before = len(f.elements) f.elements = [e for e in f.elements if e.id != element_id] removed = len(f.elements) < before # 若有抠图文件也删 if removed: for ext in ("jpg", "png"): cutout = job_dir(job_id) / "elements" / f"{idx:03d}_{element_id}.{ext}" if cutout.exists(): try: cutout.unlink() except OSError: pass new_frames.append(f) if not removed: raise HTTPException(404, "element not found") update(job, frames=new_frames, message=f"删除元素 · 分镜 {idx + 1}") return job class CutoutReq(BaseModel): background: Literal["white", "black"] = "white" @app.post("/jobs/{job_id}/frames/{idx}/elements/{element_id}/cutout", response_model=Job) def cutout_element(job_id: str, idx: int, element_id: str, req: CutoutReq | None = None) -> Job: """单元素抠图:调 nano-banana image edit,输出纯白底 / 纯黑底元素图。 注:Gemini 输出 JPEG 不支持真 alpha,因此用纯色背景。""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") el = next((e for e in frame.elements if e.id == element_id), None) if not el: raise HTTPException(404, "element not found") # 优先用 cleaned 版作 reference(已去掉 logo / 水印干扰),fallback 原图 cleaned_path = job_dir(job_id) / "cleaned" / f"{idx:03d}.jpg" src = cleaned_path if cleaned_path.exists() else job_dir(job_id) / "frames" / f"{idx:03d}.jpg" if not src.exists(): raise HTTPException(404, "source frame file missing") background = (req.background if req else "white") or "white" bg_phrase = f"pure {background}" target = (el.name_en or el.name_zh).strip() region_phrase = _region_to_phrase(el.region) if el.region else "" if region_phrase: prompt = ( f"Extract whatever is in the {region_phrase} part of the image as a standalone asset. " f"Place it on a {bg_phrase} background, isolated, no other objects." ) else: position_hint = f" Located in the {el.position} area." if el.position else "" prompt = ( f"Extract the {target} from this image as a standalone asset.{position_hint} " f"Place it on a {bg_phrase} background, isolated, no other objects." ) models = [ IMAGE_MODEL, "gemini-3.1-flash-image-preview", "gemini-2.5-flash-image", ] try: img_bytes, _mode = _image_edit_call( src, prompt, models=models, fallback_text=False, max_attempts=3, ) except RuntimeError as e: raise HTTPException(500, f"cutout failed: {e}") out_dir = job_dir(job_id) / "elements" out_dir.mkdir(parents=True, exist_ok=True) # 实际是 JPEG bytes,文件用 .jpg 真名 out_path = out_dir / f"{idx:03d}_{element_id}.jpg" out_path.write_bytes(img_bytes) # 旧版的 .png 文件(错命名为 .png 的 JPEG)也清理掉 old_png = out_dir / f"{idx:03d}_{element_id}.png" if old_png.exists(): try: old_png.unlink() except OSError: pass new_frames = [] for f in job.frames: if f.index == idx: for e in f.elements: if e.id == element_id: e.cutout_id = element_id e.cutout_background = background new_frames.append(f) update(job, frames=new_frames, message=f"抠图完成 · {el.name_zh}({background} 底)") return job @app.get("/jobs/{job_id}/frames/{idx}/elements/{element_id}/cutout.jpg") def get_cutout(job_id: str, idx: int, element_id: str): p = job_dir(job_id) / "elements" / f"{idx:03d}_{element_id}.jpg" if not p.exists(): # 兼容老数据 legacy = job_dir(job_id) / "elements" / f"{idx:03d}_{element_id}.png" if legacy.exists(): return FileResponse(legacy, media_type="image/jpeg") raise HTTPException(404, "cutout not found") return FileResponse(p, media_type="image/jpeg") # ---------- 删除:关键帧 / 单张生成图 ---------- @app.delete("/jobs/{job_id}/frames/{idx}", response_model=Job) def delete_frame(job_id: str, idx: int) -> Job: """删除整张关键帧,清理所有附属文件(原图 / 干净版 / 元素抠图 / 生成图)""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") target = next((f for f in job.frames if f.index == idx), None) if not target: raise HTTPException(404, "frame not found") d = job_dir(job_id) # 删文件 — 静默错误,文件可能不存在 paths = [ d / "frames" / f"{idx:03d}.jpg", d / "cleaned" / f"{idx:03d}.jpg", ] for p in paths: if p.exists(): try: p.unlink() except OSError: pass # 该帧的所有元素抠图(命名前缀 {idx:03d}_) elements_dir = d / "elements" if elements_dir.exists(): for ext in ("png", "jpg"): for p in elements_dir.glob(f"{idx:03d}_*.{ext}"): try: p.unlink() except OSError: pass # 该帧的所有生成图 gen_dir = d / "gen" if gen_dir.exists(): for p in gen_dir.glob(f"{idx:03d}_*.jpg"): try: p.unlink() except OSError: pass new_frames = [f for f in job.frames if f.index != idx] update(job, frames=new_frames, message=f"删除分镜 {idx + 1}") return job @app.delete("/jobs/{job_id}/frames/{idx}/gen/{gen_id}", response_model=Job) def delete_generated(job_id: str, idx: int, gen_id: str) -> Job: """删除该 frame 的某张生成图(文件 + 列表)""" job = JOBS.get(job_id) if not job: raise HTTPException(404, "job not found") frame = next((f for f in job.frames if f.index == idx), None) if not frame: raise HTTPException(404, "frame not found") p = job_dir(job_id) / "gen" / f"{idx:03d}_{gen_id}.jpg" if p.exists(): try: p.unlink() except OSError: pass new_frames = [] found = False for f in job.frames: if f.index == idx: before = len(f.generated_images) f.generated_images = [g for g in f.generated_images if g.id != gen_id] found = len(f.generated_images) < before new_frames.append(f) if not found: raise HTTPException(404, "generated image not found") update(job, frames=new_frames, message=f"删除生成图 · 分镜 {idx + 1}") return job