api.py
13.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
"""MiniMax-H3 文生视频 HTTP API（异步任务式）。
为什么不是同步接口：模型加载 238s、单条视频 7 分钟（50 步 90 分钟），
同步 HTTP 必然超时。所以是「提交任务拿 id → 轮询状态 → 下载文件」。
为什么是串行单 worker：一条生成占 ~12.3G 显存 + ~32G CPU 内存，
一张 22G 卡同时只能跑一条。并发提交会排队，不会 OOM。
"""
import os, sys, time, uuid, threading, queue, traceback, io, urllib.request
from typing import Optional, Literal
from pathlib import Path
BASE = Path("/data/models/minimax-h3")
OUT = BASE / "outputs"; OUT.mkdir(exist_ok=True)
# 一个进程只跑一个分区。两套权重的 DiT、processor、Turbo LoRA 三样都不同，
# 同进程装两套要么重建 pipe（每次重读 16G 权重），要么常驻 77G 内存。
# 每卡一个进程本来就是这台机器唯一的并发方式（H3 不支持模型并行），顺手就把分区分开了。
PARTITION = os.environ.get("H3_PARTITION", "fl2va")
# 没上 GPU 的那部分权重存哪。cpu 是每进程一份 40G 匿名内存，并发被 125G 内存卡在 3 个；
# disk 让权重留在 safetensors 文件里（layers.py 里 offload() 就是 .to("meta")），
# 每次前向从文件读进显存，走 OS 页缓存——多个进程读同一个文件，物理内存里只有一份。
OFFLOAD = os.environ.get("H3_OFFLOAD", "cpu")
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "torch")
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(BASE / "models"))
os.environ.setdefault("MODELSCOPE_CACHE", str(BASE / ".mscache"))
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field, model_validator
def _rows(width: int, height: int, frames: int) -> float:
"""序列行数：潜帧数 x 宽 x 高 / 1024。真正决定显存/算力的量，帧数和分辨率是同一笔预算。"""
latent_frames = (max(frames, 5) - 5) // 17 * 5 + 2
return latent_frames * width * height / 1024
# 22G 卡的显存墙按 rows 算，不是单独卡 frames——实测边界（2026-08-20 long15s.sh）：
# 832x480x124f rows=14430 稳定基准；448x256x362f rows=11984、576x320x362f rows=19260 都能跑
# （后者峰值 14.5G，换入换出吃紧但没崩）；640x384x362f rows=25680 OOM（撞到 20.6G）。
# 20000 卡在 19260（能跑）和 25680（OOM）之间，留安全余量。
ROWS_LIMIT = 20000
# ---------------- 任务模型 ----------------
class GenReq(BaseModel):
prompt: str = Field(min_length=1, max_length=2000)
# 上界取模型自己的默认档 1344x768。pipeline 本身不限分辨率，只把长宽向上取整到 32 的倍数，
# 真正的墙是这张 22G 卡：序列行数按 h*w/1024 每潜帧算，832x480 是 390 行/帧 x 37 帧 = 14430 行，
# 1344x768 是 1008 行/帧 = 37296 行，注意力是平方级，涨 2.6 倍行数就是涨约 6.7 倍算力。
# 832x480 带三张参考图实测峰值 13.7G，往上哪一档撞墙没测过，超了就是这条任务 failed，不会拖死服务。
width: int = Field(832, ge=64, le=1344)
height: int = Field(480, ge=64, le=768)
# 上界 362=15s（协议规格上限），是否真能生成看下面 _check_rows——832x480 撞了还是会被拒。
frames: int = Field(124, ge=8, le=362)
steps: int = Field(4, ge=1, le=50)
turbo: bool = True
seed: int = Field(0, ge=0, le=2**31 - 1)
flow_shift: Optional[float] = None
# 首末帧条件（FL2VA）。给了首帧，这一段就从那张画面长出去——同一间屋子、同一个人。
# 不给就是纯文生视频，屋子和人由模型自己编，跨条对不上。
first_frame_url: Optional[str] = None
last_frame_url: Optional[str] = None
# 参考图（Ref2VA）。首帧只管这一段从哪开始，参考图全程参与注意力，管的是「中间别乱改」。
# 张数没有硬上限，代价是序列变长：一张按 ref_image_short_edge 缩放后折成 token，
# 短边 480 时约 390 行，而目标视频本身才 14430 行。
reference_image_urls: list[str] = []
ref_image_short_edge: int = Field(480, ge=64, le=2048)
@model_validator(mode="after")
def _check_rows(self):
rows = _rows(self.width, self.height, self.frames)
if rows > ROWS_LIMIT:
raise ValueError(f"分辨率×帧数超出显存安全上限（rows={rows:.0f} > {ROWS_LIMIT}）；"
f"降低 frames 或分辨率，例如 448x256/576x320 可到 362 帧，832x480 建议 ≤124 帧")
return self
class Job(BaseModel):
id: str
status: Literal["queued", "loading", "running", "done", "failed"]
req: GenReq
created_at: float
started_at: Optional[float] = None
finished_at: Optional[float] = None
step: int = 0
total_steps: int = 0
file: Optional[str] = None
size_bytes: Optional[int] = None
error: Optional[str] = None
peak_vram_gb: Optional[float] = None
JOBS: dict[str, Job] = {}
Q: "queue.Queue[str]" = queue.Queue()
LOCK = threading.Lock()
STATE = {"model": "loading", "loaded_at": None, "gpu": os.environ.get("CUDA_VISIBLE_DEVICES", "?")}
# ---------------- 推理 worker ----------------
# DiT、processor、Turbo LoRA 三样成套，跨分区混用会对不上层。
# ref2v 的 Turbo LoRA 只有 v0.1，fl2v 那个是 v1.0 768p——这一项两边成色不同。
PARTITIONS = {
"fl2va": {"dit": "minimax-h3-fl2va-nf4.safetensors",
"processor": "FL2VA/processor/",
"lora": "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"},
# ref2v 的 bf16 版实测出 124 帧噪点：同尺寸同 seed，挂上它是噪点、摘掉就出画面，
# 换 flow_shift 无效（见 docs/H3 的二分表）。改用 comfyui 版试另一套键名。
"ref2va": {"dit": "minimax-h3-ref2va-nf4.safetensors",
"processor": "Ref2VA/processor/",
"lora": "minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors"},
}
def _turbo_lora_config():
from diffsynth.pipelines.minimax_h3_audio_video import ModelConfig
return ModelConfig(model_id="lightx2v/Minimax-h3-Turbo",
origin_file_pattern=PARTITIONS[PARTITION]["lora"])
def _build_pipe():
from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig
part = PARTITIONS[PARTITION]
DT = torch.float32 # fp16 在 sm_75 上出全黑视频，见 README
# 常驻显存里留给权重多少 GB，剩下的全是激活。默认公式是按 832x480 调的（留 14G 给激活，实测峰值 13.7G）。
# 分辨率往上走就得调低它：1280x720 在默认值下要 22.3G、卡只有 21.49G，差约 1G 撞 OOM；
# 调低后权重多从内存流几趟，慢一点，换来激活空间。用 H3_VRAM_LIMIT 覆盖。
free = torch.cuda.mem_get_info("cuda")[0] / 1024**3
vram_limit = float(os.environ.get("H3_VRAM_LIMIT", max(4.0, free - 14.0)))
offload_dt = "disk" if OFFLOAD == "disk" else DT
vc = dict(offload_dtype=offload_dt, offload_device=OFFLOAD,
onload_dtype=offload_dt, onload_device=OFFLOAD,
preparing_dtype=DT, preparing_device="cuda",
computation_dtype=DT, computation_device="cuda")
NF4 = "DiffSynth-Studio/MiniMax-H3-NF4"
pipe = MiniMaxH3Pipeline.from_pretrained(
torch_dtype=DT, device="cuda",
model_configs=[ModelConfig(model_id=NF4, origin_file_pattern=f, **vc) for f in
[part["dit"], "minimax-h3-text-encoder-nf4.safetensors",
"video_vae_nf4.safetensors", "audio_vae_nf4.safetensors"]],
processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3",
origin_file_pattern=part["processor"]),
vram_limit=vram_limit)
pipe.load_lora(pipe.dit, _turbo_lora_config())
return pipe
def _fetch_image(url):
from PIL import Image
with urllib.request.urlopen(url, timeout=60) as resp:
return Image.open(io.BytesIO(resp.read())).convert("RGB")
def _keyframes(req):
"""取首末帧图，缩到出片尺寸。都没给就返回 (None, None)，pipeline 退回纯文生视频。"""
from PIL import Image
images, indices = [], []
for index, url in ((0, req.first_frame_url), (-1, req.last_frame_url)):
if url is None:
continue
images.append(_fetch_image(url).resize((req.width, req.height), Image.LANCZOS))
indices.append(index)
return (images, indices) if images else (None, None)
def _references(req):
"""参考图按请求顺序取下来。尺寸交给 pipeline 按 ref_image_short_edge 归一，这里不动。"""
if not req.reference_image_urls:
return None
return [{"type": "image", "image": _fetch_image(url)} for url in req.reference_image_urls]
def worker():
from diffsynth.utils.data.audio_video import write_video_audio
import numpy as np
t0 = time.time()
pipe = _build_pipe()
STATE["model"] = "ready"; STATE["loaded_at"] = time.time()
print(f"[api] 模型就绪，耗时 {time.time()-t0:.0f}s", flush=True)
while True:
jid = Q.get()
j = JOBS[jid]
try:
j.status = "running"; j.started_at = time.time()
j.total_steps = j.req.steps
torch.cuda.reset_peak_memory_stats()
fs = j.req.flow_shift if j.req.flow_shift is not None else (6.0 if j.req.turbo else 12.0)
if not j.req.turbo:
pipe.clear_lora()
keyframes, keyframe_indices = _keyframes(j.req)
video, audio = pipe(prompt=j.req.prompt, height=j.req.height, width=j.req.width,
num_frames=j.req.frames, num_inference_steps=j.req.steps,
seed=j.req.seed, flow_shift=fs,
keyframes=keyframes, keyframe_indices=keyframe_indices,
references=_references(j.req),
ref_image_short_edge=j.req.ref_image_short_edge)
# 产物必须验过才算成功——全黑视频也能写出文件（踩过）
a = np.stack([np.asarray(f.convert("RGB"), dtype=np.float32)
for f in video[:: max(1, len(video)//8)]])
if a.std() < 1.0:
raise RuntimeError(f"生成结果近乎纯色 (std={a.std():.3f})，判为失败")
if torch.is_tensor(audio) and not torch.isfinite(audio).all():
raise RuntimeError("音频含非有限值，判为失败")
out = OUT / f"api_{jid}.mp4"
write_video_audio(video=video, audio=audio, output_path=str(out),
fps=24, audio_sample_rate=32000)
j.file = str(out); j.size_bytes = out.stat().st_size
j.peak_vram_gb = round(torch.cuda.max_memory_allocated()/1024**3, 2)
j.status = "done"
except Exception as e:
j.status = "failed"; j.error = f"{type(e).__name__}: {e}"
traceback.print_exc()
finally:
j.finished_at = time.time()
if not j.req.turbo: # 下一个任务可能要 turbo，把 LoRA 挂回去
try:
pipe.load_lora(pipe.dit, _turbo_lora_config())
except Exception:
traceback.print_exc()
Q.task_done()
# ---------------- HTTP ----------------
app = FastAPI(title="MiniMax-H3 Video API", version="1.0")
@app.on_event("startup")
def _start():
threading.Thread(target=worker, daemon=True).start()
@app.get("/health")
def health():
return {"status": "ok" if STATE["model"] == "ready" else "loading",
"model": STATE["model"], "gpu": STATE["gpu"], "partition": PARTITION,
"offload": OFFLOAD, "queued": Q.qsize(),
"running": sum(1 for j in JOBS.values() if j.status == "running")}
@app.post("/v1/videos", status_code=202)
def create(req: GenReq):
if req.turbo and req.steps > 8:
raise HTTPException(400, "turbo 模式步数应 ≤8（LoRA 是 4 步蒸馏的）；要多步请传 turbo=false")
jid = uuid.uuid4().hex[:12]
with LOCK:
JOBS[jid] = Job(id=jid, status="queued", req=req, created_at=time.time())
Q.put(jid)
# 耗时的自变量是序列行数（潜帧数 x 宽 x 高 / 1024），不是分辨率——分辨率和帧数是同一笔预算。
# 基准点 832x480x124f = 14430 行，4 步实测 692 秒 = 4 x 132 + 163
#（163 是文本编码 + 参考图编码 + VAE 解码，它也跟帧数走，所以一起按行数缩放）。
# 回代：1280x736x56f（15640 行）估 749 / 实测 736；832x480x22f（2730 行）估 131 / 实测 96。
# 注意力是平方级，行数越大这个线性估计偏得越小。
rows = _rows(req.width, req.height, req.frames)
est = round((req.steps * 132 + 163) * rows / 14430)
return {"id": jid, "status": "queued", "queue_position": Q.qsize(),
"eta_seconds": est, "poll": f"/v1/videos/{jid}"}
@app.get("/v1/videos/{jid}")
def get(jid: str):
j = JOBS.get(jid)
if not j: raise HTTPException(404, "无此任务")
d = j.model_dump()
if j.status == "running" and j.started_at:
d["elapsed_seconds"] = round(time.time() - j.started_at)
if j.status == "done":
d["download"] = f"/v1/videos/{jid}/file"
return d
@app.get("/v1/videos/{jid}/file")
def download(jid: str):
j = JOBS.get(jid)
if not j: raise HTTPException(404, "无此任务")
if j.status != "done": raise HTTPException(409, f"任务尚未完成（当前 {j.status}）")
return FileResponse(j.file, media_type="video/mp4", filename=f"{jid}.mp4")
@app.get("/v1/videos")
def listjobs(limit: int = 50):
js = sorted(JOBS.values(), key=lambda x: x.created_at, reverse=True)[:limit]
return {"jobs": [{"id": j.id, "status": j.status, "created_at": j.created_at,
"size_bytes": j.size_bytes, "error": j.error} for j in js]}