175 lines
6.9 KiB
Python
175 lines
6.9 KiB
Python
"""Wan 2.2 TI2V-5B text/image→video via the resident ComfyUI (same pattern as
|
|
comfyui_sd: pure stdlib, ComfyUI's venv does the heavy lifting, model stays
|
|
cached in RAM between jobs). Frames come back as PNGs; ffmpeg muxes the MP4."""
|
|
import argparse
|
|
import json
|
|
import mimetypes
|
|
import os
|
|
import random
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
# job env has a minimal PATH; we shell out to ffmpeg for the mux
|
|
os.environ["PATH"] = "/opt/homebrew/bin:/usr/local/bin:" + os.environ.get("PATH", "")
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
COMFY = ROOT / "vendor" / "comfyui"
|
|
API = "http://127.0.0.1:8188"
|
|
FPS = 24 # Wan 2.2
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--input", action="append", default=[])
|
|
ap.add_argument("--outdir", required=True)
|
|
ap.add_argument("--params", default="{}")
|
|
a = ap.parse_args()
|
|
p = json.loads(a.params)
|
|
outdir = Path(a.outdir)
|
|
|
|
prompt = (p.get("prompt") or "").strip()
|
|
if not prompt:
|
|
print("ERROR: prompt is required")
|
|
sys.exit(1)
|
|
if not (COMFY / "main.py").exists():
|
|
print(f"ERROR: ComfyUI not installed at {COMFY}. Run scripts/install_comfyui.sh")
|
|
sys.exit(1)
|
|
model_file = COMFY / "models" / "diffusion_models" / "wan2.2_ti2v_5B_fp16.safetensors"
|
|
if not model_file.exists():
|
|
print("ERROR: Wan weights missing — see manifest description for the three files")
|
|
sys.exit(1)
|
|
|
|
|
|
def alive():
|
|
try:
|
|
urllib.request.urlopen(f"{API}/", timeout=3)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def ensure_server():
|
|
if alive():
|
|
print("comfyui: already resident", flush=True)
|
|
return
|
|
print("comfyui: starting resident server ...", flush=True)
|
|
log = open("/tmp/comfyui.log", "ab")
|
|
subprocess.Popen(
|
|
[str(COMFY / ".venv" / "bin" / "python"), "main.py", "--port", "8188"],
|
|
cwd=str(COMFY), stdout=log, stderr=log, start_new_session=True)
|
|
for _ in range(90):
|
|
if alive():
|
|
print("comfyui: up", flush=True)
|
|
return
|
|
time.sleep(2)
|
|
print("ERROR: ComfyUI did not come up — see /tmp/comfyui.log")
|
|
sys.exit(1)
|
|
|
|
|
|
def upload_image(path):
|
|
"""POST multipart to /upload/image, return server-side filename."""
|
|
boundary = uuid.uuid4().hex
|
|
fname = Path(path).name
|
|
ctype = mimetypes.guess_type(fname)[0] or "application/octet-stream"
|
|
body = (f"--{boundary}\r\nContent-Disposition: form-data; name=\"image\"; "
|
|
f"filename=\"{fname}\"\r\nContent-Type: {ctype}\r\n\r\n").encode()
|
|
body += Path(path).read_bytes()
|
|
body += f"\r\n--{boundary}--\r\n".encode()
|
|
req = urllib.request.Request(f"{API}/upload/image", data=body, headers={
|
|
"Content-Type": f"multipart/form-data; boundary={boundary}"})
|
|
return json.load(urllib.request.urlopen(req))["name"]
|
|
|
|
|
|
def build(seed, start_image):
|
|
length = int(p.get("length", 49))
|
|
if length % 4 != 1:
|
|
length = (length // 4) * 4 + 1 # model requires 4n+1 frames
|
|
lat = {"vae": ["v", 0], "width": int(p.get("width", 832)),
|
|
"height": int(p.get("height", 480)), "length": length, "batch_size": 1}
|
|
g = {
|
|
"u": {"class_type": "UNETLoader", "inputs": {
|
|
"unet_name": "wan2.2_ti2v_5B_fp16.safetensors", "weight_dtype": "default"}},
|
|
"c": {"class_type": "CLIPLoader", "inputs": {
|
|
"clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", "type": "wan",
|
|
"device": "default"}},
|
|
"v": {"class_type": "VAELoader", "inputs": {"vae_name": "wan2.2_vae.safetensors"}},
|
|
"mo": {"class_type": "ModelSamplingSD3", "inputs": {"model": ["u", 0], "shift": 8.0}},
|
|
"pos": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["c", 0]}},
|
|
"neg": {"class_type": "CLIPTextEncode", "inputs": {
|
|
"text": p.get("negative", "blurry, distorted, low quality, static image, watermark, text"),
|
|
"clip": ["c", 0]}},
|
|
"lat": {"class_type": "Wan22ImageToVideoLatent", "inputs": lat},
|
|
"ks": {"class_type": "KSampler", "inputs": {
|
|
"seed": seed, "steps": int(p.get("steps", 20)), "cfg": float(p.get("cfg", 5.0)),
|
|
"sampler_name": "uni_pc", "scheduler": "simple", "denoise": 1.0,
|
|
"model": ["mo", 0], "positive": ["pos", 0], "negative": ["neg", 0],
|
|
"latent_image": ["lat", 0]}},
|
|
"dec": {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["v", 0]}},
|
|
"sav": {"class_type": "SaveImage", "inputs": {
|
|
"filename_prefix": "mb_wan", "images": ["dec", 0]}},
|
|
}
|
|
if start_image:
|
|
g["img"] = {"class_type": "LoadImage", "inputs": {"image": start_image}}
|
|
lat["start_image"] = ["img", 0]
|
|
return g
|
|
|
|
|
|
ensure_server()
|
|
seed = int(p.get("seed", -1))
|
|
if seed < 0:
|
|
seed = random.randint(0, 2**31 - 1)
|
|
start_image = upload_image(a.input[0]) if a.input else None
|
|
print(f"seed={seed} mode={'i2v' if start_image else 't2v'} "
|
|
f"{p.get('width', 832)}x{p.get('height', 480)}x{p.get('length', 49)}f", flush=True)
|
|
|
|
body = json.dumps({"prompt": build(seed, start_image)}).encode()
|
|
req = urllib.request.Request(f"{API}/prompt", data=body,
|
|
headers={"Content-Type": "application/json"})
|
|
try:
|
|
pid = json.load(urllib.request.urlopen(req))["prompt_id"]
|
|
except urllib.error.HTTPError as e:
|
|
print(f"ERROR: ComfyUI rejected the workflow: {e.read().decode()[:600]}")
|
|
sys.exit(1)
|
|
|
|
t0 = time.time()
|
|
images = []
|
|
while time.time() - t0 < 5400:
|
|
h = json.load(urllib.request.urlopen(f"{API}/history/{pid}"))
|
|
if pid in h:
|
|
st = h[pid].get("status", {})
|
|
if st.get("status_str") == "error":
|
|
print("ERROR: generation failed — check /tmp/comfyui.log")
|
|
print(json.dumps(st)[:400])
|
|
sys.exit(1)
|
|
if h[pid].get("outputs"):
|
|
for node in h[pid]["outputs"].values():
|
|
images += node.get("images", [])
|
|
break
|
|
time.sleep(5)
|
|
|
|
if not images:
|
|
print("ERROR: no frames produced (timeout after 90min)")
|
|
sys.exit(1)
|
|
|
|
frames_dir = outdir / "work" # subdir: keep raw frames out of asset registration
|
|
frames_dir.mkdir(exist_ok=True)
|
|
for i, im in enumerate(images):
|
|
q = urllib.parse.urlencode({"filename": im["filename"],
|
|
"subfolder": im.get("subfolder", ""),
|
|
"type": im.get("type", "output")})
|
|
(frames_dir / f"f{i:05d}.png").write_bytes(
|
|
urllib.request.urlopen(f"{API}/view?{q}").read())
|
|
|
|
out_mp4 = outdir / f"wan_{seed}.mp4"
|
|
r = subprocess.run(["ffmpeg", "-y", "-framerate", str(FPS),
|
|
"-i", str(frames_dir / "f%05d.png"),
|
|
"-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "18",
|
|
str(out_mp4)], capture_output=True, text=True)
|
|
if r.returncode != 0 or not out_mp4.exists():
|
|
print("ERROR: ffmpeg mux failed:", r.stderr[-800:])
|
|
sys.exit(1)
|
|
print(f"done: {len(images)} frames → {out_mp4.name} in {time.time() - t0:.1f}s", flush=True)
|