generate.py
| 1 | """MiniMax-H3 Turbo LoRA — 4-step text-to-audio-video generation. |
| 2 | |
| 3 | A lightweight LoRA that lets MiniMax-H3 render joint video + stereo audio in |
| 4 | **4 sampling steps** instead of the usual ~20, at a fraction of the wall-clock |
| 5 | cost. This single file is a self-contained generator: it loads the base H3 DiT |
| 6 | plus this LoRA, encodes the prompt with the Qwen3-VL text encoder, runs the |
| 7 | model's native dual-schedule sampler for 4 steps, decodes both streams and muxes |
| 8 | a playable mp4. |
| 9 | |
| 10 | The audio stream runs on its own shifted flow schedule (video shift 12, audio |
| 11 | shift 3); each stream is integrated on its own clock, which is the schedule |
| 12 | semantics MiniMax-H3 was designed around. That is the only non-obvious part of |
| 13 | sampling — everything else is a plain Euler flow sampler. |
| 14 | |
| 15 | Dependencies (see requirements.txt), plus a ComfyUI checkout for the H3 model / |
| 16 | VAE / text-encoder module definitions: |
| 17 | |
| 18 | git clone https://github.com/comfyanonymous/ComfyUI |
| 19 | cd ComfyUI && git checkout 14b05228cef127ce529bc0c08660770d4af3e9a8 |
| 20 | |
| 21 | Base weights come from the official MiniMax-H3 release |
| 22 | (Comfy-Org/MiniMax-H3 on the Hugging Face Hub): the bf16 DiT, the int8 Qwen3-VL |
| 23 | text encoder, and the video + audio VAEs. |
| 24 | |
| 25 | Usage: |
| 26 | python generate.py \ |
| 27 | --comfyui /path/to/ComfyUI \ |
| 28 | --base models/diffusion_models/minimax_h3_fl2va_bf16.safetensors \ |
| 29 | --lora minimax_h3_turbo_4step.safetensors \ |
| 30 | --te models/text_encoders/qwen3vl_32b_minimax_h3_int8_convrot.safetensors \ |
| 31 | --video-vae models/vae/minimax_h3_video_vae_fp16.safetensors \ |
| 32 | --audio-vae models/vae/minimax_h3_audio_vae_fp32.safetensors \ |
| 33 | --prompt "A corgi in a tiny chef hat flipping a pancake, sizzling sounds." \ |
| 34 | --width 1344 --height 768 --frames 124 --out corgi.mp4 |
| 35 | |
| 36 | `minimax_h3_turbo_4step.safetensors` is the trained LoRA; the accompanying |
| 37 | `minimax_h3_turbo_4step_ema.safetensors` is a time-averaged variant — try both, |
| 38 | the trained one tends to be crisper on fast motion, the averaged one smoother. |
| 39 | """ |
| 40 | |
| 41 | import argparse |
| 42 | import math |
| 43 | import os |
| 44 | import subprocess |
| 45 | import sys |
| 46 | import time |
| 47 | import wave |
| 48 | |
| 49 | import torch |
| 50 | import torch.nn.functional as F |
| 51 | |
| 52 | |
| 53 | def log(msg): |
| 54 | print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) |
| 55 | |
| 56 | |
| 57 | # ====================================================================== |
| 58 | # Flow schedule (video shift 12 / audio shift 3, closed-form dual mapping) |
| 59 | # ====================================================================== |
| 60 | SHIFT_VIDEO = 12.0 |
| 61 | SHIFT_AUDIO = 3.0 |
| 62 | |
| 63 | |
| 64 | def shift_sigma(u, shift): |
| 65 | return shift * u / (1.0 + (shift - 1.0) * u) |
| 66 | |
| 67 | |
| 68 | def time_shift_sigma(sigma, from_shift, to_shift): |
| 69 | base = sigma / (from_shift + sigma * (1.0 - from_shift)) |
| 70 | return to_shift * base / (1.0 + (to_shift - 1.0) * base) |
| 71 | |
| 72 | |
| 73 | def time_shift_slope(sigma, from_shift, to_shift): |
| 74 | base = sigma / (from_shift + sigma * (1.0 - from_shift)) |
| 75 | return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / ( |
| 76 | from_shift * (1.0 + (to_shift - 1.0) * base) ** 2) |
| 77 | |
| 78 | |
| 79 | def timesteps(n, shift=SHIFT_VIDEO): |
| 80 | """n-step video sigma grid: ts[0]=1 (pure noise) > ... > ts[n]=0.""" |
| 81 | return [shift_sigma(1.0 - i / n, shift) for i in range(n + 1)] |
| 82 | |
| 83 | |
| 84 | def audio_sigma(sigma_v): |
| 85 | return time_shift_sigma(sigma_v, SHIFT_VIDEO, SHIFT_AUDIO) |
| 86 | |
| 87 | |
| 88 | def audio_slope(sigma_v): |
| 89 | return time_shift_slope(sigma_v, SHIFT_VIDEO, SHIFT_AUDIO) |
| 90 | |
| 91 | |
| 92 | @torch.no_grad() |
| 93 | def sample(vfn, xv, xa, ts): |
| 94 | """4-step Euler on the joint flow. The model returns the audio velocity |
| 95 | already scaled by d(sigma_a)/d(sigma_v), so video steps on its own sigma |
| 96 | delta while audio steps on its own schedule's delta (recovering the raw |
| 97 | audio velocity by dividing out the slope). This dual-clock stepping is the |
| 98 | schedule MiniMax-H3 expects; a single flat step on the video clock would |
| 99 | over/under-shoot the audio stream badly at 4 steps. |
| 100 | """ |
| 101 | for i in range(len(ts) - 1): |
| 102 | ov, oa = vfn(xv, xa, ts[i]) |
| 103 | hv = ts[i + 1] - ts[i] |
| 104 | sl = audio_slope(max(ts[i], 1e-6)) |
| 105 | ha = audio_sigma(ts[i + 1]) - audio_sigma(ts[i]) |
| 106 | xv = xv + hv * ov |
| 107 | xa = xa + ha * (oa / sl) |
| 108 | return xv, xa |
| 109 | |
| 110 | |
| 111 | # ====================================================================== |
| 112 | # Functional forward (out-of-place, mirrors the reference module math) |
| 113 | # ====================================================================== |
| 114 | def _rms(x, weight, eps): |
| 115 | return F.rms_norm(x, (x.shape[-1],), weight, eps) |
| 116 | |
| 117 | |
| 118 | def _attn(attn, x, rope_cos, rope_sin): |
| 119 | s = x.shape[0] |
| 120 | heads, hd = attn.heads, attn.head_dim |
| 121 | q, k, v = attn.qkv_proj(x).split(heads * hd, dim=-1) |
| 122 | q = _rms(q.view(s, heads, hd), attn.q_norm.weight, attn.q_norm.eps) |
| 123 | k = _rms(k.view(s, heads, hd), attn.k_norm.weight, attn.k_norm.eps) |
| 124 | v = v.view(s, heads, hd) |
| 125 | if rope_cos is not None: |
| 126 | c, si = rope_cos[:, None, :], rope_sin[:, None, :] |
| 127 | |
| 128 | def rot(t): |
| 129 | t96 = t[..., :96].float() |
| 130 | x1, x2 = t96[..., :48], t96[..., 48:] |
| 131 | return torch.cat([(x1 * c - x2 * si).to(t.dtype), |
| 132 | (x1 * si + x2 * c).to(t.dtype), |
| 133 | t[..., 96:]], dim=-1) |
| 134 | |
| 135 | q, k = rot(q), rot(k) |
| 136 | q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v)) |
| 137 | out = F.scaled_dot_product_attention(q, k, v) |
| 138 | return attn.out_proj(out.squeeze(0).transpose(0, 1).reshape(s, heads * hd)) |
| 139 | |
| 140 | |
| 141 | def _mlp(mlp, x): |
| 142 | x1, x2 = mlp.fc1(x).chunk(2, dim=-1) |
| 143 | return mlp.fc2(F.silu(x1) * x2) |
| 144 | |
| 145 | |
| 146 | def _refiner(refiner, x): |
| 147 | for blk in refiner.blocks: |
| 148 | x = x + _attn(blk.attn, _rms(x, blk.norm1.weight, blk.norm1.eps), |
| 149 | None, None) |
| 150 | x = x + _mlp(blk.mlp, _rms(x, blk.norm2.weight, blk.norm2.eps)) |
| 151 | return _rms(x, refiner.final_norm.weight, refiner.final_norm.eps) |
| 152 | |
| 153 | |
| 154 | def _apply_mod(h, shift, scale, segments): |
| 155 | parts = [] |
| 156 | for a, b, row in segments: |
| 157 | parts.append(h[a:b] * (1.0 + scale[row].to(h.dtype)) + shift[row].to(h.dtype)) |
| 158 | return torch.cat(parts) |
| 159 | |
| 160 | |
| 161 | def _apply_gate(x, gate, other, segments): |
| 162 | parts = [] |
| 163 | for a, b, row in segments: |
| 164 | parts.append(x[a:b] + other[a:b] * gate[row].to(x.dtype)) |
| 165 | return torch.cat(parts) |
| 166 | |
| 167 | |
| 168 | def _block(blk, h, mods, segments, rope_cos, rope_sin): |
| 169 | sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = mods.unbind(dim=1) |
| 170 | hn = _apply_mod(_rms(h, blk.norm1.weight, blk.norm1.eps), sh_msa, sc_msa, segments) |
| 171 | h = _apply_gate(h, g_msa, _attn(blk.attn, hn, rope_cos, rope_sin), segments) |
| 172 | hn = _apply_mod(_rms(h, blk.norm2.weight, blk.norm2.eps), sh_mlp, sc_mlp, segments) |
| 173 | return _apply_gate(h, g_mlp, _mlp(blk.mlp, hn), segments) |
| 174 | |
| 175 | |
| 176 | class LoRALinear(torch.nn.Module): |
| 177 | """Applies the low-rank update at run time in activation space: |
| 178 | y = base(x) + B(A(x)). Folding it into the (bf16) base weight instead would |
| 179 | round most of the update away when it is small relative to the weight, so we |
| 180 | keep it as a separate matmul — same as how the update is meant to act.""" |
| 181 | |
| 182 | def __init__(self, base, a, b): |
| 183 | super().__init__() |
| 184 | self.base = base |
| 185 | self.a, self.b = a, b # [rank, in], [out, rank]; alpha == rank -> scale 1 |
| 186 | |
| 187 | def forward(self, x): |
| 188 | return self.base(x) + F.linear(F.linear(x, self.a), self.b) |
| 189 | |
| 190 | |
| 191 | # ====================================================================== |
| 192 | # Model load + LoRA (applied at run time, not merged) |
| 193 | # ====================================================================== |
| 194 | def load_model(comfyui, base_path, lora_path, device, offload_adaln): |
| 195 | import comfy.ldm.minimax.model as h3ref |
| 196 | import comfy.ops |
| 197 | import comfy.utils |
| 198 | from safetensors.torch import load_file |
| 199 | |
| 200 | log(f"loading base DiT: {base_path}") |
| 201 | sd = comfy.utils.load_torch_file(base_path) |
| 202 | model = h3ref.MiniMaxH3Model(dtype=torch.bfloat16, device="cpu", |
| 203 | operations=comfy.ops.disable_weight_init) |
| 204 | missing, unexpected = model.load_state_dict(sd, strict=True, assign=True) |
| 205 | assert not missing and not unexpected, (missing[:3], unexpected[:3]) |
| 206 | model.requires_grad_(False) |
| 207 | model.eval() |
| 208 | |
| 209 | for i, blk in enumerate(model.blocks): |
| 210 | blk.to(device) |
| 211 | for mod in (model.token_refiner, model.final_layer, model.condition_proj, |
| 212 | model.video_patch_proj, model.audio_patch_proj, |
| 213 | model.time_embedder, model.rope): |
| 214 | mod.to(device) |
| 215 | |
| 216 | log(f"applying LoRA: {lora_path}") |
| 217 | lora = load_file(lora_path) |
| 218 | names = sorted({k.rsplit(".lora_", 1)[0] for k in lora}) |
| 219 | # adaLN projections are read weight-first (bypassing their module), so their |
| 220 | # LoRA can't ride a wrapper — stash it and add the delta where adaLN is built. |
| 221 | model._adaln_lora = {} # block index -> (a, b) |
| 222 | model._final_adaln_lora = None |
| 223 | n_wrap = 0 |
| 224 | for name in names: |
| 225 | a = lora[name + ".lora_A.weight"].to(device, torch.bfloat16) |
| 226 | b = lora[name + ".lora_B.weight"].to(device, torch.bfloat16) |
| 227 | if name.endswith("adaln_proj.linear"): |
| 228 | if name.startswith("final_layer"): |
| 229 | model._final_adaln_lora = (a, b) |
| 230 | else: |
| 231 | model._adaln_lora[int(name.split(".")[1])] = (a, b) |
| 232 | else: |
| 233 | parent = model.get_submodule(name.rsplit(".", 1)[0]) |
| 234 | setattr(parent, name.rsplit(".", 1)[1], |
| 235 | LoRALinear(model.get_submodule(name), a, b)) |
| 236 | n_wrap += 1 |
| 237 | log(f"LoRA: {n_wrap} wrapped + {len(model._adaln_lora)} adaLN " |
| 238 | f"+ {1 if model._final_adaln_lora else 0} final") |
| 239 | |
| 240 | if offload_adaln: |
| 241 | # The per-layer adaLN projection is huge (2688 -> 96768) but depends only |
| 242 | # on the timestep, of which there are a handful per denoise. Keep it in |
| 243 | # CPU fp32 to save ~13 GB of VRAM; the matmul is cheap at 4 steps. |
| 244 | for blk in model.blocks: |
| 245 | lin = blk.adaln_proj.linear |
| 246 | lin.weight.data = lin.weight.data.float().cpu() |
| 247 | lin.bias.data = lin.bias.data.float().cpu() |
| 248 | return model, h3ref |
| 249 | |
| 250 | |
| 251 | VISUAL_COND_T = 0.999 |
| 252 | |
| 253 | |
| 254 | def timestep_rows(model, sigma_v): |
| 255 | sigma_v = float(max(sigma_v, 1e-6)) |
| 256 | t_v = 1.0 - sigma_v |
| 257 | t_a = 1.0 - time_shift_sigma(sigma_v, model.sigma_shift_video, |
| 258 | model.sigma_shift_audio) |
| 259 | seg_t = {"text": t_v, "video": t_v, "audio": t_a} |
| 260 | unique_t = sorted({t_v, t_a}) |
| 261 | return seg_t, unique_t, {t: i for i, t in enumerate(unique_t)} |
| 262 | |
| 263 | |
| 264 | def adaln_mods(model, unique_t, device, offload, cache): |
| 265 | key = tuple(round(t, 9) for t in unique_t) |
| 266 | if key in cache: |
| 267 | return cache[key] |
| 268 | ts = torch.tensor(unique_t, dtype=torch.float32, device=device) |
| 269 | with torch.no_grad(): |
| 270 | temb = model.time_embedder(ts).float() # [M, 2688] GPU |
| 271 | si = F.silu(temb) |
| 272 | si_base = si.cpu() if offload else si.to(torch.bfloat16) |
| 273 | outs = torch.stack([F.linear(si_base, b.adaln_proj.linear.weight, |
| 274 | b.adaln_proj.linear.bias) |
| 275 | for b in model.blocks]) # [50, M, 96768] |
| 276 | mods = outs.to(device, torch.bfloat16) |
| 277 | if getattr(model, "_adaln_lora", None): |
| 278 | # run-time low-rank delta, on GPU (base built in CPU fp32 under offload) |
| 279 | si_g = si.to(torch.bfloat16) |
| 280 | for idx, (a, b) in model._adaln_lora.items(): |
| 281 | mods[idx] = mods[idx] + F.linear(F.linear(si_g, a), b) |
| 282 | M, H = len(unique_t), model.hidden_size |
| 283 | mods = mods.view(len(model.blocks), M, 3, 6, H).reshape( |
| 284 | len(model.blocks), M * 3, 6, H) |
| 285 | temb_bf = temb.to(torch.bfloat16) |
| 286 | cache[key] = (mods, temb_bf) |
| 287 | return mods, temb_bf |
| 288 | |
| 289 | |
| 290 | class Prepared: |
| 291 | """Static packed-sequence structure for one (text_len, shape) signature.""" |
| 292 | |
| 293 | def __init__(self, model, h3ref, text_len, video_shape, audio_t, tags, |
| 294 | device): |
| 295 | _, _, lt, lh, lw = video_shape |
| 296 | self.video_shape = tuple(video_shape) |
| 297 | self.lat_pad = ((lh + 1) // 2 * 2, (lw + 1) // 2 * 2) |
| 298 | self.layout = h3ref.PackedLayout(text_len, lt, *self.lat_pad, audio_t) |
| 299 | pos = self.layout.position_ids.to(torch.float32).to(device) |
| 300 | inv = model.rope.inv_freq.to(device) |
| 301 | ang = (pos.unsqueeze(-1) * inv.view(1, 1, -1)).flatten(1) |
| 302 | self.rope_cos, self.rope_sin = torch.cos(ang), torch.sin(ang) |
| 303 | |
| 304 | segs = [] |
| 305 | for a, b, kind in self.layout.segments: |
| 306 | if kind == "text" and tags is not None: |
| 307 | tg = tags.view(-1).tolist() |
| 308 | run = 0 |
| 309 | for i in range(1, b - a + 1): |
| 310 | if i == b - a or tg[i] != tg[run]: |
| 311 | segs.append((a + run, a + i, int(tg[run]), kind)) |
| 312 | run = i |
| 313 | else: |
| 314 | tag = {"text": 1, "video": 0, "audio": 2}[kind] |
| 315 | segs.append((a, b, tag, kind)) |
| 316 | self.seg_template = segs |
| 317 | (self.video_seg,) = [(a, b) for a, b, k in self.layout.segments if k == "video"] |
| 318 | (self.audio_seg,) = [(a, b) for a, b, k in self.layout.segments if k == "audio"] |
| 319 | |
| 320 | |
| 321 | @torch.no_grad() |
| 322 | def forward(model, h3ref, prep, video_x, audio_x, sigma_v, context, device, |
| 323 | offload, cache): |
| 324 | """One denoise evaluation in the sigma_v domain. Returns |
| 325 | (video_velocity, audio_velocity * slope), matching what the sampler wants.""" |
| 326 | import comfy.ldm.common_dit |
| 327 | video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, model.patch_size) |
| 328 | orig_t, orig_h, orig_w = prep.video_shape[2:] |
| 329 | |
| 330 | sigma_v = float(max(sigma_v, 1e-6)) |
| 331 | seg_t, unique_t, t_row = timestep_rows(model, sigma_v) |
| 332 | segments = [(a, b, t_row[seg_t[k]] * 3 + tag) |
| 333 | for a, b, tag, k in prep.seg_template] |
| 334 | |
| 335 | base_mods, t_emb = adaln_mods(model, unique_t, device, offload, cache) |
| 336 | silu_temb = F.silu(t_emb) |
| 337 | |
| 338 | video_rows = h3ref.patchify_video(video_x.to(torch.float32), model.patch_size) |
| 339 | audio_rows = h3ref.pack_audio(audio_x.to(torch.float32)) |
| 340 | video_embed = model.video_patch_proj(video_rows).to(torch.bfloat16) |
| 341 | audio_embed = model.audio_patch_proj(audio_rows).to(torch.bfloat16) |
| 342 | |
| 343 | with torch.autocast("cuda", dtype=torch.bfloat16): |
| 344 | text_states = context[0] |
| 345 | if text_states.shape[-1] != model.hidden_size: |
| 346 | text_states = _refiner(model.token_refiner, |
| 347 | model.condition_proj(text_states)) |
| 348 | pieces = [] |
| 349 | for a, b, kind in prep.layout.segments: |
| 350 | if kind == "text": |
| 351 | pieces.append(text_states) |
| 352 | elif kind == "video": |
| 353 | pieces.append(video_embed) |
| 354 | else: |
| 355 | pieces.append(audio_embed) |
| 356 | h = torch.cat(pieces) |
| 357 | for i, blk in enumerate(model.blocks): |
| 358 | h = _block(blk, h, base_mods[i], segments, |
| 359 | prep.rope_cos, prep.rope_sin) |
| 360 | |
| 361 | fl = model.final_layer |
| 362 | with torch.autocast("cuda", dtype=torch.bfloat16): |
| 363 | si_t = F.silu(t_emb) |
| 364 | f_mod = fl.adaln_proj.linear(si_t) |
| 365 | if getattr(model, "_final_adaln_lora", None): |
| 366 | a, b = model._final_adaln_lora |
| 367 | f_mod = f_mod + F.linear(F.linear(si_t, a), b) |
| 368 | f_shift, f_scale = f_mod.view(len(unique_t), 2, model.hidden_size).unbind(1) |
| 369 | (va, vb), (aa, ab) = prep.video_seg, prep.audio_seg |
| 370 | vrow, arow = t_row[seg_t["video"]], t_row[seg_t["audio"]] |
| 371 | hn = _rms(h, fl.norm.weight, fl.norm.eps) |
| 372 | hv = (hn[va:vb] * (1.0 + f_scale[vrow]) + f_shift[vrow]).to(torch.float32) |
| 373 | ha = (hn[aa:ab] * (1.0 + f_scale[arow]) + f_shift[arow]).to(torch.float32) |
| 374 | v_rows, a_rows = fl.video_out(hv), fl.audio_out(ha) |
| 375 | |
| 376 | lt = video_x.shape[2] |
| 377 | video_out = h3ref.unpatchify_video(v_rows, lt, prep.lat_pad[0] // 2, |
| 378 | prep.lat_pad[1] // 2, model.latents_dim, |
| 379 | model.patch_size)[:, :, :orig_t, :orig_h, :orig_w] |
| 380 | audio_out = h3ref.unpack_audio(a_rows) |
| 381 | slope_a = time_shift_slope(sigma_v, model.sigma_shift_video, |
| 382 | model.sigma_shift_audio) |
| 383 | return -video_out.to(video_x.dtype), (-slope_a) * audio_out.to(audio_x.dtype) |
| 384 | |
| 385 | |
| 386 | # ====================================================================== |
| 387 | # Text encode / decode / mux |
| 388 | # ====================================================================== |
| 389 | def encode_prompt(comfyui, te_path, prompt, device): |
| 390 | import comfy.model_management |
| 391 | import comfy.sd |
| 392 | log(f"loading text encoder: {te_path}") |
| 393 | clip = comfy.sd.load_clip([te_path], clip_type=comfy.sd.CLIPType.MINIMAX) |
| 394 | cond = clip.encode_from_tokens_scheduled(clip.tokenize(prompt)) |
| 395 | ca, ex = cond[0][0], cond[0][1] |
| 396 | tags = ex.get("minimax_token_tags") |
| 397 | ctx = ca.to(device, torch.bfloat16) |
| 398 | tags = tags.to(device) if torch.is_tensor(tags) else tags |
| 399 | del clip |
| 400 | comfy.model_management.unload_all_models() |
| 401 | comfy.model_management.soft_empty_cache() |
| 402 | return ctx, tags |
| 403 | |
| 404 | |
| 405 | def _write_wav(path, waveform, sr): |
| 406 | w = waveform.detach().cpu().float() |
| 407 | if w.ndim == 3: |
| 408 | w = w[0] |
| 409 | w = w.clamp(-1.0, 1.0) |
| 410 | ch = w.shape[0] |
| 411 | pcm = (w.transpose(0, 1).contiguous().numpy() * 32767.0).astype("<i2") |
| 412 | with wave.open(path, "wb") as f: |
| 413 | f.setnchannels(ch) |
| 414 | f.setsampwidth(2) |
| 415 | f.setframerate(int(sr)) |
| 416 | f.writeframes(pcm.tobytes()) |
| 417 | |
| 418 | |
| 419 | def save_mp4(images, waveform, sr, fps, out_path): |
| 420 | import imageio.v2 as imageio |
| 421 | import imageio_ffmpeg |
| 422 | frames = images.detach().cpu().float().clamp(0, 1).mul(255).round().to( |
| 423 | torch.uint8).numpy() |
| 424 | tv, ta = out_path + ".v.mp4", out_path + ".a.wav" |
| 425 | writer = imageio.get_writer(tv, fps=fps, codec="libx264", quality=8, |
| 426 | pixelformat="yuv420p", macro_block_size=1, |
| 427 | ffmpeg_log_level="error") |
| 428 | for fr in frames: |
| 429 | writer.append_data(fr) |
| 430 | writer.close() |
| 431 | _write_wav(ta, waveform, sr) |
| 432 | ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() |
| 433 | subprocess.run([ffmpeg, "-y", "-loglevel", "error", "-i", tv, "-i", ta, |
| 434 | "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest", |
| 435 | out_path], check=True) |
| 436 | os.remove(tv) |
| 437 | os.remove(ta) |
| 438 | |
| 439 | |
| 440 | # ====================================================================== |
| 441 | # Main |
| 442 | # ====================================================================== |
| 443 | def main(): |
| 444 | ap = argparse.ArgumentParser(description="MiniMax-H3 Turbo LoRA 4-step generator") |
| 445 | ap.add_argument("--comfyui", required=True, help="path to a ComfyUI checkout @14b05228") |
| 446 | ap.add_argument("--base", required=True, help="H3 bf16 DiT safetensors") |
| 447 | ap.add_argument("--lora", required=True, help="turbo LoRA safetensors") |
| 448 | ap.add_argument("--te", required=True, help="Qwen3-VL text encoder safetensors") |
| 449 | ap.add_argument("--video-vae", required=True) |
| 450 | ap.add_argument("--audio-vae", required=True) |
| 451 | ap.add_argument("--prompt", required=True) |
| 452 | ap.add_argument("--out", default="out.mp4") |
| 453 | ap.add_argument("--width", type=int, default=1344, help="multiple of 16 (canvas is 32-based)") |
| 454 | ap.add_argument("--height", type=int, default=768) |
| 455 | ap.add_argument("--frames", type=int, default=124, help="24 fps; snaps to the 17k+5 grid") |
| 456 | ap.add_argument("--steps", type=int, default=4) |
| 457 | ap.add_argument("--seed", type=int, default=42) |
| 458 | ap.add_argument("--offload-adaln", action="store_true", |
| 459 | help="keep the timestep-projection weights in CPU fp32 (saves ~13GB VRAM)") |
| 460 | args = ap.parse_args() |
| 461 | sys.path.insert(0, args.comfyui) # ComfyUI supplies the H3 module definitions |
| 462 | |
| 463 | dev = "cuda" |
| 464 | frames = args.frames |
| 465 | while frames % 17 != 5: |
| 466 | frames += 1 |
| 467 | lt = (frames - 5) // 17 * 5 + 2 |
| 468 | lh, lw = args.height // 16, args.width // 16 |
| 469 | audio_t = round(frames / 24 * 40) |
| 470 | v_shape, a_shape = (1, 24, lt, lh, lw), (1, 32, 2, audio_t) |
| 471 | ts = timesteps(args.steps) |
| 472 | log(f"{args.width}x{args.height}x{frames}f ({frames/24:.1f}s) -> " |
| 473 | f"video{v_shape} audio{a_shape}; {args.steps}-step grid " |
| 474 | f"{['%.3f' % t for t in ts]}") |
| 475 | |
| 476 | ctx, tags = encode_prompt(args.comfyui, args.te, args.prompt, dev) |
| 477 | model, h3ref = load_model(args.comfyui, args.base, args.lora, dev, |
| 478 | args.offload_adaln) |
| 479 | prep = Prepared(model, h3ref, ctx.shape[1], v_shape, audio_t, tags, dev) |
| 480 | |
| 481 | g = torch.Generator(dev).manual_seed(args.seed) |
| 482 | ga = torch.Generator(dev).manual_seed(args.seed + 1) |
| 483 | nv = torch.randn(v_shape, generator=g, device=dev, dtype=torch.bfloat16) |
| 484 | na = torch.randn(a_shape, generator=ga, device=dev, dtype=torch.bfloat16) |
| 485 | |
| 486 | cache = {} |
| 487 | |
| 488 | def vfn(xv, xa, sv): |
| 489 | return forward(model, h3ref, prep, xv, xa, sv, ctx, dev, |
| 490 | args.offload_adaln, cache) |
| 491 | |
| 492 | log("sampling ...") |
| 493 | t0 = time.time() |
| 494 | with torch.inference_mode(): |
| 495 | zv, za = sample(vfn, nv, na, ts) |
| 496 | log(f"sampled in {time.time()-t0:.1f}s") |
| 497 | |
| 498 | import comfy.sd |
| 499 | import comfy.utils |
| 500 | video_vae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(args.video_vae)) |
| 501 | audio_vae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(args.audio_vae)) |
| 502 | with torch.inference_mode(): |
| 503 | images = video_vae.decode(zv.float()) |
| 504 | if images.ndim == 5: |
| 505 | images = images.reshape(-1, *images.shape[-3:]) |
| 506 | waveform = audio_vae.decode(za.float()).movedim(-1, 1) |
| 507 | std = torch.std(waveform, dim=[1, 2], keepdim=True) * 5.0 |
| 508 | std[std < 1.0] = 1.0 |
| 509 | waveform = waveform / std |
| 510 | sr = getattr(audio_vae, "audio_sample_rate_output", |
| 511 | getattr(audio_vae, "audio_sample_rate", 44100)) |
| 512 | save_mp4(images, waveform, sr, 24, args.out) |
| 513 | log(f"done -> {args.out} ({os.path.getsize(args.out)/2**20:.1f}MB)") |
| 514 | |
| 515 | |
| 516 | if __name__ == "__main__": |
| 517 | main() |
| 518 | |