Spaces:
Running on Zero
Running on Zero
| """MiniMax-H3 image-to-video (I2VA / FLF2VA) — slim ZeroGPU Space for protocol-generator. | |
| Stripped from RioShiina/MiniMax-H3: one Gradio endpoint (`generate_video`) that runs | |
| ComfyUI's `MiniMaxH3ImageToVideo` graph on the NVFP4-pruned FL2VA checkpoint. A start | |
| frame is required; an optional last frame turns the job into FLF2VA. No T2VA, REF2VA, | |
| keyframe guides, ControlNet, LoRA chains or MCP server — see README.md. | |
| """ | |
| import spaces # noqa: F401 — must be imported before torch (ZeroGPU patches CUDA) | |
| import asyncio | |
| import glob | |
| import json | |
| import math | |
| import os | |
| import pathlib | |
| import random | |
| import shutil | |
| import subprocess | |
| import sys | |
| import time | |
| import uuid | |
| from typing import Any | |
| import gradio as gr | |
| from huggingface_hub import HfApi, hf_hub_download | |
| from PIL import Image, ImageOps | |
| # ── Paths / pins ───────────────────────────────────────────────────────────── | |
| APP_DIR = pathlib.Path(__file__).resolve().parent | |
| COMFY = APP_DIR / "ComfyUI" | |
| COMFY_MODELS = COMFY / "models" | |
| COMFY_INPUT = COMFY / "input" | |
| COMFY_OUTPUT = COMFY / "output" | |
| # ComfyUI master on 2026-09-12; requirements.txt mirrors this commit's own pins. | |
| COMFYUI_COMMIT = os.environ.get("COMFYUI_COMMIT", "c75d8c966c29cb0392259af791f43373315b72db") | |
| # Diffusion model conversion. The community NVFP4 file the Space started on lost prompt | |
| # adherence against Wavespeed's open-weights H3 (same prompt, image and settings, | |
| # 2026-10-16); Comfy-Org's own pruned int8_convrot is the file its I2V template | |
| # recommends. "pruned" drops the precomputable AdaLN branches, not quality. The model | |
| # card says to prefer int8_convrot where torch has cu130 and to use fp8_scaled only | |
| # when it can't — switch with the Space variable H3_DIFFUSION_MODEL (int8 | fp8 | nvfp4). | |
| DIFFUSION_MODELS = { | |
| "int8": ("minimax_h3_fl2va_pruned_int8_convrot.safetensors", | |
| "Comfy-Org/MiniMax-H3", "diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors"), | |
| "fp8": ("minimax_h3_fl2va_pruned_fp8_scaled.safetensors", | |
| "Comfy-Org/MiniMax-H3", "diffusion_models/minimax_h3_fl2va_pruned_fp8_scaled.safetensors"), | |
| "nvfp4": ("minimax_h3_fl2va_pruned_nvfp4.safetensors", | |
| "lilcheaty/MiniMax-H3-NVFP4", "minimax_h3_fl2va_pruned_nvfp4.safetensors"), | |
| } | |
| DIFFUSION_MODEL = os.environ.get("H3_DIFFUSION_MODEL", "int8").strip().lower() | |
| if DIFFUSION_MODEL not in DIFFUSION_MODELS: | |
| print(f"[setup] unknown H3_DIFFUSION_MODEL={DIFFUSION_MODEL!r}, using int8", flush=True) | |
| DIFFUSION_MODEL = "int8" | |
| # (models/<subdir>, local filename, HF repo, path in repo) | |
| DOWNLOADS = [ | |
| ("diffusion_models", *DIFFUSION_MODELS[DIFFUSION_MODEL]), | |
| ("text_encoders", "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", | |
| "Comfy-Org/MiniMax-H3", "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"), | |
| ("vae", "minimax_h3_video_vae_fp16.safetensors", | |
| "Comfy-Org/MiniMax-H3", "vae/minimax_h3_video_vae_fp16.safetensors"), | |
| ("vae", "minimax_h3_audio_vae_fp32.safetensors", | |
| "Comfy-Org/MiniMax-H3", "vae/minimax_h3_audio_vae_fp32.safetensors"), | |
| ] | |
| UNET_NAME = DOWNLOADS[0][1] | |
| CLIP_NAME = DOWNLOADS[1][1] | |
| VIDEO_VAE_NAME = DOWNLOADS[2][1] | |
| AUDIO_VAE_NAME = DOWNLOADS[3][1] | |
| # Override with Space secret/variable H3_JOB_DATASET. Requires HF_TOKEN (write) on the Space. | |
| H3_JOB_DATASET = os.environ.get("H3_JOB_DATASET", "igitagoz/h3-job-outputs").strip() | |
| def _flag(name: str, default: str) -> bool: | |
| return os.environ.get(name, default).strip().lower() in ("1", "true", "yes") | |
| # Load the four checkpoints into CPU RAM once at startup so the forked GPU worker | |
| # inherits them instead of re-reading ~34 GB from disk inside every job's GPU window. | |
| # On 2026-10-16 that read, plus the text encode, took 94 s of a 248 s job — 38% of the | |
| # quota. Default on; H3_PRELOAD_MODELS=0 goes back to loading inside the window. | |
| PRELOAD_MODELS = _flag("H3_PRELOAD_MODELS", "1") | |
| # Go one step further and place the preloaded weights on `cuda` at module level, the | |
| # way ZeroGPU's docs say to (the CUDA emulation outside @spaces.GPU holds them and | |
| # restores them on allocation — this is what the Wan Space does with its diffusers | |
| # pipeline). Here it means asking ComfyUI's model management to load the four patchers | |
| # in the parent process. Opt-in until it has been seen to work with ComfyUI under the | |
| # emulation; a failure falls back to the CPU preload. | |
| GPU_RESIDENT = PRELOAD_MODELS and _flag("H3_GPU_RESIDENT", "0") | |
| # SageAttention is quantised attention (INT8 QK) — faster, near-lossless, not bit-exact. | |
| # Off unless asked for, so the default output stays ComfyUI's reference attention. | |
| SAGE_ATTENTION = _flag("H3_SAGE_ATTENTION", "0") | |
| # ── Model geometry (mirrored in src/shared/pipelineCost.ts — keep in lockstep) ── | |
| FPS = 24 | |
| CANVAS_MULTIPLE = 32 | |
| MIN_SECONDS = 1.0 | |
| MAX_SECONDS = 15.0 | |
| MIN_STEPS = 1 | |
| MAX_STEPS = 50 | |
| MIN_MAX_DIM = 256 | |
| MAX_MAX_DIM = 1920 | |
| # ZeroGPU declaration = kill timeout. Calibrated on 2026-10-16 from a 960×640 / 124 | |
| # frames / 20 steps job that held the GPU for 248 s: 94 s before the first sampling | |
| # step (worker start, checkpoint load, text encode) and 151 s for the 20 steps plus | |
| # the video + audio VAE decode. The old anchor (upstream's "544p / 20 steps / 5 s fits | |
| # in 120 s") declared 205 s for that job. BASE covers the load phase while the models | |
| # are loaded inside the window; once the preload is confirmed to cut it, lower BASE — | |
| # and `estimateH3SpaceAllocSecsFromDims` in src/shared/pipelineCost.ts in lockstep. | |
| # The app fits a residual on top for the *cost* figure. | |
| MIN_GPU_SECONDS = int(os.environ.get("MIN_GPU_SECONDS", "60")) | |
| MAX_GPU_SECONDS = int(os.environ.get("MAX_GPU_SECONDS", "1800")) | |
| H3_REF_WORK = 960 * 544 * 124 | |
| H3_BASE_SECS = 100.0 | |
| H3_PER_STEP_SECS = 6.5 | |
| H3_WORK_EXPONENT = 1.3 | |
| H3_SAFETY = 1.25 | |
| def frame_count(seconds: float) -> int: | |
| """24 fps frame count snapped *up* to the model's 17k+5 grid: 5, 22, …, 124 (≈5s), 362 (≈15s).""" | |
| raw = int(round(float(seconds) * FPS)) | |
| if raw <= 5: | |
| return 5 | |
| return 5 + 17 * math.ceil((raw - 5) / 17) | |
| def fit_canvas(image_w: int, image_h: int, max_dim: int) -> tuple[int, int]: | |
| """Longest side = max_dim, short side follows the input aspect ratio, both snapped to 32.""" | |
| long_side = max(CANVAS_MULTIPLE, int(max_dim)) | |
| scale = long_side / max(image_w, image_h) | |
| def snap(v: float) -> int: | |
| return max(CANVAS_MULTIPLE, round(v * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE) | |
| return snap(image_w), snap(image_h) | |
| def estimate_gpu_seconds(width: int, height: int, frames: int, steps: int) -> int: | |
| work = (width * height * frames) / H3_REF_WORK | |
| est = (H3_BASE_SECS + H3_PER_STEP_SECS * steps * work ** H3_WORK_EXPONENT) * H3_SAFETY | |
| return max(MIN_GPU_SECONDS, int(math.ceil(est))) | |
| # ── SageAttention (optional; wheel from RioShiina's ZeroGPU build) ─────────── | |
| if SAGE_ATTENTION and "--use-sage-attention" not in sys.argv: | |
| sys.argv.append("--use-sage-attention") | |
| _sage_patched = False | |
| def apply_sage_attention_patch() -> str: | |
| global _sage_patched | |
| if not SAGE_ATTENTION: | |
| return "sage attention off (H3_SAGE_ATTENTION=0) — default attention" | |
| if _sage_patched: | |
| return "sage attention already enabled" | |
| try: | |
| from comfy import model_management | |
| import sageattention # noqa: F401 | |
| model_management.sage_attention_enabled = lambda: True | |
| model_management.pytorch_attention_enabled = lambda: False | |
| _sage_patched = True | |
| return "sage attention enabled" | |
| except ImportError: | |
| return "sageattention not installed — default attention" | |
| except Exception as exc: # pragma: no cover | |
| return f"sage attention patch failed: {exc}" | |
| # ── ComfyUI bootstrap ──────────────────────────────────────────────────────── | |
| def _run(cmd: list[str], cwd: pathlib.Path | None = None, check: bool = True) -> None: | |
| print(f"[setup] $ {' '.join(cmd)}", flush=True) | |
| subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=check) | |
| def ensure_comfy() -> None: | |
| if not COMFY.exists(): | |
| _run(["git", "clone", "--depth", "1", "https://github.com/comfyanonymous/ComfyUI.git", str(COMFY)]) | |
| _run(["git", "fetch", "--depth", "1", "origin", COMFYUI_COMMIT], cwd=COMFY, check=False) | |
| _run(["git", "checkout", COMFYUI_COMMIT], cwd=COMFY, check=False) | |
| for sub in {d[0] for d in DOWNLOADS}: | |
| (COMFY_MODELS / sub).mkdir(parents=True, exist_ok=True) | |
| COMFY_INPUT.mkdir(parents=True, exist_ok=True) | |
| COMFY_OUTPUT.mkdir(parents=True, exist_ok=True) | |
| def ensure_models(progress: gr.Progress | None = None) -> None: | |
| token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") | |
| for index, (sub, name, repo, repo_path) in enumerate(DOWNLOADS): | |
| dest = COMFY_MODELS / sub / name | |
| if dest.is_symlink() and not dest.exists(): | |
| dest.unlink() | |
| if dest.exists(): | |
| continue | |
| if progress: | |
| progress(index / len(DOWNLOADS), desc=f"downloading {name}") | |
| print(f"[setup] downloading {repo}/{repo_path}", flush=True) | |
| cached = hf_hub_download(repo_id=repo, filename=repo_path, token=token) | |
| dest.parent.mkdir(parents=True, exist_ok=True) | |
| os.symlink(cached, dest) | |
| _nodes_ready = False | |
| NODE_CLASS_MAPPINGS: dict[str, Any] = {} | |
| def init_comfy_nodes() -> None: | |
| """Import ComfyUI in-process (cwd = ComfyUI/ so folder_paths resolves models/, input/, output/).""" | |
| global _nodes_ready, NODE_CLASS_MAPPINGS | |
| if _nodes_ready: | |
| return | |
| comfy_path = str(COMFY) | |
| sys.path = [p for p in sys.path if p != comfy_path] | |
| sys.path.insert(0, comfy_path) | |
| for module_name in list(sys.modules): | |
| if module_name == "utils" or module_name.startswith("utils."): | |
| del sys.modules[module_name] | |
| os.chdir(COMFY) | |
| import execution | |
| import nodes | |
| import server | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| server_instance = server.PromptServer(loop) | |
| execution.PromptQueue(server_instance) | |
| loop.run_until_complete(nodes.init_extra_nodes()) | |
| NODE_CLASS_MAPPINGS = nodes.NODE_CLASS_MAPPINGS | |
| for required in ("MiniMaxH3ImageToVideo", "VAEDecodeAudio", "CreateVideo", "SaveVideo"): | |
| if required not in NODE_CLASS_MAPPINGS: | |
| raise RuntimeError(f"ComfyUI @ {COMFYUI_COMMIT} has no '{required}' node") | |
| _nodes_ready = True | |
| print("[setup] ComfyUI nodes ready", flush=True) | |
| # ── Workflow (the FL2VA recipe, inlined) ──────────────────────────────────── | |
| class _DummyHidden: | |
| extra_pnginfo = None | |
| prompt = None | |
| def node_template(class_type: str) -> dict[str, Any]: | |
| """Node dict pre-filled with every input's declared default, like the ComfyUI API format.""" | |
| if class_type not in NODE_CLASS_MAPPINGS: | |
| raise ValueError(f"node class '{class_type}' not found") | |
| node_class = NODE_CLASS_MAPPINGS[class_type] | |
| input_types = node_class.INPUT_TYPES() | |
| inputs: dict[str, Any] = {} | |
| for name, details in {**input_types.get("required", {}), **input_types.get("optional", {})}.items(): | |
| config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {} | |
| inputs[name] = config.get("default") | |
| return {"inputs": inputs, "class_type": class_type, "_meta": {"title": class_type}} | |
| def build_workflow( | |
| *, | |
| prompt: str, | |
| width: int, | |
| height: int, | |
| frames: int, | |
| steps: int, | |
| seed: int, | |
| first_frame_name: str, | |
| last_frame_name: str | None, | |
| filename_prefix: str, | |
| ) -> dict[str, Any]: | |
| wf: dict[str, Any] = {} | |
| def add(node_id: str, class_type: str, title: str, **inputs: Any) -> str: | |
| node = node_template(class_type) | |
| node["_meta"]["title"] = title | |
| node["inputs"].update(inputs) | |
| wf[node_id] = node | |
| return node_id | |
| add("unet", "UNETLoader", "Load Diffusion Model", unet_name=UNET_NAME, weight_dtype="default") | |
| add("clip", "CLIPLoader", "Load CLIP", clip_name=CLIP_NAME, type="minimax", device="default") | |
| add("vvae", "VAELoader", "Load Video VAE", vae_name=VIDEO_VAE_NAME) | |
| add("avae", "VAELoader", "Load Audio VAE", vae_name=AUDIO_VAE_NAME) | |
| add("ff_load", "LoadImage", "Load First Frame", image=first_frame_name) | |
| # lanczos, not nearest-exact: the keyframes are usually downscaled to the canvas and | |
| # nearest sampling aliases them (the Wan Space resizes its keyframes with LANCZOS). | |
| add("ff_scale", "ImageScale", "Scale First Frame", image=["ff_load", 0], | |
| upscale_method="lanczos", width=width, height=height, crop="center") | |
| h3_inputs: dict[str, Any] = dict( | |
| clip=["clip", 0], vae=["vvae", 0], prompt=prompt, | |
| width=width, height=height, length=frames, first_frame=["ff_scale", 0], | |
| ) | |
| if last_frame_name: | |
| add("lf_load", "LoadImage", "Load Last Frame", image=last_frame_name) | |
| add("lf_scale", "ImageScale", "Scale Last Frame", image=["lf_load", 0], | |
| upscale_method="lanczos", width=width, height=height, crop="center") | |
| h3_inputs["last_frame"] = ["lf_scale", 0] | |
| add("h3", "MiniMaxH3ImageToVideo", "MiniMax H3 Image to Video", **h3_inputs) | |
| add("guider", "BasicGuider", "Basic Guider", model=["unet", 0], conditioning=["h3", 0]) | |
| add("sampler", "KSamplerSelect", "KSamplerSelect", sampler_name="res_multistep") | |
| add("sigmas", "BasicScheduler", "BasicScheduler", model=["unet", 0], | |
| scheduler="simple", steps=steps, denoise=1.0) | |
| add("noise", "RandomNoise", "RandomNoise", noise_seed=seed) | |
| add("sample", "SamplerCustomAdvanced", "SamplerCustomAdvanced", noise=["noise", 0], | |
| guider=["guider", 0], sampler=["sampler", 0], sigmas=["sigmas", 0], latent_image=["h3", 1]) | |
| add("vdec", "VAEDecode", "VAE Decode", samples=["sample", 0], vae=["vvae", 0]) | |
| add("adec", "VAEDecodeAudio", "VAE Decode Audio", samples=["sample", 0], vae=["avae", 0]) | |
| add("video", "CreateVideo", "Create Video", images=["vdec", 0], audio=["adec", 0], fps=FPS) | |
| add("save", "SaveVideo", "Save Video", video=["video", 0], filename_prefix=filename_prefix, | |
| format="mp4", codec={"codec": "h264"}) | |
| return wf | |
| # ── Executor (upstream's minimal topological runner — no PromptServer round-trip) ── | |
| def _value_at(obj: Any, index: int) -> Any: | |
| try: | |
| return obj[index] | |
| except (KeyError, IndexError, TypeError): | |
| try: | |
| return obj["result"][index] | |
| except (KeyError, IndexError, TypeError): | |
| return None | |
| def _topological_order(workflow: dict[str, Any]) -> list[str]: | |
| from collections import defaultdict, deque | |
| graph: dict[str, list[str]] = defaultdict(list) | |
| in_degree = {node_id: 0 for node_id in workflow} | |
| for node_id, node in workflow.items(): | |
| for value in node.get("inputs", {}).values(): | |
| if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in workflow: | |
| graph[value[0]].append(node_id) | |
| in_degree[node_id] += 1 | |
| queue = deque(n for n, d in in_degree.items() if d == 0) | |
| ordered: list[str] = [] | |
| while queue: | |
| current = queue.popleft() | |
| ordered.append(current) | |
| for nxt in graph[current]: | |
| in_degree[nxt] -= 1 | |
| if in_degree[nxt] == 0: | |
| queue.append(nxt) | |
| if len(ordered) != len(workflow): | |
| raise RuntimeError("workflow contains a cycle") | |
| return ordered | |
| def execute_workflow(workflow: dict[str, Any], initial_objects: dict[str, Any] | None = None) -> Any: | |
| import torch | |
| computed: dict[str, Any] = dict(initial_objects or {}) | |
| with torch.no_grad(): | |
| order = _topological_order(workflow) | |
| for node_id in order: | |
| if node_id in computed: | |
| continue | |
| node = workflow[node_id] | |
| node_class = NODE_CLASS_MAPPINGS.get(node["class_type"]) | |
| if node_class is None: | |
| raise RuntimeError(f"unknown node class '{node['class_type']}'") | |
| instance = node_class() | |
| for target in (node_class, instance): | |
| if getattr(target, "hidden", None) is None: | |
| try: | |
| setattr(target, "hidden", _DummyHidden()) | |
| except Exception: | |
| pass | |
| kwargs: dict[str, Any] = {} | |
| for name, value in node["inputs"].items(): | |
| if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in workflow: | |
| src_id, out_idx = value | |
| if src_id not in computed: | |
| raise RuntimeError(f"output of node {src_id} needed by {node_id} but not computed") | |
| actual = _value_at(computed[src_id], out_idx) | |
| else: | |
| actual = value | |
| if "." in name: | |
| parent, child = name.split(".", 1) | |
| kwargs.setdefault(parent, {})[child] = actual | |
| else: | |
| kwargs[name] = actual | |
| fn = getattr(instance, getattr(node_class, "FUNCTION")) | |
| print(f"[run] {node_id} ({node['class_type']})", flush=True) | |
| computed[node_id] = fn(**kwargs) | |
| return computed.get("save") | |
| def extract_video_path(out: Any) -> str | None: | |
| """Locate the mp4 SaveVideo wrote (ui payload, object attrs, or newest file in output/).""" | |
| found: list[str] = [] | |
| def visit(item: Any) -> None: | |
| if item is None or found: | |
| return | |
| if isinstance(item, (list, tuple)): | |
| for sub in item: | |
| visit(sub) | |
| return | |
| if isinstance(item, str): | |
| if os.path.exists(item): | |
| found.append(item) | |
| return | |
| if isinstance(item, dict): | |
| for key in ("images", "videos", "video"): | |
| for sub in item.get("ui", {}).get(key, []): | |
| if isinstance(sub, dict) and sub.get("filename"): | |
| candidate = COMFY_OUTPUT / sub.get("subfolder", "") / sub["filename"] | |
| if candidate.exists(): | |
| found.append(str(candidate)) | |
| return | |
| return | |
| ui = getattr(item, "ui", None) | |
| if ui is not None: | |
| visit(ui if isinstance(ui, dict) else getattr(ui, "as_dict", lambda: None)()) | |
| for attr in ("path", "saved_path", "filepath", "filename", "full_path"): | |
| val = getattr(item, attr, None) | |
| if isinstance(val, str): | |
| for candidate in (pathlib.Path(val), COMFY_OUTPUT / val): | |
| if candidate.exists(): | |
| found.append(str(candidate)) | |
| return | |
| visit(out) | |
| if found: | |
| return found[0] | |
| files = glob.glob(str(COMFY_OUTPUT / "**" / "*.mp4"), recursive=True) | |
| if not files: | |
| return None | |
| files.sort(key=os.path.getmtime, reverse=True) | |
| return files[0] | |
| # ── Preload of the four loaders (see PRELOAD_MODELS / GPU_RESIDENT) ───────── | |
| PRELOADED: dict[str, Any] = {} | |
| LOADER_NODES = ("unet", "clip", "vvae", "avae") | |
| def preload_models() -> None: | |
| """Run the four loader nodes in the parent process (outside any GPU window).""" | |
| template = build_workflow(prompt="", width=960, height=544, frames=124, steps=1, seed=0, | |
| first_frame_name="", last_frame_name=None, filename_prefix="") | |
| t0 = time.time() | |
| for node_id in LOADER_NODES: | |
| node = template[node_id] | |
| node_class = NODE_CLASS_MAPPINGS[node["class_type"]] | |
| instance = node_class() | |
| fn = getattr(instance, getattr(node_class, "FUNCTION")) | |
| print(f"[setup] preloading {node_id} ({node['class_type']})", flush=True) | |
| PRELOADED[node_id] = fn(**node["inputs"]) | |
| print(f"[setup] preloaded {len(PRELOADED)} loaders into CPU RAM in {time.time() - t0:.0f}s", flush=True) | |
| def make_models_gpu_resident() -> None: | |
| """Ask ComfyUI to place the preloaded patchers on the (emulated) GPU at module level. | |
| ZeroGPU then holds those tensors and restores them into the real GPU at allocation, | |
| so the worker starts with the weights in place instead of copying them up from the | |
| parent's RAM inside the billed window. | |
| """ | |
| from comfy import model_management | |
| patchers = [] | |
| for node_id in LOADER_NODES: | |
| obj = PRELOADED[node_id][0] | |
| patchers.append(getattr(obj, "patcher", obj)) | |
| t0 = time.time() | |
| try: | |
| model_management.load_models_gpu(patchers, force_full_load=True) | |
| except TypeError: # older signature without force_full_load | |
| model_management.load_models_gpu(patchers) | |
| print(f"[setup] {len(patchers)} patchers placed on cuda in {time.time() - t0:.0f}s", flush=True) | |
| # ── Durable outputs (same scheme as the Wan / LTX Spaces) ─────────────────── | |
| def persist_job_output(job_id: str, video_path: str, seed: int) -> None: | |
| """Upload the finished mp4 (+ tiny sidecar) to a Hub dataset keyed by job_id so a | |
| client that lost the SSE stream can later download `{job_id}.mp4`.""" | |
| job_id = (job_id or "").strip() | |
| if not job_id or not video_path or not os.path.isfile(video_path): | |
| return | |
| token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") | |
| if not token: | |
| print(f"[persist] skip {job_id}: no HF_TOKEN on Space", flush=True) | |
| return | |
| if not H3_JOB_DATASET: | |
| print(f"[persist] skip {job_id}: H3_JOB_DATASET empty", flush=True) | |
| return | |
| try: | |
| api = HfApi(token=token) | |
| api.upload_file(path_or_fileobj=video_path, path_in_repo=f"{job_id}.mp4", | |
| repo_id=H3_JOB_DATASET, repo_type="dataset") | |
| meta = json.dumps({"job_id": job_id, "seed": int(seed)}, separators=(",", ":")) | |
| api.upload_file(path_or_fileobj=meta.encode("utf-8"), path_in_repo=f"{job_id}.json", | |
| repo_id=H3_JOB_DATASET, repo_type="dataset") | |
| print(f"[persist] uploaded {job_id} → dataset {H3_JOB_DATASET}", flush=True) | |
| except Exception as exc: | |
| print(f"[persist] failed for {job_id}: {exc}", flush=True) | |
| # ── Job helpers ────────────────────────────────────────────────────────────── | |
| def stage_input_image(path: str, tag: str) -> tuple[str, int, int]: | |
| """Copy an upload into ComfyUI/input as an oriented RGB PNG; returns (name, w, h).""" | |
| with Image.open(path) as img: | |
| img = ImageOps.exif_transpose(img).convert("RGB") | |
| name = f"h3_{tag}_{uuid.uuid4().hex}.png" | |
| img.save(COMFY_INPUT / name, "PNG") | |
| return name, img.width, img.height | |
| def prune_old_files(directory: pathlib.Path, max_age_secs: int) -> None: | |
| cutoff = time.time() - max_age_secs | |
| for p in directory.rglob("*"): | |
| try: | |
| if p.is_file() and p.stat().st_mtime < cutoff: | |
| p.unlink() | |
| except OSError: | |
| pass | |
| def _gpu_budget(workflow: dict[str, Any], budget: int, progress: gr.Progress | None = None) -> int: | |
| return int(budget) | |
| def run_workflow_gpu(workflow: dict[str, Any], budget: int, progress: gr.Progress | None = None) -> str | None: | |
| # PRELOADED is read from the module global, never passed in: ZeroGPU pickles every | |
| # argument through a queue to the worker, and ComfyUI's model objects can't be pickled | |
| # ("Can't get local object 'mixed_precision_ops.<locals>.MixedPrecisionOps.Linear'", | |
| # 2026-10-16). The worker is forked from this process, so it inherits the global. | |
| print(f"[gpu] {apply_sage_attention_patch()} — budget {budget}s, " | |
| f"preloaded={','.join(PRELOADED) or 'none'}", flush=True) | |
| out = execute_workflow(workflow, PRELOADED or None) | |
| return extract_video_path(out) | |
| def generate_video( | |
| first_frame: str | None, | |
| last_frame: str | None, | |
| prompt: str, | |
| max_dim: float, | |
| seconds: float, | |
| steps: float, | |
| seed: float, | |
| randomize_seed: bool, | |
| job_id: str = "", | |
| progress: gr.Progress = gr.Progress(track_tqdm=True), | |
| ) -> tuple[str, int, str, int]: | |
| """api_name=generate_video — returns (video, used_seed, job_id, duration_ms).""" | |
| t0 = time.time() | |
| if not first_frame: | |
| raise gr.Error("upload a start image first") | |
| if not str(prompt or "").strip(): | |
| raise gr.Error("write a prompt first") | |
| try: | |
| max_dim_i = int(max_dim) | |
| seconds_f = float(seconds) | |
| steps_i = int(steps) | |
| except (TypeError, ValueError): | |
| raise gr.Error("max_dim, seconds and steps must be numbers") | |
| if not (MIN_MAX_DIM <= max_dim_i <= MAX_MAX_DIM): | |
| raise gr.Error(f"max_dim must be between {MIN_MAX_DIM} and {MAX_MAX_DIM}") | |
| if not (MIN_SECONDS <= seconds_f <= MAX_SECONDS): | |
| raise gr.Error(f"seconds must be between {MIN_SECONDS:g} and {MAX_SECONDS:g}") | |
| if not (MIN_STEPS <= steps_i <= MAX_STEPS): | |
| raise gr.Error(f"steps must be between {MIN_STEPS} and {MAX_STEPS}") | |
| used_seed = random.randint(0, 2**32 - 1) if randomize_seed or int(seed) < 0 else int(seed) | |
| progress(0.02, desc="Preparing inputs") | |
| prune_old_files(COMFY_INPUT, 6 * 3600) | |
| prune_old_files(COMFY_OUTPUT, 6 * 3600) | |
| ff_name, in_w, in_h = stage_input_image(first_frame, "first") | |
| lf_name = stage_input_image(last_frame, "last")[0] if last_frame else None | |
| width, height = fit_canvas(in_w, in_h, max_dim_i) | |
| frames = frame_count(seconds_f) | |
| budget = estimate_gpu_seconds(width, height, frames, steps_i) | |
| if budget > MAX_GPU_SECONDS: | |
| raise gr.Error( | |
| f"this job needs ~{budget}s of GPU time, above the Space cap of {MAX_GPU_SECONDS}s — " | |
| f"lower the resolution ({width}×{height}), duration ({frames} frames) or steps ({steps_i})" | |
| ) | |
| print(f"[job] {job_id or '-'} {width}x{height} frames={frames} steps={steps_i} seed={used_seed} " | |
| f"last_frame={'yes' if lf_name else 'no'} budget={budget}s", flush=True) | |
| progress(0.05, desc="Checking models") | |
| ensure_models(progress) | |
| workflow = build_workflow( | |
| prompt=str(prompt), width=width, height=height, frames=frames, steps=steps_i, seed=used_seed, | |
| first_frame_name=ff_name, last_frame_name=lf_name, | |
| filename_prefix=f"video/h3_{job_id.strip() or uuid.uuid4().hex}", | |
| ) | |
| progress(0.1, desc=f"Generating ({width}×{height}, {frames} frames, {steps_i} steps)") | |
| try: | |
| video = run_workflow_gpu(workflow, budget, progress=progress) | |
| finally: | |
| for name in (ff_name, lf_name): | |
| if name: | |
| try: | |
| (COMFY_INPUT / name).unlink() | |
| except OSError: | |
| pass | |
| if not video: | |
| raise gr.Error("generation finished without an output video — check the Space logs") | |
| persist_job_output(job_id, video, used_seed) | |
| duration_ms = int((time.time() - t0) * 1000) | |
| print(f"[job] {job_id or '-'} done in {duration_ms}ms → {video}", flush=True) | |
| return video, int(used_seed), job_id or "", duration_ms | |
| # ── Startup ────────────────────────────────────────────────────────────────── | |
| print(f"[setup] ComfyUI bootstrap — diffusion model {DIFFUSION_MODEL} ({UNET_NAME})", flush=True) | |
| ensure_comfy() | |
| ensure_models() | |
| init_comfy_nodes() | |
| print(f"[setup] {apply_sage_attention_patch()}", flush=True) | |
| if PRELOAD_MODELS: | |
| try: | |
| preload_models() | |
| except Exception as exc: # fall back to loading inside the GPU window | |
| PRELOADED.clear() | |
| print(f"[setup] preload failed, loaders will run on GPU: {exc}", flush=True) | |
| if PRELOADED and GPU_RESIDENT: | |
| try: | |
| make_models_gpu_resident() | |
| except Exception as exc: # keep the CPU preload | |
| print(f"[setup] cuda placement failed, weights stay in CPU RAM: {exc}", flush=True) | |
| else: | |
| print("[setup] H3_PRELOAD_MODELS=0 — loaders run inside each job's GPU window", flush=True) | |
| # ── UI (the visible form IS the API — inputs are positional in this order) ── | |
| with gr.Blocks(title="MiniMax-H3 I2V") as demo: | |
| gr.Markdown( | |
| "# MiniMax-H3 · image-to-video\n" | |
| "Start image required; add a last image for first-to-last-frame. Native audio. " | |
| "Slim fork of [RioShiina/MiniMax-H3](https://huggingface.co/spaces/RioShiina/MiniMax-H3) " | |
| "for protocol-generator — I2VA / FLF2VA only." | |
| ) | |
| with gr.Row(): | |
| first_frame = gr.Image(type="filepath", label="Start image", height=240) | |
| last_frame = gr.Image(type="filepath", label="Last image (optional)", height=240) | |
| prompt = gr.Textbox(label="Prompt", lines=4) | |
| with gr.Row(): | |
| max_dim = gr.Number(label="Longest side (px, ×32)", value=960, precision=0, | |
| info="960 ≈ 544p at 16:9 · 1344 ≈ 768p · 1920 ≈ 1080p") | |
| seconds = gr.Slider(label="Duration (s)", minimum=MIN_SECONDS, maximum=MAX_SECONDS, step=0.5, value=5.0) | |
| steps = gr.Slider(label="Steps", minimum=MIN_STEPS, maximum=MAX_STEPS, step=1, value=20) | |
| with gr.Row(): | |
| seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| job_id = gr.Textbox(label="job_id", value="", visible=False) | |
| run_btn = gr.Button("Generate", variant="primary") | |
| video = gr.Video(label="Result", interactive=False) | |
| used_seed = gr.Number(label="used_seed", precision=0, visible=False) | |
| job_id_out = gr.Textbox(label="job_id", visible=False) | |
| duration_ms = gr.Number(label="duration_ms", precision=0, visible=False) | |
| run_btn.click( | |
| fn=generate_video, | |
| inputs=[first_frame, last_frame, prompt, max_dim, seconds, steps, seed, randomize_seed, job_id], | |
| outputs=[video, used_seed, job_id_out, duration_ms], | |
| api_name="generate_video", | |
| ) | |
| demo.queue() | |
| if __name__ == "__main__": | |
| demo.launch() | |