330 lines
16 KiB
Python
330 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
agent_short_drama.py — 苍耳 · 一键短剧 Agent
|
||
剧本 → 分镜(豆包) → 路由(豆包) → 逐镜出图(本地Z-Image)
|
||
→ 逐镜视频(LTX I2V, 本地) → 拼接(FFmpeg) → 成片
|
||
|
||
用法:
|
||
python agent_short_drama.py <剧本> -e 1 --until render # 只到出图
|
||
python agent_short_drama.py <剧本> -e 1 --until video # 到视频
|
||
python agent_short_drama.py <剧本> -e 1 --until compose # 到成片
|
||
python agent_short_drama.py <分镜.json> --from render # 从出图续跑
|
||
python agent_short_drama.py <manifest> --from video # 从视频续跑
|
||
"""
|
||
import os, sys, json, time, subprocess, urllib.request, argparse, glob, shutil
|
||
|
||
COMFY = "http://127.0.0.1:8188"
|
||
VIDEO_AI = os.path.dirname(os.path.abspath(__file__))
|
||
COMFY_OUT = os.path.expanduser("~/comfy/ComfyUI/output")
|
||
COMFY_IN = os.path.expanduser("~/comfy/ComfyUI/input")
|
||
|
||
LTX_NEG = ("static, frozen, no movement, flickering, jittery, choppy motion, morphing, "
|
||
"deformed, twisted, contorted, distorted, blurry, low resolution, extra limbs, "
|
||
"unnatural body, ugly, bad anatomy, disfigured, malformed, warped spine")
|
||
|
||
# ---------- 工具 ----------
|
||
def run_script(script, args):
|
||
cmd = [sys.executable, os.path.join(VIDEO_AI, script)] + args
|
||
print(" $", " ".join(cmd))
|
||
r = subprocess.run(cmd)
|
||
if r.returncode != 0:
|
||
raise RuntimeError(f"{script} 退出码 {r.returncode}")
|
||
|
||
def comfy_submit(prompt):
|
||
req = urllib.request.Request(f"{COMFY}/prompt",
|
||
data=json.dumps({"prompt": prompt}).encode(),
|
||
headers={"Content-Type": "application/json"})
|
||
return json.load(urllib.request.urlopen(req, timeout=15))["prompt_id"]
|
||
|
||
def comfy_wait(pid, timeout=900):
|
||
t0 = time.time()
|
||
while time.time() - t0 < timeout:
|
||
time.sleep(3)
|
||
try:
|
||
h = json.load(urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=8))
|
||
except Exception:
|
||
continue
|
||
if pid in h and h[pid].get("outputs"):
|
||
return h[pid]["outputs"]
|
||
raise TimeoutError(f"ComfyUI 任务 {pid} 超时")
|
||
|
||
def wait_file(path, timeout=60):
|
||
t0 = time.time()
|
||
while time.time() - t0 < timeout:
|
||
if os.path.isfile(path) and os.path.getsize(path) > 0:
|
||
return True
|
||
time.sleep(2)
|
||
return False
|
||
|
||
# ---------- Z-Image 出图 ----------
|
||
def zimage_workflow(prompt, seed, prefix, width=1024, height=1024):
|
||
return {
|
||
"1": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
|
||
"2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}},
|
||
"3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
|
||
"4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
|
||
"5": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}},
|
||
"6": {"class_type": "EmptySD3LatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
|
||
"7": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": 3.0}},
|
||
"8": {"class_type": "KSampler", "inputs": {"model": ["7", 0], "seed": seed, "steps": 8, "cfg": 1.0,
|
||
"sampler_name": "res_multistep", "scheduler": "simple",
|
||
"positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "denoise": 1.0}},
|
||
"9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["3", 0]}},
|
||
"10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": prefix}},
|
||
}
|
||
|
||
# ---------- LTX I2V 出视频 ----------
|
||
def ltx_i2v_workflow(image_name, prompt, seed, prefix, frames=49, width=768, height=512):
|
||
return {
|
||
"1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf"}},
|
||
"2": {"class_type": "LTXAVTextEncoderLoader", "inputs": {"text_encoder": "gemma_3_12B_it.safetensors", "ckpt_name": "ltx-2.3-22b-embconn.safetensors", "device": "cpu"}},
|
||
"3": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
|
||
"4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": LTX_NEG}},
|
||
"5": {"class_type": "LTXVConditioning", "inputs": {"positive": ["3", 0], "negative": ["4", 0], "frame_rate": 24.0}},
|
||
"6": {"class_type": "LoadImage", "inputs": {"image": image_name}},
|
||
"7": {"class_type": "VAELoader", "inputs": {"vae_name": "LTX23_video_vae_bf16.safetensors"}},
|
||
"8": {"class_type": "LTXVImgToVideo", "inputs": {"positive": ["5", 0], "negative": ["5", 1], "vae": ["7", 0], "image": ["6", 0], "width": width, "height": height, "length": frames, "batch_size": 1, "strength": 1.0}},
|
||
"9": {"class_type": "RandomNoise", "inputs": {"noise_seed": seed}},
|
||
"10": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler_ancestral_cfg_pp"}},
|
||
"11": {"class_type": "CFGGuider", "inputs": {"model": ["1", 0], "positive": ["8", 0], "negative": ["8", 1], "cfg": 1.0}},
|
||
"12": {"class_type": "ManualSigmas", "inputs": {"sigmas": "1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0"}},
|
||
"13": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["9", 0], "guider": ["11", 0], "sampler": ["10", 0], "sigmas": ["12", 0], "latent_image": ["8", 2]}},
|
||
"14": {"class_type": "LTXVTiledVAEDecode", "inputs": {"vae": ["7", 0], "latents": ["13", 0], "horizontal_tiles": 2, "vertical_tiles": 2, "overlap": 6, "last_frame_fix": False, "working_device": "cpu", "working_dtype": "float16"}},
|
||
"15": {"class_type": "CreateVideo", "inputs": {"images": ["14", 0], "fps": 24.0}},
|
||
"16": {"class_type": "SaveVideo", "inputs": {"video": ["15", 0], "filename_prefix": prefix, "format": "mp4", "codec": "h264"}},
|
||
}
|
||
|
||
def shot_to_prompt(shot, style=""):
|
||
parts = []
|
||
cam = shot.get('camera', '中景')
|
||
scenes = shot.get('scenes', [])
|
||
chars = shot.get('characters', [])
|
||
desc = shot.get('description', '')
|
||
if cam: parts.append(cam)
|
||
if scenes: parts.append("in " + scenes[0])
|
||
if chars: parts.append("with " + ", ".join(chars[:3]))
|
||
if desc: parts.append(desc[:120])
|
||
if style: parts.append(style)
|
||
return ", ".join(parts)
|
||
|
||
def shot_to_ltx_prompt(desc, camera=""):
|
||
"""分镜 → LTX I2V 运动提示词（四段式简化：起势→动作→环境→收尾）"""
|
||
base = (desc or "").strip()
|
||
action = base if base else "the subject moves naturally"
|
||
parts = [
|
||
f"Opening on the first frame scene, {camera or 'medium shot'}. The subject {action}, "
|
||
"with gentle flowing motion, grass and light drifting softly around. "
|
||
"Camera holds steady, shallow depth of field, warm natural light. "
|
||
"Slow, fluid, seamless motion, high detail. "
|
||
"Ends with the subject settling into a calm final pose."
|
||
]
|
||
return " ".join(parts)
|
||
|
||
# ---------- 各阶段 ----------
|
||
def stage_storyboard(script, episode, pro=False):
|
||
args = [script, "-e", str(episode)]
|
||
if pro: args.append("--pro")
|
||
run_script("tools/run_storyboard.py", args)
|
||
d = os.path.dirname(os.path.abspath(script))
|
||
m = glob.glob(os.path.join(d, "STORYBOARD*.json"))
|
||
if not m:
|
||
raise RuntimeError("未找到分镜 JSON")
|
||
return m[0]
|
||
|
||
def stage_route(sb_json):
|
||
run_script("tools/route_shots.py", [sb_json])
|
||
return sb_json
|
||
|
||
def stage_render(sb_json, out_dir=None, style="", max_shots=0):
|
||
sb = json.load(open(sb_json, encoding="utf-8"))
|
||
shots = sb.get("shots", [])
|
||
if max_shots: shots = shots[:max_shots]
|
||
if not shots: raise RuntimeError("分镜无 shots")
|
||
if not out_dir:
|
||
out_dir = os.path.join(os.path.dirname(os.path.abspath(sb_json)), "renders")
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
manifest = os.path.join(out_dir, "shot_manifest.json")
|
||
entries = []
|
||
for i, sh in enumerate(shots):
|
||
p = shot_to_prompt(sh, style)
|
||
seed = 1000 + i * 17
|
||
prefix = f"eed_shot_{i+1:04d}"
|
||
pid = comfy_submit(zimage_workflow(p, seed, prefix))
|
||
print(f" 🎬 [{i+1}/{len(shots)}] {sh.get('shot_number', f'S{i+1:02d}')}: {p[:50]}...")
|
||
comfy_wait(pid)
|
||
# 等文件落盘
|
||
img = os.path.join(COMFY_OUT, f"{prefix}_00001_.png")
|
||
wait_file(img)
|
||
entries.append({"shot": sh.get("shot_number", f"S{i+1:02d}"), "prompt": p, "ltx_desc": sh.get("description",""),
|
||
"camera": sh.get("camera",""), "seed": seed, "prefix": prefix})
|
||
with open(manifest, "w", encoding="utf-8") as f:
|
||
json.dump({"out_dir": out_dir, "shots": entries}, f, ensure_ascii=False, indent=2)
|
||
print(f" 📦 出图清单: {manifest}")
|
||
return manifest
|
||
|
||
def stage_video(manifest, frames=49, max_shots=0):
|
||
data = json.load(open(manifest, encoding="utf-8"))
|
||
entries = data["shots"]
|
||
if max_shots: entries = entries[:max_shots]
|
||
vdir = os.path.join(data["out_dir"], "videos")
|
||
os.makedirs(vdir, exist_ok=True)
|
||
vlist = []
|
||
for i, e in enumerate(entries):
|
||
src = os.path.join(COMFY_OUT, f"{e['prefix']}_00001_.png")
|
||
if not os.path.isfile(src):
|
||
print(f" ⚠️ 缺起始图 {src}，跳过 {e['shot']}"); continue
|
||
iname = f"agent_shot_{i+1:04d}.png"
|
||
shutil.copy(src, os.path.join(COMFY_IN, iname))
|
||
prompt = shot_to_ltx_prompt(e.get("ltx_desc", ""), e.get("camera", ""))
|
||
seed = 5000 + i * 29
|
||
prefix = f"eed_vid_{i+1:04d}"
|
||
pid = comfy_submit(ltx_i2v_workflow(iname, prompt, seed, prefix, frames=frames))
|
||
print(f" 🎥 [{i+1}/{len(entries)}] {e['shot']} LTX I2V {frames}帧 提交，预计 {frames*4}s...")
|
||
comfy_wait(pid, timeout=900)
|
||
vout = os.path.join(COMFY_OUT, f"{prefix}_00001_.mp4")
|
||
if not wait_file(vout, timeout=60):
|
||
print(f" ⚠️ 视频未落盘 {vout}")
|
||
continue
|
||
dest = os.path.join(vdir, f"shot_{i+1:04d}.mp4")
|
||
shutil.copy(vout, dest)
|
||
vlist.append({"shot": e["shot"], "file": dest, "prompt": prompt})
|
||
print(f" ✅ {e['shot']} → {dest}")
|
||
vman = os.path.join(vdir, "video_manifest.json")
|
||
with open(vman, "w", encoding="utf-8") as f:
|
||
json.dump({"video_dir": vdir, "videos": vlist}, f, ensure_ascii=False, indent=2)
|
||
print(f" 📦 视频清单: {vman}")
|
||
return vman
|
||
|
||
def stage_compose(vman_or_dir, output="", title="", credits=""):
|
||
if os.path.isfile(vman_or_dir) and vman_or_dir.endswith("video_manifest.json"):
|
||
d = json.load(open(vman_or_dir, encoding="utf-8"))
|
||
folder = d["video_dir"]
|
||
else:
|
||
folder = vman_or_dir
|
||
vids = sorted(glob.glob(os.path.join(folder, "shot_*.mp4")))
|
||
if not vids:
|
||
raise RuntimeError(f"没有可拼接的视频片段: {folder}")
|
||
if not output:
|
||
output = os.path.join(os.path.dirname(folder), "EP01.mp4")
|
||
if len(vids) == 1 and not title and not credits:
|
||
# 单段且无片头片尾：直接拷贝（LTX 视频无声，video_composer 方案B 会因缺音频流失败）
|
||
shutil.copy(vids[0], output)
|
||
print(f" 🎬 单段成片(直接拷贝): {output}")
|
||
return output
|
||
args = [folder, "-o", output]
|
||
if title: args += ["--title", title]
|
||
if credits: args += ["--credits", credits]
|
||
run_script("tools/video_composer.py", args)
|
||
print(f" 🎬 成片: {output}")
|
||
return output
|
||
|
||
def _video_duration(path):
|
||
try:
|
||
r = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
|
||
"-of", "csv=p=0", path], capture_output=True, text=True)
|
||
return float(r.stdout.strip() or 0)
|
||
except Exception:
|
||
return 0
|
||
|
||
def stage_audio_pipeline(sb_json, composed, ep_dir):
|
||
"""⑥ 声画装配：配音 → 字幕 → BGM → 混音 → 有声终成片"""
|
||
from tools import audio_pipeline as ap
|
||
print(" 🎙 配音 (Edge-TTS)...")
|
||
voice = ap.generate_voice(sb_json, ep_dir)
|
||
print(" 📝 字幕...")
|
||
ap.generate_srt(sb_json, ep_dir, voice)
|
||
dur = _video_duration(composed) or 20
|
||
print(f" 🎵 BGM ({dur:.0f}s)...")
|
||
ap.generate_bgm(ep_dir, int(dur) + 2)
|
||
final = os.path.join(ep_dir, "FINAL.mp4")
|
||
print(" 🎛 混音...")
|
||
ap.mix_final(composed, os.path.join(ep_dir, "voice"),
|
||
os.path.join(ep_dir, "bgm.mp3"), final)
|
||
print(f" 🎉 有声终成片: {final}")
|
||
return final
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="苍耳 · 一键短剧 Agent")
|
||
ap.add_argument("input", help="剧本 / 分镜JSON(--from render,route) / 清单(--from video,compose)")
|
||
ap.add_argument("-e", "--episode", type=int, default=1)
|
||
ap.add_argument("--pro", action="store_true", help="分镜用Pro模型")
|
||
ap.add_argument("--style", default="", help="全局风格")
|
||
ap.add_argument("--until", choices=["storyboard","route","render","video","compose","audio"], default="render")
|
||
ap.add_argument("--from", dest="from_stage", choices=["route","render","video","compose","audio"])
|
||
ap.add_argument("--max-shots", type=int, default=0, help="只跑前 N 镜")
|
||
ap.add_argument("--frames", type=int, default=49, help="LTX 视频帧数(默认49)")
|
||
ap.add_argument("-o", "--output", default="", help="成片路径")
|
||
args = ap.parse_args()
|
||
|
||
sb_json = None
|
||
if args.from_stage in ("video", "compose"):
|
||
manifest = args.input
|
||
print(f"📖 从清单续跑: {manifest}")
|
||
if args.from_stage == "video":
|
||
vman = stage_video(manifest, frames=args.frames, max_shots=args.max_shots)
|
||
if args.from_stage == "compose":
|
||
out = stage_compose(manifest, args.output)
|
||
print("🎯 成片完成:", out); return
|
||
if args.from_stage == "audio":
|
||
out = stage_compose(manifest, args.output)
|
||
sb = glob.glob(os.path.join(os.path.dirname(os.path.dirname(out)), "STORYBOARD*.json"))
|
||
if sb:
|
||
ep_dir = os.path.dirname(os.path.abspath(out))
|
||
final = stage_audio_pipeline(sb[0], out, ep_dir)
|
||
print("🎯 有声终成片:", final)
|
||
return
|
||
print("🎯 阶段完成。续跑: --from compose")
|
||
return
|
||
|
||
if args.from_stage:
|
||
sb_json = args.input
|
||
if not os.path.isfile(sb_json):
|
||
print(f"❌ 找不到分镜: {sb_json}"); sys.exit(1)
|
||
print(f"📖 从分镜续跑: {sb_json}")
|
||
else:
|
||
if not os.path.isfile(args.input):
|
||
print(f"❌ 找不到剧本: {args.input}"); sys.exit(1)
|
||
print(f"🚀 苍耳 · 一键短剧 Agent | 剧本: {args.input} | EP{args.episode:02d} | 到: {args.until}")
|
||
print("═" * 44); print("① 剧本→分镜 (豆包, ¥0.01/集 级)"); print("═" * 44)
|
||
sb_json = stage_storyboard(args.input, args.episode, args.pro)
|
||
print(f"✅ 分镜: {sb_json}")
|
||
|
||
if not sb_json: sys.exit(1)
|
||
|
||
if args.from_stage in (None, "storyboard") and args.until not in ("storyboard",):
|
||
print("═" * 44); print("② 逐镜路由分配 (豆包)"); print("═" * 44)
|
||
sb_json = stage_route(sb_json)
|
||
print("✅ 路由完成")
|
||
|
||
if args.until in ("render", "video", "compose"):
|
||
print("═" * 44); print("③ 逐镜出图 (本地 Z-Image, 零成本)"); print("═" * 44)
|
||
manifest = stage_render(sb_json, style=args.style, max_shots=args.max_shots)
|
||
else:
|
||
print("🎯 阶段完成。"); return
|
||
|
||
if args.until in ("video", "compose"):
|
||
print("═" * 44); print("④ 逐镜视频 (本地 LTX I2V, 零成本)"); print("═" * 44)
|
||
vman = stage_video(manifest, frames=args.frames, max_shots=args.max_shots)
|
||
else:
|
||
print("🎯 出图完成。续跑: --from video <清单> 或 --until video")
|
||
return
|
||
|
||
if args.until in ("compose", "audio"):
|
||
print("═" * 44); print("⑤ 剪辑拼接 (FFmpeg)"); print("═" * 44)
|
||
out = stage_compose(vman, args.output)
|
||
print(f"\n🎬 成片: {out}")
|
||
if args.until == "audio":
|
||
print("═" * 44); print("⑥ 声画装配 (配音+字幕+BGM+混音)"); print("═" * 44)
|
||
sb = sb_json or (glob.glob(os.path.join(os.path.dirname(os.path.dirname(out)), "STORYBOARD*.json")) or [None])[0]
|
||
if sb:
|
||
ep_dir = os.path.dirname(os.path.abspath(out))
|
||
final = stage_audio_pipeline(sb, out, ep_dir)
|
||
print(f"\n🎉🎉 有声终成片: {final}")
|
||
else:
|
||
print("⚠️ 未找到分镜JSON，跳过声画装配")
|
||
elif args.until != "compose":
|
||
print("🎯 视频完成。续跑: --from compose <清单> 或 --until compose/audio")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|