Deploy Blender Headless API

app.py

@@ -1,239 +1,124 @@
+import base64
+import json
 import os
-import re
+import shutil
 import subprocess
-import tempfile
+import uuid
 from pathlib import Path
+from typing import Any, List, Optional
+
+from fastapi import FastAPI, Header, HTTPException
+from fastapi.responses import FileResponse, JSONResponse
+from pydantic import BaseModel
+
+APP_DIR = Path(__file__).parent
+JOB_DIR = Path(os.environ.get("JOB_DIR", "/data/jobs"))
+BLENDER_BIN = os.environ.get("BLENDER_BIN", "/opt/blender/blender")
+JOB_TIMEOUT = int(os.environ.get("JOB_TIMEOUT_SECONDS", "180"))
+WORKER_SCRIPT = APP_DIR / "worker.py"
+
+# 可选的简单鉴权：设置了 API_KEY 环境变量后，/execute 需要带 X-API-Key 头
+API_KEY = os.environ.get("API_KEY")
+
+JOB_DIR.mkdir(parents=True, exist_ok=True)
+
+app = FastAPI(title="Blender Headless API")
+
+
+class ExecuteRequest(BaseModel):
+    ops: List[dict[str, Any]]
+    scene_b64: Optional[str] = None   # 可选：输入 .blend 文件（base64）
+
+
+def check_api_key(x_api_key: Optional[str]):
+    if API_KEY and x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="invalid or missing X-API-Key")
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "blender_bin": BLENDER_BIN,
+        "blender_exists": os.path.exists(BLENDER_BIN),
+    }
+
+
+@app.post("/execute")
+def execute(req: ExecuteRequest, x_api_key: Optional[str] = Header(default=None)):
+    check_api_key(x_api_key)
+
+    job_id = uuid.uuid4().hex[:12]
+    job_path = JOB_DIR / job_id
+    input_dir = job_path / "input"
+    output_dir = job_path / "output"
+    input_dir.mkdir(parents=True)
+    output_dir.mkdir(parents=True)
+
+    cmd = [BLENDER_BIN, "--background"]
+
+    if req.scene_b64:
+        scene_in_path = input_dir / "input.blend"
+        scene_in_path.write_bytes(base64.b64decode(req.scene_b64))
+        cmd.append(str(scene_in_path))
+
+    ops_path = input_dir / "ops.json"
+    ops_path.write_text(json.dumps(req.ops))
+
+    cmd += ["--python", str(WORKER_SCRIPT), "--", str(ops_path), str(output_dir)]
 
-import gradio as gr
-import torch
-import soundfile as sf
-import torchaudio
-from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import VitsModel
-
-ASR_MODEL = os.getenv("ASR_MODEL", "openai/whisper-large-v3-turbo")
-TRANSLATION_MODEL = os.getenv("TRANSLATION_MODEL", "Helsinki-NLP/opus-mt-zh-ms")
-TTS_MODEL = os.getenv("TTS_MODEL", "facebook/mms-tts-zlm")
-
-torch.set_num_threads(2)
-DEVICE = "cpu"
-
-_asr_processor = _asr_model = None
-_tr_tokenizer = _tr_model = None
-_tts_tokenizer = _tts_model = None
-
-def run(cmd):
-    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-    if p.returncode:
-        raise RuntimeError(p.stderr[-5000:])
-    return p.stdout
-
-def video_duration(path):
-    return float(run(["ffprobe","-v","error","-show_entries","format=duration",
-                      "-of","default=noprint_wrappers=1:nokey=1",str(path)]).strip())
-
-def extract_audio(video, wav):
-    run(["ffmpeg","-y","-i",str(video),"-vn","-ac","1","-ar","16000",str(wav)])
-
-def load_asr():
-    global _asr_processor, _asr_model
-    if _asr_model is None:
-        _asr_processor = AutoProcessor.from_pretrained(ASR_MODEL)
-        _asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(
-            ASR_MODEL, low_cpu_mem_usage=True
-        ).to(DEVICE).eval()
-    return _asr_processor, _asr_model
-
-def transcribe(wav):
-    processor, model = load_asr()
-    audio, sr = sf.read(str(wav))
-    if len(audio.shape) > 1:
-        audio = audio.mean(axis=1)
-    if sr != 16000:
-        audio = torchaudio.functional.resample(
-            torch.from_numpy(audio).float(), sr, 16000
-        ).numpy()
-
-    total = len(audio) / 16000
-    pieces = []
-    chunk = 30
-
-    for start in range(0, int(total + 0.999), chunk):
-        x = audio[start*16000:min(int(total*16000),(start+chunk)*16000)]
-        if len(x) < 400:
-            continue
-        inputs = processor(x, sampling_rate=16000, return_tensors="pt")
-        with torch.inference_mode():
-            ids = model.generate(
-                inputs.input_features,
-                language="zh",
-                task="transcribe",
-                max_new_tokens=256,
-                prompt="无人机 架 大疆 北约 大毛 小前线 杨栋梁"
-            )
-        text = processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
-        if text:
-            pieces.append(text)
-    return " ".join(pieces)
-
-def load_translation():
-    global _tr_tokenizer, _tr_model
-    if _tr_model is None:
-        _tr_tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL)
-        _tr_model = AutoModelForSeq2SeqLM.from_pretrained(
-            TRANSLATION_MODEL, low_cpu_mem_usage=True
-        ).to(DEVICE).eval()
-    return _tr_tokenizer, _tr_model
-
-def translate(text):
-    tokenizer, model = load_translation()
-    parts = [x.strip() for x in re.split(r'(?<=[。！��!?])\s*', text) if x.strip()]
-    if not parts:
-        parts = [text]
-
-    result = []
-    for part in parts:
-        inputs = tokenizer(part, return_tensors="pt", truncation=True, max_length=256)
-        with torch.inference_mode():
-            ids = model.generate(**inputs, max_new_tokens=256, num_beams=4,
-                                 renormalize_logits=True)
-        result.append(tokenizer.decode(ids[0], skip_special_tokens=True).strip())
-    return " ".join(result)
-
-def load_tts():
-    global _tts_tokenizer, _tts_model
-    if _tts_model is None:
-        _tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL)
-        _tts_model = VitsModel.from_pretrained(
-            TTS_MODEL, low_cpu_mem_usage=True
-        ).to(DEVICE).eval()
-    return _tts_tokenizer, _tts_model
-
-def tts(text, wav):
-    tokenizer, model = load_tts()
-    inputs = tokenizer(text, return_tensors="pt", padding=True)
-    with torch.inference_mode():
-        audio = model(**inputs).waveform.cpu()
-    sf.write(str(wav), audio.squeeze().numpy(), model.config.sampling_rate)
-
-def fit_audio(src, dst, target):
-    target = max(float(target), 0.1)
-    ratio = video_duration(src) / target
-    filters = []
-    while ratio > 2:
-        filters.append("atempo=2.0")
-        ratio /= 2
-    while ratio < 0.5:
-        filters.append("atempo=0.5")
-        ratio /= 0.5
-    filters.append(f"atempo={ratio:.6f}")
-    run(["ffmpeg","-y","-i",str(src),"-filter:a",",".join(filters),
-         "-t",f"{target:.3f}","-ar","24000","-ac","1",str(dst)])
-
-def make_dub(ms, total, work):
-    parts = [x.strip() for x in re.split(r'(?<=[。！？!?])\s*|(?<=[，,])\s+', ms) if x.strip()]
-    if not parts:
-        parts = [ms]
-    weights = [max(1,len(re.sub(r'\s+','',x))) for x in parts]
-    total_w = sum(weights)
-
-    # Create silent base
-    silent = work/"silent.wav"
-    run(["ffmpeg","-y","-f","lavfi","-i","anullsrc=r=24000:cl=mono",
-         "-t",f"{total:.3f}","-ar","24000","-ac","1",str(silent)])
-
-    # Generate and fit each TTS segment
-    inputs = ["-i", str(silent)]
-    filter_parts = []
-    start = 0.0
-
-    for i, part in enumerate(parts, 1):
-        slot = total * weights[i-1] / total_w
-        raw, fitted = work/f"tts_{i}.wav", work/f"fit_{i}.wav"
-        tts(part, raw)
-        fit_audio(raw, fitted, slot)
-        inputs += ["-i", str(fitted)]
-        delay = int(start*1000)
-        filter_parts.append(f"[{i}:a]adelay={delay}|{delay}[d{i}]")
-        start += slot
-
-    # Build filter_complex: delay each TTS segment, then mix with silent base
-    # Input 0 = silent, Input 1..n = TTS segments
-    mix_inputs = "[0:a]" + "".join(f"[d{i}]" for i in range(1, len(parts)+1))
-    filt = ";".join(filter_parts) + f";{mix_inputs}amix=inputs={len(parts)+1}:duration=first[a]"
-    
-    out = work/"dub.wav"
-    run(["ffmpeg","-y",*inputs,"-filter_complex",filt,"-map","[a]",
-         "-ar","24000","-ac","1",str(out)])
-    return out
-
-def process(video, keep_original, original_volume, progress=gr.Progress()):
-    if not video:
-        raise gr.Error("请上传中文视频。")
-    work = Path(tempfile.mkdtemp(prefix="zhms_cpu_"))
     try:
-        src = Path(video)
-        audio = work/"source.wav"
-        progress(0.05, desc="提取声音…")
-        extract_audio(src, audio)
-        total = video_duration(src)
-
-        progress(0.15, desc="中文语音识别（CPU）…")
-        zh = transcribe(audio)
-        if not zh:
-            raise RuntimeError("没有识别到中文语音。")
-
-        progress(0.42, desc="翻译成 Bahasa Melayu…")
-        ms = translate(zh)
-
-        progress(0.58, desc="生成马来文配音…")
-        dub = make_dub(ms, total, work)
-
-        out = work/"malay_dub.mp4"
-        progress(0.9, desc="合成 MP4…")
-        if keep_original:
-            filt = f"[0:a]volume={float(original_volume):.3f}[o];[o][1:a]amix=inputs=2:duration=first[a]"
-            run(["ffmpeg","-y","-i",str(src),"-i",str(dub),"-filter_complex",filt,
-                 "-map","0:v:0","-map","[a]","-c:v","copy","-c:a","aac","-b:a","160k",
-                 "-shortest",str(out)])
-        else:
-            run(["ffmpeg","-y","-i",str(src),"-i",str(dub),
-                 "-map","0:v:0","-map","1:a:0","-c:v","copy","-c:a","aac",
-                 "-b:a","160k","-shortest",str(out)])
-
-        txt = work/"translation.txt"
-        txt.write_text("=== 中文识别 ===\n"+zh+"\n\n=== Bahasa Melayu ===\n"+ms,encoding="utf-8")
-        progress(1, desc="完成")
-        return str(out), str(txt), zh, ms
-    except Exception as e:
-        raise gr.Error(str(e))
-
-with gr.Blocks(title="中文 → Bahasa Melayu（CPU Basic）") as demo:
-    gr.Markdown("""# 🇨🇳 → 🇲🇾 中文视频 → Bahasa Melayu
-**Hugging Face CPU Basic 专用版本**
-
-上传中文视频 → 中文语音识别 → 马来文翻译 → 马来文 TTS → 输出 MP4。
-
-- 不需要 GPU
-- 不需要 API Key
-- 不使用声音克隆
-- 使用专门的 Chinese → Malay 翻译模型
-- 使用 Meta MMS Malay TTS
-
-⚠️ CPU Basic 只有 2 vCPU / 16 GB RAM，建议先测试 1–3 分钟视频。
-""")
-    video = gr.Video(label="上传中文视频", sources=["upload"])
-    with gr.Row():
-        keep = gr.Checkbox(label="保留原声音作为背景", value=False)
-        volume = gr.Slider(0,0.3,value=0.06,step=0.01,label="原声音量")
-    button = gr.Button("🚀 开始翻译成马来文", variant="primary")
-    output = gr.Video(label="🇲🇾 马来文视频")
-    text_file = gr.File(label="识别 + 翻译文本")
-    with gr.Row():
-        zh_box = gr.Textbox(label="中文识别结果", lines=8)
-        ms_box = gr.Textbox(label="Bahasa Melayu", lines=8)
-    button.click(process,[video,keep,volume],[output,text_file,zh_box,ms_box])
-
-if __name__ == "__main__":
-    demo.queue(max_size=2).launch()
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=JOB_TIMEOUT,
+        )
+    except subprocess.TimeoutExpired:
+        shutil.rmtree(job_path, ignore_errors=True)
+        raise HTTPException(status_code=504, detail="Blender job timed out")
+
+    (job_path / "blender.log").write_text(
+        (result.stdout or "") + "\n---STDERR---\n" + (result.stderr or "")
+    )
+
+    result_json_path = output_dir / "result.json"
+    if result.returncode != 0 or not result_json_path.exists():
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "message": "Blender execution failed",
+                "returncode": result.returncode,
+                "log_tail": (result.stdout[-2000:] + result.stderr[-2000:]),
+            },
+        )
+
+    op_result = json.loads(result_json_path.read_text())
+    files = sorted(p.name for p in output_dir.iterdir() if p.name != "result.json")
+
+    return JSONResponse(
+        {
+            "job_id": job_id,
+            "ops_result": op_result,
+            "files": files,
+            "download_base": f"/files/{job_id}/",
+        }
+    )
+
+
+@app.get("/files/{job_id}/{filename}")
+def get_file(job_id: str, filename: str):
+    file_path = (JOB_DIR / job_id / "output" / filename).resolve()
+    if JOB_DIR.resolve() not in file_path.parents or not file_path.exists():
+        raise HTTPException(status_code=404, detail="File not found")
+    return FileResponse(file_path)
+
+
+@app.delete("/jobs/{job_id}")
+def delete_job(job_id: str, x_api_key: Optional[str] = Header(default=None)):
+    check_api_key(x_api_key)
+    job_path = JOB_DIR / job_id
+    if job_path.exists():
+        shutil.rmtree(job_path)
+        return {"deleted": job_id}
+    raise HTTPException(status_code=404, detail="Job not found")