Spaces:

FastVideo
/

FastH3-Omni

Running on Zero

FastH3-Omni / fasth3_engine.py

Read the kernel cache anonymously from the now-public dataset

1f01791 verified 2 days ago

16.9 kB

"""In-process FastH3 OmniRef engine — FP8 variant (benchmark against NVFP4).
Same Ref2VA pipeline as the NVFP4 Space, but the DiT linears are FP8 W8A8
(FastH3-OmniRef-FP8) instead of packed NVFP4. The sparse attention runs on
FastVideo's Triton tile-64 block-sparse kernel instead of the profile's bf16
tile-128 emulator.
FP8 specifics:
  - FastVideo's `rtx-4090` profile carries the FP8 numerics (FP8 linears, bf16
    VSA emulator, dense light VAE), but its offload policy streams the DiT
    layerwise for a 24 GB card. On the 96 GB ZeroGPU slice we override the
    offload to keep everything resident.
  - The FP8 repo ships `adaln_table.pt`; using it drops 24 GB of AdaLN
    projections from the load.
  - The text encoder is still the NVFP4 Qwen3-VL encoder (shared across the
    consumer releases), so the one-time flashinfer FP4 JIT still runs for it.
"""
from __future__ import annotations

import contextlib
import os
import tempfile
import threading
import time
from pathlib import Path

MODEL_REPO = os.getenv("FASTH3_MODEL_REPO", "FastVideo/FastH3-OmniRef-FP8")
PROFILE_NAME = "rtx-4090"  # FP8 numerics; offload and fine attention overridden below
ADALN_SUBPATH = os.getenv("FASTH3_ADALN_SUBPATH", "adaln_table.pt")
FINE_ATTENTION = os.getenv("FASTH3_FINE_ATTENTION", "triton64")
TILE64_CONFIG = (5, 4)
KERNEL_CACHE_DIR = os.getenv("FASTH3_KERNEL_CACHE_DIR", "/tmp/fasth3-kernels")
KERNEL_CACHE_REPO = os.getenv("FASTH3_KERNEL_CACHE_REPO", "multimodalart/fasth3-omni-kernel-cache")


def kernel_cache_env() -> None:
    os.environ.setdefault("TRITON_CACHE_DIR", f"{KERNEL_CACHE_DIR}/triton")
    os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{KERNEL_CACHE_DIR}/inductor")
    os.environ.setdefault("FLASHINFER_WORKSPACE_BASE", KERNEL_CACHE_DIR)
    os.environ.setdefault("CUDA_CACHE_PATH", f"{KERNEL_CACHE_DIR}/nv")
    os.environ.setdefault("CUDA_CACHE_MAXSIZE", str(4 << 30))
    os.environ.setdefault("PYTORCH_KERNEL_CACHE_PATH", f"{KERNEL_CACHE_DIR}/torch")
    os.makedirs(os.environ["PYTORCH_KERNEL_CACHE_PATH"], exist_ok=True)


def kernel_cache_name() -> str:
    from importlib.metadata import version

    packages = ("torch", "triton", "flashinfer-python", "fastvideo", "fastvideo-kernel")
    return "-".join(version(p) for p in packages).replace("+", "_") + f"-{FINE_ATTENTION}.tar"


def fetch_kernel_cache() -> bool:
    import tarfile

    from huggingface_hub import hf_hub_download

    t0 = time.time()
    try:
        archive = hf_hub_download(KERNEL_CACHE_REPO, kernel_cache_name(), repo_type="dataset", token=False)
        with tarfile.open(archive) as tar:
            tar.extractall(KERNEL_CACHE_DIR, filter="tar")
    except Exception as error:
        print(f"[fasth3-fp8] no kernel cache ({type(error).__name__}); compiling on first use", flush=True)
        return False
    print(f"[fasth3-fp8] kernel cache {kernel_cache_name()} restored in {time.time() - t0:.1f}s", flush=True)
    return True


kernel_cache_env()


def _hf_token() -> str | None:
    return os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or None


def _patch_mig_device_ids() -> None:
    """ZeroGPU exposes a MIG slice via a UUID in CUDA_VISIBLE_DEVICES; FastVideo
    does int() on it and crashes. Fall back to the logical index for UUIDs."""
    import fastvideo.platforms.cuda as cuda

    if getattr(cuda, "_mig_patched", False):
        return
    orig = cuda.device_id_to_physical_device_id

    def safe(device_id: int = 0) -> int:
        try:
            return orig(device_id)
        except (ValueError, TypeError):
            return device_id

    cuda.device_id_to_physical_device_id = safe
    cuda._mig_patched = True
    print("[fasth3-fp8] patched device_id_to_physical_device_id for MIG UUID", flush=True)


_CUDA_ORIGINALS: dict = {}
_DEFERRED_VALIDATIONS: list = []


def _is_cpu(device) -> bool:
    return device is None or str(device) == "cpu"


class _CapturedSlice:
    def __init__(self, inner, device) -> None:
        self._inner = inner
        self._device = device

    def __getitem__(self, key):
        tensor = self._inner[key]
        return tensor if _is_cpu(self._device) else tensor.to(self._device)

    def __getattr__(self, name):
        return getattr(self._inner, name)


def _captured_safe_open(original):
    class CapturedSafeOpen:
        def __init__(self, filename, framework="pt", device="cpu", **kwargs) -> None:
            self._inner = original(filename, framework=framework, device="cpu", **kwargs)
            self._device = device

        def __enter__(self):
            self._inner.__enter__()
            return self

        def __exit__(self, *exc):
            return self._inner.__exit__(*exc)

        def get_tensor(self, name):
            tensor = self._inner.get_tensor(name)
            return tensor if _is_cpu(self._device) else tensor.to(self._device)

        def get_tensors(self):
            tensors = self._inner.get_tensors()
            return tensors if _is_cpu(self._device) else {k: v.to(self._device) for k, v in tensors.items()}

        def get_slice(self, name):
            return _CapturedSlice(self._inner.get_slice(name), self._device)

        def __getattr__(self, name):
            return getattr(self._inner, name)

    return CapturedSafeOpen


def zerogpu_main_process_shims() -> None:
    """Let FastVideo's in-process worker initialize under ZeroGPU's CUDA emulation.
    `spaces` emulates CUDA in the main process but leaves `set_device` and
    `mem_get_info` pointing at the real runtime, which has no device there, and
    NVML sees no GPU, so FastVideo's platform probe would fall back to MPS.
    safetensors' `safe_open(device="cuda")` allocates on the real runtime, so files
    are read on CPU and moved with `.to()`, which ZeroGPU captures and packs.
    FastVideo probes CU_DEVICE_ATTRIBUTE_INTEGRATED with a ctypes cuInit; a failed
    cuInit in the GPU-less parent is inherited by every forked GPU worker, so the
    answer (a discrete card) is pinned instead. The NVFP4 text encoder's runtime
    check imports flashinfer, whose CuTe DSL loads the CUTLASS runtime against the
    driver, so that check (and its FP4 op registration) runs in the first GPU call.
    NCCL refuses to start without a GPU; with one rank every FastVideo collective
    short-circuits, so a gloo group is equivalent, and a one-rank barrier (which
    torch would otherwise route through CUDA device selection) is skipped.
    """
    import torch

    if _CUDA_ORIGINALS:
        return
    import fastvideo.platforms as platforms
    from fastvideo.platforms.cuda import NonNvmlCudaPlatform

    platforms._current_platform = NonNvmlCudaPlatform()

    import safetensors
    import safetensors.torch

    captured = _captured_safe_open(safetensors.safe_open)
    safetensors.safe_open = captured
    safetensors.torch.safe_open = captured
    import importlib

    for module_name in (
        "fastvideo.models.loader.weight_utils",
        "fastvideo.models.loader.component_loader",
        "fastvideo.models.loader.lora_patch",
        "fastvideo.models.loader.text_encoder_quantization",
    ):
        importlib.import_module(module_name).safe_open = captured
    _CUDA_ORIGINALS["set_device"] = torch.cuda.set_device
    _CUDA_ORIGINALS["mem_get_info"] = torch.cuda.mem_get_info
    total = torch.cuda.get_device_properties(0).total_memory
    torch.cuda.set_device = lambda *args, **kwargs: None
    torch.cuda.mem_get_info = lambda *args, **kwargs: (total, total)

    import torch.distributed as dist

    def gloo_instead_of_nccl(function):
        def wrapped(*args, backend=None, **kwargs):
            return function(*args, backend="gloo" if backend in (None, "nccl") else backend, **kwargs)

        return wrapped

    dist.init_process_group = gloo_instead_of_nccl(dist.init_process_group)
    dist.new_group = gloo_instead_of_nccl(dist.new_group)
    single_rank_barrier = dist.barrier

    def barrier(*args, **kwargs):
        if dist.is_initialized() and dist.get_world_size() == 1:
            return None
        return single_rank_barrier(*args, **kwargs)

    dist.barrier = barrier

    import fastvideo.platforms.cuda as fastvideo_cuda

    fastvideo_cuda._cuda_driver_device_is_integrated = lambda device_id: False

    from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config

    validate_runtime = MiniMaxH3SerializedNVFP4Config.validate_runtime

    def deferred_validate_runtime(self, device):
        _DEFERRED_VALIDATIONS.append((validate_runtime, self, device))

    MiniMaxH3SerializedNVFP4Config.validate_runtime = deferred_validate_runtime


def restore_cuda() -> None:
    """Inside a GPU call the device is real: put the runtime functions back."""
    import torch

    for name, function in _CUDA_ORIGINALS.items():
        setattr(torch.cuda, name, function)
    while _DEFERRED_VALIDATIONS:
        validate, config, device = _DEFERRED_VALIDATIONS.pop(0)
        validate(config, device)


def _find_pipeline():
    import gc

    from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline

    found = [o for o in gc.get_objects() if isinstance(o, MiniMaxH3BasePipeline)]
    return found[0] if len(found) == 1 else None


def _reuse_built_flashinfer_modules() -> None:
    from flashinfer.jit.core import JitSpecNvcc

    build = JitSpecNvcc.build
    if getattr(build, "reuses_built", False):
        return

    def build_unless_built(self, *args, **kwargs):
        if self.jit_library_path.exists():
            return None
        return build(self, *args, **kwargs)

    build_unless_built.reuses_built = True
    JitSpecNvcc.build = build_unless_built


def _ensure_tile64_kernel() -> None:
    if FINE_ATTENTION != "triton64":
        return
    import fastvideo.attention.backends.video_sparse_attn_h3 as vsa

    if vsa.block_sparse_attn_64_bhsd is None:
        from fastvideo_kernel.block_sparse_attn import block_sparse_attn

        vsa.block_sparse_attn_64_bhsd = block_sparse_attn
    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import _attn_fwd_sparse

    pinned = [c for c in _attn_fwd_sparse.configs if (c.num_stages, c.num_warps) == TILE64_CONFIG]
    if pinned:
        _attn_fwd_sparse.configs = pinned


STAGE_PROGRESS = {
    "conditioning_stage": (0.02, "Encoding the prompt and references"),
    "latent_preparation_stage": (0.08, "Preparing latents"),
    "video_decoding_stage": (0.86, "Decoding video"),
    "audio_decoding_stage": (0.94, "Decoding audio"),
}


@contextlib.contextmanager
def _progress_hooks(pipeline, num_steps, on_progress):
    if on_progress is None or pipeline is None:
        yield
        return
    stages = getattr(pipeline, "_stage_name_mapping", {})
    transformer = getattr(stages.get("denoising_stage"), "transformer", None)
    step = 0
    patched = []

    def report(fraction, desc):
        try:
            on_progress(fraction, desc)
        except Exception as error:
            print(f"[fasth3-fp8] progress report failed: {error!r}", flush=True)

    def before_step(module, args, kwargs):
        nonlocal step
        step += 1
        report(0.1 + 0.75 * (step - 1) / num_steps, f"Denoising step {step} of {num_steps}")

    def wrap(name, stage):
        forward = stage.forward

        def staged(*args, **kwargs):
            if name in STAGE_PROGRESS:
                report(*STAGE_PROGRESS[name])
            out = forward(*args, **kwargs)
            if name == "audio_decoding_stage":
                report(0.97, "Muxing video and audio")
            return out

        stage.forward = staged
        patched.append(stage)

    for name, stage in stages.items():
        wrap(name, stage)
    handle = transformer.register_forward_pre_hook(before_step, with_kwargs=True) if transformer is not None else None
    try:
        yield
    finally:
        if handle is not None:
            handle.remove()
        for stage in patched:
            del stage.forward


def prefetch() -> str:
    from huggingface_hub import snapshot_download

    return snapshot_download(MODEL_REPO, token=_hf_token())


class FastH3OmniRefEngine:
    def __init__(self) -> None:
        self.generator = None
        self.num_steps = 8
        self._lock = threading.Lock()

    def initialize(self) -> None:
        if self.generator is not None:
            return
        with self._lock:
            if self.generator is not None:
                return
            self._build()

    def _build(self) -> None:
        import dataclasses

        import torch  # noqa: F401

        try:
            p = torch.cuda.get_device_properties(0)
            print(f"[fasth3-fp8] GPU: {p.name}, total VRAM {p.total_memory / 1e9:.1f} GB", flush=True)
        except Exception as e:  # pragma: no cover
            print(f"[fasth3-fp8] device props unavailable: {e!r}", flush=True)

        import fasth3_ref as ref

        _patch_mig_device_ids()

        snapshot = prefetch()
        adaln = os.path.join(snapshot, ADALN_SUBPATH)
        if not os.path.isfile(adaln):
            print(f"[fasth3-fp8] no adaln table at {adaln}; loading full AdaLN proj", flush=True)
            adaln = None

        composed = os.path.join(tempfile.gettempdir(), "fasth3_omniref_fp8_composed")
        model_dir = ref.compose_model_dir(Path(snapshot), Path(snapshot), Path(composed))
        contract = ref.load_contract(model_dir)
        self.num_steps = int(contract["num_inference_steps"])

        from fastvideo.api import OffloadConfig

        # FP8 numerics from rtx-4090, but resident on 96 GB.
        profile = ref.DEVICE_PROFILES[PROFILE_NAME]
        profile = dataclasses.replace(
            profile,
            offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False,
                                  vae=False, pin_cpu_memory=False),
            sequential_load=False,
        )
        if FINE_ATTENTION == "triton64":
            env = {k: v for k, v in profile.env.items() if k != "FASTVIDEO_H3_VSA_FINE_ATTENTION"}
            profile = dataclasses.replace(profile, env=env)
            os.environ.pop("FASTVIDEO_H3_VSA_FINE_ATTENTION", None)
        ref.apply_profile(profile, None, adaln)  # dense decoder (no NVFP4 decoder on FP8)
        ref.apply_attention_prefix("default")
        print(f"[fasth3-fp8] profile: FP8 resident on {profile.label} numerics", flush=True)

        from fastvideo import VideoGenerator

        cfg = ref.build_generator_config(model_dir, num_gpus=1, lossless_accel=False, profile=profile)
        t0 = time.time()
        self.generator = VideoGenerator.from_config(cfg)
        print(f"[fasth3-fp8] generator ready in {time.time() - t0:.1f}s", flush=True)
        self.pipeline = _find_pipeline()
        print(f"[fasth3-fp8] fine attention: {FINE_ATTENTION}", flush=True)

    def generate(self, prompt, references, height, width, num_frames, seed, on_progress=None) -> str:
        if self.generator is None:
            raise RuntimeError("engine not initialized")
        _reuse_built_flashinfer_modules()
        restore_cuda()
        _ensure_tile64_kernel()
        if self.pipeline is not None and not self.pipeline.post_init_called:
            t0 = time.time()
            self.pipeline.post_init()
            print(f"[fasth3-fp8] pipeline stages ready in {time.time() - t0:.1f}s", flush=True)
        from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig
        from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference

        references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in references]

        out_dir = tempfile.mkdtemp(prefix="fasth3_fp8_out_")
        out_path = os.path.join(out_dir, "fasth3_omniref.mp4")
        t0 = time.time()
        with _progress_hooks(self.pipeline, self.num_steps, on_progress):
            result = self.generator.generate(
                GenerationRequest(
                    prompt=prompt,
                    negative_prompt="",
                    inputs=InputConfig(references=references),
                    sampling=SamplingConfig(height=height, width=width, num_frames=num_frames, fps=24,
                                            num_inference_steps=self.num_steps, guidance_scale=1.0,
                                            batch_cfg=False, seed=seed),
                    output=OutputConfig(output_path=out_path, save_video=True, return_frames=False),
                )
            )
        dt = time.time() - t0
        path = getattr(result, "video_path", None) or out_path
        peak = getattr(result, "peak_memory_mb", None)
        print(f"[fasth3-fp8] generated {width}x{height} {num_frames}f in {dt:.1f}s"
              + (f" | peak {peak / 1024:.1f} GiB" if peak else ""), flush=True)
        return path

	"""In-process FastH3 OmniRef engine — FP8 variant (benchmark against NVFP4).

	Same Ref2VA pipeline as the NVFP4 Space, but the DiT linears are FP8 W8A8
	(FastH3-OmniRef-FP8) instead of packed NVFP4. The sparse attention runs on
	FastVideo's Triton tile-64 block-sparse kernel instead of the profile's bf16
	tile-128 emulator.

	FP8 specifics:
	- FastVideo's `rtx-4090` profile carries the FP8 numerics (FP8 linears, bf16
	VSA emulator, dense light VAE), but its offload policy streams the DiT
	layerwise for a 24 GB card. On the 96 GB ZeroGPU slice we override the
	offload to keep everything resident.
	- The FP8 repo ships `adaln_table.pt`; using it drops 24 GB of AdaLN
	projections from the load.
	- The text encoder is still the NVFP4 Qwen3-VL encoder (shared across the
	consumer releases), so the one-time flashinfer FP4 JIT still runs for it.
	"""
	from __future__ import annotations

	import contextlib
	import os
	import tempfile
	import threading
	import time
	from pathlib import Path

	MODEL_REPO = os.getenv("FASTH3_MODEL_REPO", "FastVideo/FastH3-OmniRef-FP8")
	PROFILE_NAME = "rtx-4090" # FP8 numerics; offload and fine attention overridden below
	ADALN_SUBPATH = os.getenv("FASTH3_ADALN_SUBPATH", "adaln_table.pt")
	FINE_ATTENTION = os.getenv("FASTH3_FINE_ATTENTION", "triton64")
	TILE64_CONFIG = (5, 4)
	KERNEL_CACHE_DIR = os.getenv("FASTH3_KERNEL_CACHE_DIR", "/tmp/fasth3-kernels")
	KERNEL_CACHE_REPO = os.getenv("FASTH3_KERNEL_CACHE_REPO", "multimodalart/fasth3-omni-kernel-cache")


	def kernel_cache_env() -> None:
	os.environ.setdefault("TRITON_CACHE_DIR", f"{KERNEL_CACHE_DIR}/triton")
	os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{KERNEL_CACHE_DIR}/inductor")
	os.environ.setdefault("FLASHINFER_WORKSPACE_BASE", KERNEL_CACHE_DIR)
	os.environ.setdefault("CUDA_CACHE_PATH", f"{KERNEL_CACHE_DIR}/nv")
	os.environ.setdefault("CUDA_CACHE_MAXSIZE", str(4 << 30))
	os.environ.setdefault("PYTORCH_KERNEL_CACHE_PATH", f"{KERNEL_CACHE_DIR}/torch")
	os.makedirs(os.environ["PYTORCH_KERNEL_CACHE_PATH"], exist_ok=True)


	def kernel_cache_name() -> str:
	from importlib.metadata import version

	packages = ("torch", "triton", "flashinfer-python", "fastvideo", "fastvideo-kernel")
	return "-".join(version(p) for p in packages).replace("+", "_") + f"-{FINE_ATTENTION}.tar"


	def fetch_kernel_cache() -> bool:
	import tarfile

	from huggingface_hub import hf_hub_download

	t0 = time.time()
	try:
	archive = hf_hub_download(KERNEL_CACHE_REPO, kernel_cache_name(), repo_type="dataset", token=False)
	with tarfile.open(archive) as tar:
	tar.extractall(KERNEL_CACHE_DIR, filter="tar")
	except Exception as error:
	print(f"[fasth3-fp8] no kernel cache ({type(error).__name__}); compiling on first use", flush=True)
	return False
	print(f"[fasth3-fp8] kernel cache {kernel_cache_name()} restored in {time.time() - t0:.1f}s", flush=True)
	return True


	kernel_cache_env()


	def _hf_token() -> str \| None:
	return os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN") or None


	def _patch_mig_device_ids() -> None:
	"""ZeroGPU exposes a MIG slice via a UUID in CUDA_VISIBLE_DEVICES; FastVideo
	does int() on it and crashes. Fall back to the logical index for UUIDs."""
	import fastvideo.platforms.cuda as cuda

	if getattr(cuda, "_mig_patched", False):
	return
	orig = cuda.device_id_to_physical_device_id

	def safe(device_id: int = 0) -> int:
	try:
	return orig(device_id)
	except (ValueError, TypeError):
	return device_id

	cuda.device_id_to_physical_device_id = safe
	cuda._mig_patched = True
	print("[fasth3-fp8] patched device_id_to_physical_device_id for MIG UUID", flush=True)


	_CUDA_ORIGINALS: dict = {}
	_DEFERRED_VALIDATIONS: list = []


	def _is_cpu(device) -> bool:
	return device is None or str(device) == "cpu"


	class _CapturedSlice:
	def __init__(self, inner, device) -> None:
	self._inner = inner
	self._device = device

	def __getitem__(self, key):
	tensor = self._inner[key]
	return tensor if _is_cpu(self._device) else tensor.to(self._device)

	def __getattr__(self, name):
	return getattr(self._inner, name)


	def _captured_safe_open(original):
	class CapturedSafeOpen:
	def __init__(self, filename, framework="pt", device="cpu", **kwargs) -> None:
	self._inner = original(filename, framework=framework, device="cpu", **kwargs)
	self._device = device

	def __enter__(self):
	self._inner.__enter__()
	return self

	def __exit__(self, *exc):
	return self._inner.__exit__(*exc)

	def get_tensor(self, name):
	tensor = self._inner.get_tensor(name)
	return tensor if _is_cpu(self._device) else tensor.to(self._device)

	def get_tensors(self):
	tensors = self._inner.get_tensors()
	return tensors if _is_cpu(self._device) else {k: v.to(self._device) for k, v in tensors.items()}

	def get_slice(self, name):
	return _CapturedSlice(self._inner.get_slice(name), self._device)

	def __getattr__(self, name):
	return getattr(self._inner, name)

	return CapturedSafeOpen


	def zerogpu_main_process_shims() -> None:
	"""Let FastVideo's in-process worker initialize under ZeroGPU's CUDA emulation.

	`spaces` emulates CUDA in the main process but leaves `set_device` and
	`mem_get_info` pointing at the real runtime, which has no device there, and
	NVML sees no GPU, so FastVideo's platform probe would fall back to MPS.
	safetensors' `safe_open(device="cuda")` allocates on the real runtime, so files
	are read on CPU and moved with `.to()`, which ZeroGPU captures and packs.
	FastVideo probes CU_DEVICE_ATTRIBUTE_INTEGRATED with a ctypes cuInit; a failed
	cuInit in the GPU-less parent is inherited by every forked GPU worker, so the
	answer (a discrete card) is pinned instead. The NVFP4 text encoder's runtime
	check imports flashinfer, whose CuTe DSL loads the CUTLASS runtime against the
	driver, so that check (and its FP4 op registration) runs in the first GPU call.
	NCCL refuses to start without a GPU; with one rank every FastVideo collective
	short-circuits, so a gloo group is equivalent, and a one-rank barrier (which
	torch would otherwise route through CUDA device selection) is skipped.
	"""
	import torch

	if _CUDA_ORIGINALS:
	return
	import fastvideo.platforms as platforms
	from fastvideo.platforms.cuda import NonNvmlCudaPlatform

	platforms._current_platform = NonNvmlCudaPlatform()

	import safetensors
	import safetensors.torch

	captured = _captured_safe_open(safetensors.safe_open)
	safetensors.safe_open = captured
	safetensors.torch.safe_open = captured
	import importlib

	for module_name in (
	"fastvideo.models.loader.weight_utils",
	"fastvideo.models.loader.component_loader",
	"fastvideo.models.loader.lora_patch",
	"fastvideo.models.loader.text_encoder_quantization",
	):
	importlib.import_module(module_name).safe_open = captured
	_CUDA_ORIGINALS["set_device"] = torch.cuda.set_device
	_CUDA_ORIGINALS["mem_get_info"] = torch.cuda.mem_get_info
	total = torch.cuda.get_device_properties(0).total_memory
	torch.cuda.set_device = lambda args, *kwargs: None
	torch.cuda.mem_get_info = lambda args, *kwargs: (total, total)

	import torch.distributed as dist

	def gloo_instead_of_nccl(function):
	def wrapped(args, backend=None, *kwargs):
	return function(args, backend="gloo" if backend in (None, "nccl") else backend, *kwargs)

	return wrapped

	dist.init_process_group = gloo_instead_of_nccl(dist.init_process_group)
	dist.new_group = gloo_instead_of_nccl(dist.new_group)
	single_rank_barrier = dist.barrier

	def barrier(args, *kwargs):
	if dist.is_initialized() and dist.get_world_size() == 1:
	return None
	return single_rank_barrier(args, *kwargs)

	dist.barrier = barrier

	import fastvideo.platforms.cuda as fastvideo_cuda

	fastvideo_cuda._cuda_driver_device_is_integrated = lambda device_id: False

	from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config

	validate_runtime = MiniMaxH3SerializedNVFP4Config.validate_runtime

	def deferred_validate_runtime(self, device):
	_DEFERRED_VALIDATIONS.append((validate_runtime, self, device))

	MiniMaxH3SerializedNVFP4Config.validate_runtime = deferred_validate_runtime


	def restore_cuda() -> None:
	"""Inside a GPU call the device is real: put the runtime functions back."""
	import torch

	for name, function in _CUDA_ORIGINALS.items():
	setattr(torch.cuda, name, function)
	while _DEFERRED_VALIDATIONS:
	validate, config, device = _DEFERRED_VALIDATIONS.pop(0)
	validate(config, device)


	def _find_pipeline():
	import gc

	from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3BasePipeline

	found = [o for o in gc.get_objects() if isinstance(o, MiniMaxH3BasePipeline)]
	return found[0] if len(found) == 1 else None


	def _reuse_built_flashinfer_modules() -> None:
	from flashinfer.jit.core import JitSpecNvcc

	build = JitSpecNvcc.build
	if getattr(build, "reuses_built", False):
	return

	def build_unless_built(self, args, *kwargs):
	if self.jit_library_path.exists():
	return None
	return build(self, args, *kwargs)

	build_unless_built.reuses_built = True
	JitSpecNvcc.build = build_unless_built


	def _ensure_tile64_kernel() -> None:
	if FINE_ATTENTION != "triton64":
	return
	import fastvideo.attention.backends.video_sparse_attn_h3 as vsa

	if vsa.block_sparse_attn_64_bhsd is None:
	from fastvideo_kernel.block_sparse_attn import block_sparse_attn

	vsa.block_sparse_attn_64_bhsd = block_sparse_attn
	from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import _attn_fwd_sparse

	pinned = [c for c in _attn_fwd_sparse.configs if (c.num_stages, c.num_warps) == TILE64_CONFIG]
	if pinned:
	_attn_fwd_sparse.configs = pinned


	STAGE_PROGRESS = {
	"conditioning_stage": (0.02, "Encoding the prompt and references"),
	"latent_preparation_stage": (0.08, "Preparing latents"),
	"video_decoding_stage": (0.86, "Decoding video"),
	"audio_decoding_stage": (0.94, "Decoding audio"),
	}


	@contextlib.contextmanager
	def _progress_hooks(pipeline, num_steps, on_progress):
	if on_progress is None or pipeline is None:
	yield
	return
	stages = getattr(pipeline, "_stage_name_mapping", {})
	transformer = getattr(stages.get("denoising_stage"), "transformer", None)
	step = 0
	patched = []

	def report(fraction, desc):
	try:
	on_progress(fraction, desc)
	except Exception as error:
	print(f"[fasth3-fp8] progress report failed: {error!r}", flush=True)

	def before_step(module, args, kwargs):
	nonlocal step
	step += 1
	report(0.1 + 0.75 * (step - 1) / num_steps, f"Denoising step {step} of {num_steps}")

	def wrap(name, stage):
	forward = stage.forward

	def staged(args, *kwargs):
	if name in STAGE_PROGRESS:
	report(*STAGE_PROGRESS[name])
	out = forward(args, *kwargs)
	if name == "audio_decoding_stage":
	report(0.97, "Muxing video and audio")
	return out

	stage.forward = staged
	patched.append(stage)

	for name, stage in stages.items():
	wrap(name, stage)
	handle = transformer.register_forward_pre_hook(before_step, with_kwargs=True) if transformer is not None else None
	try:
	yield
	finally:
	if handle is not None:
	handle.remove()
	for stage in patched:
	del stage.forward


	def prefetch() -> str:
	from huggingface_hub import snapshot_download

	return snapshot_download(MODEL_REPO, token=_hf_token())


	class FastH3OmniRefEngine:
	def __init__(self) -> None:
	self.generator = None
	self.num_steps = 8
	self._lock = threading.Lock()

	def initialize(self) -> None:
	if self.generator is not None:
	return
	with self._lock:
	if self.generator is not None:
	return
	self._build()

	def _build(self) -> None:
	import dataclasses

	import torch # noqa: F401

	try:
	p = torch.cuda.get_device_properties(0)
	print(f"[fasth3-fp8] GPU: {p.name}, total VRAM {p.total_memory / 1e9:.1f} GB", flush=True)
	except Exception as e: # pragma: no cover
	print(f"[fasth3-fp8] device props unavailable: {e!r}", flush=True)

	import fasth3_ref as ref

	_patch_mig_device_ids()

	snapshot = prefetch()
	adaln = os.path.join(snapshot, ADALN_SUBPATH)
	if not os.path.isfile(adaln):
	print(f"[fasth3-fp8] no adaln table at {adaln}; loading full AdaLN proj", flush=True)
	adaln = None

	composed = os.path.join(tempfile.gettempdir(), "fasth3_omniref_fp8_composed")
	model_dir = ref.compose_model_dir(Path(snapshot), Path(snapshot), Path(composed))
	contract = ref.load_contract(model_dir)
	self.num_steps = int(contract["num_inference_steps"])

	from fastvideo.api import OffloadConfig

	# FP8 numerics from rtx-4090, but resident on 96 GB.
	profile = ref.DEVICE_PROFILES[PROFILE_NAME]
	profile = dataclasses.replace(
	profile,
	offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False,
	vae=False, pin_cpu_memory=False),
	sequential_load=False,
	)
	if FINE_ATTENTION == "triton64":
	env = {k: v for k, v in profile.env.items() if k != "FASTVIDEO_H3_VSA_FINE_ATTENTION"}
	profile = dataclasses.replace(profile, env=env)
	os.environ.pop("FASTVIDEO_H3_VSA_FINE_ATTENTION", None)
	ref.apply_profile(profile, None, adaln) # dense decoder (no NVFP4 decoder on FP8)
	ref.apply_attention_prefix("default")
	print(f"[fasth3-fp8] profile: FP8 resident on {profile.label} numerics", flush=True)

	from fastvideo import VideoGenerator

	cfg = ref.build_generator_config(model_dir, num_gpus=1, lossless_accel=False, profile=profile)
	t0 = time.time()
	self.generator = VideoGenerator.from_config(cfg)
	print(f"[fasth3-fp8] generator ready in {time.time() - t0:.1f}s", flush=True)
	self.pipeline = _find_pipeline()
	print(f"[fasth3-fp8] fine attention: {FINE_ATTENTION}", flush=True)

	def generate(self, prompt, references, height, width, num_frames, seed, on_progress=None) -> str:
	if self.generator is None:
	raise RuntimeError("engine not initialized")
	_reuse_built_flashinfer_modules()
	restore_cuda()
	_ensure_tile64_kernel()
	if self.pipeline is not None and not self.pipeline.post_init_called:
	t0 = time.time()
	self.pipeline.post_init()
	print(f"[fasth3-fp8] pipeline stages ready in {time.time() - t0:.1f}s", flush=True)
	from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig
	from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference

	references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in references]

	out_dir = tempfile.mkdtemp(prefix="fasth3_fp8_out_")
	out_path = os.path.join(out_dir, "fasth3_omniref.mp4")
	t0 = time.time()
	with _progress_hooks(self.pipeline, self.num_steps, on_progress):
	result = self.generator.generate(
	GenerationRequest(
	prompt=prompt,
	negative_prompt="",
	inputs=InputConfig(references=references),
	sampling=SamplingConfig(height=height, width=width, num_frames=num_frames, fps=24,
	num_inference_steps=self.num_steps, guidance_scale=1.0,
	batch_cfg=False, seed=seed),
	output=OutputConfig(output_path=out_path, save_video=True, return_frames=False),
	)
	)
	dt = time.time() - t0
	path = getattr(result, "video_path", None) or out_path
	peak = getattr(result, "peak_memory_mb", None)
	print(f"[fasth3-fp8] generated {width}x{height} {num_frames}f in {dt:.1f}s"
	+ (f" \| peak {peak / 1024:.1f} GiB" if peak else ""), flush=True)
	return path