From 19e37595033d802c7dedfed77f2a8f42b1d4a992 Mon Sep 17 00:00:00 2001 From: yifeif-nv <277870278+yifeif-nv@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:39:25 -0700 Subject: [PATCH 1/8] feat(minimax-h3): unify native audiovisual delivery Build T2VA, FL2VA, and Ref2VA together from the pinned full Comfy INT8 denoisers and NVFP4 AWQ text checkpoint. Keep normal generation as the default and enable fixed-base super resolution only through an explicit build option. Preserve dynamic request shapes, synchronized audio, and native C++ execution. Keep model semantics in the H3 family and provide the generic runtime, bundle loading, and Windows media support needed to consume it. Document clean setup, CLI and C++ consumption, checkpoint semantics, cache settings, and artifact rebuild requirements on the model page. Signed-off-by: yifeif-nv <277870278+yifeif-nv@users.noreply.github.com> --- CMakeLists.txt | 234 +- apps/benchmark/native/benchmark_worker.cpp | 16 +- .../native/test_benchmark_worker_e2e.cpp | 15 + apps/cli/cli.cpp | 159 +- apps/cli/cli.h | 6 + apps/cli/io.cpp | 25 +- apps/cli/main.cpp | 28 +- apps/cli/tests/test_cli.cpp | 103 +- apps/cli/tests/test_windows_media.cpp | 559 ++++ apps/cli/tests/test_windows_utf8_argv.cpp | 137 + apps/cli/trtmc_windows_utf8.manifest | 13 + apps/cli/windows_media.cpp | 1434 +++++++++++ apps/cli/windows_media.h | 76 + apps/cli/windows_utf8_argv.cpp | 73 + apps/cli/windows_utf8_argv.h | 33 + core/builder/tensorrt_model_connect/build.py | 19 + .../tensorrt_model_connect/build_cli.py | 40 + .../tensorrt_model_connect/bundle_writer.py | 47 +- core/builder/tests/test_build.py | 21 + core/builder/tests/test_build_cli.py | 22 + core/builder/tests/test_bundle_writer.py | 130 +- core/runtime/bundle/bundle_format.cpp | 10 +- core/runtime/include/trtmc/bundle.h | 4 + .../include/trtmc/runtime/family_factory.h | 3 + .../include/trtmc/runtime/trt_backend.h | 33 + core/runtime/include/trtmc/task.h | 84 + core/runtime/loader/family_loader.cpp | 57 +- core/runtime/platform/dynamic_library.cpp | 184 ++ core/runtime/platform/dynamic_library.h | 31 + core/runtime/tensorrt/rtx_backend.cpp | 1041 +++++++- .../tensorrt/runtime_cache_persistence.cpp | 153 ++ .../tensorrt/runtime_cache_persistence.h | 45 + core/runtime/tensorrt/trt_backend.cpp | 5 +- core/runtime/tensorrt/trt_logger.cpp | 7 +- core/runtime/tensorrt/trt_module_impl.cpp | 302 ++- core/runtime/tensorrt/trt_module_impl.h | 35 +- core/runtime/tests/test_bundle_format_v1.cpp | 11 +- core/runtime/tests/test_dynamic_library.cpp | 57 + core/runtime/tests/test_family_loader.cpp | 17 +- .../tests/test_runtime_cache_persistence.cpp | 122 + core/runtime/tests/test_task_api.cpp | 35 + families/minimax_h3/adaln_builder.py | 35 +- families/minimax_h3/audio_vae_builder.py | 771 ++++++ families/minimax_h3/checkpoint.py | 75 +- families/minimax_h3/config.py | 118 +- families/minimax_h3/delivery.py | 70 + families/minimax_h3/dit_builder.py | 339 +-- families/minimax_h3/fl2va_contract.py | 665 +++++ .../minimax_h3/fl2va_vae_encoder_builder.py | 454 ++++ families/minimax_h3/graph_ops.py | 412 ++- families/minimax_h3/model.py | 515 ++-- .../multimodal_text_encoder_builder.py | 448 ++++ .../minimax_h3/multimodal_vision_builder.py | 446 ++++ families/minimax_h3/nvfp4_text_checkpoint.py | 296 +++ families/minimax_h3/provenance.py | 265 ++ families/minimax_h3/quantized_checkpoint.py | 738 ++++++ .../ref2va_audio_encoder_builder.py | 524 ++++ families/minimax_h3/ref2va_bundle_contract.py | 318 +++ families/minimax_h3/ref2va_checkpoint.py | 282 +++ families/minimax_h3/ref2va_contract.py | 1114 ++++++++ families/minimax_h3/ref2va_dit_builder.py | 653 +++++ families/minimax_h3/ref2va_qwen_builder.py | 70 + families/minimax_h3/ref2va_qwen_contract.py | 84 + .../ref2va_video_encoder_builder.py | 527 ++++ families/minimax_h3/runtime/CMakeLists.txt | 56 +- families/minimax_h3/runtime/conditioning.cpp | 508 ++++ families/minimax_h3/runtime/conditioning.h | 57 + families/minimax_h3/runtime/fl2va_runtime.cpp | 765 ++++++ families/minimax_h3/runtime/fl2va_runtime.h | 108 + .../minimax_h3/runtime/hot_engine_policy.h | 55 + families/minimax_h3/runtime/pipeline.cpp | 2243 +++++++++++++---- families/minimax_h3/runtime/pipeline.h | 137 +- families/minimax_h3/runtime/plugin.cpp | 323 ++- families/minimax_h3/runtime/public_profile.h | 141 ++ .../minimax_h3/runtime/ref2va_runtime.cpp | 1812 +++++++++++++ families/minimax_h3/runtime/ref2va_runtime.h | 287 +++ .../runtime/super_resolution_runtime.cpp | 139 + .../runtime/super_resolution_runtime.h | 39 + .../minimax_h3/runtime/torch_cuda_normal.cu | 276 +- .../minimax_h3/runtime/torch_cuda_normal.h | 14 +- families/minimax_h3/runtime_config_schema.py | 56 + families/minimax_h3/staged_build.py | 941 +++++++ .../minimax_h3/super_resolution_builder.py | 491 ++++ .../cpp/test_minimax_h3_conditioning.cpp | 156 ++ .../tests/cpp/test_minimax_h3_cuda_rng.cpp | 223 +- .../cpp/test_minimax_h3_fl2va_runtime.cpp | 453 ++++ .../tests/cpp/test_minimax_h3_math.cpp | 480 +++- .../cpp/test_minimax_h3_ref2va_runtime.cpp | 734 ++++++ .../tests/prompts/t2va-example-1.json | 2 +- .../tests/test_audio_vae_builder.py | 259 ++ families/minimax_h3/tests/test_checkpoint.py | 50 + families/minimax_h3/tests/test_config.py | 14 +- families/minimax_h3/tests/test_delivery.py | 154 ++ .../minimax_h3/tests/test_dynamic_shapes.py | 9 +- families/minimax_h3/tests/test_e2e.py | 109 +- .../tests/test_fl2va_native_builders.py | 331 +++ families/minimax_h3/tests/test_model.py | 222 ++ .../tests/test_native_runtime_boundary.py | 45 + .../tests/test_nvfp4_text_checkpoint.py | 186 ++ .../tests/test_quantized_checkpoint.py | 348 +++ .../tests/test_ref2va_bundle_contract.py | 184 ++ .../tests/test_ref2va_first_block_cache.py | 384 +++ .../tests/test_ref2va_native_contract.py | 615 +++++ .../tests/test_ref2va_native_encoders.py | 168 ++ .../minimax_h3/tests/test_staged_build.py | 719 ++++++ .../tests/test_super_resolution_builder.py | 157 ++ .../tests/test_super_resolution_bundle.py | 127 + .../minimax_h3/tests/test_trt_builders.py | 439 +++- .../tests/test_vae_builder_dynamic.py | 72 + families/minimax_h3/text_encoder_builder.py | 42 +- families/minimax_h3/trt_compat.py | 191 ++ families/minimax_h3/vae_builder.py | 153 +- tools/legal_headers.py | 1 + tools/tests/test_architecture.py | 36 +- website/docs/models-recipes/minimax-h3.md | 340 +++ website/sidebars.js | 1 + 116 files changed, 28811 insertions(+), 1686 deletions(-) create mode 100644 apps/cli/tests/test_windows_media.cpp create mode 100644 apps/cli/tests/test_windows_utf8_argv.cpp create mode 100644 apps/cli/trtmc_windows_utf8.manifest create mode 100644 apps/cli/windows_media.cpp create mode 100644 apps/cli/windows_media.h create mode 100644 apps/cli/windows_utf8_argv.cpp create mode 100644 apps/cli/windows_utf8_argv.h create mode 100644 core/runtime/platform/dynamic_library.cpp create mode 100644 core/runtime/platform/dynamic_library.h create mode 100644 core/runtime/tensorrt/runtime_cache_persistence.cpp create mode 100644 core/runtime/tensorrt/runtime_cache_persistence.h create mode 100644 core/runtime/tests/test_dynamic_library.cpp create mode 100644 core/runtime/tests/test_runtime_cache_persistence.cpp create mode 100644 families/minimax_h3/audio_vae_builder.py create mode 100644 families/minimax_h3/delivery.py create mode 100644 families/minimax_h3/fl2va_contract.py create mode 100644 families/minimax_h3/fl2va_vae_encoder_builder.py create mode 100644 families/minimax_h3/multimodal_text_encoder_builder.py create mode 100644 families/minimax_h3/multimodal_vision_builder.py create mode 100644 families/minimax_h3/nvfp4_text_checkpoint.py create mode 100644 families/minimax_h3/provenance.py create mode 100644 families/minimax_h3/quantized_checkpoint.py create mode 100644 families/minimax_h3/ref2va_audio_encoder_builder.py create mode 100644 families/minimax_h3/ref2va_bundle_contract.py create mode 100644 families/minimax_h3/ref2va_checkpoint.py create mode 100644 families/minimax_h3/ref2va_contract.py create mode 100644 families/minimax_h3/ref2va_dit_builder.py create mode 100644 families/minimax_h3/ref2va_qwen_builder.py create mode 100644 families/minimax_h3/ref2va_qwen_contract.py create mode 100644 families/minimax_h3/ref2va_video_encoder_builder.py create mode 100644 families/minimax_h3/runtime/conditioning.cpp create mode 100644 families/minimax_h3/runtime/conditioning.h create mode 100644 families/minimax_h3/runtime/fl2va_runtime.cpp create mode 100644 families/minimax_h3/runtime/fl2va_runtime.h create mode 100644 families/minimax_h3/runtime/hot_engine_policy.h create mode 100644 families/minimax_h3/runtime/public_profile.h create mode 100644 families/minimax_h3/runtime/ref2va_runtime.cpp create mode 100644 families/minimax_h3/runtime/ref2va_runtime.h create mode 100644 families/minimax_h3/runtime/super_resolution_runtime.cpp create mode 100644 families/minimax_h3/runtime/super_resolution_runtime.h create mode 100644 families/minimax_h3/runtime_config_schema.py create mode 100644 families/minimax_h3/staged_build.py create mode 100644 families/minimax_h3/super_resolution_builder.py create mode 100644 families/minimax_h3/tests/cpp/test_minimax_h3_conditioning.cpp create mode 100644 families/minimax_h3/tests/cpp/test_minimax_h3_fl2va_runtime.cpp create mode 100644 families/minimax_h3/tests/cpp/test_minimax_h3_ref2va_runtime.cpp create mode 100644 families/minimax_h3/tests/test_audio_vae_builder.py create mode 100644 families/minimax_h3/tests/test_checkpoint.py create mode 100644 families/minimax_h3/tests/test_delivery.py create mode 100644 families/minimax_h3/tests/test_fl2va_native_builders.py create mode 100644 families/minimax_h3/tests/test_model.py create mode 100644 families/minimax_h3/tests/test_native_runtime_boundary.py create mode 100644 families/minimax_h3/tests/test_nvfp4_text_checkpoint.py create mode 100644 families/minimax_h3/tests/test_quantized_checkpoint.py create mode 100644 families/minimax_h3/tests/test_ref2va_bundle_contract.py create mode 100644 families/minimax_h3/tests/test_ref2va_first_block_cache.py create mode 100644 families/minimax_h3/tests/test_ref2va_native_contract.py create mode 100644 families/minimax_h3/tests/test_ref2va_native_encoders.py create mode 100644 families/minimax_h3/tests/test_staged_build.py create mode 100644 families/minimax_h3/tests/test_super_resolution_builder.py create mode 100644 families/minimax_h3/tests/test_super_resolution_bundle.py create mode 100644 families/minimax_h3/tests/test_vae_builder_dynamic.py create mode 100644 families/minimax_h3/trt_compat.py create mode 100644 website/docs/models-recipes/minimax-h3.md diff --git a/CMakeLists.txt b/CMakeLists.txt index 2d653e3281..1ceeb5872b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,11 +10,20 @@ set(CMAKE_CXX_EXTENSIONS OFF) set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) -set(CMAKE_BUILD_RPATH_USE_ORIGIN TRUE) +if(WIN32) + set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) +else() + set(CMAKE_BUILD_RPATH_USE_ORIGIN TRUE) +endif() + +set(TRTMC_CXX_WARNINGS + "$<$:-Wall;-Wextra;-Wpedantic>" +) include(GNUInstallDirs) find_package(CUDAToolkit REQUIRED) find_package(nlohmann_json 3.11 REQUIRED) +option(TRTMC_BUILD_BACKEND_TRT "Build the standard TensorRT backend DSO" ON) set(_trtmc_dependency_roots /usr @@ -22,26 +31,34 @@ set(_trtmc_dependency_roots /usr/local/cuda /opt/tensorrt ) +if(DEFINED ENV{CUDA_PATH}) + list(PREPEND _trtmc_dependency_roots "$ENV{CUDA_PATH}") +endif() if(DEFINED ENV{TRT_ROOT}) list(PREPEND _trtmc_dependency_roots "$ENV{TRT_ROOT}") endif() -find_path(TRTMC_TRT_INCLUDE_DIR - NAMES NvInferRuntime.h - HINTS ${_trtmc_dependency_roots} - PATH_SUFFIXES include include/zapped_headers - REQUIRED -) -find_library(TRTMC_TRT_LIBRARY - NAMES nvinfer libnvinfer.so.11 - HINTS ${_trtmc_dependency_roots} - PATH_SUFFIXES lib lib64 lib/aarch64-linux-gnu lib/x86_64-linux-gnu - REQUIRED -) +if(TRTMC_BUILD_BACKEND_TRT) + find_path(TRTMC_TRT_INCLUDE_DIR + NAMES NvInferRuntime.h + HINTS ${_trtmc_dependency_roots} + PATH_SUFFIXES include include/zapped_headers + REQUIRED + ) + find_library(TRTMC_TRT_LIBRARY + NAMES nvinfer libnvinfer.so.11 nvinfer_11 + HINTS ${_trtmc_dependency_roots} + PATH_SUFFIXES lib lib64 lib/x64 lib/aarch64-linux-gnu lib/x86_64-linux-gnu + REQUIRED + ) +endif() option(TRTMC_ENABLE_BYOK "Enable the optional TVM-FFI BYOK bridge" ON) set(TRTMC_HAS_TVM_FFI OFF) if(TRTMC_ENABLE_BYOK) + if(NOT TRTMC_BUILD_BACKEND_TRT) + message(FATAL_ERROR "TRTMC_ENABLE_BYOK=ON requires TRTMC_BUILD_BACKEND_TRT=ON") + endif() find_package(Python3 COMPONENTS Interpreter QUIET) if(Python3_Interpreter_FOUND) execute_process( @@ -77,7 +94,11 @@ endif() # Family CMake files use these two variables to link the shared Engine API. set(TRTMC_CUDA_INCLUDE_DIR ${CUDAToolkit_INCLUDE_DIRS}) -set(TRTMC_CUDART_LIBRARY CUDA::cudart) +if(WIN32 AND CMAKE_CUDA_RUNTIME_LIBRARY STREQUAL "Static") + set(TRTMC_CUDART_LIBRARY CUDA::cudart_static) +else() + set(TRTMC_CUDART_LIBRARY CUDA::cudart) +endif() add_library(trtmc_core SHARED core/runtime/bundle/bundle_format.cpp @@ -95,18 +116,22 @@ target_include_directories(trtmc_core target_include_directories(trtmc_core SYSTEM PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) target_link_libraries(trtmc_core PUBLIC - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} PRIVATE nlohmann_json::nlohmann_json ) -target_compile_options(trtmc_core PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_core PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_core PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) add_library(trtmc_runtime SHARED core/runtime/loader/family_loader.cpp + core/runtime/platform/dynamic_library.cpp ) target_include_directories(trtmc_runtime PUBLIC @@ -120,12 +145,16 @@ target_link_libraries(trtmc_runtime trtmc_core ${CMAKE_DL_LIBS} ) -target_compile_options(trtmc_runtime PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_runtime PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_runtime PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) +if(TRTMC_BUILD_BACKEND_TRT) add_library(trtmc_backend_trt SHARED core/runtime/tensorrt/trt_backend.cpp core/runtime/tensorrt/trt_logger.cpp @@ -144,16 +173,21 @@ target_link_libraries(trtmc_backend_trt PRIVATE trtmc_core ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${CMAKE_DL_LIBS} ) -target_compile_options(trtmc_backend_trt PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc_backend_trt PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_backend_trt PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) +endif() option(TRTMC_BUILD_BACKEND_RTX "Build TensorRT-RTX backend DSO" OFF) +set(TRTMC_RTX_RUNTIME_DIR "" CACHE PATH "TensorRT-RTX runtime DLL directory on Windows") if(TRTMC_BUILD_BACKEND_RTX) if(NOT IS_DIRECTORY "${TRTMC_RTX_INCLUDE_DIR}" OR NOT EXISTS "${TRTMC_RTX_INCLUDE_DIR}/NvInfer.h") @@ -167,7 +201,7 @@ if(TRTMC_BUILD_BACKEND_RTX) ) endif() find_library(TRTMC_RTX_LIBRARY - NAMES tensorrt_rtx + NAMES tensorrt_rtx tensorrt_rtx_1_6 PATHS "${TRTMC_RTX_LIBRARY_DIR}" NO_DEFAULT_PATH ) @@ -177,8 +211,22 @@ if(TRTMC_BUILD_BACKEND_RTX) ) endif() + if(WIN32) + find_file(TRTMC_RTX_RUNTIME_DLL + NAMES tensorrt_rtx_1_6.dll tensorrt_rtx.dll + PATHS "${TRTMC_RTX_RUNTIME_DIR}" + NO_DEFAULT_PATH + ) + if(NOT TRTMC_RTX_RUNTIME_DLL) + message(FATAL_ERROR + "TRTMC_BUILD_BACKEND_RTX=ON requires TRTMC_RTX_RUNTIME_DIR with the runtime DLL" + ) + endif() + endif() + add_library(trtmc_backend_rtx SHARED core/runtime/tensorrt/rtx_backend.cpp + core/runtime/tensorrt/runtime_cache_persistence.cpp core/runtime/tensorrt/trt_logger.cpp core/runtime/tensorrt/trt_module_impl.cpp ) @@ -195,14 +243,30 @@ if(TRTMC_BUILD_BACKEND_RTX) PRIVATE trtmc_core ${TRTMC_RTX_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} + ) + target_compile_options(trtmc_backend_rtx PRIVATE ${TRTMC_CXX_WARNINGS}) + target_compile_definitions(trtmc_backend_rtx PRIVATE + TRTMC_TRT_DELETE_OBJECTS=1 + TRTMC_TRT_HAS_TENSOR_ALIAS_API=1 + TRTMC_TRT_HAS_COMMUNICATOR_API=1 ) - target_compile_options(trtmc_backend_rtx PRIVATE -Wall -Wextra -Wpedantic) set_target_properties(trtmc_backend_rtx PROPERTIES OUTPUT_NAME trtmc_backend_trt_rtx + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN" ) + if(WIN32) + add_custom_command(TARGET trtmc_backend_rtx POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${TRTMC_RTX_RUNTIME_DLL}" + "$" + COMMENT "Copying TensorRT-RTX runtime DLL next to the backend" + ) + endif() endif() if(TRTMC_HAS_TVM_FFI) @@ -226,7 +290,7 @@ if(TRTMC_HAS_TVM_FFI) ) target_link_libraries(trtmc_byok_tvm_ffi PRIVATE ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} nlohmann_json::nlohmann_json ) @@ -234,7 +298,7 @@ if(TRTMC_HAS_TVM_FFI) TRTMC_HAS_TRT=1 TRTMC_HAS_TVM_FFI=1 ) - target_compile_options(trtmc_byok_tvm_ffi PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_byok_tvm_ffi PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_byok_tvm_ffi PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN;${TRTMC_TVM_FFI_LIBRARY_DIR}" @@ -250,6 +314,9 @@ endif() # A family owns its target, sources, dependencies, warnings, and installation. # The root knows only the directory convention; adding a family never changes # a central source or target list. +set(TRTMC_RUNTIME_MODELS "" CACHE STRING + "Semicolon-separated family runtimes to build; empty builds every family" +) file(GLOB _trtmc_family_runtime_cmake CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/families/*/runtime/CMakeLists.txt" ) @@ -257,6 +324,9 @@ foreach(_trtmc_runtime_cmake IN LISTS _trtmc_family_runtime_cmake) get_filename_component(_trtmc_runtime_dir "${_trtmc_runtime_cmake}" DIRECTORY) get_filename_component(_trtmc_family_dir "${_trtmc_runtime_dir}" DIRECTORY) get_filename_component(_trtmc_family "${_trtmc_family_dir}" NAME) + if(TRTMC_RUNTIME_MODELS AND NOT _trtmc_family IN_LIST TRTMC_RUNTIME_MODELS) + continue() + endif() add_subdirectory( "${_trtmc_runtime_dir}" "${CMAKE_BINARY_DIR}/families/${_trtmc_family}" @@ -266,10 +336,12 @@ endforeach() add_library(trtmc_cli STATIC apps/cli/cli.cpp apps/cli/io.cpp + apps/cli/windows_media.cpp ) target_include_directories(trtmc_cli PUBLIC ${PROJECT_SOURCE_DIR}/apps PRIVATE + ${PROJECT_SOURCE_DIR}/core ${PROJECT_SOURCE_DIR}/core/runtime/include ${PROJECT_SOURCE_DIR}/third_party/stb ) @@ -280,17 +352,29 @@ target_link_libraries(trtmc_cli nlohmann_json::nlohmann_json ${CMAKE_DL_LIBS} ) +if(WIN32) + target_link_libraries(trtmc_cli PRIVATE mfplat mfreadwrite mfuuid ole32) +endif() target_compile_definitions(trtmc_cli PRIVATE TRTMC_VERSION_STRING="${PROJECT_VERSION}") -target_compile_options(trtmc_cli PRIVATE -Wall -Wextra -Wpedantic) -set_source_files_properties(apps/cli/io.cpp PROPERTIES - COMPILE_OPTIONS "-Wno-missing-field-initializers;-Wno-pedantic" -) +target_compile_options(trtmc_cli PRIVATE ${TRTMC_CXX_WARNINGS}) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + set_source_files_properties(apps/cli/io.cpp PROPERTIES + COMPILE_OPTIONS "-Wno-missing-field-initializers;-Wno-pedantic" + ) +endif() add_executable(trtmc apps/cli/main.cpp) +if(WIN32) + target_sources(trtmc PRIVATE + apps/cli/windows_utf8_argv.cpp + apps/cli/trtmc_windows_utf8.manifest + ) +endif() target_include_directories(trtmc PRIVATE ${PROJECT_SOURCE_DIR}/apps) target_link_libraries(trtmc PRIVATE trtmc_cli) -target_compile_options(trtmc PRIVATE -Wall -Wextra -Wpedantic) +target_compile_options(trtmc PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/$<0:>" BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" ) @@ -305,7 +389,7 @@ if(TRTMC_BUILD_EXAMPLES) trtmc_runtime nlohmann_json::nlohmann_json ) - target_compile_options(trtmc_benchmark_worker PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_benchmark_worker PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_benchmark_worker PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" @@ -320,7 +404,7 @@ if(TRTMC_BUILD_EXAMPLES) trtmc_runtime nlohmann_json::nlohmann_json ) - target_compile_options(trtmc_dataset_benchmark PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(trtmc_dataset_benchmark PRIVATE ${TRTMC_CXX_WARNINGS}) set_target_properties(trtmc_dataset_benchmark PROPERTIES BUILD_RPATH "\$ORIGIN" INSTALL_RPATH "\$ORIGIN/../${CMAKE_INSTALL_LIBDIR}" @@ -331,12 +415,12 @@ if(TRTMC_BUILD_TESTS) add_executable(test_bundle_format_v1 core/runtime/tests/test_bundle_format_v1.cpp) target_include_directories(test_bundle_format_v1 PRIVATE ${PROJECT_SOURCE_DIR}/core) target_link_libraries(test_bundle_format_v1 PRIVATE trtmc_core) - target_compile_options(test_bundle_format_v1 PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_bundle_format_v1 PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME bundle_format_v1 COMMAND test_bundle_format_v1) add_executable(test_task_api core/runtime/tests/test_task_api.cpp) target_include_directories(test_task_api PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include) - target_compile_options(test_task_api PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_task_api PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME task_api COMMAND test_task_api) add_executable(test_cli apps/cli/tests/test_cli.cpp) @@ -345,29 +429,33 @@ if(TRTMC_BUILD_TESTS) ${PROJECT_SOURCE_DIR}/apps ) target_link_libraries(test_cli PRIVATE trtmc_cli) - target_compile_options(test_cli PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_cli PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME cli COMMAND test_cli) set(_trtmc_test_runtime_root "${CMAKE_BINARY_DIR}/tests/runtime") add_library(trtmc_test_backend_fake SHARED core/runtime/tests/fake_backend.cpp) target_include_directories(trtmc_test_backend_fake PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include) - target_link_libraries(trtmc_test_backend_fake PRIVATE CUDA::cudart) + target_link_libraries(trtmc_test_backend_fake PRIVATE ${TRTMC_CUDART_LIBRARY}) set_target_properties(trtmc_test_backend_fake PROPERTIES OUTPUT_NAME trtmc_backend_fake LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" ) add_library(trtmc_test_backend_fake_rtx SHARED core/runtime/tests/fake_backend.cpp) target_include_directories(trtmc_test_backend_fake_rtx PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include ) - target_link_libraries(trtmc_test_backend_fake_rtx PRIVATE CUDA::cudart) + target_link_libraries(trtmc_test_backend_fake_rtx PRIVATE ${TRTMC_CUDART_LIBRARY}) target_compile_definitions(trtmc_test_backend_fake_rtx PRIVATE TRTMC_FAKE_BACKEND_NAME="trt_rtx" ) set_target_properties(trtmc_test_backend_fake_rtx PROPERTIES OUTPUT_NAME trtmc_backend_trt_rtx LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" ) add_library(trtmc_test_family_fake SHARED core/runtime/tests/fake_family.cpp) @@ -376,6 +464,8 @@ if(TRTMC_BUILD_TESTS) set_target_properties(trtmc_test_family_fake PROPERTIES OUTPUT_NAME trtmc_model_fake LIBRARY_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}" + RUNTIME_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" + ARCHIVE_OUTPUT_DIRECTORY "${_trtmc_test_runtime_root}/$<0:>" BUILD_RPATH "\$ORIGIN/../.." ) @@ -385,7 +475,7 @@ if(TRTMC_BUILD_TESTS) ${PROJECT_SOURCE_DIR}/core ) target_link_libraries(test_family_loader PRIVATE trtmc_runtime) - target_compile_options(test_family_loader PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_family_loader PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_family_loader trtmc_test_backend_fake trtmc_test_backend_fake_rtx @@ -393,6 +483,43 @@ if(TRTMC_BUILD_TESTS) ) add_test(NAME family_loader COMMAND test_family_loader "${_trtmc_test_runtime_root}") + add_executable(test_dynamic_library core/runtime/tests/test_dynamic_library.cpp) + target_include_directories(test_dynamic_library PRIVATE ${PROJECT_SOURCE_DIR}/core) + target_link_libraries(test_dynamic_library PRIVATE trtmc_runtime) + target_compile_definitions(test_dynamic_library PRIVATE + TRTMC_TEST_LIBRARY="$" + ) + target_compile_options(test_dynamic_library PRIVATE ${TRTMC_CXX_WARNINGS}) + add_dependencies(test_dynamic_library trtmc_test_backend_fake) + add_test(NAME dynamic_library COMMAND test_dynamic_library) + + add_executable(test_runtime_cache_persistence + core/runtime/tests/test_runtime_cache_persistence.cpp + core/runtime/tensorrt/runtime_cache_persistence.cpp + ) + target_include_directories(test_runtime_cache_persistence PRIVATE ${PROJECT_SOURCE_DIR}/core) + target_compile_options(test_runtime_cache_persistence PRIVATE ${TRTMC_CXX_WARNINGS}) + add_test(NAME runtime_cache_persistence COMMAND test_runtime_cache_persistence) + + if(WIN32) + add_executable(test_windows_utf8_argv + apps/cli/tests/test_windows_utf8_argv.cpp + apps/cli/windows_utf8_argv.cpp + apps/cli/trtmc_windows_utf8.manifest + ) + target_include_directories(test_windows_utf8_argv PRIVATE ${PROJECT_SOURCE_DIR}/apps) + add_test(NAME windows_utf8_argv COMMAND test_windows_utf8_argv) + + add_executable(test_windows_media apps/cli/tests/test_windows_media.cpp) + target_include_directories(test_windows_media PRIVATE + ${PROJECT_SOURCE_DIR}/apps + ${PROJECT_SOURCE_DIR}/core/runtime/include + ) + target_link_libraries(test_windows_media PRIVATE trtmc_cli) + add_test(NAME windows_media COMMAND test_windows_media) + endif() + + if(TRTMC_BUILD_BACKEND_TRT) add_executable(test_trt_module_dynamic_input core/runtime/tests/test_trt_module_dynamic_input.cpp ) @@ -408,25 +535,26 @@ if(TRTMC_BUILD_TESTS) trtmc_backend_trt trtmc_core ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ) - target_compile_options(test_trt_module_dynamic_input PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_trt_module_dynamic_input PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME trt_module_dynamic_input COMMAND test_trt_module_dynamic_input) set_tests_properties(trt_module_dynamic_input PROPERTIES SKIP_RETURN_CODE 77 LABELS gpu) + endif() if(TRTMC_BUILD_EXAMPLES) add_executable(test_dataset_answer apps/benchmark/tests/native/test_dataset_answer.cpp) target_include_directories(test_dataset_answer PRIVATE ${PROJECT_SOURCE_DIR}/apps/benchmark/native ) - target_compile_options(test_dataset_answer PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_dataset_answer PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME dataset_answer COMMAND test_dataset_answer) add_executable(test_benchmark_worker_e2e apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp ) target_link_libraries(test_benchmark_worker_e2e PRIVATE nlohmann_json::nlohmann_json) - target_compile_options(test_benchmark_worker_e2e PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_benchmark_worker_e2e PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_benchmark_worker_e2e trtmc_benchmark_worker trtmc_test_backend_fake @@ -453,7 +581,7 @@ if(TRTMC_BUILD_TESTS AND TRTMC_HAS_TVM_FFI) ${TRTMC_TRT_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) - target_compile_options(test_byok_shape_spec PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_byok_shape_spec PRIVATE ${TRTMC_CXX_WARNINGS}) add_test(NAME byok_shape_spec COMMAND test_byok_shape_spec) endif() @@ -466,7 +594,7 @@ if(TRTMC_BUILD_EXAMPLES AND TRTMC_HAS_TVM_FFI) ${TRTMC_TVM_FFI_INCLUDE_DIR} ) target_link_libraries(trtmc_byok_identity_copy PRIVATE - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) set_target_properties(trtmc_byok_identity_copy PROPERTIES @@ -488,14 +616,14 @@ if(TRTMC_BUILD_EXAMPLES AND TRTMC_HAS_TVM_FFI) target_link_libraries(test_byok_tvm_ffi PRIVATE trtmc_byok_tvm_ffi ${TRTMC_TRT_LIBRARY} - CUDA::cudart + ${TRTMC_CUDART_LIBRARY} ${TRTMC_TVM_FFI_LIBRARY} ) target_compile_definitions(test_byok_tvm_ffi PRIVATE TRTMC_HAS_TRT=1 TRTMC_HAS_TVM_FFI=1 ) - target_compile_options(test_byok_tvm_ffi PRIVATE -Wall -Wextra -Wpedantic) + target_compile_options(test_byok_tvm_ffi PRIVATE ${TRTMC_CXX_WARNINGS}) add_dependencies(test_byok_tvm_ffi trtmc_byok_identity_copy) add_test(NAME byok_tvm_ffi COMMAND test_byok_tvm_ffi $ @@ -510,15 +638,25 @@ install(TARGETS trtmc_core trtmc_runtime EXPORT trtmcTargets ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} ) -install(TARGETS trtmc_backend_trt - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} -) +if(TARGET trtmc_backend_trt) + install(TARGETS trtmc_backend_trt + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) +endif() if(TARGET trtmc_backend_rtx) install(TARGETS trtmc_backend_rtx + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} ) + if(WIN32) + install(FILES "${TRTMC_RTX_RUNTIME_DLL}" DESTINATION ${CMAKE_INSTALL_BINDIR}) + endif() endif() if(TRTMC_HAS_TVM_FFI) install(TARGETS trtmc_byok_tvm_ffi diff --git a/apps/benchmark/native/benchmark_worker.cpp b/apps/benchmark/native/benchmark_worker.cpp index 0abbf93e7a..533ccf2c31 100644 --- a/apps/benchmark/native/benchmark_worker.cpp +++ b/apps/benchmark/native/benchmark_worker.cpp @@ -421,10 +421,10 @@ Json run_generate_audio(trtmc::ITask& task, const Json& request, const Timing& t return measure( timing, [&]() { return interface.generate_audio(prompt, config); }, [](const trtmc::AudioResult& result) { - const double seconds = - result.sample_rate > 0 - ? static_cast(result.samples.size()) / result.sample_rate - : 0.0; + const double seconds = result.sample_rate > 0 && result.channels > 0 + ? static_cast(result.samples.size()) / + result.channels / result.sample_rate + : 0.0; return Json{{"output_samples", result.samples.size()}, {"num_samples", result.samples.size()}, {"output_audio_seconds", seconds}, @@ -458,10 +458,10 @@ Json run_speak(trtmc::ITask& task, const Json& request, const Timing& timing) { [](const auto& value) { const auto& result = value.first; return Json{{"input_audio_seconds", value.second}, - {"output_audio_seconds", - result.sample_rate > 0 - ? static_cast(result.samples.size()) / result.sample_rate - : 0.0}, + {"output_audio_seconds", result.sample_rate > 0 && result.channels > 0 + ? static_cast(result.samples.size()) / + result.channels / result.sample_rate + : 0.0}, {"output_samples", result.samples.size()}, {"num_samples", result.samples.size()}, {"sample_rate", result.sample_rate}}; diff --git a/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp b/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp index 0313a15f51..50c5231867 100644 --- a/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp +++ b/apps/benchmark/tests/native/test_benchmark_worker_e2e.cpp @@ -11,6 +11,10 @@ #include #include +#if defined(_WIN32) +#include +#endif + namespace { using Json = nlohmann::json; @@ -39,12 +43,14 @@ void write_bundle(const std::filesystem::path& path) { throw std::runtime_error("failed to write fake bundle"); } +#if !defined(_WIN32) std::string shell_quote(const std::string& value) { std::string result{"'"}; for (const char character : value) result += character == '\'' ? "'\\''" : std::string(1, character); return result + "'"; } +#endif } // namespace @@ -81,10 +87,19 @@ int main(int argc, char** argv) { throw std::runtime_error("failed to write worker request"); } +#if defined(_WIN32) + const std::string worker_argument = '"' + std::string(argv[1]) + '"'; + const std::string request_argument = '"' + request_path.string() + '"'; + const std::string output_argument = '"' + output_path.string() + '"'; + check(_spawnl(_P_WAIT, argv[1], worker_argument.c_str(), "--request", + request_argument.c_str(), "--output", output_argument.c_str(), nullptr) == 0, + "worker process completed"); +#else const std::string command = shell_quote(argv[1]) + " --request " + shell_quote(request_path.string()) + " --output " + shell_quote(output_path.string()); check(std::system(command.c_str()) == 0, "worker process completed"); +#endif std::ifstream output_file(output_path); Json result; diff --git a/apps/cli/cli.cpp b/apps/cli/cli.cpp index 7cf58f10d7..8785ee174d 100644 --- a/apps/cli/cli.cpp +++ b/apps/cli/cli.cpp @@ -6,6 +6,8 @@ #include "cli/cli.h" #include "cli/io.h" +#include "cli/windows_media.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_loader.h" #include @@ -14,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -123,8 +124,10 @@ const std::unordered_map& command_specs() { "--num-steps", "--guidance-scale", "--cfg-scale"}}}, {"generate-video", {CommandKind::kGenerateVideo, - {"--prompt", "--image", "--output", "--negative-prompt", "--height", "--width", - "--num-steps", "--seed", "--guidance-scale", "--cfg-scale", "--initial-latents-raw"}}}, + {"--prompt", "--output", "--negative-prompt", "--height", "--width", "--num-frames", + "--num-steps", "--seed", "--guidance-scale", "--cfg-scale", "--initial-latents-raw", + "--first-frame", "--last-frame", "--reference-image", "--reference-video", + "--reference-audio"}}}, {"solve", {CommandKind::kSolve, {"--branch", "--trunk"}}}, {"forecast", {CommandKind::kForecast, {"--input", "--mask", "--frequency"}}}, {"control", {CommandKind::kControl, {"--image", "--state", "--output"}}}, @@ -142,21 +145,24 @@ bool is_byok_option(const std::string& option) { void load_byok_extension(const Command& command) { using LoadKernelFn = const char* (*)(const char*, const char*, const char*) noexcept; - const fs::path extension = fs::path(command.runtime_root) / "libtrtmc_byok_tvm_ffi.so"; - dlerror(); - void* handle = dlopen(extension.c_str(), RTLD_NOW | RTLD_LOCAL); + const fs::path extension = + fs::path(command.runtime_root) / internal::dynamic_library_filename("trtmc_byok_tvm_ffi"); + std::string loader_error; + auto handle = internal::open_dynamic_library( + extension, internal::DynamicLibraryVisibility::local, &loader_error); if (handle == nullptr) { - const char* error = dlerror(); throw std::runtime_error("unable to load BYOK extension '" + extension.string() + - "': " + (error != nullptr ? error : "unknown dlopen error")); + "': " + loader_error); } - static auto* handles = new std::vector; - handles->push_back(handle); - dlerror(); - auto load = reinterpret_cast(dlsym(handle, "trtmc_load_byok_kernel")); - if (const char* error = dlerror(); error != nullptr || load == nullptr) { - throw std::runtime_error("BYOK extension is missing trtmc_load_byok_kernel"); + auto load = reinterpret_cast( + internal::dynamic_library_symbol(handle, "trtmc_load_byok_kernel", &loader_error)); + if (load == nullptr) { + (void)internal::close_dynamic_library(handle); + throw std::runtime_error("BYOK extension is missing trtmc_load_byok_kernel: " + + loader_error); } + static auto* handles = new std::vector; + handles->push_back(handle); if (const char* error = load(command.options.at("--byok-library").c_str(), command.options.at("--byok-function").c_str(), command.options.at("--byok-name").c_str())) { @@ -464,6 +470,7 @@ ImageGenerationConfig image_config(const Command& command) { config.negative_prompt = command.options.at("--negative-prompt"); config.height = int_option(command, "--height", 0, 1); config.width = int_option(command, "--width", 0, 1); + config.video_num_frames = int_option(command, "--num-frames", 0, 1); config.num_steps = int_option(command, "--num-steps", -1, 1); config.seed = int_option(command, "--seed", -1); config.guidance_scale = float_option(command, "--guidance-scale", -1.0F); @@ -538,6 +545,102 @@ ImageResult generate_image(const Command& command, ITask& task) { return require_interface(task).generate_image(prompt, config); } +VideoImageInput load_video_image(const std::string& path) { + auto decoded = read_image(path); + VideoImageInput result; + result.pixels = std::move(decoded.pixels); + result.height = decoded.height; + result.width = decoded.width; + result.channels = 3; + return result; +} + +VideoGenerationRequest video_request(const Command& command, IVideoGeneration& generator) { + VideoGenerationRequest request; + request.prompt = require_option(command, "--prompt"); + request.config = image_config(command); + + const bool has_first = has_option(command, "--first-frame"); + const bool has_last = has_option(command, "--last-frame"); + if ((has_first || has_last) && !command.video_references.empty()) + throw std::invalid_argument("key frames cannot be combined with media references"); + if (has_first || has_last) { + request.mode = VideoGenerationMode::kFirstLastFrameToVideoAudio; + if (has_first) + request.first_frame = load_video_image(command.options.at("--first-frame")); + if (has_last) + request.last_frame = load_video_image(command.options.at("--last-frame")); + return request; + } + if (command.video_references.empty()) + return request; + + request.mode = VideoGenerationMode::kReferenceToVideoAudio; + const auto policy = generator.reference_media_decode_policy(); + request.references.reserve(command.video_references.size()); + for (const auto& argument : command.video_references) { + VideoReferenceInput reference; + reference.kind = argument.kind; + switch (argument.kind) { + case VideoReferenceKind::kImage: + reference.image = load_video_image(argument.path); + break; + case VideoReferenceKind::kVideo: + if (!policy) + throw std::runtime_error( + "loaded family does not provide a reference-media decode policy"); + reference.video = read_video_file(argument.path, *policy); + break; + case VideoReferenceKind::kAudio: + if (!policy) + throw std::runtime_error( + "loaded family does not provide a reference-media decode policy"); + reference.audio = read_audio_file(argument.path, *policy); + break; + } + request.references.push_back(std::move(reference)); + } + return request; +} + +nlohmann::json write_generated_video(const VideoResult& result, const std::string& path) { + if (result.frames.pixels.empty() && result.frames.height == 0 && result.frames.width == 0 && + result.frames.channels == 3 && result.frames.num_frames == 0 && + result.audio.samples.empty() && result.fps == 0) { + return {{"worker", true}}; + } + validate_image_result(result.frames); + if (result.fps <= 0) + throw std::runtime_error("video result has a non-positive frame rate"); + if (!result.audio.samples.empty()) { + if (result.audio.sample_rate <= 0 || + (result.audio.channels != 1 && result.audio.channels != 2) || + result.audio.samples.size() % static_cast(result.audio.channels) != 0) { + throw std::runtime_error("video result has invalid interleaved audio metadata"); + } + require_finite(result.audio.samples, "video audio result"); + } + if (is_mp4_path(path)) { + write_mp4(result, path); + return {{"output", path}, + {"frames", result.frames.num_frames}, + {"fps", result.fps}, + {"height", result.frames.height}, + {"width", result.frames.width}, + {"audio_channels", result.audio.samples.empty() ? 0 : result.audio.channels}, + {"audio_sample_rate", result.audio.samples.empty() ? 0 : result.audio.sample_rate}}; + } + + auto payload = write_video(result.frames, path); + payload["fps"] = result.fps; + if (!result.audio.samples.empty()) { + const fs::path audio_path = fs::path(path) / "audio.wav"; + io::write_wav(result.audio, audio_path.string()); + payload["audio"] = audio_path.string(); + } + return payload; +} + const char* event_kind_name(SpeechSessionEventKind kind) { switch (kind) { case SpeechSessionEventKind::kAgentAudio: @@ -665,17 +768,17 @@ Command parse_args(int argc, char** argv) { if (name == "help") { if (argc != 2) throw std::invalid_argument("help does not accept arguments"); - return {CommandKind::kHelp, name, {}, {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kHelp, name, {}, {}, {}, {}, {}, {}, 0, {}, false}; } if (name == "version") { if (argc != 2) throw std::invalid_argument("version does not accept arguments"); - return {CommandKind::kVersion, name, {}, {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kVersion, name, {}, {}, {}, {}, {}, {}, 0, {}, false}; } if (name == "inspect") { if (argc != 3 || std::string(argv[2]).empty()) throw std::invalid_argument("inspect requires exactly one BUNDLE path"); - return {CommandKind::kInspect, name, argv[2], {}, {}, {}, {}, 0, {}, false}; + return {CommandKind::kInspect, name, argv[2], {}, {}, {}, {}, {}, 0, {}, false}; } const auto spec = command_specs().find(name); @@ -684,7 +787,7 @@ Command parse_args(int argc, char** argv) { if (argc < 3 || std::string(argv[2]).empty()) throw std::invalid_argument(name + " requires a BUNDLE path"); - Command command{spec->second.kind, name, argv[2], {}, {}, {}, {}, 0, {}, false}; + Command command{spec->second.kind, name, argv[2], {}, {}, {}, {}, {}, 0, {}, false}; for (int index = 3; index < argc; ++index) { const std::string option = argv[index]; if (option == "--runtime-root") { @@ -721,6 +824,16 @@ Command parse_args(int argc, char** argv) { command.inputs.push_back(take_value(argc, argv, index, option)); continue; } + if (option == "--reference-image" || option == "--reference-video" || + option == "--reference-audio") { + VideoReferenceKind kind = VideoReferenceKind::kImage; + if (option == "--reference-video") + kind = VideoReferenceKind::kVideo; + else if (option == "--reference-audio") + kind = VideoReferenceKind::kAudio; + command.video_references.push_back({kind, take_value(argc, argv, index, option)}); + continue; + } if (command.options.count(option) != 0) throw std::invalid_argument(option + " may be specified only once"); command.options.emplace(option, take_value(argc, argv, index, option)); @@ -1195,8 +1308,10 @@ int dispatch(const Command& command, ITask& task, std::ostream& output) { return EXIT_SUCCESS; } case CommandKind::kGenerateVideo: { - const std::string directory = require_option(command, "--output"); - write_json(output, write_video(generate_image(command, task), directory)); + const std::string path = require_option(command, "--output"); + auto& generator = require_interface(task); + auto request = video_request(command, generator); + write_json(output, write_generated_video(generator.generate_video(request), path)); return EXIT_SUCCESS; } case CommandKind::kSolve: { @@ -1305,6 +1420,10 @@ void print_usage(std::ostream& output) { " [--kv-cache-size BYTES|GB|GiB]\n\n" "TensorRT-RTX runtime options:\n" " [--runtime-cache PATH] [--cuda-graphs]\n\n" + "Video generation options:\n" + " --prompt TEXT --output OUTPUT.mp4 [--num-frames N] [--height N] [--width N]\n" + " [--first-frame IMAGE] [--last-frame IMAGE]\n" + " [--reference-image IMAGE|--reference-video VIDEO|--reference-audio AUDIO]...\n\n" "Execution never searches for runtimes; --runtime-root is always required.\n"; } diff --git a/apps/cli/cli.h b/apps/cli/cli.h index 9fe114557f..13bb778da2 100644 --- a/apps/cli/cli.h +++ b/apps/cli/cli.h @@ -48,6 +48,11 @@ enum class CommandKind { kGenerateWorld, }; +struct VideoReferenceArgument { + VideoReferenceKind kind{VideoReferenceKind::kImage}; + std::string path; +}; + struct Command { CommandKind kind{CommandKind::kHelp}; std::string name; @@ -56,6 +61,7 @@ struct Command { std::unordered_map options; std::vector frames; std::vector inputs; + std::vector video_references; std::uint64_t kv_cache_size_bytes{0}; std::string runtime_cache_path; bool cuda_graphs{false}; diff --git a/apps/cli/io.cpp b/apps/cli/io.cpp index 472c86e721..d061039846 100644 --- a/apps/cli/io.cpp +++ b/apps/cli/io.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -24,18 +25,34 @@ namespace trtmc::cli::io { void write_wav(const AudioResult& audio, const std::string& path) { if (audio.samples.empty()) throw std::runtime_error("write_wav: empty audio"); + if (audio.sample_rate <= 0) + throw std::runtime_error("write_wav: sample rate must be positive"); + if (audio.channels != 1 && audio.channels != 2) + throw std::runtime_error("write_wav: channel count must be mono or stereo"); + if (audio.samples.size() % static_cast(audio.channels) != 0) + throw std::runtime_error("write_wav: interleaved sample count is not channel-aligned"); + if (audio.samples.size() > + static_cast(std::numeric_limits::max() / sizeof(float))) + throw std::runtime_error("write_wav: sample buffer is too large for RIFF/WAVE"); + if (audio.num_samples != 0 && + audio.num_samples != static_cast(audio.samples.size())) + throw std::runtime_error("write_wav: num_samples does not match the sample buffer"); std::ofstream output(path, std::ios::binary); if (!output) throw std::runtime_error("write_wav: cannot open " + path); - const auto num_samples = static_cast(audio.samples.size()); const std::int32_t sample_rate = audio.sample_rate; - const std::int16_t num_channels = 1; + const auto num_channels = static_cast(audio.channels); const std::int16_t bits_per_sample = 32; - const std::int32_t byte_rate = sample_rate * num_channels * (bits_per_sample / 8); const auto block_align = static_cast(num_channels * (bits_per_sample / 8)); - const std::int32_t data_size = num_samples * block_align; + const auto byte_rate_64 = static_cast(sample_rate) * block_align; + if (byte_rate_64 > std::numeric_limits::max()) + throw std::runtime_error("write_wav: byte rate exceeds the RIFF/WAVE range"); + const auto byte_rate = static_cast(byte_rate_64); + const auto data_size = static_cast(audio.samples.size() * sizeof(float)); + if (data_size > std::numeric_limits::max() - 36) + throw std::runtime_error("write_wav: output exceeds the RIFF/WAVE range"); const std::int32_t chunk_size = 36 + data_size; const std::int32_t format_size = 16; const std::int16_t audio_format = 3; diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index dbb5c93bb8..c853030487 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -4,9 +4,35 @@ */ #include "cli/cli.h" +#if defined(_WIN32) +#include "cli/windows_utf8_argv.h" +#endif +#include +#include #include -int main(int argc, char** argv) { +namespace { + +int run_cli(int argc, char** argv) { return trtmc::cli::run(argc, argv, std::cout, std::cerr); } + +} // namespace + +#if defined(_WIN32) +int wmain(int argc, wchar_t** argv) { + try { + trtmc::cli::Utf8CommandLine command_line(argc, argv); + return run_cli(command_line.argc(), command_line.argv()); + } catch (const std::exception& error) { + std::cerr << "Error: unable to decode the Windows command line as UTF-8: " << error.what() + << '\n'; + return EXIT_FAILURE; + } +} +#else +int main(int argc, char** argv) { + return run_cli(argc, argv); +} +#endif diff --git a/apps/cli/tests/test_cli.cpp b/apps/cli/tests/test_cli.cpp index 0824f5c86c..605f0529c7 100644 --- a/apps/cli/tests/test_cli.cpp +++ b/apps/cli/tests/test_cli.cpp @@ -5,7 +5,9 @@ #include "cli/cli.h" #include "cli/io.h" +#include "cli/windows_media.h" +#include #include #include #include @@ -21,6 +23,31 @@ namespace { int failures = 0; +class TestTempDirectory { + public: + TestTempDirectory() { + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + path_ = + std::filesystem::temp_directory_path() / ("trtmc-cli-test-" + std::to_string(nonce)); + std::filesystem::create_directory(path_); + } + + ~TestTempDirectory() { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + std::filesystem::path file(const char* name) const { return path_ / name; } + + private: + std::filesystem::path path_; +}; + +const TestTempDirectory& test_temp_directory() { + static const TestTempDirectory directory; + return directory; +} + void check(bool condition, const char* name) { if (!condition) { std::cerr << "FAIL: " << name << '\n'; @@ -89,7 +116,7 @@ class BareTextTask final : public trtmc::ITask { const char* task() const noexcept override { return trtmc::ITextGeneration::kTask; } }; -class EmptyImageWorker final : public trtmc::IImageGeneration { +class EmptyImageWorker final : public trtmc::IImageGeneration, public trtmc::IVideoGeneration { public: trtmc::ImageGenerationConfig seen; @@ -100,6 +127,18 @@ class EmptyImageWorker final : public trtmc::IImageGeneration { result.num_frames = 0; return result; } + + trtmc::VideoResult generate_video(const std::string&, + const trtmc::ImageGenerationConfig& config) override { + seen = config; + trtmc::VideoResult result; + result.frames.num_frames = 0; + return result; + } + + trtmc::VideoResult generate_video(const trtmc::VideoGenerationRequest& request) override { + return generate_video(request.prompt, request.config); + } }; class FakeStreamingAudio final : public trtmc::IAudioGeneration, @@ -388,9 +427,25 @@ int main() { "image replay latents option is retained"); const auto video_generation = parse({"trtmc", "generate-video", "model.bundle", "--runtime-root", "lib", "--prompt", - "cat", "--output", "cat.mp4", "--initial-latents-raw", "initial.f32"}); - check(video_generation.options.at("--initial-latents-raw") == "initial.f32", - "video replay latents option is retained"); + "cat", "--output", "cat.mp4", "--initial-latents-raw", "initial.f32", "--num-frames", + "124", "--first-frame", "first.png", "--last-frame", "last.png"}); + check(video_generation.options.at("--initial-latents-raw") == "initial.f32" && + video_generation.options.at("--num-frames") == "124" && + video_generation.options.at("--first-frame") == "first.png" && + video_generation.options.at("--last-frame") == "last.png", + "structured first/last-frame video options are retained"); + const auto reference_generation = + parse({"trtmc", "generate-video", "model.bundle", "--runtime-root", "lib", "--prompt", + "cat", "--output", "cat.mp4", "--reference-video", "clip.mp4", "--reference-image", + "still.png", "--reference-audio", "sound.wav"}); + check(reference_generation.video_references.size() == 3 && + reference_generation.video_references[0].kind == trtmc::VideoReferenceKind::kVideo && + reference_generation.video_references[0].path == "clip.mp4" && + reference_generation.video_references[1].kind == trtmc::VideoReferenceKind::kImage && + reference_generation.video_references[1].path == "still.png" && + reference_generation.video_references[2].kind == trtmc::VideoReferenceKind::kAudio && + reference_generation.video_references[2].path == "sound.wav", + "mixed video references preserve CLI order and media kind"); const auto batch = parse({"trtmc", "generate-image-batch", "model.bundle", "--runtime-root", "lib", "--prompts", "prompts.txt", "--seeds", "1,2", "--output", "images"}); @@ -441,9 +496,10 @@ int main() { run_command.options.emplace("--threshold", "0.75"); run_command.options.emplace("--use-chat-template", "true"); run_command.options.emplace("--enable-thinking", "false"); - run_command.options.emplace("--lora-adapter", "/tmp/adapter"); + const auto adapter_path = test_temp_directory().file("adapter").string(); + run_command.options.emplace("--lora-adapter", adapter_path); run_command.options.emplace("--lora-adapter-id", "demo"); - const std::filesystem::path replay_values_path = "/tmp/trtmc-cli-replay-values.f32"; + const auto replay_values_path = test_temp_directory().file("replay-values.f32"); const float replay_values[] = {0.25F, -0.5F}; { std::ofstream replay_file(replay_values_path, std::ios::binary); @@ -478,7 +534,7 @@ int main() { text.seen.condition_mask == text.seen.initial_latents && text.seen.sampling_steps == text.seen.initial_latents && text.seen.sde_noises == text.seen.initial_latents && - text.loaded_adapter_id == "demo" && text.loaded_adapter_path == "/tmp/adapter", + text.loaded_adapter_id == "demo" && text.loaded_adapter_path == adapter_path, "text sampling options reach the Task API"); auto invalid_block = run_command; invalid_block.options["--block-length"] = "invalid"; @@ -500,7 +556,7 @@ int main() { text.seen.text_generation_mode == "auto" && text.seen.block_length == 0 && text.seen.confidence_threshold == -1.0F && text.seen.temperature == 1.0F, "text diffusion options preserve Task API defaults"); - const std::filesystem::path unsupported_image_path = "/tmp/trtmc-cli-unsupported.ppm"; + const auto unsupported_image_path = test_temp_directory().file("unsupported.ppm"); { std::ofstream image_file(unsupported_image_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -520,7 +576,7 @@ int main() { edit_command.name = "generate-image"; edit_command.options.emplace("--prompt", "make it night"); edit_command.options.emplace("--image", unsupported_image_path.string()); - edit_command.options.emplace("--output", "/tmp/trtmc-cli-unused.png"); + edit_command.options.emplace("--output", test_temp_directory().file("unused.png").string()); EmptyImageWorker image_only; check(dispatch_throws(edit_command, image_only), "image-generation task rejects an edit image instead of dropping it"); @@ -568,7 +624,7 @@ int main() { trtmc::cli::Command audio_command; audio_command.kind = trtmc::cli::CommandKind::kGenerateAudio; audio_command.name = "generate-audio"; - const std::filesystem::path audio_path = "/tmp/trtmc-cli-stream-test.raw"; + const auto audio_path = test_temp_directory().file("stream-test.raw"); audio_command.options.emplace("--prompt", "hello"); audio_command.options.emplace("--output", audio_path.string()); audio_command.options.emplace("--max-new-tokens", "9"); @@ -587,7 +643,7 @@ int main() { "streaming audio output format is explicit"); std::filesystem::remove(audio_path); - const std::filesystem::path transcription_path = "/tmp/trtmc-cli-transcription-stream.wav"; + const auto transcription_path = test_temp_directory().file("transcription-stream.wav"); trtmc::AudioResult transcription_audio; transcription_audio.samples = {0.25F, -0.5F, 0.75F}; transcription_audio.num_samples = 3; @@ -722,7 +778,7 @@ int main() { usage.str().find("--cuda-graphs") != std::string::npos, "help lists direct Task and backend options"); - const std::filesystem::path video_path = "/tmp/trtmc-cli-video-frame.ppm"; + const auto video_path = test_temp_directory().file("video-frame.ppm"); { std::ofstream image_file(video_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -738,8 +794,8 @@ int main() { check(video_task.prompt.empty(), "missing video prompt reaches the family as empty"); std::filesystem::remove(video_path); - const std::filesystem::path forecast_values_path = "/tmp/trtmc-cli-forecast-values.f32"; - const std::filesystem::path forecast_mask_path = "/tmp/trtmc-cli-forecast-mask.f32"; + const auto forecast_values_path = test_temp_directory().file("forecast-values.f32"); + const auto forecast_mask_path = test_temp_directory().file("forecast-mask.f32"); const float forecast_values[] = {1.0F, 2.0F}; const float forecast_mask[] = {1.0F, 1.0F}; { @@ -762,9 +818,9 @@ int main() { std::filesystem::remove(forecast_values_path); std::filesystem::remove(forecast_mask_path); - const std::filesystem::path control_image_path = "/tmp/trtmc-cli-control.ppm"; - const std::filesystem::path control_state_path = "/tmp/trtmc-cli-control-state.f32"; - const std::filesystem::path control_output_path = "/tmp/trtmc-cli-control-actions.f32"; + const auto control_image_path = test_temp_directory().file("control.ppm"); + const auto control_state_path = test_temp_directory().file("control-state.f32"); + const auto control_output_path = test_temp_directory().file("control-actions.f32"); { std::ofstream image_file(control_image_path, std::ios::binary); image_file << "P6\n1 1\n255\n"; @@ -801,7 +857,16 @@ int main() { return true; } }; - const std::filesystem::path wav_path = "/tmp/trtmc-cli-io.wav"; + const auto wav_path = test_temp_directory().file("io.wav"); + trtmc::ReferenceMediaDecodePolicy canvas_policy{15, 24, 240, 1, 1, 32, 0.25, 4.0}; + check(trtmc::cli::detail::reference_video_decode_size(canvas_policy, 1, 1) == + std::make_pair(32U, 32U), + "canvas axes clamp to one alignment unit on every platform"); + canvas_policy.canvas_short_edge = 80; + canvas_policy.canvas_max_pixels = 80 * 80; + check(trtmc::cli::detail::reference_video_decode_size(canvas_policy, 1, 1) == + std::make_pair(64U, 64U), + "canvas rounding uses ties-to-even on every platform"); trtmc::AudioResult wav; wav.samples = {-1.0F, 0.25F, 1.0F}; wav.num_samples = 3; @@ -816,7 +881,7 @@ int main() { check(throws_runtime([&] { trtmc::cli::io::write_wav({}, wav_path.string()); }), "empty WAV is rejected"); - const std::filesystem::path png_path = "/tmp/trtmc-cli-io.png"; + const auto png_path = test_temp_directory().file("io.png"); const std::vector rgb{-0.5F, 0.5F, 1.5F}; trtmc::cli::io::save_png(png_path.string(), rgb, 1, 1); const auto loaded_image = trtmc::cli::io::read_image(png_path.string()); diff --git a/apps/cli/tests/test_windows_media.cpp b/apps/cli/tests/test_windows_media.cpp new file mode 100644 index 0000000000..715e928d34 --- /dev/null +++ b/apps/cli/tests/test_windows_media.cpp @@ -0,0 +1,559 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/io.h" +#include "cli/windows_media.h" + +#define WIN32_LEAN_AND_MEAN +#define NOMINMAX +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Microsoft::WRL::ComPtr; + +class TempDirectory { + public: + TempDirectory() { + path_ = std::filesystem::temp_directory_path() / + ("trtmc-windows-media-test-" + std::to_string(GetCurrentProcessId()) + "-" + + std::to_string(GetTickCount64())); + std::filesystem::create_directory(path_); + } + + ~TempDirectory() { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + + const std::filesystem::path& path() const { return path_; } + + private: + std::filesystem::path path_; +}; + +void require(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} + +void require_success(HRESULT result) { + require(SUCCEEDED(result), "Media Foundation call failed"); +} + +template +void require_throws_with(Fn&& fn, const char* needle, const char* message) { + try { + fn(); + } catch (const std::exception& error) { + require(std::string(error.what()).find(needle) != std::string::npos, message); + return; + } + throw std::runtime_error(message); +} + +} // namespace + +int run_test() { + constexpr trtmc::ReferenceMediaDecodePolicy policy{ + 15, 24, 240, 768, 768ULL * 1344, 32, 0.25, 4.0, + }; + constexpr trtmc::ReferenceMediaDecodePolicy compact_policy{ + 2, 12, 60, 256, 256ULL * 448, 16, 0.5, 2.0, + }; + require(trtmc::cli::detail::reference_video_frame_ceiling(compact_policy, 30'000, 1'001) == 60, + "frame ceiling must follow a non-default duration policy"); + require(trtmc::cli::detail::reference_video_decode_size(compact_policy, 3840, 2160) == + std::pair{448, 256}, + "decode canvas must follow non-default size and alignment policy values"); + require_throws_with( + [&] { (void)trtmc::cli::detail::reference_video_decode_size(compact_policy, 2560, 1080); }, + "configured range", "decode aspect bounds must follow the supplied policy"); + require(trtmc::cli::detail::reference_video_frame_ceiling(policy, 30'000, 1'001) == 450, + "configured video ceiling must round up fractional frame rates"); + require(trtmc::cli::detail::reference_video_decode_size(policy, 3840, 2160) == + std::pair{1344, 768}, + "4K reference video must decode onto the policy's bounded canvas"); + require(trtmc::cli::detail::reference_timeline_within_limit(policy, 0, 150'000'000), + "an exact configured presentation timeline must be accepted"); + require(!trtmc::cli::detail::reference_timeline_within_limit(policy, 150'000'000, 1), + "a non-empty sample starting at the duration limit must be rejected"); + require(!trtmc::cli::detail::reference_timeline_within_limit(policy, 149'000'000, 2'000'000), + "a sparse sample crossing the duration limit must be rejected"); + + std::vector timeline; + const std::vector first_audio{1.0F, -1.0F, 2.0F, -2.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, first_audio.data(), first_audio.size() / 2, 0, first_audio.size() / 2, 2, 20'000, + 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F}), + "a positive first audio PTS must preserve leading silence"); + const std::vector gapped_audio{3.0F, -3.0F, 4.0F, -4.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, gapped_audio.data(), gapped_audio.size() / 2, 0, gapped_audio.size() / 2, 2, + 60'000, 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F, 0.0F, + 0.0F, 0.0F, 0.0F, 3.0F, -3.0F, 4.0F, -4.0F}), + "an audio PTS discontinuity must preserve inter-sample silence"); + const std::vector overlapped_audio{9.0F, -9.0F, 10.0F, -10.0F, + 11.0F, -11.0F, 12.0F, -12.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, overlapped_audio.data(), overlapped_audio.size() / 2, 0, + overlapped_audio.size() / 2, 2, 70'000, 1'000, 64); + require(timeline == std::vector({0.0F, 0.0F, 0.0F, 0.0F, 1.0F, -1.0F, 2.0F, -2.0F, + 0.0F, 0.0F, 0.0F, 0.0F, 3.0F, -3.0F, 4.0F, -4.0F, + 10.0F, -10.0F, 11.0F, -11.0F, 12.0F, -12.0F}), + "a later overlapping audio sample must deterministically retain only its new tail"); + const auto timeline_before_contained_overlap = timeline; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + timeline, first_audio.data(), first_audio.size() / 2, 0, first_audio.size() / 2, 2, 0, + 1'000, 64); + require(timeline == timeline_before_contained_overlap, + "a wholly overlapped later audio sample must not alter the existing timeline"); + + std::vector trimmed_codec_timeline; + const std::vector primed_audio{99.0F, -99.0F, 1.0F, -1.0F, 2.0F, -2.0F}; + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + trimmed_codec_timeline, primed_audio.data(), primed_audio.size() / 2, 1, 3, 2, -10'000, + 1'000, 16); + require(trimmed_codec_timeline == std::vector({1.0F, -1.0F, 2.0F, -2.0F}), + "negative-PTS codec priming must start at zero after using the retained source offset"); + require_throws_with( + [&] { + trtmc::cli::detail::append_reference_audio_frames_on_timeline( + trimmed_codec_timeline, primed_audio.data(), 2, 0, 3, 2, 0, 1'000, 16); + }, + "invalid timeline layout", + "a retained audio window outside its source buffer must fail closed"); + + trtmc::cli::detail::validate_reference_video_soundtrack_format(32'000, 1); + trtmc::cli::detail::validate_reference_video_soundtrack_format(48'000, 2); + require_throws_with( + [] { trtmc::cli::detail::validate_reference_video_soundtrack_format(0, 2); }, + "invalid sample rate", "a video soundtrack with an invalid sample rate must fail closed"); + require_throws_with( + [] { trtmc::cli::detail::validate_reference_video_soundtrack_format(32'000, 3); }, + "mono or stereo", "a multichannel video soundtrack must fail closed"); + + constexpr std::uint32_t synthetic_codec_rate = 32'000; + constexpr std::uint64_t synthetic_mp3_access_unit_frames = 1'152; + constexpr std::uint64_t synthetic_mp3_padding_frames = 3 * synthetic_mp3_access_unit_frames; + constexpr std::uint64_t synthetic_mp3_padding_ticks = + (synthetic_mp3_padding_frames * 10'000'000 + synthetic_codec_rate - 1) / + synthetic_codec_rate; + constexpr std::uint64_t synthetic_public_audio_frames = 15 * synthetic_codec_rate; + require(trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, -static_cast(synthetic_mp3_padding_ticks), + synthetic_mp3_padding_ticks) && + trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, 150'000'000 + static_cast(synthetic_mp3_padding_ticks), + synthetic_mp3_padding_ticks), + "compressed empty-event timestamps must accept the exact codec-padding window"); + require(!trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, -36'000'000'000, synthetic_mp3_padding_ticks) && + !trtmc::cli::detail::reference_audio_event_timestamp_within_padding( + policy, 36'000'000'000, synthetic_mp3_padding_ticks), + "compressed empty-event timestamps at negative or positive hours must be rejected"); + trtmc::cli::detail::ReferenceAudioDecodeState boundary_padding_state; + require(trtmc::cli::detail::account_reference_audio_decode( + policy, 0, synthetic_public_audio_frames + synthetic_mp3_padding_frames, + synthetic_codec_rate, synthetic_mp3_padding_frames, boundary_padding_state), + "an exact three-access-unit MP3 decoded tail must be accepted"); + require(boundary_padding_state.decoded_padding_frames == synthetic_mp3_padding_frames, + "synthetic MP3 tail accounting must reach the codec-padding boundary"); + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'001, 1, synthetic_codec_rate, synthetic_mp3_padding_frames, + boundary_padding_state), + "one decoded frame beyond the MP3 padding boundary must be rejected"); + + trtmc::cli::detail::ReferenceAudioDecodeState falsified_duration_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'000, synthetic_mp3_padding_frames + 1, synthetic_codec_rate, + synthetic_mp3_padding_frames, falsified_duration_state), + "a short or falsified presentation duration must not hide a decoded tail over three " + "MP3 access units"); + + trtmc::cli::detail::ReferenceAudioDecodeState future_timestamp_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 36'000'000'000, 1, synthetic_codec_rate, synthetic_mp3_padding_frames, + future_timestamp_state), + "a tiny decoded sample with a far-future timestamp must not spend only one padding " + "frame"); + trtmc::cli::detail::ReferenceAudioDecodeState excessive_leading_state; + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, -130'000'000, synthetic_public_audio_frames, synthetic_codec_rate, + synthetic_mp3_padding_frames, excessive_leading_state), + "excessive negative-timestamp decoded PCM must not be hidden as leading padding"); + + constexpr std::uint64_t synthetic_aac_padding_frames = 3 * 1'024; + trtmc::cli::detail::ReferenceAudioDecodeState aac_padding_state; + require(trtmc::cli::detail::account_reference_audio_decode( + policy, 0, synthetic_public_audio_frames + synthetic_aac_padding_frames, + synthetic_codec_rate, synthetic_aac_padding_frames, aac_padding_state), + "an exact three-access-unit AAC decoded tail must be accepted"); + require(!trtmc::cli::detail::account_reference_audio_decode( + policy, 150'000'001, 1, synthetic_codec_rate, synthetic_aac_padding_frames, + aac_padding_state), + "one decoded frame beyond the AAC padding boundary must be rejected"); + + trtmc::VideoResult result; + result.frames.width = 64; + result.frames.height = 64; + result.frames.channels = 3; + result.frames.num_frames = 24; + result.fps = 24; + const std::size_t pixel_count = static_cast(result.frames.width) * + result.frames.height * result.frames.channels * + result.frames.num_frames; + result.frames.pixels.resize(pixel_count); + for (int frame = 0; frame < result.frames.num_frames; ++frame) { + for (int row = 0; row < result.frames.height; ++row) { + for (int column = 0; column < result.frames.width; ++column) { + const auto offset = + ((static_cast(frame) * result.frames.height + row) * + result.frames.width + + column) * + 3; + result.frames.pixels[offset] = static_cast(column) / 63.0F; + result.frames.pixels[offset + 1] = static_cast(row) / 63.0F; + result.frames.pixels[offset + 2] = static_cast(frame) / 23.0F; + } + } + } + + result.audio.sample_rate = 32000; + result.audio.channels = 2; + result.audio.samples.resize(static_cast(result.audio.sample_rate) * 2); + for (int frame = 0; frame < result.audio.sample_rate; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / result.audio.sample_rate); + result.audio.samples[static_cast(frame) * 2] = value; + result.audio.samples[static_cast(frame) * 2 + 1] = value; + } + result.audio.num_samples = static_cast(result.audio.samples.size()); + + TempDirectory temporary; + const auto& temporary_root = temporary.path(); + const auto path = temporary_root / "video.mp4"; + const auto wav_path = temporary_root / "audio.wav"; + const auto mp3_path = temporary_root / "audio.mp3"; + const auto aac_path = temporary_root / "audio.m4a"; + const auto mp3_44k_path = temporary_root / "audio-44k.mp3"; + const auto aac_44k_path = temporary_root / "audio-44k.m4a"; + const auto over_limit_mp3_path = temporary_root / "over-limit.mp3"; + const auto boundary_wav_path = temporary_root / "boundary.wav"; + const auto over_limit_wav_path = temporary_root / "over-limit.wav"; + const auto over_limit_video_path = temporary_root / "over-limit.mp4"; + trtmc::cli::write_mp4(result, path.string()); + require(std::filesystem::is_regular_file(path), "MP4 output file is missing"); + require(std::filesystem::file_size(path) > 1024, "MP4 output file is empty"); + + const auto decoded = trtmc::cli::read_video_file(path.string(), policy); + require(decoded.width == 768, "decoded MP4 width must use the configured resolver"); + require(decoded.height == 768, "decoded MP4 height must use the configured resolver"); + require(decoded.channels == 3, "decoded MP4 channel mismatch"); + require(decoded.num_frames == result.frames.num_frames, "decoded MP4 frame-count mismatch"); + require(decoded.fps_numerator == result.fps, "decoded MP4 frame-rate mismatch"); + require(decoded.fps_denominator == 1, "decoded MP4 frame-rate denominator mismatch"); + require(decoded.pixels.size() == static_cast(decoded.width) * decoded.height * + decoded.num_frames * decoded.channels, + "decoded MP4 pixel-count mismatch"); + require(decoded.soundtrack.sample_rate == result.audio.sample_rate, + "decoded MP4 audio-rate mismatch"); + require(decoded.soundtrack.channels == result.audio.channels, + "decoded MP4 audio-channel mismatch"); + require(!decoded.soundtrack.samples.empty(), "decoded MP4 has no soundtrack samples"); + + const auto compact_decoded = trtmc::cli::read_video_file(path.string(), compact_policy); + require(compact_decoded.width == 256 && compact_decoded.height == 256, + "decoded MP4 canvas must follow a second policy"); + require(compact_decoded.num_frames == 12 && compact_decoded.fps_numerator == 12 && + compact_decoded.fps_denominator == 1, + "decoded MP4 sampling rate must follow a second policy"); + + const auto extracted_audio = trtmc::cli::read_audio_file(path.string(), policy); + require(extracted_audio.sample_rate == result.audio.sample_rate, + "standalone media audio-rate mismatch"); + require(extracted_audio.channels == result.audio.channels, + "standalone media audio-channel mismatch"); + require(!extracted_audio.samples.empty(), "standalone media audio decode is empty"); + + trtmc::cli::io::write_wav(result.audio, wav_path.string()); + const auto decoded_wav = trtmc::cli::read_audio_file(wav_path.string(), policy); + require(decoded_wav.sample_rate == result.audio.sample_rate, + "Media Foundation WAV audio-rate mismatch"); + require(decoded_wav.channels == result.audio.channels, + "Media Foundation WAV audio-channel mismatch"); + require(decoded_wav.num_samples == static_cast(decoded_wav.samples.size()) && + !decoded_wav.samples.empty(), + "Media Foundation WAV decode returned invalid samples"); + + trtmc::AudioResult boundary_audio; + boundary_audio.sample_rate = 8000; + boundary_audio.channels = 1; + boundary_audio.samples.resize(static_cast(boundary_audio.sample_rate) * 15); + boundary_audio.num_samples = static_cast(boundary_audio.samples.size()); + trtmc::cli::io::write_wav(boundary_audio, boundary_wav_path.string()); + const auto decoded_boundary_audio = + trtmc::cli::read_audio_file(boundary_wav_path.string(), policy); + require(decoded_boundary_audio.samples.size() == boundary_audio.samples.size(), + "configured reference audio boundary must decode successfully"); + + boundary_audio.samples.push_back(0.0F); + boundary_audio.num_samples = static_cast(boundary_audio.samples.size()); + trtmc::cli::io::write_wav(boundary_audio, over_limit_wav_path.string()); + require_throws_with( + [&] { (void)trtmc::cli::read_audio_file(over_limit_wav_path.string(), policy); }, + "configured duration limit", + "reference audio over the policy limit must fail during Media Foundation decode"); + + trtmc::VideoResult over_limit_video; + over_limit_video.frames.width = 64; + over_limit_video.frames.height = 64; + over_limit_video.frames.channels = 3; + over_limit_video.frames.num_frames = 361; + over_limit_video.fps = 24; + over_limit_video.frames.pixels.resize( + static_cast(over_limit_video.frames.width) * over_limit_video.frames.height * + over_limit_video.frames.channels * over_limit_video.frames.num_frames); + trtmc::cli::write_mp4(over_limit_video, over_limit_video_path.string()); + require_throws_with( + [&] { (void)trtmc::cli::read_video_file(over_limit_video_path.string(), policy); }, + "configured duration limit", + "reference video over the policy limit must fail during Media Foundation decode"); + + const auto decoded_pixel = [&](int frame, int row, int column, int channel) { + return decoded + .pixels[((static_cast(frame) * decoded.height + row) * decoded.width + + column) * + 3 + + channel]; + }; + require(decoded_pixel(0, 384, 672, 0) > decoded_pixel(0, 384, 96, 0), + "decoded MP4 red axis is reversed"); + require(decoded_pixel(0, 672, 384, 1) > decoded_pixel(0, 96, 384, 1), + "decoded MP4 green axis is reversed"); + require(decoded_pixel(20, 384, 384, 2) > decoded_pixel(3, 384, 384, 2), + "decoded MP4 frame order is reversed"); + + const HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + const bool owns_com = SUCCEEDED(com_result); + require(owns_com || com_result == RPC_E_CHANGED_MODE, "CoInitializeEx failed"); + require_success(MFStartup(MF_VERSION, MFSTARTUP_FULL)); + + ComPtr empty_audio_buffer; + require_success(MFCreateMemoryBuffer(sizeof(float), &empty_audio_buffer)); + DWORD empty_audio_length = 1; + require_success(empty_audio_buffer->GetCurrentLength(&empty_audio_length)); + require(empty_audio_length == 0, + "synthetic Media Foundation audio buffer must start with zero current length"); + ComPtr empty_audio_sample; + require_success(MFCreateSample(&empty_audio_sample)); + require_success(empty_audio_sample->AddBuffer(empty_audio_buffer.Get())); + DWORD empty_audio_sample_length = 1; + require_success(empty_audio_sample->GetTotalLength(&empty_audio_sample_length)); + require(empty_audio_sample_length == 0, + "synthetic non-null audio sample must contain no payload"); + require_throws_with( + [&] { + trtmc::cli::detail::require_nonempty_decoded_buffer(empty_audio_sample_length, "audio"); + }, + "empty decoded audio buffer", + "a non-null sample with a zero-length audio buffer must fail closed"); + + ComPtr empty_video_sample; + require_success(MFCreateSample(&empty_video_sample)); + ComPtr empty_video_buffer; + require_success(MFCreateMemoryBuffer(4, &empty_video_buffer)); + require_success(empty_video_sample->AddBuffer(empty_video_buffer.Get())); + DWORD empty_video_length = 1; + require_success(empty_video_sample->GetTotalLength(&empty_video_length)); + require(empty_video_length == 0, + "synthetic Media Foundation video sample must have zero total length"); + require_throws_with( + [&] { trtmc::cli::detail::require_nonempty_decoded_buffer(empty_video_length, "video"); }, + "empty decoded video buffer", + "a non-null sample with a zero-length video buffer must fail closed"); + trtmc::cli::detail::require_nonempty_decoded_buffer(1, "audio"); + + ComPtr reader; + require_success(MFCreateSourceReaderFromURL(path.wstring().c_str(), nullptr, &reader)); + + ComPtr video_type; + require_success( + reader->GetNativeMediaType(MF_SOURCE_READER_FIRST_VIDEO_STREAM, 0, &video_type)); + GUID video_subtype{}; + require_success(video_type->GetGUID(MF_MT_SUBTYPE, &video_subtype)); + require(video_subtype == MFVideoFormat_H264, "MP4 video track is not H.264"); + + ComPtr audio_type; + require_success( + reader->GetNativeMediaType(MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, &audio_type)); + GUID audio_subtype{}; + require_success(audio_type->GetGUID(MF_MT_SUBTYPE, &audio_subtype)); + require(audio_subtype == MFAudioFormat_AAC, "MP4 audio track is not AAC"); + + constexpr int kCompressedBoundarySeconds = 15; + const std::size_t boundary_frames = + static_cast(result.audio.sample_rate) * kCompressedBoundarySeconds; + std::vector pcm(boundary_frames * result.audio.channels); + for (std::size_t frame = 0; frame < boundary_frames; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / result.audio.sample_rate); + const auto quantized = static_cast(std::lround(value * 32767.0F)); + pcm[frame * 2] = quantized; + pcm[frame * 2 + 1] = quantized; + } + constexpr std::uint32_t boundary_44k_rate = 44'100; + const std::size_t boundary_44k_frames = + static_cast(boundary_44k_rate) * kCompressedBoundarySeconds; + std::vector pcm_44k(boundary_44k_frames * result.audio.channels); + for (std::size_t frame = 0; frame < boundary_44k_frames; ++frame) { + const float value = 0.1F * std::sin(2.0F * 3.14159265358979323846F * 440.0F * + static_cast(frame) / boundary_44k_rate); + const auto quantized = static_cast(std::lround(value * 32767.0F)); + pcm_44k[frame * 2] = quantized; + pcm_44k[frame * 2 + 1] = quantized; + } + const auto write_compressed_boundary = [&](const std::filesystem::path& output_path, + const GUID& subtype, + const std::vector& input_pcm, + std::uint32_t sample_rate, LONGLONG duration) { + ComPtr writer; + require_success( + MFCreateSinkWriterFromURL(output_path.wstring().c_str(), nullptr, nullptr, &writer)); + ComPtr output; + require_success(MFCreateMediaType(&output)); + require_success(output->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio)); + require_success(output->SetGUID(MF_MT_SUBTYPE, subtype)); + require_success(output->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, result.audio.channels)); + require_success(output->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate)); + require_success(output->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + subtype == MFAudioFormat_AAC ? 24'000 : 16'000)); + if (subtype == MFAudioFormat_AAC) { + require_success(output->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16)); + require_success(output->SetUINT32(MF_MT_AAC_PAYLOAD_TYPE, 0)); + require_success(output->SetUINT32(MF_MT_AAC_AUDIO_PROFILE_LEVEL_INDICATION, 0x29)); + } + DWORD stream = 0; + require_success(writer->AddStream(output.Get(), &stream)); + ComPtr input; + require_success(MFCreateMediaType(&input)); + require_success(input->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio)); + require_success(input->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_PCM)); + require_success(input->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, result.audio.channels)); + require_success(input->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate)); + require_success(input->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16)); + const auto block_alignment = static_cast(result.audio.channels * 2); + require_success(input->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment)); + require_success( + input->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, sample_rate * block_alignment)); + require_success(writer->SetInputMediaType(stream, input.Get(), nullptr)); + require_success(writer->BeginWriting()); + ComPtr buffer; + require_success(MFCreateMemoryBuffer( + static_cast(input_pcm.size() * sizeof(std::int16_t)), &buffer)); + BYTE* bytes = nullptr; + DWORD capacity = 0; + require_success(buffer->Lock(&bytes, &capacity, nullptr)); + require(capacity >= input_pcm.size() * sizeof(std::int16_t), + "Media Foundation compressed-audio input buffer is undersized"); + std::memcpy(bytes, input_pcm.data(), input_pcm.size() * sizeof(std::int16_t)); + require_success(buffer->Unlock()); + require_success( + buffer->SetCurrentLength(static_cast(input_pcm.size() * sizeof(std::int16_t)))); + ComPtr sample; + require_success(MFCreateSample(&sample)); + require_success(sample->AddBuffer(buffer.Get())); + require_success(sample->SetSampleTime(0)); + require_success(sample->SetSampleDuration(duration)); + require_success(writer->WriteSample(stream, sample.Get())); + require_success(writer->Finalize()); + }; + write_compressed_boundary(mp3_path, MFAudioFormat_MP3, pcm, result.audio.sample_rate, + 150'000'000); + write_compressed_boundary(aac_path, MFAudioFormat_AAC, pcm, result.audio.sample_rate, + 150'000'000); + write_compressed_boundary(mp3_44k_path, MFAudioFormat_MP3, pcm_44k, boundary_44k_rate, + 150'000'000); + write_compressed_boundary(aac_44k_path, MFAudioFormat_AAC, pcm_44k, boundary_44k_rate, + 150'000'000); + auto over_limit_pcm = pcm; + over_limit_pcm.resize(static_cast(result.audio.sample_rate) * + result.audio.channels * 16); + write_compressed_boundary(over_limit_mp3_path, MFAudioFormat_MP3, over_limit_pcm, + result.audio.sample_rate, 160'000'000); + + reader.Reset(); + require_success(MFShutdown()); + if (owns_com) + CoUninitialize(); + + require(std::filesystem::is_regular_file(mp3_path), "MP3 output file is missing"); + require(std::filesystem::file_size(mp3_path) > 1024, "MP3 output file is empty"); + const auto decoded_mp3 = trtmc::cli::read_audio_file(mp3_path.string(), policy); + require(decoded_mp3.sample_rate == result.audio.sample_rate, + "Media Foundation MP3 audio-rate mismatch"); + require(decoded_mp3.channels == result.audio.channels, + "Media Foundation MP3 audio-channel mismatch"); + require(decoded_mp3.num_samples == static_cast(decoded_mp3.samples.size()) && + !decoded_mp3.samples.empty(), + "Media Foundation MP3 decode returned invalid samples"); + require(decoded_mp3.samples.size() <= pcm.size() && + decoded_mp3.samples.size() >= pcm.size() - 4096, + "configured MP3 boundary must trim only codec padding"); + + require(std::filesystem::is_regular_file(aac_path), "AAC output file is missing"); + require(std::filesystem::file_size(aac_path) > 1024, "AAC output file is empty"); + const auto decoded_aac = trtmc::cli::read_audio_file(aac_path.string(), policy); + require(decoded_aac.sample_rate == result.audio.sample_rate, + "Media Foundation AAC audio-rate mismatch"); + require(decoded_aac.channels == result.audio.channels, + "Media Foundation AAC audio-channel mismatch"); + require(decoded_aac.samples.size() <= pcm.size() && + decoded_aac.samples.size() >= pcm.size() - 4096, + "configured AAC boundary must trim only codec padding"); + + const auto decoded_mp3_44k = trtmc::cli::read_audio_file(mp3_44k_path.string(), policy); + require(decoded_mp3_44k.sample_rate == static_cast(boundary_44k_rate), + "44.1 kHz MP3 boundary audio-rate mismatch"); + require(decoded_mp3_44k.samples.size() <= pcm_44k.size() && + decoded_mp3_44k.samples.size() >= pcm_44k.size() - 4096, + "configured 44.1 kHz MP3 boundary must trim only codec padding"); + const auto decoded_aac_44k = trtmc::cli::read_audio_file(aac_44k_path.string(), policy); + require(decoded_aac_44k.sample_rate == static_cast(boundary_44k_rate), + "44.1 kHz AAC boundary audio-rate mismatch"); + require(decoded_aac_44k.samples.size() <= pcm_44k.size() && + decoded_aac_44k.samples.size() >= pcm_44k.size() - 4096, + "configured 44.1 kHz AAC boundary must trim only codec padding"); + require_throws_with( + [&] { (void)trtmc::cli::read_audio_file(over_limit_mp3_path.string(), policy); }, + "configured duration", "a true 16-second MP3 must not be accepted as codec padding"); + + return 0; +} + +int main() { + try { + return run_test(); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/apps/cli/tests/test_windows_utf8_argv.cpp b/apps/cli/tests/test_windows_utf8_argv.cpp new file mode 100644 index 0000000000..3d0c524497 --- /dev/null +++ b/apps/cli/tests/test_windows_utf8_argv.cpp @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_utf8_argv.h" + +#define NOMINMAX +#define WIN32_LEAN_AND_MEAN +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +void test_multilingual_utf16_converts_to_utf8() { + const std::wstring multilingual = + L"Arabic \u0627\u0644\u0639\u0631\u0628\u064a\u0629 | Chinese \u4e2d\u6587 | English | " + L"French fran\u00e7ais | German Stra\u00dfe | Italian italiano | Japanese " + L"\u65e5\u672c\u8a9e | " + L"Korean \ud55c\uad6d\uc5b4 | Portuguese portugu\u00eas | Russian " + L"\u0440\u0443\u0441\u0441\u043a\u0438\u0439 | Spanish espa\u00f1ol"; + const std::string expected = + u8"Arabic \u0627\u0644\u0639\u0631\u0628\u064a\u0629 | Chinese \u4e2d\u6587 | English | " + u8"French fran\u00e7ais | German Stra\u00dfe | Italian italiano | Japanese " + u8"\u65e5\u672c\u8a9e | " + u8"Korean \ud55c\uad6d\uc5b4 | Portuguese portugu\u00eas | Russian " + u8"\u0440\u0443\u0441\u0441\u043a\u0438\u0439 | Spanish espa\u00f1ol"; + check(trtmc::cli::utf8_from_utf16(multilingual) == expected, + "11-language UTF-16 prompt converts exactly to UTF-8"); +} + +void test_command_line_preserves_prompt_and_path() { + wchar_t program[] = L"trtmc.exe"; + wchar_t prompt_flag[] = L"--prompt"; + wchar_t prompt[] = L"\u4e2d\u6587\u63d0\u793a"; + wchar_t output_flag[] = L"--output"; + wchar_t path[] = L"C:\\video\\\u6a21\u578b\\\u8f93\u5165 \u56fe\u7247.mp4"; + wchar_t* wide_argv[] = {program, prompt_flag, prompt, output_flag, path}; + + trtmc::cli::Utf8CommandLine command_line(5, wide_argv); + check(command_line.argc() == 5, "UTF-8 command line preserves argc"); + check(command_line.argv()[5] == nullptr, "UTF-8 command line is null terminated"); + check(std::string(command_line.argv()[2]) == u8"\u4e2d\u6587\u63d0\u793a", + "UTF-8 command line preserves Chinese prompt"); + check(std::string(command_line.argv()[4]) == + u8"C:\\video\\\u6a21\u578b\\\u8f93\u5165 \u56fe\u7247.mp4", + "UTF-8 command line preserves non-ASCII media path"); +} + +void test_invalid_utf16_is_rejected() { + const wchar_t invalid[] = {static_cast(0xD800), L'\0'}; + bool rejected = false; + try { + (void)trtmc::cli::utf8_from_utf16(std::wstring_view(invalid, 1)); + } catch (const std::runtime_error&) { + rejected = true; + } + check(rejected, "malformed UTF-16 command-line text is rejected"); +} + +void test_utf8_process_manifest_and_narrow_filesystem_path() { + HRSRC resource = FindResourceW(nullptr, MAKEINTRESOURCEW(1), MAKEINTRESOURCEW(24)); + check(resource != nullptr, "test executable contains an embedded process manifest"); + if (resource != nullptr) { + HGLOBAL loaded = LoadResource(nullptr, resource); + const DWORD size = SizeofResource(nullptr, resource); + const void* data = loaded == nullptr ? nullptr : LockResource(loaded); + const std::string manifest = + data == nullptr ? std::string{} : std::string(static_cast(data), size); + check(manifest.find("activeCodePage") != std::string::npos && + manifest.find(">UTF-8") != std::string::npos, + "PE manifest declares UTF-8 as the active process code page"); + } + check(GetACP() == CP_UTF8, "embedded manifest activates the UTF-8 process code page"); + + const auto wide_directory = + std::filesystem::temp_directory_path() / + (std::wstring(L"trtmc_utf8_path_\u6a21\u578b_") + std::to_wstring(GetCurrentProcessId())); + const auto wide_file = wide_directory / L"\u8f93\u5165 \u56fe\u7247.txt"; + std::error_code error; + (void)std::filesystem::remove(wide_file, error); + error.clear(); + (void)std::filesystem::remove(wide_directory, error); + error.clear(); + check(std::filesystem::create_directory(wide_directory, error) && !error, + "wide setup creates a non-ASCII test directory"); + + const std::string utf8_directory = trtmc::cli::utf8_from_utf16(wide_directory.native()); + const std::filesystem::path narrow_file = + std::filesystem::path(utf8_directory) / u8"\u8f93\u5165 \u56fe\u7247.txt"; + { + std::ofstream output(narrow_file, std::ios::binary | std::ios::trunc); + output << "utf8-path"; + check(static_cast(output), + "UTF-8 narrow std::filesystem path opens a non-ASCII output file"); + } + check(std::filesystem::exists(wide_file), + "UTF-8 narrow path resolves to the intended UTF-16 filesystem path"); + { + std::ifstream input(narrow_file, std::ios::binary); + std::string value; + input >> value; + check(value == "utf8-path", "UTF-8 narrow path reads the non-ASCII file back"); + } + + error.clear(); + check(std::filesystem::remove(wide_file, error) && !error, + "non-ASCII path test removes its file"); + error.clear(); + check(std::filesystem::remove(wide_directory, error) && !error, + "non-ASCII path test removes its directory"); +} + +} // namespace + +int main() { + test_multilingual_utf16_converts_to_utf8(); + test_command_line_preserves_prompt_and_path(); + test_invalid_utf16_is_rejected(); + test_utf8_process_manifest_and_narrow_filesystem_path(); + if (failures == 0) + std::cout << "All Windows UTF-8 argv tests passed\n"; + return failures == 0 ? 0 : 1; +} diff --git a/apps/cli/trtmc_windows_utf8.manifest b/apps/cli/trtmc_windows_utf8.manifest new file mode 100644 index 0000000000..afccd0d9ca --- /dev/null +++ b/apps/cli/trtmc_windows_utf8.manifest @@ -0,0 +1,13 @@ + + + + + + UTF-8 + true + + + diff --git a/apps/cli/windows_media.cpp b/apps/cli/windows_media.cpp new file mode 100644 index 0000000000..1166e8f040 --- /dev/null +++ b/apps/cli/windows_media.cpp @@ -0,0 +1,1434 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_media.h" + +#include "cli/io.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#define NOMINMAX +#include +#include +#include +#include +#include +#include +#include +#endif + +namespace trtmc::cli { +namespace { + +std::string lowercase_extension(std::string_view path) { + const auto extension = std::filesystem::path(std::string(path)).extension().string(); + std::string lowered(extension); + std::transform(lowered.begin(), lowered.end(), lowered.begin(), + [](unsigned char value) { return static_cast(std::tolower(value)); }); + return lowered; +} + +constexpr std::uint64_t kAudioMediaTicksPerSecond = 10'000'000; + +bool reference_limit_ticks(const ReferenceMediaDecodePolicy& policy, + std::uint64_t& ticks) noexcept { + if (policy.maximum_duration_seconds == 0 || + policy.maximum_duration_seconds > + std::numeric_limits::max() / kAudioMediaTicksPerSecond) { + return false; + } + ticks = static_cast(policy.maximum_duration_seconds) * kAudioMediaTicksPerSecond; + return true; +} + +void validate_reference_media_decode_policy(const ReferenceMediaDecodePolicy& policy) { + std::uint64_t ignored_ticks = 0; + if (!reference_limit_ticks(policy, ignored_ticks) || policy.target_video_fps == 0 || + policy.maximum_source_video_fps == 0 || policy.canvas_short_edge == 0 || + policy.canvas_max_pixels == 0 || policy.canvas_multiple == 0 || + !std::isfinite(policy.minimum_aspect_ratio) || + !std::isfinite(policy.maximum_aspect_ratio) || policy.minimum_aspect_ratio <= 0.0 || + policy.maximum_aspect_ratio < policy.minimum_aspect_ratio) { + throw std::invalid_argument("reference-media decode policy values must be positive"); + } +} + +struct ReferenceVideoSize { + std::uint32_t width{0}; + std::uint32_t height{0}; +}; + +std::uint32_t round_canvas_axis(double value, const ReferenceMediaDecodePolicy& policy) { + const double scaled = value / policy.canvas_multiple; + const double lower = std::floor(scaled); + const double fraction = scaled - lower; + double rounded = lower; + if (fraction > 0.5 || (fraction == 0.5 && static_cast(lower) % 2 != 0)) + rounded += 1.0; + if (rounded > std::numeric_limits::max() / policy.canvas_multiple) + throw std::runtime_error("reference video canvas rounding overflow"); + return std::max(policy.canvas_multiple, + static_cast(rounded) * policy.canvas_multiple); +} + +ReferenceVideoSize resolve_reference_video_size(std::uint32_t source_width, + std::uint32_t source_height, + const ReferenceMediaDecodePolicy& policy) { + if (source_width == 0 || source_height == 0) + throw std::runtime_error("reference video has invalid source dimensions"); + const double ratio = static_cast(source_width) / source_height; + if (!std::isfinite(ratio) || ratio < policy.minimum_aspect_ratio || + ratio > policy.maximum_aspect_ratio) { + throw std::runtime_error("reference video aspect is outside the configured range"); + } + double width = ratio >= 1.0 ? policy.canvas_short_edge * ratio + : static_cast(policy.canvas_short_edge); + double height = ratio >= 1.0 ? static_cast(policy.canvas_short_edge) + : policy.canvas_short_edge / ratio; + const double area = width * height; + if (area > static_cast(policy.canvas_max_pixels)) { + const double scale = std::sqrt(static_cast(policy.canvas_max_pixels) / area); + width *= scale; + height *= scale; + } + return {round_canvas_axis(width, policy), round_canvas_axis(height, policy)}; +} + +struct ReferenceAudioFrameWindow { + std::uint64_t first_frame{0}; + std::uint64_t end_frame{0}; +}; + +std::uint64_t ticks_to_frames(std::uint64_t ticks, std::uint32_t sample_rate, + bool round_up) noexcept { + const std::uint64_t whole_seconds = ticks / kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = ticks % kAudioMediaTicksPerSecond; + if (whole_seconds > std::numeric_limits::max() / sample_rate) + return std::numeric_limits::max(); + std::uint64_t frames = whole_seconds * sample_rate; + const std::uint64_t partial_product = partial_ticks * sample_rate; + const std::uint64_t partial_frames = + round_up ? (partial_product + kAudioMediaTicksPerSecond - 1) / kAudioMediaTicksPerSecond + : partial_product / kAudioMediaTicksPerSecond; + if (partial_frames > std::numeric_limits::max() - frames) + return std::numeric_limits::max(); + return frames + partial_frames; +} + +std::uint64_t ticks_to_frames_nearest(std::uint64_t ticks, std::uint32_t sample_rate) noexcept { + const std::uint64_t whole_seconds = ticks / kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = ticks % kAudioMediaTicksPerSecond; + if (whole_seconds > std::numeric_limits::max() / sample_rate) + return std::numeric_limits::max(); + std::uint64_t frames = whole_seconds * sample_rate; + const std::uint64_t partial_product = partial_ticks * sample_rate; + const std::uint64_t partial_frames = + (partial_product + kAudioMediaTicksPerSecond / 2) / kAudioMediaTicksPerSecond; + if (partial_frames > std::numeric_limits::max() - frames) + return std::numeric_limits::max(); + return frames + partial_frames; +} + +bool frames_to_ticks_rounded_up(std::uint64_t frames, std::uint32_t sample_rate, + std::uint64_t& ticks) noexcept { + const std::uint64_t whole_seconds = frames / sample_rate; + const std::uint64_t partial_frames = frames % sample_rate; + if (whole_seconds > std::numeric_limits::max() / kAudioMediaTicksPerSecond) { + return false; + } + ticks = whole_seconds * kAudioMediaTicksPerSecond; + const std::uint64_t partial_product = partial_frames * kAudioMediaTicksPerSecond; + const std::uint64_t partial_ticks = (partial_product + sample_rate - 1) / sample_rate; + if (partial_ticks > std::numeric_limits::max() - ticks) + return false; + ticks += partial_ticks; + return true; +} + +bool account_reference_audio_decode_impl(std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, + std::uint64_t maximum_padding_frames, + detail::ReferenceAudioDecodeState& state, + ReferenceAudioFrameWindow* retained_window, + const ReferenceMediaDecodePolicy& policy) noexcept { + if (sample_rate == 0 || frame_count == 0) + return false; + + const std::uint64_t public_decoded_frames = + static_cast(policy.maximum_duration_seconds) * sample_rate; + if (maximum_padding_frames > + std::numeric_limits::max() - public_decoded_frames) { + return false; + } + const std::uint64_t maximum_decoded_frames = public_decoded_frames + maximum_padding_frames; + if (state.decoded_frames > maximum_decoded_frames || + frame_count > maximum_decoded_frames - state.decoded_frames) { + return false; + } + + std::uint64_t padding_ticks = 0; + std::uint64_t sample_duration_ticks = 0; + if (!frames_to_ticks_rounded_up(maximum_padding_frames, sample_rate, padding_ticks) || + !frames_to_ticks_rounded_up(frame_count, sample_rate, sample_duration_ticks)) { + return false; + } + std::uint64_t reference_limit = 0; + if (!reference_limit_ticks(policy, reference_limit) || + padding_ticks > std::numeric_limits::max() - reference_limit) { + return false; + } + const std::uint64_t latest_end_tick = reference_limit + padding_ticks; + + std::uint64_t leading_ticks = 0; + if (timestamp < 0) { + leading_ticks = static_cast(-(timestamp + 1)) + 1; + if (leading_ticks > padding_ticks) + return false; + if (sample_duration_ticks > leading_ticks && + sample_duration_ticks - leading_ticks > latest_end_tick) { + return false; + } + } else { + const std::uint64_t start_tick = static_cast(timestamp); + if (start_tick > latest_end_tick || sample_duration_ticks > latest_end_tick - start_tick) { + return false; + } + } + + const std::uint64_t first_frame = + std::min(frame_count, ticks_to_frames(leading_ticks, sample_rate, true)); + std::uint64_t end_frame = first_frame; + if (timestamp < static_cast(reference_limit)) { + const std::uint64_t remaining_ticks = + timestamp < 0 ? reference_limit + leading_ticks + : reference_limit - static_cast(timestamp); + end_frame = std::min(frame_count, ticks_to_frames(remaining_ticks, sample_rate, false)); + end_frame = std::max(end_frame, first_frame); + } + const std::uint64_t padding_frames = frame_count - (end_frame - first_frame); + if (state.decoded_padding_frames > maximum_padding_frames || + padding_frames > maximum_padding_frames - state.decoded_padding_frames) { + return false; + } + + state.decoded_frames += frame_count; + state.decoded_padding_frames += padding_frames; + if (retained_window != nullptr) + *retained_window = {first_frame, end_frame}; + return true; +} + +#if defined(_WIN32) + +using Microsoft::WRL::ComPtr; + +constexpr LONGLONG kMediaTicksPerSecond = 10'000'000; +constexpr std::size_t kMaxConsecutiveEmptyReads = 256; +constexpr std::size_t kMaxTotalEmptyReads = 16'384; + +[[noreturn]] void throw_hresult(const char* operation, HRESULT result) { + std::ostringstream message; + message << operation << " failed with HRESULT 0x" << std::hex << std::setw(8) + << std::setfill('0') << static_cast(result); + throw std::runtime_error(message.str()); +} + +void check_hresult(HRESULT result, const char* operation) { + if (FAILED(result)) + throw_hresult(operation, result); +} + +class MediaFoundationSession { + public: + MediaFoundationSession() { + const HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + if (SUCCEEDED(com_result)) + owns_com_ = true; + else if (com_result != RPC_E_CHANGED_MODE) + throw_hresult("CoInitializeEx", com_result); + const HRESULT media_result = MFStartup(MF_VERSION, MFSTARTUP_FULL); + if (FAILED(media_result)) { + if (owns_com_) + CoUninitialize(); + owns_com_ = false; + throw_hresult("MFStartup", media_result); + } + media_started_ = true; + } + + ~MediaFoundationSession() { + if (media_started_) + (void)MFShutdown(); + if (owns_com_) + CoUninitialize(); + } + + MediaFoundationSession(const MediaFoundationSession&) = delete; + MediaFoundationSession& operator=(const MediaFoundationSession&) = delete; + + private: + bool owns_com_{false}; + bool media_started_{false}; +}; + +std::uint32_t checked_u32(std::uint64_t value, const char* label) { + if (value > std::numeric_limits::max()) + throw std::runtime_error(std::string(label) + " exceeds the Media Foundation limit"); + return static_cast(value); +} + +LONGLONG media_time(std::uint64_t numerator, std::uint32_t denominator) { + constexpr std::uint64_t kTicksPerSecond = 10'000'000; + if (denominator == 0 || numerator > std::numeric_limits::max() / kTicksPerSecond) + throw std::runtime_error("media timestamp overflow"); + return static_cast((numerator * kTicksPerSecond) / denominator); +} + +BYTE quantize(float value) { + return static_cast(std::clamp(static_cast(std::lround(value)), 0, 255)); +} + +void rgb_to_nv12(const float* rgb, std::uint32_t width, std::uint32_t height, + std::vector& nv12) { + const std::size_t y_size = static_cast(width) * height; + nv12.resize(y_size + y_size / 2); + auto* y_plane = nv12.data(); + auto* uv_plane = nv12.data() + y_size; + + const auto component = [](const float* pixel, int index) { + return std::clamp(pixel[index], 0.0F, 1.0F); + }; + for (std::uint32_t row = 0; row < height; ++row) { + for (std::uint32_t column = 0; column < width; ++column) { + const float* pixel = rgb + (static_cast(row) * width + column) * 3; + const float red = component(pixel, 0); + const float green = component(pixel, 1); + const float blue = component(pixel, 2); + y_plane[static_cast(row) * width + column] = + quantize(16.0F + 219.0F * (0.2126F * red + 0.7152F * green + 0.0722F * blue)); + } + } + + for (std::uint32_t row = 0; row < height; row += 2) { + for (std::uint32_t column = 0; column < width; column += 2) { + float cb = 0.0F; + float cr = 0.0F; + for (std::uint32_t dy = 0; dy < 2; ++dy) { + for (std::uint32_t dx = 0; dx < 2; ++dx) { + const float* pixel = + rgb + (static_cast(row + dy) * width + column + dx) * 3; + const float red = component(pixel, 0); + const float green = component(pixel, 1); + const float blue = component(pixel, 2); + cb += -0.114572F * red - 0.385428F * green + 0.5F * blue; + cr += 0.5F * red - 0.454153F * green - 0.045847F * blue; + } + } + const auto uv_offset = static_cast(row / 2) * width + column; + uv_plane[uv_offset] = quantize(128.0F + 224.0F * cb / 4.0F); + uv_plane[uv_offset + 1] = quantize(128.0F + 224.0F * cr / 4.0F); + } + } +} + +ComPtr make_sample(const void* bytes, std::size_t byte_count, LONGLONG timestamp, + LONGLONG duration) { + const auto media_bytes = checked_u32(byte_count, "media sample size"); + ComPtr buffer; + check_hresult(MFCreateMemoryBuffer(media_bytes, &buffer), "MFCreateMemoryBuffer"); + BYTE* destination = nullptr; + DWORD maximum_length = 0; + check_hresult(buffer->Lock(&destination, &maximum_length, nullptr), "IMFMediaBuffer::Lock"); + if (maximum_length < media_bytes) { + (void)buffer->Unlock(); + throw std::runtime_error("Media Foundation returned an undersized sample buffer"); + } + std::memcpy(destination, bytes, byte_count); + check_hresult(buffer->Unlock(), "IMFMediaBuffer::Unlock"); + check_hresult(buffer->SetCurrentLength(media_bytes), "IMFMediaBuffer::SetCurrentLength"); + + ComPtr sample; + check_hresult(MFCreateSample(&sample), "MFCreateSample"); + check_hresult(sample->AddBuffer(buffer.Get()), "IMFSample::AddBuffer"); + check_hresult(sample->SetSampleTime(timestamp), "IMFSample::SetSampleTime"); + check_hresult(sample->SetSampleDuration(duration), "IMFSample::SetSampleDuration"); + return sample; +} + +ComPtr make_video_output_type(std::uint32_t width, std::uint32_t height, + std::uint32_t fps) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(video output)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "video output major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_H264), "video output subtype"); + const auto bitrate = checked_u32( + std::max(8'000'000, static_cast(width) * height * fps / 3), + "H.264 bitrate"); + check_hresult(type->SetUINT32(MF_MT_AVG_BITRATE, bitrate), "video output bitrate"); + check_hresult(type->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive), + "video output interlace mode"); + check_hresult(type->SetUINT32(MF_MT_MPEG2_PROFILE, eAVEncH264VProfile_Main), + "video output H.264 profile"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "video output frame size"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_FRAME_RATE, fps, 1), + "video output frame rate"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_PIXEL_ASPECT_RATIO, 1, 1), + "video output pixel aspect ratio"); + return type; +} + +ComPtr make_video_input_type(std::uint32_t width, std::uint32_t height, + std::uint32_t fps) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(video input)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "video input major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_NV12), "video input subtype"); + check_hresult(type->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive), + "video input interlace mode"); + check_hresult(type->SetUINT32(MF_MT_FIXED_SIZE_SAMPLES, TRUE), "video input fixed samples"); + check_hresult(type->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE), + "video input independent samples"); + check_hresult(type->SetUINT32(MF_MT_SAMPLE_SIZE, + checked_u32(static_cast(width) * height * 3 / 2, + "NV12 sample size")), + "video input sample size"); + check_hresult(type->SetUINT32(MF_MT_DEFAULT_STRIDE, width), "video input stride"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "video input frame size"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_FRAME_RATE, fps, 1), + "video input frame rate"); + check_hresult(MFSetAttributeRatio(type.Get(), MF_MT_PIXEL_ASPECT_RATIO, 1, 1), + "video input pixel aspect ratio"); + return type; +} + +ComPtr make_audio_output_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(audio output)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "audio output major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_AAC), "audio output subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "audio output channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "audio output sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16), "audio output bit depth"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, 24'000), + "audio output bitrate"); + check_hresult(type->SetUINT32(MF_MT_AAC_PAYLOAD_TYPE, 0), "AAC payload type"); + check_hresult(type->SetUINT32(MF_MT_AAC_AUDIO_PROFILE_LEVEL_INDICATION, 0x29), + "AAC profile level"); + return type; +} + +ComPtr make_audio_input_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(audio input)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "audio input major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_PCM), "audio input subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "audio input channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "audio input sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 16), "audio input bit depth"); + const auto block_alignment = + checked_u32(static_cast(channels) * 2, "audio block alignment"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment), + "audio input block alignment"); + check_hresult( + type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + checked_u32(static_cast(sample_rate) * block_alignment, + "PCM byte rate")), + "audio input byte rate"); + check_hresult(type->SetUINT32(MF_MT_ALL_SAMPLES_INDEPENDENT, TRUE), + "audio input independent samples"); + return type; +} + +void validate_result(const VideoResult& result) { + const auto& frames = result.frames; + if (frames.width <= 0 || frames.height <= 0 || frames.num_frames <= 0 || frames.channels != 3 || + result.fps <= 0) + throw std::runtime_error( + "write_mp4 requires valid THWC RGB video and a positive frame rate"); + if ((frames.width & 1) != 0 || (frames.height & 1) != 0) + throw std::runtime_error("write_mp4 requires even frame dimensions for NV12/H.264"); + const auto pixels_per_frame = static_cast(frames.width) * frames.height * 3; + const auto required_pixels = pixels_per_frame * static_cast(frames.num_frames); + if (required_pixels > frames.pixels.size()) + throw std::runtime_error("write_mp4 frame storage is smaller than its THWC metadata"); + const auto& audio = result.audio; + if (!audio.samples.empty()) { + if (audio.sample_rate <= 0 || (audio.channels != 1 && audio.channels != 2) || + audio.samples.size() % static_cast(audio.channels) != 0) + throw std::runtime_error("write_mp4 requires valid mono or stereo interleaved audio"); + } +} + +std::uint64_t rounded_reference_frame_slot(std::uint64_t frame, std::uint32_t fps_numerator, + std::uint32_t fps_denominator, + const ReferenceMediaDecodePolicy& policy) { + if (fps_numerator == 0 || fps_denominator == 0 || + frame > + std::numeric_limits::max() / policy.target_video_fps / fps_denominator) { + throw std::runtime_error("reference video frame-slot arithmetic overflow"); + } + const std::uint64_t numerator = frame * policy.target_video_fps * fps_denominator; + if (numerator > (std::numeric_limits::max() - fps_numerator) / 2) + throw std::runtime_error("reference video frame-slot rounding overflow"); + return (2 * numerator + fps_numerator) / (2ULL * fps_numerator); +} + +ComPtr source_reader_video_type(std::uint32_t width, std::uint32_t height) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(source video)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video), "source video major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_RGB32), "source video subtype"); + check_hresult(MFSetAttributeSize(type.Get(), MF_MT_FRAME_SIZE, width, height), + "source video decoded frame size"); + return type; +} + +ComPtr source_reader_audio_type(std::uint32_t sample_rate, std::uint32_t channels) { + ComPtr type; + check_hresult(MFCreateMediaType(&type), "MFCreateMediaType(source audio)"); + check_hresult(type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Audio), "source audio major type"); + check_hresult(type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_Float), "source audio subtype"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_NUM_CHANNELS, channels), "source audio channels"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, sample_rate), + "source audio sample rate"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, 32), "source audio bit depth"); + const auto block_alignment = + checked_u32(static_cast(channels) * 4, "source audio block alignment"); + check_hresult(type->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_alignment), + "source audio block alignment"); + check_hresult( + type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, + checked_u32(static_cast(sample_rate) * block_alignment, + "source audio byte rate")), + "source audio byte rate"); + return type; +} + +void append_rgb32_frame(IMFSample* sample, std::uint32_t width, std::uint32_t height, LONG stride, + std::vector& pixels, std::size_t maximum_scalars) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null video sample"); + ComPtr buffer; + check_hresult(sample->ConvertToContiguousBuffer(&buffer), + "IMFSample::ConvertToContiguousBuffer(video)"); + BYTE* data = nullptr; + DWORD maximum_length = 0; + DWORD current_length = 0; + check_hresult(buffer->Lock(&data, &maximum_length, ¤t_length), + "IMFMediaBuffer::Lock(video)"); + const auto unlock = [&]() { (void)buffer->Unlock(); }; + if (current_length == 0) { + unlock(); + detail::require_nonempty_decoded_buffer(current_length, "video"); + } + const auto absolute_stride = + static_cast(stride < 0 ? -static_cast(stride) : stride); + const auto required_bytes = absolute_stride * height; + if (absolute_stride < static_cast(width) * 4 || + required_bytes > current_length) { + unlock(); + throw std::runtime_error("Media Foundation returned an invalid RGB32 stride or buffer"); + } + const auto old_size = pixels.size(); + const auto frame_scalars = static_cast(width) * height * 3; + if (frame_scalars > maximum_scalars || old_size > maximum_scalars - frame_scalars) { + unlock(); + throw std::runtime_error("decoded reference video exceeds its bounded RGB allocation"); + } + if (frame_scalars > std::numeric_limits::max() - old_size) { + unlock(); + throw std::runtime_error("decoded video exceeds host address space"); + } + pixels.resize(old_size + static_cast(frame_scalars)); + float* destination = pixels.data() + old_size; + for (std::uint32_t row = 0; row < height; ++row) { + const auto source_row = stride >= 0 ? row : height - 1 - row; + const BYTE* source = data + static_cast(source_row) * absolute_stride; + for (std::uint32_t column = 0; column < width; ++column) { + const BYTE* bgra = source + static_cast(column) * 4; + float* rgb = destination + (static_cast(row) * width + column) * 3; + rgb[0] = static_cast(bgra[2]) / 255.0F; + rgb[1] = static_cast(bgra[1]) / 255.0F; + rgb[2] = static_cast(bgra[0]) / 255.0F; + } + } + unlock(); +} + +void append_float_audio(IMFSample* sample, std::uint32_t sample_rate, std::uint32_t channels, + LONGLONG timestamp, std::uint64_t maximum_codec_padding_frames, + detail::ReferenceAudioDecodeState& decode_state, + std::vector& samples, std::size_t maximum_scalars, + const ReferenceMediaDecodePolicy& policy) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null audio sample"); + ComPtr buffer; + check_hresult(sample->ConvertToContiguousBuffer(&buffer), + "IMFSample::ConvertToContiguousBuffer(audio)"); + BYTE* data = nullptr; + DWORD maximum_length = 0; + DWORD current_length = 0; + check_hresult(buffer->Lock(&data, &maximum_length, ¤t_length), + "IMFMediaBuffer::Lock(audio)"); + if (current_length == 0) { + (void)buffer->Unlock(); + detail::require_nonempty_decoded_buffer(current_length, "audio"); + } + if (current_length % (sizeof(float) * channels) != 0) { + (void)buffer->Unlock(); + throw std::runtime_error("Media Foundation returned misaligned float audio"); + } + const auto scalar_count = current_length / sizeof(float); + const std::size_t frame_count = scalar_count / channels; + const bool trim_codec_padding = maximum_codec_padding_frames > 0; + ReferenceAudioFrameWindow retained_window{0, frame_count}; + if (trim_codec_padding) { + if (!account_reference_audio_decode_impl(timestamp, frame_count, sample_rate, + maximum_codec_padding_frames, decode_state, + &retained_window, policy)) { + (void)buffer->Unlock(); + throw std::runtime_error( + "decoded reference audio exceeds the configured duration plus codec padding"); + } + } + if (!trim_codec_padding) { + LONGLONG declared_sample_duration = 0; + if (FAILED(sample->GetSampleDuration(&declared_sample_duration)) || + declared_sample_duration < 0) { + declared_sample_duration = 0; + } + const auto inferred_duration = static_cast( + (static_cast(frame_count) * kMediaTicksPerSecond + sample_rate - 1) / + sample_rate); + const LONGLONG timeline_duration = std::max(declared_sample_duration, inferred_duration); + if (!detail::reference_timeline_within_limit(policy, timestamp, timeline_duration)) { + (void)buffer->Unlock(); + throw std::runtime_error( + "decoded reference audio sample crosses the configured reference timeline"); + } + } + try { + detail::append_reference_audio_frames_on_timeline( + samples, data, frame_count, retained_window.first_frame, retained_window.end_frame, + channels, timestamp, sample_rate, maximum_scalars); + } catch (...) { + (void)buffer->Unlock(); + throw; + } + check_hresult(buffer->Unlock(), "IMFMediaBuffer::Unlock(audio)"); +} + +std::optional presentation_duration(IMFSourceReader& reader) { + PROPVARIANT value; + PropVariantInit(&value); + const HRESULT status = + reader.GetPresentationAttribute(MF_SOURCE_READER_MEDIASOURCE, MF_PD_DURATION, &value); + std::optional result; + if (SUCCEEDED(status)) { + if (value.vt == VT_UI8 && + value.uhVal.QuadPart <= static_cast(std::numeric_limits::max())) { + result = static_cast(value.uhVal.QuadPart); + } else if (value.vt == VT_I8 && value.hVal.QuadPart >= 0) { + result = value.hVal.QuadPart; + } + } + (void)PropVariantClear(&value); + return result; +} + +struct CodecPaddingAllowance { + std::uint64_t frames{0}; + LONGLONG duration{0}; +}; + +CodecPaddingAllowance codec_padding_allowance(IMFMediaType* native_audio) { + if (native_audio == nullptr) + return {}; + GUID subtype{}; + UINT32 sample_rate = 0; + if (FAILED(native_audio->GetGUID(MF_MT_SUBTYPE, &subtype)) || + FAILED(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &sample_rate)) || + sample_rate == 0) { + return {}; + } + std::uint64_t padding_frames = 0; + if (subtype == MFAudioFormat_MP3) + padding_frames = 3ULL * 1152; + else if (subtype == MFAudioFormat_AAC) + padding_frames = 3ULL * 1024; + if (padding_frames == 0) + return {}; + return {padding_frames, + static_cast((padding_frames * kMediaTicksPerSecond + sample_rate - 1) / + sample_rate)}; +} + +std::size_t maximum_reference_audio_scalars(std::uint32_t sample_rate, std::uint32_t channels, + const ReferenceMediaDecodePolicy& policy) { + const std::uint64_t public_limit = + static_cast(sample_rate) * channels * policy.maximum_duration_seconds; + return static_cast(std::min( + public_limit, static_cast(std::numeric_limits::max()))); +} + +void reject_source_reader_error(DWORD flags, const char* label) { + if ((flags & MF_SOURCE_READERF_ERROR) != 0) + throw std::runtime_error(std::string(label) + " reported MF_SOURCE_READERF_ERROR"); +} + +void reject_empty_decoded_sample(IMFSample* sample, std::string_view media_kind) { + if (sample == nullptr) + throw std::runtime_error("Media Foundation returned a null decoded sample"); + DWORD total_length = 0; + check_hresult(sample->GetTotalLength(&total_length), "IMFSample::GetTotalLength(decoded)"); + detail::require_nonempty_decoded_buffer(total_length, media_kind); +} + +void reject_reference_timestamp(LONGLONG timestamp, const char* label, + const ReferenceMediaDecodePolicy& policy) { + const auto limit = + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond; + if (timestamp >= static_cast(limit)) + throw std::runtime_error(std::string(label) + + " timestamp exceeds the configured reference duration"); +} + +void reject_video_sample_timeline(IMFSample* sample, LONGLONG timestamp, + std::uint32_t fps_numerator, std::uint32_t fps_denominator, + const ReferenceMediaDecodePolicy& policy) { + LONGLONG duration = 0; + if (sample == nullptr || FAILED(sample->GetSampleDuration(&duration)) || duration <= 0) { + duration = static_cast( + (static_cast(kMediaTicksPerSecond) * fps_denominator + fps_numerator - + 1) / + fps_numerator); + } + if (!detail::reference_timeline_within_limit(policy, timestamp, duration)) + throw std::runtime_error( + "decoded reference video sample crosses the configured reference timeline"); +} + +void note_empty_source_reader_event(DWORD flags, LONGLONG timestamp, LONGLONG& last_tick_timestamp, + std::size_t& consecutive_empty_reads, + std::size_t& total_empty_reads, + std::uint64_t maximum_codec_padding_ticks, const char* label, + const ReferenceMediaDecodePolicy& policy) { + // A stream tick deliberately carries no sample. Treat an advancing tick as + // progress, but bound its timestamp to the public reference duration. All + // other empty reads (including repeated/non-advancing ticks) are bounded so + // a malformed source cannot spin forever. + if (++total_empty_reads > kMaxTotalEmptyReads) + throw std::runtime_error(std::string(label) + " returned too many empty stream events"); + if (maximum_codec_padding_ticks > 0 && !detail::reference_audio_event_timestamp_within_padding( + policy, timestamp, maximum_codec_padding_ticks)) { + throw std::runtime_error(std::string(label) + + " event timestamp exceeds the configured duration plus codec " + "padding"); + } + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (maximum_codec_padding_ticks == 0) + reject_reference_timestamp(timestamp, label, policy); + if (timestamp > last_tick_timestamp) { + last_tick_timestamp = timestamp; + consecutive_empty_reads = 0; + return; + } + } + if (++consecutive_empty_reads > kMaxConsecutiveEmptyReads) + throw std::runtime_error(std::string(label) + " made no progress while decoding"); +} + +#endif + +} // namespace + +void detail::append_reference_audio_frames_on_timeline( + std::vector& destination, const void* source, std::uint64_t source_frame_count, + std::uint64_t first_source_frame, std::uint64_t end_source_frame, std::uint32_t channels, + std::int64_t timestamp, std::uint32_t sample_rate, std::size_t maximum_scalars) { + if (source == nullptr || channels == 0 || sample_rate == 0 || + first_source_frame > end_source_frame || end_source_frame > source_frame_count || + destination.size() % channels != 0) { + throw std::runtime_error("decoded reference audio has an invalid timeline layout"); + } + + const std::uint64_t retained_frames = end_source_frame - first_source_frame; + if (retained_frames == 0) + return; + + // Media Foundation timestamps are expressed in 100 ns units. Map each sample + // independently to the nearest PCM frame so timestamp rounding does not + // accumulate across access units. Negative codec-priming samples have already + // been trimmed to the public timeline and therefore begin at frame zero. + const std::uint64_t target_frame = + timestamp <= 0 + ? 0 + : ticks_to_frames_nearest(static_cast(timestamp), sample_rate); + if (target_frame == std::numeric_limits::max()) + throw std::runtime_error("decoded reference audio timestamp exceeds host range"); + + const std::uint64_t destination_frames = destination.size() / channels; + const std::uint64_t gap_frames = + target_frame > destination_frames ? target_frame - destination_frames : 0; + const std::uint64_t overlap_frames = + target_frame < destination_frames + ? std::min(destination_frames - target_frame, retained_frames) + : 0; + const std::uint64_t append_frames = retained_frames - overlap_frames; + if (gap_frames > std::numeric_limits::max() - append_frames) + throw std::runtime_error("decoded reference audio timeline exceeds host range"); + const std::uint64_t added_frames = gap_frames + append_frames; + if (added_frames > std::numeric_limits::max() / channels) + throw std::runtime_error("decoded audio exceeds host address space"); + const std::size_t added_scalars = static_cast(added_frames) * channels; + const std::size_t old_size = destination.size(); + if (added_scalars > maximum_scalars || old_size > maximum_scalars - added_scalars) + throw std::runtime_error("decoded reference audio exceeds the configured duration limit"); + if (added_scalars > std::numeric_limits::max() - old_size) + throw std::runtime_error("decoded audio exceeds host address space"); + + std::size_t source_byte_offset = 0; + std::size_t append_bytes = 0; + if (append_frames != 0) { + const std::uint64_t source_frame = first_source_frame + overlap_frames; + if (source_frame > std::numeric_limits::max() / channels || + append_frames > std::numeric_limits::max() / channels) { + throw std::runtime_error("decoded audio exceeds host address space"); + } + const std::size_t source_scalar = static_cast(source_frame) * channels; + const std::size_t append_scalars = static_cast(append_frames) * channels; + if (source_scalar > std::numeric_limits::max() / sizeof(float) || + append_scalars > std::numeric_limits::max() / sizeof(float)) { + throw std::runtime_error("decoded audio exceeds host address space"); + } + source_byte_offset = source_scalar * sizeof(float); + append_bytes = append_scalars * sizeof(float); + } + + destination.resize(old_size + added_scalars, 0.0F); + if (append_frames != 0) { + const auto* source_bytes = static_cast(source); + std::memcpy(destination.data() + old_size + static_cast(gap_frames) * channels, + source_bytes + source_byte_offset, append_bytes); + } +} + +void detail::validate_reference_video_soundtrack_format(std::uint32_t sample_rate, + std::uint32_t channels) { + if (sample_rate == 0 || + sample_rate > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("reference video soundtrack has an invalid sample rate"); + } + if (channels != 1 && channels != 2) { + throw std::runtime_error("reference video soundtrack must contain mono or stereo audio"); + } +} + +std::uint64_t detail::reference_video_frame_ceiling(const ReferenceMediaDecodePolicy& policy, + std::uint32_t fps_numerator, + std::uint32_t fps_denominator) { + if (policy.maximum_duration_seconds == 0 || fps_numerator == 0 || fps_denominator == 0) + throw std::invalid_argument("reference video frame rate must be positive"); + const std::uint64_t numerator = + static_cast(policy.maximum_duration_seconds) * fps_numerator; + if (numerator > std::numeric_limits::max() - (fps_denominator - 1ULL)) + throw std::invalid_argument("reference video frame ceiling overflow"); + return (numerator + fps_denominator - 1) / fps_denominator; +} + +std::pair +detail::reference_video_decode_size(const ReferenceMediaDecodePolicy& policy, + std::uint32_t source_width, std::uint32_t source_height) { + validate_reference_media_decode_policy(policy); + const auto size = resolve_reference_video_size(source_width, source_height, policy); + return {size.width, size.height}; +} + +bool detail::reference_timeline_within_limit(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, + std::int64_t duration) noexcept { + std::uint64_t unsigned_limit = 0; + if (!reference_limit_ticks(policy, unsigned_limit) || + unsigned_limit > static_cast(std::numeric_limits::max())) { + return false; + } + const auto limit = static_cast(unsigned_limit); + return timestamp >= 0 && duration >= 0 && timestamp <= limit && duration <= limit - timestamp; +} + +bool detail::reference_audio_event_timestamp_within_padding( + const ReferenceMediaDecodePolicy& policy, std::int64_t timestamp, + std::uint64_t maximum_padding_ticks) noexcept { + std::uint64_t reference_limit = 0; + if (!reference_limit_ticks(policy, reference_limit) || + maximum_padding_ticks > std::numeric_limits::max() - reference_limit) { + return false; + } + if (timestamp < 0) { + const std::uint64_t leading_ticks = static_cast(-(timestamp + 1)) + 1; + return leading_ticks <= maximum_padding_ticks; + } + return static_cast(timestamp) <= reference_limit + maximum_padding_ticks; +} + +bool detail::account_reference_audio_decode(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, + std::uint64_t maximum_padding_frames, + ReferenceAudioDecodeState& state) noexcept { + return account_reference_audio_decode_impl(timestamp, frame_count, sample_rate, + maximum_padding_frames, state, nullptr, policy); +} + +void detail::require_nonempty_decoded_buffer(std::uint32_t current_length, + std::string_view media_kind) { + if (current_length == 0) { + throw std::runtime_error("Media Foundation returned an empty decoded " + + std::string(media_kind) + " buffer"); + } +} + +bool is_mp4_path(std::string_view path) { + return lowercase_extension(path) == ".mp4"; +} + +void write_mp4(const VideoResult& result, const std::string& path) { +#if defined(_WIN32) + validate_result(result); + if (!is_mp4_path(path)) + throw std::runtime_error("write_mp4 output path must end in .mp4"); + const auto output_path = std::filesystem::path(path); + if (!output_path.parent_path().empty()) + std::filesystem::create_directories(output_path.parent_path()); + + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 2), "MFCreateAttributes(sink writer)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable Media Foundation hardware transforms"); + check_hresult(attributes->SetUINT32(MF_LOW_LATENCY, FALSE), + "configure Media Foundation latency"); + + ComPtr writer; + check_hresult(MFCreateSinkWriterFromURL(output_path.wstring().c_str(), nullptr, + attributes.Get(), &writer), + "MFCreateSinkWriterFromURL"); + + const auto width = static_cast(result.frames.width); + const auto height = static_cast(result.frames.height); + const auto fps = static_cast(result.fps); + DWORD video_stream = 0; + const auto video_output = make_video_output_type(width, height, fps); + check_hresult(writer->AddStream(video_output.Get(), &video_stream), + "IMFSinkWriter::AddStream(video)"); + const auto video_input = make_video_input_type(width, height, fps); + check_hresult(writer->SetInputMediaType(video_stream, video_input.Get(), nullptr), + "IMFSinkWriter::SetInputMediaType(video)"); + + const bool has_audio = !result.audio.samples.empty(); + DWORD audio_stream = 0; + std::uint32_t audio_rate = 0; + std::uint32_t audio_channels = 0; + if (has_audio) { + audio_rate = static_cast(result.audio.sample_rate); + audio_channels = static_cast(result.audio.channels); + const auto audio_output = make_audio_output_type(audio_rate, audio_channels); + check_hresult(writer->AddStream(audio_output.Get(), &audio_stream), + "IMFSinkWriter::AddStream(audio)"); + const auto audio_input = make_audio_input_type(audio_rate, audio_channels); + check_hresult(writer->SetInputMediaType(audio_stream, audio_input.Get(), nullptr), + "IMFSinkWriter::SetInputMediaType(audio)"); + } + + check_hresult(writer->BeginWriting(), "IMFSinkWriter::BeginWriting"); + + const std::size_t pixels_per_frame = static_cast(width) * height * 3; + const std::uint64_t audio_frames = has_audio ? result.audio.samples.size() / audio_channels : 0; + constexpr std::uint64_t kAudioChunkFrames = 1024; + std::uint64_t video_index = 0; + std::uint64_t audio_offset = 0; + std::vector nv12; + std::vector pcm; + + while (video_index < static_cast(result.frames.num_frames) || + audio_offset < audio_frames) { + const auto next_video_time = + video_index < static_cast(result.frames.num_frames) + ? media_time(video_index, fps) + : std::numeric_limits::max(); + const auto next_audio_time = audio_offset < audio_frames + ? media_time(audio_offset, audio_rate) + : std::numeric_limits::max(); + if (next_video_time <= next_audio_time) { + const float* source = result.frames.pixels.data() + + static_cast(video_index) * pixels_per_frame; + rgb_to_nv12(source, width, height, nv12); + const auto end_time = media_time(video_index + 1, fps); + auto sample = + make_sample(nv12.data(), nv12.size(), next_video_time, end_time - next_video_time); + check_hresult(writer->WriteSample(video_stream, sample.Get()), + "IMFSinkWriter::WriteSample(video)"); + ++video_index; + } else { + const auto chunk_frames = std::min(kAudioChunkFrames, audio_frames - audio_offset); + pcm.resize(static_cast(chunk_frames) * audio_channels); + const auto scalar_offset = static_cast(audio_offset) * audio_channels; + for (std::size_t index = 0; index < pcm.size(); ++index) { + const float value = + std::clamp(result.audio.samples[scalar_offset + index], -1.0F, 1.0F); + pcm[index] = static_cast( + std::lround(value * (value < 0.0F ? 32768.0F : 32767.0F))); + } + const auto end_time = media_time(audio_offset + chunk_frames, audio_rate); + auto sample = make_sample(pcm.data(), pcm.size() * sizeof(std::int16_t), + next_audio_time, end_time - next_audio_time); + check_hresult(writer->WriteSample(audio_stream, sample.Get()), + "IMFSinkWriter::WriteSample(audio)"); + audio_offset += chunk_frames; + } + } + + check_hresult(writer->Finalize(), "IMFSinkWriter::Finalize"); +#else + (void)result; + (void)path; + throw std::runtime_error("native MP4 output is available on Windows through Media Foundation"); +#endif +} + +VideoClipInput read_video_file(const std::string& path, const ReferenceMediaDecodePolicy& policy) { +#if defined(_WIN32) + validate_reference_media_decode_policy(policy); + if (path.empty() || std::filesystem::is_directory(path)) + throw std::runtime_error("read_video_file requires a media file path"); + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 2), "MFCreateAttributes(source reader)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable source-reader hardware transforms"); + check_hresult(attributes->SetUINT32(MF_SOURCE_READER_ENABLE_ADVANCED_VIDEO_PROCESSING, TRUE), + "enable source-reader advanced video processing"); + + ComPtr reader; + check_hresult(MFCreateSourceReaderFromURL(std::filesystem::path(path).wstring().c_str(), + attributes.Get(), &reader), + "MFCreateSourceReaderFromURL"); + const auto declared_duration = presentation_duration(*reader.Get()); + check_hresult(reader->SetStreamSelection(MF_SOURCE_READER_ALL_STREAMS, FALSE), + "disable source-reader streams"); + + DWORD video_stream_index = MAXDWORD; + DWORD audio_stream_index = MAXDWORD; + ComPtr native_video; + ComPtr native_audio; + for (DWORD stream = 0; stream < 64; ++stream) { + ComPtr native_type; + const HRESULT type_result = reader->GetNativeMediaType(stream, 0, &native_type); + if (type_result == MF_E_INVALIDSTREAMNUMBER) + break; + if (FAILED(type_result)) + continue; + GUID major_type{}; + if (FAILED(native_type->GetGUID(MF_MT_MAJOR_TYPE, &major_type))) + continue; + if (major_type == MFMediaType_Video && video_stream_index == MAXDWORD) { + video_stream_index = stream; + native_video = native_type; + } else if (major_type == MFMediaType_Audio && audio_stream_index == MAXDWORD) { + audio_stream_index = stream; + native_audio = native_type; + } + } + if (video_stream_index == MAXDWORD || !native_video) + throw std::runtime_error("reference media contains no video stream"); + const CodecPaddingAllowance padding = codec_padding_allowance(native_audio.Get()); + if (declared_duration && + *declared_duration > + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond + + padding.duration) { + throw std::runtime_error( + "reference media presentation exceeds the configured duration limit"); + } + const bool trim_codec_padding = padding.frames > 0; + UINT32 fps_numerator = 0; + UINT32 fps_denominator = 0; + check_hresult( + MFGetAttributeRatio(native_video.Get(), MF_MT_FRAME_RATE, &fps_numerator, &fps_denominator), + "read source frame rate"); + if (fps_numerator == 0 || fps_denominator == 0 || + static_cast(fps_numerator) > + static_cast(policy.maximum_source_video_fps) * fps_denominator) { + throw std::runtime_error("reference video source frame rate exceeds the configured limit"); + } + UINT32 source_width = 0; + UINT32 source_height = 0; + check_hresult( + MFGetAttributeSize(native_video.Get(), MF_MT_FRAME_SIZE, &source_width, &source_height), + "read source frame size"); + const auto [requested_width, requested_height] = + detail::reference_video_decode_size(policy, source_width, source_height); + const ReferenceVideoSize requested_size{requested_width, requested_height}; + check_hresult(reader->SetStreamSelection(video_stream_index, TRUE), "select source video"); + const auto requested_video = + source_reader_video_type(requested_size.width, requested_size.height); + check_hresult(reader->SetCurrentMediaType(video_stream_index, nullptr, requested_video.Get()), + "SetCurrentMediaType(video RGB32)"); + ComPtr decoded_video; + check_hresult(reader->GetCurrentMediaType(video_stream_index, &decoded_video), + "GetCurrentMediaType(video)"); + UINT32 width = 0; + UINT32 height = 0; + check_hresult(MFGetAttributeSize(decoded_video.Get(), MF_MT_FRAME_SIZE, &width, &height), + "read decoded frame size"); + if (width != requested_size.width || height != requested_size.height) { + throw std::runtime_error( + "Media Foundation did not honor the bounded reference-video decode canvas"); + } + UINT32 raw_stride = 0; + LONG stride = static_cast(width * 4); + if (SUCCEEDED(decoded_video->GetUINT32(MF_MT_DEFAULT_STRIDE, &raw_stride))) + stride = static_cast(raw_stride); + + bool has_audio = false; + std::uint32_t audio_rate = 0; + std::uint32_t audio_channels = 0; + if (audio_stream_index != MAXDWORD && native_audio) { + UINT32 rate = 0; + UINT32 channels = 0; + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &rate), + "read source soundtrack sample rate"); + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &channels), + "read source soundtrack channel count"); + detail::validate_reference_video_soundtrack_format(rate, channels); + audio_rate = rate; + audio_channels = channels; + check_hresult(reader->SetStreamSelection(audio_stream_index, TRUE), "select source audio"); + const auto requested_audio = source_reader_audio_type(audio_rate, audio_channels); + check_hresult( + reader->SetCurrentMediaType(audio_stream_index, nullptr, requested_audio.Get()), + "SetCurrentMediaType(audio float)"); + ComPtr decoded_audio; + check_hresult(reader->GetCurrentMediaType(audio_stream_index, &decoded_audio), + "GetCurrentMediaType(audio)"); + UINT32 decoded_rate = 0; + UINT32 decoded_channels = 0; + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &decoded_rate), + "read decoded soundtrack sample rate"); + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &decoded_channels), + "read decoded soundtrack channel count"); + detail::validate_reference_video_soundtrack_format(decoded_rate, decoded_channels); + if (decoded_rate != audio_rate || decoded_channels != audio_channels) { + throw std::runtime_error( + "Media Foundation did not honor the requested soundtrack PCM format"); + } + has_audio = true; + } + + if (width == 0 || height == 0 || fps_numerator == 0 || fps_denominator == 0 || + width > static_cast(std::numeric_limits::max()) || + height > static_cast(std::numeric_limits::max()) || + fps_numerator > static_cast(std::numeric_limits::max()) || + fps_denominator > static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded video metadata is outside the public C++ value range"); + + VideoClipInput result; + result.width = static_cast(width); + result.height = static_cast(height); + result.channels = 3; + const bool downsample_video = + static_cast(fps_numerator) > + static_cast(policy.target_video_fps) * fps_denominator; + result.fps_numerator = + static_cast(downsample_video ? policy.target_video_fps : fps_numerator); + result.fps_denominator = static_cast(downsample_video ? 1 : fps_denominator); + result.soundtrack.sample_rate = static_cast(audio_rate); + result.soundtrack.channels = static_cast(audio_channels); + + const std::uint64_t maximum_source_video_frames = std::min( + detail::reference_video_frame_ceiling(policy, fps_numerator, fps_denominator), + static_cast(std::numeric_limits::max())); + const std::uint64_t maximum_video_frames = + downsample_video + ? static_cast(policy.maximum_duration_seconds) * policy.target_video_fps + : maximum_source_video_frames; + const std::uint64_t frame_scalars = static_cast(width) * height * 3; + if (maximum_video_frames > + static_cast(std::numeric_limits::max()) / frame_scalars) + throw std::runtime_error("bounded reference-video RGB allocation exceeds host range"); + const std::size_t maximum_rgb_scalars = + static_cast(maximum_video_frames * frame_scalars); + const std::size_t maximum_soundtrack_scalars = + has_audio ? maximum_reference_audio_scalars(audio_rate, audio_channels, policy) : 0; + + bool video_done = false; + bool audio_done = !has_audio; + LONGLONG last_video_tick = std::numeric_limits::min(); + LONGLONG last_audio_tick = std::numeric_limits::min(); + std::size_t consecutive_empty_reads = 0; + std::size_t total_empty_reads = 0; + std::uint64_t source_video_frames = 0; + detail::ReferenceAudioDecodeState audio_decode_state; + while (!video_done || !audio_done) { + DWORD stream_index = 0; + DWORD flags = 0; + LONGLONG timestamp = 0; + ComPtr sample; + check_hresult(reader->ReadSample(MF_SOURCE_READER_ANY_STREAM, 0, &stream_index, &flags, + ×tamp, &sample), + "IMFSourceReader::ReadSample"); + reject_source_reader_error(flags, "reference media source reader"); + if ((flags & (MF_SOURCE_READERF_NATIVEMEDIATYPECHANGED | + MF_SOURCE_READERF_CURRENTMEDIATYPECHANGED)) != 0) + throw std::runtime_error("reference video changes media type mid-stream"); + if ((flags & MF_SOURCE_READERF_ENDOFSTREAM) != 0) { + if (stream_index == video_stream_index) + video_done = true; + else if (stream_index == audio_stream_index) + audio_done = true; + else + throw std::runtime_error( + "reference media ended an unexpected source-reader stream"); + continue; + } + if (stream_index != video_stream_index && stream_index != audio_stream_index) + throw std::runtime_error("reference media returned an unexpected source-reader stream"); + const bool audio_padding_event = trim_codec_padding && stream_index == audio_stream_index; + const std::uint64_t event_padding_ticks = + audio_padding_event ? static_cast(padding.duration) : 0; + if (!audio_padding_event) + reject_reference_timestamp(timestamp, "reference media source reader", policy); + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (sample) + throw std::runtime_error( + "reference media returned a sample for an empty stream tick"); + auto& last_tick = + stream_index == video_stream_index ? last_video_tick : last_audio_tick; + note_empty_source_reader_event(flags, timestamp, last_tick, consecutive_empty_reads, + total_empty_reads, event_padding_ticks, + "reference media source reader", policy); + continue; + } + if (!sample) { + auto& last_tick = + stream_index == video_stream_index ? last_video_tick : last_audio_tick; + note_empty_source_reader_event(flags, timestamp, last_tick, consecutive_empty_reads, + total_empty_reads, event_padding_ticks, + "reference media source reader", policy); + continue; + } + reject_empty_decoded_sample(sample.Get(), + stream_index == video_stream_index ? "video" : "audio"); + if (stream_index == video_stream_index) { + if (source_video_frames >= maximum_source_video_frames) + throw std::runtime_error( + "decoded reference video exceeds the configured duration limit"); + reject_video_sample_timeline(sample.Get(), timestamp, fps_numerator, fps_denominator, + policy); + bool retain = true; + if (downsample_video) { + const auto begin = rounded_reference_frame_slot(source_video_frames, fps_numerator, + fps_denominator, policy); + const auto end = rounded_reference_frame_slot( + source_video_frames + 1, fps_numerator, fps_denominator, policy); + retain = end > begin; + } + ++source_video_frames; + if (retain) { + if (static_cast(result.num_frames) >= maximum_video_frames) + throw std::runtime_error( + "decoded reference video exceeds the bounded target frame count"); + append_rgb32_frame(sample.Get(), width, height, stride, result.pixels, + maximum_rgb_scalars); + ++result.num_frames; + } + } else if (stream_index == audio_stream_index) { + append_float_audio(sample.Get(), audio_rate, audio_channels, timestamp, padding.frames, + audio_decode_state, result.soundtrack.samples, + maximum_soundtrack_scalars, policy); + } + consecutive_empty_reads = 0; + } + + if (result.num_frames <= 0) + throw std::runtime_error("reference video contains no decoded frames"); + if (!result.soundtrack.samples.empty()) { + if (result.soundtrack.samples.size() > + static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded soundtrack exceeds the public C++ value range"); + result.soundtrack.num_samples = static_cast(result.soundtrack.samples.size()); + } + return result; +#else + (void)path; + (void)policy; + throw std::runtime_error( + "native media-file input is available on Windows through Media Foundation"); +#endif +} + +AudioResult read_audio_file(const std::string& path, const ReferenceMediaDecodePolicy& policy) { +#if defined(_WIN32) + validate_reference_media_decode_policy(policy); + if (path.empty() || std::filesystem::is_directory(path)) + throw std::runtime_error("read_audio_file requires a media file path"); + MediaFoundationSession session; + ComPtr attributes; + check_hresult(MFCreateAttributes(&attributes, 1), "MFCreateAttributes(audio source reader)"); + check_hresult(attributes->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, TRUE), + "enable audio source-reader hardware transforms"); + + ComPtr reader; + check_hresult(MFCreateSourceReaderFromURL(std::filesystem::path(path).wstring().c_str(), + attributes.Get(), &reader), + "MFCreateSourceReaderFromURL(audio)"); + const auto declared_duration = presentation_duration(*reader.Get()); + check_hresult(reader->SetStreamSelection(MF_SOURCE_READER_ALL_STREAMS, FALSE), + "disable audio source-reader streams"); + + DWORD audio_stream_index = MAXDWORD; + ComPtr native_audio; + for (DWORD stream = 0; stream < 64; ++stream) { + ComPtr native_type; + const HRESULT type_result = reader->GetNativeMediaType(stream, 0, &native_type); + if (type_result == MF_E_INVALIDSTREAMNUMBER) + break; + if (FAILED(type_result)) + continue; + GUID major_type{}; + if (SUCCEEDED(native_type->GetGUID(MF_MT_MAJOR_TYPE, &major_type)) && + major_type == MFMediaType_Audio) { + audio_stream_index = stream; + native_audio = native_type; + break; + } + } + if (audio_stream_index == MAXDWORD || !native_audio) + throw std::runtime_error("reference audio file contains no audio stream"); + + UINT32 native_rate = 0; + UINT32 native_channels = 0; + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &native_rate), + "read source audio sample rate"); + check_hresult(native_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &native_channels), + "read source audio channel count"); + if (native_rate == 0 || + native_rate > static_cast(std::numeric_limits::max()) || + (native_channels != 1 && native_channels != 2)) + throw std::runtime_error( + "reference audio metadata is outside the supported mono/stereo C++ value range"); + const CodecPaddingAllowance padding = codec_padding_allowance(native_audio.Get()); + if (declared_duration && + *declared_duration > + static_cast(policy.maximum_duration_seconds) * kMediaTicksPerSecond + + padding.duration) { + throw std::runtime_error( + "reference audio presentation exceeds the configured duration limit"); + } + const bool trim_codec_padding = padding.frames > 0; + + check_hresult(reader->SetStreamSelection(audio_stream_index, TRUE), "select source audio"); + const auto requested_audio = source_reader_audio_type(native_rate, native_channels); + check_hresult(reader->SetCurrentMediaType(audio_stream_index, nullptr, requested_audio.Get()), + "SetCurrentMediaType(audio float)"); + ComPtr decoded_audio; + check_hresult(reader->GetCurrentMediaType(audio_stream_index, &decoded_audio), + "GetCurrentMediaType(audio)"); + UINT32 decoded_rate = 0; + UINT32 decoded_channels = 0; + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &decoded_rate), + "read decoded audio sample rate"); + check_hresult(decoded_audio->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &decoded_channels), + "read decoded audio channel count"); + if (decoded_rate == 0 || + decoded_rate > static_cast(std::numeric_limits::max()) || + (decoded_channels != 1 && decoded_channels != 2)) + throw std::runtime_error( + "decoded audio metadata is outside the supported mono/stereo C++ value range"); + if (decoded_rate != native_rate || decoded_channels != native_channels) { + throw std::runtime_error("Media Foundation did not honor the requested decoded PCM format"); + } + + AudioResult result; + result.sample_rate = static_cast(decoded_rate); + result.channels = static_cast(decoded_channels); + const std::size_t maximum_scalars = + maximum_reference_audio_scalars(decoded_rate, decoded_channels, policy); + LONGLONG last_tick_timestamp = std::numeric_limits::min(); + std::size_t consecutive_empty_reads = 0; + std::size_t total_empty_reads = 0; + detail::ReferenceAudioDecodeState decode_state; + while (true) { + DWORD stream_index = 0; + DWORD flags = 0; + LONGLONG timestamp = 0; + ComPtr sample; + check_hresult( + reader->ReadSample(audio_stream_index, 0, &stream_index, &flags, ×tamp, &sample), + "IMFSourceReader::ReadSample(audio)"); + reject_source_reader_error(flags, "reference audio source reader"); + if ((flags & (MF_SOURCE_READERF_NATIVEMEDIATYPECHANGED | + MF_SOURCE_READERF_CURRENTMEDIATYPECHANGED)) != 0) + throw std::runtime_error("reference audio changes media type mid-stream"); + if ((flags & MF_SOURCE_READERF_ENDOFSTREAM) != 0) + break; + if (!trim_codec_padding) + reject_reference_timestamp(timestamp, "reference audio source reader", policy); + if (stream_index != audio_stream_index) + throw std::runtime_error("reference audio returned an unexpected source-reader stream"); + if ((flags & MF_SOURCE_READERF_STREAMTICK) != 0) { + if (sample) + throw std::runtime_error( + "reference audio returned a sample for an empty stream tick"); + note_empty_source_reader_event(flags, timestamp, last_tick_timestamp, + consecutive_empty_reads, total_empty_reads, + static_cast(padding.duration), + "reference audio source reader", policy); + continue; + } + if (!sample) { + note_empty_source_reader_event(flags, timestamp, last_tick_timestamp, + consecutive_empty_reads, total_empty_reads, + static_cast(padding.duration), + "reference audio source reader", policy); + continue; + } + reject_empty_decoded_sample(sample.Get(), "audio"); + append_float_audio(sample.Get(), decoded_rate, decoded_channels, timestamp, padding.frames, + decode_state, result.samples, maximum_scalars, policy); + consecutive_empty_reads = 0; + } + if (result.samples.empty()) + throw std::runtime_error("reference audio file contains no decoded samples"); + if (result.samples.size() > static_cast(std::numeric_limits::max())) + throw std::runtime_error("decoded audio exceeds the public C++ value range"); + result.num_samples = static_cast(result.samples.size()); + return result; +#else + (void)policy; + if (lowercase_extension(path) == ".wav") + return io::read_wav(path); + throw std::runtime_error( + "compressed native media-file input is available on Windows through Media Foundation"); +#endif +} + +} // namespace trtmc::cli diff --git a/apps/cli/windows_media.h b/apps/cli/windows_media.h new file mode 100644 index 0000000000..23784a3443 --- /dev/null +++ b/apps/cli/windows_media.h @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "trtmc/task.h" + +#include +#include +#include +#include +#include +#include + +namespace trtmc::cli { + +namespace detail { +// Overflow-safe decoded-frame allocation ceiling. Presentation timestamps and +// duration remain authoritative for the configured validity decision. +std::uint64_t reference_video_frame_ceiling(const ReferenceMediaDecodePolicy& policy, + std::uint32_t fps_numerator, + std::uint32_t fps_denominator); +std::pair +reference_video_decode_size(const ReferenceMediaDecodePolicy& policy, std::uint32_t source_width, + std::uint32_t source_height); +bool reference_timeline_within_limit(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::int64_t duration) noexcept; +bool reference_audio_event_timestamp_within_padding(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, + std::uint64_t maximum_padding_ticks) noexcept; +struct ReferenceAudioDecodeState { + std::uint64_t decoded_frames{0}; + std::uint64_t decoded_padding_frames{0}; +}; +// Account decoded PCM independently of container duration metadata. Compressed +// audio may decode at most maximum_padding_frames beyond the configured sample +// budget, and no more than that many decoded frames may fall outside its timeline. +bool account_reference_audio_decode(const ReferenceMediaDecodePolicy& policy, + std::int64_t timestamp, std::uint64_t frame_count, + std::uint32_t sample_rate, std::uint64_t maximum_padding_frames, + ReferenceAudioDecodeState& state) noexcept; +// Append retained PCM frames at their presentation timestamp. Timeline gaps +// become silence; a later overlapping sample contributes only its new tail. +void append_reference_audio_frames_on_timeline( + std::vector& destination, const void* source, std::uint64_t source_frame_count, + std::uint64_t first_source_frame, std::uint64_t end_source_frame, std::uint32_t channels, + std::int64_t timestamp, std::uint32_t sample_rate, std::size_t maximum_scalars); +// A discovered video soundtrack must fail closed when its native or decoded +// format cannot be represented by the public mono/stereo audio value type. +void validate_reference_video_soundtrack_format(std::uint32_t sample_rate, std::uint32_t channels); +// Reject a decoded Media Foundation sample/buffer that carries no payload. +// media_kind is a trusted internal label such as "audio" or "video". +void require_nonempty_decoded_buffer(std::uint32_t current_length, std::string_view media_kind); +} // namespace detail + +bool is_mp4_path(std::string_view path); + +// Write synchronized H.264 video and optional AAC audio using the Windows +// Media Foundation codecs shipped with the operating system. No FFmpeg or +// other runtime media dependency is involved. +void write_mp4(const VideoResult& result, const std::string& path); + +// Decode a Windows-supported media container into the public THWC RGB video +// and interleaved float-audio value types. Decode fails closed +// as soon as the pipeline's reference-media policy is exceeded. +VideoClipInput read_video_file(const std::string& path, const ReferenceMediaDecodePolicy& policy); + +// Decode the first audio stream in a Windows-supported media file (including +// MP3 and WAV) into interleaved float32 samples. The operating-system Media +// Foundation codecs are the only runtime dependency. Decode fails closed at +// the pipeline's reference-media duration limit. +AudioResult read_audio_file(const std::string& path, const ReferenceMediaDecodePolicy& policy); + +} // namespace trtmc::cli diff --git a/apps/cli/windows_utf8_argv.cpp b/apps/cli/windows_utf8_argv.cpp new file mode 100644 index 0000000000..226cf207bb --- /dev/null +++ b/apps/cli/windows_utf8_argv.cpp @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cli/windows_utf8_argv.h" + +#if !defined(_WIN32) +#error "windows_utf8_argv.cpp is Windows-only" +#endif + +#define NOMINMAX +#define WIN32_LEAN_AND_MEAN +#include +#include +#include +#include + +namespace trtmc::cli { + +namespace { + +int checked_utf16_size(std::wstring_view value) { + if (value.size() > static_cast(std::numeric_limits::max())) + throw std::length_error("Windows command-line argument is too large to convert to UTF-8"); + return static_cast(value.size()); +} + +[[noreturn]] void throw_conversion_error(DWORD error) { + throw std::runtime_error("WideCharToMultiByte(CP_UTF8) rejected command-line UTF-16 (Windows " + "error " + + std::to_string(error) + ")"); +} + +} // namespace + +std::string utf8_from_utf16(std::wstring_view value) { + if (value.empty()) + return {}; + + const int input_size = checked_utf16_size(value); + const int output_size = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, value.data(), + input_size, nullptr, 0, nullptr, nullptr); + if (output_size <= 0) + throw_conversion_error(GetLastError()); + + std::string result(static_cast(output_size), '\0'); + const int converted = + WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, value.data(), input_size, result.data(), + output_size, nullptr, nullptr); + if (converted != output_size) + throw_conversion_error(GetLastError()); + return result; +} + +Utf8CommandLine::Utf8CommandLine(int argc, wchar_t* const* argv) : argc_(argc) { + if (argc < 0 || (argc > 0 && argv == nullptr)) + throw std::invalid_argument("Windows command line has invalid argc/argv metadata"); + + storage_.reserve(static_cast(argc)); + for (int index = 0; index < argc; ++index) { + if (argv[index] == nullptr) + throw std::invalid_argument("Windows command line contains a null argument"); + storage_.push_back(utf8_from_utf16(argv[index])); + } + + pointers_.reserve(static_cast(argc) + 1U); + for (auto& argument : storage_) + pointers_.push_back(argument.data()); + pointers_.push_back(nullptr); +} + +} // namespace trtmc::cli diff --git a/apps/cli/windows_utf8_argv.h b/apps/cli/windows_utf8_argv.h new file mode 100644 index 0000000000..1c457c73a9 --- /dev/null +++ b/apps/cli/windows_utf8_argv.h @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc::cli { + +// Convert one UTF-16 command-line value to UTF-8. The Windows implementation +// rejects malformed surrogate sequences instead of replacing them. +std::string utf8_from_utf16(std::wstring_view value); + +// Own the UTF-8 storage and the mutable pointer array expected by the existing +// platform-neutral CLI parser. argv()[argc()] is always null. +class Utf8CommandLine { + public: + Utf8CommandLine(int argc, wchar_t* const* argv); + + int argc() const noexcept { return argc_; } + char** argv() noexcept { return pointers_.data(); } + + private: + int argc_{0}; + std::vector storage_; + std::vector pointers_; +}; + +} // namespace trtmc::cli diff --git a/core/builder/tensorrt_model_connect/build.py b/core/builder/tensorrt_model_connect/build.py index d0d5d68638..8884dce612 100644 --- a/core/builder/tensorrt_model_connect/build.py +++ b/core/builder/tensorrt_model_connect/build.py @@ -42,6 +42,7 @@ class BuildRequest: dynamic_kv_cache: bool = False verbose: bool = False graph_transform: GraphTransform | None = None + family_options: tuple[tuple[str, str | int | float | bool | None], ...] = () def __post_init__(self) -> None: if not self.precision: @@ -70,6 +71,19 @@ def __post_init__(self) -> None: raise ValueError("dynamic_kv_cache must be a bool") if self.graph_transform is not None and not callable(self.graph_transform): raise ValueError("graph_transform must be callable when provided") + if not isinstance(self.family_options, tuple): + raise ValueError("family_options must be an immutable tuple") + names: set[str] = set() + for item in self.family_options: + if not isinstance(item, tuple) or len(item) != 2: + raise ValueError("family_options entries must be (name, value) tuples") + name, value = item + _validate_id("family option", name) + if name in names: + raise ValueError(f"duplicate family option: {name}") + if not isinstance(value, (str, int, float, bool, type(None))): + raise ValueError("family option values must be JSON scalar values") + names.add(name) def _validate_id(field: str, value: object) -> str: @@ -136,6 +150,11 @@ def build(request: BuildRequest) -> None: family = _resolve_family(request) _select_backend(request.backend) family_module = _load_family(family) + if request.family_options: + validate_options = getattr(family_module, "validate_build_options", None) + if validate_options is None: + raise ValueError(f"family {family!r} does not accept family_options") + validate_options(dict(request.family_options)) writer = BundleWriter(request.output_path) try: with graph_transform(request.graph_transform): diff --git a/core/builder/tensorrt_model_connect/build_cli.py b/core/builder/tensorrt_model_connect/build_cli.py index 9042a68320..c909bb213d 100644 --- a/core/builder/tensorrt_model_connect/build_cli.py +++ b/core/builder/tensorrt_model_connect/build_cli.py @@ -7,6 +7,7 @@ import argparse import json +import re from pathlib import Path from typing import Sequence @@ -14,6 +15,9 @@ from .model_support import load_model_metadata, resolve_family +_OPTION = re.compile(r"([a-z][a-z0-9_]*)\.([a-z][a-z0-9_]*)=(.*)\Z", re.DOTALL) + + def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="trtmc") commands = parser.add_subparsers(dest="command", required=True) @@ -35,6 +39,14 @@ def _parser() -> argparse.ArgumentParser: build_parser.add_argument("--fp32-layer", type=int, action="append", default=[]) build_parser.add_argument("--dynamic-kv-cache", action="store_true") build_parser.add_argument("--verbose", action="store_true") + build_parser.add_argument( + "--set", + dest="family_options", + action="append", + default=[], + metavar="FAMILY.KEY=VALUE", + help="Set one family-owned scalar build option", + ) prepare_parser = commands.add_parser( "prepare-structure", help="Prepare one structure request without rebuilding its model bundle", @@ -47,6 +59,33 @@ def _parser() -> argparse.ArgumentParser: return parser +def _parse_family_options( + values: Sequence[str], family: str +) -> tuple[tuple[str, str | int | float | bool | None], ...]: + result: list[tuple[str, str | int | float | bool | None]] = [] + names: set[str] = set() + for value in values: + match = _OPTION.fullmatch(value) + if match is None: + raise ValueError("--set must use FAMILY.KEY=VALUE") + namespace, name, raw = match.groups() + if namespace != family: + raise ValueError( + f"--set namespace {namespace!r} does not match resolved family {family!r}" + ) + if name in names: + raise ValueError(f"duplicate --set option: {family}.{name}") + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + parsed = raw + if not isinstance(parsed, (str, int, float, bool, type(None))): + raise ValueError("--set values must be JSON scalars") + result.append((name, parsed)) + names.add(name) + return tuple(result) + + def main(argv: Sequence[str] | None = None) -> int: args = _parser().parse_args(argv) model_dir = _resolve_model(args.model, args.revision) @@ -93,6 +132,7 @@ def main(argv: Sequence[str] | None = None) -> int: fp32_layers=tuple(args.fp32_layer), dynamic_kv_cache=args.dynamic_kv_cache, verbose=args.verbose, + family_options=_parse_family_options(args.family_options, family), ) ) return 0 diff --git a/core/builder/tensorrt_model_connect/bundle_writer.py b/core/builder/tensorrt_model_connect/bundle_writer.py index fc0944d807..5836009899 100644 --- a/core/builder/tensorrt_model_connect/bundle_writer.py +++ b/core/builder/tensorrt_model_connect/bundle_writer.py @@ -38,6 +38,11 @@ def _validate_nonempty_string(field: str, value: object) -> str: return value +def _file_identity(path: Path) -> tuple[int, int, int, int]: + info = path.stat() + return info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns + + class BundleWriter: """Stage named sections and atomically publish one bundle.""" @@ -50,6 +55,7 @@ def __init__(self, destination: str | Path) -> None: self._header: dict[str, Any] | None = None self._sections: list[tuple[str, Path]] = [] self._section_names: set[str] = set() + self._borrowed_files: dict[str, tuple[int, int, int, int]] = {} self._staging_dir: Path | None = None self._open_sections = 0 self._failed_section = False @@ -89,15 +95,18 @@ def set_header(self, *, family: str, task: str, backend: str) -> None: "backend": _validate_id("backend", backend), } - @contextmanager - def open_section(self, name: str) -> Iterator[BinaryIO]: - """Open one file-backed section for incremental binary writes.""" - + def _validate_section_name(self, name: str) -> str: self._ensure_writable() name = _validate_nonempty_string("section name", name) if name in self._section_names: raise ValueError(f"duplicate bundle section name: {name!r}") + return name + @contextmanager + def open_section(self, name: str) -> Iterator[BinaryIO]: + """Open one file-backed section for incremental binary writes.""" + + name = self._validate_section_name(name) section_path = self._ensure_staging_dir() / f"section-{len(self._sections)}" self._section_names.add(name) self._sections.append((name, section_path)) @@ -111,6 +120,31 @@ def open_section(self, name: str) -> Iterator[BinaryIO]: finally: self._open_sections -= 1 + def add_file(self, name: str, path: str | Path) -> None: + """Borrow an existing file without copying it into section staging. + + The caller must keep the file unchanged until finish or abort. The + writer only reads it; cleanup never removes borrowed source files. + """ + + name = self._validate_section_name(name) + source = Path(path).absolute() + if not source.is_file(): + raise FileNotFoundError(f"bundle section source is not a file: {source}") + resolved = source.resolve() + if resolved == self._destination.resolve() or ( + self._staging_dir is not None and self._staging_dir.resolve() in resolved.parents + ): + raise ValueError("borrowed bundle section must be outside writer-owned paths") + self._borrowed_files[name] = _file_identity(source) + self._section_names.add(name) + self._sections.append((name, source)) + + def _check_borrowed_file(self, name: str, path: Path) -> None: + identity = self._borrowed_files.get(name) + if identity is not None and _file_identity(path) != identity: + raise RuntimeError(f"borrowed bundle section source changed: {name!r}") + def add_bytes(self, name: str, data: bytes) -> None: """Add a complete in-memory binary section.""" @@ -139,6 +173,7 @@ def finish(self) -> None: section_table: dict[str, dict[str, int]] = {} offset = 0 for name, path in self._sections: + self._check_borrowed_file(name, path) length = path.stat().st_size if length > _MAX_UINT64 - offset: raise OverflowError("bundle section table exceeds uint64 range") @@ -162,9 +197,11 @@ def finish(self) -> None: output.write(BUNDLE_MAGIC) output.write(struct.pack(" None: assert request.context_parallel_size == 3 assert request.backend == "trt" assert request.dynamic_kv_cache is False + assert request.family_options == () @pytest.mark.parametrize( @@ -57,6 +58,10 @@ def test_build_request_is_a_plain_frozen_dataclass(tmp_path: Path) -> None: ("dynamic_kv_cache", 1), ("graph_transform", object()), ("backend", "unknown"), + ("family_options", [("option", True)]), + ("family_options", (("bad-name", True),)), + ("family_options", (("option", []),)), + ("family_options", (("option", True), ("option", False))), ], ) def test_build_request_rejects_invalid_direct_inputs( @@ -235,6 +240,22 @@ def transform(network: object, engine_index: int) -> None: assert fake_trt.Builder is FakeTrtBuilder +def test_nonempty_family_options_require_family_validation(monkeypatch, tmp_path: Path) -> None: + request = replace(_request(tmp_path), family_options=(("example_option", True),)) + family = SimpleNamespace(build=lambda _request, _writer: pytest.fail("unexpected build")) + monkeypatch.setattr(build_core, "_load_family", lambda _family: family) + with pytest.raises(ValueError, match="does not accept family_options"): + build_core.build(request) + + def reject_options(options): + assert options == {"example_option": True} + raise ValueError("unsupported example option") + + family.validate_build_options = reject_options + with pytest.raises(ValueError, match="unsupported example option"): + build_core.build(request) + + def test_build_aborts_and_preserves_family_error(monkeypatch, tmp_path: Path) -> None: events: list[str] = [] family_error = RuntimeError("family failed") diff --git a/core/builder/tests/test_build_cli.py b/core/builder/tests/test_build_cli.py index e082ff1ddc..453467c303 100644 --- a/core/builder/tests/test_build_cli.py +++ b/core/builder/tests/test_build_cli.py @@ -56,6 +56,10 @@ def test_build_command_forwards_only_direct_inputs(monkeypatch, tmp_path: Path) "5", "--dynamic-kv-cache", "--verbose", + "--set", + "patchtsmixer.feature=true", + "--set", + "patchtsmixer.path=C:\\models\\weights.bin", ] ) == 0 @@ -78,6 +82,24 @@ def test_build_command_forwards_only_direct_inputs(monkeypatch, tmp_path: Path) assert request.fp32_layers == (2, 5) assert request.dynamic_kv_cache is True assert request.verbose is True + assert request.family_options == ( + ("feature", True), + ("path", "C:\\models\\weights.bin"), + ) + + +@pytest.mark.parametrize( + "values", + [ + ["wrong.option=true"], + ["gpt2.option={}"], + ["gpt2.option=true", "gpt2.option=false"], + ["missing_namespace=true"], + ], +) +def test_family_options_reject_invalid_or_cross_family_values(values) -> None: + with pytest.raises(ValueError): + build_cli._parse_family_options(values, "gpt2") def test_build_command_uses_the_family_owned_default_task(monkeypatch, tmp_path: Path) -> None: diff --git a/core/builder/tests/test_bundle_writer.py b/core/builder/tests/test_bundle_writer.py index b5406b4cec..94f58a354f 100644 --- a/core/builder/tests/test_bundle_writer.py +++ b/core/builder/tests/test_bundle_writer.py @@ -5,6 +5,7 @@ import json import os +import shutil import struct from pathlib import Path @@ -66,9 +67,7 @@ def test_writer_rejects_duplicate_and_empty_section_names(tmp_path: Path) -> Non ("field", "value"), [("family", "../family"), ("task", ""), ("backend", "")], ) -def test_writer_rejects_unsafe_header_ids( - tmp_path: Path, field: str, value: str -) -> None: +def test_writer_rejects_unsafe_header_ids(tmp_path: Path, field: str, value: str) -> None: header = {"family": "family", "task": "text_generation", "backend": "trt"} header[field] = value writer = BundleWriter(tmp_path / "model.bundle") @@ -111,9 +110,7 @@ def test_abort_discards_staging_and_preserves_destination(tmp_path: Path) -> Non writer.finish() -def test_failed_atomic_replace_preserves_destination( - monkeypatch, tmp_path: Path -) -> None: +def test_failed_atomic_replace_preserves_destination(monkeypatch, tmp_path: Path) -> None: destination = tmp_path / "model.bundle" destination.write_bytes(b"previous") writer = BundleWriter(destination) @@ -131,3 +128,124 @@ def fail_replace(_source: Path, _destination: Path) -> None: assert not list(tmp_path.glob(".model.bundle.*.tmp")) writer.abort() + + +def test_borrowed_file_mixes_with_streamed_sections_without_staging_copy(tmp_path: Path) -> None: + source = tmp_path / "existing.plan" + source.write_bytes(b"engine-bytes") + before = source.stat() + destination = tmp_path / "model.bundle" + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + assert list(tmp_path.iterdir()) == [source] + writer.add_json("config.json", {"size": 7}) + with writer.open_section("tokenizer.model") as section: + section.write(b"tokens") + writer.finish() + + header, payload = _read_bundle(destination) + assert header["sections"] == { + "engine.plan": {"offset": 0, "length": 12}, + "config.json": {"offset": 12, "length": 10}, + "tokenizer.model": {"offset": 22, "length": 6}, + } + assert payload == b'engine-bytes{"size":7}tokens' + assert source.read_bytes() == b"engine-bytes" + assert source.stat().st_mtime_ns == before.st_mtime_ns + assert set(tmp_path.iterdir()) == {source, destination} + with pytest.raises(RuntimeError, match="already finished"): + writer.add_file("other", source) + + +def test_borrowed_files_reuse_section_guards_and_require_external_files(tmp_path: Path) -> None: + source = tmp_path / "existing.plan" + source.write_bytes(b"plan") + writer = BundleWriter(tmp_path / "model.bundle") + writer.add_file("borrowed", source) + with pytest.raises(ValueError, match="duplicate"): + writer.add_bytes("borrowed", b"other") + writer.add_bytes("staged", b"bytes") + with pytest.raises(ValueError, match="duplicate"): + writer.add_file("staged", source) + with pytest.raises(ValueError, match="section name"): + writer.add_file("", source) + for missing in (tmp_path, tmp_path / "missing.plan"): + with pytest.raises(FileNotFoundError, match="source is not a file"): + writer.add_file("invalid", missing) + with writer.open_section("owned") as section: + section.write(b"owned") + owned_path = Path(section.name) + with pytest.raises(ValueError, match="writer-owned paths"): + writer.add_file("invalid", owned_path) + writer.abort() + assert source.read_bytes() == b"plan" + with pytest.raises(RuntimeError, match="aborted"): + writer.add_file("other", source) + + +def test_borrowing_destination_is_rejected_without_modifying_it(tmp_path: Path) -> None: + destination = tmp_path / "model.bundle" + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + with pytest.raises(ValueError, match="writer-owned paths"): + writer.add_file("previous", destination) + writer.abort() + assert destination.read_bytes() == b"previous" + + +@pytest.mark.parametrize("failure", ("abort", "copy", "replace")) +def test_borrowed_source_and_old_destination_survive_abort_or_finish_failure( + tmp_path: Path, monkeypatch, failure: str +) -> None: + source, destination = tmp_path / "existing.plan", tmp_path / "model.bundle" + source.write_bytes(b"plan") + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + writer.add_bytes("metadata", b"metadata") + + def fail(*_args, **_kwargs): + raise OSError("publication failed") + + if failure != "abort": + monkeypatch.setattr( + shutil if failure == "copy" else os, + "copyfileobj" if failure == "copy" else "replace", + fail, + ) + with pytest.raises(OSError, match="publication failed"): + writer.finish() + assert not list(tmp_path.glob(".model.bundle.*.tmp")) + writer.abort() + assert source.read_bytes() == b"plan" + assert destination.read_bytes() == b"previous" + assert set(tmp_path.iterdir()) == {source, destination} + + +@pytest.mark.parametrize("change", ("size", "mtime", "replacement")) +def test_changed_borrowed_source_fails_before_publishing(tmp_path: Path, change: str) -> None: + source, destination = tmp_path / "existing.plan", tmp_path / "model.bundle" + source.write_bytes(b"original") + destination.write_bytes(b"previous") + writer = BundleWriter(destination) + writer.set_header(family="family", task="text_generation", backend="trt") + writer.add_file("engine.plan", source) + before = source.stat() + if change == "size": + source.write_bytes(b"longer source content") + elif change == "mtime": + source.write_bytes(b"modified") + os.utime(source, ns=(before.st_atime_ns, before.st_mtime_ns + 1_000_000_000)) + else: + replacement = tmp_path / "replacement.plan" + replacement.write_bytes(b"modified") + os.utime(replacement, ns=(before.st_atime_ns, before.st_mtime_ns)) + os.replace(replacement, source) + with pytest.raises(RuntimeError, match="borrowed bundle section source changed"): + writer.finish() + writer.abort() + assert source.is_file() + assert destination.read_bytes() == b"previous" + assert set(tmp_path.iterdir()) == {source, destination} diff --git a/core/runtime/bundle/bundle_format.cpp b/core/runtime/bundle/bundle_format.cpp index eb11f20674..f6c48301fe 100644 --- a/core/runtime/bundle/bundle_format.cpp +++ b/core/runtime/bundle/bundle_format.cpp @@ -190,12 +190,18 @@ const BundleSectionInfo* BundleReader::find_section(std::string_view name) const return nullptr; } +std::uint64_t BundleReader::section_file_offset(std::string_view name) const { + const auto* section = find_section(name); + if (section == nullptr) + throw std::runtime_error("Bundle section not found: " + std::string(name)); + return checked_section_file_offset(*section, data_offset_, file_size_, path_); +} + std::vector BundleReader::read_section(std::string_view name) const { const auto* section = find_section(name); if (section == nullptr) throw std::runtime_error("Bundle section not found: " + std::string(name)); - const std::uint64_t file_offset = - checked_section_file_offset(*section, data_offset_, file_size_, path_); + const std::uint64_t file_offset = section_file_offset(name); if (file_offset > static_cast(std::numeric_limits::max())) throw std::runtime_error("Bundle section '" + section->name + "' has an unsupported file offset: " + path_); diff --git a/core/runtime/include/trtmc/bundle.h b/core/runtime/include/trtmc/bundle.h index 154d697ef1..eedcf2c75c 100644 --- a/core/runtime/include/trtmc/bundle.h +++ b/core/runtime/include/trtmc/bundle.h @@ -42,6 +42,10 @@ class BundleReader { const std::string& path() const noexcept { return path_; } const BundleInfo& info() const noexcept { return info_; } const BundleSectionInfo* find_section(std::string_view name) const noexcept; + // Return the validated absolute byte offset of a section in the bundle. + // Large TensorRT plans can be stream-deserialized from this range without + // first duplicating the entire section in host memory. + std::uint64_t section_file_offset(std::string_view name) const; std::vector read_section(std::string_view name) const; private: diff --git a/core/runtime/include/trtmc/runtime/family_factory.h b/core/runtime/include/trtmc/runtime/family_factory.h index c0927b50b0..0a728dd008 100644 --- a/core/runtime/include/trtmc/runtime/family_factory.h +++ b/core/runtime/include/trtmc/runtime/family_factory.h @@ -9,6 +9,7 @@ #include "trtmc/task.h" #include +#include namespace trtmc { @@ -18,6 +19,8 @@ struct FamilyContext { const BundleReader& reader; IBackend& backend; std::uint64_t kv_cache_size_bytes{0}; + std::string runtime_cache_path; + bool cuda_graphs{false}; }; using CreateFamilyFn = ITask* (*)(const FamilyContext& context); diff --git a/core/runtime/include/trtmc/runtime/trt_backend.h b/core/runtime/include/trtmc/runtime/trt_backend.h index ace743553e..d852e3b6c3 100644 --- a/core/runtime/include/trtmc/runtime/trt_backend.h +++ b/core/runtime/include/trtmc/runtime/trt_backend.h @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -23,6 +24,7 @@ struct ModuleCreateOptions { std::shared_ptr distributed_owner; // keeps communicator alive const char* runtime_cache_path{""}; // TensorRT-RTX JIT cache, optional bool cuda_graphs{false}; // TensorRT-RTX whole-graph capture + std::int32_t optimization_profile{0}; // context profile selected by the family }; struct ModuleExternalBinding { @@ -59,6 +61,37 @@ class IBackend { create_dual_profile_modules(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) = 0; + // Optional TensorRT-RTX capability for multi-GiB bundle sections. The + // backend validates the exact file range and lets TensorRT stream it + // directly instead of materializing a duplicate host buffer. + virtual std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) { + (void)plan_path; + (void)plan_offset; + (void)plan_size; + (void)options; + (void)external_bindings; + (void)weight_streaming_budget_bytes; + (void)retain_engine; + (void)serial_execution_context; + throw std::runtime_error("backend does not support file-backed TensorRT plans"); + } + + // Optional explicit lifetime for one process-shared runtime cache. A + // family releases its lease after all lazy modules have been destroyed. + virtual std::uint64_t acquire_runtime_cache_lease(const char* path) { + (void)path; + throw std::runtime_error("backend does not support runtime-cache leases"); + } + virtual void release_runtime_cache_lease(std::uint64_t lease) { + (void)lease; + throw std::runtime_error("backend does not support runtime-cache leases"); + } + // Backend identity written into the bundle header (currently "trt"). virtual const char* name() const = 0; }; diff --git a/core/runtime/include/trtmc/task.h b/core/runtime/include/trtmc/task.h index 087c9f02c5..e4e1bd1972 100644 --- a/core/runtime/include/trtmc/task.h +++ b/core/runtime/include/trtmc/task.h @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -107,6 +108,15 @@ struct AudioResult { std::vector samples; std::int32_t num_samples{0}; std::int32_t sample_rate{24000}; + // Interleaved channel count. Existing mono producers keep the default. + std::int32_t channels{1}; +}; + +// Synchronized decoded video and audio produced by an audiovisual model. +struct VideoResult { + ImageResult frames; + AudioResult audio; + std::int32_t fps{0}; }; struct TranscriptionStreamConfig { @@ -379,6 +389,67 @@ struct ImageGenerationConfig { std::string negative_prompt; std::int32_t height{0}; std::int32_t width{0}; + // Requested video frame count. Zero selects the family/bundle default. + std::int32_t video_num_frames{0}; +}; + +// Decoded host-resident inputs for native video conditioning. Image pixels are +// contiguous HWC float32 in [0, 1]; video pixels are contiguous THWC. +struct VideoImageInput { + std::vector pixels; + std::int32_t height{0}; + std::int32_t width{0}; + std::int32_t channels{3}; +}; + +struct VideoClipInput { + std::vector pixels; + std::int32_t num_frames{0}; + std::int32_t height{0}; + std::int32_t width{0}; + std::int32_t channels{3}; + std::int32_t fps_numerator{0}; + std::int32_t fps_denominator{1}; + AudioResult soundtrack; +}; + +struct ReferenceMediaDecodePolicy { + std::uint32_t maximum_duration_seconds{0}; + std::uint32_t target_video_fps{0}; + std::uint32_t maximum_source_video_fps{0}; + std::uint32_t canvas_short_edge{0}; + std::uint64_t canvas_max_pixels{0}; + std::uint32_t canvas_multiple{0}; + double minimum_aspect_ratio{0.0}; + double maximum_aspect_ratio{0.0}; +}; + +enum class VideoReferenceKind { + kImage, + kVideo, + kAudio, +}; + +struct VideoReferenceInput { + VideoReferenceKind kind{VideoReferenceKind::kImage}; + VideoImageInput image; + VideoClipInput video; + AudioResult audio; +}; + +enum class VideoGenerationMode { + kTextToVideoAudio, + kFirstLastFrameToVideoAudio, + kReferenceToVideoAudio, +}; + +struct VideoGenerationRequest { + std::string prompt; + ImageGenerationConfig config; + VideoGenerationMode mode{VideoGenerationMode::kTextToVideoAudio}; + std::optional first_frame; + std::optional last_frame; + std::vector references; }; struct WorldModelRequest { @@ -524,6 +595,19 @@ class IImageGeneration : public virtual ITask { const ImageGenerationConfig& config = {}) = 0; }; +// Optional capability for image-generation families that return synchronized +// video and audio or accept structured video-conditioning inputs. +class IVideoGeneration { + public: + virtual ~IVideoGeneration() = default; + virtual VideoResult generate_video(const std::string& prompt, + const ImageGenerationConfig& config = {}) = 0; + virtual VideoResult generate_video(const VideoGenerationRequest& request) = 0; + virtual std::optional reference_media_decode_policy() const { + return std::nullopt; + } +}; + class IImageEditing : public virtual ITask { public: static constexpr const char* kTask = "image_edit"; diff --git a/core/runtime/loader/family_loader.cpp b/core/runtime/loader/family_loader.cpp index 76c419e633..25daee3c64 100644 --- a/core/runtime/loader/family_loader.cpp +++ b/core/runtime/loader/family_loader.cpp @@ -6,10 +6,10 @@ #include "trtmc/runtime/family_loader.h" #include "runtime/bundle/bundle_format.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_factory.h" #include "trtmc/runtime/trt_backend.h" -#include #include #include #include @@ -62,12 +62,11 @@ fs::path explicit_runtime_root(const std::string& runtime_root) { class SharedLibrary { public: explicit SharedLibrary(const fs::path& path) : path_(path.string()) { - dlerror(); - handle_ = dlopen(path_.c_str(), RTLD_NOW | RTLD_LOCAL); + std::string error; + handle_ = + internal::open_dynamic_library(path, internal::DynamicLibraryVisibility::local, &error); if (handle_ == nullptr) { - const char* error = dlerror(); - throw std::runtime_error("Unable to load '" + path_ + - "': " + (error != nullptr ? error : "unknown dlopen error")); + throw std::runtime_error("Unable to load '" + path_ + "': " + error); } } @@ -76,29 +75,29 @@ class SharedLibrary { ~SharedLibrary() { if (handle_ != nullptr) - dlclose(handle_); + (void)internal::close_dynamic_library(handle_); } void* require_symbol(const char* name) const { - dlerror(); - void* symbol = dlsym(handle_, name); - const char* error = dlerror(); - if (error != nullptr || symbol == nullptr) { + std::string error; + void* symbol = internal::dynamic_library_symbol(handle_, name, &error); + if (symbol == nullptr) { throw std::runtime_error("Library '" + path_ + "' is missing required symbol '" + name + - "'"); + "': " + error); } return symbol; } private: std::string path_; - void* handle_{nullptr}; + internal::DynamicLibraryHandle handle_{nullptr}; }; class BackendLibrary { public: BackendLibrary(const fs::path& runtime_root, const std::string& backend_id) - : library_(runtime_root / ("libtrtmc_backend_" + backend_id + ".so")) { + : library_(runtime_root / + internal::dynamic_library_filename("trtmc_backend_" + backend_id)) { const auto create = reinterpret_cast(library_.require_symbol("trtmc_create_backend")); destroy_ = @@ -136,7 +135,7 @@ class BackendLibrary { class FamilyLibrary { public: FamilyLibrary(const fs::path& runtime_root, const std::string& family_id) - : library_(runtime_root / ("libtrtmc_model_" + family_id + ".so")), + : library_(runtime_root / internal::dynamic_library_filename("trtmc_model_" + family_id)), create_(reinterpret_cast(library_.require_symbol(kCreateFamilySymbol))) {} FamilyLibrary(const FamilyLibrary&) = delete; @@ -175,6 +174,25 @@ class RuntimeOptionsBackend final : public IBackend { with_runtime_options(options)); } + std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) override { + return backend_.create_module_from_file( + plan_path, plan_offset, plan_size, with_runtime_options(options), bindings, + weight_streaming_budget_bytes, retain_engine, serial_execution_context); + } + + std::uint64_t acquire_runtime_cache_lease(const char* path) override { + return backend_.acquire_runtime_cache_lease(path); + } + + void release_runtime_cache_lease(std::uint64_t lease) override { + backend_.release_runtime_cache_lease(lease); + } + const char* name() const override { return backend_.name(); } private: @@ -228,7 +246,8 @@ RuntimeLibraryCache& runtime_library_cache() { } IBackend& cached_backend(const fs::path& runtime_root, const std::string& backend_id) { - const std::string path = (runtime_root / ("libtrtmc_backend_" + backend_id + ".so")).string(); + const std::string path = + (runtime_root / internal::dynamic_library_filename("trtmc_backend_" + backend_id)).string(); auto& cache = runtime_library_cache(); std::lock_guard lock(cache.mutex); const auto found = cache.backends.find(path); @@ -258,7 +277,8 @@ IBackend& cached_configured_backend(IBackend& backend, const std::string& runtim } FamilyLibrary& cached_family(const fs::path& runtime_root, const std::string& family_id) { - const std::string path = (runtime_root / ("libtrtmc_model_" + family_id + ".so")).string(); + const std::string path = + (runtime_root / internal::dynamic_library_filename("trtmc_model_" + family_id)).string(); auto& cache = runtime_library_cache(); std::lock_guard lock(cache.mutex); const auto found = cache.families.find(path); @@ -301,7 +321,8 @@ std::unique_ptr load_task(const std::string& bundle_path, const std::stri FamilyLibrary& family = cached_family(root, info.family); IBackend& configured_backend = cached_configured_backend(backend, runtime_cache_path, cuda_graphs); - FamilyContext context{reader, configured_backend, kv_cache_size_bytes}; + FamilyContext context{reader, configured_backend, kv_cache_size_bytes, runtime_cache_path, + cuda_graphs}; std::unique_ptr task(family.create(context)); if (task == nullptr) throw std::runtime_error("trtmc_create_family returned nullptr"); diff --git a/core/runtime/platform/dynamic_library.cpp b/core/runtime/platform/dynamic_library.cpp new file mode 100644 index 0000000000..1b0b7df552 --- /dev/null +++ b/core/runtime/platform/dynamic_library.cpp @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/platform/dynamic_library.h" + +#include +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + +namespace trtmc::internal { +namespace { + +namespace fs = std::filesystem; + +void assign_error(std::string* output, std::string message) noexcept { + if (output == nullptr) + return; + try { + *output = std::move(message); + } catch (...) { + try { + output->clear(); + } catch (...) { + } + } +} + +#if defined(_WIN32) + +std::string utf8_from_wide(const wchar_t* value, int length) { + if (value == nullptr || length <= 0) + return {}; + const int required = + WideCharToMultiByte(CP_UTF8, 0, value, length, nullptr, 0, nullptr, nullptr); + if (required <= 0) + return {}; + std::string result(static_cast(required), '\0'); + if (WideCharToMultiByte(CP_UTF8, 0, value, length, result.data(), required, nullptr, nullptr) <= + 0) { + return {}; + } + return result; +} + +std::string windows_error_message(DWORD code) { + wchar_t* buffer = nullptr; + const DWORD length = FormatMessageW( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, + nullptr, code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + reinterpret_cast(&buffer), 0, nullptr); + std::string message = utf8_from_wide(buffer, static_cast(length)); + if (buffer != nullptr) + LocalFree(buffer); + while (!message.empty() && + (message.back() == '\r' || message.back() == '\n' || message.back() == ' ')) { + message.pop_back(); + } + return message.empty() ? "Windows error " + std::to_string(code) : message; +} + +#else + +std::string loader_error() { + const char* error = dlerror(); + return error == nullptr ? std::string("unknown dynamic-loader error") : std::string(error); +} + +#endif + +} // namespace + +DynamicLibraryHandle open_dynamic_library(const fs::path& path, DynamicLibraryVisibility visibility, + std::string* error) { + if (error != nullptr) + error->clear(); +#if defined(_WIN32) + (void)visibility; + try { + fs::path load_path = path; + DWORD flags = LOAD_LIBRARY_SEARCH_DEFAULT_DIRS; + if (path.is_absolute() || path.has_parent_path()) { + std::error_code ec; + const fs::path absolute = fs::absolute(path, ec); + if (!ec) + load_path = absolute; + flags |= LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR; + } + HMODULE handle = LoadLibraryExW(load_path.c_str(), nullptr, flags); + if (handle == nullptr) { + assign_error(error, windows_error_message(GetLastError())); + return nullptr; + } + return reinterpret_cast(handle); + } catch (const std::exception& exception) { + assign_error(error, exception.what()); + return nullptr; + } +#else + dlerror(); + const int flags = + RTLD_NOW | (visibility == DynamicLibraryVisibility::global ? RTLD_GLOBAL : RTLD_LOCAL); + DynamicLibraryHandle handle = dlopen(path.c_str(), flags); + if (handle == nullptr) + assign_error(error, loader_error()); + return handle; +#endif +} + +void* dynamic_library_symbol(DynamicLibraryHandle handle, const char* name, std::string* error) { + if (error != nullptr) + error->clear(); + if (handle == nullptr || name == nullptr) { + assign_error(error, "invalid dynamic-library handle or symbol name"); + return nullptr; + } +#if defined(_WIN32) + FARPROC symbol = GetProcAddress(reinterpret_cast(handle), name); + if (symbol == nullptr) { + assign_error(error, windows_error_message(GetLastError())); + return nullptr; + } + return reinterpret_cast(symbol); +#else + dlerror(); + void* symbol = dlsym(handle, name); + const char* message = dlerror(); + if (message != nullptr) { + assign_error(error, message); + return nullptr; + } + return symbol; +#endif +} + +bool close_dynamic_library(DynamicLibraryHandle handle, std::string* error) noexcept { + if (error != nullptr) + error->clear(); + if (handle == nullptr) + return true; +#if defined(_WIN32) + if (FreeLibrary(reinterpret_cast(handle)) != 0) + return true; + try { + assign_error(error, windows_error_message(GetLastError())); + } catch (...) { + assign_error(error, "unknown Windows dynamic-loader error"); + } + return false; +#else + if (dlclose(handle) == 0) + return true; + try { + assign_error(error, loader_error()); + } catch (...) { + assign_error(error, "unknown dynamic-loader error"); + } + return false; +#endif +} + +std::string dynamic_library_filename(std::string_view stem) { +#if defined(_WIN32) + return std::string(stem) + ".dll"; +#elif defined(__APPLE__) + return "lib" + std::string(stem) + ".dylib"; +#else + return "lib" + std::string(stem) + ".so"; +#endif +} + +} // namespace trtmc::internal diff --git a/core/runtime/platform/dynamic_library.h b/core/runtime/platform/dynamic_library.h new file mode 100644 index 0000000000..53c8455a00 --- /dev/null +++ b/core/runtime/platform/dynamic_library.h @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace trtmc::internal { + +using DynamicLibraryHandle = void*; + +enum class DynamicLibraryVisibility { + local, + global, +}; + +DynamicLibraryHandle +open_dynamic_library(const std::filesystem::path& path, + DynamicLibraryVisibility visibility = DynamicLibraryVisibility::local, + std::string* error = nullptr); +void* dynamic_library_symbol(DynamicLibraryHandle handle, const char* name, + std::string* error = nullptr); +bool close_dynamic_library(DynamicLibraryHandle handle, std::string* error = nullptr) noexcept; + +std::string dynamic_library_filename(std::string_view stem); + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/rtx_backend.cpp b/core/runtime/tensorrt/rtx_backend.cpp index 5bb004cd14..7da89beb1f 100644 --- a/core/runtime/tensorrt/rtx_backend.cpp +++ b/core/runtime/tensorrt/rtx_backend.cpp @@ -4,122 +4,709 @@ */ // RtxBackend: IBackend implementation for TensorRT-RTX. -// Compiled into libtrtmc_backend_trt_rtx.so. Links libtensorrt_rtx.so. +// Compiled into libtrtmc_backend_rtx.so. Links libtensorrt_rtx.so. +// +// Uses the RTX-specific NvInfer.h headers which declare IRuntimeCache, +// CudaGraphStrategy, and DynamicShapesKernelSpecializationStrategy. #include "runtime/primitives/cuda_common.h" +#include "runtime/tensorrt/runtime_cache_persistence.h" #include "runtime/tensorrt/trt_logger.h" #include "trt_module_impl.h" #include "trtmc/runtime/trt_backend.h" #include +#include +#include +#include +#include #include +#include #include +#include #include #include #include #include +#include +#include #include #include +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + namespace trtmc { namespace { +namespace fs = std::filesystem; + struct StreamSetup { cudaStream_t stream{nullptr}; std::shared_ptr owner; }; +#if defined(_WIN32) +// Some Windows CUDA configurations report memory-pool support while rejecting +// the asynchronous allocation path used by TensorRT-RTX. IGpuAsyncAllocator +// permits a synchronizing implementation, so retain the callback surface while +// backing it with cudaMalloc/cudaFree on Windows only. +class SynchronousGpuAllocator final : public nvinfer1::IGpuAsyncAllocator { + public: + void* allocateAsync(std::uint64_t size, std::uint64_t /*alignment*/, + nvinfer1::AllocatorFlags /*flags*/, + cudaStream_t /*stream*/) noexcept override { + if (size == 0 || size > static_cast(std::numeric_limits::max())) + return nullptr; + void* memory = nullptr; + return cudaMalloc(&memory, static_cast(size)) == cudaSuccess ? memory + : nullptr; + } + + bool deallocateAsync(void* memory, cudaStream_t /*stream*/) noexcept override { + return memory == nullptr || cudaFree(memory) == cudaSuccess; + } +}; +#endif + StreamSetup resolve_stream(cudaStream_t requested_stream) { - if (requested_stream) - return {requested_stream, {}}; + if (requested_stream) { + return StreamSetup{requested_stream, {}}; + } auto owned = std::make_shared(); - if (!owned->ok()) + if (!owned->ok()) { throw std::runtime_error("[trtmc] Failed to create CUDA stream"); - return {owned->get(), owned}; + } + + return StreamSetup{owned->get(), owned}; } -std::shared_ptr deserialize_engine(nvinfer1::IRuntime& runtime, - const void* plan_data, size_t plan_size) { - auto* engine = runtime.deserializeCudaEngine(plan_data, plan_size); - if (!engine) - throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); - return {engine, [](nvinfer1::ICudaEngine* value) { delete value; }}; +void append_plan_cache_identity_field(std::string& identity, const std::string& value) { + identity += std::to_string(value.size()); + identity.push_back(':'); + identity += value; } -class RuntimeCacheState { +class PlanFileMutationGuard final { public: - RuntimeCacheState(nvinfer1::IRuntimeConfig& config, std::string path) - : cache_(config.createRuntimeCache()), path_(std::move(path)) { - if (cache_ == nullptr) - throw std::runtime_error("[trtmc] Failed to create TensorRT-RTX runtime cache"); + explicit PlanFileMutationGuard(const char* path) { + if (path == nullptr || path[0] == '\0') + throw std::invalid_argument("[trtmc] RTX plan file path must not be empty"); + try { + path_ = fs::canonical(fs::path(path)); + } catch (const fs::filesystem_error& error) { + throw std::runtime_error(std::string("[trtmc] Failed to resolve RTX plan file: ") + + error.what()); + } +#if defined(_WIN32) + // Deny write/delete sharing while the exact file is streamed and deserialized. + handle_ = CreateFileW(path_.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OPEN_REPARSE_POINT, nullptr); + if (handle_ == INVALID_HANDLE_VALUE) { + const std::error_code error(static_cast(GetLastError()), std::system_category()); + throw std::runtime_error("[trtmc] Failed to lock RTX plan file: " + error.message()); + } + BY_HANDLE_FILE_INFORMATION info{}; + if (!GetFileInformationByHandle(handle_, &info)) { + const std::error_code error(static_cast(GetLastError()), std::system_category()); + (void)CloseHandle(handle_); + handle_ = INVALID_HANDLE_VALUE; + throw std::runtime_error("[trtmc] Failed to identify RTX plan file: " + + error.message()); + } + const auto path_utf8 = path_.u8string(); + append_plan_cache_identity_field(file_identity_, path_utf8); + append_plan_cache_identity_field( + file_identity_, std::to_string(static_cast(info.dwVolumeSerialNumber))); + append_plan_cache_identity_field( + file_identity_, + std::to_string((static_cast(info.nFileIndexHigh) << 32U) | + static_cast(info.nFileIndexLow))); + append_plan_cache_identity_field( + file_identity_, std::to_string((static_cast(info.nFileSizeHigh) << 32U) | + static_cast(info.nFileSizeLow))); + append_plan_cache_identity_field( + file_identity_, + std::to_string( + (static_cast(info.ftLastWriteTime.dwHighDateTime) << 32U) | + static_cast(info.ftLastWriteTime.dwLowDateTime))); +#else + struct stat info{}; + if (::stat(path_.c_str(), &info) != 0) { + const std::error_code error(errno, std::generic_category()); + throw std::runtime_error("[trtmc] Failed to identify RTX plan file: " + + error.message()); + } + initial_device_ = static_cast(info.st_dev); + initial_inode_ = static_cast(info.st_ino); + initial_size_ = fs::file_size(path_); + initial_write_time_ = fs::last_write_time(path_); + const auto path_utf8 = path_.u8string(); + append_plan_cache_identity_field(file_identity_, path_utf8); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_device_)); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_inode_)); + append_plan_cache_identity_field(file_identity_, std::to_string(initial_size_)); + append_plan_cache_identity_field( + file_identity_, std::to_string(initial_write_time_.time_since_epoch().count())); +#endif + } + + ~PlanFileMutationGuard() { +#if defined(_WIN32) + if (handle_ != INVALID_HANDLE_VALUE) + (void)CloseHandle(handle_); +#endif + } + + PlanFileMutationGuard(const PlanFileMutationGuard&) = delete; + PlanFileMutationGuard& operator=(const PlanFileMutationGuard&) = delete; + + const fs::path& path() const noexcept { return path_; } + const std::string& cache_identity() const noexcept { return file_identity_; } + + void verify_unchanged() const { +#if !defined(_WIN32) + struct stat info{}; + if (::stat(path_.c_str(), &info) != 0) { + throw std::runtime_error("[trtmc] RTX plan file disappeared while it was parsed"); + } + if (static_cast(info.st_dev) != initial_device_ || + static_cast(info.st_ino) != initial_inode_ || + fs::file_size(path_) != initial_size_ || + fs::last_write_time(path_) != initial_write_time_) { + throw std::runtime_error("[trtmc] RTX plan file changed while it was parsed"); + } +#endif + } + + private: + fs::path path_; + std::string file_identity_; +#if defined(_WIN32) + HANDLE handle_{INVALID_HANDLE_VALUE}; +#else + std::uintmax_t initial_device_{0}; + std::uintmax_t initial_inode_{0}; + std::uintmax_t initial_size_{0}; + fs::file_time_type initial_write_time_{}; +#endif +}; + +void validate_plan_description(std::uint64_t size) { + if (size == 0 || size > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("[trtmc] RTX plan file range is invalid"); + } +} + +void validate_plan_file_range(std::uint64_t offset, std::uint64_t size, std::uint64_t file_size) { + const auto max_offset = static_cast(std::numeric_limits::max()); + if (offset > file_size || size > file_size - offset || offset > max_offset || + size > max_offset - offset) { + throw std::invalid_argument("[trtmc] RTX plan section is outside its file"); + } +} + +enum class ReaderDestination { + host, + device, + invalid, +}; + +ReaderDestination classify_reader_destination(void* destination) noexcept { + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, destination); + if (status == cudaErrorInvalidValue) { + (void)cudaGetLastError(); + return ReaderDestination::host; + } + if (status != cudaSuccess) + return ReaderDestination::invalid; + if (attributes.type == cudaMemoryTypeUnregistered || attributes.type == cudaMemoryTypeHost) + return ReaderDestination::host; + if (attributes.type == cudaMemoryTypeDevice || attributes.type == cudaMemoryTypeManaged) + return ReaderDestination::device; + return ReaderDestination::invalid; +} + +void read_host_bytes(std::ifstream& file, void* destination, std::uint64_t requested) { + file.read(static_cast(destination), static_cast(requested)); + if (file.gcount() != static_cast(requested)) + throw std::runtime_error("[trtmc] Failed to read RTX plan bytes into host memory"); +} + +void read_device_bytes(std::ifstream& file, std::vector& staging, void* destination, + std::uint64_t requested) { + std::uint64_t copied = 0; + while (copied < requested) { + const auto chunk = std::min(staging.size(), requested - copied); + file.read(staging.data(), static_cast(chunk)); + if (file.gcount() != static_cast(chunk)) + throw std::runtime_error("[trtmc] Failed to stage RTX plan bytes"); + if (cudaMemcpy(static_cast(destination) + copied, staging.data(), + static_cast(chunk), cudaMemcpyHostToDevice) != cudaSuccess) { + throw std::runtime_error("[trtmc] Failed to copy RTX plan bytes to the device"); + } + copied += chunk; + } +} + +void read_destination_bytes(std::ifstream& file, std::vector& staging, void* destination, + std::uint64_t requested) { + switch (classify_reader_destination(destination)) { + case ReaderDestination::host: + read_host_bytes(file, destination, requested); + return; + case ReaderDestination::device: + read_device_bytes(file, staging, destination, requested); + return; + case ReaderDestination::invalid: + throw std::runtime_error("[trtmc] Unsupported RTX plan reader destination"); + } +} - std::ifstream input(path_, std::ios::binary | std::ios::ate); - if (!input) +// TensorRT-RTX may request plan bytes into host or device memory and may seek +// within the stream. Restrict every operation to the validated bundle section. +class BoundedPlanStreamReader final : public nvinfer1::IStreamReaderV2 { + public: + BoundedPlanStreamReader(const char* path, std::uint64_t offset, std::uint64_t size) + : mutation_guard_(path), offset_(offset), size_(size), staging_(4U << 20) { + validate_plan_description(size_); + file_.open(mutation_guard_.path(), std::ios::binary | std::ios::ate); + if (!file_) + throw std::runtime_error("[trtmc] Failed to open RTX plan file"); + const auto end = file_.tellg(); + if (end < 0) + throw std::runtime_error("[trtmc] Failed to determine RTX plan file size"); + const auto file_size = static_cast(static_cast(end)); + validate_plan_file_range(offset_, size_, file_size); + } + + void verify_unchanged() const { mutation_guard_.verify_unchanged(); } + + std::string cache_identity() const { + std::string identity = mutation_guard_.cache_identity(); + append_plan_cache_identity_field(identity, std::to_string(offset_)); + append_plan_cache_identity_field(identity, std::to_string(size_)); + return identity; + } + + int64_t read(void* destination, int64_t nb_bytes, cudaStream_t stream) noexcept override { + (void)stream; + if (destination == nullptr || nb_bytes < 0) + return -1; + if (nb_bytes == 0) + return 0; + try { + return read_from_current_position(destination, nb_bytes); + } catch (...) { + return -1; + } + } + + bool seek(int64_t offset, nvinfer1::SeekPosition where) noexcept override { + std::uint64_t base = 0; + switch (where) { + case nvinfer1::SeekPosition::kSET: + break; + case nvinfer1::SeekPosition::kCUR: + base = cursor_; + break; + case nvinfer1::SeekPosition::kEND: + base = size_; + break; + default: + return false; + } + if (offset >= 0) { + const auto delta = static_cast(offset); + if (base > size_ || delta > size_ - base) + return false; + cursor_ = base + delta; + return true; + } + const auto magnitude = static_cast(-(offset + 1)) + 1; + if (magnitude > base) + return false; + cursor_ = base - magnitude; + return true; + } + + private: + int64_t read_from_current_position(void* destination, int64_t nb_bytes) { + const auto requested = + std::min(static_cast(nb_bytes), size_ - cursor_); + if (requested == 0) + return 0; + file_.clear(); + file_.seekg(static_cast(offset_ + cursor_), std::ios::beg); + if (!file_) + throw std::runtime_error("[trtmc] Failed to seek in the RTX plan section"); + read_destination_bytes(file_, staging_, destination, requested); + cursor_ += requested; + return static_cast(requested); + } + + PlanFileMutationGuard mutation_guard_; + std::ifstream file_; + std::uint64_t offset_{0}; + std::uint64_t size_{0}; + std::uint64_t cursor_{0}; + std::vector staging_; +}; + +void apply_weight_streaming_budget(nvinfer1::ICudaEngine& engine, std::int64_t budget, + bool cuda_graphs) { + if (budget < 0) + return; + if (cuda_graphs) + throw std::invalid_argument( + "[trtmc] RTX weight streaming is incompatible with CUDA graph capture"); + const std::int64_t streamable = engine.getStreamableWeightsSize(); + if (streamable <= 0) + throw std::runtime_error("[trtmc] Engine was not built with TensorRT-RTX weight streaming"); + const std::int64_t applied = std::min(budget, streamable); + if (!engine.setWeightStreamingBudgetV2(applied) || + engine.getWeightStreamingBudgetV2() != applied) { + throw std::runtime_error("[trtmc] TensorRT-RTX rejected the weight streaming budget"); + } + std::cerr << "[trtmc.rtx_weight_budget] requested_bytes=" << budget + << " streamable_bytes=" << streamable << " applied_bytes=" << applied + << " streaming_scratch_bytes=" << engine.getWeightStreamingScratchMemorySize() + << '\n'; +} + +void validate_optimization_profile(const nvinfer1::ICudaEngine& engine, int32_t profile) { + if (profile < 0 || profile >= engine.getNbOptimizationProfiles()) + throw std::invalid_argument("[trtmc] Invalid optimization profile index"); +} + +void configure_cuda_graphs(nvinfer1::IRuntimeConfig& config, bool enabled) { + if (enabled && + !config.setCudaGraphStrategy(nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE)) { + throw std::runtime_error("[trtmc] Failed to enable RTX CUDA graph capture"); + } +} + +void validate_serial_execution_context(const ModuleCreateOptions& options, + bool serial_execution_context) { + if (!serial_execution_context) + return; + if (options.stream == nullptr) { + throw std::invalid_argument( + "[trtmc] Shared RTX activation memory requires an explicit CUDA stream"); + } + if (options.cuda_graphs) { + throw std::invalid_argument( + "[trtmc] Shared RTX activation memory is incompatible with CUDA graph capture"); + } +} + +struct ActivationArenaKey { + int device{-1}; + cudaStream_t stream{nullptr}; + + bool operator==(const ActivationArenaKey& other) const noexcept { + return device == other.device && stream == other.stream; + } +}; + +struct ActivationArenaKeyHash { + std::size_t operator()(const ActivationArenaKey& key) const noexcept { + const auto stream = reinterpret_cast(key.stream); + return std::hash{}(stream) ^ + (std::hash{}(key.device) + 0x9e3779b9U + (stream << 6U) + (stream >> 2U)); + } +}; + +// One arena is shared only by explicitly marked contexts on one CUDA stream. +// Context creation records each profile upper bound. Allocation is deferred to +// the first enqueue, after that context has received its request shapes. Each +// serial context is shape-sized once during its lifetime so repeated forwards +// avoid repartitioning TensorRT's internal activation memory. +class RtxActivationArena final : public ITrtActivationArena { + public: + RtxActivationArena(int device, cudaStream_t stream) : device_(device), stream_(stream) {} + + ~RtxActivationArena() override { + if (memory_ == nullptr) return; - const auto end = input.tellg(); - if (end <= 0) + int previous_device = -1; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess && + previous_device >= 0 && previous_device != device_; + if (restore_device && cudaSetDevice(device_) != cudaSuccess) return; - std::vector bytes(static_cast(end)); - input.seekg(0); - input.read(bytes.data(), static_cast(bytes.size())); - if (!input) - throw std::runtime_error("[trtmc] Failed to read TensorRT-RTX runtime cache"); - if (!cache_->deserialize(bytes.data(), bytes.size())) - throw std::runtime_error("[trtmc] TensorRT-RTX rejected the runtime cache"); - } - - RuntimeCacheState(const RuntimeCacheState&) = delete; - RuntimeCacheState& operator=(const RuntimeCacheState&) = delete; - - ~RuntimeCacheState() { - auto* serialized = cache_ != nullptr ? cache_->serialize() : nullptr; - if (serialized != nullptr && serialized->size() > 0) { - std::ofstream output(path_, std::ios::binary | std::ios::trunc); - if (output) { - output.write(static_cast(serialized->data()), - static_cast(serialized->size())); - } else { - std::cerr << "[trtmc] Failed to write TensorRT-RTX runtime cache: " << path_ - << '\n'; + (void)cudaStreamSynchronize(stream_); + if (memory_) + (void)cudaFree(memory_); + if (restore_device) + (void)cudaSetDevice(previous_device); + } + + void attach(nvinfer1::IExecutionContext* context, cudaStream_t stream, + std::int64_t required_bytes) override { + if (context == nullptr) + throw std::invalid_argument("[trtmc] Cannot attach a null RTX execution context"); + if (stream == nullptr || stream != stream_) + throw std::invalid_argument("[trtmc] RTX activation arena stream mismatch"); + if (required_bytes < 0 || + static_cast(required_bytes) > + static_cast(std::numeric_limits::max())) { + throw std::overflow_error("[trtmc] RTX activation memory requirement is invalid"); + } + require_current_device(); + + std::lock_guard lock(mutex_); + if (contexts_.count(context) != 0) + throw std::logic_error("[trtmc] RTX execution context is already attached to an arena"); + contexts_.emplace(context, required_bytes); + required_bytes_ = std::max(required_bytes_, required_bytes); + } + + void begin_enqueue(nvinfer1::IExecutionContext* context) override { + mutex_.lock(); + try { + require_current_device(); + const auto found = contexts_.find(context); + if (found == contexts_.end()) { + throw std::logic_error( + "[trtmc] RTX execution context is not attached to its activation arena"); } + auto shaped = shaped_requirements_.find(context); + if (shaped == shaped_requirements_.end()) { + const std::int64_t requested = shaped_requirement(*context, found->second); + shaped = shaped_requirements_.emplace(context, requested).first; + } + ensure_capacity_locked(shaped->second); + // setDeviceMemoryV2 has no status return. Rebind on every enqueue + // because a later serial context may have raised the arena's + // high-water mark and replaced its allocation. + context->setDeviceMemoryV2(memory_, capacity_bytes_); + } catch (...) { + mutex_.unlock(); + throw; } - delete serialized; - delete cache_; } - nvinfer1::IRuntimeCache& cache() const { return *cache_; } - const std::string& path() const { return path_; } + void end_enqueue() noexcept override { mutex_.unlock(); } + + void detach(nvinfer1::IExecutionContext* context) noexcept override { + if (context == nullptr) + return; + std::lock_guard lock(mutex_); + const auto found = contexts_.find(context); + if (found == contexts_.end()) + return; + // TensorRT permits the activation buffer and execution context to be + // released only after the enqueued work completes. + const cudaError_t status = cudaStreamSynchronize(stream_); + if (status != cudaSuccess) { + std::cerr << "[trtmc] Failed to synchronize shared RTX activation memory during " + "context teardown: " + << cudaGetErrorString(status) << '\n'; + } + contexts_.erase(found); + shaped_requirements_.erase(context); + if (memory_ == nullptr) { + required_bytes_ = 0; + for (const auto& [attached_context, attached_bytes] : contexts_) { + (void)attached_context; + required_bytes_ = std::max(required_bytes_, attached_bytes); + } + } + } private: - nvinfer1::IRuntimeCache* cache_{nullptr}; - std::string path_; + int device_{-1}; + cudaStream_t stream_{nullptr}; + std::mutex mutex_; + std::unordered_map contexts_; + std::unordered_map shaped_requirements_; + void* memory_{nullptr}; + std::int64_t capacity_bytes_{0}; + std::int64_t required_bytes_{0}; + + void require_current_device() const { + int current_device = -1; + const cudaError_t status = cudaGetDevice(¤t_device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } + if (current_device != device_) + throw std::runtime_error("[trtmc] RTX activation arena CUDA device mismatch"); + } + + static std::int64_t shaped_requirement(nvinfer1::IExecutionContext& context, + std::int64_t profile_upper_bound) { + const std::size_t shaped_bytes = context.updateDeviceMemorySizeForShapes(); + // Zero reports an error (including incomplete shapes). Preserve the + // validated profile bound as a correctness-first fallback. + if (shaped_bytes == 0) + return profile_upper_bound; + if (shaped_bytes > static_cast(profile_upper_bound)) { + throw std::runtime_error( + "[trtmc] RTX shape-sized activation memory exceeds the profile bound"); + } + return static_cast(shaped_bytes); + } + + void ensure_capacity_locked(std::int64_t requested_bytes) { + if (requested_bytes <= capacity_bytes_) + return; + if (memory_ != nullptr) { + // Serial work using the old high-water allocation is ordered on + // this stream. Synchronize once before replacing it for a larger + // context encountered later in the same request. + const cudaError_t synchronize_status = cudaStreamSynchronize(stream_); + if (synchronize_status != cudaSuccess) { + throw std::runtime_error( + std::string("[trtmc] Failed to synchronize before growing the RTX activation " + "arena: ") + + cudaGetErrorString(synchronize_status)); + } + const cudaError_t free_status = cudaFree(memory_); + if (free_status != cudaSuccess) { + throw std::runtime_error( + std::string("[trtmc] Failed to release the previous RTX activation arena: ") + + cudaGetErrorString(free_status)); + } + memory_ = nullptr; + capacity_bytes_ = 0; + } + + const cudaError_t allocation_status = + cudaMalloc(&memory_, static_cast(requested_bytes)); + if (allocation_status != cudaSuccess) { + memory_ = nullptr; + throw std::runtime_error( + std::string("[trtmc] Failed to allocate shared RTX activation memory (") + + std::to_string(requested_bytes) + + " bytes): " + cudaGetErrorString(allocation_status)); + } + capacity_bytes_ = requested_bytes; + std::cerr << "[trtmc.rtx_activation_arena] device=" << device_ + << " contexts=" << contexts_.size() << " capacity_bytes=" << capacity_bytes_ + << " profile_capacity_bytes=" << required_bytes_ << '\n'; + } }; -void keep_backend_resources(ITrtModule& module, - const std::shared_ptr& engine, - const StreamSetup& stream_setup, - const std::shared_ptr& distributed_owner) { - module.keep_alive(engine); - if (stream_setup.owner) - module.keep_alive(stream_setup.owner); - if (distributed_owner) - module.keep_alive(distributed_owner); +std::ifstream open_runtime_cache_file(const char* path) { + std::error_code exists_error; + const bool cache_exists = fs::exists(path, exists_error); + if (exists_error) { + throw std::system_error(exists_error, + "failed to inspect RTX runtime cache " + std::string(path)); + } + + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (cache_exists && !file) { + throw std::runtime_error("[trtmc] Failed to open existing RTX runtime cache: " + + std::string(path)); + } + return file; +} + +std::streamoff runtime_cache_file_size(std::ifstream& file, const char* path) { + const std::streamoff size = file.tellg(); + if (size < 0) { + throw std::runtime_error("[trtmc] Failed to size RTX runtime cache: " + std::string(path)); + } + if (static_cast(size) > + static_cast(std::numeric_limits::max()) || + size > std::numeric_limits::max()) { + throw std::overflow_error("[trtmc] RTX runtime cache is too large to load: " + + std::string(path)); + } + return size; +} + +std::vector read_runtime_cache_bytes(std::ifstream& file, std::streamoff size, + const char* path) { + std::vector bytes(static_cast(size)); + file.seekg(0, std::ios::beg); + if (!file) { + throw std::runtime_error("[trtmc] Failed to seek RTX runtime cache: " + std::string(path)); + } + file.read(bytes.data(), static_cast(size)); + if (!file || file.gcount() != static_cast(size)) { + throw std::runtime_error("[trtmc] Failed to read RTX runtime cache: " + std::string(path)); + } + return bytes; +} + +void load_runtime_cache(nvinfer1::IRuntimeCache& runtime_cache, const char* path) { + std::ifstream file = open_runtime_cache_file(path); + if (!file) + return; + + const std::streamoff size = runtime_cache_file_size(file, path); + if (size == 0) + return; + + const std::vector bytes = read_runtime_cache_bytes(file, size, path); + if (!runtime_cache.deserialize(bytes.data(), bytes.size())) { + throw std::runtime_error("[trtmc] TensorRT-RTX rejected the serialized runtime cache: " + + std::string(path)); + } + std::cerr << "[trtmc] RTX runtime cache loaded: " << path << " (" << size << " bytes)\n"; +} + +std::unique_ptr create_runtime_cache(nvinfer1::IRuntimeConfig& config, + const char* path) { + std::unique_ptr runtime_cache(config.createRuntimeCache()); + if (runtime_cache == nullptr) + throw std::runtime_error("[trtmc] Failed to create RTX runtime cache"); + load_runtime_cache(*runtime_cache, path); + return runtime_cache; } } // namespace class RtxBackend final : public IBackend { public: - RtxBackend() : runtime_(create_trt_runtime()) { + RtxBackend() + : runtime_(create_trt_runtime()) +#if defined(_WIN32) + , + staged_runtime_(create_trt_runtime()) +#endif + { if (!runtime_) throw std::runtime_error("[trtmc] Failed to create TRT-RTX runtime"); +#if defined(_WIN32) + if (!staged_runtime_) + throw std::runtime_error("[trtmc] Failed to create staged TRT-RTX runtime"); + // This dedicated runtime serves only file-backed staged plans, so the + // synchronizing allocator cannot change ordinary in-memory RTX paths. + staged_runtime_->setGpuAllocator(&gpu_allocator_); +#endif + } + + ~RtxBackend() override { + if (!runtime_cache_leases_.empty()) { + std::cerr << "[trtmc] RTX runtime cache not persisted: " << runtime_cache_leases_.size() + << " pipeline lease(s) are still active\n"; + } + delete runtime_cache_; } std::unique_ptr create_module(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) override { - return create_module_impl(plan_data, plan_size, options, {}); + auto* engine = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + return create_single_module(engine, options, {}, -1); } std::unique_ptr @@ -128,91 +715,297 @@ class RtxBackend final : public IBackend { const std::vector& external_bindings) override { if (external_bindings.empty()) throw std::invalid_argument("[trtmc] External I/O prebindings must not be empty"); - return create_module_impl(plan_data, plan_size, options, external_bindings); + auto* engine = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + return create_single_module(engine, options, external_bindings, -1); + } + + std::unique_ptr + create_module_from_file(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) override { + return create_module_from_file_impl(plan_path, plan_offset, plan_size, options, + external_bindings, weight_streaming_budget_bytes, + retain_engine, serial_execution_context); } BackendDualProfileModules create_dual_profile_modules(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options) override { - auto engine = deserialize_engine(*runtime_, plan_data, plan_size); - const auto stream_setup = resolve_stream(options.stream); - if (engine->getNbOptimizationProfiles() < 2) - throw std::runtime_error("[trtmc] Dual-profile engine requires two profiles"); + auto* engine_raw = runtime_->deserializeCudaEngine(plan_data, plan_size); + if (!engine_raw) + throw std::runtime_error("[trtmc] Failed to deserialize engine (RTX)"); + std::shared_ptr engine(engine_raw, + [](nvinfer1::ICudaEngine* p) { delete p; }); + + StreamSetup stream_setup = resolve_stream(options.stream); - BackendDualProfileModules modules; - modules.prefill = create_context_module(engine, stream_setup, options, 0, {}); - modules.decode = create_context_module(engine, stream_setup, options, 1, {}); - return modules; + const int32_t nprofiles = engine->getNbOptimizationProfiles(); + if (nprofiles < 2) + throw std::runtime_error("[trtmc] Dual-profile modules require at least two profiles"); + auto make_ctx_module = [&](int32_t profile_idx) -> std::unique_ptr { + auto profile_options = options; + profile_options.optimization_profile = profile_idx; + auto config = create_runtime_config(*engine, profile_options, false); + return create_execution_module(engine, config, stream_setup, profile_options, {}, false, + false); + }; + + BackendDualProfileModules out; + out.prefill = make_ctx_module(0); + out.decode = make_ctx_module(1); + return out; } const char* name() const override { return "trt_rtx"; } + std::uint64_t acquire_runtime_cache_lease(const char* path) override { + std::lock_guard lock(runtime_cache_mutex_); + return runtime_cache_leases_.acquire(path); + } + + void release_runtime_cache_lease(std::uint64_t lease) override { + std::lock_guard lock(runtime_cache_mutex_); + runtime_cache_leases_.release(lease, runtime_cache_ != nullptr, + [this] { persist_runtime_cache_locked(); }); + } + private: +#if defined(_WIN32) + // Declared before runtime_ so it outlives the runtime and every engine + // deserialized through it (members are destroyed in reverse order). + SynchronousGpuAllocator gpu_allocator_; +#endif + TrtUniquePtr runtime_; +#if defined(_WIN32) + TrtUniquePtr staged_runtime_; +#endif + nvinfer1::IRuntimeCache* runtime_cache_{nullptr}; + std::mutex runtime_cache_mutex_; + internal::RuntimeCacheLeaseState runtime_cache_leases_; + std::mutex retained_engines_mutex_; + std::unordered_map> retained_engines_; + std::mutex activation_arenas_mutex_; + std::unordered_map, + ActivationArenaKeyHash> + activation_arenas_; + std::unique_ptr - create_module_impl(const void* plan_data, size_t plan_size, const ModuleCreateOptions& options, - const std::vector& external_bindings) { - auto engine = deserialize_engine(*runtime_, plan_data, plan_size); - return create_context_module(engine, resolve_stream(options.stream), options, 0, - external_bindings); + create_module_from_file_impl(const char* plan_path, std::uint64_t plan_offset, + std::uint64_t plan_size, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, bool retain_engine, + bool serial_execution_context) { + validate_serial_execution_context(options, serial_execution_context); + if (weight_streaming_budget_bytes >= 0 && options.cuda_graphs) { + throw std::invalid_argument( + "[trtmc] RTX weight streaming is incompatible with CUDA graph capture"); + } + BoundedPlanStreamReader reader(plan_path, plan_offset, plan_size); + // Stable file and section identity prevents retained-engine reuse across + // bundle files or after an in-place plan replacement. + const std::string cache_key = + retain_engine ? retained_engine_key(reader, weight_streaming_budget_bytes) + : std::string{}; + std::unique_lock retained_lock; + if (retain_engine) { + // Serialize retained-engine creation. Two concurrent deserializations + // of the same multi-GiB plan can exceed the device-memory envelope + // before the losing insertion is released. + retained_lock = std::unique_lock(retained_engines_mutex_); + const auto hit = retained_engines_.find(cache_key); + if (hit != retained_engines_.end()) { + reader.verify_unchanged(); + std::cerr << "[trtmc.rtx_engine_cache] hit=1\n"; + return create_single_module_from_engine(hit->second, options, external_bindings, + true, serial_execution_context); + } + } + TrtUniquePtr engine( + file_backed_runtime().deserializeCudaEngine(reader)); + reader.verify_unchanged(); + if (!engine) + throw std::runtime_error("[trtmc] Failed to stream-deserialize engine (RTX)"); + if (retain_engine) { + std::shared_ptr retained( + engine.release(), [](nvinfer1::ICudaEngine* value) { delete value; }); + apply_weight_streaming_budget(*retained, weight_streaming_budget_bytes, + options.cuda_graphs); + auto module = create_single_module_from_engine(retained, options, external_bindings, + true, serial_execution_context); + retained_engines_.emplace(cache_key, retained); + std::cerr << "[trtmc.rtx_engine_cache] hit=0 retained=1\n"; + return module; + } + return create_single_module(engine.release(), options, external_bindings, + weight_streaming_budget_bytes, true, serial_execution_context); + } + + nvinfer1::IRuntime& file_backed_runtime() { +#if defined(_WIN32) + return *staged_runtime_; +#else + return *runtime_; +#endif + } + + static std::string retained_engine_key(const BoundedPlanStreamReader& reader, + std::int64_t weight_streaming_budget_bytes) { + int device = -1; + const cudaError_t status = cudaGetDevice(&device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } + return std::to_string(device) + ":" + reader.cache_identity() + ":" + + std::to_string(weight_streaming_budget_bytes); + } + + std::unique_ptr create_single_module_from_engine( + const std::shared_ptr& engine, const ModuleCreateOptions& options, + const std::vector& external_bindings, + bool use_synchronous_allocator = false, bool serial_execution_context = false) { + validate_optimization_profile(*engine, options.optimization_profile); + const StreamSetup stream_setup = resolve_stream(options.stream); + auto config = create_runtime_config(*engine, options, serial_execution_context); + return create_execution_module(engine, config, stream_setup, options, external_bindings, + use_synchronous_allocator, serial_execution_context); } std::unique_ptr - create_context_module(const std::shared_ptr& engine, - const StreamSetup& stream_setup, const ModuleCreateOptions& options, - int32_t profile_idx, - const std::vector& external_bindings) { - TrtUniquePtr runtime_config(engine->createRuntimeConfig()); - if (!runtime_config) - throw std::runtime_error("[trtmc] Failed to create RTX runtime config"); + create_single_module(nvinfer1::ICudaEngine* engine_raw, const ModuleCreateOptions& options, + const std::vector& external_bindings, + std::int64_t weight_streaming_budget_bytes, + bool use_synchronous_allocator = false, + bool serial_execution_context = false) { + std::shared_ptr engine( + engine_raw, [](nvinfer1::ICudaEngine* value) { delete value; }); + apply_weight_streaming_budget(*engine, weight_streaming_budget_bytes, options.cuda_graphs); + return create_single_module_from_engine(engine, options, external_bindings, + use_synchronous_allocator, + serial_execution_context); + } - auto runtime_cache = resolve_runtime_cache(*runtime_config, options.runtime_cache_path); - if (runtime_cache && !runtime_config->setRuntimeCache(runtime_cache->cache())) - throw std::runtime_error("[trtmc] TensorRT-RTX rejected the runtime cache"); - if (options.cuda_graphs) { - if (!runtime_config->setCudaGraphStrategy( - nvinfer1::CudaGraphStrategy::kWHOLE_GRAPH_CAPTURE)) { - throw std::runtime_error("[trtmc] TensorRT-RTX rejected whole-graph capture"); + std::shared_ptr + create_runtime_config(nvinfer1::ICudaEngine& engine, const ModuleCreateOptions& options, + bool serial_execution_context) { + std::shared_ptr config( + engine.createRuntimeConfig(), [](nvinfer1::IRuntimeConfig* value) { delete value; }); + if (!config) + throw std::runtime_error("[trtmc] Failed to create RTX runtime config"); + if (options.runtime_cache_path && options.runtime_cache_path[0] != '\0') + ensure_runtime_cache(config.get(), options.runtime_cache_path); + if (serial_execution_context) { + config->setExecutionContextAllocationStrategy( + nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED); + if (config->getExecutionContextAllocationStrategy() != + nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED) { + throw std::runtime_error( + "[trtmc] TensorRT-RTX rejected user-managed activation memory"); } } + configure_cuda_graphs(*config, options.cuda_graphs); + return config; + } + + std::shared_ptr resolve_activation_arena(bool serial_execution_context, + cudaStream_t stream) { + if (!serial_execution_context) + return {}; + int device = -1; + const cudaError_t status = cudaGetDevice(&device); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("[trtmc] Failed to identify the CUDA device: ") + + cudaGetErrorString(status)); + } + + const ActivationArenaKey key{device, stream}; + std::lock_guard lock(activation_arenas_mutex_); + const auto found = activation_arenas_.find(key); + if (found != activation_arenas_.end()) { + if (auto arena = found->second.lock()) + return arena; + activation_arenas_.erase(found); + } + auto arena = std::make_shared(device, stream); + activation_arenas_.emplace(key, arena); + return arena; + } - TrtUniquePtr context( - engine->createExecutionContext(runtime_config.get())); + std::unique_ptr + create_execution_module(const std::shared_ptr& engine, + const std::shared_ptr& config, + const StreamSetup& stream_setup, const ModuleCreateOptions& options, + const std::vector& external_bindings, + bool use_synchronous_allocator, bool serial_execution_context) { + const auto activation_arena = + resolve_activation_arena(serial_execution_context, stream_setup.stream); + const std::int64_t activation_memory_bytes = + activation_arena ? engine->getDeviceMemorySizeForProfileV2(options.optimization_profile) + : 0; + if (activation_memory_bytes < 0) + throw std::runtime_error("[trtmc] Invalid RTX activation memory requirement"); + std::unique_ptr context( + engine->createExecutionContext(config.get())); if (!context) throw std::runtime_error("[trtmc] Failed to create RTX execution context"); - +#if defined(_WIN32) + if (use_synchronous_allocator && !context->setTemporaryStorageAllocator(&gpu_allocator_)) + throw std::runtime_error( + "[trtmc] Failed to set synchronous RTX temporary-storage allocator"); +#else + (void)use_synchronous_allocator; +#endif auto module = std::make_unique( - engine.get(), context.release(), stream_setup.stream, profile_idx, - options.distributed_communicator, external_bindings, true); + engine.get(), context.get(), stream_setup.stream, options.optimization_profile, nullptr, + external_bindings, options.cuda_graphs, activation_arena, activation_memory_bytes); + // A completed TrtModuleImpl owns (or has already rejected and deleted) + // the context. Before that point the local owner handles exceptions. + (void)context.release(); if (!module->ok()) throw std::runtime_error("[trtmc] TrtModuleImpl creation failed (RTX)"); - - keep_backend_resources(*module, engine, stream_setup, options.distributed_owner); - if (runtime_cache) - module->keep_alive(runtime_cache); + module->keep_alive(engine); + module->keep_alive(config); + if (stream_setup.owner) + module->keep_alive(stream_setup.owner); return module; } - std::shared_ptr resolve_runtime_cache(nvinfer1::IRuntimeConfig& config, - const char* raw_path) { - if (raw_path == nullptr || *raw_path == '\0') - return {}; - const std::string path(raw_path); + void ensure_runtime_cache(nvinfer1::IRuntimeConfig* cfg, const char* path) { std::lock_guard lock(runtime_cache_mutex_); - if (auto existing = runtime_cache_.lock()) { - if (existing->path() != path) { - throw std::invalid_argument( - "one TensorRT-RTX backend cannot use two runtime cache paths concurrently"); - } - return existing; + if (runtime_cache_leases_.empty()) { + throw std::runtime_error( + "[trtmc] RTX runtime cache use requires an active pipeline lease"); + } + if (runtime_cache_leases_.path() != path) { + throw std::invalid_argument( + "[trtmc] RTX backend cannot share one runtime cache across different paths"); } - auto created = std::make_shared(config, path); - runtime_cache_ = created; - return created; + if (!runtime_cache_) + runtime_cache_ = create_runtime_cache(*cfg, path).release(); + if (!cfg->setRuntimeCache(*runtime_cache_)) + throw std::runtime_error("[trtmc] Failed to attach the TensorRT-RTX runtime cache"); } - TrtUniquePtr runtime_; - std::mutex runtime_cache_mutex_; - std::weak_ptr runtime_cache_; + void persist_runtime_cache_locked() { + auto* serialized = runtime_cache_->serialize(); + if (serialized == nullptr) + throw std::runtime_error( + "[trtmc] TensorRT-RTX returned a null serialized runtime cache"); + std::unique_ptr memory(serialized); + if (memory->size() > + static_cast(std::numeric_limits::max())) { + throw std::overflow_error("[trtmc] RTX runtime cache is too large to persist"); + } + + internal::persist_runtime_cache_file(runtime_cache_leases_.path(), memory->data(), + memory->size()); + std::cerr << "[trtmc] RTX runtime cache saved: " << runtime_cache_leases_.path() << " (" + << memory->size() << " bytes)\n"; + } }; } // namespace trtmc @@ -220,12 +1013,26 @@ class RtxBackend final : public IBackend { extern "C" trtmc::IBackend* trtmc_create_backend() { try { return new trtmc::RtxBackend(); - } catch (const std::exception& error) { - std::cerr << "[trtmc] RTX backend init failed: " << error.what() << std::endl; + } catch (const std::exception& e) { + std::cerr << "[trtmc] RTX backend init failed: " << e.what() << std::endl; return nullptr; } } -extern "C" void trtmc_destroy_backend(trtmc::IBackend* backend) { - delete backend; +extern "C" void trtmc_destroy_backend(trtmc::IBackend* b) { + delete b; +} + +extern "C" const char* trtmc_backend_abi() { + static const std::string abi = + std::to_string(getInferLibMajorVersion()) + "." + std::to_string(getInferLibMinorVersion()); + return abi.c_str(); +} + +extern "C" const char* trtmc_backend_runtime_version() { + static const std::string version = std::to_string(getInferLibMajorVersion()) + "." + + std::to_string(getInferLibMinorVersion()) + "." + + std::to_string(getInferLibPatchVersion()) + "." + + std::to_string(getInferLibBuildVersion()); + return version.c_str(); } diff --git a/core/runtime/tensorrt/runtime_cache_persistence.cpp b/core/runtime/tensorrt/runtime_cache_persistence.cpp new file mode 100644 index 0000000000..c168859871 --- /dev/null +++ b/core/runtime/tensorrt/runtime_cache_persistence.cpp @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/tensorrt/runtime_cache_persistence.h" + +#include +#include +#include +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#endif + +namespace trtmc::internal { + +namespace { + +namespace fs = std::filesystem; + +std::uint64_t current_process_id() { +#if defined(_WIN32) + return static_cast(GetCurrentProcessId()); +#else + return static_cast(getpid()); +#endif +} + +fs::path runtime_cache_temporary_path(const fs::path& target) { + fs::path temporary = target; + temporary += ".tmp." + std::to_string(current_process_id()); + return temporary; +} + +void durable_flush(const fs::path& path) { +#if defined(_WIN32) + HANDLE handle = CreateFileW(path.c_str(), GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (handle == INVALID_HANDLE_VALUE) { + const auto error = static_cast(GetLastError()); + throw std::system_error(error, std::system_category(), + "failed to open RTX runtime cache for durable flush " + + path.string()); + } + DWORD error = ERROR_SUCCESS; + if (!FlushFileBuffers(handle)) + error = GetLastError(); + if (!CloseHandle(handle) && error == ERROR_SUCCESS) + error = GetLastError(); + if (error != ERROR_SUCCESS) { + throw std::system_error(static_cast(error), std::system_category(), + "failed to durably flush RTX runtime cache " + path.string()); + } +#else + // POSIX rename retains the existing best-effort durability contract. + (void)path; +#endif +} + +void atomic_replace(const fs::path& temporary, const fs::path& target) { +#if defined(_WIN32) + if (!MoveFileExW(temporary.c_str(), target.c_str(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH)) { + const auto error = static_cast(GetLastError()); + throw std::system_error(error, std::system_category(), + "failed to atomically replace RTX runtime cache " + + target.string()); + } +#else + std::error_code error; + fs::rename(temporary, target, error); + if (error) { + throw std::system_error(error, "failed to atomically replace RTX runtime cache " + + target.string()); + } +#endif +} + +} // namespace + +void persist_runtime_cache_file(const fs::path& target, const void* data, std::size_t size) { + if (target.empty()) + throw std::invalid_argument("[trtmc] RTX runtime cache path must not be empty"); + if (size != 0 && data == nullptr) + throw std::invalid_argument("[trtmc] Serialized RTX runtime cache data is null"); + if (size > static_cast(std::numeric_limits::max())) + throw std::overflow_error("[trtmc] RTX runtime cache is too large to persist"); + + const fs::path temporary = runtime_cache_temporary_path(target); + try { + std::ofstream output(temporary, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("[trtmc] Failed to open RTX runtime cache temporary file: " + + temporary.string()); + } + if (size != 0) + output.write(static_cast(data), static_cast(size)); + output.flush(); + if (!output) { + throw std::runtime_error("[trtmc] Failed to write or flush RTX runtime cache: " + + temporary.string()); + } + output.close(); + if (!output) { + throw std::runtime_error("[trtmc] Failed to close RTX runtime cache: " + + temporary.string()); + } + durable_flush(temporary); + atomic_replace(temporary, target); + } catch (...) { + std::error_code ignored; + fs::remove(temporary, ignored); + throw; + } +} + +std::uint64_t RuntimeCacheLeaseState::acquire(const char* path) { + if (path == nullptr || path[0] == '\0') + throw std::invalid_argument("[trtmc] RTX runtime cache lease requires a path"); + if (path_.empty()) { + path_ = path; + } else if (path_ != path) { + throw std::invalid_argument( + "[trtmc] RTX backend cannot share one runtime cache across different paths"); + } + const std::uint64_t lease = next_lease_++; + if (lease == 0 || !active_.insert(lease).second) + throw std::runtime_error("[trtmc] RTX runtime cache lease space exhausted"); + return lease; +} + +void RuntimeCacheLeaseState::release(std::uint64_t lease, bool cache_materialized, + const std::function& persist_final_cache) { + const auto active = active_.find(lease); + if (lease == 0 || active == active_.end()) + throw std::invalid_argument("[trtmc] RTX runtime cache lease is invalid or inactive"); + if (active_.size() == 1 && cache_materialized) + persist_final_cache(); + active_.erase(active); +} + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/runtime_cache_persistence.h b/core/runtime/tensorrt/runtime_cache_persistence.h new file mode 100644 index 0000000000..fce4196a8c --- /dev/null +++ b/core/runtime/tensorrt/runtime_cache_persistence.h @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace trtmc::internal { + +// Write a complete serialized cache to a same-directory temporary file, make +// its contents durable, then atomically replace the target. The temporary is +// removed on every reported failure. +void persist_runtime_cache_file(const std::filesystem::path& target, const void* data, + std::size_t size); + +// Lease bookkeeping is deliberately independent of TensorRT-RTX so the +// failure contract can be regression-tested without constructing an engine. +// Callers must serialize access to this object. +class RuntimeCacheLeaseState final { + public: + std::uint64_t acquire(const char* path); + + // The final active lease is erased only after persistence succeeds. If the + // callback throws, the exact lease remains active and can be retried. + void release(std::uint64_t lease, bool cache_materialized, + const std::function& persist_final_cache); + + bool empty() const noexcept { return active_.empty(); } + std::size_t size() const noexcept { return active_.size(); } + const std::string& path() const noexcept { return path_; } + + private: + std::string path_; + std::uint64_t next_lease_{1}; + std::unordered_set active_; +}; + +} // namespace trtmc::internal diff --git a/core/runtime/tensorrt/trt_backend.cpp b/core/runtime/tensorrt/trt_backend.cpp index caf20f03a0..44b29f761e 100644 --- a/core/runtime/tensorrt/trt_backend.cpp +++ b/core/runtime/tensorrt/trt_backend.cpp @@ -124,8 +124,9 @@ class TrtBackend final : public IBackend { stream_owner = owned; } - auto module = std::make_unique( - engine, ctx, stream, 0, options.distributed_communicator, external_bindings); + auto module = + std::make_unique(engine, ctx, stream, options.optimization_profile, + options.distributed_communicator, external_bindings); if (!module->ok()) { delete engine; throw std::runtime_error("[trtmc] TrtModuleImpl creation failed"); diff --git a/core/runtime/tensorrt/trt_logger.cpp b/core/runtime/tensorrt/trt_logger.cpp index acd902a8b7..976a4fd363 100644 --- a/core/runtime/tensorrt/trt_logger.cpp +++ b/core/runtime/tensorrt/trt_logger.cpp @@ -54,8 +54,11 @@ void TrtLogger::clear_error() { } TrtUniquePtr create_trt_runtime() { - static TrtLogger logger; - return TrtUniquePtr(nvinfer1::createInferRuntime(logger)); + // Backends may be released from an atexit handler after function-static + // destruction has begun. Keep the small logger process-lifetime so it + // outlives every TensorRT runtime and engine on every platform. + static TrtLogger* logger = new TrtLogger(); + return TrtUniquePtr(nvinfer1::createInferRuntime(*logger)); } } // namespace trtmc diff --git a/core/runtime/tensorrt/trt_module_impl.cpp b/core/runtime/tensorrt/trt_module_impl.cpp index 6aaaef4f23..60de6025ce 100644 --- a/core/runtime/tensorrt/trt_module_impl.cpp +++ b/core/runtime/tensorrt/trt_module_impl.cpp @@ -62,37 +62,58 @@ TrtModuleImpl::TrtModuleImpl(nvinfer1::ICudaEngine* engine, nvinfer1::IExecution cudaStream_t stream, int32_t profile_idx, void* distributed_communicator, const std::vector& external_bindings, - bool backend_managed_cuda_graph) + bool backend_managed_cuda_graph, + std::shared_ptr activation_arena, + std::int64_t activation_memory_bytes) : engine_(engine), ctx_(ctx), stream_(stream), profile_idx_(profile_idx), distributed_communicator_(distributed_communicator), - backend_managed_cuda_graph_(backend_managed_cuda_graph), - cuda_graph_(std::make_unique()) { + activation_arena_(std::move(activation_arena)), + backend_managed_cuda_graph_(backend_managed_cuda_graph) { if (!ctx_) return; try { - discover_tensor_aliases(engine); - validate_initial_external_bindings(engine, external_bindings); + initialize_module(engine, external_bindings, activation_memory_bytes); } catch (const std::exception& error) { - std::cerr << "[trt_module] Invalid external binding: " << error.what() << '\n'; - delete ctx_; - ctx_ = nullptr; - return; + std::cerr << "[trt_module] Module initialization failed: " << error.what() << '\n'; + cleanup_failed_initialization(); + } catch (...) { + std::cerr << "[trt_module] Module initialization failed\n"; + cleanup_failed_initialization(); } - if (!attach_distributed_communicator()) { - delete ctx_; - ctx_ = nullptr; +} + +void TrtModuleImpl::initialize_module(nvinfer1::ICudaEngine* engine, + const std::vector& external_bindings, + std::int64_t activation_memory_bytes) { + cuda_graph_ = std::make_unique(); + discover_tensor_aliases(engine); + validate_initial_external_bindings(engine, external_bindings); + if (!attach_distributed_communicator()) + throw std::runtime_error("failed to attach the distributed communicator"); + select_optimization_profile(); + if (activation_arena_) + activation_arena_->attach(ctx_, stream_, activation_memory_bytes); + allocate_buffers(engine); +} + +void TrtModuleImpl::select_optimization_profile() { + if (profile_idx_ <= 0) return; + if (!ctx_->setOptimizationProfileAsync(profile_idx_, stream_)) + throw std::runtime_error("failed to select the optimization profile"); + const cudaError_t status = cudaStreamSynchronize(stream_); + if (status != cudaSuccess) { + throw std::runtime_error(std::string("failed to synchronize the optimization ") + + "profile selection: " + cudaGetErrorString(status)); } - if (profile_idx_ > 0) { - if (!ctx_->setOptimizationProfileAsync(profile_idx_, stream_)) { - std::cerr << "[trt_module] Failed to set optimization profile " << profile_idx_ << "\n"; - delete ctx_; - ctx_ = nullptr; - return; - } - cudaStreamSynchronize(stream_); +} + +void TrtModuleImpl::cleanup_failed_initialization() noexcept { + try { + free_buffers(); + } catch (...) { } - allocate_buffers(engine); + destroy_execution_context(); } void TrtModuleImpl::discover_tensor_aliases(nvinfer1::ICudaEngine* engine) { @@ -122,37 +143,94 @@ void TrtModuleImpl::discover_tensor_aliases(nvinfer1::ICudaEngine* engine) { void TrtModuleImpl::validate_initial_external_bindings( nvinfer1::ICudaEngine* engine, const std::vector& external_bindings) { for (const auto& binding : external_bindings) { - if (binding.tensor_name.empty()) - throw std::invalid_argument("tensor name must not be empty"); - if (binding.device_ptr == nullptr) - throw std::invalid_argument("buffer for '" + binding.tensor_name + "' is null"); - if (initial_external_bindings_.count(binding.tensor_name) != 0) - throw std::invalid_argument("duplicate tensor '" + binding.tensor_name + "'"); - - if (!engine_has_io_tensor(engine, binding.tensor_name)) - throw std::invalid_argument("unknown tensor '" + binding.tensor_name + "'"); - if (alias_input_by_output_.count(binding.tensor_name) != 0 || - alias_outputs_by_input_.count(binding.tensor_name) != 0) { - throw std::invalid_argument("TensorRT alias tensor '" + binding.tensor_name + - "' must be bound after module creation"); - } - - const auto dims = engine->getTensorShape(binding.tensor_name.c_str()); - if (dims_are_dynamic(dims)) { - throw std::invalid_argument("tensor '" + binding.tensor_name + - "' is dynamic; prebinding requires a static shape"); - } - std::vector shape; - const auto required_bytes = compute_alloc_bytes( - dims, from_trt_dtype(engine->getTensorDataType(binding.tensor_name.c_str())), shape); - if (binding.capacity_bytes < required_bytes) { - throw std::invalid_argument("buffer for '" + binding.tensor_name + "' has " + - std::to_string(binding.capacity_bytes) + - " bytes; expected at least " + - std::to_string(required_bytes)); - } + validate_initial_external_binding_identity(engine, binding); + const auto required_bytes = initial_external_binding_required_bytes(engine, binding); + validate_external_binding_capacity(binding.tensor_name, binding.capacity_bytes, + required_bytes); initial_external_bindings_.emplace(binding.tensor_name, binding.device_ptr); + initial_external_binding_capacities_.emplace(binding.tensor_name, binding.capacity_bytes); + } +} + +void TrtModuleImpl::validate_initial_external_binding_identity( + nvinfer1::ICudaEngine* engine, const ModuleExternalBinding& binding) const { + if (binding.tensor_name.empty()) + throw std::invalid_argument("tensor name must not be empty"); + if (binding.device_ptr == nullptr) + throw std::invalid_argument("buffer for '" + binding.tensor_name + "' is null"); + if (initial_external_bindings_.count(binding.tensor_name) != 0) + throw std::invalid_argument("duplicate tensor '" + binding.tensor_name + "'"); + if (!engine_has_io_tensor(engine, binding.tensor_name)) + throw std::invalid_argument("unknown tensor '" + binding.tensor_name + "'"); + if (alias_input_by_output_.count(binding.tensor_name) != 0 || + alias_outputs_by_input_.count(binding.tensor_name) != 0) { + throw std::invalid_argument("TensorRT alias tensor '" + binding.tensor_name + + "' must be bound after module creation"); + } +} + +std::size_t +TrtModuleImpl::initial_external_binding_required_bytes(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const { + auto dims = engine->getTensorShape(binding.tensor_name.c_str()); + const bool is_input = + engine->getTensorIOMode(binding.tensor_name.c_str()) == nvinfer1::TensorIOMode::kINPUT; + if (dims_are_dynamic(dims) && is_input) { + if (profile_idx_ < 0 || engine->getNbOptimizationProfiles() <= profile_idx_) + throw std::invalid_argument("dynamic tensor '" + binding.tensor_name + + "' has no selected optimization profile"); + dims = engine->getProfileShape(binding.tensor_name.c_str(), profile_idx_, + nvinfer1::OptProfileSelector::kMAX); + } + if (dims_are_dynamic(dims)) + return 0; + std::vector shape; + return compute_alloc_bytes( + dims, from_trt_dtype(engine->getTensorDataType(binding.tensor_name.c_str())), shape); +} + +void TrtModuleImpl::validate_external_binding_capacity(const std::string& name, + std::size_t available, + std::size_t required) { + if (required > 0 && available < required) { + throw std::invalid_argument("buffer for '" + name + "' has " + std::to_string(available) + + " bytes; expected at least " + std::to_string(required)); + } +} + +bool TrtModuleImpl::attach_initial_external_binding(const std::string& name, BufferEntry& entry) { + const auto external = initial_external_bindings_.find(name); + if (external == initial_external_bindings_.end()) + return false; + const auto capacity = initial_external_binding_capacities_.find(name); + const auto available = + capacity == initial_external_binding_capacities_.end() ? std::size_t{0} : capacity->second; + validate_external_binding_capacity(name, available, entry.nbytes); + entry.d_ptr = external->second; + entry.is_external = true; + return true; +} + +void TrtModuleImpl::assign_initial_input_storage(const std::string& name, BufferEntry& entry) { + if (attach_initial_external_binding(name, entry)) + return; + if (entry.is_dynamic || !should_allocate_input(name, entry.nbytes)) + return; + void* allocation = nullptr; + if (cudaMalloc(&allocation, entry.nbytes) == cudaSuccess) { + entry.d_ptr = allocation; + cudaMemsetAsync(entry.d_ptr, 0, entry.nbytes, stream_); + } +} + +void TrtModuleImpl::assign_initial_output_storage(const std::string& name, BufferEntry& entry) { + if (attach_initial_external_binding(name, entry) || entry.nbytes == 0) + return; + if (cudaMalloc(&entry.d_ptr, entry.nbytes) != cudaSuccess) { + entry.d_ptr = nullptr; + return; } + cudaMemsetAsync(entry.d_ptr, 0, entry.nbytes, stream_); } void TrtModuleImpl::bind_external(const std::string& name, void* ptr, @@ -213,9 +291,22 @@ TrtModuleImpl::~TrtModuleImpl() { // CUDA Graphs may contain TensorRT collective launches that retain the // distributed communicator. Destroy the captured graph before member // teardown releases distributed_owner from keep_alive_. - cuda_graph_->reset(); + if (cuda_graph_) + cuda_graph_->reset(); + auto activation_arena = activation_arena_; + if (activation_arena && ctx_) + activation_arena->detach(ctx_); free_buffers(); delete ctx_; + ctx_ = nullptr; + activation_arena_.reset(); +} + +void TrtModuleImpl::destroy_execution_context() noexcept { + if (activation_arena_ && ctx_) + activation_arena_->detach(ctx_); + delete ctx_; + ctx_ = nullptr; } void TrtModuleImpl::keep_alive(std::shared_ptr resource) { @@ -257,7 +348,8 @@ void TrtModuleImpl::update_dynamic_shape(const std::string& name, BufferEntry& e dims.nbDims = static_cast(new_shape.size()); for (int32_t d = 0; d < dims.nbDims; ++d) dims.d[d] = new_shape[d]; - ctx_->setInputShape(name.c_str(), dims); + if (!ctx_->setInputShape(name.c_str(), dims)) + throw std::runtime_error("TensorRT rejected dynamic input shape for '" + name + "'"); entry.shape = new_shape; } @@ -298,7 +390,9 @@ void TrtModuleImpl::set_dynamic_input_shapes(nvinfer1::ICudaEngine* engine, int3 continue; if (dims_are_dynamic(engine->getTensorShape(name.c_str()))) { auto dims = engine->getProfileShape(name.c_str(), profile_idx_, selector); - ctx_->setInputShape(name.c_str(), dims); + if (!ctx_->setInputShape(name.c_str(), dims)) + throw std::runtime_error("TensorRT rejected profile input shape for '" + name + + "'"); } } } @@ -330,25 +424,16 @@ void TrtModuleImpl::allocate_single_input(nvinfer1::ICudaEngine* engine, const s entry.is_dynamic = is_dynamic; entry.shape = is_dynamic ? dims_to_shape(init_dims) : shape; - const auto external = initial_external_bindings_.find(name); - if (external != initial_external_bindings_.end()) { - entry.d_ptr = external->second; - entry.is_external = true; - } else if (!is_dynamic && should_allocate_input(name, nbytes)) { - void* allocation = nullptr; - if (cudaMalloc(&allocation, nbytes) == cudaSuccess) { - entry.d_ptr = allocation; - cudaMemsetAsync(entry.d_ptr, 0, nbytes, stream_); - } - } + assign_initial_input_storage(name, entry); if (entry.d_ptr) bind_tensor_address(name, entry); - if (is_dynamic) - ctx_->setInputShape(name.c_str(), init_dims); - buffers_[name] = std::move(entry); + + if (is_dynamic && !ctx_->setInputShape(name.c_str(), init_dims)) + throw std::runtime_error("TensorRT rejected initial dynamic input shape for '" + name + + "'"); } void TrtModuleImpl::ensure_input_buffer(const std::string& name, BufferEntry& entry) { @@ -386,48 +471,41 @@ void TrtModuleImpl::allocate_output_buffers(nvinfer1::ICudaEngine* engine, int32 const std::string name(raw_name); if (engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT) continue; + allocate_single_output(engine, name); + } +} - auto dtype = from_trt_dtype(engine->getTensorDataType(name.c_str())); +void TrtModuleImpl::allocate_single_output(nvinfer1::ICudaEngine* engine, const std::string& name) { + auto dtype = from_trt_dtype(engine->getTensorDataType(name.c_str())); - // For dynamic engines, query the context for inferred output shape - // (based on the max input shapes set by the caller). - // For static engines, use the engine shape directly. - nvinfer1::Dims out_dims = has_dynamic_shapes_ ? ctx_->getTensorShape(name.c_str()) - : engine->getTensorShape(name.c_str()); + // For dynamic engines, query the context for inferred output shape + // (based on the max input shapes set by the caller). + // For static engines, use the engine shape directly. + nvinfer1::Dims out_dims = has_dynamic_shapes_ ? ctx_->getTensorShape(name.c_str()) + : engine->getTensorShape(name.c_str()); - std::vector shape; - std::size_t nbytes = compute_alloc_bytes(out_dims, dtype, shape); + std::vector shape; + std::size_t nbytes = compute_alloc_bytes(out_dims, dtype, shape); - BufferEntry entry; - entry.shape = shape; - entry.dtype = dtype; - entry.nbytes = nbytes; - entry.is_input = false; + BufferEntry entry; + entry.shape = shape; + entry.dtype = dtype; + entry.nbytes = nbytes; + entry.is_input = false; - if (alias_input_by_output_.count(name) != 0) { - buffers_[name] = std::move(entry); - continue; - } + if (alias_input_by_output_.count(name) != 0) { + buffers_[name] = std::move(entry); + return; + } - const auto external = initial_external_bindings_.find(name); - if (external != initial_external_bindings_.end()) { - entry.d_ptr = external->second; - entry.is_external = true; - } else if (nbytes > 0) { - auto err = cudaMalloc(&entry.d_ptr, nbytes); - if (err != cudaSuccess) - entry.d_ptr = nullptr; - else - cudaMemsetAsync(entry.d_ptr, 0, nbytes, stream_); - } + assign_initial_output_storage(name, entry); - if (entry.d_ptr) - bind_tensor_address(name, entry); + if (entry.d_ptr) + bind_tensor_address(name, entry); - allocate_host_output_staging(host_output_staging_, name, nbytes, entry.is_external); + allocate_host_output_staging(host_output_staging_, name, nbytes, entry.is_external); - buffers_[name] = std::move(entry); - } + buffers_[name] = std::move(entry); } // --- Buffer allocation --- @@ -451,6 +529,7 @@ void TrtModuleImpl::allocate_buffers(nvinfer1::ICudaEngine* engine) { set_dynamic_input_shapes(engine, num_io, nvinfer1::OptProfileSelector::kOPT); initial_external_bindings_.clear(); + initial_external_binding_capacities_.clear(); cudaStreamSynchronize(stream_); } @@ -479,9 +558,9 @@ void TrtModuleImpl::bind_alias_outputs_or_invalidate(const std::vectorreset(); - delete ctx_; - ctx_ = nullptr; + if (cuda_graph_) + cuda_graph_->reset(); + destroy_execution_context(); throw std::runtime_error("TensorRT rejected external alias output '" + output_name + "'"); } @@ -599,6 +678,10 @@ TensorMap TrtModuleImpl::forward(const TensorMap& inputs) { void TrtModuleImpl::enable_cuda_graph() { if (backend_managed_cuda_graph_) return; + if (activation_arena_) { + throw std::logic_error( + "CUDA graph capture is incompatible with a shared TensorRT activation arena"); + } use_cuda_graph_ = true; cuda_graph_->reset(); } @@ -635,7 +718,18 @@ void TrtModuleImpl::execute_enqueue() { validate_alias_groups_bound(); alias_groups_ready_ = true; } - record_timed_enqueue(); + if (!activation_arena_) { + record_timed_enqueue(); + return; + } + activation_arena_->begin_enqueue(ctx_); + try { + record_timed_enqueue(); + } catch (...) { + activation_arena_->end_enqueue(); + throw; + } + activation_arena_->end_enqueue(); } bool TrtModuleImpl::cuda_graph_captured() const { diff --git a/core/runtime/tensorrt/trt_module_impl.h b/core/runtime/tensorrt/trt_module_impl.h index 58f11a33f2..544c8dbbf1 100644 --- a/core/runtime/tensorrt/trt_module_impl.h +++ b/core/runtime/tensorrt/trt_module_impl.h @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -24,6 +25,18 @@ namespace trtmc { class CudaGraphExec; class TrtModuleImplTestPeer; +// Optional TensorRT-RTX coordinator for execution contexts which run serially +// on one stream and share a user-managed activation allocation. +class ITrtActivationArena { + public: + virtual ~ITrtActivationArena() = default; + virtual void attach(nvinfer1::IExecutionContext* context, cudaStream_t stream, + std::int64_t required_bytes) = 0; + virtual void begin_enqueue(nvinfer1::IExecutionContext* context) = 0; + virtual void end_enqueue() noexcept = 0; + virtual void detach(nvinfer1::IExecutionContext* context) noexcept = 0; +}; + class TrtModuleImpl final : public ITrtModule { public: // Backend creates engine + context, passes them in. @@ -32,7 +45,9 @@ class TrtModuleImpl final : public ITrtModule { cudaStream_t stream, int32_t profile_idx = 0, void* distributed_communicator = nullptr, const std::vector& external_bindings = {}, - bool backend_managed_cuda_graph = false); + bool backend_managed_cuda_graph = false, + std::shared_ptr activation_arena = {}, + std::int64_t activation_memory_bytes = 0); ~TrtModuleImpl() override; TrtModuleImpl(const TrtModuleImpl&) = delete; @@ -91,6 +106,7 @@ class TrtModuleImpl final : public ITrtModule { cudaStream_t stream_{nullptr}; int32_t profile_idx_{0}; void* distributed_communicator_{nullptr}; + std::shared_ptr activation_arena_; bool has_dynamic_shapes_{false}; bool use_cuda_graph_{false}; bool backend_managed_cuda_graph_{false}; @@ -98,6 +114,7 @@ class TrtModuleImpl final : public ITrtModule { std::unique_ptr cuda_graph_; std::vector> keep_alive_; std::unordered_map initial_external_bindings_; + std::unordered_map initial_external_binding_capacities_; std::unordered_map buffers_; std::unordered_map alias_input_by_output_; std::unordered_map> alias_outputs_by_input_; @@ -106,11 +123,25 @@ class TrtModuleImpl final : public ITrtModule { std::string timing_label_{"engine"}; std::vector timing_events_; + void initialize_module(nvinfer1::ICudaEngine* engine, + const std::vector& external_bindings, + std::int64_t activation_memory_bytes); + void select_optimization_profile(); + void cleanup_failed_initialization() noexcept; void allocate_buffers(nvinfer1::ICudaEngine* engine); void discover_tensor_aliases(nvinfer1::ICudaEngine* engine); void validate_initial_external_bindings(nvinfer1::ICudaEngine* engine, const std::vector& external_bindings); + void validate_initial_external_binding_identity(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const; + std::size_t initial_external_binding_required_bytes(nvinfer1::ICudaEngine* engine, + const ModuleExternalBinding& binding) const; + static void validate_external_binding_capacity(const std::string& name, std::size_t available, + std::size_t required); + bool attach_initial_external_binding(const std::string& name, BufferEntry& entry); + void assign_initial_input_storage(const std::string& name, BufferEntry& entry); + void assign_initial_output_storage(const std::string& name, BufferEntry& entry); void bind_alias_group(const std::string& input_name, void* ptr); bool should_allocate_input(const std::string& name, std::size_t nbytes) const; void validate_alias_outputs_exist(const std::vector& output_names) const; @@ -125,11 +156,13 @@ class TrtModuleImpl final : public ITrtModule { int32_t num_profiles); void ensure_input_buffer(const std::string& name, BufferEntry& entry); void allocate_output_buffers(nvinfer1::ICudaEngine* engine, int32_t num_io); + void allocate_single_output(nvinfer1::ICudaEngine* engine, const std::string& name); void set_dynamic_input_shapes(nvinfer1::ICudaEngine* engine, int32_t num_io, nvinfer1::OptProfileSelector selector); void update_dynamic_shape(const std::string& name, BufferEntry& entry, const std::vector& new_shape); void execute_enqueue(); + void destroy_execution_context() noexcept; void flush_timing_events(); bool begin_timing_event(TimingEvent& event); void finish_timing_event(TimingEvent event); diff --git a/core/runtime/tests/test_bundle_format_v1.cpp b/core/runtime/tests/test_bundle_format_v1.cpp index 0da636b49b..f61cca4d0a 100644 --- a/core/runtime/tests/test_bundle_format_v1.cpp +++ b/core/runtime/tests/test_bundle_format_v1.cpp @@ -5,13 +5,13 @@ #include "runtime/bundle/bundle_format.h" +#include #include #include #include #include #include #include -#include namespace { @@ -25,10 +25,11 @@ void check(bool condition, const char* name) { } std::filesystem::path temp_dir() { - char pattern[] = "/tmp/trtmc_bundle_v1_XXXXXX"; - char* path = mkdtemp(pattern); - if (path == nullptr) - throw std::runtime_error("mkdtemp failed"); + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + const auto path = + std::filesystem::temp_directory_path() / ("trtmc-bundle-v1-" + std::to_string(nonce)); + if (!std::filesystem::create_directory(path)) + throw std::runtime_error("failed to create temporary test directory"); return path; } diff --git a/core/runtime/tests/test_dynamic_library.cpp b/core/runtime/tests/test_dynamic_library.cpp new file mode 100644 index 0000000000..6cc0043970 --- /dev/null +++ b/core/runtime/tests/test_dynamic_library.cpp @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/platform/dynamic_library.h" + +#include +#include +#include + +namespace { + +int failures = 0; + +void check(bool condition, const char* name) { + if (!condition) { + std::cerr << "FAIL: " << name << '\n'; + ++failures; + } +} + +} // namespace + +int main() { +#if defined(_WIN32) + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "trtmc_example.dll", + "Windows library filename"); +#elif defined(__APPLE__) + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "libtrtmc_example.dylib", + "macOS library filename"); +#else + check(trtmc::internal::dynamic_library_filename("trtmc_example") == "libtrtmc_example.so", + "Linux library filename"); +#endif + + std::string error; + auto missing = trtmc::internal::open_dynamic_library( + std::filesystem::temp_directory_path() / "trtmc-missing-dynamic-library", + trtmc::internal::DynamicLibraryVisibility::local, &error); + check(missing == nullptr, "missing dynamic library is rejected"); + check(!error.empty(), "missing dynamic library reports an error"); + + auto handle = trtmc::internal::open_dynamic_library( + TRTMC_TEST_LIBRARY, trtmc::internal::DynamicLibraryVisibility::local, &error); + check(handle != nullptr, "test dynamic library opens"); + if (handle != nullptr) { + auto symbol = + trtmc::internal::dynamic_library_symbol(handle, "trtmc_create_backend", &error); + check(symbol != nullptr, "exported C ABI symbol resolves"); + check(trtmc::internal::close_dynamic_library(handle, &error), + "test dynamic library closes"); + } + + std::cerr << (failures == 0 ? "ALL PASSED" : "SOME FAILED") << '\n'; + return failures; +} diff --git a/core/runtime/tests/test_family_loader.cpp b/core/runtime/tests/test_family_loader.cpp index 1f77d55db2..e7ae6a62fe 100644 --- a/core/runtime/tests/test_family_loader.cpp +++ b/core/runtime/tests/test_family_loader.cpp @@ -4,10 +4,10 @@ */ #include "runtime/bundle/bundle_format.h" +#include "runtime/platform/dynamic_library.h" #include "trtmc/runtime/family_loader.h" #include -#include #include #include #include @@ -62,18 +62,19 @@ bool rtx_options_throw(const std::filesystem::path& bundle, const std::string& r void check_rtx_options(const std::filesystem::path& runtime_root, const std::string& expected_cache_path, bool expected_cuda_graphs) { - const auto library_path = runtime_root / "libtrtmc_backend_trt_rtx.so"; - void* handle = dlopen(library_path.c_str(), RTLD_NOW | RTLD_LOCAL); + const auto library_path = + runtime_root / trtmc::internal::dynamic_library_filename("trtmc_backend_trt_rtx"); + void* handle = trtmc::internal::open_dynamic_library(library_path); check(handle != nullptr, "fake RTX backend remains loaded"); if (handle == nullptr) return; using CachePathFn = const char* (*)(); using CudaGraphsFn = bool (*)(); - const auto cache_path = - reinterpret_cast(dlsym(handle, "trtmc_test_backend_last_runtime_cache_path")); - const auto cuda_graphs = - reinterpret_cast(dlsym(handle, "trtmc_test_backend_last_cuda_graphs")); + const auto cache_path = reinterpret_cast(trtmc::internal::dynamic_library_symbol( + handle, "trtmc_test_backend_last_runtime_cache_path")); + const auto cuda_graphs = reinterpret_cast( + trtmc::internal::dynamic_library_symbol(handle, "trtmc_test_backend_last_cuda_graphs")); check(cache_path != nullptr, "fake RTX cache-path probe is exported"); check(cuda_graphs != nullptr, "fake RTX CUDA-graphs probe is exported"); if (cache_path != nullptr) @@ -81,7 +82,7 @@ void check_rtx_options(const std::filesystem::path& runtime_root, if (cuda_graphs != nullptr) check(cuda_graphs() == expected_cuda_graphs, "CUDA-graphs option reaches delayed module creation"); - dlclose(handle); + trtmc::internal::close_dynamic_library(handle); } } // namespace diff --git a/core/runtime/tests/test_runtime_cache_persistence.cpp b/core/runtime/tests/test_runtime_cache_persistence.cpp new file mode 100644 index 0000000000..9fcc442166 --- /dev/null +++ b/core/runtime/tests/test_runtime_cache_persistence.cpp @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "runtime/tensorrt/runtime_cache_persistence.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +namespace fs = std::filesystem; + +class TempDirectory { + public: + TempDirectory() { + const auto nonce = std::chrono::steady_clock::now().time_since_epoch().count(); + path_ = fs::temp_directory_path() / ("trtmc-runtime-cache-test-" + std::to_string(nonce)); + fs::create_directory(path_); + } + + ~TempDirectory() { + std::error_code error; + fs::remove_all(path_, error); + } + + const fs::path& path() const { return path_; } + + private: + fs::path path_; +}; + +void check(bool condition, const char* message) { + if (!condition) + throw std::runtime_error(message); +} + +std::string read_file(const fs::path& path) { + std::ifstream input(path, std::ios::binary); + if (!input) + throw std::runtime_error("failed to read test file"); + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +void test_only_final_lease_persists_with_durable_atomic_replace() { + TempDirectory directory; + const fs::path target = directory.path() / "cache.bin"; + { + std::ofstream output(target, std::ios::binary); + output << "old-cache"; + } + + trtmc::internal::RuntimeCacheLeaseState leases; + const std::string target_string = target.string(); + const auto first = leases.acquire(target_string.c_str()); + const auto final = leases.acquire(target_string.c_str()); + int serialization_attempts = 0; + const std::string new_cache = "new-serialized-cache"; + const auto persist = [&] { + ++serialization_attempts; + trtmc::internal::persist_runtime_cache_file(target, new_cache.data(), new_cache.size()); + }; + + leases.release(first, true, persist); + check(serialization_attempts == 0, "non-final lease unexpectedly persisted the cache"); + check(read_file(target) == "old-cache", "non-final lease changed the cache file"); + check(leases.size() == 1, "final cache lease was lost before persistence"); + + leases.release(final, true, persist); + check(serialization_attempts == 1, "final lease did not serialize exactly once"); + check(leases.empty(), "successful final persistence retained its lease"); + check(read_file(target) == new_cache, "atomic cache replacement wrote incorrect bytes"); +} + +void test_write_failure_preserves_final_lease_for_successful_retry() { + TempDirectory directory; + const fs::path parent = directory.path() / "not-created"; + const fs::path target = parent / "cache.bin"; + const std::string target_string = target.string(); + const std::string serialized = "cache-after-write-retry"; + + trtmc::internal::RuntimeCacheLeaseState leases; + const auto lease = leases.acquire(target_string.c_str()); + bool failure_observed = false; + try { + leases.release(lease, true, [&] { + trtmc::internal::persist_runtime_cache_file(target, serialized.data(), + serialized.size()); + }); + } catch (const std::runtime_error& error) { + failure_observed = std::string(error.what()).find("temporary file") != std::string::npos; + } + + check(failure_observed, "runtime-cache write failure was not reported"); + check(leases.size() == 1, "runtime-cache write failure consumed the final lease"); + fs::create_directory(parent); + leases.release(lease, true, [&] { + trtmc::internal::persist_runtime_cache_file(target, serialized.data(), serialized.size()); + }); + check(leases.empty(), "successful write retry retained the final lease"); + check(read_file(target) == serialized, "write retry persisted incorrect cache bytes"); +} + +} // namespace + +int main() { + try { + test_only_final_lease_persists_with_durable_atomic_replace(); + test_write_failure_preserves_final_lease_for_successful_retry(); + std::cout << "Runtime-cache persistence tests passed\n"; + return 0; + } catch (const std::exception& error) { + std::cerr << "Runtime-cache persistence test failed: " << error.what() << '\n'; + return 1; + } +} diff --git a/core/runtime/tests/test_task_api.cpp b/core/runtime/tests/test_task_api.cpp index f302c9fa0e..f7544c8211 100644 --- a/core/runtime/tests/test_task_api.cpp +++ b/core/runtime/tests/test_task_api.cpp @@ -14,6 +14,7 @@ static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); +static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); static_assert(std::is_abstract_v); @@ -70,6 +71,39 @@ void test_robot_observation() { throw std::runtime_error("RobotObservation did not preserve its input spans"); } +void test_structured_video_generation_contract() { + trtmc::AudioResult stereo; + stereo.samples = {0.25F, -0.25F, 0.5F, -0.5F}; + stereo.num_samples = static_cast(stereo.samples.size()); + stereo.sample_rate = 32000; + stereo.channels = 2; + + trtmc::VideoGenerationRequest request; + request.prompt = "ordered references"; + request.mode = trtmc::VideoGenerationMode::kReferenceToVideoAudio; + request.config.video_num_frames = 124; + + trtmc::VideoReferenceInput image; + image.kind = trtmc::VideoReferenceKind::kImage; + trtmc::VideoReferenceInput audio; + audio.kind = trtmc::VideoReferenceKind::kAudio; + audio.audio = stereo; + trtmc::VideoReferenceInput video; + video.kind = trtmc::VideoReferenceKind::kVideo; + video.video.soundtrack = stereo; + request.references = {image, audio, video}; + + if (stereo.channels != 2 || stereo.num_samples != 4 || + request.mode != trtmc::VideoGenerationMode::kReferenceToVideoAudio || + request.config.video_num_frames != 124 || request.references.size() != 3 || + request.references[0].kind != trtmc::VideoReferenceKind::kImage || + request.references[1].kind != trtmc::VideoReferenceKind::kAudio || + request.references[2].kind != trtmc::VideoReferenceKind::kVideo || + request.references[2].video.soundtrack.channels != 2) { + throw std::runtime_error("structured video request did not preserve mode or media order"); + } +} + } // namespace int main() { @@ -87,6 +121,7 @@ int main() { return 1; test_robot_observation(); + test_structured_video_generation_contract(); std::unique_ptr task = std::make_unique(); if (std::string(task->task()) != trtmc::ITextGeneration::kTask) diff --git a/families/minimax_h3/adaln_builder.py b/families/minimax_h3/adaln_builder.py index fc543c3cdd..05c9e823b8 100644 --- a/families/minimax_h3/adaln_builder.py +++ b/families/minimax_h3/adaln_builder.py @@ -12,8 +12,10 @@ import gc import sys +from pathlib import Path import tensorrt as trt +from . import trt_compat from . import graph_ops as op from .config import ( @@ -41,6 +43,7 @@ def checkpoint_keys( return tuple(names) +@op.cleanup_failed_build def build_adaln_precompute_engine( weights: dict, profile: MiniMaxH3Config, @@ -48,19 +51,24 @@ def build_adaln_precompute_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: profile.validate() logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) config = builder.create_builder_config() config.builder_optimization_level = 1 - op.configure_builder(config) - op.configure_workspace( - config, - workspace_bytes, - default_bytes=ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, - ) + op.configure_builder(config, weight_streaming=weight_streaming) + # Let TensorRT-RTX use its native maximum pools by default. An explicit + # override remains useful for constrained build hosts and unit tests. + if workspace_bytes is not None: + op.configure_workspace( + config, + workspace_bytes, + default_bytes=ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, + ) # Match PyTorch's default FP32 matmul policy used by the reference. config.clear_flag(trt.BuilderFlag.TF32) @@ -126,14 +134,21 @@ def build_adaln_precompute_engine( f"timesteps={profile.max_timestep_count}", file=sys.stderr, ) + plan = None + record = None try: - plan = builder.build_serialized_network(network, config) + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) finally: op.release_weight_buffers(network) if consume_weights: weights.clear() - if plan is None: + if output_path is None and plan is None: raise RuntimeError("TensorRT failed to build MiniMax-H3 AdaLN precompute engine") del network, config, builder gc.collect() - return bytes(plan) + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/audio_vae_builder.py b/families/minimax_h3/audio_vae_builder.py new file mode 100644 index 0000000000..2e25705786 --- /dev/null +++ b/families/minimax_h3/audio_vae_builder.py @@ -0,0 +1,771 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native FP32 TensorRT decoder for the MiniMax-H3 audio VAE. + +The released checkpoint uses a BigVGAN decoder. Stereo is represented as two +batch items, so this plan maps denormalized ``[2, 32, frames]`` latents to +``[2, 1, frames * 800]`` mono waveforms. Interleaving the two batch items and +applying the checkpoint's per-channel latent normalization belong to the +pipeline boundary, not to this decoder. + +TensorRT-RTX convolution layers are two-dimensional. The graph therefore +keeps a singleton height dimension and expresses every Conv1d/ConvTranspose1d +as a native 2-D convolution with a ``(1, kernel)`` kernel. No plugin or +framework runtime is required by the serialized engine. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import gc +import math +import sys +from pathlib import Path +from typing import Mapping + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .config import ( + AUDIO_LATENT_FRAMES_MAX, + AUDIO_LATENT_FRAMES_MIN, + AUDIO_LATENT_FRAMES_OPT, + AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, +) + + +trt = trt_compat.get_trt() + + +@dataclass(frozen=True) +class AudioVAEDecoderConfig: + """Static architecture and input shape of one audio decoder plan.""" + + batch_size: int = 2 + latent_channels: int = 32 + latent_frames: int = AUDIO_LATENT_FRAMES_OPT + min_latent_frames: int = AUDIO_LATENT_FRAMES_MIN + max_latent_frames: int = AUDIO_LATENT_FRAMES_MAX + latent_dim: int = 2048 + decoder_dim: int = 1024 + decoder_rates: tuple[int, ...] = (5, 5, 2, 2, 2, 2, 2) + decoder_kernel_sizes: tuple[int, ...] = (9, 9, 4, 4, 4, 4, 4) + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11) + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ( + (1, 3, 5), + (1, 3, 5), + (1, 3, 5), + ) + sampling_rate: int = 32000 + + @property + def num_upsamples(self) -> int: + return len(self.decoder_rates) + + @property + def num_kernels(self) -> int: + return len(self.resblock_kernel_sizes) + + @property + def hop_length(self) -> int: + return math.prod(self.decoder_rates) + + @property + def output_samples(self) -> int: + return self.latent_frames * self.hop_length + + @property + def min_output_samples(self) -> int: + return self.min_latent_frames * self.hop_length + + @property + def max_output_samples(self) -> int: + return self.max_latent_frames * self.hop_length + + def validate(self) -> None: + integer_fields = { + "batch_size": self.batch_size, + "latent_channels": self.latent_channels, + "latent_frames": self.latent_frames, + "min_latent_frames": self.min_latent_frames, + "max_latent_frames": self.max_latent_frames, + "latent_dim": self.latent_dim, + "decoder_dim": self.decoder_dim, + "sampling_rate": self.sampling_rate, + } + if any( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + for value in integer_fields.values() + ): + raise ValueError("MiniMax-H3 AudioVAE dimensions must be positive integers") + if self.batch_size != 2: + raise ValueError("MiniMax-H3 AudioVAE stereo requires batch_size=2") + if not self.min_latent_frames <= self.latent_frames <= self.max_latent_frames: + raise ValueError( + "MiniMax-H3 AudioVAE latent frames must satisfy min <= opt <= max" + ) + if len(self.decoder_rates) != len(self.decoder_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE decoder rates and kernels must align") + if not self.decoder_rates or not self.resblock_kernel_sizes: + raise ValueError("MiniMax-H3 AudioVAE decoder stages and residual kernels are required") + if len(self.resblock_kernel_sizes) != len(self.resblock_dilation_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual kernels and dilations must align") + if any(rate <= 0 for rate in self.decoder_rates): + raise ValueError("MiniMax-H3 AudioVAE decoder rates must be positive") + if any(kernel <= 0 for kernel in self.decoder_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE decoder kernels must be positive") + if any(kernel <= 0 or kernel % 2 == 0 for kernel in self.resblock_kernel_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual kernels must be positive and odd") + if any(not dilations or any(value <= 0 for value in dilations) for dilations in self.resblock_dilation_sizes): + raise ValueError("MiniMax-H3 AudioVAE residual dilations must be positive") + divisor = 1 << self.num_upsamples + if self.decoder_dim % divisor: + raise ValueError("MiniMax-H3 AudioVAE decoder_dim must halve at every stage") + if self.decoder_dim // divisor <= 0: + raise ValueError("MiniMax-H3 AudioVAE final decoder width must be positive") + + +DEFAULT_AUDIO_VAE_DECODER_CONFIG = AudioVAEDecoderConfig() + + +def decoder_config_from_checkpoint( + raw: Mapping[str, object], + *, + latent_frames: int, + min_latent_frames: int | None = None, + max_latent_frames: int | None = None, +) -> AudioVAEDecoderConfig: + """Validate and convert the public Diffusers audio-VAE config.""" + + required = ( + "latent_channels", + "latent_dim", + "decoder_dim", + "decoder_rates", + "decoder_kernel_sizes", + "resblock_kernel_sizes", + "resblock_dilation_sizes", + "sampling_rate", + ) + missing = [name for name in required if name not in raw] + if missing: + raise ValueError(f"MiniMax-H3 AudioVAE config is missing fields: {missing}") + + try: + profile = AudioVAEDecoderConfig( + latent_channels=int(raw["latent_channels"]), + latent_frames=latent_frames, + min_latent_frames=( + latent_frames if min_latent_frames is None else min_latent_frames + ), + max_latent_frames=( + latent_frames if max_latent_frames is None else max_latent_frames + ), + latent_dim=int(raw["latent_dim"]), + decoder_dim=int(raw["decoder_dim"]), + decoder_rates=tuple(int(value) for value in raw["decoder_rates"]), + decoder_kernel_sizes=tuple(int(value) for value in raw["decoder_kernel_sizes"]), + resblock_kernel_sizes=tuple(int(value) for value in raw["resblock_kernel_sizes"]), + resblock_dilation_sizes=tuple( + tuple(int(value) for value in values) + for values in raw["resblock_dilation_sizes"] + ), + sampling_rate=int(raw["sampling_rate"]), + ) + except (TypeError, ValueError) as error: + raise ValueError("MiniMax-H3 AudioVAE config has invalid decoder fields") from error + profile.validate() + + encoder_rates = raw.get("encoder_rates") + if not isinstance(encoder_rates, (list, tuple)): + raise ValueError("MiniMax-H3 AudioVAE config is missing encoder_rates") + try: + encoder_hop = math.prod(int(value) for value in encoder_rates) + except (TypeError, ValueError) as error: + raise ValueError("MiniMax-H3 AudioVAE encoder_rates are invalid") from error + if encoder_hop != profile.hop_length: + raise ValueError( + "MiniMax-H3 AudioVAE encoder/decoder hop mismatch: " + f"encoder={encoder_hop}, decoder={profile.hop_length}" + ) + return profile + + +def checkpoint_keys(profile: AudioVAEDecoderConfig = DEFAULT_AUDIO_VAE_DECODER_CONFIG) -> tuple[str, ...]: + """Return exactly the checkpoint tensors consumed by the decoder plan.""" + + profile.validate() + names = [ + "dec_in_proj.weight", + "dec_in_proj.bias", + "decoder.conv_pre.weight_g", + "decoder.conv_pre.weight_v", + "decoder.conv_pre.bias", + ] + for stage in range(profile.num_upsamples): + prefix = f"decoder.ups.{stage}.0" + names.extend((f"{prefix}.weight_g", f"{prefix}.weight_v", f"{prefix}.bias")) + + for block in range(profile.num_upsamples * profile.num_kernels): + dilation_count = len( + profile.resblock_dilation_sizes[block % profile.num_kernels] + ) + prefix = f"decoder.resblocks.{block}" + for activation in range(2 * dilation_count): + activation_prefix = f"{prefix}.activations.{activation}" + names.extend( + ( + f"{activation_prefix}.act.alpha", + f"{activation_prefix}.act.beta", + f"{activation_prefix}.upsample.filter", + f"{activation_prefix}.downsample.lowpass.filter", + ) + ) + for group in ("convs1", "convs2"): + for index in range(dilation_count): + convolution_prefix = f"{prefix}.{group}.{index}" + names.extend( + ( + f"{convolution_prefix}.weight_g", + f"{convolution_prefix}.weight_v", + f"{convolution_prefix}.bias", + ) + ) + + names.extend( + ( + "decoder.activation_post.act.alpha", + "decoder.activation_post.act.beta", + "decoder.activation_post.upsample.filter", + "decoder.activation_post.downsample.lowpass.filter", + "decoder.conv_post.weight_g", + "decoder.conv_post.weight_v", + ) + ) + return tuple(names) + + +def _require_array(weights: Mapping[str, object], name: str, shape: tuple[int, ...]) -> np.ndarray: + try: + value = np.asarray(weights[name]) + except KeyError as error: + raise ValueError(f"MiniMax-H3 AudioVAE checkpoint is missing tensor: {name}") from error + if tuple(value.shape) != shape: + raise ValueError( + f"MiniMax-H3 AudioVAE tensor {name} has shape {tuple(value.shape)}, expected {shape}" + ) + if value.dtype != np.float32: + raise ValueError( + f"MiniMax-H3 AudioVAE tensor {name} must remain FP32, got {value.dtype}" + ) + return np.ascontiguousarray(value) + + +def fold_weight_norm(weight_g: object, weight_v: object) -> np.ndarray: + """Fold PyTorch ``weight_norm(..., dim=0)`` without changing FP32 dtype.""" + + gain = np.asarray(weight_g) + direction = np.asarray(weight_v) + if gain.dtype != np.float32 or direction.dtype != np.float32: + raise ValueError("MiniMax-H3 AudioVAE weight-normalized tensors must be FP32") + if gain.ndim != direction.ndim or gain.shape[0] != direction.shape[0]: + raise ValueError("MiniMax-H3 AudioVAE weight-normalized tensor shapes do not align") + if gain.shape[1:] != (1,) * (direction.ndim - 1): + raise ValueError("MiniMax-H3 AudioVAE weight_g must normalize dimension zero") + axes = tuple(range(1, direction.ndim)) + norm = np.sqrt( + np.sum(direction * direction, axis=axes, keepdims=True, dtype=np.float32) + ) + if np.any(norm == 0.0): + raise ValueError("MiniMax-H3 AudioVAE weight_v contains a zero-norm filter") + return np.ascontiguousarray(direction * (gain / norm), dtype=np.float32) + + +def _fold_checkpoint_weight( + weights: Mapping[str, object], prefix: str, shape: tuple[int, ...] +) -> np.ndarray: + gain_shape = (shape[0],) + (1,) * (len(shape) - 1) + gain = _require_array(weights, f"{prefix}.weight_g", gain_shape) + direction = _require_array(weights, f"{prefix}.weight_v", shape) + return fold_weight_norm(gain, direction) + + +def _conv1d( + network, + hidden, + weight: np.ndarray, + bias: np.ndarray | None, + *, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + name: str, + owned_weights: list[np.ndarray], +): + out_channels, in_channels_per_group, kernel_size = weight.shape + kernel = np.ascontiguousarray(weight[:, :, None, :], dtype=np.float32) + owned_weights.append(kernel) + if bias is not None: + bias = np.ascontiguousarray(bias, dtype=np.float32) + owned_weights.append(bias) + layer = network.add_convolution_nd( + hidden, + out_channels, + (1, kernel_size), + kernel, + bias, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE convolution {name}") + layer.name = name + layer.stride_nd = (1, stride) + layer.padding_nd = (0, padding) + layer.dilation_nd = (1, dilation) + layer.num_groups = groups + expected_inputs = int(tuple(hidden.shape)[1]) // groups + if in_channels_per_group != expected_inputs: + raise ValueError( + f"MiniMax-H3 AudioVAE convolution {name} expects {in_channels_per_group} " + f"channels per group, got {expected_inputs}" + ) + return layer.get_output(0) + + +def _deconv1d( + network, + hidden, + weight: np.ndarray, + bias: np.ndarray | None, + *, + out_channels: int, + stride: int = 1, + padding: int = 0, + groups: int = 1, + name: str, + owned_weights: list[np.ndarray], +): + in_channels, out_channels_per_group, kernel_size = weight.shape + kernel = np.ascontiguousarray(weight[:, :, None, :], dtype=np.float32) + owned_weights.append(kernel) + if bias is not None: + bias = np.ascontiguousarray(bias, dtype=np.float32) + owned_weights.append(bias) + if in_channels != int(tuple(hidden.shape)[1]): + raise ValueError( + f"MiniMax-H3 AudioVAE deconvolution {name} expects {in_channels} input channels, " + f"got {tuple(hidden.shape)[1]}" + ) + if out_channels_per_group * groups != out_channels: + raise ValueError( + f"MiniMax-H3 AudioVAE deconvolution {name} has an invalid grouped output width" + ) + layer = network.add_deconvolution_nd( + hidden, + out_channels, + (1, kernel_size), + kernel, + bias, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE deconvolution {name}") + layer.name = name + layer.stride_nd = (1, stride) + layer.padding_nd = (0, padding) + layer.num_groups = groups + return layer.get_output(0) + + +def _replicate_pad(network, hidden, left: int, right: int, *, name: str): + batch, channels, height, _width = tuple(int(value) for value in hidden.shape) + layer = network.add_slice( + hidden, + (0, 0, 0, -left), + (batch, channels, height, 1), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE replicate pad {name}") + layer.name = name + layer.mode = trt.SampleMode.CLAMP + shape = network.add_shape(hidden).get_output(0) + delta = op.constant( + network, + np.asarray([0, 0, 0, left + right], dtype=np.int64), + dtype=np.int64, + ) + size = network.add_elementwise(shape, delta, trt.ElementWiseOperation.SUM).get_output(0) + layer.set_input(2, size) + return layer.get_output(0) + + +def _crop_time(network, hidden, left: int, right: int, *, name: str): + batch, channels, height, _width = tuple(int(value) for value in hidden.shape) + layer = network.add_slice( + hidden, + (0, 0, 0, left), + (batch, channels, height, 1), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 AudioVAE crop {name}") + layer.name = name + shape = network.add_shape(hidden).get_output(0) + delta = op.constant( + network, + np.asarray([0, 0, 0, -(left + right)], dtype=np.int64), + dtype=np.int64, + ) + size = network.add_elementwise(shape, delta, trt.ElementWiseOperation.SUM).get_output(0) + layer.set_input(2, size) + return layer.get_output(0) + + +def _snake_beta(network, hidden, weights: Mapping[str, object], prefix: str): + channels = int(tuple(hidden.shape)[1]) + alpha = _require_array(weights, f"{prefix}.alpha", (channels,)).reshape(1, channels, 1, 1) + beta = _require_array(weights, f"{prefix}.beta", (channels,)).reshape(1, channels, 1, 1) + alpha = network.add_unary(op.weight_constant(network, alpha), trt.UnaryOperation.EXP).get_output(0) + beta = network.add_unary(op.weight_constant(network, beta), trt.UnaryOperation.EXP).get_output(0) + phase = network.add_elementwise(hidden, alpha, trt.ElementWiseOperation.PROD).get_output(0) + sine = network.add_unary(phase, trt.UnaryOperation.SIN).get_output(0) + sine_square = network.add_elementwise(sine, sine, trt.ElementWiseOperation.PROD).get_output(0) + epsilon = op.constant(network, np.full((1, 1, 1, 1), 1.0e-9, np.float32)) + denominator = network.add_elementwise(beta, epsilon, trt.ElementWiseOperation.SUM).get_output(0) + periodic = network.add_elementwise( + sine_square, denominator, trt.ElementWiseOperation.DIV + ).get_output(0) + return network.add_elementwise(hidden, periodic, trt.ElementWiseOperation.SUM).get_output(0) + + +def _alias_free_activation( + network, + hidden, + weights: Mapping[str, object], + prefix: str, + *, + owned_weights: list[np.ndarray], +): + ratio = 2 + kernel_size = 12 + channels = int(tuple(hidden.shape)[1]) + pad = kernel_size // ratio - 1 + pad_left = pad * ratio + (kernel_size - ratio) // 2 + pad_right = pad * ratio + (kernel_size - ratio + 1) // 2 + + padded = _replicate_pad(network, hidden, pad, pad, name=f"{prefix}.upsample.pad") + upsample_filter = _require_array(weights, f"{prefix}.upsample.filter", (1, 1, kernel_size)) + upsample_filter = np.broadcast_to( + upsample_filter, (channels, 1, kernel_size) + ).copy() + hidden = _deconv1d( + network, + padded, + upsample_filter, + None, + out_channels=channels, + stride=ratio, + groups=channels, + name=f"{prefix}.upsample", + owned_weights=owned_weights, + ) + scale = op.constant(network, np.full((1, 1, 1, 1), float(ratio), np.float32)) + hidden = network.add_elementwise(hidden, scale, trt.ElementWiseOperation.PROD).get_output(0) + hidden = _crop_time( + network, + hidden, + pad_left, + pad_right, + name=f"{prefix}.upsample.crop", + ) + hidden = _snake_beta(network, hidden, weights, f"{prefix}.act") + + even = kernel_size % 2 == 0 + down_left = kernel_size // 2 - int(even) + down_right = kernel_size // 2 + hidden = _replicate_pad( + network, + hidden, + down_left, + down_right, + name=f"{prefix}.downsample.pad", + ) + downsample_filter = _require_array( + weights, f"{prefix}.downsample.lowpass.filter", (1, 1, kernel_size) + ) + downsample_filter = np.broadcast_to( + downsample_filter, (channels, 1, kernel_size) + ).copy() + return _conv1d( + network, + hidden, + downsample_filter, + None, + stride=ratio, + groups=channels, + name=f"{prefix}.downsample", + owned_weights=owned_weights, + ) + + +def _amp_block( + network, + hidden, + weights: Mapping[str, object], + prefix: str, + *, + kernel_size: int, + dilations: tuple[int, ...], + owned_weights: list[np.ndarray], +): + channels = int(tuple(hidden.shape)[1]) + for index, dilation in enumerate(dilations): + residual = _alias_free_activation( + network, + hidden, + weights, + f"{prefix}.activations.{2 * index}", + owned_weights=owned_weights, + ) + conv1_prefix = f"{prefix}.convs1.{index}" + conv1_weight = _fold_checkpoint_weight( + weights, conv1_prefix, (channels, channels, kernel_size) + ) + conv1_bias = _require_array(weights, f"{conv1_prefix}.bias", (channels,)) + residual = _conv1d( + network, + residual, + conv1_weight, + conv1_bias, + padding=(kernel_size * dilation - dilation) // 2, + dilation=dilation, + name=conv1_prefix, + owned_weights=owned_weights, + ) + residual = _alias_free_activation( + network, + residual, + weights, + f"{prefix}.activations.{2 * index + 1}", + owned_weights=owned_weights, + ) + conv2_prefix = f"{prefix}.convs2.{index}" + conv2_weight = _fold_checkpoint_weight( + weights, conv2_prefix, (channels, channels, kernel_size) + ) + conv2_bias = _require_array(weights, f"{conv2_prefix}.bias", (channels,)) + residual = _conv1d( + network, + residual, + conv2_weight, + conv2_bias, + padding=(kernel_size - 1) // 2, + name=conv2_prefix, + owned_weights=owned_weights, + ) + hidden = network.add_elementwise( + hidden, residual, trt.ElementWiseOperation.SUM + ).get_output(0) + return hidden + + +@op.cleanup_failed_build +def build_audio_vae_decoder_engine( + weights: dict, + profile: AudioVAEDecoderConfig = DEFAULT_AUDIO_VAE_DECODER_CONFIG, + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the released MiniMax-H3 BigVGAN decoder as a native FP32 plan.""" + + profile.validate() + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, + ) + missing = sorted(set(checkpoint_keys(profile)) - set(weights)) + if missing: + raise ValueError(f"MiniMax-H3 AudioVAE checkpoint is missing tensors: {missing}") + owned_weights: list[np.ndarray] = [] + + latents = network.add_input( + "audio_latents", + trt.float32, + (profile.batch_size, profile.latent_channels, -1), + ) + shape_profile = builder.create_optimization_profile() + shape_profile.set_shape( + "audio_latents", + (profile.batch_size, profile.latent_channels, profile.min_latent_frames), + (profile.batch_size, profile.latent_channels, profile.latent_frames), + (profile.batch_size, profile.latent_channels, profile.max_latent_frames), + ) + profile_index = config.add_optimization_profile(shape_profile) + if profile_index != 0: + raise RuntimeError("TensorRT rejected MiniMax-H3 AudioVAE dynamic shape profile") + expanded = network.add_shuffle(latents) + expanded.reshape_dims = ( + profile.batch_size, + profile.latent_channels, + 1, + -1, + ) + hidden = expanded.get_output(0) + + dec_in_weight = _require_array( + weights, + "dec_in_proj.weight", + (profile.latent_dim, profile.latent_channels, 1), + ) + dec_in_bias = _require_array(weights, "dec_in_proj.bias", (profile.latent_dim,)) + hidden = _conv1d( + network, + hidden, + dec_in_weight, + dec_in_bias, + name="dec_in_proj", + owned_weights=owned_weights, + ) + + conv_pre_weight = _fold_checkpoint_weight( + weights, + "decoder.conv_pre", + (profile.decoder_dim, profile.latent_dim, 7), + ) + conv_pre_bias = _require_array(weights, "decoder.conv_pre.bias", (profile.decoder_dim,)) + hidden = _conv1d( + network, + hidden, + conv_pre_weight, + conv_pre_bias, + padding=3, + name="decoder.conv_pre", + owned_weights=owned_weights, + ) + + for stage, (rate, kernel_size) in enumerate( + zip(profile.decoder_rates, profile.decoder_kernel_sizes) + ): + input_channels = profile.decoder_dim // (1 << stage) + output_channels = profile.decoder_dim // (1 << (stage + 1)) + upsample_prefix = f"decoder.ups.{stage}.0" + upsample_weight = _fold_checkpoint_weight( + weights, + upsample_prefix, + (input_channels, output_channels, kernel_size), + ) + upsample_bias = _require_array(weights, f"{upsample_prefix}.bias", (output_channels,)) + hidden = _deconv1d( + network, + hidden, + upsample_weight, + upsample_bias, + out_channels=output_channels, + stride=rate, + padding=(kernel_size - rate) // 2, + name=upsample_prefix, + owned_weights=owned_weights, + ) + + block_outputs = [] + for kernel_index, (resblock_kernel, dilations) in enumerate( + zip(profile.resblock_kernel_sizes, profile.resblock_dilation_sizes) + ): + block_index = stage * profile.num_kernels + kernel_index + block_outputs.append( + _amp_block( + network, + hidden, + weights, + f"decoder.resblocks.{block_index}", + kernel_size=resblock_kernel, + dilations=dilations, + owned_weights=owned_weights, + ) + ) + combined = block_outputs[0] + for block_output in block_outputs[1:]: + combined = network.add_elementwise( + combined, block_output, trt.ElementWiseOperation.SUM + ).get_output(0) + divisor = op.constant( + network, np.full((1, 1, 1, 1), float(profile.num_kernels), np.float32) + ) + hidden = network.add_elementwise( + combined, divisor, trt.ElementWiseOperation.DIV + ).get_output(0) + + hidden = _alias_free_activation( + network, + hidden, + weights, + "decoder.activation_post", + owned_weights=owned_weights, + ) + final_channels = profile.decoder_dim // (1 << profile.num_upsamples) + conv_post_weight = _fold_checkpoint_weight( + weights, + "decoder.conv_post", + (1, final_channels, 7), + ) + hidden = _conv1d( + network, + hidden, + conv_post_weight, + None, + padding=3, + name="decoder.conv_post", + owned_weights=owned_weights, + ) + clamp = network.add_activation(hidden, trt.ActivationType.CLIP) + if clamp is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 AudioVAE output clamp") + clamp.alpha = -1.0 + clamp.beta = 1.0 + output = network.add_shuffle(clamp.get_output(0)) + output.reshape_dims = (profile.batch_size, 1, -1) + decoded = output.get_output(0) + decoded.name = "decoded_audio" + network.mark_output(decoded) + op.validate_native_network(network, expected_attentions=0, label="AudioVAE decoder") + + print( + "[minimax-h3] building native FP32 AudioVAE decoder: " + f"latents={profile.min_latent_frames}/{profile.latent_frames}/" + f"{profile.max_latent_frames}, samples={profile.min_output_samples}/" + f"{profile.output_samples}/{profile.max_output_samples}, " + f"upsamples={profile.num_upsamples}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 AudioVAE decoder") + del owned_weights, network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/checkpoint.py b/families/minimax_h3/checkpoint.py index 8221812c45..c9f489c548 100644 --- a/families/minimax_h3/checkpoint.py +++ b/families/minimax_h3/checkpoint.py @@ -24,56 +24,45 @@ def __array_finalize__(self, source) -> None: self._tensor_owner = getattr(source, "_tensor_owner", None) -def load_component_state_dict(component_dir: str | Path) -> dict[str, Any]: - """Load a safetensors component without materializing duplicate tensors.""" - - from safetensors.torch import load_file - - root = Path(component_dir) - indexes = sorted(root.glob("*.safetensors.index.json")) - if indexes: - if len(indexes) != 1: - raise ValueError(f"Expected one safetensors index in {root}, found {len(indexes)}") - weight_map = json.loads(indexes[0].read_text())["weight_map"] - paths = [root / name for name in sorted(set(weight_map.values()))] - else: - paths = sorted(root.glob("*.safetensors")) - if not paths: - raise FileNotFoundError(f"No safetensors checkpoint found in {root}") - - state: dict[str, Any] = {} - for path in paths: - for name, tensor in load_file(path, device="cpu").items(): - if name in state: - raise ValueError(f"Duplicate MiniMax-H3 tensor {name!r}") - state[name] = tensor - return state - - def load_selected_component_state_dict( component_dir: str | Path, names: Iterable[str] ) -> dict[str, Any]: - """Load only selected indexed tensors, avoiding unused H3 language/vision weights.""" + """Load selected tensors from one indexed or unsharded component.""" from safetensors import safe_open root = Path(component_dir) indexes = sorted(root.glob("*.safetensors.index.json")) - if len(indexes) != 1: - raise ValueError(f"Selective loading requires one safetensors index in {root}") - weight_map = json.loads(indexes[0].read_text())["weight_map"] requested = tuple(names) - missing = sorted(set(requested) - set(weight_map)) - if missing: - raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") - by_file: dict[str, list[str]] = {} - for name in requested: - by_file.setdefault(weight_map[name], []).append(name) state: dict[str, Any] = {} - for filename, tensor_names in sorted(by_file.items()): - with safe_open(root / filename, framework="pt", device="cpu") as reader: - for name in tensor_names: - state[name] = reader.get_tensor(name) + if indexes: + if len(indexes) != 1: + raise ValueError(f"Selective loading requires one safetensors index in {root}") + weight_map = json.loads(indexes[0].read_text())["weight_map"] + missing = sorted(set(requested) - set(weight_map)) + if missing: + raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") + by_file: dict[str, list[str]] = {} + for name in requested: + by_file.setdefault(weight_map[name], []).append(name) + for filename, tensor_names in sorted(by_file.items()): + with safe_open(root / filename, framework="pt", device="cpu") as reader: + for name in tensor_names: + state[name] = reader.get_tensor(name) + return state + + paths = sorted(root.glob("*.safetensors")) + if len(paths) != 1: + raise ValueError( + f"Selective loading requires one safetensors index or one unsharded file in {root}" + ) + with safe_open(paths[0], framework="pt", device="cpu") as reader: + available = set(reader.keys()) + missing = sorted(set(requested) - available) + if missing: + raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") + for name in requested: + state[name] = reader.get_tensor(name) return state @@ -104,12 +93,6 @@ def validate_component_key_partition( ) -def require_keys(state: dict[str, Any], names: Iterable[str]) -> None: - missing = sorted(set(names) - set(state)) - if missing: - raise ValueError(f"MiniMax-H3 checkpoint is missing tensors: {missing}") - - def numpy_state(state: dict[str, Any]) -> dict[str, Any]: """Expose CPU tensors to NumPy without expanding checkpoint-native BF16. diff --git a/families/minimax_h3/config.py b/families/minimax_h3/config.py index d038603de5..efdd8c067c 100644 --- a/families/minimax_h3/config.py +++ b/families/minimax_h3/config.py @@ -3,11 +3,12 @@ """Validated native TensorRT profiles for MiniMax-H3. -The default profile is the 124-frame, 1344x768 shape used by the public -Sol-Engine H3 benchmark. Structural row counts are explicit because prompt -packing is part of the engine ABI and must match the Hugging Face reference. -Text rows are dynamic; ``text_rows`` remains the maximum for compatible -bundle metadata. +Production uses one dynamic profile for the complete released 5--15 second +T2VA and FL2VA media envelope. It supports dynamic prompt lengths, the public +continuous 1:4--4:1 canvas resolver, and the documented explicit performance +canvases, including the optional landscape 480p super-resolution source geometry. +Structural row counts are explicit because prompt/media packing is part of +the engine ABI and must match the Hugging Face reference. """ from __future__ import annotations @@ -16,15 +17,49 @@ TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +VISION_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES = 64 << 30 DENOISER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 VAE_TILE_DECODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES = 96 << 30 +AUDIO_LATENT_FRAMES_MIN = 207 +AUDIO_LATENT_FRAMES_OPT = 207 +AUDIO_LATENT_FRAMES_MAX = 575 +VIDEO_NUM_FRAMES_MIN = 124 +VIDEO_NUM_FRAMES_OPT = 124 +VIDEO_NUM_FRAMES_MAX = 345 +VIDEO_ROWS_MIN = 14_985 +VIDEO_ROWS_OPT = 37_296 +VIDEO_ROWS_MAX = 108_576 +CANVAS_MULTIPLE = 32 +CANVAS_SHORT_EDGE = 768 +CANVAS_MAX_PIXELS = 768 * 1344 +CANVAS_MIN_ASPECT_RATIO = 0.25 +CANVAS_MAX_ASPECT_RATIO = 4.0 +# Extra explicit performance canvases, stored as (height, width). The TensorRT +# runtime remains a finite allowlist: these orientations are in addition to, +# not a replacement for, the 95 resolver-produced canvases. The landscape +# 480p canvas is the source geometry for the optional native SR stage. +NATIVE_EXPLICIT_CANVAS_SIZES = ( + (480, 864), + (544, 960), + (960, 544), +) +# The RTX path builds each plan in a fresh process, so one conservative +# workspace and runtime budget cover every stage without coupling the public +# artifact to a particular workstation identity. +RTX_STAGED_WORKSPACE_BYTES = 16 << 30 +RTX_WEIGHT_STREAMING_BUDGET_BYTES = 32 << 30 +TRT_DEFAULT_WORKSPACE_POLICY = "trt_default_max" DEFAULT_WORKSPACE_LIMIT_BYTES = { "text_encoder.plan": TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES, + "vision_encoder.plan": VISION_ENCODER_DEFAULT_WORKSPACE_BYTES, "adaln_precompute.plan": ADALN_PRECOMPUTE_DEFAULT_WORKSPACE_BYTES, - "denoiser.plan": DENOISER_DEFAULT_WORKSPACE_BYTES, + "fl2va_keyframe_vae_encoder.plan": KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, "vae_tile_decoder.plan": VAE_TILE_DECODER_DEFAULT_WORKSPACE_BYTES, + "audio_vae_decoder.plan": AUDIO_VAE_DECODER_DEFAULT_WORKSPACE_BYTES, } FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES = ( @@ -34,32 +69,30 @@ ) -def native_plan_filenames(*, first_block_cache: bool) -> tuple[str, ...]: - """Return the exact plan set selected by the native denoiser profile.""" +def native_plan_filenames() -> tuple[str, ...]: + """Return the singular original-weight dense FirstBlockCache plan set.""" - if not isinstance(first_block_cache, bool): - raise ValueError("MiniMax-H3 first_block_cache must be a boolean") - denoiser_plans = ( - FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES if first_block_cache else ("denoiser.plan",) - ) return ( "text_encoder.plan", + "vision_encoder.plan", "adaln_precompute.plan", - *denoiser_plans, + *FIRST_BLOCK_CACHE_DENOISER_PLAN_FILENAMES, + "fl2va_keyframe_vae_encoder.plan", "vae_tile_decoder.plan", + "audio_vae_decoder.plan", ) -def default_workspace_limit_bytes(*, first_block_cache: bool) -> dict[str, int]: - """Return per-plan tactic workspace limits for one denoiser layout.""" +def default_workspace_limit_bytes() -> dict[str, int | str]: + """Return per-plan tactic workspace limits for the dense FBC layout.""" return { filename: ( - DENOISER_DEFAULT_WORKSPACE_BYTES - if filename.startswith("denoiser_") or filename == "denoiser.plan" + TRT_DEFAULT_WORKSPACE_POLICY + if filename.startswith("denoiser_") or filename.startswith("adaln_") else DEFAULT_WORKSPACE_LIMIT_BYTES[filename] ) - for filename in native_plan_filenames(first_block_cache=first_block_cache) + for filename in native_plan_filenames() } @@ -89,7 +122,11 @@ class MiniMaxH3Config: timestep_embed_dim: int = 2688 rope_freq_dim: int = 16 norm_eps: float = 1.0e-5 + min_video_rows: int = 37296 + opt_video_rows: int = 37296 video_rows: int = 37296 + min_audio_rows: int = 414 + opt_audio_rows: int = 414 audio_rows: int = 414 min_text_rows: int = 1 opt_text_rows: int = 128 @@ -97,7 +134,7 @@ class MiniMaxH3Config: padded_sequence_length: int = 38247 max_timestep_count: int = 4 context_parallel_size: int = 1 - first_block_cache: bool = False + first_block_cache: bool = True @property def sequence_length(self) -> int: @@ -105,11 +142,27 @@ def sequence_length(self) -> int: @property def min_sequence_length(self) -> int: - return self.video_rows + self.audio_rows + self.min_text_rows + return self.min_video_rows + self.min_audio_rows + self.min_text_rows @property def opt_sequence_length(self) -> int: - return self.video_rows + self.audio_rows + self.opt_text_rows + return self.opt_video_rows + self.opt_audio_rows + self.opt_text_rows + + @property + def video_row_profile(self) -> tuple[int, int, int]: + return self.min_video_rows, self.opt_video_rows, self.video_rows + + @property + def audio_row_profile(self) -> tuple[int, int, int]: + return self.min_audio_rows, self.opt_audio_rows, self.audio_rows + + @property + def packed_row_profile(self) -> tuple[int, int, int]: + return self.min_sequence_length, self.opt_sequence_length, self.sequence_length + + @property + def text_row_profile(self) -> tuple[int, int, int]: + return self.min_text_rows, self.opt_text_rows, self.text_rows @property def padding_rows(self) -> int: @@ -137,6 +190,12 @@ def validate(self) -> None: raise ValueError("MiniMax-H3 attention width must exceed its residual width") if not 1 <= self.min_text_rows <= self.opt_text_rows <= self.text_rows: raise ValueError("MiniMax-H3 text rows must satisfy 1 <= min <= opt <= max") + if not 0 < self.min_video_rows <= self.opt_video_rows <= self.video_rows: + raise ValueError("MiniMax-H3 video rows must satisfy 0 < min <= opt <= max") + if not 0 < self.min_audio_rows <= self.opt_audio_rows <= self.audio_rows: + raise ValueError("MiniMax-H3 audio rows must satisfy 0 < min <= opt <= max") + if any(rows % 2 for rows in self.audio_row_profile): + raise ValueError("MiniMax-H3 audio_rows must contain two equal stereo channels") if self.sequence_length != self.padded_sequence_length: raise ValueError( "MiniMax-H3 requires no packed-sequence padding: " @@ -149,3 +208,18 @@ def validate(self) -> None: SOL_ENGINE_1344X768_124F = MiniMaxH3Config() + +# Each production FirstBlockCache engine has one profile covering the complete +# public 5--15 second envelope, including dynamic prompts and keyframes. +# The released local pipeline aligns requested frame counts to ``17 * n + 5``. +# At 24 fps its supported 5--15 second endpoints are therefore 124 and 345 +# frames. Video tokens use 1,008 rows per latent frame at 1344x768; audio is +# packed as two stereo row groups of 207 through 575 latent frames. +SOL_ENGINE_1344X768_124_TO_345F = MiniMaxH3Config( + min_video_rows=VIDEO_ROWS_MIN, + opt_video_rows=VIDEO_ROWS_OPT, + video_rows=VIDEO_ROWS_MAX, + audio_rows=1150, + text_rows=2641, + padded_sequence_length=112367, +) diff --git a/families/minimax_h3/delivery.py b/families/minimax_h3/delivery.py new file mode 100644 index 0000000000..d3d04603b2 --- /dev/null +++ b/families/minimax_h3/delivery.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time source resolution for the single quantized H3 delivery.""" + +from __future__ import annotations + +from pathlib import Path + + +COMFY_REPOSITORY = "Comfy-Org/MiniMax-H3" +COMFY_REVISION = "4cc1d817b6184899b41293954329f576cb5ae86b" +SR_SOURCE_SHAPE = (480, 864) +QUANTIZED_SOURCES = { + "quantized_transformer": "diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors", + "quantized_ref_transformer": "diffusion_models/minimax_h3_ref2va_int8_convrot.safetensors", + "quantized_text_encoder": "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", +} + + +def resolve_quantized_sources(model_dir: Path, options: dict) -> dict[str, Path]: + """Prefer explicit/offline files, otherwise use the pinned public HF cache.""" + + from huggingface_hub import hf_hub_download + + for option in QUANTIZED_SOURCES: + if options.get(option) and not Path(options[option]).is_file(): + raise FileNotFoundError(f"MiniMax-H3 {option} checkpoint is missing: {options[option]}") + sources = {} + for option, filename in QUANTIZED_SOURCES.items(): + explicit = options.get(option) + local = Path(model_dir) / filename + if explicit: + source = Path(explicit).absolute() + elif local.is_file(): + source = local.absolute() + else: + source = Path(hf_hub_download(COMFY_REPOSITORY, filename, revision=COMFY_REVISION)) + if not source.is_file(): + raise FileNotFoundError(f"MiniMax-H3 {option} checkpoint is missing: {source}") + # Preserve the public filename when HF cache entries are symlinks. + sources[option] = source.absolute() + return sources + + +def resolve_super_resolution_sources(options: dict) -> tuple[Path | None, Path | None]: + """Only the explicit build flag may enable the fixed-base SR delivery.""" + + enabled = options.get("super_resolution", False) + keys = ("super_resolution_model", "super_resolution_weak_model") + if not enabled: + if any(options.get(key) for key in keys): + raise ValueError("MiniMax-H3 SR checkpoint overrides require super_resolution=true") + return None, None + from torch.hub import download_url_to_file, get_dir + + result = [] + for key, filename in zip(keys, ("realesr-general-x4v3.pth", "realesr-general-wdn-x4v3.pth")): + explicit = options.get(key) + path = Path(explicit).absolute() if explicit else Path(get_dir()) / "checkpoints" / filename + if explicit and not path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 SR checkpoint is missing: {path}") + if not path.is_file(): + path.parent.mkdir(parents=True, exist_ok=True) + download_url_to_file( + "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/" + filename, + str(path), + ) + result.append(path.absolute()) + return result[0], result[1] diff --git a/families/minimax_h3/dit_builder.py b/families/minimax_h3/dit_builder.py index 0113411375..94e781df8e 100644 --- a/families/minimax_h3/dit_builder.py +++ b/families/minimax_h3/dit_builder.py @@ -8,10 +8,12 @@ import gc import math import sys +from pathlib import Path import numpy as np import tensorrt as trt +from . import trt_compat from . import graph_ops as op from .config import ( @@ -85,9 +87,7 @@ def tail_checkpoint_keys( return _block_checkpoint_keys(range(1, profile.num_layers)) -def finish_checkpoint_keys( - profile: MiniMaxH3Config = SOL_ENGINE_1344X768_124F, -) -> tuple[str, ...]: +def finish_checkpoint_keys() -> tuple[str, ...]: """Weights used by the final norm and modality-specific projections.""" return ( @@ -107,11 +107,11 @@ def checkpoint_keys( return ( *head_checkpoint_keys(profile), *tail_checkpoint_keys(profile), - *finish_checkpoint_keys(profile), + *finish_checkpoint_keys(), ) -def _slice_modulation(network, selected, index: int, rows: int, width: int): +def _slice_modulation(network, selected, index: int, width: int): value = op.dynamic_slice(network, selected, (0, index, 0), (None, 1, width)) reshape = network.add_shuffle(value) reshape.reshape_dims = (-1, width) @@ -129,7 +129,7 @@ def _per_head_norm(network, tensor, weight, profile: MiniMaxH3Config, rows: int) return flatten.get_output(0) -def _rope_tables(network, position_ids, profile: MiniMaxH3Config, rows: int): +def _rope_tables(network, position_ids, profile: MiniMaxH3Config): positions = op.cast(network, position_ids, trt.float32) position_shape = network.add_shuffle(positions) position_shape.reshape_dims = (-1, 3, 1) @@ -152,9 +152,9 @@ def _rope_tables(network, position_ids, profile: MiniMaxH3Config, rows: int): def _native_attention(network, q, k, v, *, rows: int, profile: MiniMaxH3Config, name: str): - q4 = op.rows_to_heads(network, q, rows, profile.num_heads, profile.head_dim) - k4 = op.rows_to_heads(network, k, rows, profile.num_heads, profile.head_dim) - v4 = op.rows_to_heads(network, v, rows, profile.num_heads, profile.head_dim) + q4 = op.rows_to_heads(network, q, profile.num_heads, profile.head_dim) + k4 = op.rows_to_heads(network, k, profile.num_heads, profile.head_dim) + v4 = op.rows_to_heads(network, v, profile.num_heads, profile.head_dim) scale = op.constant( network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(profile.head_dim), dtype=np.float32), @@ -168,7 +168,7 @@ def _native_attention(network, q, k, v, *, rows: int, profile: MiniMaxH3Config, attention.metadata = f"trtmc.native_op=IAttention;source={name}" attention.get_output(0).name = f"{name}.output" attention.decomposable = False - return op.heads_to_rows(network, attention.get_output(0), rows, profile.attention_size) + return op.heads_to_rows(network, attention.get_output(0), profile.attention_size) def _attention_block( @@ -181,8 +181,15 @@ def _attention_block( *, cos=None, sin=None, + consume_weights: bool = False, ): - q, k, v = op.fused_qkv(network, hidden, weights, f"{prefix}.attn") + q, k, v = op.fused_qkv( + network, + hidden, + weights, + f"{prefix}.attn", + consume_weights=consume_weights, + ) q = _per_head_norm(network, q, weights[f"{prefix}.attn.norm_q.weight"], profile, rows) k = _per_head_norm(network, k, weights[f"{prefix}.attn.norm_k.weight"], profile, rows) if cos is not None: @@ -192,7 +199,6 @@ def _attention_block( q, cos, sin, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, rotary_dim=rotary_dim, @@ -202,7 +208,6 @@ def _attention_block( k, cos, sin, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, rotary_dim=rotary_dim, @@ -212,7 +217,6 @@ def _attention_block( q, k, v, - rows=rows, heads=profile.num_heads, head_dim=profile.head_dim, name=f"{prefix}.attn.native_attention", @@ -230,7 +234,14 @@ def _attention_block( return op.linear(network, attended, weights[f"{prefix}.attn.to_out.0.weight"]) -def _refine_text(network, text, weights, profile: MiniMaxH3Config): +def _refine_text( + network, + text, + weights, + profile: MiniMaxH3Config, + *, + consume_weights: bool = False, +): hidden = op.linear( network, text, weights["context_embedder.weight"], weights["context_embedder.bias"] ) @@ -244,7 +255,15 @@ def _refine_text(network, text, weights, profile: MiniMaxH3Config): profile.hidden_size, profile.norm_eps, ) - update = _attention_block(network, normalized, weights, prefix, profile, rows) + update = _attention_block( + network, + normalized, + weights, + prefix, + profile, + rows, + consume_weights=consume_weights, + ) hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) normalized = op.rms_norm( network, @@ -270,10 +289,25 @@ def _refine_text(network, text, weights, profile: MiniMaxH3Config): ) -def _packed_hidden(network, video, audio, text, weights, profile: MiniMaxH3Config): +def _packed_hidden( + network, + video, + audio, + text, + weights, + profile: MiniMaxH3Config, + *, + consume_weights: bool = False, +): """Project and pack text | audio | video exactly like the Diffusers model.""" - text_hidden = _refine_text(network, text, weights, profile) + text_hidden = _refine_text( + network, + text, + weights, + profile, + consume_weights=consume_weights, + ) audio_hidden = op.linear( network, audio, weights["audio_proj_in.weight"], weights["audio_proj_in.bias"], bf16=False ) @@ -297,6 +331,8 @@ def _transformer_block( weights, profile: MiniMaxH3Config, index: int, + *, + consume_weights: bool = False, ): """Add one native H3 transformer block and return its residual stream.""" @@ -304,7 +340,7 @@ def _transformer_block( prefix = f"transformer_blocks.{index}" selected = op.gather_rows(network, block_modulation, adaln_indices) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - _slice_modulation(network, selected, part, rows, profile.hidden_size) for part in range(6) + _slice_modulation(network, selected, part, profile.hidden_size) for part in range(6) ) normalized = op.rms_norm( network, @@ -323,6 +359,7 @@ def _transformer_block( rows, cos=cos, sin=sin, + consume_weights=consume_weights, ) hidden = op.gated_residual(network, hidden, update, gate_msa) @@ -345,10 +382,9 @@ def _transformer_block( def _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile): - rows = -1 selected = op.gather_rows(network, final_modulation, timestep_indices) - final_shift = _slice_modulation(network, selected, 0, rows, profile.hidden_size) - final_scale = _slice_modulation(network, selected, 1, rows, profile.hidden_size) + final_shift = _slice_modulation(network, selected, 0, profile.hidden_size) + final_scale = _slice_modulation(network, selected, 1, profile.hidden_size) hidden = op.rms_norm( network, hidden, weights["norm_out.norm.weight"], profile.hidden_size, profile.norm_eps ) @@ -356,37 +392,13 @@ def _final_hidden(network, hidden, timestep_indices, final_modulation, weights, return op.cast(network, hidden, trt.float32) -def _mark_full_velocity_outputs(network, hidden, weights): - """Preserve the original monolithic plan's full packed-sequence outputs.""" - - video_all = op.linear( - network, hidden, weights["proj_out.weight"], weights["proj_out.bias"], bf16=False - ) - audio_all = op.linear( - network, - hidden, - weights["audio_proj_out.weight"], - weights["audio_proj_out.bias"], - bf16=False, - ) - video_all.name = "video_velocity" - audio_all.name = "audio_velocity" - network.mark_output(video_all) - network.mark_output(audio_all) - - -def _mark_sliced_velocity_outputs(network, hidden, weights, profile: MiniMaxH3Config): +def _mark_sliced_velocity_outputs(network, hidden, weights, video_reference, audio_reference): """Project only rows consumed by the audio and video scheduler updates.""" - audio_hidden = op.slice_rows_from_end( - network, - hidden, - offset=profile.audio_rows + profile.video_rows, - rows=profile.audio_rows, - ) - video_hidden = op.slice_rows_from_end( - network, hidden, offset=profile.video_rows, rows=profile.video_rows + audio_hidden = op.slice_rows_like_from_end( + network, hidden, audio_reference, trailing_reference=video_reference ) + video_hidden = op.slice_rows_like_from_end(network, hidden, video_reference) video = op.linear( network, video_hidden, @@ -407,18 +419,27 @@ def _mark_sliced_velocity_outputs(network, hidden, weights, profile: MiniMaxH3Co network.mark_output(audio) -def _native_builder(verbose: bool, workspace_bytes: int | None): +def _native_builder( + verbose: bool, + workspace_bytes: int | None, + *, + weight_streaming: bool = False, +): logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) config = builder.create_builder_config() config.builder_optimization_level = 1 - op.configure_builder(config) - op.configure_workspace( - config, - workspace_bytes, - default_bytes=DENOISER_DEFAULT_WORKSPACE_BYTES, - ) + op.configure_builder(config, weight_streaming=weight_streaming) + # Leave TensorRT-RTX at its device-derived maximum unless the caller + # explicitly asks for a smaller tactic workspace. The largest dense plans + # use the TensorRT default workspace limit. + if workspace_bytes is not None: + op.configure_workspace( + config, + workspace_bytes, + default_bytes=DENOISER_DEFAULT_WORKSPACE_BYTES, + ) # Hugging Face keeps TF32 disabled for the FP32 input/output projections. config.clear_flag(trt.BuilderFlag.TF32) return logger, builder, network, config @@ -429,6 +450,8 @@ def _add_dynamic_profile( config, profile: MiniMaxH3Config, *, + video_inputs: tuple[str, ...] = (), + audio_inputs: tuple[str, ...] = (), text_inputs: tuple[str, ...] = (), packed_inputs: tuple[str, ...] = (), ) -> None: @@ -443,6 +466,20 @@ def _add_dynamic_profile( profile.opt_sequence_length, profile.sequence_length, ) + for name in video_inputs: + optimization.set_shape( + name, + min=(profile.min_video_rows, profile.video_patch_dim), + opt=(profile.opt_video_rows, profile.video_patch_dim), + max=(profile.video_rows, profile.video_patch_dim), + ) + for name in audio_inputs: + optimization.set_shape( + name, + min=(profile.min_audio_rows, profile.audio_in_channels), + opt=(profile.opt_audio_rows, profile.audio_in_channels), + max=(profile.audio_rows, profile.audio_in_channels), + ) for name in text_inputs: optimization.set_shape( name, @@ -460,6 +497,17 @@ def _add_dynamic_profile( config.add_optimization_profile(optimization) +def _add_first_block_cache_profiles( + builder, + config, + profile: MiniMaxH3Config, + **binding_groups, +) -> None: + """Use one dynamic profile for the complete requested media envelope.""" + + _add_dynamic_profile(builder, config, profile, **binding_groups) + + def _serialize( *, logger, @@ -469,109 +517,26 @@ def _serialize( weights: dict, consume_weights: bool, label: str, -) -> bytes: + output_path: str | Path | None, +) -> bytes | dict[str, int | str]: + plan = None + record = None try: - plan = builder.build_serialized_network(network, config) + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) finally: op.release_weight_buffers(network) if consume_weights: weights.clear() - if plan is None: + if output_path is None and plan is None: raise RuntimeError(f"TensorRT failed to build MiniMax-H3 {label} engine") del network, config, builder, logger gc.collect() - return bytes(plan) - - -def build_dit_engine( - weights: dict, - profile: MiniMaxH3Config, - *, - verbose: bool = False, - consume_weights: bool = False, - workspace_bytes: int | None = None, -) -> bytes: - """Build the full-sequence single-device H3 TensorRT plan.""" - - profile.validate() - if profile.first_block_cache: - raise ValueError("MiniMax-H3 first_block_cache profile requires the split DiT builders") - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) - - video = network.add_input( - "video_hidden_states", trt.float32, (profile.video_rows, profile.video_patch_dim) - ) - audio = network.add_input( - "audio_hidden_states", trt.float32, (profile.audio_rows, profile.audio_in_channels) - ) - text = network.add_input("encoder_hidden_states", trt.float32, (-1, profile.text_dim)) - positions = network.add_input("position_ids", trt.float32, (-1, 3)) - adaln_indices = network.add_input("adaln_indices", trt.int32, (-1,)) - timestep_indices = network.add_input("timestep_indices", trt.int32, (-1,)) - _add_dynamic_profile( - builder, - config, - profile, - text_inputs=("encoder_hidden_states",), - packed_inputs=("position_ids", "adaln_indices", "timestep_indices"), - ) - block_modulations = [ - network.add_input( - f"block_modulation_{index}", - trt.bfloat16, - (profile.adaln_table_rows, 6, profile.hidden_size), - ) - for index in range(profile.num_layers) - ] - final_modulation = network.add_input( - "final_modulation", trt.bfloat16, (profile.max_timestep_count, 2, profile.hidden_size) - ) - # The public single-device FL2VA profile is packed as text | audio | video. - # Projection, text refinement, packing, and full-sequence attention all - # remain native TensorRT operations on one device. - hidden = _packed_hidden(network, video, audio, text, weights, profile) - - cos, sin = _rope_tables(network, positions, profile, rows) - # The dynamic packed sequence contains live rows only, like Diffusers, so - # its attention mask remains None for every supported prompt length. - for index in range(profile.num_layers): - hidden = _transformer_block( - network, - hidden, - block_modulations[index], - adaln_indices, - cos, - sin, - weights, - profile, - index, - ) - - hidden = _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile) - _mark_sliced_velocity_outputs(network, hidden, weights, profile) - - op.validate_native_network( - network, - expected_attentions=profile.num_refiner_layers + profile.num_layers, - label="DiT", - ) - - print( - f"[minimax-h3] building native DiT: layers={profile.num_layers}, " - f"packed={profile.min_sequence_length}..{profile.sequence_length} " - f"(opt={profile.opt_sequence_length}), devices=1", - file=sys.stderr, - ) - return _serialize( - logger=logger, - builder=builder, - network=network, - config=config, - weights=weights, - consume_weights=consume_weights, - label="DiT", - ) + return record if record is not None else bytes(plan) def _require_first_block_cache_profile(profile: MiniMaxH3Config) -> None: @@ -580,6 +545,7 @@ def _require_first_block_cache_profile(profile: MiniMaxH3Config) -> None: raise ValueError("MiniMax-H3 split DiT plans require profile.first_block_cache=True") +@op.cleanup_failed_build def build_dit_head_engine( weights: dict, profile: MiniMaxH3Config, @@ -587,17 +553,20 @@ def build_dit_head_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Build packing, text refinement, block zero, and the native cache metric.""" _require_first_block_cache_profile(profile) - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) video = network.add_input( - "video_hidden_states", trt.float32, (profile.video_rows, profile.video_patch_dim) + "video_hidden_states", trt.float32, (-1, profile.video_patch_dim) ) audio = network.add_input( - "audio_hidden_states", trt.float32, (profile.audio_rows, profile.audio_in_channels) + "audio_hidden_states", trt.float32, (-1, profile.audio_in_channels) ) text = network.add_input("encoder_hidden_states", trt.float32, (-1, profile.text_dim)) positions = network.add_input("position_ids", trt.float32, (-1, 3)) @@ -610,16 +579,26 @@ def build_dit_head_engine( previous_head_residual = network.add_input( "previous_head_residual", trt.bfloat16, (-1, profile.hidden_size) ) - _add_dynamic_profile( + _add_first_block_cache_profiles( builder, config, profile, + video_inputs=("video_hidden_states",), + audio_inputs=("audio_hidden_states",), text_inputs=("encoder_hidden_states",), packed_inputs=("position_ids", "adaln_indices", "previous_head_residual"), ) - pre_block_hidden = _packed_hidden(network, video, audio, text, weights, profile) - cos, sin = _rope_tables(network, positions, profile, rows) + pre_block_hidden = _packed_hidden( + network, + video, + audio, + text, + weights, + profile, + consume_weights=consume_weights, + ) + cos, sin = _rope_tables(network, positions, profile) head_hidden = _transformer_block( network, pre_block_hidden, @@ -630,6 +609,7 @@ def build_dit_head_engine( weights, profile, 0, + consume_weights=consume_weights, ) head_residual = network.add_elementwise( head_hidden, pre_block_hidden, trt.ElementWiseOperation.SUB @@ -687,9 +667,11 @@ def build_dit_head_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache head", + output_path=output_path, ) +@op.cleanup_failed_build def build_dit_tail_engine( weights: dict, profile: MiniMaxH3Config, @@ -697,16 +679,19 @@ def build_dit_tail_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Build blocks one through 49 and expose their reusable total residual.""" _require_first_block_cache_profile(profile) - rows = -1 - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) head_hidden = network.add_input("head_hidden", trt.bfloat16, (-1, profile.hidden_size)) positions = network.add_input("position_ids", trt.float32, (-1, 3)) adaln_indices = network.add_input("adaln_indices", trt.int32, (-1,)) - _add_dynamic_profile( + _add_first_block_cache_profiles( builder, config, profile, @@ -720,7 +705,7 @@ def build_dit_tail_engine( ) for index in range(1, profile.num_layers) } - cos, sin = _rope_tables(network, positions, profile, rows) + cos, sin = _rope_tables(network, positions, profile) hidden = head_hidden for index in range(1, profile.num_layers): hidden = _transformer_block( @@ -733,6 +718,7 @@ def build_dit_tail_engine( weights, profile, index, + consume_weights=consume_weights, ) tail_residual = network.add_elementwise( hidden, head_hidden, trt.ElementWiseOperation.SUB @@ -757,9 +743,11 @@ def build_dit_tail_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache tail", + output_path=output_path, ) +@op.cleanup_failed_build def build_dit_finish_engine( weights: dict, profile: MiniMaxH3Config, @@ -767,18 +755,30 @@ def build_dit_finish_engine( verbose: bool = False, consume_weights: bool = False, workspace_bytes: int | None = None, -) -> bytes: + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: """Apply a selected tail residual, final norm, and consumed-row projections.""" _require_first_block_cache_profile(profile) - logger, builder, network, config = _native_builder(verbose, workspace_bytes) + logger, builder, network, config = _native_builder( + verbose, workspace_bytes, weight_streaming=weight_streaming + ) head_hidden = network.add_input("head_hidden", trt.bfloat16, (-1, profile.hidden_size)) tail_residual = network.add_input("tail_residual", trt.bfloat16, (-1, profile.hidden_size)) timestep_indices = network.add_input("timestep_indices", trt.int32, (-1,)) - _add_dynamic_profile( + video = network.add_input( + "video_hidden_states", trt.float32, (-1, profile.video_patch_dim) + ) + audio = network.add_input( + "audio_hidden_states", trt.float32, (-1, profile.audio_in_channels) + ) + _add_first_block_cache_profiles( builder, config, profile, + video_inputs=("video_hidden_states",), + audio_inputs=("audio_hidden_states",), packed_inputs=("head_hidden", "tail_residual", "timestep_indices"), ) final_modulation = network.add_input( @@ -788,7 +788,7 @@ def build_dit_finish_engine( head_hidden, tail_residual, trt.ElementWiseOperation.SUM ).get_output(0) hidden = _final_hidden(network, hidden, timestep_indices, final_modulation, weights, profile) - _mark_sliced_velocity_outputs(network, hidden, weights, profile) + _mark_sliced_velocity_outputs(network, hidden, weights, video, audio) op.validate_native_network( network, expected_attentions=0, @@ -807,4 +807,5 @@ def build_dit_finish_engine( weights=weights, consume_weights=consume_weights, label="DiT FirstBlockCache finish", + output_path=output_path, ) diff --git a/families/minimax_h3/fl2va_contract.py b/families/minimax_h3/fl2va_contract.py new file mode 100644 index 0000000000..9e4d78d4c3 --- /dev/null +++ b/families/minimax_h3/fl2va_contract.py @@ -0,0 +1,665 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time contracts for MiniMax-H3 first/last-frame conditioning. + +The runtime consumes the plans described here without importing this module. +Keeping the public FL2VA geometry, Qwen presentation, and row accounting in a +small module makes it possible to validate a bundle before any multi-gigabyte +engine is built. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + + +CANVAS_MULTIPLE = 32 +CANVAS_SHORT_EDGE = 768 +CANVAS_MAX_PIXELS = 768 * 1344 +VAE_SPATIAL_COMPRESSION = 16 +VAE_TILE_SIZE = 256 +VAE_TILE_MIN_OVERLAP = 64 +VAE_TILE_OPT_BATCH = 28 +VAE_TILE_MAX_BATCH = 33 +VIDEO_PATCH_SIZE = (1, 2, 2) +QWEN_VISION_PATCH_SIZE = 16 +QWEN_VISION_TEMPORAL_PATCH_SIZE = 2 +QWEN_VISION_MERGE_SIZE = 2 +QWEN_VISION_PATCH_WIDTH = 3 * QWEN_VISION_TEMPORAL_PATCH_SIZE * QWEN_VISION_PATCH_SIZE**2 +QWEN_VISION_HIDDEN_SIZE = 1152 +QWEN_TEXT_HIDDEN_SIZE = 5120 +QWEN_LABEL_TOKENS = 6 +QWEN_VISION_BOUNDARY_TOKENS = 2 +QWEN_VISION_START_TOKEN_ID = 151652 +QWEN_VISION_END_TOKEN_ID = 151653 +QWEN_IMAGE_PAD_TOKEN_ID = 151655 +MIN_FRAMES = 124 +MAX_FRAMES = 345 +FRAMES_PER_CHUNK = 17 +LATENTS_PER_CHUNK = 5 +AUDIO_LATENTS_PER_SECOND = 40 +FPS = 24 +LATENTS_MEAN = ( + 0.858090341091156, + -0.9606591463088989, + 1.0661640167236328, + -0.5090325474739075, + -0.2727581858634949, + -1.3675414323806763, + -0.2553254961967468, + -0.26907554268836975, + -0.5376840829849243, + -0.0464097298681736, + 0.6657370328903198, + 0.19690127670764923, + -0.5460608005523682, + -0.4035342037677765, + -0.23683024942874908, + 0.25928452610969543, + -0.30133944749832153, + 0.211341992020607, + -1.1206848621368408, + 0.3581933379173279, + -0.04225143790245056, + 0.2604829967021942, + 0.22864092886447906, + 0.7056031823158264, +) +LATENTS_STD = ( + 1.2223774194717407, + 1.2767263650894165, + 1.6831774711608887, + 1.7549455165863037, + 1.5636216402053833, + 2.194143533706665, + 0.9653137922286987, + 1.0569885969161987, + 0.841948926448822, + 0.7729952931404114, + 1.8955937623977661, + 0.946841835975647, + 0.7996809482574463, + 0.44988900423049927, + 0.7197399735450745, + 0.6936293244361877, + 2.961095094680786, + 2.7694199085235596, + 3.0496184825897217, + 2.1088054180145264, + 3.276226282119751, + 3.1627357006073, + 2.2816812992095947, + 2.6127843856811523, +) + + +@dataclass(frozen=True) +class TensorAbi: + """One TensorRT binding, including every optimization-profile shape.""" + + name: str + dtype: str + min_shape: tuple[int, ...] + opt_shape: tuple[int, ...] + max_shape: tuple[int, ...] + + +@dataclass(frozen=True) +class PlanAbi: + """The complete externally visible binding contract of one native plan.""" + + filename: str + inputs: tuple[TensorAbi, ...] + outputs: tuple[TensorAbi, ...] + + +@dataclass(frozen=True) +class VisionEncoderProfile: + """Patch-row profile for one Qwen3-VL image or temporal video block.""" + + min_patches: int = 1620 + opt_patches: int = 4032 + max_patches: int = 4176 + + def validate(self) -> None: + values = (self.min_patches, self.opt_patches, self.max_patches) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 vision profile dimensions must be integers") + if not 4 <= self.min_patches <= self.opt_patches <= self.max_patches: + raise ValueError("MiniMax-H3 vision profile must satisfy 4 <= min <= opt <= max") + if any(value % (QWEN_VISION_MERGE_SIZE**2) for value in values): + raise ValueError("MiniMax-H3 vision profile dimensions must be divisible by four") + + +@dataclass(frozen=True) +class MultimodalTextProfile: + """Sequence profile shared by T2VA, FL2VA, and Ref2VA conditioning.""" + + min_sequence_length: int = 1 + opt_sequence_length: int = 1144 + max_sequence_length: int = 2641 + # TensorRT profiles cannot bind a zero-length tensor. Text-only T2VA uses + # one ignored dummy visual row with ``vision_count == 0``. + min_vision_rows: int = 1 + opt_vision_rows: int = 1008 + max_vision_rows: int = 2088 + + def validate(self) -> None: + values = ( + self.min_sequence_length, + self.opt_sequence_length, + self.max_sequence_length, + self.min_vision_rows, + self.opt_vision_rows, + self.max_vision_rows, + ) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 text profile dimensions must be integers") + if ( + not 1 + <= self.min_sequence_length + <= self.opt_sequence_length + <= self.max_sequence_length + ): + raise ValueError("MiniMax-H3 text profile must satisfy 1 <= min <= opt <= max") + if not 1 <= self.min_vision_rows <= self.opt_vision_rows <= self.max_vision_rows: + raise ValueError("MiniMax-H3 visual-row profile must satisfy 1 <= min <= opt <= max") + if self.min_vision_rows > self.min_sequence_length: + raise ValueError("MiniMax-H3 minimum visual rows cannot exceed minimum sequence rows") + if self.opt_vision_rows > self.opt_sequence_length: + raise ValueError("MiniMax-H3 optimum visual rows cannot exceed optimum sequence rows") + if self.max_vision_rows > self.max_sequence_length: + raise ValueError("MiniMax-H3 maximum visual rows cannot exceed maximum sequence rows") + + +@dataclass(frozen=True) +class TileAxis: + starts: tuple[int, ...] + lengths: tuple[int, ...] + overlaps: tuple[int, ...] + + +@dataclass(frozen=True) +class KeyframeResize: + """Integer geometry for one exact LANCZOS resize/crop operation.""" + + mode: str + resized_height: int + resized_width: int + crop_top: int = 0 + crop_left: int = 0 + + +@dataclass(frozen=True) +class PackedRows: + text: int + condition_video: int + target_audio: int + target_video: int + + @property + def total(self) -> int: + return self.text + self.condition_video + self.target_audio + self.target_video + + +def resolve_canvas_size( + aspect_width: float, + aspect_height: float, + *, + multiple: int = CANVAS_MULTIPLE, + short_edge: int = CANVAS_SHORT_EDGE, + max_pixels: int = CANVAS_MAX_PIXELS, +) -> tuple[int, int]: + """Reproduce the public canvas resolver and return ``(height, width)``.""" + + values = (aspect_width, aspect_height, multiple, short_edge, max_pixels) + if any(isinstance(value, bool) for value in values): + raise ValueError("MiniMax-H3 canvas geometry cannot use booleans") + if not all(isinstance(value, (int, float)) for value in (aspect_width, aspect_height)): + raise ValueError("MiniMax-H3 aspect dimensions must be numeric") + if not all(isinstance(value, int) for value in (multiple, short_edge, max_pixels)): + raise ValueError("MiniMax-H3 canvas limits must be integers") + if not math.isfinite(aspect_width) or not math.isfinite(aspect_height): + raise ValueError("MiniMax-H3 aspect dimensions must be finite") + if ( + aspect_width <= 0 + or aspect_height <= 0 + or multiple <= 0 + or short_edge <= 0 + or max_pixels <= 0 + ): + raise ValueError("MiniMax-H3 canvas geometry must be positive") + + ratio = aspect_width / aspect_height + if not 0.25 <= ratio <= 4.0: + raise ValueError( + f"MiniMax-H3 supports aspect ratios from 1:4 to 4:1, got {aspect_width}:{aspect_height}" + ) + if ratio >= 1.0: + width, height = short_edge * ratio, float(short_edge) + else: + width, height = float(short_edge), short_edge / ratio + area = width * height + if area > max_pixels: + scale = math.sqrt(max_pixels / area) + width *= scale + height *= scale + return ( + max(multiple, round(height / multiple) * multiple), + max(multiple, round(width / multiple) * multiple), + ) + + +def validate_canvas(height: int, width: int) -> None: + if isinstance(height, bool) or isinstance(width, bool): + raise ValueError("MiniMax-H3 canvas dimensions must be integers") + if not isinstance(height, int) or not isinstance(width, int) or height <= 0 or width <= 0: + raise ValueError(f"MiniMax-H3 canvas must be positive, got {height}x{width}") + if height % CANVAS_MULTIPLE or width % CANVAS_MULTIPLE: + raise ValueError( + f"MiniMax-H3 canvas must be a multiple of {CANVAS_MULTIPLE}, got {height}x{width}" + ) + if width > 4 * height or height > 4 * width: + raise ValueError(f"MiniMax-H3 canvas must be within 1:4 and 4:1, got {height}x{width}") + + +def resolve_keyframe_anchors(*, has_first: bool, has_last: bool) -> tuple[str, ...]: + """Return the reference's packed keyframe order and endpoint anchors.""" + + if not isinstance(has_first, bool) or not isinstance(has_last, bool): + raise ValueError("MiniMax-H3 keyframe presence flags must be booleans") + anchors = tuple( + anchor for anchor, present in (("first", has_first), ("last", has_last)) if present + ) + if not anchors: + raise ValueError("MiniMax-H3 FL2VA needs a first frame, a last frame, or both") + return anchors + + +def keyframe_resize_geometry( + source_height: int, + source_width: int, + target_height: int, + target_width: int, + *, + packed_index: int, +) -> KeyframeResize: + """Resolve the exact stretch/cover-crop geometry for one packed keyframe. + + Packed index zero is always the geometry anchor, including a last-only + request. It is stretched to the canvas. Only the second keyframe is a + cover-cropped follower. + """ + + values = (source_height, source_width, target_height, target_width, packed_index) + if any(isinstance(value, bool) or not isinstance(value, int) for value in values): + raise ValueError("MiniMax-H3 resize geometry must use integer values") + if min(source_height, source_width, target_height, target_width) <= 0: + raise ValueError("MiniMax-H3 resize dimensions must be positive") + if packed_index not in (0, 1): + raise ValueError("MiniMax-H3 FL2VA has at most two packed keyframes") + validate_canvas(target_height, target_width) + if (source_height, source_width) == (target_height, target_width): + return KeyframeResize("identity", target_height, target_width) + if packed_index == 0: + return KeyframeResize("stretch", target_height, target_width) + + scale = max(target_width / source_width, target_height / source_height) + resized_width = max(target_width, round(source_width * scale)) + resized_height = max(target_height, round(source_height * scale)) + return KeyframeResize( + "cover_crop", + resized_height, + resized_width, + crop_top=max(0, (resized_height - target_height) // 2), + crop_left=max(0, (resized_width - target_width) // 2), + ) + + +def split_tile_axis( + length: int, + tile_size: int = VAE_TILE_SIZE, + min_overlap: int = VAE_TILE_MIN_OVERLAP, + alignment: int = VAE_SPATIAL_COMPRESSION, +) -> TileAxis: + """Reproduce ``AutoencoderKLMiniMaxH3._split_tiles`` exactly.""" + + if any(isinstance(value, bool) for value in (length, tile_size, min_overlap, alignment)): + raise ValueError("MiniMax-H3 tile geometry must use integer values") + if not all(isinstance(value, int) for value in (length, tile_size, min_overlap, alignment)): + raise ValueError("MiniMax-H3 tile geometry must use integer values") + if length <= 0 or tile_size <= 0 or alignment <= 0 or not 0 <= min_overlap < tile_size: + raise ValueError("invalid MiniMax-H3 tile geometry") + if length % alignment: + raise ValueError(f"MiniMax-H3 tile axis {length} is not latent-aligned to {alignment}") + if tile_size >= length: + return TileAxis((0,), (length,), ()) + + num_tiles = math.ceil(length / tile_size) + while tile_size * num_tiles - min_overlap * (num_tiles - 1) - length < 0: + num_tiles += 1 + + overlaps = [min_overlap] * (num_tiles - 1) + remaining = tile_size * num_tiles - sum(overlaps) - length + if remaining % alignment: + raise ValueError("MiniMax-H3 tile slack is not latent-aligned") + for index in range(remaining // alignment): + overlaps[index % (num_tiles - 1)] += alignment + + starts = [0] + for overlap in overlaps: + starts.append(starts[-1] + tile_size - overlap) + result = TileAxis(tuple(starts), (tile_size,) * num_tiles, tuple(overlaps)) + if result.starts[-1] + result.lengths[-1] != length: + raise ValueError("MiniMax-H3 tile split does not cover the canvas exactly") + return result + + +def latent_tile_axis(length: int) -> TileAxis: + """Map an exact pixel-space tile split onto the VAE latent grid. + + Stitching blends each overlap of ``n`` latent cells at positions + ``i = 0..n-1`` as ``previous * (1 - i/n) + current * (i/n)``. Vertical + blending is applied before horizontal blending, then every non-final tile + drops its trailing overlap exactly as the reference ``_stitch_tiles`` does. + """ + + pixels = split_tile_axis(length) + return TileAxis( + tuple(start // VAE_SPATIAL_COMPRESSION for start in pixels.starts), + tuple(size // VAE_SPATIAL_COMPRESSION for size in pixels.lengths), + tuple(overlap // VAE_SPATIAL_COMPRESSION for overlap in pixels.overlaps), + ) + + +def keyframe_tile_count(height: int, width: int) -> int: + validate_canvas(height, width) + if height < VAE_TILE_SIZE or width < VAE_TILE_SIZE: + raise ValueError( + f"MiniMax-H3 keyframe plan requires both canvas axes to be at least " + f"{VAE_TILE_SIZE}, got {height}x{width}" + ) + count = len(split_tile_axis(height).starts) * len(split_tile_axis(width).starts) + if count > VAE_TILE_MAX_BATCH: + raise ValueError( + f"MiniMax-H3 keyframe needs {count} VAE tiles, plan capacity is {VAE_TILE_MAX_BATCH}" + ) + return count + + +def qwen_vision_patch_rows(height: int, width: int) -> int: + validate_canvas(height, width) + if height % QWEN_VISION_PATCH_SIZE or width % QWEN_VISION_PATCH_SIZE: + raise ValueError("MiniMax-H3 Qwen vision canvas is not patch-aligned") + return (height // QWEN_VISION_PATCH_SIZE) * (width // QWEN_VISION_PATCH_SIZE) + + +def qwen_vision_token_rows(height: int, width: int) -> int: + patches = qwen_vision_patch_rows(height, width) + merge_unit = QWEN_VISION_MERGE_SIZE**2 + if patches % merge_unit: + raise ValueError("MiniMax-H3 Qwen vision patches are not merge-aligned") + return patches // merge_unit + + +def fl2va_text_rows(prompt_tokens: int, keyframes: int, *, height: int, width: int) -> int: + if isinstance(prompt_tokens, bool) or not isinstance(prompt_tokens, int) or prompt_tokens < 0: + raise ValueError("MiniMax-H3 prompt token count must be a non-negative integer") + if keyframes not in (1, 2): + raise ValueError(f"MiniMax-H3 FL2VA needs one or two keyframes, got {keyframes}") + per_keyframe = ( + QWEN_LABEL_TOKENS + QWEN_VISION_BOUNDARY_TOKENS + qwen_vision_token_rows(height, width) + ) + return prompt_tokens + keyframes * per_keyframe + + +def fl2va_mrope_position_ids( + prompt_tokens: int, + keyframes: int, + *, + height: int, + width: int, +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """Build Qwen3-VL's exact temporal/height/width position IDs. + + Qwen marks only ``image_pad`` tokens as image modality. Each label and + ``vision_start`` therefore forms a text run before the image grid, while a + preceding ``vision_end`` joins the next label when two images are present. + This mirrors ``Qwen3VLModel.get_rope_index`` without requiring Transformers + at bundle-build time. + """ + + expected_rows = fl2va_text_rows(prompt_tokens, keyframes, height=height, width=width) + llm_height = height // (QWEN_VISION_PATCH_SIZE * QWEN_VISION_MERGE_SIZE) + llm_width = width // (QWEN_VISION_PATCH_SIZE * QWEN_VISION_MERGE_SIZE) + axes: tuple[list[int], list[int], list[int]] = ([], [], []) + current_position = 0 + + def append_text(length: int) -> None: + nonlocal current_position + positions = range(current_position, current_position + length) + for axis in axes: + axis.extend(positions) + current_position += length + + for index in range(keyframes): + # The second text run begins with the previous image's vision_end. + append_text(QWEN_LABEL_TOKENS + 1 + int(index > 0)) + for row in range(llm_height): + for column in range(llm_width): + axes[0].append(current_position) + axes[1].append(current_position + row) + axes[2].append(current_position + column) + current_position += max(llm_height, llm_width) + append_text(1 + prompt_tokens) # final vision_end followed by the prompt + + result = tuple(tuple(axis) for axis in axes) + if any(len(axis) != expected_rows for axis in result): + raise RuntimeError("MiniMax-H3 FL2VA MRoPE row accounting mismatch") + return result + + +def align_num_frames(num_frames: int) -> int: + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 frame count must be a positive integer") + aligned = num_frames + (LATENTS_PER_CHUNK - num_frames % FRAMES_PER_CHUNK) % FRAMES_PER_CHUNK + if not MIN_FRAMES <= aligned <= MAX_FRAMES: + raise ValueError( + f"MiniMax-H3 aligned frame count must be {MIN_FRAMES}..{MAX_FRAMES}, got {aligned}" + ) + return aligned + + +def video_latent_frames(num_frames: int) -> int: + aligned = align_num_frames(num_frames) + return ((aligned - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * LATENTS_PER_CHUNK + 2 + + +def audio_latent_frames(num_frames: int) -> int: + aligned = align_num_frames(num_frames) + return int(round(aligned / FPS * AUDIO_LATENTS_PER_SECOND)) + + +def packed_rows( + *, + prompt_tokens: int, + keyframes: int, + height: int, + width: int, + num_frames: int, +) -> PackedRows: + text = fl2va_text_rows(prompt_tokens, keyframes, height=height, width=width) + rows_per_frame = (height // VAE_SPATIAL_COMPRESSION // VIDEO_PATCH_SIZE[1]) * ( + width // VAE_SPATIAL_COMPRESSION // VIDEO_PATCH_SIZE[2] + ) + return PackedRows( + text=text, + condition_video=keyframes * rows_per_frame, + target_audio=2 * audio_latent_frames(num_frames), + target_video=video_latent_frames(num_frames) * rows_per_frame, + ) + + +def last_keyframe_rotary_time(text_rows: int, num_frames: int) -> float: + """The reference's NumPy-summed last-frame rotary anchor.""" + + latent_frames = video_latent_frames(num_frames) + spans = [5.0 / 3.0 * (1, 4, 4, 4, 4)[index % 5] for index in range(latent_frames)] + # Python's sum is deliberately not used: NumPy pairwise summation is the + # reference. The native runtime implements the same reduction order. + import numpy as np + + return float(text_rows) + float(np.asarray(spans, dtype=np.float64).sum()) - 5.0 / 3.0 + + +def keyframe_vae_encoder_abi() -> PlanAbi: + """Tile-plan ABI for the exact released keyframe posterior recipe. + + The caller supplies row-major tiles after ImageNet normalization. It must + stitch all 48 posterior-parameter channels before splitting mean/logvar, + clamp logvar to ``[-30, 20]``, sample with a fresh CPU generator seeded 42, + round the sample through FP16, and apply ``LATENTS_MEAN/LATENTS_STD``. Each + keyframe repeats that fresh-generator recipe independently. + """ + + return PlanAbi( + filename="fl2va_keyframe_vae_encoder.plan", + inputs=( + TensorAbi( + "pixel_tiles", + "float32", + (1, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + (VAE_TILE_OPT_BATCH, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + (VAE_TILE_MAX_BATCH, 3, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + ), + ), + outputs=( + TensorAbi( + "posterior_parameter_tiles", + "float32", + ( + 1, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ( + VAE_TILE_OPT_BATCH, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ( + VAE_TILE_MAX_BATCH, + 48, + 1, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION, + ), + ), + ), + ) + + +def vision_encoder_abi(profile: VisionEncoderProfile = VisionEncoderProfile()) -> PlanAbi: + """One-image/frame dynamic-aspect Qwen3-VL vision-tower plan. + + A reference video is evaluated one temporal patch block at a time. This + preserves Qwen3-VL's ``cu_seqlens`` boundary (vision attention never crosses + temporal blocks) while allowing the same weights and engine to serve FL2VA + images, Ref2VA images, and Ref2VA video frames. + """ + + profile.validate() + min_patches = profile.min_patches + max_patches = profile.max_patches + opt_patches = profile.opt_patches + bindings = ( + TensorAbi( + "pixel_values", + "float32", + (min_patches, QWEN_VISION_PATCH_WIDTH), + (opt_patches, QWEN_VISION_PATCH_WIDTH), + (max_patches, QWEN_VISION_PATCH_WIDTH), + ), + TensorAbi("interp_indices", "int32", (min_patches, 4), (opt_patches, 4), (max_patches, 4)), + TensorAbi( + "interp_weights", "float32", (min_patches, 4), (opt_patches, 4), (max_patches, 4) + ), + TensorAbi( + "vision_position_ids", "int32", (min_patches, 2), (opt_patches, 2), (max_patches, 2) + ), + ) + + def output(name: str) -> TensorAbi: + return TensorAbi( + name, + "float32", + (min_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + (opt_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + (max_patches // 4, QWEN_TEXT_HIDDEN_SIZE), + ) + + return PlanAbi( + filename="vision_encoder.plan", + inputs=bindings, + outputs=tuple( + output(name) for name in ("vision_embeds", "deepstack_0", "deepstack_1", "deepstack_2") + ), + ) + + +def text_encoder_abi( + profile: MultimodalTextProfile = MultimodalTextProfile(), +) -> PlanAbi: + """Unified Qwen3-VL language-stack ABI for all three public workflows. + + Visual features remain compact at the ABI and are scattered through + ``vision_row_indices``. Text-only requests bind one dummy row, set + ``vision_count`` to zero, and clear ``vision_mask``. Thus the plan owns one + copy of language weights instead of shipping separate T2VA and FL2VA + language plans. + """ + + profile.validate() + min_rows = profile.min_sequence_length + opt_rows = profile.opt_sequence_length + max_rows = profile.max_sequence_length + min_vision = profile.min_vision_rows + opt_vision = profile.opt_vision_rows + max_vision = profile.max_vision_rows + + def rows(name: str, dtype: str, width: int | None = None) -> TensorAbi: + suffix = () if width is None else (width,) + return TensorAbi(name, dtype, (min_rows, *suffix), (opt_rows, *suffix), (max_rows, *suffix)) + + def visual_rows(name: str, dtype: str, width: int | None = None) -> TensorAbi: + suffix = () if width is None else (width,) + return TensorAbi( + name, + dtype, + (min_vision, *suffix), + (opt_vision, *suffix), + (max_vision, *suffix), + ) + + inputs = ( + rows("input_ids", "int32"), + TensorAbi("mrope_position_ids", "int32", (3, min_rows), (3, opt_rows), (3, max_rows)), + rows("vision_mask", "float32", 1), + TensorAbi("vision_count", "int32", (1,), (1,), (1,)), + visual_rows("vision_row_indices", "int32"), + visual_rows("vision_embeds", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_0", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_1", "float32", QWEN_TEXT_HIDDEN_SIZE), + visual_rows("deepstack_2", "float32", QWEN_TEXT_HIDDEN_SIZE), + ) + return PlanAbi( + filename="text_encoder.plan", + inputs=inputs, + outputs=(rows("encoder_hidden_states", "float32", QWEN_TEXT_HIDDEN_SIZE),), + ) diff --git a/families/minimax_h3/fl2va_vae_encoder_builder.py b/families/minimax_h3/fl2va_vae_encoder_builder.py new file mode 100644 index 0000000000..0bb885d399 --- /dev/null +++ b/families/minimax_h3/fl2va_vae_encoder_builder.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native TensorRT encoder for MiniMax-H3 FL2VA keyframe tiles. + +The released VAE always encodes keyframes as one temporal frame. Causal 3-D +convolution therefore has one non-zero temporal tap: the last tap of every +3x3x3 kernel (or the only tap of a 1x1x1 kernel). Expressing that exact case +as 2-D convolutions avoids manufacturing zero temporal frames while retaining +the released spatial reflection padding, isolated group norm, and FP32 math. + +The plan deliberately emits posterior parameters. Posterior sampling, +FP16-rounding, per-channel latent normalization, tile stitching, and the +request-generator draw order belong to the native runtime rather than the +engine and are stated in :mod:`fl2va_contract`. +""" + +from __future__ import annotations + +import gc +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .fl2va_contract import ( + VAE_SPATIAL_COMPRESSION, + VAE_TILE_MAX_BATCH, + VAE_TILE_SIZE, + keyframe_vae_encoder_abi, +) + + +trt = trt_compat.get_trt() + +IN_CHANNELS = 3 +MOMENT_CHANNELS = 48 +BLOCK_OUT_CHANNELS = (128, 256, 256, 512, 512, 1024) +SPATIAL_DOWNSAMPLE_FACTORS = (2, 2, 2, 2, 1, 1) +TEMPORAL_DOWNSAMPLE_FACTORS = (1, 2, 2, 1, 1, 1) +LAYERS_PER_BLOCK = 2 +NORM_GROUPS = 32 +NORM_EPS = 1.0e-6 +KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 + + +def checkpoint_keys() -> tuple[str, ...]: + """Return the exhaustive VAE checkpoint partition consumed by this plan.""" + + names = ["encoder.conv_in.weight", "encoder.conv_in.bias"] + block_in_channels = (BLOCK_OUT_CHANNELS[0], *BLOCK_OUT_CHANNELS[:-1]) + for block_index, (in_channels, out_channels) in enumerate( + zip(block_in_channels, BLOCK_OUT_CHANNELS) + ): + for layer_index in range(LAYERS_PER_BLOCK): + prefix = f"encoder.down_blocks.{block_index}.resnets.{layer_index}" + names.extend( + [ + f"{prefix}.norm1.weight", + f"{prefix}.norm1.bias", + f"{prefix}.conv1.weight", + f"{prefix}.conv1.bias", + f"{prefix}.norm2.weight", + f"{prefix}.norm2.bias", + f"{prefix}.conv2.weight", + f"{prefix}.conv2.bias", + ] + ) + if layer_index == 0 and in_channels != out_channels: + names.extend([f"{prefix}.conv_shortcut.weight", f"{prefix}.conv_shortcut.bias"]) + if SPATIAL_DOWNSAMPLE_FACTORS[block_index] * TEMPORAL_DOWNSAMPLE_FACTORS[block_index] > 1: + prefix = f"encoder.down_blocks.{block_index}.downsamplers.0.conv" + names.extend([f"{prefix}.weight", f"{prefix}.bias"]) + names.extend( + [ + "encoder.norm_out.weight", + "encoder.norm_out.bias", + "encoder.conv_out.weight", + "encoder.conv_out.bias", + "quant_conv.weight", + "quant_conv.bias", + ] + ) + return tuple(names) + + +def _shape_dim(network, tensor, axis: int): + shape = network.add_shape(tensor).get_output(0) + return network.add_slice(shape, (axis,), (1,), (1,)).get_output(0) + + +def _shape_vector(network, values): + tensors = [] + for value in values: + if isinstance(value, int): + tensors.append( + op.constant(network, np.asarray([value], dtype=np.int64), dtype=np.int64) + ) + else: + tensors.append(value) + if len(tensors) == 1: + return tensors[0] + layer = network.add_concatenation(tensors) + layer.axis = 0 + return layer.get_output(0) + + +def _reflect_pad_2d( + network, + hidden, + *, + channels: int, + height: int, + width: int, + top: int = 0, + bottom: int = 0, + left: int = 0, + right: int = 0, +): + """PyTorch-compatible reflection padding for static CHW and dynamic batch.""" + + if min(top, bottom, left, right) < 0: + raise ValueError("MiniMax-H3 reflect padding cannot be negative") + if top >= height or bottom >= height or left >= width or right >= width: + raise ValueError("MiniMax-H3 reflect padding must be smaller than its input axis") + if top == bottom == left == right == 0: + return hidden + output_height = height + top + bottom + output_width = width + left + right + layer = network.add_slice( + hidden, + (0, 0, -top, -left), + (1, channels, output_height, output_width), + (1, 1, 1, 1), + ) + if layer is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 VAE reflection padding") + layer.set_input( + 2, + _shape_vector( + network, + (_shape_dim(network, hidden, 0), channels, output_height, output_width), + ), + ) + layer.mode = trt.SampleMode.REFLECT + return layer.get_output(0) + + +def _spatial_kernel(weight: np.ndarray) -> np.ndarray: + """Select the only non-zero causal temporal tap for a one-frame clip.""" + + value = np.asarray(weight) + if value.ndim != 5 or value.shape[2] not in (1, 3): + raise ValueError(f"MiniMax-H3 VAE convolution weight has invalid shape {value.shape}") + return np.ascontiguousarray(value[:, :, -1]) + + +def _conv( + network, + hidden, + weights: dict[str, np.ndarray], + prefix: str, + *, + in_channels: int, + out_channels: int, + height: int, + width: int, + stride: int = 1, + symmetric_reflect: int = 0, + bottom_right_reflect: int = 0, +): + if symmetric_reflect and bottom_right_reflect: + raise ValueError("MiniMax-H3 VAE convolution has ambiguous padding") + if symmetric_reflect: + hidden = _reflect_pad_2d( + network, + hidden, + channels=in_channels, + height=height, + width=width, + top=symmetric_reflect, + bottom=symmetric_reflect, + left=symmetric_reflect, + right=symmetric_reflect, + ) + height += 2 * symmetric_reflect + width += 2 * symmetric_reflect + elif bottom_right_reflect: + hidden = _reflect_pad_2d( + network, + hidden, + channels=in_channels, + height=height, + width=width, + bottom=bottom_right_reflect, + right=bottom_right_reflect, + ) + height += bottom_right_reflect + width += bottom_right_reflect + + kernel = _spatial_kernel(weights[f"{prefix}.weight"]) + layer = network.add_convolution_nd( + hidden, + out_channels, + tuple(kernel.shape[-2:]), + kernel, + np.ascontiguousarray(weights[f"{prefix}.bias"]), + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE convolution {prefix}") + layer.name = prefix + layer.stride_nd = (stride, stride) + output_height = (height - kernel.shape[-2]) // stride + 1 + output_width = (width - kernel.shape[-1]) // stride + 1 + return layer.get_output(0), output_height, output_width + + +def _group_norm(network, hidden, weights, prefix: str, channels: int): + rank = len(tuple(hidden.shape)) + gamma = op.weight_constant( + network, np.asarray(weights[f"{prefix}.weight"]).reshape(1, channels, 1, 1) + ) + beta = op.weight_constant( + network, np.asarray(weights[f"{prefix}.bias"]).reshape(1, channels, 1, 1) + ) + gamma = op.cast(network, gamma, hidden.dtype) + beta = op.cast(network, beta, hidden.dtype) + # TensorRT group normalization treats the second NCHW dimension as the + # grouped channel axis. When num_groups != 1, the reduction mask must + # therefore contain only dimensions after C (NvInfer.h, INormalizationLayer + # contract), rather than C itself. Including bit 1 is rejected by + # TensorRT-RTX before the graph can be serialized. + axes = sum(1 << axis for axis in range(2, rank)) + layer = network.add_normalization_v2(hidden, gamma, beta, axes) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE group norm {prefix}") + layer.name = prefix + layer.epsilon = NORM_EPS + layer.num_groups = NORM_GROUPS + return layer.get_output(0) + + +def _resnet( + network, + hidden, + weights, + prefix: str, + *, + in_channels: int, + out_channels: int, + height: int, + width: int, +): + residual = hidden + update = _group_norm(network, hidden, weights, f"{prefix}.norm1", in_channels) + update = op.silu(network, update) + update, update_height, update_width = _conv( + network, + update, + weights, + f"{prefix}.conv1", + in_channels=in_channels, + out_channels=out_channels, + height=height, + width=width, + symmetric_reflect=1, + ) + if (update_height, update_width) != (height, width): + raise RuntimeError("MiniMax-H3 VAE resnet conv1 changed spatial geometry") + update = _group_norm(network, update, weights, f"{prefix}.norm2", out_channels) + update = op.silu(network, update) + update, update_height, update_width = _conv( + network, + update, + weights, + f"{prefix}.conv2", + in_channels=out_channels, + out_channels=out_channels, + height=height, + width=width, + symmetric_reflect=1, + ) + if in_channels != out_channels: + residual, shortcut_height, shortcut_width = _conv( + network, + residual, + weights, + f"{prefix}.conv_shortcut", + in_channels=in_channels, + out_channels=out_channels, + height=height, + width=width, + ) + if (shortcut_height, shortcut_width) != (height, width): + raise RuntimeError("MiniMax-H3 VAE shortcut changed spatial geometry") + result = network.add_elementwise(residual, update, trt.ElementWiseOperation.SUM) + if result is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 VAE residual {prefix}") + return result.get_output(0), update_height, update_width + + +@op.cleanup_failed_build +def build_keyframe_vae_encoder_engine( + weights: dict[str, np.ndarray], + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the dynamic-tile FL2VA VAE encoder plan.""" + + missing = sorted(set(checkpoint_keys()) - set(weights)) + unexpected = sorted(set(weights) - set(checkpoint_keys())) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 FL2VA VAE encoder checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=KEYFRAME_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = keyframe_vae_encoder_abi() + binding = abi.inputs[0] + pixels = network.add_input( + binding.name, + trt.float32, + (-1, IN_CHANNELS, 1, VAE_TILE_SIZE, VAE_TILE_SIZE), + ) + profile = builder.create_optimization_profile() + profile.set_shape(binding.name, binding.min_shape, binding.opt_shape, binding.max_shape) + config.add_optimization_profile(profile) + + # T=1 is removed only after the ABI boundary. Every causal 3-D + # convolution below selects its last temporal kernel tap. + squeeze = network.add_shuffle(pixels) + squeeze.reshape_dims = (-1, IN_CHANNELS, VAE_TILE_SIZE, VAE_TILE_SIZE) + hidden = squeeze.get_output(0) + height = width = VAE_TILE_SIZE + hidden, height, width = _conv( + network, + hidden, + weights, + "encoder.conv_in", + in_channels=IN_CHANNELS, + out_channels=BLOCK_OUT_CHANNELS[0], + height=height, + width=width, + symmetric_reflect=1, + ) + + current_channels = BLOCK_OUT_CHANNELS[0] + for block_index, out_channels in enumerate(BLOCK_OUT_CHANNELS): + for layer_index in range(LAYERS_PER_BLOCK): + prefix = f"encoder.down_blocks.{block_index}.resnets.{layer_index}" + hidden, height, width = _resnet( + network, + hidden, + weights, + prefix, + in_channels=current_channels, + out_channels=out_channels, + height=height, + width=width, + ) + current_channels = out_channels + if SPATIAL_DOWNSAMPLE_FACTORS[block_index] * TEMPORAL_DOWNSAMPLE_FACTORS[block_index] > 1: + prefix = f"encoder.down_blocks.{block_index}.downsamplers.0.conv" + spatial_stride = SPATIAL_DOWNSAMPLE_FACTORS[block_index] + hidden, height, width = _conv( + network, + hidden, + weights, + prefix, + in_channels=current_channels, + out_channels=current_channels, + height=height, + width=width, + stride=spatial_stride, + bottom_right_reflect=int(spatial_stride == 2), + ) + + expected_side = VAE_TILE_SIZE // VAE_SPATIAL_COMPRESSION + if (height, width) != (expected_side, expected_side): + raise RuntimeError( + f"MiniMax-H3 VAE encoder produced {height}x{width}, expected {expected_side}x{expected_side}" + ) + hidden = _group_norm(network, hidden, weights, "encoder.norm_out", current_channels) + hidden = op.silu(network, hidden) + hidden, height, width = _conv( + network, + hidden, + weights, + "encoder.conv_out", + in_channels=current_channels, + out_channels=MOMENT_CHANNELS, + height=height, + width=width, + symmetric_reflect=1, + ) + hidden, height, width = _conv( + network, + hidden, + weights, + "quant_conv", + in_channels=MOMENT_CHANNELS, + out_channels=MOMENT_CHANNELS, + height=height, + width=width, + ) + expand = network.add_shuffle(hidden) + expand.reshape_dims = (-1, MOMENT_CHANNELS, 1, expected_side, expected_side) + output = expand.get_output(0) + output.name = abi.outputs[0].name + network.mark_output(output) + + op.validate_native_network(network, expected_attentions=0, label="FL2VA keyframe VAE encoder") + print( + f"[minimax-h3] building native FL2VA VAE encoder: tiles=1..{VAE_TILE_MAX_BATCH}, " + f"tile={VAE_TILE_SIZE}x{VAE_TILE_SIZE}, moments={MOMENT_CHANNELS}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 FL2VA VAE encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/graph_ops.py b/families/minimax_h3/graph_ops.py index d81c8ab133..6948560141 100644 --- a/families/minimax_h3/graph_ops.py +++ b/families/minimax_h3/graph_ops.py @@ -10,6 +10,8 @@ from __future__ import annotations from collections import Counter +from contextvars import ContextVar +from functools import lru_cache, wraps import math import ml_dtypes @@ -18,18 +20,55 @@ import tensorrt as trt from .config import resolve_workspace_bytes +from .quantized_checkpoint import ConvRotInt8Weight # TensorRT's explicit BF16 ``Weights`` constructor stores a pointer rather # than owning its input. Retain every backing array until the associated # network has finished building, including temporary packed QKV buffers. _WEIGHT_BUFFER_KEEPALIVE: dict[int, list[np.ndarray]] = {} +_ACTIVE_BUILD_NETWORKS: ContextVar[set[int] | None] = ContextVar( + "minimax_h3_active_build_networks", default=None +) -def configure_builder(config) -> None: - """Retain enough engine metadata to audit native TensorRT lowering.""" +def cleanup_failed_build(function): + """Release retained arrays and consumed weights after graph-build failures.""" + + @wraps(function) + def wrapped(weights, *args, **kwargs): + network_ids: set[int] = set() + token = _ACTIVE_BUILD_NETWORKS.set(network_ids) + try: + return function(weights, *args, **kwargs) + except BaseException: + for network_id in network_ids: + _WEIGHT_BUFFER_KEEPALIVE.pop(network_id, None) + if kwargs.get("consume_weights", False): + weights.clear() + raise + finally: + _ACTIVE_BUILD_NETWORKS.reset(token) + + return wrapped + + +def configure_builder(config, *, weight_streaming: bool = False) -> None: + """Retain graph metadata and enable RTX weight streaming when requested.""" config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED + if not isinstance(weight_streaming, bool): + raise ValueError("MiniMax-H3 weight_streaming must be a boolean") + if not weight_streaming: + return + flag = getattr(trt.BuilderFlag, "WEIGHT_STREAMING", None) + if flag is None: + raise RuntimeError( + "Selected TensorRT-RTX bindings do not expose BuilderFlag.WEIGHT_STREAMING" + ) + config.set_flag(flag) + if not config.get_flag(flag): + raise RuntimeError("TensorRT-RTX did not enable MiniMax-H3 weight streaming") def configure_workspace(config, workspace_bytes: int | None, *, default_bytes: int) -> int: @@ -47,8 +86,13 @@ def configure_workspace(config, workspace_bytes: int | None, *, default_bytes: i return applied -def validate_native_network(network, *, expected_attentions: int, label: str) -> dict[str, int]: - """Fail closed if a native H3 graph gains plugins or collectives.""" +def validate_native_network( + network, + *, + expected_attentions: int, + label: str, +) -> dict[str, int]: + """Fail closed on any layer outside the selected native H3 contract.""" counts = Counter(network.get_layer(index).type for index in range(network.num_layers)) expected = { @@ -79,7 +123,11 @@ def validate_native_network(network, *, expected_attentions: int, label: str) -> def _add_constant(network, array: np.ndarray): array = np.ascontiguousarray(array) - _WEIGHT_BUFFER_KEEPALIVE.setdefault(id(network), []).append(array) + network_id = id(network) + _WEIGHT_BUFFER_KEEPALIVE.setdefault(network_id, []).append(array) + active_networks = _ACTIVE_BUILD_NETWORKS.get() + if active_networks is not None: + active_networks.add(network_id) if array.dtype == np.dtype(ml_dtypes.bfloat16): weights = trt.Weights(trt.bfloat16, array.ctypes.data, array.size) else: @@ -113,6 +161,247 @@ def cast(network, tensor, dtype): return network.add_cast(tensor, dtype).get_output(0) +def _name_convrot_layer(network, layer, role: str) -> None: + """Assign a stable, network-unique diagnostic name to a ConvRot layer.""" + + layer.name = f"minimax_h3.convrot.{role}.{network.num_layers - 1}" + + +@lru_cache(maxsize=None) +def _regular_hadamard(group_size: int) -> np.ndarray: + """Return Comfy's normalized regular Hadamard matrix in checkpoint dtype. + + Comfy ConvRot starts from its symmetric 4x4 basis and takes Kronecker + powers. H3 currently uses groups of 64 and 256. Keeping the coefficients + in BF16 exactly matches the activation dtype used by the published + checkpoint and avoids introducing a second activation quantizer. + """ + + if not isinstance(group_size, int) or isinstance(group_size, bool) or group_size < 4: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + basis = np.asarray( + ( + (1.0, 1.0, 1.0, -1.0), + (1.0, 1.0, -1.0, 1.0), + (1.0, -1.0, 1.0, 1.0), + (-1.0, 1.0, 1.0, 1.0), + ), + dtype=np.float32, + ) + matrix = basis + while matrix.shape[0] < group_size: + matrix = np.kron(matrix, basis) + if matrix.shape != (group_size, group_size): + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + return np.ascontiguousarray( + matrix / math.sqrt(float(group_size)), dtype=ml_dtypes.bfloat16 + ) + + +def _convrot_activation(network, tensor, *, in_features: int, group_size: int): + """Apply the checkpoint's grouped right-Hadamard rotation with native TRT.""" + + if in_features % group_size: + raise ValueError( + "MiniMax-H3 ConvRot input width must be divisible by group_size: " + f"width={in_features}, group_size={group_size}" + ) + if len(tuple(tensor.shape)) != 2: + raise ValueError("MiniMax-H3 ConvRot linears require rank-2 activations") + static_width = int(tensor.shape[1]) + if static_width >= 0 and static_width != in_features: + raise ValueError( + "MiniMax-H3 ConvRot activation/weight width mismatch: " + f"activation={static_width}, weight={in_features}" + ) + + value = cast(network, tensor, trt.bfloat16) + grouped = network.add_shuffle(value) + # Collapse rows and groups into one GEMM M dimension. This is materially + # faster than issuing one tiny batched GEMM per sequence row while keeping + # the exact same contiguous group boundaries. + grouped.reshape_dims = (-1, group_size) + hadamard = weight_constant(network, _regular_hadamard(group_size)) + rotation = network.add_matrix_multiply( + grouped.get_output(0), + trt.MatrixOperation.NONE, + hadamard, + trt.MatrixOperation.NONE, + ) + if rotation is None: + raise RuntimeError("TensorRT rejected the MiniMax-H3 ConvRot activation rotation") + _name_convrot_layer(network, rotation, f"group_{group_size}") + rotation.metadata = ( + "trtmc.quantization=int8_tensorwise_convrot;" + f"activation=bf16;group_size={group_size}" + ) + restored = network.add_shuffle(rotation.get_output(0)) + restored.reshape_dims = (-1, in_features) + return restored.get_output(0) + + +def _convrot_int8_linear(network, tensor, weight: ConvRotInt8Weight, bias=None): + """Build Comfy-compatible dynamic-rowwise ConvRot W8A8 with native TRT. + + TensorRT's ordinary Quantize layer requires a build-time-constant scale. + To retain the published checkpoint's exact dynamic semantics without a + plugin, express the BF16 rowwise quantizer as native math, then place unit + Q/DQ markers around the resulting INT8 tensors. TensorRT recognizes those + markers and lowers the MatrixMultiply to an INT8 GEMM; the dynamic row and + per-output weight scales are applied explicitly in FP32 before BF16 output. + """ + + qweight = np.asarray(weight.qweight) + scale = np.asarray(weight.scale) + if qweight.dtype != np.int8 or qweight.ndim != 2: + raise ValueError("MiniMax-H3 ConvRot qweight must be rank-2 INT8") + if scale.dtype != np.float32 or scale.shape not in { + (qweight.shape[0],), + (qweight.shape[0], 1), + }: + raise ValueError( + "MiniMax-H3 ConvRot scale must be FP32 [out] or [out,1]: " + f"weight={qweight.shape}, scale={scale.shape}" + ) + + rotated = _convrot_activation( + network, + tensor, + in_features=int(qweight.shape[1]), + group_size=int(weight.group_size), + ) + + # Match comfy-kitchen >= 0.2.15 exactly: absmax is computed at the source + # BF16 dtype, scale storage is FP32, scale math and division return to BF16, + # and ROUND is round-to-nearest-even before the signed INT8 clamp. + absolute = network.add_unary(rotated, trt.UnaryOperation.ABS) + if absolute is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation abs") + _name_convrot_layer(network, absolute, "activation_abs_bf16") + row_max = network.add_reduce( + absolute.get_output(0), + trt.ReduceOperation.MAX, + 1 << 1, + True, + ) + if row_max is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot rowwise maximum") + _name_convrot_layer(network, row_max, "activation_row_max_bf16") + row_max_f32 = cast(network, row_max.get_output(0), trt.float32) + divisor = constant(network, np.full((1, 1), 127.0, dtype=np.float32)) + row_scale = network.add_elementwise( + row_max_f32, + divisor, + trt.ElementWiseOperation.DIV, + ) + if row_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot row scale") + _name_convrot_layer(network, row_scale, "activation_scale_f32") + minimum_scale = constant(network, np.full((1, 1), 1.0e-30, dtype=np.float32)) + clamped_scale = network.add_elementwise( + row_scale.get_output(0), + minimum_scale, + trt.ElementWiseOperation.MAX, + ) + if clamped_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot row scale clamp") + _name_convrot_layer(network, clamped_scale, "activation_scale_clamped_f32") + row_scale_f32 = clamped_scale.get_output(0) + row_scale_bf16 = cast(network, row_scale_f32, trt.bfloat16) + divided = network.add_elementwise( + rotated, + row_scale_bf16, + trt.ElementWiseOperation.DIV, + ) + if divided is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot BF16 quantization division") + _name_convrot_layer(network, divided, "activation_divide_bf16") + rounded = network.add_unary(divided.get_output(0), trt.UnaryOperation.ROUND) + if rounded is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation rounding") + _name_convrot_layer(network, rounded, "activation_round_bf16") + minimum_code = weight_constant( + network, + np.full((1, 1), -128.0, dtype=ml_dtypes.bfloat16), + ) + maximum_code = weight_constant( + network, + np.full((1, 1), 127.0, dtype=ml_dtypes.bfloat16), + ) + clipped_minimum = network.add_elementwise( + rounded.get_output(0), + minimum_code, + trt.ElementWiseOperation.MAX, + ) + if clipped_minimum is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation lower clamp") + clipped = network.add_elementwise( + clipped_minimum.get_output(0), + maximum_code, + trt.ElementWiseOperation.MIN, + ) + if clipped is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot activation upper clamp") + activation_int8 = cast(network, clipped.get_output(0), trt.int8) + + quantized = weight_constant(network, qweight) + unit_scale = constant(network, np.asarray(1.0, dtype=np.float32)) + activation_dequantize = network.add_dequantize( + activation_int8, + unit_scale, + trt.float32, + ) + weight_dequantize = network.add_dequantize( + quantized, + unit_scale, + trt.float32, + ) + if activation_dequantize is None or weight_dequantize is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot unit INT8 Q/DQ markers") + _name_convrot_layer(network, activation_dequantize, "activation_unit_dequantize") + _name_convrot_layer(network, weight_dequantize, "weight_unit_dequantize") + matmul = network.add_matrix_multiply( + activation_dequantize.get_output(0), + trt.MatrixOperation.NONE, + weight_dequantize.get_output(0), + trt.MatrixOperation.TRANSPOSE, + ) + if matmul is None: + raise RuntimeError("TensorRT rejected the MiniMax-H3 ConvRot INT8 linear") + _name_convrot_layer(network, matmul, "dynamic_rowwise_int8_gemm") + matmul.metadata = ( + "trtmc.quantization=int8_tensorwise_convrot;" + "activation=dynamic_int8_rowwise;weight=int8;accumulator=fp32" + ) + accumulator = cast(network, matmul.get_output(0), trt.float32) + weight_scale = weight_constant(network, scale.reshape(1, -1)) + output_scale = network.add_elementwise( + row_scale_f32, + weight_scale, + trt.ElementWiseOperation.PROD, + ) + if output_scale is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot combined output scale") + _name_convrot_layer(network, output_scale, "combined_scale_f32") + scaled = network.add_elementwise( + accumulator, + output_scale.get_output(0), + trt.ElementWiseOperation.PROD, + ) + if scaled is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 ConvRot output scaling") + _name_convrot_layer(network, scaled, "output_scale_f32") + output = scaled.get_output(0) + if bias is not None: + shape = (1,) * (len(tuple(output.shape)) - 1) + (-1,) + bias_tensor = weight_constant(network, np.asarray(bias).reshape(shape)) + bias_tensor = cast(network, bias_tensor, trt.float32) + output = network.add_elementwise( + output, bias_tensor, trt.ElementWiseOperation.SUM + ).get_output(0) + return cast(network, output, trt.bfloat16) + + def linear( network, tensor, @@ -124,6 +413,11 @@ def linear( ): """PyTorch ``[out, in]`` linear expressed as native TensorRT GEMM.""" + if isinstance(weight, ConvRotInt8Weight): + if compute_dtype not in (None, trt.bfloat16) or not bf16: + raise ValueError("MiniMax-H3 ConvRot linears require BF16 activation compute") + return _convrot_int8_linear(network, tensor, weight, bias) + tensor_rank = len(tuple(tensor.shape)) rhs_value = np.asarray(weight) if tensor_rank > 2: @@ -218,26 +512,50 @@ def dynamic_slice(network, tensor, starts: tuple[int, ...], sizes: tuple[int | N if len(starts) != len(sizes): raise ValueError("MiniMax-H3 dynamic slice rank mismatch") - runtime_sizes = [ - _shape_dim(network, tensor, axis) if size is None else size - for axis, size in enumerate(sizes) - ] + runtime_sizes = [] + for axis, size in enumerate(sizes): + if size is not None: + runtime_sizes.append(size) + continue + static_size = int(tensor.shape[axis]) + runtime_sizes.append( + static_size if static_size >= 0 else _shape_dim(network, tensor, axis) + ) + if all(isinstance(size, (int, np.integer)) for size in runtime_sizes): + return network.add_slice( + tensor, + starts, + tuple(int(size) for size in runtime_sizes), + (1,) * len(sizes), + ).get_output(0) initial_sizes = tuple(1 if size is None else size for size in sizes) layer = network.add_slice(tensor, starts, initial_sizes, (1,) * len(sizes)) layer.set_input(2, _shape_vector(network, runtime_sizes)) return layer.get_output(0) -def slice_rows_from_end(network, tensor, *, offset: int, rows: int): - """Take fixed rows from a dynamic 2-D tensor, measured from its end.""" +def slice_rows_like_from_end(network, tensor, reference, *, trailing_reference=None): + """Take runtime-sized rows from ``tensor`` using modality input shapes. + + ``reference`` supplies the number of rows to return. When a trailing + modality is supplied, its runtime row count is included in the offset + measured from the packed sequence end. Only shape tensors participate in + this operation, so the reference contents are never copied or consumed. + """ total_rows = _shape_dim(network, tensor, 0) - offset_tensor = constant(network, np.asarray([offset], dtype=np.int64), dtype=np.int64) + rows = _shape_dim(network, reference, 0) + offset = rows + if trailing_reference is not None: + trailing_rows = _shape_dim(network, trailing_reference, 0) + offset = network.add_elementwise( + offset, trailing_rows, trt.ElementWiseOperation.SUM + ).get_output(0) start_row = network.add_elementwise( - total_rows, offset_tensor, trt.ElementWiseOperation.SUB + total_rows, offset, trt.ElementWiseOperation.SUB ).get_output(0) width = int(tensor.shape[1]) - layer = network.add_slice(tensor, (0, 0), (rows, width), (1, 1)) + layer = network.add_slice(tensor, (0, 0), (1, width), (1, 1)) layer.set_input(1, _shape_vector(network, (start_row, 0))) layer.set_input(2, _shape_vector(network, (rows, width))) return layer.get_output(0) @@ -283,20 +601,50 @@ def swiglu(network, tensor, weight_in, weight_out, ffn_dim: int): return linear(network, hidden, weight_out) -def fused_qkv(network, tensor, weights: dict, prefix: str): +def fused_qkv( + network, + tensor, + weights: dict, + prefix: str, + *, + consume_weights: bool = False, +): """Pack Q/K/V into one TensorRT GEMM, matching Sol-Engine's lossless path.""" - packed_weight = np.concatenate( - [weights[f"{prefix}.to_{name}.weight"] for name in ("q", "k", "v")], axis=0 - ) + keys = tuple(f"{prefix}.to_{name}.weight" for name in ("q", "k", "v")) + sources = tuple(weights[key] for key in keys) + if all(isinstance(source, ConvRotInt8Weight) for source in sources): + parents = tuple(source.packed_parent for source in sources) + if parents[0] is None or not all(parent is parents[0] for parent in parents): + raise ValueError( + "MiniMax-H3 quantized Q/K/V must share one packed parent" + ) + packed_weight = parents[0] + if not packed_weight.is_full_fused_qkv: + raise ValueError("MiniMax-H3 quantized QKV parent is not marked as full fused QKV") + expected_width = int(packed_weight.qweight.shape[0]) // 3 + expected_slices = tuple( + (index * expected_width, (index + 1) * expected_width) for index in range(3) + ) + if tuple(source.row_slice for source in sources) != expected_slices: + raise ValueError("MiniMax-H3 quantized Q/K/V row slices are not q-k-v ordered") + elif any(isinstance(source, ConvRotInt8Weight) for source in sources): + raise ValueError("MiniMax-H3 fused QKV cannot mix quantized and BF16 weights") + else: + packed_weight = np.concatenate(sources, axis=0) packed = linear(network, tensor, packed_weight) + if consume_weights: + # The packed array is retained by the network. Its three sources are + # no longer referenced and can be released before serialization. + for key in keys: + weights.pop(key) width = int(packed.shape[1]) // 3 return tuple( dynamic_slice(network, packed, (0, index * width), (None, width)) for index in range(3) ) -def rows_to_heads(network, tensor, rows: int, heads: int, head_dim: int): +def rows_to_heads(network, tensor, heads: int, head_dim: int): reshape = network.add_shuffle(tensor) reshape.reshape_dims = (-1, heads, head_dim) reshape.second_transpose = trt.Permutation([1, 0, 2]) @@ -305,7 +653,7 @@ def rows_to_heads(network, tensor, rows: int, heads: int, head_dim: int): return batch.get_output(0) -def heads_to_rows(network, tensor, rows: int, width: int): +def heads_to_rows(network, tensor, width: int): reshape = network.add_shuffle(tensor) reshape.first_transpose = trt.Permutation([0, 2, 1, 3]) reshape.reshape_dims = (-1, width) @@ -318,7 +666,6 @@ def partial_rope( cos_half, sin_half, *, - rows: int, heads: int, head_dim: int, rotary_dim: int, @@ -326,7 +673,7 @@ def partial_rope( ): """Apply H3's 96-channel rotate-half MM-RoPE with native layers.""" - value = rows_to_heads(network, tensor, rows, heads, head_dim) + value = rows_to_heads(network, tensor, heads, head_dim) if interleaved: raise ValueError("MiniMax-H3 uses rotate-half, non-interleaved RoPE") rotary = dynamic_slice(network, value, (0, 0, 0, 0), (1, heads, None, rotary_dim)) @@ -357,21 +704,17 @@ def duplicate_table(table): rotated = network.add_elementwise(left, right, trt.ElementWiseOperation.SUM).get_output(0) result = network.add_concatenation((rotated, passthrough)) result.axis = 3 - return heads_to_rows(network, result.get_output(0), rows, heads * head_dim) + return heads_to_rows(network, result.get_output(0), heads * head_dim) -def native_attention(network, q, k, v, *, rows: int, heads: int, head_dim: int, name: str): +def native_attention(network, q, k, v, *, heads: int, head_dim: int, name: str): """Full-sequence single-device fused TensorRT attention.""" - q = rows_to_heads(network, q, rows, heads, head_dim) - k = rows_to_heads(network, k, rows, heads, head_dim) - v = rows_to_heads(network, v, rows, heads, head_dim) - # TensorRT's fused FP16 attention has three additional mantissa bits over - # BF16. Keep the surrounding block in checkpoint-native BF16 while - # reducing recurrent numerical drift across 49 denoising evaluations. - q = cast(network, q, trt.float16) - k = cast(network, k, trt.float16) - v = cast(network, v, trt.float16) + q = rows_to_heads(network, q, heads, head_dim) + k = rows_to_heads(network, k, heads, head_dim) + v = rows_to_heads(network, v, heads, head_dim) + # Preserve BF16's exponent range. H3 residuals can exceed FP16's finite + # limit before a later block normalizes them. scale = constant( network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(head_dim), dtype=np.float32), @@ -385,5 +728,4 @@ def native_attention(network, q, k, v, *, rows: int, heads: int, head_dim: int, layer.metadata = f"trtmc.native_op=IAttention;source={name}" layer.get_output(0).name = f"{name}.output" layer.decomposable = False - context = cast(network, layer.get_output(0), trt.bfloat16) - return heads_to_rows(network, context, rows, heads * head_dim) + return heads_to_rows(network, layer.get_output(0), heads * head_dim) diff --git a/families/minimax_h3/model.py b/families/minimax_h3/model.py index e33e746f5a..c6f6cbb697 100644 --- a/families/minimax_h3/model.py +++ b/families/minimax_h3/model.py @@ -5,37 +5,61 @@ from __future__ import annotations -import gc +from fractions import Fraction import json +import math from dataclasses import replace from pathlib import Path from types import SimpleNamespace from typing import TYPE_CHECKING -from .checkpoint import ( - load_selected_component_state_dict, - numpy_state, - validate_component_key_partition, -) from .config import ( - SOL_ENGINE_1344X768_124F, - default_workspace_limit_bytes, + CANVAS_MAX_ASPECT_RATIO, + CANVAS_MAX_PIXELS, + CANVAS_MIN_ASPECT_RATIO, + CANVAS_MULTIPLE, + CANVAS_SHORT_EDGE, + NATIVE_EXPLICIT_CANVAS_SIZES, + SOL_ENGINE_1344X768_124_TO_345F, + VIDEO_NUM_FRAMES_MAX, + VIDEO_NUM_FRAMES_MIN, + VIDEO_NUM_FRAMES_OPT, ) +from .runtime_config_schema import normalize_build_options as validate_build_options +from .delivery import SR_SOURCE_SHAPE if TYPE_CHECKING: from tensorrt_model_connect.build import BuildRequest from tensorrt_model_connect.bundle_writer import BundleWriter -def _fixed_profile(raw: dict): +def _effective_build_config(raw: dict) -> dict: + family_options = raw.get("_family_build_options", {}) + options = family_options.get("minimax_h3", {}) if isinstance(family_options, dict) else {} + if not isinstance(options, dict): + raise ValueError("minimax_h3 build options must be an object") + return {**raw, **validate_build_options(options)} + + +def _public_dynamic_profile(raw: dict): + profile = SOL_ENGINE_1344X768_124_TO_345F expected = { - "text_rows": SOL_ENGINE_1344X768_124F.text_rows, - "text_rows_min": SOL_ENGINE_1344X768_124F.min_text_rows, - "text_rows_opt": SOL_ENGINE_1344X768_124F.opt_text_rows, - "text_rows_max": SOL_ENGINE_1344X768_124F.text_rows, - "audio_rows": SOL_ENGINE_1344X768_124F.audio_rows, - "video_rows": SOL_ENGINE_1344X768_124F.video_rows, - "padded_sequence_length": SOL_ENGINE_1344X768_124F.padded_sequence_length, + "text_rows": profile.text_rows, + "text_rows_min": profile.min_text_rows, + "text_rows_opt": profile.opt_text_rows, + "text_rows_max": profile.text_rows, + "audio_rows": profile.opt_audio_rows, + "audio_rows_min": profile.min_audio_rows, + "audio_rows_opt": profile.opt_audio_rows, + "audio_rows_max": profile.audio_rows, + "video_rows": profile.opt_video_rows, + "video_rows_min": profile.min_video_rows, + "video_rows_opt": profile.opt_video_rows, + "video_rows_max": profile.video_rows, + "packed_sequence_length_min": profile.min_sequence_length, + "packed_sequence_length_opt": profile.opt_sequence_length, + "packed_sequence_length_max": profile.sequence_length, + "padded_sequence_length": profile.padded_sequence_length, } mismatches = { name: (raw[name], value) @@ -44,28 +68,141 @@ def _fixed_profile(raw: dict): } if mismatches: raise ValueError(f"Unsupported MiniMax-H3 packed-row profile: {mismatches}") - explicit_flag = raw.get("first_block_cache") - mode = raw.get( - "denoiser_cache_mode", - "first_block" if explicit_flag is True else "monolithic", - ) - if mode not in ("monolithic", "first_block"): - raise ValueError(f"Unsupported MiniMax-H3 denoiser_cache_mode: {mode!r}") - if explicit_flag is not None and not isinstance(explicit_flag, bool): - raise ValueError("MiniMax-H3 first_block_cache must be a boolean") - mode_flag = mode == "first_block" - if explicit_flag is not None and explicit_flag != mode_flag: - raise ValueError("MiniMax-H3 cache mode and first_block_cache flag disagree") - if not mode_flag: - return SOL_ENGINE_1344X768_124F - return replace(SOL_ENGINE_1344X768_124F, first_block_cache=True) - - -class _MiniMaxH3Model: - def load_weights(self, model_dir: str, config) -> dict: + explicit_flag = raw.get("first_block_cache", True) + if explicit_flag is not True: + raise ValueError("MiniMax-H3 only supports the dense FirstBlockCache build") + mode = raw.get("denoiser_cache_mode", "first_block") + if mode != "first_block": + raise ValueError("MiniMax-H3 only supports denoiser_cache_mode='first_block'") + return replace(profile, first_block_cache=True) + + +def _default_num_frames(raw: dict) -> int: + value = raw.get("video_num_frames", raw.get("num_frames", VIDEO_NUM_FRAMES_OPT)) + if isinstance(value, bool): + raise ValueError("MiniMax-H3 video_num_frames must be a valid 5--15 second geometry") + try: + frames = int(value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 video_num_frames must be a valid 5--15 second geometry" + ) from error + if not VIDEO_NUM_FRAMES_MIN <= frames <= VIDEO_NUM_FRAMES_MAX or frames % 17 != 5: + raise ValueError("MiniMax-H3 video_num_frames must be a valid 5--15 second geometry") + return frames + + +def _resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: + """Mirror the public H3 resolver with Python ties-to-even 32 rounding.""" + + if not math.isfinite(aspect_width) or not math.isfinite(aspect_height): + raise ValueError("MiniMax-H3 canvas aspect must be finite and positive") + if aspect_width <= 0.0 or aspect_height <= 0.0: + raise ValueError("MiniMax-H3 canvas aspect must be finite and positive") + ratio = aspect_width / aspect_height + if not CANVAS_MIN_ASPECT_RATIO <= ratio <= CANVAS_MAX_ASPECT_RATIO: + raise ValueError("MiniMax-H3 canvas aspect must be within 1:4 through 4:1") + if ratio >= 1.0: + height = float(CANVAS_SHORT_EDGE) + width = height * ratio + else: + width = float(CANVAS_SHORT_EDGE) + height = width / ratio + pixels = width * height + if pixels > CANVAS_MAX_PIXELS: + scale = math.sqrt(CANVAS_MAX_PIXELS / pixels) + width *= scale + height *= scale + resolved_width = int(round(width / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE + resolved_height = int(round(height / CANVAS_MULTIPLE)) * CANVAS_MULTIPLE + return resolved_height, resolved_width + + +def _reachable_canvas_sizes() -> tuple[tuple[int, int], ...]: + """Enumerate the exact finite image of the continuous 32-rounded resolver.""" + + landscape = { + (CANVAS_SHORT_EDGE, width) + for width in range( + CANVAS_SHORT_EDGE, + int(CANVAS_SHORT_EDGE * 1.75) + CANVAS_MULTIPLE, + CANVAS_MULTIPLE, + ) + } + half = CANVAS_MULTIPLE // 2 + maximum_dimension = CANVAS_SHORT_EDGE * int(CANVAS_MAX_ASPECT_RATIO) + for height in range(CANVAS_MULTIPLE, CANVAS_SHORT_EDGE + 1, CANVAS_MULTIPLE): + for width in range(CANVAS_SHORT_EDGE, maximum_dimension + 1, CANVAS_MULTIPLE): + # In the area-limited landscape branch raw dimensions satisfy + # h*w=max_pixels and r=w/h. Intersect the two nearest-multiple + # rounding cells with the resolver's exact r interval. + lower = max( + Fraction(7, 4), + Fraction(CANVAS_MAX_PIXELS, (height + half) ** 2), + Fraction((width - half) ** 2, CANVAS_MAX_PIXELS), + ) + upper = min( + Fraction(4, 1), + Fraction(CANVAS_MAX_PIXELS, (height - half) ** 2), + Fraction((width + half) ** 2, CANVAS_MAX_PIXELS), + ) + if lower >= upper: + continue + sample = float((lower + upper) / 2) + if _resolve_canvas_size(sample, 1.0) == (height, width): + landscape.add((height, width)) + reachable = landscape | {(width, height) for height, width in landscape} + return tuple(sorted(reachable)) + + +def _default_canvas_size(raw: dict) -> tuple[int, int]: + sr = raw.get("super_resolution", False) + default_height, default_width = SR_SOURCE_SHAPE if sr else (768, 1344) + height_value = raw.get("video_height", raw.get("height", default_height)) + width_value = raw.get("video_width", raw.get("width", default_width)) + if isinstance(height_value, bool) or isinstance(width_value, bool): + raise ValueError("MiniMax-H3 video dimensions must match the public canvas resolver") + try: + height = int(height_value) + width = int(width_value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 video dimensions must match the public canvas resolver" + ) from error + if sr and (height, width) != SR_SOURCE_SHAPE: + raise ValueError("MiniMax-H3 super_resolution=true requires height=480 and width=864") + if ( + height <= 0 + or width <= 0 + or ( + (height, width) not in NATIVE_EXPLICIT_CANVAS_SIZES + and (height, width) != _resolve_canvas_size(width, height) + ) + ): + raise ValueError("MiniMax-H3 video dimensions must match the public canvas resolver") + return height, width + + +def _first_block_cache_threshold(raw: dict, *, default: float = 0.08) -> float: + value = raw.get("first_block_cache_threshold", default) + if isinstance(value, bool): + raise ValueError("MiniMax-H3 first_block_cache_threshold must be finite and positive") + try: + threshold = float(value) + except (TypeError, ValueError) as error: + raise ValueError( + "MiniMax-H3 first_block_cache_threshold must be finite and positive" + ) from error + if not math.isfinite(threshold) or threshold <= 0.0: + raise ValueError("MiniMax-H3 first_block_cache_threshold must be finite and positive") + return threshold + + +class MiniMaxH3Plugin: + def load_weights(self, model_dir: str, config, **_kwargs) -> dict: del config root = Path(model_dir) - required_dirs = ("transformer", "text_encoder", "vae", "audio_vae", "tokenizer") + required_dirs = ("transformer", "vae", "audio_vae", "tokenizer") missing = [str(root / name) for name in required_dirs if not (root / name).is_dir()] if missing: raise FileNotFoundError( @@ -86,241 +223,131 @@ def load_weights(self, model_dir: str, config) -> dict: } if mismatches: raise ValueError(f"Unsupported MiniMax-H3 transformer architecture: {mismatches}") - return { - "_model_dir": str(root), - "_transformer_dir": str(root / "transformer"), - "_text_encoder_dir": str(root / "text_encoder"), - "_vae_dir": str(root / "vae"), - "_audio_vae_dir": str(root / "audio_vae"), - "_tokenizer_dir": str(root / "tokenizer"), - } + return {"_model_dir": str(root)} - def build_components( + def build_staged_bundle( self, model_dir: str, + writer: "BundleWriter", config, weights: dict, *, - precision: str = "bf16", + plans_dir: str | Path, + precision: str, verbose: bool = False, - ) -> dict: - del model_dir - if precision.lower() != "bf16": - raise ValueError("MiniMax-H3 native builds require BF16 checkpoint weights") - raw = getattr(config, "raw", {}) - profile = _fixed_profile(raw) - profile.validate() - workspace_limits = default_workspace_limit_bytes( - first_block_cache=profile.first_block_cache - ) - from .adaln_builder import build_adaln_precompute_engine - from .adaln_builder import checkpoint_keys as adaln_checkpoint_keys - from .dit_builder import ( - build_dit_engine, - build_dit_finish_engine, - build_dit_head_engine, - build_dit_tail_engine, - checkpoint_keys as dit_checkpoint_keys, - finish_checkpoint_keys, - head_checkpoint_keys, - tail_checkpoint_keys, - ) - from .text_encoder_builder import ( - build_text_encoder_engine, - checkpoint_keys as text_encoder_checkpoint_keys, - ) - - if profile.first_block_cache: - denoiser_specs = ( - ( - "denoiser_head", - "denoiser_head.plan", - build_dit_head_engine, - head_checkpoint_keys(profile), - ), - ( - "denoiser_tail", - "denoiser_tail.plan", - build_dit_tail_engine, - tail_checkpoint_keys(profile), - ), - ( - "denoiser_finish", - "denoiser_finish.plan", - build_dit_finish_engine, - finish_checkpoint_keys(profile), - ), - ) - checkpoint_groups = ( - adaln_checkpoint_keys(profile), - *(spec[3] for spec in denoiser_specs), - ) - else: - denoiser_specs = ( - ( - "denoiser", - "denoiser.plan", - build_dit_engine, - dit_checkpoint_keys(profile), - ), - ) - checkpoint_groups = ( - adaln_checkpoint_keys(profile), - dit_checkpoint_keys(profile), - ) - validate_component_key_partition(weights["_transformer_dir"], checkpoint_groups) + parallel_config=None, + max_batch_size: int = 1, + ) -> None: + """Build a staged RTX bundle without retaining serialized plans in RAM.""" - text_state = load_selected_component_state_dict( - weights["_text_encoder_dir"], text_encoder_checkpoint_keys() - ) - text_weights = numpy_state(text_state) - del text_state - text_encoder_plan = build_text_encoder_engine( - text_weights, - sequence_length=profile.text_rows, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["text_encoder.plan"], - ) - del text_weights - gc.collect() + if precision.lower() != "bf16": + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require BF16") + if max_batch_size != 1: + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require max_batch_size=1") + mode = str(getattr(parallel_config, "mode", "single")) + if mode != "single": + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds require one GPU") - adaln_state = load_selected_component_state_dict( - weights["_transformer_dir"], adaln_checkpoint_keys(profile) - ) - adaln_weights = numpy_state(adaln_state) - del adaln_state - adaln_plan = build_adaln_precompute_engine( - adaln_weights, - profile, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["adaln_precompute.plan"], - ) - del adaln_weights - gc.collect() + raw = _effective_build_config(getattr(config, "raw", {})) + if raw.get("_fp32_layers"): + raise ValueError("MiniMax-H3 TensorRT-RTX staged builds do not support FP32 layers") + from .delivery import resolve_quantized_sources, resolve_super_resolution_sources - denoiser_components = {} - for component_name, filename, denoiser_builder, selected_keys in denoiser_specs: - dit_state = load_selected_component_state_dict( - weights["_transformer_dir"], selected_keys - ) - dit_weights = numpy_state(dit_state) - del dit_state - denoiser_plan = denoiser_builder( - dit_weights, - profile, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits[filename], - ) - del dit_weights - gc.collect() - denoiser_components[component_name] = denoiser_plan + root = Path(weights.get("_model_dir", model_dir)) + staged_raw = dict(raw) + staged_raw.setdefault("first_block_cache", True) + staged_raw.setdefault("denoiser_cache_mode", "first_block") + _public_dynamic_profile(staged_raw) + expected_request = { + "num_inference_steps": 50, + "seed": 0, + } + mismatches = { + name: (raw[name], value) + for name, value in expected_request.items() + if name in raw and int(raw[name]) != value + } + if mismatches: + raise ValueError(f"Unsupported MiniMax-H3 staged profile: {mismatches}") + _default_canvas_size(raw) + _default_num_frames(raw) - from .vae_builder import ( - build_vae_tile_decoder_engine, - checkpoint_keys as vae_checkpoint_keys, - ) + from .staged_build import build_staged_bundle - vae_state = load_selected_component_state_dict(weights["_vae_dir"], vae_checkpoint_keys()) - vae_weights = numpy_state(vae_state) - del vae_state - vae_decoder_plan = build_vae_tile_decoder_engine( - vae_weights, - verbose=verbose, - consume_weights=True, - workspace_bytes=workspace_limits["vae_tile_decoder.plan"], - ) - tokenizer_json = (Path(weights["_tokenizer_dir"]) / "tokenizer.json").read_bytes() - - return { - "text_encoder": text_encoder_plan, - "adaln_precompute": adaln_plan, - **denoiser_components, - "vae_decoder": vae_decoder_plan, - "profile": profile, - # Text/VAE paths remain explicit so follow-on native component - # builders cannot silently substitute a different checkpoint. - "vae_dir": weights["_vae_dir"], - "audio_vae_dir": weights["_audio_vae_dir"], - "tokenizer_dir": weights["_tokenizer_dir"], - "tokenizer_json": tokenizer_json, + # All public builds use the same quantized checkpoints for all three + # workflows. Only the explicit SR flag changes the delivery mode. + staged_options = {"verbose": verbose, **resolve_quantized_sources(root, raw)} + super_resolution_model, super_resolution_weak_model = resolve_super_resolution_sources(raw) + staged_options["runtime_defaults"] = { + "height": _default_canvas_size(raw)[0], + "width": _default_canvas_size(raw)[1], + "num_frames": _default_num_frames(raw), + "first_block_cache_threshold": _first_block_cache_threshold(raw), } + staged_options["runtime_defaults"].update( + ref2va_first_block_cache=raw.get("ref2va_first_block_cache", True), + ref2va_first_block_cache_threshold=raw.get("ref2va_first_block_cache_threshold", 0.08), + ) + if super_resolution_model is not None: + staged_options["super_resolution_model"] = super_resolution_model + if super_resolution_weak_model is not None: + staged_options["super_resolution_weak_model"] = super_resolution_weak_model + build_staged_bundle(root, writer, plans_dir=plans_dir, **staged_options) -def build(request: "BuildRequest", writer: "BundleWriter") -> None: - """Build one MiniMax-H3 image-generation bundle.""" - if request.dynamic_kv_cache: - raise NotImplementedError("minimax_h3 does not support dynamic_kv_cache") +plugin = MiniMaxH3Plugin() - if request.max_sequence_length is not None: - raise NotImplementedError("minimax_h3 does not support max_sequence_length") - if request.context_parallel_size != 1: - raise ValueError("this family does not support context parallelism") +def build(request: "BuildRequest", writer: "BundleWriter") -> None: + """Build the unified T2VA/FL2VA/Ref2VA bundle through TensorRT-RTX.""" if request.task != "image_generation": raise ValueError("minimax_h3 supports only task=image_generation") + if request.backend != "trt_rtx": + raise ValueError("MiniMax-H3 requires backend=trt_rtx") if request.precision != "bf16": raise ValueError("MiniMax-H3 requires precision=bf16") - if request.tensor_parallel_size != 1: - raise NotImplementedError("MiniMax-H3 requires tensor_parallel_size=1") + if request.dynamic_kv_cache: + raise NotImplementedError("minimax_h3 does not support dynamic_kv_cache") + if request.max_sequence_length is not None: + raise NotImplementedError("minimax_h3 does not support max_sequence_length") + if request.tensor_parallel_size != 1 or request.context_parallel_size != 1: + raise ValueError("MiniMax-H3 requires single-device parallel settings") if request.max_batch_size != 1: - raise NotImplementedError("MiniMax-H3 requires max_batch_size=1") - if request.quantization not in {None, "none"}: - raise NotImplementedError("MiniMax-H3 does not support quantization") + raise ValueError("MiniMax-H3 requires max_batch_size=1") if request.fp32_layers: raise NotImplementedError("MiniMax-H3 does not support fp32_layers") - if int(request.image_height or 768) != 768 or int(request.image_width or 1344) != 1344: - raise ValueError("MiniMax-H3 requires image_height=768 and image_width=1344") - if int(request.video_num_frames or 124) != 124: - raise ValueError("MiniMax-H3 requires video_num_frames=124") - - config = SimpleNamespace(raw={}) - model_dir = Path(request.model_dir) - model = _MiniMaxH3Model() - weights = model.load_weights(str(model_dir), config) - components = model.build_components( - str(model_dir), + + family_options = validate_build_options(dict(getattr(request, "family_options", ()))) + if request.quantization not in {None, "int8_tensorwise_convrot"}: + raise ValueError("MiniMax-H3 delivery requires Comfy INT8 denoisers and the NVFP4 text checkpoint") + if (request.image_height is None) != (request.image_width is None): + raise ValueError("MiniMax-H3 requires both image height and width, or neither") + sr = family_options.get("super_resolution", False) + default_height, default_width = SR_SOURCE_SHAPE if sr else (768, 1344) + + raw = { + "_family_build_options": {"minimax_h3": family_options}, + "height": int(request.image_height or default_height), + "width": int(request.image_width or default_width), + "video_num_frames": int(request.video_num_frames or VIDEO_NUM_FRAMES_OPT), + "num_inference_steps": 50, + "seed": 0, + } + config = SimpleNamespace(raw=raw) + _default_canvas_size(_effective_build_config(raw)) + _default_num_frames(raw) + weights = plugin.load_weights(str(request.model_dir), config) + + writer.set_header(family="minimax_h3", task=request.task, backend=request.backend) + plans_dir = request.output_path.with_name(f"{request.output_path.name}.plans") + plugin.build_staged_bundle( + str(request.model_dir), + writer, config, weights, + plans_dir=plans_dir, precision=request.precision, verbose=request.verbose, - ) - profile = components["profile"] - if profile.first_block_cache: - raise RuntimeError("MiniMax-H3 minimal build uses the monolithic denoiser profile") - - writer.set_header(family="minimax_h3", task=request.task, backend=request.backend) - writer.add_bytes("text_encoder.plan", components["text_encoder"]) - writer.add_bytes("adaln.plan", components["adaln_precompute"]) - writer.add_bytes("denoiser.plan", components["denoiser"]) - writer.add_bytes("vae.plan", components["vae_decoder"]) - writer.add_bytes("tokenizer.json", components["tokenizer_json"]) - writer.add_json( - "runtime.json", - { - "height": 768, - "width": 1344, - "num_frames": 124, - "fps": 24, - "num_inference_steps": 50, - "seed": 0, - "first_block_cache": False, - "denoiser_cache_mode": "monolithic", - "first_block_cache_threshold": 0.025, - "text_rows": profile.text_rows, - "text_rows_min": profile.min_text_rows, - "text_rows_opt": profile.opt_text_rows, - "text_rows_max": profile.text_rows, - "audio_rows": profile.audio_rows, - "video_rows": profile.video_rows, - "padded_sequence_length": profile.padded_sequence_length, - "max_timestep_count": profile.max_timestep_count, - "context_parallel_size": profile.context_parallel_size, - "vae_tile_batch": 28, - "vae_tile_size": 256, - "vae_tile_overlap": 64, - }, + max_batch_size=request.max_batch_size, ) diff --git a/families/minimax_h3/multimodal_text_encoder_builder.py b/families/minimax_h3/multimodal_text_encoder_builder.py new file mode 100644 index 0000000000..dd9ae604b4 --- /dev/null +++ b/families/minimax_h3/multimodal_text_encoder_builder.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared native Qwen3-VL language stack for MiniMax-H3 public workflows. + +The plan contains one copy of the Qwen embedding and decoder layers 0..49. +T2VA binds one disabled dummy visual row; FL2VA and Ref2VA bind compact vision +features plus their presentation-row indices. The engine scatters the main +features before layer zero, injects DeepStack features after layers 0..2, and +emits ``hidden_states[50]`` before the full model's final normalization. +""" + +from __future__ import annotations + +import gc +import math +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .config import TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES +from .fl2va_contract import MultimodalTextProfile, QWEN_TEXT_HIDDEN_SIZE, text_encoder_abi + + +trt = trt_compat.get_trt() + +HIDDEN_SIZE = QWEN_TEXT_HIDDEN_SIZE +NUM_LAYERS = 50 +NUM_HEADS = 64 +NUM_KV_HEADS = 8 +HEAD_DIM = 128 +NORM_EPS = 1.0e-6 +# Bit-exact output of Qwen3VLTextRotaryEmbedding's Torch FP32 construction. +# Recomputing ``theta**exponent`` with NumPy differs by one ULP at index 19. +ROPE_INV_FREQ_BITS = ( + 0x3F800000, + 0x3F492C28, + 0x3F1E165D, + 0x3EF875A7, + 0x3EC33F3A, + 0x3E996E52, + 0x3E712429, + 0x3E3D7EFC, + 0x3E14E963, + 0x3DEA09DB, + 0x3DB7EA1A, + 0x3D908687, + 0x3D63251B, + 0x3D327F50, + 0x3D0C44BF, + 0x3CDC7456, + 0x3CAD3D5E, + 0x3C88230F, + 0x3C55F605, + 0x3C282311, + 0x3C042088, + 0x3BCFA8A9, + 0x3BA32F3E, + 0x3B803C3D, + 0x3B498AD3, + 0x3B1E60C3, + 0x3AF8EA93, + 0x3AC39B1C, + 0x3A99B686, + 0x3A7195A5, + 0x3A3DD829, + 0x3A152F76, + 0x39EA77FF, + 0x39B840A8, + 0x3990CA8B, + 0x39639000, + 0x3932D34F, + 0x390C86C1, + 0x38DCDC14, + 0x38AD8EE4, + 0x38886320, + 0x38565AB5, + 0x38287230, + 0x38045EB6, + 0x37D00A62, + 0x37A37C09, + 0x37807896, + 0x3749E9AB, + 0x371EAB4B, + 0x36F95FB6, + 0x36C3F72A, + 0x3699FEDC, + 0x36720755, + 0x363E3180, + 0x361575AB, + 0x35EAE655, + 0x35B8975D, + 0x35910EAE, + 0x3563FB18, + 0x35332777, + 0x350CC8E3, + 0x34DD4405, + 0x34ADE091, + 0x3488A34F, +) + + +def checkpoint_keys() -> tuple[str, ...]: + """Return the exhaustive shared language-weight partition.""" + + names = ["model.language_model.embed_tokens.weight"] + for index in range(NUM_LAYERS): + prefix = f"model.language_model.layers.{index}" + names.extend( + [ + f"{prefix}.input_layernorm.weight", + f"{prefix}.post_attention_layernorm.weight", + f"{prefix}.self_attn.q_proj.weight", + f"{prefix}.self_attn.k_proj.weight", + f"{prefix}.self_attn.v_proj.weight", + f"{prefix}.self_attn.o_proj.weight", + f"{prefix}.self_attn.q_norm.weight", + f"{prefix}.self_attn.k_norm.weight", + f"{prefix}.mlp.gate_proj.weight", + f"{prefix}.mlp.up_proj.weight", + f"{prefix}.mlp.down_proj.weight", + ] + ) + return tuple(names) + + +def _network_shape(binding) -> tuple[int, ...]: + return tuple( + minimum if minimum == maximum else -1 + for minimum, maximum in zip(binding.min_shape, binding.max_shape) + ) + + +def _per_head_norm(network, tensor, weight, heads: int): + reshape = network.add_shuffle(tensor) + reshape.reshape_dims = (-1, heads, HEAD_DIM) + normalized = op.rms_norm(network, reshape.get_output(0), weight, HEAD_DIM, NORM_EPS) + flatten = network.add_shuffle(normalized) + flatten.reshape_dims = (-1, heads * HEAD_DIM) + return flatten.get_output(0) + + +def _repeat_kv(network, tensor): + repeated = [] + repeat = NUM_HEADS // NUM_KV_HEADS + for index in range(NUM_KV_HEADS): + head = op.dynamic_slice(network, tensor, (0, index, 0, 0), (1, 1, None, HEAD_DIM)) + repeated.extend([head] * repeat) + concatenation = network.add_concatenation(repeated) + concatenation.axis = 1 + return concatenation.get_output(0) + + +def _mrope_cache(network, position_ids): + """Build Qwen3-VL interleaved ``[T,H,W]`` MRoPE in FP32.""" + + positions = op.cast(network, position_ids, trt.float32) + inverse = np.asarray(ROPE_INV_FREQ_BITS, dtype=np.uint32).view(np.float32) + inverse = op.constant(network, inverse.reshape(1, HEAD_DIM // 2)) + frequencies = [] + for axis in range(3): + coordinate = op.dynamic_slice(network, positions, (axis, 0), (1, None)) + reshape = network.add_shuffle(coordinate) + reshape.reshape_dims = (-1, 1) + frequencies.append( + network.add_elementwise( + reshape.get_output(0), inverse, trt.ElementWiseOperation.PROD + ).get_output(0) + ) + + height_mask = np.zeros((1, HEAD_DIM // 2), dtype=np.float32) + width_mask = np.zeros_like(height_mask) + height_mask[:, 1:60:3] = 1.0 + width_mask[:, 2:60:3] = 1.0 + temporal_mask = 1.0 - height_mask - width_mask + selected = [] + for frequency, mask in zip(frequencies, (temporal_mask, height_mask, width_mask)): + selected.append( + network.add_elementwise( + frequency, + op.constant(network, mask), + trt.ElementWiseOperation.PROD, + ).get_output(0) + ) + frequency = network.add_elementwise( + selected[0], selected[1], trt.ElementWiseOperation.SUM + ).get_output(0) + frequency = network.add_elementwise( + frequency, selected[2], trt.ElementWiseOperation.SUM + ).get_output(0) + cos = network.add_unary(frequency, trt.UnaryOperation.COS).get_output(0) + sin = network.add_unary(frequency, trt.UnaryOperation.SIN).get_output(0) + return op.cast(network, cos, trt.bfloat16), op.cast(network, sin, trt.bfloat16) + + +def _visual_count_active(network, vision_count): + zero = op.constant(network, np.zeros((1,), dtype=np.int32), dtype=np.int32) + active = network.add_elementwise( + vision_count, zero, trt.ElementWiseOperation.GREATER + ).get_output(0) + reshape = network.add_shuffle(active) + reshape.reshape_dims = (1, 1) + return reshape.get_output(0) + + +def _scatter_visual_rows(network, base_like, row_indices, compact, active): + """Scatter compact valid features into a dynamic sequence-sized zero base.""" + + zero = op.constant(network, np.zeros((1, 1), dtype=np.float32)) + zero = op.cast(network, zero, base_like.dtype) + base = network.add_elementwise(base_like, zero, trt.ElementWiseOperation.PROD).get_output(0) + compact = op.cast(network, compact, base_like.dtype) + compact = network.add_select(active, compact, zero).get_output(0) + indices = network.add_shuffle(row_indices) + indices.reshape_dims = (-1, 1) + scatter = network.add_scatter( + base, + indices.get_output(0), + compact, + trt.ScatterMode.ND, + ) + if scatter is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 compact visual-row scatter") + return scatter.get_output(0) + + +def _linear(network, hidden, weights, name: str): + prescale = weights.get(f"{name}.pre_quant_scale") + if prescale is not None: + prescale = op.weight_constant(network, np.asarray(prescale).reshape(1, -1)) + prescale = op.cast(network, prescale, hidden.dtype) + hidden = network.add_elementwise( + hidden, prescale, trt.ElementWiseOperation.PROD + ).get_output(0) + return op.linear(network, hidden, weights[f"{name}.weight"]) + + +@op.cleanup_failed_build +def build_multimodal_text_encoder_engine( + weights: dict[str, np.ndarray], + profile: MultimodalTextProfile = MultimodalTextProfile(), + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the unified T2VA/FL2VA/Ref2VA Qwen language plan.""" + + profile.validate() + expected_keys = set(checkpoint_keys()) + optional_keys = { + f"model.language_model.layers.{index}.{name}.pre_quant_scale" + for index in range(NUM_LAYERS) + for name in ("self_attn.o_proj", "mlp.down_proj") + } + missing = sorted(expected_keys - set(weights)) + unexpected = sorted(set(weights) - expected_keys - optional_keys) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 multimodal text checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + for name in optional_keys & set(weights): + width = np.asarray(weights[f"{name.removesuffix('.pre_quant_scale')}.weight"]).shape[1] + value = np.asarray(weights[name]) + if value.shape != (width,) or not np.isfinite(value).all(): + raise ValueError(f"Invalid MiniMax-H3 AWQ input scale: {name}") + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=TEXT_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = text_encoder_abi(profile) + dtype_map = {"float32": trt.float32, "int32": trt.int32} + inputs = { + binding.name: network.add_input( + binding.name, dtype_map[binding.dtype], _network_shape(binding) + ) + for binding in abi.inputs + } + optimization = builder.create_optimization_profile() + for binding in abi.inputs: + if binding.min_shape != binding.max_shape: + optimization.set_shape( + binding.name, binding.min_shape, binding.opt_shape, binding.max_shape + ) + config.add_optimization_profile(optimization) + + table = op.weight_constant(network, weights["model.language_model.embed_tokens.weight"]) + table = op.cast(network, table, trt.bfloat16) + hidden = network.add_gather(table, inputs["input_ids"], 0).get_output(0) + + active = _visual_count_active(network, inputs["vision_count"]) + visual = _scatter_visual_rows( + network, + hidden, + inputs["vision_row_indices"], + inputs["vision_embeds"], + active, + ) + zero_mask = op.constant(network, np.zeros((1, 1), dtype=np.float32)) + visual_mask = network.add_elementwise( + inputs["vision_mask"], zero_mask, trt.ElementWiseOperation.GREATER + ).get_output(0) + hidden = network.add_select(visual_mask, visual, hidden).get_output(0) + cos, sin = _mrope_cache(network, inputs["mrope_position_ids"]) + + for index in range(NUM_LAYERS): + prefix = f"model.language_model.layers.{index}" + normalized = op.rms_norm( + network, + hidden, + weights[f"{prefix}.input_layernorm.weight"], + HIDDEN_SIZE, + NORM_EPS, + ) + query = _linear(network, normalized, weights, f"{prefix}.self_attn.q_proj") + key = _linear(network, normalized, weights, f"{prefix}.self_attn.k_proj") + value = _linear(network, normalized, weights, f"{prefix}.self_attn.v_proj") + query = _per_head_norm( + network, query, weights[f"{prefix}.self_attn.q_norm.weight"], NUM_HEADS + ) + key = _per_head_norm( + network, key, weights[f"{prefix}.self_attn.k_norm.weight"], NUM_KV_HEADS + ) + query = op.partial_rope( + network, + query, + cos, + sin, + heads=NUM_HEADS, + head_dim=HEAD_DIM, + rotary_dim=HEAD_DIM, + ) + key = op.partial_rope( + network, + key, + cos, + sin, + heads=NUM_KV_HEADS, + head_dim=HEAD_DIM, + rotary_dim=HEAD_DIM, + ) + query_heads = op.rows_to_heads(network, query, NUM_HEADS, HEAD_DIM) + key_heads = op.rows_to_heads(network, key, NUM_KV_HEADS, HEAD_DIM) + value_heads = op.rows_to_heads(network, value, NUM_KV_HEADS, HEAD_DIM) + key_heads = _repeat_kv(network, key_heads) + value_heads = _repeat_kv(network, value_heads) + scale = op.constant( + network, + np.full((1, 1, 1, 1), 1.0 / math.sqrt(HEAD_DIM), dtype=np.float32), + ) + scale = op.cast(network, scale, query_heads.dtype) + query_heads = network.add_elementwise( + query_heads, scale, trt.ElementWiseOperation.PROD + ).get_output(0) + attention = network.add_attention( + query_heads, + key_heads, + value_heads, + trt.AttentionNormalizationOp.SOFTMAX, + True, + ) + if attention is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen attention layer {index}") + attention.name = f"{prefix}.self_attn.native_attention" + attention.metadata = f"trtmc.native_op=IAttention;source={attention.name}" + attention.get_output(0).name = f"{attention.name}.output" + attention.decomposable = False + update = op.heads_to_rows( + network, + attention.get_output(0), + NUM_HEADS * HEAD_DIM, + ) + update = _linear(network, update, weights, f"{prefix}.self_attn.o_proj") + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + normalized = op.rms_norm( + network, + hidden, + weights[f"{prefix}.post_attention_layernorm.weight"], + HIDDEN_SIZE, + NORM_EPS, + ) + gate = _linear(network, normalized, weights, f"{prefix}.mlp.gate_proj") + up = _linear(network, normalized, weights, f"{prefix}.mlp.up_proj") + gate = op.silu(network, gate) + gated = network.add_elementwise(gate, up, trt.ElementWiseOperation.PROD).get_output(0) + update = _linear(network, gated, weights, f"{prefix}.mlp.down_proj") + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + if index < 3: + deepstack = _scatter_visual_rows( + network, + hidden, + inputs["vision_row_indices"], + inputs[f"deepstack_{index}"], + active, + ) + hidden = network.add_elementwise( + hidden, deepstack, trt.ElementWiseOperation.SUM + ).get_output(0) + + output = op.cast(network, hidden, trt.float32) + output.name = abi.outputs[0].name + network.mark_output(output) + op.validate_native_network( + network, expected_attentions=NUM_LAYERS, label="multimodal text encoder" + ) + print( + "[minimax-h3] building shared Qwen3-VL language stack: " + f"layers={NUM_LAYERS}, sequence={profile.min_sequence_length}.." + f"{profile.max_sequence_length}, compact_vision={profile.min_vision_rows}.." + f"{profile.max_vision_rows}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 multimodal text encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/multimodal_vision_builder.py b/families/minimax_h3/multimodal_vision_builder.py new file mode 100644 index 0000000000..bcd80e7291 --- /dev/null +++ b/families/minimax_h3/multimodal_vision_builder.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native Qwen3-VL vision tower shared by MiniMax-H3 FL2VA and Ref2VA. + +The engine evaluates one image or one temporal video-patch block. Ref2VA video +conditioning invokes it once per temporal block and concatenates outputs, which +is exactly Qwen3-VL's packed ``cu_seqlens`` attention boundary. Interpolation +indices/weights and 2-D rotary positions are runtime inputs because reference +images have dynamic aspect and Ref2VA images use a larger 2048-short-edge grid. +The learned position table and every vision weight remain inside this one plan. +""" + +from __future__ import annotations + +import gc +import math +import sys +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .fl2va_contract import ( + QWEN_VISION_HIDDEN_SIZE, + QWEN_VISION_MERGE_SIZE, + QWEN_VISION_PATCH_WIDTH, + VisionEncoderProfile, + vision_encoder_abi, +) + + +trt = trt_compat.get_trt() + +DEPTH = 27 +NUM_HEADS = 16 +HEAD_DIM = QWEN_VISION_HIDDEN_SIZE // NUM_HEADS +INTERMEDIATE_SIZE = 4304 +POSITION_TABLE_SIDE = 48 +MERGE_UNIT = QWEN_VISION_MERGE_SIZE**2 +MERGED_HIDDEN_SIZE = QWEN_VISION_HIDDEN_SIZE * MERGE_UNIT +DEEPSTACK_VISUAL_INDEXES = (8, 16, 24) +NORM_EPS = 1.0e-6 +VISION_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 +ROPE_INV_FREQ_BITS = ( + 0x3F800000, + 0x3F1977CC, + 0x3EB800D6, + 0x3E5C9D35, + 0x3E044133, + 0x3D9E91B6, + 0x3D3E1E95, + 0x3CE3F280, + 0x3C88A69B, + 0x3C23D70A, + 0x3BC47060, + 0x3B6B8631, + 0x3B0D3169, + 0x3AA94938, + 0x3A4AF7F3, + 0x39F35A5C, + 0x3991E2E1, + 0x392EE9BF, +) + + +def checkpoint_keys() -> tuple[str, ...]: + names = ["model.visual.patch_embed.proj.weight", "model.visual.patch_embed.proj.bias"] + names.append("model.visual.pos_embed.weight") + for index in range(DEPTH): + prefix = f"model.visual.blocks.{index}" + names.extend( + [ + f"{prefix}.norm1.weight", + f"{prefix}.norm1.bias", + f"{prefix}.attn.qkv.weight", + f"{prefix}.attn.qkv.bias", + f"{prefix}.attn.proj.weight", + f"{prefix}.attn.proj.bias", + f"{prefix}.norm2.weight", + f"{prefix}.norm2.bias", + f"{prefix}.mlp.linear_fc1.weight", + f"{prefix}.mlp.linear_fc1.bias", + f"{prefix}.mlp.linear_fc2.weight", + f"{prefix}.mlp.linear_fc2.bias", + ] + ) + for prefix in [ + "model.visual.merger", + *(f"model.visual.deepstack_merger_list.{index}" for index in range(3)), + ]: + names.extend( + [ + f"{prefix}.norm.weight", + f"{prefix}.norm.bias", + f"{prefix}.linear_fc1.weight", + f"{prefix}.linear_fc1.bias", + f"{prefix}.linear_fc2.weight", + f"{prefix}.linear_fc2.bias", + ] + ) + return tuple(names) + + +def _layer_norm(network, hidden, weights, prefix: str, width: int): + rank = len(tuple(hidden.shape)) + shape = (1,) * (rank - 1) + (width,) + gamma = op.weight_constant(network, np.asarray(weights[f"{prefix}.weight"]).reshape(shape)) + beta = op.weight_constant(network, np.asarray(weights[f"{prefix}.bias"]).reshape(shape)) + gamma = op.cast(network, gamma, hidden.dtype) + beta = op.cast(network, beta, hidden.dtype) + layer = network.add_normalization_v2(hidden, gamma, beta, 1 << (rank - 1)) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen vision LayerNorm {prefix}") + layer.name = prefix + layer.epsilon = NORM_EPS + return layer.get_output(0) + + +def _gelu(network, hidden, activation_type): + layer = network.add_activation(hidden, activation_type) + if layer is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen vision GELU") + return layer.get_output(0) + + +def _interpolated_position_embeddings(network, weights, indices, interp_weights): + table = op.weight_constant(network, weights["model.visual.pos_embed.weight"]) + table = op.cast(network, table, trt.bfloat16) + gathered = network.add_gather(table, indices, 0) + if gathered is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position-table gather") + # The learned table is BF16, but Transformers multiplies it by FP32 + # bilinear weights and performs the four-way sum in FP32 before rounding + # the completed positional embedding back to the patch-embedding dtype. + gathered_value = op.cast(network, gathered.get_output(0), trt.float32) + coefficients = network.add_shuffle(interp_weights) + coefficients.reshape_dims = (-1, 4, 1) + coefficients_value = op.cast(network, coefficients.get_output(0), trt.float32) + weighted = network.add_elementwise( + gathered_value, coefficients_value, trt.ElementWiseOperation.PROD + ) + if weighted is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position interpolation") + reduced = network.add_reduce(weighted.get_output(0), trt.ReduceOperation.SUM, 1 << 1, False) + if reduced is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 Qwen position interpolation reduce") + return reduced.get_output(0) + + +def _vision_rope_cache(network, position_ids): + positions = op.cast(network, position_ids, trt.float32) + # Preserve Qwen3VLVisionRotaryEmbedding's Torch-created FP32 buffer + # without making Torch a builder dependency. + inverse = np.asarray(ROPE_INV_FREQ_BITS, dtype=np.uint32).view(np.float32) + inverse = op.constant(network, inverse.reshape(1, -1)) + axes = [] + for axis in range(2): + coordinate = op.dynamic_slice(network, positions, (0, axis), (None, 1)) + axes.append( + network.add_elementwise(coordinate, inverse, trt.ElementWiseOperation.PROD).get_output( + 0 + ) + ) + frequency = network.add_concatenation(axes) + frequency.axis = 1 + cos = network.add_unary(frequency.get_output(0), trt.UnaryOperation.COS).get_output(0) + sin = network.add_unary(frequency.get_output(0), trt.UnaryOperation.SIN).get_output(0) + # Qwen3-VL explicitly performs vision rotary embedding in FP32 and rounds + # q/k back to their source dtype only after both products are added. + return cos, sin + + +def _vision_partial_rope(network, tensor, cos_half, sin_half): + """Qwen3-VL's FP32 vision rotate-half operation.""" + + value = op.rows_to_heads(network, tensor, NUM_HEADS, HEAD_DIM) + source_dtype = value.dtype + value = op.cast(network, value, trt.float32) + first = op.dynamic_slice(network, value, (0, 0, 0, 0), (1, NUM_HEADS, None, HEAD_DIM // 2)) + second = op.dynamic_slice( + network, + value, + (0, 0, 0, HEAD_DIM // 2), + (1, NUM_HEADS, None, HEAD_DIM // 2), + ) + negative_second = network.add_unary(second, trt.UnaryOperation.NEG).get_output(0) + rotated_half = network.add_concatenation((negative_second, first)) + rotated_half.axis = 3 + + def duplicate(table): + reshape = network.add_shuffle(op.cast(network, table, trt.float32)) + reshape.reshape_dims = (1, 1, -1, HEAD_DIM // 2) + result = network.add_concatenation((reshape.get_output(0), reshape.get_output(0))) + result.axis = 3 + return result.get_output(0) + + left = network.add_elementwise( + value, duplicate(cos_half), trt.ElementWiseOperation.PROD + ).get_output(0) + right = network.add_elementwise( + rotated_half.get_output(0), duplicate(sin_half), trt.ElementWiseOperation.PROD + ).get_output(0) + result = network.add_elementwise(left, right, trt.ElementWiseOperation.SUM).get_output(0) + result = op.cast(network, result, source_dtype) + return op.heads_to_rows(network, result, QWEN_VISION_HIDDEN_SIZE) + + +def _attention(network, hidden, weights, prefix: str, cos, sin): + projected = op.linear( + network, + hidden, + weights[f"{prefix}.qkv.weight"], + weights[f"{prefix}.qkv.bias"], + ) + query, key, value = tuple( + op.dynamic_slice( + network, + projected, + (0, index * QWEN_VISION_HIDDEN_SIZE), + (None, QWEN_VISION_HIDDEN_SIZE), + ) + for index in range(3) + ) + query = _vision_partial_rope(network, query, cos, sin) + key = _vision_partial_rope(network, key, cos, sin) + query_heads = op.rows_to_heads(network, query, NUM_HEADS, HEAD_DIM) + key_heads = op.rows_to_heads(network, key, NUM_HEADS, HEAD_DIM) + value_heads = op.rows_to_heads(network, value, NUM_HEADS, HEAD_DIM) + scale = op.cast( + network, + op.constant(network, np.full((1, 1, 1, 1), 1.0 / math.sqrt(HEAD_DIM), np.float32)), + query_heads.dtype, + ) + query_heads = network.add_elementwise( + query_heads, scale, trt.ElementWiseOperation.PROD + ).get_output(0) + layer = network.add_attention( + query_heads, + key_heads, + value_heads, + trt.AttentionNormalizationOp.SOFTMAX, + False, + ) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 Qwen vision attention {prefix}") + layer.name = f"{prefix}.native_attention" + layer.metadata = f"trtmc.native_op=IAttention;source={layer.name}" + layer.get_output(0).name = f"{layer.name}.output" + layer.decomposable = False + rows = op.heads_to_rows(network, layer.get_output(0), QWEN_VISION_HIDDEN_SIZE) + return op.linear( + network, + rows, + weights[f"{prefix}.proj.weight"], + weights[f"{prefix}.proj.bias"], + ) + + +def _mlp(network, hidden, weights, prefix: str): + hidden = op.linear( + network, + hidden, + weights[f"{prefix}.linear_fc1.weight"], + weights[f"{prefix}.linear_fc1.bias"], + ) + hidden = _gelu(network, hidden, trt.ActivationType.GELU_TANH) + return op.linear( + network, + hidden, + weights[f"{prefix}.linear_fc2.weight"], + weights[f"{prefix}.linear_fc2.bias"], + ) + + +def _merge(network, hidden, weights, prefix: str, *, postshuffle_norm: bool): + if postshuffle_norm: + reshape = network.add_shuffle(hidden) + reshape.reshape_dims = (-1, MERGED_HIDDEN_SIZE) + merged = _layer_norm( + network, reshape.get_output(0), weights, f"{prefix}.norm", MERGED_HIDDEN_SIZE + ) + else: + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm", QWEN_VISION_HIDDEN_SIZE + ) + reshape = network.add_shuffle(normalized) + reshape.reshape_dims = (-1, MERGED_HIDDEN_SIZE) + merged = reshape.get_output(0) + merged = op.linear( + network, + merged, + weights[f"{prefix}.linear_fc1.weight"], + weights[f"{prefix}.linear_fc1.bias"], + ) + # Qwen3-VL patch mergers use nn.GELU(approximate="none"), unlike the + # tanh-approximate GELU in each vision MLP. + merged = _gelu(network, merged, trt.ActivationType.GELU_ERF) + return op.linear( + network, + merged, + weights[f"{prefix}.linear_fc2.weight"], + weights[f"{prefix}.linear_fc2.bias"], + ) + + +@op.cleanup_failed_build +def build_multimodal_vision_encoder_engine( + weights: dict[str, np.ndarray], + profile: VisionEncoderProfile = VisionEncoderProfile(), + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build one shared dynamic-aspect Qwen3-VL vision tower.""" + + profile.validate() + expected_keys = set(checkpoint_keys()) + missing = sorted(expected_keys - set(weights)) + unexpected = sorted(set(weights) - expected_keys) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 multimodal vision checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=VISION_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + + abi = vision_encoder_abi(profile) + inputs = { + binding.name: network.add_input( + binding.name, + { + "float32": trt.float32, + "int32": trt.int32, + }[binding.dtype], + (-1, binding.max_shape[1]), + ) + for binding in abi.inputs + } + optimization = builder.create_optimization_profile() + for binding in abi.inputs: + optimization.set_shape( + binding.name, binding.min_shape, binding.opt_shape, binding.max_shape + ) + config.add_optimization_profile(optimization) + + patch_weight = np.asarray(weights["model.visual.patch_embed.proj.weight"]).reshape( + QWEN_VISION_HIDDEN_SIZE, QWEN_VISION_PATCH_WIDTH + ) + hidden = op.linear( + network, + inputs["pixel_values"], + patch_weight, + weights["model.visual.patch_embed.proj.bias"], + ) + position = _interpolated_position_embeddings( + network, weights, inputs["interp_indices"], inputs["interp_weights"] + ) + position = op.cast(network, position, hidden.dtype) + hidden = network.add_elementwise(hidden, position, trt.ElementWiseOperation.SUM).get_output(0) + cos, sin = _vision_rope_cache(network, inputs["vision_position_ids"]) + + deepstack = [] + for index in range(DEPTH): + prefix = f"model.visual.blocks.{index}" + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm1", QWEN_VISION_HIDDEN_SIZE + ) + update = _attention(network, normalized, weights, f"{prefix}.attn", cos, sin) + update = op.cast(network, update, hidden.dtype) + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + normalized = _layer_norm( + network, hidden, weights, f"{prefix}.norm2", QWEN_VISION_HIDDEN_SIZE + ) + update = _mlp(network, normalized, weights, f"{prefix}.mlp") + update = op.cast(network, update, hidden.dtype) + hidden = network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + if index in DEEPSTACK_VISUAL_INDEXES: + merger_index = DEEPSTACK_VISUAL_INDEXES.index(index) + deepstack.append( + _merge( + network, + hidden, + weights, + f"model.visual.deepstack_merger_list.{merger_index}", + postshuffle_norm=True, + ) + ) + + main = _merge( + network, + hidden, + weights, + "model.visual.merger", + postshuffle_norm=False, + ) + for tensor, binding in zip((main, *deepstack), abi.outputs): + output = op.cast(network, tensor, trt.float32) + output.name = binding.name + network.mark_output(output) + + op.validate_native_network( + network, expected_attentions=DEPTH, label="multimodal vision encoder" + ) + print( + "[minimax-h3] building shared Qwen3-VL vision tower: " + f"layers={DEPTH}, patches={profile.min_patches}..{profile.max_patches}, " + f"merged={profile.min_patches // MERGE_UNIT}..{profile.max_patches // MERGE_UNIT}", + file=sys.stderr, + ) + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 multimodal vision encoder") + del network, config, builder + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/nvfp4_text_checkpoint.py b/families/minimax_h3/nvfp4_text_checkpoint.py new file mode 100644 index 0000000000..20767bb438 --- /dev/null +++ b/families/minimax_h3/nvfp4_text_checkpoint.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build-time import of Comfy's public Qwen3-VL NVFP4 AWQ checkpoint. + +The released checkpoint explicitly requests full-precision matrix products. +Decode its packed weights to BF16 for the native TensorRT language engine; +this is checkpoint compatibility, not FP4 activation or matrix computation. +The same file contains the unquantized BF16 vision tower. AWQ input scales +remain separate so the language graph preserves their BF16 rounding point. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +import math +import os +from pathlib import Path +from typing import Iterable + +import ml_dtypes +import numpy as np + +from .quantized_checkpoint import ( + CHECKPOINT_REVISION, + MODEL_ID, + QuantizedSourceFileIdentity, + _read_header, + _read_marker, + _source_file_identity, +) + +CHECKPOINT_FILENAME = "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors" +CHECKPOINT_BYTES = 15_687_142_551 +_NVFP4_MARKER = {"format": "nvfp4", "full_precision_matrix_mult": True} +_EMBEDDING = "model.embed_tokens" +_DTYPES = { + "BF16": np.dtype(ml_dtypes.bfloat16), + "F32": np.dtype(" dict[str, object]: + """Public source and compute identity, without local paths.""" + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "filename": self.filename, + "size_bytes": self.size_bytes, + "tensor_count": self.tensor_count, + "quantized_weight_count": self.quantized_weight_count, + "quantization": "nvfp4_awq", + "embedding_quantization": "int8_tensorwise", + "full_precision_matrix_mult": True, + "engine_weights_dtype": "bfloat16", + "matmul_dtype": "bfloat16", + "runtime_framework": None, + } + + +QUANTIZED_TEXT_CHECKPOINT_IDENTITY = QuantizedTextCheckpointIdentity() + + +def _physical_specs() -> dict[str, tuple[str, tuple[int, ...]]]: + specs = {} + + def tensor(name, shape, dtype="BF16"): + specs[name] = (dtype, tuple(shape)) + + tensor(f"{_EMBEDDING}.weight", (151936, 5120), "I8") + tensor(f"{_EMBEDDING}.weight_scale", (151936, 1), "F32") + tensor(f"{_EMBEDDING}.comfy_quant", (29,), "U8") + for index in range(50): + prefix = f"model.layers.{index}" + for name, width in (("input_layernorm", 5120), ("post_attention_layernorm", 5120), + ("self_attn.q_norm", 128), ("self_attn.k_norm", 128)): + tensor(f"{prefix}.{name}.weight", (width,)) + for name, rows, cols in ( + ("self_attn.q_proj", 8192, 5120), ("self_attn.k_proj", 1024, 5120), + ("self_attn.v_proj", 1024, 5120), ("self_attn.o_proj", 5120, 8192), + ("mlp.gate_proj", 25600, 5120), ("mlp.up_proj", 25600, 5120), + ("mlp.down_proj", 5120, 25600), + ): + base = f"{prefix}.{name}" + tensor(f"{base}.weight", (rows, cols // 2), "U8") + tensor(f"{base}.weight_scale", (rows, cols // 16), "F8_E4M3") + tensor(f"{base}.weight_scale_2", (), "F32") + tensor(f"{base}.comfy_quant", (55,), "U8") + if name in ("self_attn.o_proj", "mlp.down_proj"): + tensor(f"{base}.pre_quant_scale", (cols,)) + + def linear(base, rows, cols): + tensor(f"{base}.weight", (rows, cols)) + tensor(f"{base}.bias", (rows,)) + + def norm(base, width): + tensor(f"{base}.weight", (width,)) + tensor(f"{base}.bias", (width,)) + + tensor("visual.patch_embed.proj.weight", (1152, 3, 2, 16, 16)) + tensor("visual.patch_embed.proj.bias", (1152,)) + tensor("visual.pos_embed.weight", (2304, 1152)) + for index in range(27): + prefix = f"visual.blocks.{index}" + norm(f"{prefix}.norm1", 1152) + norm(f"{prefix}.norm2", 1152) + linear(f"{prefix}.attn.qkv", 3456, 1152) + linear(f"{prefix}.attn.proj", 1152, 1152) + linear(f"{prefix}.mlp.linear_fc1", 4304, 1152) + linear(f"{prefix}.mlp.linear_fc2", 1152, 4304) + for prefix in ("visual.merger", *(f"visual.deepstack_merger_list.{i}" for i in range(3))): + norm(f"{prefix}.norm", 1152 if prefix == "visual.merger" else 4608) + linear(f"{prefix}.linear_fc1", 4608, 4608) + linear(f"{prefix}.linear_fc2", 5120, 4608) + return specs + + +_PHYSICAL_SPECS = _physical_specs() + + +def _authenticate_source(path: Path) -> None: + identity = QUANTIZED_TEXT_CHECKPOINT_IDENTITY + relative = Path(identity.filename) + root = path.parents[len(relative.parts) - 1] + if os.path.normcase(os.path.abspath(path)) != os.path.normcase(os.path.abspath(root / relative)): + raise ValueError("MiniMax-H3 text checkpoint must retain its released text_encoders path") + if root.parent.name == "snapshots": + if root.name != identity.revision or root.parent.parent.name != "models--Comfy-Org--MiniMax-H3": + raise ValueError("MiniMax-H3 text checkpoint is not the pinned Hugging Face snapshot") + return + metadata_path = root / ".cache/huggingface/download" / f"{relative.as_posix()}.metadata" + try: + lines = metadata_path.read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError) as error: + raise ValueError("Download the MiniMax-H3 text checkpoint with hf download and pinned revision metadata") from error + if not lines or lines[0] != identity.revision: + raise ValueError("MiniMax-H3 text checkpoint does not match the pinned Hugging Face revision") + + +def _validate_inventory(header, data_offset: int, file_size: int): + tensors = {key: value for key, value in header.items() if key != "__metadata__"} + if set(tensors) != set(_PHYSICAL_SPECS): + raise ValueError("MiniMax-H3 NVFP4 text checkpoint tensor inventory mismatch") + regions = [] + for name, (dtype, shape) in _PHYSICAL_SPECS.items(): + entry = tensors[name] + if not isinstance(entry, dict) or set(entry) != {"dtype", "shape", "data_offsets"}: + raise ValueError(f"Invalid MiniMax-H3 text tensor header: {name}") + if entry["dtype"] != dtype or entry["shape"] != list(shape): + raise ValueError(f"MiniMax-H3 text tensor contract mismatch: {name}") + offsets = entry["data_offsets"] + if (not isinstance(offsets, list) or len(offsets) != 2 + or any(type(value) is not int for value in offsets)): + raise ValueError(f"Invalid MiniMax-H3 text tensor offsets: {name}") + start, end = offsets + if start < 0 or end - start != math.prod(shape) * _DTYPES[dtype].itemsize: + raise ValueError(f"Invalid MiniMax-H3 text tensor byte range: {name}") + regions.append((start, end)) + cursor = 0 + for start, end in sorted(regions): + if start != cursor: + raise ValueError("MiniMax-H3 text tensor storage has gaps or overlaps") + cursor = end + if data_offset + cursor != file_size: + raise ValueError("MiniMax-H3 text checkpoint payload size mismatch") + return tensors + + +def validate_quantized_text_checkpoint(checkpoint_file: str | Path) -> QuantizedTextCheckpointIdentity: + """Check the pinned file, header and all tiny markers; never hash its payload.""" + path = Path(checkpoint_file) + identity = QUANTIZED_TEXT_CHECKPOINT_IDENTITY + before = _source_file_identity(path) + if path.name != Path(identity.filename).name or before.size_bytes != identity.size_bytes: + raise ValueError("MiniMax-H3 NVFP4 text checkpoint filename or size mismatch") + _authenticate_source(path) + header, offset = _read_header(path) + tensors = _validate_inventory(header, offset, before.size_bytes) + for name, entry in tensors.items(): + if name.endswith(".comfy_quant"): + expected = {"format": "int8_tensorwise"} if name == f"{_EMBEDDING}.comfy_quant" else _NVFP4_MARKER + if _read_marker(path, offset, entry) != expected: + raise ValueError(f"Unsupported MiniMax-H3 text comfy_quant marker: {name}") + if _source_file_identity(path) != before: + raise ValueError("MiniMax-H3 text checkpoint changed during validation") + return replace(identity, source_file_identity=before) + + +def _unblock_scales(scales: np.ndarray, rows: int, cols: int) -> np.ndarray: + """Invert Comfy's cuBLAS 128x4 block-scale storage on the build CPU. + + Format reference: Comfy-Org/comfy-kitchen, float_utils.py::from_blocked. + """ + row_blocks, col_blocks = (rows + 127) // 128, (cols + 3) // 4 + if scales.shape != (row_blocks * 128, col_blocks * 4): + raise ValueError("MiniMax-H3 NVFP4 block-scale shape mismatch") + tiles = scales.reshape(row_blocks, col_blocks, 32, 4, 4) + plain = tiles.transpose(0, 1, 3, 2, 4).reshape(row_blocks, col_blocks, 128, 4) + plain = plain.transpose(0, 2, 1, 3).reshape(row_blocks * 128, col_blocks * 4) + return plain[:rows, :cols] + + +def _dequantize_nvfp4(packed, scales, global_scale) -> np.ndarray: + """Match Comfy CUDA decode: FP32 scale/product, then one BF16 rounding. + + Comfy stores the even element in the HIGH nibble. Its CUDA reference is + backends/cuda/ops/quantize_nvfp4.cu::dequantize_nvfp4_kernel; the eager CPU + backend has extra intermediate BF16 rounding and is not used as an oracle. + """ + if packed.dtype != np.uint8 or packed.ndim != 2 or packed.shape[1] % 8: + raise ValueError("MiniMax-H3 NVFP4 weights require packed uint8 [out,in/2], in divisible by 16") + rows, cols = packed.shape[0], packed.shape[1] * 2 + if np.asarray(global_scale).shape != () or not np.isfinite(global_scale) or global_scale <= 0: + raise ValueError("MiniMax-H3 NVFP4 global scale must be a positive finite scalar") + plain_scales = _unblock_scales(scales, rows, cols // 16) + lut = np.asarray((0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6), np.float32) + output = np.empty((rows, cols), dtype=ml_dtypes.bfloat16) + for start in range(0, rows, 128): + end = min(rows, start + 128) + scale = plain_scales[start:end].astype(np.float32) * np.float32(global_scale) + if not np.isfinite(scale).all() or (scale < 0).any(): + raise ValueError("MiniMax-H3 NVFP4 block scales must be finite and non-negative") + values = np.empty((end - start, cols), np.float32) + values[:, 0::2] = lut[packed[start:end] >> 4] + values[:, 1::2] = lut[packed[start:end] & 15] + values.reshape(end - start, cols // 16, 16)[:] *= scale[..., None] + output[start:end] = values + return output + + +def _physical_name(logical_name: str) -> str: + for logical, physical in (("model.language_model.", "model."), ("model.visual.", "visual.")): + if logical_name.startswith(logical): + name = physical + logical_name.removeprefix(logical) + if name in _PHYSICAL_SPECS and name.endswith((".weight", ".bias")): + return name + raise ValueError(f"Unsupported MiniMax-H3 quantized text logical tensor: {logical_name}") + + +def load_selected_quantized_text_weights( + checkpoint_file: str | Path, logical_names: Iterable[str] +) -> dict[str, np.ndarray]: + """Load requested language/vision tensors and their optional AWQ input scales. + + Only requested payloads are touched. Memory-mapped BF16 vision/norm tensors + retain their file owner; decoded weights use bounded per-row temporaries. + """ + path = Path(checkpoint_file) + identity = validate_quantized_text_checkpoint(path) + requested = tuple(logical_names) + if len(requested) != len(set(requested)): + raise ValueError("MiniMax-H3 text weight request contains duplicates") + names = {logical: _physical_name(logical) for logical in requested} + header, offset = _read_header(path) + + def read(name): + entry = header[name] + return np.memmap(path, mode="r", dtype=_DTYPES[entry["dtype"]], + offset=offset + entry["data_offsets"][0], shape=tuple(entry["shape"])) + + result = {} + for logical, physical in names.items(): + value = read(physical) + base = physical.removesuffix(".weight") + if f"{base}.weight_scale_2" in header: + value = _dequantize_nvfp4(value, read(f"{base}.weight_scale"), read(f"{base}.weight_scale_2")) + prescale = f"{base}.pre_quant_scale" + if prescale in header: + result[f"{logical.removesuffix('.weight')}.pre_quant_scale"] = read(prescale) + elif physical == f"{_EMBEDDING}.weight": + scale = read(f"{base}.weight_scale") + decoded = np.empty(value.shape, dtype=ml_dtypes.bfloat16) + for start in range(0, value.shape[0], 128): + chunk_scale = scale[start:start + 128] + if not np.isfinite(chunk_scale).all() or (chunk_scale < 0).any(): + raise ValueError("MiniMax-H3 embedding scales must be finite and non-negative") + decoded[start:start + 128] = value[start:start + 128].astype(np.float32) * chunk_scale + value = decoded + result[logical] = value + if _source_file_identity(path) != identity.source_file_identity: + raise ValueError("MiniMax-H3 text checkpoint changed during loading") + return result diff --git a/families/minimax_h3/provenance.py b/families/minimax_h3/provenance.py new file mode 100644 index 0000000000..90b23a9de0 --- /dev/null +++ b/families/minimax_h3/provenance.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Path-free public metadata for MiniMax-H3 native bundles.""" + +from __future__ import annotations + +import json +import math +import struct +from pathlib import Path + +from tensorrt_model_connect.bundle_writer import BUNDLE_MAGIC + +from .config import ( + NATIVE_EXPLICIT_CANVAS_SIZES, + SOL_ENGINE_1344X768_124_TO_345F, + native_plan_filenames, +) + + +CHECKPOINT_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc" +CHECKPOINT_REPOSITORY = "MiniMaxAI/MiniMax-H3" + +QUANTIZED_TRANSFORMER_CONFIG = { + "format": "int8_tensorwise_convrot", + "checkpoint_scope": "full_transformer_overlay", + "quantized_weight_scope": ( + "blocks.*.{adaln_proj.linear,attn.qkv_proj,attn.out_proj,mlp.fc1,mlp.fc2}" + ), + "quantized_weight_count": 250, + "activation_precision": "dynamic_int8_rowwise", +} + +SUPER_RESOLUTION_PLAN_FILENAME = "video_super_resolution.plan" +SUPER_RESOLUTION_SECTION = "video_super_resolution_plan" +SUPER_RESOLUTION_MODEL = "realesr-general-x4v3" +SUPER_RESOLUTION_ARCHITECTURE = "SRVGGNetCompact" +SUPER_RESOLUTION_PRIMARY_FILENAME = "realesr-general-x4v3.pth" +SUPER_RESOLUTION_WEAK_FILENAME = "realesr-general-wdn-x4v3.pth" +SUPER_RESOLUTION_SOURCE_SHAPE = [480, 864] +SUPER_RESOLUTION_TARGET_SHAPE = [720, 1296] +SUPER_RESOLUTION_BATCH_PROFILE = [1, 4, 8] +SUPER_RESOLUTION_LEARNED_RESIDUAL_STRENGTH = 0.25 + + +def checkpoint_snapshot_record( + snapshot: Path, + *, + include_transformer_weights: bool = True, +) -> dict[str, object]: + """Describe the selected public checkpoint without recording local paths.""" + + root = Path(snapshot) + required = ["text_encoder", "vae", "audio_vae", "tokenizer"] + if include_transformer_weights: + required.append("transformer") + missing = [name for name in required if not (root / name).exists()] + if missing: + raise FileNotFoundError( + "Incomplete MiniMax-H3 checkpoint; missing: " + ", ".join(missing) + ) + return { + "repository": CHECKPOINT_REPOSITORY, + "revision": CHECKPOINT_REVISION, + "includes_transformer_weights": include_transformer_weights, + } + + +def validate_checkpoint_snapshot_record( + record: object, + *, + include_transformer_weights: bool = True, +) -> dict[str, object]: + expected = { + "repository": CHECKPOINT_REPOSITORY, + "revision": CHECKPOINT_REVISION, + "includes_transformer_weights": include_transformer_weights, + } + if record != expected: + raise ValueError("MiniMax-H3 checkpoint metadata does not match the public model") + return dict(expected) + + +def _checkpoint_record(path: Path, role: str, filename: str) -> dict[str, object]: + value = Path(path) + if not value.is_file(): + raise FileNotFoundError(f"MiniMax-H3 {role} checkpoint is missing: {value}") + if value.name != filename: + raise ValueError(f"MiniMax-H3 {role} checkpoint must be named {filename}") + return {"role": role, "filename": filename, "bytes": value.stat().st_size} + + +def super_resolution_source_identity( + primary_checkpoint: Path, + weak_checkpoint: Path | None, + *, + denoise_strength: float, + learned_residual_strength: float = SUPER_RESOLUTION_LEARNED_RESIDUAL_STRENGTH, +) -> dict[str, object]: + if isinstance(denoise_strength, bool) or not isinstance(denoise_strength, (int, float)): + raise ValueError("MiniMax-H3 super-resolution denoise_strength must be numeric") + if isinstance(learned_residual_strength, bool) or not isinstance( + learned_residual_strength, (int, float) + ): + raise ValueError("MiniMax-H3 learned residual strength must be numeric") + strength = float(denoise_strength) + residual = float(learned_residual_strength) + if not math.isfinite(strength) or not 0.0 <= strength <= 1.0: + raise ValueError("MiniMax-H3 super-resolution denoise_strength must be in [0, 1]") + if not math.isfinite(residual) or not 0.0 <= residual <= 1.0: + raise ValueError("MiniMax-H3 learned residual strength must be in [0, 1]") + if weak_checkpoint is None and strength != 1.0: + raise ValueError("MiniMax-H3 denoise blending requires the weak checkpoint") + sources = [_checkpoint_record(primary_checkpoint, "primary", SUPER_RESOLUTION_PRIMARY_FILENAME)] + if weak_checkpoint is not None: + sources.append(_checkpoint_record(weak_checkpoint, "weak_denoise", SUPER_RESOLUTION_WEAK_FILENAME)) + return { + "model": SUPER_RESOLUTION_MODEL, + "architecture": SUPER_RESOLUTION_ARCHITECTURE, + "denoise_strength": strength, + "learned_residual_strength": residual, + "sources": sources, + } + + +def validate_super_resolution_source_identity(record: object) -> dict[str, object]: + if not isinstance(record, dict): + raise ValueError("MiniMax-H3 super-resolution source metadata must be an object") + if record.get("model") != SUPER_RESOLUTION_MODEL: + raise ValueError("MiniMax-H3 super-resolution model is unsupported") + if record.get("architecture") != SUPER_RESOLUTION_ARCHITECTURE: + raise ValueError("MiniMax-H3 super-resolution architecture is unsupported") + sources = record.get("sources") + if not isinstance(sources, list) or len(sources) not in (1, 2): + raise ValueError("MiniMax-H3 super-resolution sources are invalid") + expected = ( + ("primary", SUPER_RESOLUTION_PRIMARY_FILENAME), + ("weak_denoise", SUPER_RESOLUTION_WEAK_FILENAME), + ) + for source, (role, filename) in zip(sources, expected, strict=False): + if ( + not isinstance(source, dict) + or source.get("role") != role + or source.get("filename") != filename + or not isinstance(source.get("bytes"), int) + or source["bytes"] <= 0 + ): + raise ValueError("MiniMax-H3 super-resolution source metadata is invalid") + strength = record.get("denoise_strength") + residual = record.get("learned_residual_strength") + if not isinstance(strength, (int, float)) or isinstance(strength, bool): + raise ValueError("MiniMax-H3 super-resolution denoise strength is invalid") + if not isinstance(residual, (int, float)) or isinstance(residual, bool): + raise ValueError("MiniMax-H3 super-resolution residual strength is invalid") + if not 0.0 <= float(strength) <= 1.0 or not 0.0 <= float(residual) <= 1.0: + raise ValueError("MiniMax-H3 super-resolution strengths must be in [0, 1]") + if len(sources) == 1 and float(strength) != 1.0: + raise ValueError("MiniMax-H3 denoise blending requires the weak checkpoint") + return dict(record) + + +def super_resolution_bundle_config(source_identity: object) -> dict[str, object]: + source = validate_super_resolution_source_identity(source_identity) + return { + "section": SUPER_RESOLUTION_SECTION, + "mode": "explicit", + "input_name": "frames", + "output_name": "upscaled_frames", + "model": source["model"], + "architecture": source["architecture"], + "source_shape": list(SUPER_RESOLUTION_SOURCE_SHAPE), + "target_shape": list(SUPER_RESOLUTION_TARGET_SHAPE), + "batch_profile": list(SUPER_RESOLUTION_BATCH_PROFILE), + "denoise_strength": source["denoise_strength"], + "learned_residual_strength": source["learned_residual_strength"], + "sources": source["sources"], + } + + +def validate_super_resolution_bundle_config(record: object) -> dict[str, object]: + if not isinstance(record, dict) or record.get("section") != SUPER_RESOLUTION_SECTION: + raise ValueError("MiniMax-H3 super-resolution bundle metadata is invalid") + if record.get("mode") != "explicit": + raise ValueError("MiniMax-H3 legacy or unsupported SR mode; rebuild with super_resolution=true") + if record.get("input_name") != "frames" or record.get("output_name") != "upscaled_frames": + raise ValueError("MiniMax-H3 super-resolution tensor names are invalid") + if record.get("source_shape") != SUPER_RESOLUTION_SOURCE_SHAPE: + raise ValueError("MiniMax-H3 super-resolution source shape is invalid") + if record.get("target_shape") != SUPER_RESOLUTION_TARGET_SHAPE: + raise ValueError("MiniMax-H3 super-resolution target shape is invalid") + if record.get("batch_profile") != SUPER_RESOLUTION_BATCH_PROFILE: + raise ValueError("MiniMax-H3 super-resolution batch profile is invalid") + validate_super_resolution_source_identity( + { + "model": record.get("model"), + "architecture": record.get("architecture"), + "denoise_strength": record.get("denoise_strength"), + "learned_residual_strength": record.get("learned_residual_strength"), + "sources": record.get("sources"), + } + ) + return dict(record) + + +def validate_quantized_transformer_metadata(record: object) -> dict[str, object]: + from .quantized_checkpoint import QUANTIZED_CHECKPOINT_IDENTITY + + expected = QUANTIZED_CHECKPOINT_IDENTITY.bundle_metadata() + if record != expected: + raise ValueError("MiniMax-H3 quantized transformer metadata does not match the public model") + return expected + + +def validate_workspace_limit_bytes( + record: object, + *, + profile=SOL_ENGINE_1344X768_124_TO_345F, + additional_plan_filenames: tuple[str, ...] = (), + excluded_plan_filenames: tuple[str, ...] = (), +) -> dict[str, object]: + expected = tuple( + filename + for filename in native_plan_filenames() + if filename not in excluded_plan_filenames + ) + tuple(additional_plan_filenames) + if not isinstance(record, dict) or set(record) != set(expected): + raise ValueError("MiniMax-H3 workspace metadata must cover every native plan") + for filename, value in record.items(): + if value not in ("tensorrt_default", None) and ( + not isinstance(value, int) or isinstance(value, bool) or value <= 0 + ): + raise ValueError(f"MiniMax-H3 workspace value for {filename} is invalid") + del profile + return dict(record) + + +def load_bundle_config(bundle: Path) -> dict[str, object]: + """Load the canonical runtime section from a current ModelConnect bundle.""" + + with Path(bundle).open("rb") as stream: + if stream.read(8) != BUNDLE_MAGIC: + raise ValueError("MiniMax-H3 bundle has invalid magic") + raw_size = stream.read(8) + if len(raw_size) != 8: + raise ValueError("MiniMax-H3 bundle header is truncated") + header_size = struct.unpack(" dict: + config = load_bundle_config(bundle) + if config.get("checkpoint_revision") != CHECKPOINT_REVISION: + raise ValueError("MiniMax-H3 bundle checkpoint revision is invalid") + if config.get("explicit_canvas_sizes") != [list(size) for size in NATIVE_EXPLICIT_CANVAS_SIZES]: + raise ValueError("MiniMax-H3 bundle canvas profile is invalid") + del source_revision + return config diff --git a/families/minimax_h3/quantized_checkpoint.py b/families/minimax_h3/quantized_checkpoint.py new file mode 100644 index 0000000000..7c99d72597 --- /dev/null +++ b/families/minimax_h3/quantized_checkpoint.py @@ -0,0 +1,738 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict loader for the public full MiniMax-H3 INT8 ConvRot transformers. + +The Comfy-Org single-file checkpoint uses the original MiniMax module names, +contiguous Q/K/V thirds, and ``[gate; value]`` SwiGLU rows. The native +TensorRT builders consume the frozen Diffusers names, the same contiguous Q/K/V +layout, and ``[value; gate]`` rows. This module performs the lossless SwiGLU +row permutation while preserving the checkpoint's INT8 data and per-output-row +FP32 scales. + +T2VA/FL2VA and Ref2VA share this tensor layout but use distinct checkpoint +values. The explicit workflow selects the pinned source; it never substitutes +the FL2VA checkpoint for Ref2VA or accepts the pruned AdaLN variant. + +Large tensor payloads are selected through :mod:`safetensors`; validation reads +only the JSON header and the 250 tiny ``comfy_quant`` markers. The released +file size is checked without rereading the 34 GB payload before every build. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +import json +import math +import os +from pathlib import Path +import stat +import struct +from typing import Any, Iterable + +import numpy as np + +from .checkpoint import numpy_state + + +MODEL_ID = "Comfy-Org/MiniMax-H3" +CHECKPOINT_REVISION = "4cc1d817b6184899b41293954329f576cb5ae86b" +CHECKPOINT_FILENAME = "diffusion_models/minimax_h3_fl2va_int8_convrot.safetensors" +CHECKPOINT_BYTES = 34_038_892_334 +REF2VA_CHECKPOINT_FILENAME = "diffusion_models/minimax_h3_ref2va_int8_convrot.safetensors" +REF2VA_CHECKPOINT_BYTES = 34_038_894_550 + +_MAX_HEADER_BYTES = 8 << 20 +_NUM_HEADS = 56 +_HEAD_DIM = 128 +_HIDDEN_SIZE = 5_376 +_INNER_DIM = _NUM_HEADS * _HEAD_DIM +_FFN_DIM = 14_336 +_TIME_DIM = 2_688 +_QUANT_FORMAT = "int8_tensorwise" +_HF_CACHE_REPOSITORY = "models--Comfy-Org--MiniMax-H3" +_HF_LOCAL_DOWNLOAD_METADATA = Path(".cache/huggingface/download") + + +@dataclass(frozen=True) +class ConvRotInt8Weight: + """One Comfy tensorwise-INT8 weight in Diffusers logical row order. + + ``packed_parent`` and ``row_slice`` let the fused-QKV graph path recover a + single packed parent matrix when the loader returned three + logical ``to_q``/``to_k``/``to_v`` entries. No concatenation of those + child views is necessary. + """ + + qweight: np.ndarray + scale: np.ndarray + group_size: int + packed_parent: ConvRotInt8Weight | None = None + row_slice: tuple[int, int] | None = None + is_full_fused_qkv: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.qweight, np.ndarray) or self.qweight.dtype != np.int8: + raise TypeError("MiniMax-H3 ConvRot qweight must be a NumPy int8 array") + if self.qweight.ndim != 2 or not self.qweight.flags.c_contiguous: + raise ValueError("MiniMax-H3 ConvRot qweight must be contiguous [out,in]") + if not isinstance(self.scale, np.ndarray) or self.scale.dtype != np.float32: + raise TypeError("MiniMax-H3 ConvRot scale must be a NumPy float32 array") + if self.scale.shape not in ((self.qweight.shape[0],), (self.qweight.shape[0], 1)): + raise ValueError("MiniMax-H3 ConvRot scale must contain one value per output row") + if not self.scale.flags.c_contiguous: + raise ValueError("MiniMax-H3 ConvRot scale must be contiguous") + _validate_group_size(self.group_size, self.qweight.shape[1]) + if self.is_full_fused_qkv: + if self.packed_parent is not None or self.row_slice is not None: + raise ValueError("A full fused-QKV ConvRot weight cannot itself be a child view") + elif self.packed_parent is not None: + if not self.packed_parent.is_full_fused_qkv or self.row_slice is None: + raise ValueError("A fused-QKV child requires a full parent and row_slice") + start, end = self.row_slice + if not 0 <= start < end <= self.packed_parent.qweight.shape[0]: + raise ValueError("MiniMax-H3 fused-QKV child row_slice is out of bounds") + if end - start != self.qweight.shape[0]: + raise ValueError("MiniMax-H3 fused-QKV child row_slice does not match qweight") + elif self.row_slice is not None: + raise ValueError("MiniMax-H3 ConvRot row_slice requires packed_parent") + + +@dataclass(frozen=True) +class QuantizedSourceFileIdentity: + """Path-free local identity used only to make staged resume fail closed.""" + + device: int + inode: int + size_bytes: int + mtime_ns: int + ctime_ns: int + + def receipt_metadata(self) -> dict[str, int]: + return { + "device": self.device, + "inode": self.inode, + "size_bytes": self.size_bytes, + "mtime_ns": self.mtime_ns, + "ctime_ns": self.ctime_ns, + } + + +@dataclass(frozen=True) +class QuantizedCheckpointIdentity: + """Path-free identity of the pinned public Comfy-Org checkpoint.""" + + model_id: str + revision: str + filename: str + size_bytes: int + tensor_count: int + quantized_weight_count: int + source_file_identity: QuantizedSourceFileIdentity | None = None + + def bundle_metadata(self) -> dict[str, object]: + """Return stable public provenance without a workstation path.""" + + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "filename": self.filename, + "size_bytes": self.size_bytes, + "tensor_count": self.tensor_count, + "quantized_weight_count": self.quantized_weight_count, + "quantization": "int8_tensorwise_convrot", + "runtime_framework": None, + } + + +QUANTIZED_CHECKPOINT_IDENTITY = QuantizedCheckpointIdentity( + model_id=MODEL_ID, + revision=CHECKPOINT_REVISION, + filename=CHECKPOINT_FILENAME, + size_bytes=CHECKPOINT_BYTES, + tensor_count=1_035, + quantized_weight_count=250, +) +QUANTIZED_REF2VA_CHECKPOINT_IDENTITY = replace( + QUANTIZED_CHECKPOINT_IDENTITY, + filename=REF2VA_CHECKPOINT_FILENAME, + size_bytes=REF2VA_CHECKPOINT_BYTES, +) + + +def _checkpoint_identity(workflow: str) -> QuantizedCheckpointIdentity: + if workflow == "fl2va": + return QUANTIZED_CHECKPOINT_IDENTITY + if workflow == "ref2va": + return QUANTIZED_REF2VA_CHECKPOINT_IDENTITY + raise ValueError("MiniMax-H3 quantized checkpoint workflow must be 'fl2va' or 'ref2va'") + + +_EXPECTED_CONFIG: dict[str, object] = { + "transformer": { + "hidden_size": 5_376, + "num_layers": 50, + "token_refiner_num_layers": 2, + "num_attention_heads": 56, + "attention_head_dim": 128, + "ffn_hidden_size": 14_336, + "latents_dim": 24, + "audio_latents_dim": 32, + "patch_size": [1, 2, 2], + "text_dim": 5_120, + "timestep_input_dim": 256, + "time_embed_hidden_size": 5_376, + "time_embed_dim": 2_688, + "adaln_out_features": 96_768, + "final_adaln_out_features": 10_752, + "rope_inv_freq_len": 16, + "norm_eps": 1.0e-5, + "qk_norm_eps": 1.0e-5, + "final_norm_eps": 1.0e-5, + "image_model": "minimax_h3", + } +} + + +@dataclass(frozen=True) +class _TensorSpec: + dtype: str + shape: tuple[int, ...] + + +def _add_quantized( + specs: dict[str, _TensorSpec], + groups: dict[str, int], + base: str, + shape: tuple[int, int], + group_size: int, +) -> None: + specs[f"{base}.weight"] = _TensorSpec("I8", shape) + specs[f"{base}.weight_scale"] = _TensorSpec("F32", (shape[0], 1)) + marker_bytes = len( + json.dumps( + { + "format": _QUANT_FORMAT, + "convrot": True, + "convrot_groupsize": group_size, + } + ).encode("utf-8") + ) + specs[f"{base}.comfy_quant"] = _TensorSpec("U8", (marker_bytes,)) + groups[base] = group_size + + +def _physical_contract() -> tuple[dict[str, _TensorSpec], dict[str, int]]: + specs: dict[str, _TensorSpec] = { + "video_patch_proj.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 96)), + "video_patch_proj.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "audio_patch_proj.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 32)), + "audio_patch_proj.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "condition_proj.weight": _TensorSpec("BF16", (_HIDDEN_SIZE, 5_120)), + "condition_proj.bias": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "time_embedder.proj_in.weight": _TensorSpec("F32", (_HIDDEN_SIZE, 256)), + "time_embedder.proj_in.bias": _TensorSpec("F32", (_HIDDEN_SIZE,)), + "time_embedder.proj_out.weight": _TensorSpec("F32", (_TIME_DIM, _HIDDEN_SIZE)), + "time_embedder.proj_out.bias": _TensorSpec("F32", (_TIME_DIM,)), + "token_refiner.final_norm.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "final_layer.norm.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + "final_layer.adaln_proj.linear.weight": _TensorSpec( + "BF16", (2 * _HIDDEN_SIZE, _TIME_DIM) + ), + "final_layer.adaln_proj.linear.bias": _TensorSpec("BF16", (2 * _HIDDEN_SIZE,)), + "final_layer.video_out.weight": _TensorSpec("F32", (96, _HIDDEN_SIZE)), + "final_layer.video_out.bias": _TensorSpec("F32", (96,)), + "final_layer.audio_out.weight": _TensorSpec("F32", (32, _HIDDEN_SIZE)), + "final_layer.audio_out.bias": _TensorSpec("F32", (32,)), + "rope.inv_freq": _TensorSpec("F32", (16,)), + } + groups: dict[str, int] = {} + for index in range(50): + prefix = f"blocks.{index}" + specs[f"{prefix}.norm1.weight"] = _TensorSpec("BF16", (_HIDDEN_SIZE,)) + specs[f"{prefix}.norm2.weight"] = _TensorSpec("BF16", (_HIDDEN_SIZE,)) + specs[f"{prefix}.attn.q_norm.weight"] = _TensorSpec("BF16", (_HEAD_DIM,)) + specs[f"{prefix}.attn.k_norm.weight"] = _TensorSpec("BF16", (_HEAD_DIM,)) + specs[f"{prefix}.adaln_proj.linear.bias"] = _TensorSpec( + "BF16", (18 * _HIDDEN_SIZE,) + ) + _add_quantized( + specs, + groups, + f"{prefix}.adaln_proj.linear", + (18 * _HIDDEN_SIZE, _TIME_DIM), + 64, + ) + _add_quantized( + specs, + groups, + f"{prefix}.attn.qkv_proj", + (3 * _INNER_DIM, _HIDDEN_SIZE), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.attn.out_proj", + (_HIDDEN_SIZE, _INNER_DIM), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.mlp.fc1", + (2 * _FFN_DIM, _HIDDEN_SIZE), + 256, + ) + _add_quantized( + specs, + groups, + f"{prefix}.mlp.fc2", + (_HIDDEN_SIZE, _FFN_DIM), + 256, + ) + for index in range(2): + prefix = f"token_refiner.blocks.{index}" + specs.update( + { + f"{prefix}.norm1.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + f"{prefix}.norm2.weight": _TensorSpec("BF16", (_HIDDEN_SIZE,)), + f"{prefix}.attn.qkv_proj.weight": _TensorSpec( + "BF16", (3 * _INNER_DIM, _HIDDEN_SIZE) + ), + f"{prefix}.attn.q_norm.weight": _TensorSpec("BF16", (_HEAD_DIM,)), + f"{prefix}.attn.k_norm.weight": _TensorSpec("BF16", (_HEAD_DIM,)), + f"{prefix}.attn.out_proj.weight": _TensorSpec( + "BF16", (_HIDDEN_SIZE, _INNER_DIM) + ), + f"{prefix}.mlp.fc1.weight": _TensorSpec( + "BF16", (2 * _FFN_DIM, _HIDDEN_SIZE) + ), + f"{prefix}.mlp.fc2.weight": _TensorSpec("BF16", (_HIDDEN_SIZE, _FFN_DIM)), + } + ) + if len(specs) != 1_035 or len(groups) != 250: + raise RuntimeError("MiniMax-H3 quantized checkpoint contract is incomplete") + return specs, groups + + +_PHYSICAL_SPECS, _QUANT_GROUPS = _physical_contract() + + +def _logical_contract() -> dict[str, tuple[str, str]]: + direct = { + "proj_in.weight": "video_patch_proj.weight", + "proj_in.bias": "video_patch_proj.bias", + "audio_proj_in.weight": "audio_patch_proj.weight", + "audio_proj_in.bias": "audio_patch_proj.bias", + "context_embedder.weight": "condition_proj.weight", + "context_embedder.bias": "condition_proj.bias", + "time_embedder.linear_1.weight": "time_embedder.proj_in.weight", + "time_embedder.linear_1.bias": "time_embedder.proj_in.bias", + "time_embedder.linear_2.weight": "time_embedder.proj_out.weight", + "time_embedder.linear_2.bias": "time_embedder.proj_out.bias", + "token_refiner.final_norm.weight": "token_refiner.final_norm.weight", + "norm_out.norm.weight": "final_layer.norm.weight", + "norm_out.linear.weight": "final_layer.adaln_proj.linear.weight", + "norm_out.linear.bias": "final_layer.adaln_proj.linear.bias", + "proj_out.weight": "final_layer.video_out.weight", + "proj_out.bias": "final_layer.video_out.bias", + "audio_proj_out.weight": "final_layer.audio_out.weight", + "audio_proj_out.bias": "final_layer.audio_out.bias", + } + result = {logical: (physical, "direct") for logical, physical in direct.items()} + for physical_root, logical_root, count, has_adaln in ( + ("blocks", "transformer_blocks", 50, True), + ("token_refiner.blocks", "token_refiner.refiner_blocks", 2, False), + ): + for index in range(count): + physical = f"{physical_root}.{index}" + logical = f"{logical_root}.{index}" + result[f"{logical}.norm1.weight"] = (f"{physical}.norm1.weight", "direct") + result[f"{logical}.norm2.weight"] = (f"{physical}.norm2.weight", "direct") + for qkv_index, name in enumerate(("q", "k", "v")): + result[f"{logical}.attn.to_{name}.weight"] = ( + f"{physical}.attn.qkv_proj.weight", + f"qkv:{qkv_index}", + ) + result[f"{logical}.attn.norm_q.weight"] = ( + f"{physical}.attn.q_norm.weight", + "direct", + ) + result[f"{logical}.attn.norm_k.weight"] = ( + f"{physical}.attn.k_norm.weight", + "direct", + ) + result[f"{logical}.attn.to_out.0.weight"] = ( + f"{physical}.attn.out_proj.weight", + "direct", + ) + result[f"{logical}.ff.net.0.proj.weight"] = ( + f"{physical}.mlp.fc1.weight", + "swap_swiglu", + ) + result[f"{logical}.ff.net.2.weight"] = ( + f"{physical}.mlp.fc2.weight", + "direct", + ) + if has_adaln: + result[f"{logical}.adaln_proj.linear.weight"] = ( + f"{physical}.adaln_proj.linear.weight", + "direct", + ) + result[f"{logical}.adaln_proj.linear.bias"] = ( + f"{physical}.adaln_proj.linear.bias", + "direct", + ) + return result + + +_LOGICAL_MAP = _logical_contract() + + +def _validate_group_size(group_size: int, in_features: int) -> None: + if not isinstance(group_size, int) or isinstance(group_size, bool) or group_size < 4: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + value = group_size + while value > 1 and value % 4 == 0: + value //= 4 + if value != 1: + raise ValueError("MiniMax-H3 ConvRot group_size must be a power of four") + if in_features % group_size: + raise ValueError("MiniMax-H3 ConvRot group_size must divide the input width") + + +def _read_header(path: Path) -> tuple[dict[str, Any], int]: + with path.open("rb") as stream: + prefix = stream.read(8) + if len(prefix) != 8: + raise ValueError("MiniMax-H3 quantized checkpoint has an incomplete header") + (header_bytes,) = struct.unpack(" dict[str, dict[str, Any]]: + metadata = header.get("__metadata__") + if not isinstance(metadata, dict) or set(metadata) != {"config"}: + raise ValueError("MiniMax-H3 quantized checkpoint metadata schema is invalid") + try: + config = json.loads(metadata["config"]) + except (TypeError, json.JSONDecodeError) as error: + raise ValueError("MiniMax-H3 quantized checkpoint config metadata is invalid") from error + if config != _EXPECTED_CONFIG: + raise ValueError( + "MiniMax-H3 quantized checkpoint config is not the full released 50-layer architecture" + ) + + tensors = {name: value for name, value in header.items() if name != "__metadata__"} + actual = set(tensors) + expected = set(_PHYSICAL_SPECS) + if actual != expected: + raise ValueError( + "MiniMax-H3 quantized checkpoint tensor inventory mismatch: " + f"missing={sorted(expected - actual)[:8]}, unexpected={sorted(actual - expected)[:8]}" + ) + + dtype_bytes = {"BF16": 2, "F32": 4, "I8": 1, "U8": 1} + regions: list[tuple[int, int, str]] = [] + for name, spec in _PHYSICAL_SPECS.items(): + entry = tensors[name] + if not isinstance(entry, dict) or set(entry) != {"dtype", "shape", "data_offsets"}: + raise ValueError(f"MiniMax-H3 quantized tensor header is invalid: {name}") + shape = entry["shape"] + offsets = entry["data_offsets"] + if entry["dtype"] != spec.dtype or shape != list(spec.shape): + raise ValueError( + f"MiniMax-H3 quantized tensor contract mismatch for {name}: " + f"expected={spec.dtype}{spec.shape}, actual={entry.get('dtype')}{shape}" + ) + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets) + ): + raise ValueError(f"MiniMax-H3 quantized tensor offsets are invalid: {name}") + start, end = offsets + expected_bytes = math.prod(spec.shape) * dtype_bytes[spec.dtype] + if start < 0 or end - start != expected_bytes: + raise ValueError(f"MiniMax-H3 quantized tensor byte range is invalid: {name}") + regions.append((start, end, name)) + cursor = 0 + for start, end, name in sorted(regions): + if start != cursor: + raise ValueError(f"MiniMax-H3 quantized tensor storage is not contiguous before {name}") + cursor = end + if data_offset + cursor != file_size: + raise ValueError("MiniMax-H3 quantized checkpoint payload size is invalid") + return tensors + + +def _read_marker(path: Path, data_offset: int, entry: dict[str, Any]) -> dict[str, Any]: + start, end = entry["data_offsets"] + with path.open("rb") as stream: + stream.seek(data_offset + start) + payload = stream.read(end - start) + try: + marker = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError("MiniMax-H3 comfy_quant marker is invalid JSON") from error + if not isinstance(marker, dict): + raise ValueError("MiniMax-H3 comfy_quant marker must be an object") + return marker + + +def _source_file_identity(path: Path) -> QuantizedSourceFileIdentity: + try: + metadata = path.stat() + except OSError as error: + raise ValueError("MiniMax-H3 quantized checkpoint is unavailable") from error + if not stat.S_ISREG(metadata.st_mode): + raise ValueError("MiniMax-H3 quantized checkpoint must be a regular file") + return QuantizedSourceFileIdentity( + device=int(metadata.st_dev), + inode=int(metadata.st_ino), + size_bytes=int(metadata.st_size), + mtime_ns=int(metadata.st_mtime_ns), + ctime_ns=int(metadata.st_ctime_ns), + ) + + +def _authenticate_quantized_source(path: Path, *, workflow: str = "fl2va") -> None: + """Validate the pinned HF cache/local-dir layout without rereading 34 GB.""" + + identity = _checkpoint_identity(workflow) + relative = Path(identity.filename) + try: + source_root = path.parents[len(relative.parts) - 1] + except IndexError as error: + raise ValueError("MiniMax-H3 quantized checkpoint path is invalid") from error + expected_path = source_root / relative + if os.path.normcase(os.path.abspath(path)) != os.path.normcase(os.path.abspath(expected_path)): + raise ValueError( + "MiniMax-H3 quantized checkpoint must retain its released diffusion_models path" + ) + + if source_root.parent.name == "snapshots": + repository_root = source_root.parent.parent + if source_root.name != identity.revision or repository_root.name != _HF_CACHE_REPOSITORY: + raise ValueError( + "MiniMax-H3 quantized checkpoint is not from the pinned Hugging Face snapshot" + ) + return + + metadata_path = source_root / _HF_LOCAL_DOWNLOAD_METADATA / f"{relative.as_posix()}.metadata" + try: + metadata = metadata_path.read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeDecodeError) as error: + raise ValueError( + "Cannot validate the MiniMax-H3 quantized checkpoint source. " + "Download it with `hf download Comfy-Org/MiniMax-H3 --revision " + f'{identity.revision} --include "{identity.filename}" --local-dir ` ' + "so Hugging Face writes pinned download metadata." + ) from error + if not metadata or metadata[0] != identity.revision: + raise ValueError( + "MiniMax-H3 quantized checkpoint Hugging Face metadata does not match the pinned " + "revision" + ) + + +def validate_quantized_transformer_checkpoint( + checkpoint_file: str | Path, + *, + workflow: str = "fl2va", +) -> QuantizedCheckpointIdentity: + """Validate the pinned full INT8 ConvRot checkpoint for the selected workflow. + + A Hugging Face cache symlink is accepted, but the lexical filename, exact + byte size, complete safetensors inventory, architecture config, and every + embedded quantization marker are required. + """ + + identity = _checkpoint_identity(workflow) + path = Path(checkpoint_file) + if not path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 quantized checkpoint is missing: {path}") + expected_name = Path(identity.filename).name + if path.name != expected_name: + raise ValueError( + f"MiniMax-H3 quantized checkpoint filename must be {expected_name}, got {path.name}" + ) + size = path.stat().st_size + if size != identity.size_bytes: + raise ValueError( + "MiniMax-H3 quantized checkpoint size mismatch: " + f"expected={identity.size_bytes}, actual={size}" + ) + source_identity = _source_file_identity(path) + _authenticate_quantized_source(path, workflow=workflow) + header, data_offset = _read_header(path) + tensors = _validate_header_inventory(header, data_offset, size) + for base, group_size in sorted(_QUANT_GROUPS.items()): + marker_name = f"{base}.comfy_quant" + marker = _read_marker(path, data_offset, tensors[marker_name]) + expected_marker = { + "format": _QUANT_FORMAT, + "convrot": True, + "convrot_groupsize": group_size, + } + if marker != expected_marker: + raise ValueError( + f"MiniMax-H3 unsupported comfy_quant config for {base}: {marker!r}" + ) + _validate_group_size(group_size, _PHYSICAL_SPECS[f"{base}.weight"].shape[1]) + if _source_file_identity(path) != source_identity: + raise ValueError("MiniMax-H3 quantized checkpoint changed while it was validated") + return replace( + identity, + source_file_identity=source_identity, + ) + + +def _swap_swiglu_rows(array: np.ndarray) -> np.ndarray: + if array.shape[0] % 2: + raise ValueError("MiniMax-H3 fused SwiGLU weight must have an even output width") + half = array.shape[0] // 2 + return np.ascontiguousarray(np.concatenate((array[half:], array[:half]), axis=0)) + + +def _make_quantized_weight( + qweight: np.ndarray, + scale: np.ndarray, + group_size: int, + *, + transform: str, +) -> ConvRotInt8Weight: + qweight = np.asarray(qweight) + scale = np.asarray(scale) + if transform == "swap_swiglu": + qweight = _swap_swiglu_rows(qweight) + scale = _swap_swiglu_rows(scale) + elif transform != "direct": + raise RuntimeError(f"MiniMax-H3 quantized transform is invalid: {transform}") + return ConvRotInt8Weight(qweight=qweight, scale=scale, group_size=group_size) + + +def _make_qkv_values( + qweight: np.ndarray, + scale: np.ndarray | None, + group_size: int | None, +) -> tuple[Any, Any, Any]: + packed_weight = np.ascontiguousarray(qweight) + expected_rows = 3 * _NUM_HEADS * _HEAD_DIM + if packed_weight.shape[0] != expected_rows: + raise ValueError( + f"MiniMax-H3 fused QKV has {packed_weight.shape[0]} rows, expected {expected_rows}" + ) + width = packed_weight.shape[0] // 3 + if scale is None: + return tuple(packed_weight[index * width : (index + 1) * width] for index in range(3)) + if group_size is None: + raise RuntimeError("MiniMax-H3 QKV quantization group is missing") + packed_scale = np.ascontiguousarray(scale) + parent = ConvRotInt8Weight( + qweight=packed_weight, + scale=packed_scale, + group_size=group_size, + is_full_fused_qkv=True, + ) + children = [] + for index in range(3): + start, end = index * width, (index + 1) * width + children.append( + ConvRotInt8Weight( + qweight=parent.qweight[start:end], + scale=parent.scale[start:end], + group_size=group_size, + packed_parent=parent, + row_slice=(start, end), + ) + ) + return tuple(children) + + +def load_selected_quantized_transformer_weights( + checkpoint_file: str | Path, + logical_names: Iterable[str], + *, + workflow: str = "fl2va", +) -> dict[str, Any]: + """Load only requested frozen-Diffusers logical transformer weights.""" + + path = Path(checkpoint_file) + validate_quantized_transformer_checkpoint(path, workflow=workflow) + requested = tuple(logical_names) + if len(requested) != len(set(requested)): + raise ValueError("MiniMax-H3 quantized weight request contains duplicate logical names") + unknown = sorted(set(requested) - set(_LOGICAL_MAP)) + if unknown: + raise ValueError(f"Unsupported MiniMax-H3 quantized logical tensor names: {unknown}") + if not requested: + return {} + + physical_names: set[str] = set() + for logical_name in requested: + physical_weight, _transform = _LOGICAL_MAP[logical_name] + physical_names.add(physical_weight) + base = physical_weight.removesuffix(".weight") + if base in _QUANT_GROUPS: + physical_names.add(f"{base}.weight_scale") + + from safetensors import safe_open + + with safe_open(path, framework="pt", device="cpu") as reader: + state = {name: reader.get_tensor(name) for name in sorted(physical_names)} + arrays = numpy_state(state) + + result: dict[str, Any] = {} + qkv_cache: dict[str, tuple[Any, Any, Any]] = {} + quant_cache: dict[tuple[str, str], ConvRotInt8Weight] = {} + for logical_name in requested: + physical_weight, transform = _LOGICAL_MAP[logical_name] + if transform.startswith("qkv:"): + values = qkv_cache.get(physical_weight) + if values is None: + base = physical_weight.removesuffix(".weight") + scale = arrays.get(f"{base}.weight_scale") + values = _make_qkv_values( + arrays[physical_weight], + scale, + _QUANT_GROUPS.get(base), + ) + qkv_cache[physical_weight] = values + result[logical_name] = values[int(transform.removeprefix("qkv:"))] + continue + base = physical_weight.removesuffix(".weight") + if base in _QUANT_GROUPS: + cache_key = (physical_weight, transform) + value = quant_cache.get(cache_key) + if value is None: + value = _make_quantized_weight( + arrays[physical_weight], + arrays[f"{base}.weight_scale"], + _QUANT_GROUPS[base], + transform=transform, + ) + quant_cache[cache_key] = value + result[logical_name] = value + elif transform == "swap_swiglu": + result[logical_name] = _swap_swiglu_rows(arrays[physical_weight]) + elif transform == "direct": + result[logical_name] = arrays[physical_weight] + else: + raise RuntimeError(f"MiniMax-H3 logical transform is invalid: {transform}") + return result diff --git a/families/minimax_h3/ref2va_audio_encoder_builder.py b/families/minimax_h3/ref2va_audio_encoder_builder.py new file mode 100644 index 0000000000..c6919b799b --- /dev/null +++ b/families/minimax_h3/ref2va_audio_encoder_builder.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native FP32 TensorRT encoder for MiniMax-H3 Ref2VA soundtracks. + +The plan consumes stereo as two batch items, follows the released DAC encoder +and causal attention projection, and emits the posterior mean. Ref2VA never +evaluates ``logs_proj`` and never samples this posterior; C++ applies the +checkpoint's 32-channel mean/std normalization and then packs channel-major +rows. No Torch, audio framework, or process sidecar is part of the runtime. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import gc +import math +from pathlib import Path + +import numpy as np + +from . import trt_compat + +from . import graph_ops as op +from .audio_vae_builder import ( + _conv1d, + _fold_checkpoint_weight, + _require_array, +) +from .fl2va_contract import PlanAbi, TensorAbi + + +trt = trt_compat.get_trt() + +AUDIO_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES = 32 << 30 + + +@dataclass(frozen=True) +class Ref2VAAudioEncoderProfile: + batch_size: int = 2 + sampling_rate: int = 32_000 + min_samples: int = 64_000 + opt_samples: int = 165_600 + max_samples: int = 480_000 + encoder_dim: int = 64 + encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5) + latent_dim: int = 2048 + latent_channels: int = 32 + num_attention_heads: int = 8 + + @property + def hop_length(self) -> int: + return math.prod(self.encoder_rates) + + @property + def sample_profile(self) -> tuple[int, int, int]: + return self.min_samples, self.opt_samples, self.max_samples + + @property + def latent_profile(self) -> tuple[int, int, int]: + return tuple(value // self.hop_length for value in self.sample_profile) + + def validate(self) -> None: + values = ( + self.batch_size, + self.sampling_rate, + self.min_samples, + self.opt_samples, + self.max_samples, + self.encoder_dim, + self.latent_dim, + self.latent_channels, + self.num_attention_heads, + ) + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in values + ): + raise ValueError("MiniMax-H3 Ref2VA audio-encoder dimensions must be positive integers") + if self.batch_size != 2 or self.sampling_rate != 32_000: + raise ValueError("MiniMax-H3 Ref2VA audio encoder requires stereo at 32 kHz") + if not self.min_samples <= self.opt_samples <= self.max_samples: + raise ValueError("MiniMax-H3 Ref2VA audio samples must satisfy min <= opt <= max") + if any(value % self.hop_length for value in self.sample_profile): + raise ValueError( + "MiniMax-H3 Ref2VA audio inputs must be right-padded to the 800-sample hop" + ) + if self.latent_dim % self.num_attention_heads: + raise ValueError("MiniMax-H3 Ref2VA audio attention width must divide its heads") + + +DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE = Ref2VAAudioEncoderProfile() + + +def ref2va_audio_encoder_abi( + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, +) -> PlanAbi: + """Strict native ABI for explicit audio and attached video soundtracks. + + The C++ caller resamples/interleaves to stereo 32 kHz, separates the two + channels into this batch-of-two representation, and right-pads each clip + to the 800-sample hop. The plan emits the released posterior mode. The + caller then applies the 32 per-channel ``latents_mean``/``latents_std`` + values pinned in bundle metadata and preserves channel-major row order. + """ + + profile.validate() + sample_shapes = tuple((profile.batch_size, 1, samples) for samples in profile.sample_profile) + latent_shapes = tuple( + (profile.batch_size, profile.latent_channels, frames) for frames in profile.latent_profile + ) + return PlanAbi( + filename="ref2va_audio_vae_encoder.plan", + inputs=(TensorAbi("audio_samples", "float32", *sample_shapes),), + outputs=(TensorAbi("posterior_mean", "float32", *latent_shapes),), + ) + + +def checkpoint_keys( + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, +) -> tuple[str, ...]: + """Return exactly the weights evaluated by ``posterior.mode()``.""" + + profile.validate() + names = [ + "encoder.block.0.weight_g", + "encoder.block.0.weight_v", + "encoder.block.0.bias", + ] + channels = profile.encoder_dim + for stage, _stride in enumerate(profile.encoder_rates, start=1): + prefix = f"encoder.block.{stage}.block" + for residual in range(3): + residual_prefix = f"{prefix}.{residual}.block" + names.extend( + ( + f"{residual_prefix}.0.alpha", + f"{residual_prefix}.1.weight_g", + f"{residual_prefix}.1.weight_v", + f"{residual_prefix}.1.bias", + f"{residual_prefix}.2.alpha", + f"{residual_prefix}.3.weight_g", + f"{residual_prefix}.3.weight_v", + f"{residual_prefix}.3.bias", + ) + ) + names.extend( + ( + f"{prefix}.3.alpha", + f"{prefix}.4.weight_g", + f"{prefix}.4.weight_v", + f"{prefix}.4.bias", + ) + ) + channels *= 2 + names.extend( + ( + f"encoder.block.{len(profile.encoder_rates) + 1}.alpha", + f"encoder.block.{len(profile.encoder_rates) + 2}.weight_g", + f"encoder.block.{len(profile.encoder_rates) + 2}.weight_v", + f"encoder.block.{len(profile.encoder_rates) + 2}.bias", + "pre_block.norm1.weight", + "pre_block.norm1.bias", + "pre_block.attn.qkv.weight", + "pre_block.attn.q_bias", + "pre_block.attn.v_bias", + "pre_block.attn.zero_k_bias", + "pre_block.attn.proj.weight", + "pre_block.attn.proj.bias", + "pre_block.proj.weight", + "pre_block.proj.bias", + "pre_block.norm3.weight", + "pre_block.norm3.bias", + "pre_block.norm2.weight", + "pre_block.norm2.bias", + "pre_block.mlp.norm.weight", + "pre_block.mlp.norm.bias", + "pre_block.mlp.w0.weight", + "pre_block.mlp.w0.bias", + "pre_block.mlp.w1.weight", + "pre_block.mlp.w1.bias", + "pre_block.mlp.w2.weight", + "pre_block.mlp.w2.bias", + "mean_proj.weight", + "mean_proj.bias", + ) + ) + return tuple(names) + + +def _snake(network, hidden, weights, name: str): + channels = int(tuple(hidden.shape)[1]) + alpha = _require_array(weights, f"{name}.alpha", (1, channels, 1)).reshape(1, channels, 1, 1) + alpha_tensor = op.weight_constant(network, alpha) + phase = network.add_elementwise(hidden, alpha_tensor, trt.ElementWiseOperation.PROD).get_output( + 0 + ) + sine = network.add_unary(phase, trt.UnaryOperation.SIN).get_output(0) + square = network.add_elementwise(sine, sine, trt.ElementWiseOperation.PROD).get_output(0) + epsilon = op.constant(network, np.full((1, 1, 1, 1), 1.0e-9, np.float32)) + denominator = network.add_elementwise( + alpha_tensor, epsilon, trt.ElementWiseOperation.SUM + ).get_output(0) + periodic = network.add_elementwise( + square, denominator, trt.ElementWiseOperation.DIV + ).get_output(0) + return network.add_elementwise(hidden, periodic, trt.ElementWiseOperation.SUM).get_output(0) + + +def _residual_unit( + network, + hidden, + weights, + prefix: str, + *, + channels: int, + dilation: int, + owned_weights: list[np.ndarray], +): + update = _snake(network, hidden, weights, f"{prefix}.0") + conv1 = f"{prefix}.1" + update = _conv1d( + network, + update, + _fold_checkpoint_weight(weights, conv1, (channels, channels, 7)), + _require_array(weights, f"{conv1}.bias", (channels,)), + padding=3 * dilation, + dilation=dilation, + name=conv1, + owned_weights=owned_weights, + ) + update = _snake(network, update, weights, f"{prefix}.2") + conv2 = f"{prefix}.3" + update = _conv1d( + network, + update, + _fold_checkpoint_weight(weights, conv2, (channels, channels, 1)), + _require_array(weights, f"{conv2}.bias", (channels,)), + name=conv2, + owned_weights=owned_weights, + ) + return network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + +def _layer_norm(network, hidden, weights, prefix: str, width: int): + rank = len(tuple(hidden.shape)) + shape = (1,) * (rank - 1) + (width,) + gamma = op.weight_constant( + network, _require_array(weights, f"{prefix}.weight", (width,)).reshape(shape) + ) + beta = op.weight_constant( + network, _require_array(weights, f"{prefix}.bias", (width,)).reshape(shape) + ) + layer = network.add_normalization_v2(hidden, gamma, beta, 1 << (rank - 1)) + if layer is None: + raise RuntimeError(f"TensorRT rejected MiniMax-H3 audio LayerNorm {prefix}") + layer.name = prefix + layer.epsilon = 1.0e-5 + return layer.get_output(0) + + +def _linear(network, hidden, weights, prefix: str): + weight = _require_array( + weights, + f"{prefix}.weight", + tuple(np.asarray(weights[f"{prefix}.weight"]).shape), + ) + # ``mean_proj`` is a checkpoint Conv1d(k=1). The projection is evaluated + # on [B,T,C] rows here, so its singleton kernel axis is removed without + # changing either the values or their input/output ordering. + if weight.ndim == 3: + if weight.shape[-1] != 1: + raise ValueError(f"MiniMax-H3 audio linear {prefix} must have a unit kernel") + weight = np.ascontiguousarray(weight[..., 0]) + return op.linear( + network, + hidden, + weight, + _require_array( + weights, + f"{prefix}.bias", + tuple(np.asarray(weights[f"{prefix}.bias"]).shape), + ), + bf16=False, + ) + + +def _causal_attention_projection(network, hidden, weights, profile: Ref2VAAudioEncoderProfile): + batch = profile.batch_size + in_width = profile.latent_dim + out_width = profile.latent_channels + heads = profile.num_attention_heads + head_dim = in_width // heads + + normalized = _layer_norm(network, hidden, weights, "pre_block.norm1", in_width) + qkv_weight = _require_array(weights, "pre_block.attn.qkv.weight", (3 * in_width, in_width)) + q_bias = _require_array(weights, "pre_block.attn.q_bias", (in_width,)) + k_bias = _require_array(weights, "pre_block.attn.zero_k_bias", (in_width,)) + if np.any(k_bias != 0.0): + raise ValueError("MiniMax-H3 audio frozen key bias must be exactly zero") + v_bias = _require_array(weights, "pre_block.attn.v_bias", (in_width,)) + qkv = op.linear( + network, + normalized, + qkv_weight, + np.concatenate((q_bias, k_bias, v_bias)), + bf16=False, + ) + + tensors = [] + for index in range(3): + tensor = op.dynamic_slice( + network, + qkv, + (0, 0, index * in_width), + (batch, None, in_width), + ) + reshape = network.add_shuffle(tensor) + reshape.reshape_dims = (batch, -1, heads, head_dim) + reshape.second_transpose = trt.Permutation((0, 2, 1, 3)) + tensors.append(reshape.get_output(0)) + query, key, value = tensors + scale = op.constant( + network, + np.full((1, 1, 1, 1), 1.0 / math.sqrt(head_dim), np.float32), + ) + query = network.add_elementwise(query, scale, trt.ElementWiseOperation.PROD).get_output(0) + attention = network.add_attention( + query, + key, + value, + trt.AttentionNormalizationOp.SOFTMAX, + True, + ) + if attention is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 audio causal attention") + attention.name = "pre_block.attn.native_causal_attention" + attention.metadata = "trtmc.native_op=IAttention;causal=true;decomposable=true" + # TRT-RTX 1.6 has no dedicated causal kernel covering the complete + # dynamic 80..600-row public profile. Allow TRT itself to lower the + # IAttention composite into native layers; this remains inside the plan + # and introduces no plugin, framework, or runtime dependency. + attention.decomposable = True + + # [B,H,T,D] -> mean heads -> [B,T,D] -> average each exact group + # of D/out_width=8 channels, matching adaptive_avg_pool1d. + pooled_heads = network.add_reduce( + attention.get_output(0), + trt.ReduceOperation.AVG, + 1 << 1, + False, + ) + reshape = network.add_shuffle(pooled_heads.get_output(0)) + reshape.reshape_dims = (batch, -1, out_width, head_dim // out_width) + pooled_width = network.add_reduce( + reshape.get_output(0), + trt.ReduceOperation.AVG, + 1 << 3, + False, + ).get_output(0) + attended = _linear(network, pooled_width, weights, "pre_block.attn.proj") + + residual = _layer_norm(network, hidden, weights, "pre_block.norm3", in_width) + residual = _linear(network, residual, weights, "pre_block.proj") + hidden = network.add_elementwise(residual, attended, trt.ElementWiseOperation.SUM).get_output(0) + update = _layer_norm(network, hidden, weights, "pre_block.norm2", out_width) + update = _layer_norm(network, update, weights, "pre_block.mlp.norm", out_width) + left = _linear(network, update, weights, "pre_block.mlp.w0") + gelu = network.add_activation(left, trt.ActivationType.GELU_TANH) + if gelu is None: + raise RuntimeError("TensorRT rejected MiniMax-H3 audio GeGLU activation") + right = _linear(network, update, weights, "pre_block.mlp.w1") + update = network.add_elementwise( + gelu.get_output(0), right, trt.ElementWiseOperation.PROD + ).get_output(0) + update = _linear(network, update, weights, "pre_block.mlp.w2") + return network.add_elementwise(hidden, update, trt.ElementWiseOperation.SUM).get_output(0) + + +@op.cleanup_failed_build +def build_ref2va_audio_encoder_engine( + weights: dict, + profile: Ref2VAAudioEncoderProfile = DEFAULT_REF2VA_AUDIO_ENCODER_PROFILE, + *, + verbose: bool = False, + consume_weights: bool = False, + workspace_bytes: int | None = None, + weight_streaming: bool = False, + output_path: str | Path | None = None, +) -> bytes | dict[str, int | str]: + """Build the released stereo soundtrack posterior-mean encoder.""" + + profile.validate() + expected = set(checkpoint_keys(profile)) + missing = sorted(expected - set(weights)) + unexpected = sorted(set(weights) - expected) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 Ref2VA audio encoder checkpoint partition mismatch: " + f"missing={missing[:8]}, unexpected={unexpected[:8]}" + ) + logger = trt.Logger(trt.Logger.VERBOSE if verbose else trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + config = builder.create_builder_config() + config.builder_optimization_level = 1 + op.configure_builder(config, weight_streaming=weight_streaming) + op.configure_workspace( + config, + workspace_bytes, + default_bytes=AUDIO_VAE_ENCODER_DEFAULT_WORKSPACE_BYTES, + ) + abi = ref2va_audio_encoder_abi(profile) + samples = network.add_input( + abi.inputs[0].name, + trt.float32, + (profile.batch_size, 1, -1), + ) + optimization = builder.create_optimization_profile() + optimization.set_shape( + "audio_samples", + min=(profile.batch_size, 1, profile.min_samples), + opt=(profile.batch_size, 1, profile.opt_samples), + max=(profile.batch_size, 1, profile.max_samples), + ) + if config.add_optimization_profile(optimization) != 0: + raise RuntimeError("TensorRT rejected MiniMax-H3 Ref2VA audio profile") + expand = network.add_shuffle(samples) + expand.reshape_dims = (profile.batch_size, 1, 1, -1) + hidden = expand.get_output(0) + owned_weights: list[np.ndarray] = [] + + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight( + weights, + "encoder.block.0", + (profile.encoder_dim, 1, 7), + ), + _require_array(weights, "encoder.block.0.bias", (profile.encoder_dim,)), + padding=3, + name="encoder.block.0", + owned_weights=owned_weights, + ) + channels = profile.encoder_dim + for stage, stride in enumerate(profile.encoder_rates, start=1): + prefix = f"encoder.block.{stage}.block" + for residual, dilation in enumerate((1, 3, 9)): + hidden = _residual_unit( + network, + hidden, + weights, + f"{prefix}.{residual}.block", + channels=channels, + dilation=dilation, + owned_weights=owned_weights, + ) + hidden = _snake(network, hidden, weights, f"{prefix}.3") + down = f"{prefix}.4" + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight( + weights, + down, + (2 * channels, channels, 2 * stride), + ), + _require_array(weights, f"{down}.bias", (2 * channels,)), + stride=stride, + padding=math.ceil(stride / 2), + name=down, + owned_weights=owned_weights, + ) + channels *= 2 + hidden = _snake( + network, + hidden, + weights, + f"encoder.block.{len(profile.encoder_rates) + 1}", + ) + final = f"encoder.block.{len(profile.encoder_rates) + 2}" + hidden = _conv1d( + network, + hidden, + _fold_checkpoint_weight(weights, final, (profile.latent_dim, channels, 3)), + _require_array(weights, f"{final}.bias", (profile.latent_dim,)), + padding=1, + name=final, + owned_weights=owned_weights, + ) + rows = network.add_shuffle(hidden) + rows.first_transpose = trt.Permutation((0, 3, 1, 2)) + rows.reshape_dims = (profile.batch_size, -1, profile.latent_dim) + hidden = _causal_attention_projection(network, rows.get_output(0), weights, profile) + posterior_mean = _linear(network, hidden, weights, "mean_proj") + output = network.add_shuffle(posterior_mean) + output.first_transpose = trt.Permutation((0, 2, 1)) + value = output.get_output(0) + value.name = abi.outputs[0].name + network.mark_output(value) + op.validate_native_network( + network, + expected_attentions=1, + label="Ref2VA audio VAE encoder", + ) + + plan = None + record = None + try: + if output_path is None: + plan = builder.build_serialized_network(network, config) + else: + record = trt_compat.build_serialized_network_to_file( + builder, network, config, output_path + ) + finally: + op.release_weight_buffers(network) + if consume_weights: + weights.clear() + if output_path is None and plan is None: + raise RuntimeError("TensorRT failed to build MiniMax-H3 Ref2VA audio encoder") + del network, config, builder, logger, owned_weights + gc.collect() + return record if record is not None else bytes(plan) diff --git a/families/minimax_h3/ref2va_bundle_contract.py b/families/minimax_h3/ref2va_bundle_contract.py new file mode 100644 index 0000000000..2281885b31 --- /dev/null +++ b/families/minimax_h3/ref2va_bundle_contract.py @@ -0,0 +1,318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fail-closed bundle/provenance contract for native MiniMax-H3 Ref2VA.""" + +from __future__ import annotations + +import math + +from .fl2va_contract import PlanAbi, TensorAbi +from .ref2va_checkpoint import ( + CHECKPOINT_REVISION, + COMPONENT_NAME, + MODEL_ID, + TransformerRefIdentity, +) +from .ref2va_contract import ( + MAX_AUDIOS, + MAX_IMAGES, + MAX_REFERENCES, + MAX_REFERENCE_DURATION_SECONDS, + MAX_TOTAL_AUDIO_DURATION_SECONDS, + MAX_TOTAL_VIDEO_DURATION_SECONDS, + MAX_VIDEOS, + MIN_REFERENCE_DURATION_SECONDS, + Ref2VADenoiserProfile, + ref2va_denoiser_abi, + ref2va_denoiser_profiles, + ref2va_first_block_cache_abis, +) +from .ref2va_qwen_contract import ref2va_shared_qwen_profile_metadata +from .quantized_checkpoint import QuantizedCheckpointIdentity, QUANTIZED_REF2VA_CHECKPOINT_IDENTITY + + +# The Qwen language/vision plans are intentionally not listed here: Ref2VA +# binds the same ``text_encoder_plan`` and ``vision_encoder_plan`` sections +# already shared by T2VA/FL2VA. No second Qwen weight copy is permitted. +REF2VA_PLAN_SECTIONS = ( + ("ref2va_denoiser", "ref2va_denoiser.plan", "ref2va_denoiser_plan"), + ( + "ref2va_adaln_precompute", + "ref2va_adaln_precompute.plan", + "ref2va_adaln_precompute_plan", + ), + ( + "ref2va_video_vae_encoder", + "ref2va_video_vae_encoder.plan", + "ref2va_video_vae_encoder_plan", + ), + ( + "ref2va_audio_vae_encoder", + "ref2va_audio_vae_encoder.plan", + "ref2va_audio_vae_encoder_plan", + ), +) + +REF2VA_FIRST_BLOCK_CACHE_SECTIONS = ( + ("ref2va_dit_head", "ref2va_dit_head.plan", "ref2va_dit_head_plan"), + ("ref2va_dit_tail", "ref2va_dit_tail.plan", "ref2va_dit_tail_plan"), + ("ref2va_dit_finish", "ref2va_dit_finish.plan", "ref2va_dit_finish_plan"), + *REF2VA_PLAN_SECTIONS[1:], +) + + +def ref2va_plan_sections(first_block_cache: bool) -> tuple[tuple[str, str, str], ...]: + return REF2VA_FIRST_BLOCK_CACHE_SECTIONS if first_block_cache else REF2VA_PLAN_SECTIONS + + +REF2VA_SHARED_SECTIONS = { + "text_encoder": "text_encoder_plan", + "vision_encoder": "vision_encoder_plan", + "image_vae_encoder": "fl2va_keyframe_vae_encoder_plan", + "video_vae_decoder": "vae_tile_decoder_plan", + "audio_vae_decoder": "audio_vae_decoder_plan", +} + +# The public transformer_ref partition is the guidance-distilled 50-point +# rectified-flow recipe. ``num_inference_steps`` counts the terminal zero, so +# the transformer executes once for each of the first 49 grid points. Keep +# this mode-local instead of inheriting the base T2VA/FL2VA transformer schedule. +REF2VA_SCHEDULER_GRID_POINTS = 50 +REF2VA_TRANSFORMER_FORWARDS = 49 +REF2VA_VIDEO_FLOW_SHIFT = 12.0 +REF2VA_AUDIO_FLOW_SHIFT = 3.0 +REF2VA_GUIDANCE_SCALE = 1.0 + + +def _ref2va_adaln_abi() -> PlanAbi: + outputs = tuple( + TensorAbi( + f"block_modulation_{index}", + "bfloat16", + (12, 6, 5_376), + (12, 6, 5_376), + (12, 6, 5_376), + ) + for index in range(50) + ) + ( + TensorAbi( + "final_modulation", + "bfloat16", + (4, 2, 5_376), + (4, 2, 5_376), + (4, 2, 5_376), + ), + ) + return PlanAbi( + filename="ref2va_adaln_precompute.plan", + inputs=( + TensorAbi( + "timestep_features", + "float32", + (4, 256), + (4, 256), + (4, 256), + ), + ), + outputs=outputs, + ) + + +def _ref2va_video_encoder_abi() -> PlanAbi: + return PlanAbi( + filename="ref2va_video_vae_encoder.plan", + inputs=( + TensorAbi( + "pixel_tile_clip", + "float32", + (1, 3, 17, 256, 256), + (1, 3, 17, 256, 256), + (1, 3, 17, 256, 256), + ), + ), + outputs=( + TensorAbi( + "posterior_parameter_tile_clip", + "float32", + (1, 48, 5, 16, 16), + (1, 48, 5, 16, 16), + (1, 48, 5, 16, 16), + ), + ), + ) + + +def _ref2va_audio_encoder_abi() -> PlanAbi: + return PlanAbi( + filename="ref2va_audio_vae_encoder.plan", + inputs=( + TensorAbi( + "audio_samples", + "float32", + (2, 1, 64_000), + (2, 1, 165_600), + (2, 1, 480_000), + ), + ), + outputs=( + TensorAbi( + "posterior_mean", + "float32", + (2, 32, 80), + (2, 32, 207), + (2, 32, 600), + ), + ), + ) + + +def ref2va_plan_abi_metadata( + profile: Ref2VADenoiserProfile = Ref2VADenoiserProfile(), + *, + first_block_cache: bool = False, +) -> dict[str, object]: + """Serialize every dedicated plan binding without importing TensorRT.""" + + profile.validate() + + def tensor(binding: TensorAbi) -> dict[str, object]: + return { + "name": binding.name, + "dtype": binding.dtype, + "min_shape": list(binding.min_shape), + "opt_shape": list(binding.opt_shape), + "max_shape": list(binding.max_shape), + } + + denoiser = ( + tuple((f"{name}_plan", abi) for name, abi in ref2va_first_block_cache_abis(profile).items()) + if first_block_cache + else (("ref2va_denoiser_plan", ref2va_denoiser_abi(profile)),) + ) + plans = ( + *denoiser, + ("ref2va_adaln_precompute_plan", _ref2va_adaln_abi()), + ("ref2va_video_vae_encoder_plan", _ref2va_video_encoder_abi()), + ("ref2va_audio_vae_encoder_plan", _ref2va_audio_encoder_abi()), + ) + return { + section: { + "filename": abi.filename, + "inputs": [tensor(binding) for binding in abi.inputs], + "outputs": [tensor(binding) for binding in abi.outputs], + } + for section, abi in plans + } + + +def _capacity(profile: Ref2VADenoiserProfile) -> dict[str, list[int]]: + return { + "video_rows": [ + profile.min_video_rows, + profile.opt_video_rows, + profile.max_video_rows, + ], + "audio_rows": [ + profile.min_audio_rows, + profile.opt_audio_rows, + profile.max_audio_rows, + ], + "text_rows": [ + profile.min_text_rows, + profile.opt_text_rows, + profile.max_text_rows, + ], + "packed_rows": [ + profile.min_packed_rows, + profile.opt_packed_rows, + profile.max_packed_rows, + ], + } + + +def ref2va_bundle_metadata( + transformer_ref: TransformerRefIdentity | QuantizedCheckpointIdentity, + profile: Ref2VADenoiserProfile = Ref2VADenoiserProfile(), + *, + first_block_cache: bool = False, + first_block_cache_threshold: float = 0.08, +) -> dict[str, object]: + """Produce path-free metadata only after strict checkpoint validation.""" + + profile.validate() + if not isinstance(first_block_cache, bool): + raise ValueError("MiniMax-H3 ref2va_first_block_cache must be a boolean") + if ( + not isinstance(first_block_cache_threshold, (int, float)) + or isinstance(first_block_cache_threshold, bool) + or not math.isfinite(first_block_cache_threshold) + or first_block_cache_threshold < 0.0 + ): + raise ValueError("MiniMax-H3 Ref2VA cache threshold must be finite and nonnegative") + if isinstance(transformer_ref, QuantizedCheckpointIdentity): + if transformer_ref.bundle_metadata() != QUANTIZED_REF2VA_CHECKPOINT_IDENTITY.bundle_metadata(): + raise ValueError("MiniMax-H3 Ref2VA requires the distinct full Comfy INT8 checkpoint") + elif not isinstance(transformer_ref, TransformerRefIdentity): + raise TypeError("MiniMax-H3 Ref2VA metadata requires validated transformer_ref identity") + elif ( + transformer_ref.model_id != MODEL_ID + or transformer_ref.revision != CHECKPOINT_REVISION + or transformer_ref.component != COMPONENT_NAME + or transformer_ref.tensor_count != 638 + ): + raise ValueError("MiniMax-H3 Ref2VA transformer_ref provenance is incompatible") + denoiser_profiles = ref2va_denoiser_profiles(profile) + denoiser_profile_names = ( + ("five_second_common", "public_dynamic") + if len(denoiser_profiles) == 2 + else ("public_dynamic",) + ) + return { + "ref2va_schema_version": 5 if first_block_cache else 4, + "ref2va_supported": True, + "ref2va_first_block_cache": { + "enabled": first_block_cache, + "threshold": float(first_block_cache_threshold), + }, + "ref2va_scheduler": { + "sigma_grid_points": REF2VA_SCHEDULER_GRID_POINTS, + "transformer_forwards": REF2VA_TRANSFORMER_FORWARDS, + "video_shift": REF2VA_VIDEO_FLOW_SHIFT, + "audio_shift": REF2VA_AUDIO_FLOW_SHIFT, + "guidance_scale": REF2VA_GUIDANCE_SCALE, + "guidance_distilled": True, + }, + "ref2va_transformer_ref": transformer_ref.bundle_metadata(), + "ref2va_plan_sections": { + component: section + for component, _filename, section in ref2va_plan_sections(first_block_cache) + }, + "ref2va_plan_abis": ref2va_plan_abi_metadata(profile, first_block_cache=first_block_cache), + "ref2va_denoiser_profile_count": len(denoiser_profiles), + "ref2va_denoiser_profile_layout": "_then_".join(denoiser_profile_names), + "ref2va_denoiser_profiles": [ + {"name": name, **_capacity(capacity)} + for name, capacity in zip( + denoiser_profile_names, + denoiser_profiles, + strict=True, + ) + ], + "ref2va_shared_sections": dict(REF2VA_SHARED_SECTIONS), + "ref2va_shared_qwen_profiles": ref2va_shared_qwen_profile_metadata(), + "ref2va_limits": { + "max_images": MAX_IMAGES, + "max_videos": MAX_VIDEOS, + "max_explicit_audios": MAX_AUDIOS, + "max_reference_files": MAX_REFERENCES, + "min_seconds_each_video_or_audio": MIN_REFERENCE_DURATION_SECONDS, + "max_seconds_each_video_or_audio": MAX_REFERENCE_DURATION_SECONDS, + "max_total_video_seconds": MAX_TOTAL_VIDEO_DURATION_SECONDS, + "max_total_video_soundtrack_seconds": MAX_TOTAL_VIDEO_DURATION_SECONDS, + "max_total_explicit_audio_seconds": MAX_TOTAL_AUDIO_DURATION_SECONDS, + "audio_can_be_sole_input": True, + "video_soundtrack_stays_attached": True, + }, + "ref2va_capacity": _capacity(profile), + } diff --git a/families/minimax_h3/ref2va_checkpoint.py b/families/minimax_h3/ref2va_checkpoint.py new file mode 100644 index 0000000000..cba5e81ccf --- /dev/null +++ b/families/minimax_h3/ref2va_checkpoint.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Strict provenance and partitioning for the public H3 ``transformer_ref``. + +Ref2VA is not a mode bit on the T2VA transformer. The released repository has +a second 66.28 GB checkpoint partition with the same architecture and tensor +names but different values. This module deliberately refuses to build from +``transformer/`` or to fall back when ``transformer_ref/`` is absent. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path + + +MODEL_ID = "MiniMaxAI/MiniMax-H3" +CHECKPOINT_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc" +COMPONENT_NAME = "transformer_ref" +TOTAL_TENSOR_BYTES = 66_280_430_080 +CONFIG_BYTES = 546 +INDEX_BYTES = 64_488 + +SHARDS: tuple[tuple[str, int], ...] = ( + ("diffusion_pytorch_model-00001-of-00014.safetensors", 4_825_958_704), + ("diffusion_pytorch_model-00002-of-00014.safetensors", 4_702_158_032), + ("diffusion_pytorch_model-00003-of-00014.safetensors", 4_933_368_192), + ("diffusion_pytorch_model-00004-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00005-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00006-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00007-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00008-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00009-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00010-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00011-of-00014.safetensors", 4_702_158_080), + ("diffusion_pytorch_model-00012-of-00014.safetensors", 4_933_368_232), + ("diffusion_pytorch_model-00013-of-00014.safetensors", 4_567_069_608), + ("diffusion_pytorch_model-00014-of-00014.safetensors", 4_644_161_920), +) + +EXPECTED_CONFIG = { + "_class_name": "MiniMaxH3Transformer3DModel", + "_diffusers_version": "0.36.0.dev0", + "num_attention_heads": 56, + "attention_head_dim": 128, + "hidden_size": 5376, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "in_channels": 24, + "audio_in_channels": 32, + "patch_size": [1, 2, 2], + "text_dim": 5120, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "rope_freq_dim": 16, + "rope_theta": 10_000.0, + "norm_eps": 1.0e-5, + "qk_norm_eps": 1.0e-5, + "final_norm_eps": 1.0e-5, +} + + +def _dit_checkpoint_keys() -> tuple[str, ...]: + names = [ + "proj_in.weight", + "proj_in.bias", + "audio_proj_in.weight", + "audio_proj_in.bias", + "context_embedder.weight", + "context_embedder.bias", + "token_refiner.final_norm.weight", + ] + for index in range(2): + prefix = f"token_refiner.refiner_blocks.{index}" + names.extend( + ( + f"{prefix}.norm1.weight", + f"{prefix}.norm2.weight", + *(f"{prefix}.attn.to_{name}.weight" for name in ("q", "k", "v")), + f"{prefix}.attn.norm_q.weight", + f"{prefix}.attn.norm_k.weight", + f"{prefix}.attn.to_out.0.weight", + f"{prefix}.ff.net.0.proj.weight", + f"{prefix}.ff.net.2.weight", + ) + ) + for index in range(50): + prefix = f"transformer_blocks.{index}" + names.extend( + ( + f"{prefix}.norm1.weight", + f"{prefix}.norm2.weight", + *(f"{prefix}.attn.to_{name}.weight" for name in ("q", "k", "v")), + f"{prefix}.attn.norm_q.weight", + f"{prefix}.attn.norm_k.weight", + f"{prefix}.attn.to_out.0.weight", + f"{prefix}.ff.net.0.proj.weight", + f"{prefix}.ff.net.2.weight", + ) + ) + names.extend( + ( + "norm_out.norm.weight", + "proj_out.weight", + "proj_out.bias", + "audio_proj_out.weight", + "audio_proj_out.bias", + ) + ) + return tuple(names) + + +def _adaln_checkpoint_keys() -> tuple[str, ...]: + names = [ + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + ] + for index in range(50): + prefix = f"transformer_blocks.{index}.adaln_proj.linear" + names.extend((f"{prefix}.weight", f"{prefix}.bias")) + names.extend(("norm_out.linear.weight", "norm_out.linear.bias")) + return tuple(names) + + +REF2VA_DENOISER_KEYS = _dit_checkpoint_keys() +REF2VA_ADALN_KEYS = _adaln_checkpoint_keys() +REF2VA_ALL_KEYS = REF2VA_DENOISER_KEYS + REF2VA_ADALN_KEYS +REF2VA_FINISH_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name.startswith(("norm_out.", "proj_out.", "audio_proj_out.")) +) +REF2VA_TAIL_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name.startswith("transformer_blocks.") and not name.startswith("transformer_blocks.0.") +) +REF2VA_HEAD_KEYS = tuple( + name + for name in REF2VA_DENOISER_KEYS + if name not in REF2VA_FINISH_KEYS and name not in REF2VA_TAIL_KEYS +) + + +def download_command(local_dir: str = "MiniMax-H3") -> str: + """Return the build-time-only command for the missing public partition.""" + + if not isinstance(local_dir, str) or not local_dir.strip(): + raise ValueError("MiniMax-H3 checkpoint destination must be a non-empty string") + return ( + f"hf download {MODEL_ID} --revision {CHECKPOINT_REVISION} " + f'--include "transformer_ref/*" --local-dir "{local_dir}"' + ) + + +def _require_file(path: Path, *, size: int) -> dict[str, int]: + if not path.is_file() or path.is_symlink(): + raise FileNotFoundError(f"MiniMax-H3 Ref2VA checkpoint file is missing: {path.name}") + actual_size = path.stat().st_size + if actual_size != size: + raise ValueError( + f"MiniMax-H3 Ref2VA checkpoint size mismatch for {path.name}: " + f"expected={size}, actual={actual_size}" + ) + return {"bytes": size} + + +@dataclass(frozen=True) +class TransformerRefIdentity: + model_id: str + revision: str + component: str + tensor_bytes: int + tensor_count: int + files: dict[str, dict[str, int | str]] + + def bundle_metadata(self) -> dict[str, object]: + """Path-free provenance safe to persist in the native bundle.""" + + return { + "schema_version": 1, + "model_id": self.model_id, + "revision": self.revision, + "component": self.component, + "tensor_bytes": self.tensor_bytes, + "tensor_count": self.tensor_count, + "files": self.files, + "runtime_framework": None, + } + + +def validate_transformer_ref_checkpoint(component_dir: str | Path) -> TransformerRefIdentity: + """Validate the released ``transformer_ref`` layout and plan partition. + + The pinned config and index are hashed; large shards are checked against + the exact index names and released sizes without rereading 66 GiB solely + for an integrity pass. + """ + root = Path(component_dir) + if root.name != COMPONENT_NAME: + candidate = root / COMPONENT_NAME + if candidate.is_dir(): + root = candidate + else: + raise FileNotFoundError( + "MiniMax-H3 Ref2VA requires the distinct transformer_ref checkpoint. " + "The T2VA transformer is not a valid fallback. Download it with: " + f"{download_command(str(root))}" + ) + if not root.is_dir() or root.is_symlink(): + raise FileNotFoundError( + "MiniMax-H3 transformer_ref is absent; no T2VA fallback is permitted. " + f"Download it with: {download_command(str(root.parent))}" + ) + + files: dict[str, dict[str, int | str]] = {} + config_path = root / "config.json" + files["config.json"] = _require_file( + config_path, + size=CONFIG_BYTES, + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + if config != EXPECTED_CONFIG: + raise ValueError( + "MiniMax-H3 transformer_ref config does not match the released architecture" + ) + + index_name = "diffusion_pytorch_model.safetensors.index.json" + index_path = root / index_name + files[index_name] = _require_file( + index_path, + size=INDEX_BYTES, + ) + index = json.loads(index_path.read_text(encoding="utf-8")) + if not isinstance(index, dict) or set(index) != {"metadata", "weight_map"}: + raise ValueError("MiniMax-H3 transformer_ref index has an invalid schema") + if index["metadata"] != {"total_size": TOTAL_TENSOR_BYTES}: + raise ValueError("MiniMax-H3 transformer_ref index total_size is invalid") + weight_map = index["weight_map"] + if not isinstance(weight_map, dict): + raise ValueError("MiniMax-H3 transformer_ref weight_map is invalid") + expected_keys = set(REF2VA_ALL_KEYS) + if len(REF2VA_ALL_KEYS) != 638 or len(expected_keys) != 638: + raise RuntimeError("MiniMax-H3 transformer_ref partition is not exhaustive") + actual_keys = set(weight_map) + if actual_keys != expected_keys: + raise ValueError( + "MiniMax-H3 transformer_ref tensor partition mismatch: " + f"missing={sorted(expected_keys - actual_keys)[:8]}, " + f"unexpected={sorted(actual_keys - expected_keys)[:8]}" + ) + expected_shards = {name for name, _size in SHARDS} + if set(weight_map.values()) != expected_shards: + raise ValueError("MiniMax-H3 transformer_ref index does not reference the exact 14 shards") + + for name, size in SHARDS: + files[name] = _require_file( + root / name, + size=size, + ) + allowed = {"config.json", index_name, *expected_shards} + unexpected_files = sorted( + path.name for path in root.iterdir() if path.is_file() and path.name not in allowed + ) + if unexpected_files: + raise ValueError( + f"MiniMax-H3 transformer_ref contains unexpected files: {unexpected_files}" + ) + return TransformerRefIdentity( + model_id=MODEL_ID, + revision=CHECKPOINT_REVISION, + component=COMPONENT_NAME, + tensor_bytes=TOTAL_TENSOR_BYTES, + tensor_count=len(weight_map), + files=files, + ) diff --git a/families/minimax_h3/ref2va_contract.py b/families/minimax_h3/ref2va_contract.py new file mode 100644 index 0000000000..b988226a59 --- /dev/null +++ b/families/minimax_h3/ref2va_contract.py @@ -0,0 +1,1114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native-runtime contracts for the public MiniMax-H3 Ref2VA workflow. + +This module contains no framework runtime. It is the executable specification +used by the bundle builder and by the C++ port for request validation, media +geometry, Qwen3-VL presentation/MRoPE, packed Omni-DiT rows, and per-row noise +levels. Context-IR and Regenerate-2K are intentionally outside this contract. + +The Python is build/test tooling only. A serialized ModelConnect bundle must +implement the same operations in C++/CUDA and TensorRT-RTX; importing this file +at inference time is neither required nor supported. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +import math +from typing import Literal + +import numpy as np + +from .fl2va_contract import PlanAbi, TensorAbi, resolve_canvas_size, validate_canvas + + +ReferenceKind = Literal["image", "video", "audio"] +QwenModality = Literal["text", "image", "video"] + +MAX_IMAGES = 9 +MAX_VIDEOS = 3 +MAX_AUDIOS = 3 +MAX_REFERENCES = 12 +MIN_REFERENCE_DURATION_SECONDS = 2.0 +MAX_REFERENCE_DURATION_SECONDS = 15.0 +MAX_TOTAL_VIDEO_DURATION_SECONDS = 15.0 +MAX_TOTAL_AUDIO_DURATION_SECONDS = 15.0 + +TARGET_FPS = 24.0 +AUDIO_HOP_LENGTH = 800 +AUDIO_CHANNELS = 2 +VIDEO_PATCH_HEIGHT = 2 +VIDEO_PATCH_WIDTH = 2 +VIDEO_FRAMES_PER_CHUNK = 17 +VIDEO_LATENTS_PER_CHUNK = 5 +REFERENCE_IMAGE_SHORT_EDGE = 2048 +CANVAS_MULTIPLE = 32 +QWEN_VISION_PATCH_SIZE = 16 +QWEN_VISION_MERGE_SIZE = 2 +QWEN_VISION_TEMPORAL_PATCH_SIZE = 2 +QWEN_VIDEO_SAMPLE_FPS = 2.0 +QWEN_MAX_POSITION_EMBEDDINGS = 262_144 + +QWEN_VISION_START_TOKEN_ID = 151_652 +QWEN_VISION_END_TOKEN_ID = 151_653 +QWEN_IMAGE_PAD_TOKEN_ID = 151_655 +QWEN_VIDEO_PAD_TOKEN_ID = 151_656 + +H3_VIDEO_TAG = 0 +H3_TEXT_TAG = 1 +H3_AUDIO_TAG = 2 +CONDITION_VIDEO_TIMESTEP = 0.999 +CONDITION_AUDIO_TIMESTEP = 1.0 + +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32.0 + + +@dataclass(frozen=True) +class ReferenceSpec: + """Decoded public reference metadata, in request order. + + ``duration_seconds`` belongs to video and explicit-audio files. A video + soundtrack remains attached to its video and therefore does not increment + the explicit-audio count or the twelve-file count. + """ + + kind: ReferenceKind + duration_seconds: float | None = None + width: int | None = None + height: int | None = None + has_audio: bool = False + + def validate(self) -> None: + if self.kind not in ("image", "video", "audio"): + raise ValueError(f"MiniMax-H3 reference kind is invalid: {self.kind!r}") + if not isinstance(self.has_audio, bool): + raise ValueError("MiniMax-H3 has_audio must be a boolean") + if self.has_audio and self.kind != "video": + raise ValueError("Only a MiniMax-H3 video reference can carry a soundtrack") + + if self.kind in ("image", "video"): + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (self.width, self.height) + ): + raise ValueError("MiniMax-H3 visual references need positive integer geometry") + assert self.width is not None and self.height is not None + if self.width > 4 * self.height or self.height > 4 * self.width: + raise ValueError( + "MiniMax-H3 visual references must be within the public 1:4..4:1 aspect range" + ) + elif self.width is not None or self.height is not None: + raise ValueError("MiniMax-H3 audio references cannot carry pixel geometry") + + if self.kind == "image": + if self.duration_seconds is not None: + raise ValueError("MiniMax-H3 image references do not have a duration") + return + duration = self.duration_seconds + if ( + isinstance(duration, bool) + or not isinstance(duration, (int, float)) + or not math.isfinite(duration) + or not MIN_REFERENCE_DURATION_SECONDS <= duration <= MAX_REFERENCE_DURATION_SECONDS + ): + raise ValueError("MiniMax-H3 video/audio references must each be 2..15 seconds long") + + +@dataclass(frozen=True) +class ReferenceRequestSummary: + image_count: int + video_count: int + audio_count: int + total_video_seconds: float + total_audio_seconds: float + audio_bearing_count: int + + +def validate_reference_request( + references: Sequence[ReferenceSpec], +) -> ReferenceRequestSummary: + """Fail closed on every public H3-Base Ref2VA file/duration limit.""" + + if not isinstance(references, Sequence) or isinstance(references, (str, bytes)): + raise ValueError("MiniMax-H3 references must be an ordered sequence") + if not references: + raise ValueError("MiniMax-H3 Ref2VA needs at least one reference") + if len(references) > MAX_REFERENCES: + raise ValueError( + f"MiniMax-H3 accepts at most {MAX_REFERENCES} reference files, got {len(references)}" + ) + for reference in references: + if not isinstance(reference, ReferenceSpec): + raise ValueError("MiniMax-H3 references must use ReferenceSpec metadata") + reference.validate() + + image_count = sum(reference.kind == "image" for reference in references) + video_count = sum(reference.kind == "video" for reference in references) + audio_count = sum(reference.kind == "audio" for reference in references) + if image_count > MAX_IMAGES: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_IMAGES} image references") + if video_count > MAX_VIDEOS: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_VIDEOS} video references") + if audio_count > MAX_AUDIOS: + raise ValueError(f"MiniMax-H3 accepts at most {MAX_AUDIOS} explicit audio references") + total_video = math.fsum( + float(reference.duration_seconds) for reference in references if reference.kind == "video" + ) + total_audio = math.fsum( + float(reference.duration_seconds) for reference in references if reference.kind == "audio" + ) + if total_video > MAX_TOTAL_VIDEO_DURATION_SECONDS: + raise ValueError("MiniMax-H3 reference-video duration total exceeds 15 seconds") + if total_audio > MAX_TOTAL_AUDIO_DURATION_SECONDS: + raise ValueError("MiniMax-H3 explicit-audio duration total exceeds 15 seconds") + return ReferenceRequestSummary( + image_count=image_count, + video_count=video_count, + audio_count=audio_count, + total_video_seconds=total_video, + total_audio_seconds=total_audio, + audio_bearing_count=audio_count + + sum(reference.kind == "video" and reference.has_audio for reference in references), + ) + + +def resolve_reference_image_size(width: int, height: int) -> tuple[int, int]: + """Return ``(height, width)`` at short-edge 2048 with no area cap. + + The round is Python/IEEE round-half-to-even, matching the released setup + block. The 1:4 and 4:1 endpoints resolve to 2048x8192 and therefore to the + Qwen image processor's exact 16,777,216-pixel ceiling. + """ + + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (width, height) + ): + raise ValueError("MiniMax-H3 reference image geometry must use positive integers") + if width > 4 * height or height > 4 * width: + raise ValueError("MiniMax-H3 reference image must be within 1:4 and 4:1") + scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + target_height = max(CANVAS_MULTIPLE, round(height * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE) + target_width = max(CANVAS_MULTIPLE, round(width * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE) + return target_height, target_width + + +def resolve_reference_video_size(width: int, height: int) -> tuple[int, int]: + """Put a video reference on the public 768p canvas of its own aspect.""" + + return resolve_canvas_size(width, height) + + +def video_resample_source_indices( + source_frames: int, + source_fps: float, + *, + target_frames: int, + target_fps: float = TARGET_FPS, +) -> tuple[int, ...]: + """Reproduce the reference's constant-frame-rate drop/duplicate pass. + + The returned source row for each output slot is applied before LANCZOS + spatial resize and is truncated to the generated frame count. + """ + + if ( + isinstance(source_frames, bool) + or not isinstance(source_frames, int) + or source_frames <= 0 + or isinstance(target_frames, bool) + or not isinstance(target_frames, int) + or target_frames <= 0 + ): + raise ValueError("MiniMax-H3 video frame counts must be positive integers") + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + for value in (source_fps, target_fps) + ): + raise ValueError("MiniMax-H3 video frame rates must be finite and positive") + if float(source_fps) == float(target_fps): + return tuple(range(min(source_frames, target_frames))) + + scale = float(target_fps) / float(source_fps) + slots = np.floor(np.arange(source_frames, dtype=np.float64) * scale + 0.5).astype(np.int64) + end = math.floor(source_frames * scale + 0.5) + counts = np.diff(slots, append=np.int64(end)) + if np.any(counts < 0): + raise RuntimeError("MiniMax-H3 constant-frame-rate slots are not monotonic") + result = np.repeat(np.arange(source_frames, dtype=np.int64), counts) + return tuple(int(value) for value in result[:target_frames]) + + +def snap_reference_video_frames_down(num_frames: int) -> int: + """Snap a normalized reference down to ``17*n+5`` before VAE encode.""" + + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 reference video needs a positive frame count") + chunks = max(1, (num_frames - VIDEO_LATENTS_PER_CHUNK) // VIDEO_FRAMES_PER_CHUNK) + return chunks * VIDEO_FRAMES_PER_CHUNK + VIDEO_LATENTS_PER_CHUNK + + +def reference_video_latent_frames(num_frames: int) -> int: + aligned = snap_reference_video_frames_down(num_frames) + return ( + (aligned - VIDEO_LATENTS_PER_CHUNK) // VIDEO_FRAMES_PER_CHUNK + ) * VIDEO_LATENTS_PER_CHUNK + 2 + + +@dataclass(frozen=True) +class ReferenceVideoEncodeSchedule: + """Native clip schedule matching ``AutoencoderKLMiniMaxH3._encode``.""" + + snapped_frames: int + clip_count: int + repeated_tail_frames: int + raw_posterior_frames: int + dropped_tail_latents: int + output_latent_frames: int + + +def reference_video_encode_schedule(num_frames: int) -> ReferenceVideoEncodeSchedule: + """Plan exact 17-frame invocations and the one global three-token drop. + + The final reference clip has five real frames and twelve copies of the + final frame. The VideoVAE produces five posterior frames per clip, then + the released wrapper drops three frames once after concatenating *all* + clips (never once per clip). + """ + + snapped = snap_reference_video_frames_down(num_frames) + if snapped > num_frames: + raise ValueError("MiniMax-H3 reference video is too short for the released 17*n+5 schedule") + clip_count = (snapped + VIDEO_FRAMES_PER_CHUNK - 1) // VIDEO_FRAMES_PER_CHUNK + repeated = clip_count * VIDEO_FRAMES_PER_CHUNK - snapped + raw = clip_count * VIDEO_LATENTS_PER_CHUNK + dropped = 3 + output = raw - dropped + if repeated != VIDEO_FRAMES_PER_CHUNK - VIDEO_LATENTS_PER_CHUNK: + raise RuntimeError("MiniMax-H3 reference video clip padding invariant failed") + if output != reference_video_latent_frames(num_frames): + raise RuntimeError("MiniMax-H3 reference video latent schedule mismatch") + return ReferenceVideoEncodeSchedule( + snapped_frames=snapped, + clip_count=clip_count, + repeated_tail_frames=repeated, + raw_posterior_frames=raw, + dropped_tail_latents=dropped, + output_latent_frames=output, + ) + + +def audio_latent_frames(num_samples: int) -> int: + if isinstance(num_samples, bool) or not isinstance(num_samples, int) or num_samples <= 0: + raise ValueError("MiniMax-H3 reference audio needs a positive sample count") + return (num_samples + AUDIO_HOP_LENGTH - 1) // AUDIO_HOP_LENGTH + + +def qwen_video_condition_sample( + num_frames: int, + *, + fps: float = TARGET_FPS, + sample_fps: float = QWEN_VIDEO_SAMPLE_FPS, + temporal_patch: int = QWEN_VISION_TEMPORAL_PATCH_SIZE, +) -> tuple[tuple[int, ...], tuple[float, ...]]: + """Return the sampled frame indices and one timestamp per merged pair.""" + + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError("MiniMax-H3 Qwen video input needs a positive frame count") + if ( + not math.isfinite(fps) + or not math.isfinite(sample_fps) + or fps <= 0 + or sample_fps <= 0 + or sample_fps > fps + ): + raise ValueError("MiniMax-H3 Qwen video sample rates are invalid") + if ( + isinstance(temporal_patch, bool) + or not isinstance(temporal_patch, int) + or temporal_patch <= 0 + ): + raise ValueError("MiniMax-H3 Qwen temporal patch must be a positive integer") + + stride = fps / sample_fps + indices: list[int] = [] + cursor = 0.0 + while round(cursor) < num_frames: + index = round(cursor) + if not indices or index > indices[-1]: + indices.append(index) + cursor += stride + if len(indices) < temporal_patch: + minimum = round((temporal_patch - 1) * stride) + 1 + raise ValueError( + f"MiniMax-H3 Qwen reads a reference video only when it has at least {minimum} frames" + ) + timestamps = [index / sample_fps for index in range(len(indices))] + timestamps.extend([timestamps[-1]] * (-len(timestamps) % temporal_patch)) + blocks = tuple( + (timestamps[index] + timestamps[index + temporal_patch - 1]) / 2.0 + for index in range(0, len(timestamps), temporal_patch) + ) + return tuple(indices), blocks + + +def qwen_patch_grid(height: int, width: int) -> tuple[int, int]: + validate_canvas(height, width) + if height % QWEN_VISION_PATCH_SIZE or width % QWEN_VISION_PATCH_SIZE: + raise ValueError("MiniMax-H3 Qwen vision input is not patch aligned") + return height // QWEN_VISION_PATCH_SIZE, width // QWEN_VISION_PATCH_SIZE + + +def qwen_merged_rows(height: int, width: int) -> int: + grid_h, grid_w = qwen_patch_grid(height, width) + merge = QWEN_VISION_MERGE_SIZE + if grid_h % merge or grid_w % merge: + raise ValueError("MiniMax-H3 Qwen vision input is not merge aligned") + return (grid_h // merge) * (grid_w // merge) + + +@dataclass(frozen=True) +class PresentationPiece: + """One literal text run or one Qwen visual-pad run.""" + + modality: QwenModality + text: str = "" + height: int = 0 + width: int = 0 + + @property + def vision_rows(self) -> int: + return 0 if self.modality == "text" else qwen_merged_rows(self.height, self.width) + + +def ref2va_presentation_blueprint( + prompt: str, + references: Sequence[ReferenceSpec], + *, + normalized_visual_sizes: Sequence[tuple[int, int]], + normalized_video_frames: Sequence[int], +) -> tuple[PresentationPiece, ...]: + """Build the ordered labels/timestamps/vision blocks before tokenization.""" + + if not isinstance(prompt, str): + raise ValueError("MiniMax-H3 prompt must be one string") + validate_reference_request(references) + if len(normalized_visual_sizes) != sum(ref.kind != "audio" for ref in references): + raise ValueError("MiniMax-H3 visual-size metadata count does not match references") + if len(normalized_video_frames) != sum(ref.kind == "video" for ref in references): + raise ValueError("MiniMax-H3 video-frame metadata count does not match references") + visual_sizes = iter(normalized_visual_sizes) + video_frames = iter(normalized_video_frames) + pieces: list[PresentationPiece] = [] + counts = {"image": 0, "video": 0, "audio": 0} + for reference in references: + if reference.has_audio or reference.kind == "audio": + counts["audio"] += 1 + pieces.append(PresentationPiece("text", f"