diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4351b31..96e71cd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev + apt-get install -y --no-install-recommends ca-certificates clang clang-tools cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 @@ -97,7 +97,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev + apt-get install -y --no-install-recommends ca-certificates clang clang-tools cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 diff --git a/CMakeLists.txt b/CMakeLists.txt index 05e7abf..0a6e04c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,6 +44,10 @@ set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY "${DIN_RUNTIME_DIR}") list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") +# Load optional local URL overrides when present. +if (EXISTS "${CMAKE_CURRENT_LIST_DIR}/cmake/internal_urls.cmake") + include("${CMAKE_CURRENT_LIST_DIR}/cmake/internal_urls.cmake") +endif () include(onnxruntime) include(nlohmann_json) include(argparse) diff --git a/CMakePresets.json b/CMakePresets.json index 584a048..eef10b4 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -46,7 +46,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "AMD64", "ONNXRUNTIME_ARCH": "x64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-amd64-cuda-13.4-Release-external.zip" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-windows-amd64-cuda-13.4-release-external.zip" }, "vendor": { "jetbrains.com/clion": { @@ -64,7 +64,7 @@ "CMAKE_SYSTEM_NAME": "Windows", "CMAKE_SYSTEM_PROCESSOR": "ARM64", "ONNXRUNTIME_ARCH": "arm64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-windows-arm64-cuda-13.4-release-external.zip" }, "vendor": { "jetbrains.com/clion": { @@ -81,7 +81,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "x86_64", "ONNXRUNTIME_ARCH": "x64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Linux-x86_64-cuda-13.4-Release-external.tar.zst" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-linux-x86_64-cuda-13.4-release-external.tar.zst" }, "vendor": { "jetbrains.com/clion": { @@ -98,7 +98,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "aarch64", "ONNXRUNTIME_ARCH": "arm64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Linux-aarch64-cuda-13.4-Release-external.tar.zst" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-linux-aarch64-cuda-13.4-release-external.tar.zst" }, "vendor": { "jetbrains.com/clion": { diff --git a/cmake/imgui.cmake b/cmake/imgui.cmake new file mode 100644 index 0000000..599e65a --- /dev/null +++ b/cmake/imgui.cmake @@ -0,0 +1,41 @@ +include_guard(GLOBAL) +include(FetchContent) + +# Attach the shared GLFW/OpenGL ImGui implementation to an existing UI target. +# Dependencies are fetched only when a UI target is requested. +function(din_link_imgui TARGET) + if (NOT TARGET din_imgui) + find_package(OpenGL REQUIRED GLOBAL) + find_package(Threads REQUIRED GLOBAL) + set(GLFW_BUILD_DOCS OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_TESTS OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) + set(GLFW_INSTALL OFF CACHE BOOL "" FORCE) + if (UNIX AND NOT APPLE) + set(GLFW_BUILD_WAYLAND OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_X11 ON CACHE BOOL "" FORCE) + endif () + FetchContent_Declare(glfw + GIT_REPOSITORY https://github.com/glfw/glfw.git GIT_TAG 3.4 GIT_SHALLOW TRUE) + FetchContent_Declare(imgui + GIT_REPOSITORY https://github.com/ocornut/imgui.git GIT_TAG v1.91.5 GIT_SHALLOW TRUE) + FetchContent_MakeAvailable(glfw imgui) + add_library(din_imgui STATIC + ${imgui_SOURCE_DIR}/imgui.cpp + ${imgui_SOURCE_DIR}/imgui_draw.cpp + ${imgui_SOURCE_DIR}/imgui_tables.cpp + ${imgui_SOURCE_DIR}/imgui_widgets.cpp + ${imgui_SOURCE_DIR}/backends/imgui_impl_glfw.cpp + ${imgui_SOURCE_DIR}/backends/imgui_impl_opengl3.cpp) + target_include_directories(din_imgui PUBLIC ${imgui_SOURCE_DIR} ${imgui_SOURCE_DIR}/backends) + target_link_libraries(din_imgui PUBLIC glfw OpenGL::GL Threads::Threads) + endif () + target_link_libraries(${TARGET} PRIVATE din_imgui) +endfunction() + +# Usage: add_din_ui_executable(my_app main.cpp [other sources...]) +# App-specific libraries and settings belong in the caller's CMakeLists.txt. +function(add_din_ui_executable TARGET) + add_din_executable(${TARGET} ${ARGN}) + din_link_imgui(${TARGET}) +endfunction() diff --git a/cmake/miniaudio.cmake b/cmake/miniaudio.cmake index 398643a..0c2cf2e 100644 --- a/cmake/miniaudio.cmake +++ b/cmake/miniaudio.cmake @@ -3,7 +3,7 @@ include_guard(GLOBAL) include(FetchContent) FetchContent_Declare(din_miniaudio GIT_REPOSITORY https://github.com/mackron/miniaudio.git - GIT_TAG f40cf03f80cdb7e741d43e53b7e706e8c1394bcf + GIT_TAG "0.11.25" SOURCE_SUBDIR header-only) FetchContent_MakeAvailable(din_miniaudio) diff --git a/cmake/onnxruntime_trt_rtx_ep.cmake b/cmake/onnxruntime_trt_rtx_ep.cmake index ccd5109..aecef0e 100644 --- a/cmake/onnxruntime_trt_rtx_ep.cmake +++ b/cmake/onnxruntime_trt_rtx_ep.cmake @@ -9,7 +9,7 @@ set(ONNXRUNTIME_TRT_RTX_EP_GIT_REPOSITORY "https://github.com/NVIDIA/TensorRT-RTX-EP-ABI.git" CACHE STRING "Git repository for the ONNX Runtime TensorRT RTX Execution Provider") set(ONNXRUNTIME_TRT_RTX_EP_GIT_TAG - "v0.4.0" + "v0.5.0" CACHE STRING "Git tag, branch, or commit for the ONNX Runtime TensorRT RTX Execution Provider") set(ONNXRUNTIME_TRT_RTX_EP_GIT_HTTP_TOKEN "" @@ -155,14 +155,53 @@ if(_onnxruntime_trt_rtx_ep_runtime_dependencies) endif() +# Resolve at build time, after ExternalProject has produced the library. +# Keep the compatibility logic inline; generate its executable script in the build tree. +set(_onnxruntime_trt_rtx_ep_copy_script "${CMAKE_CURRENT_BINARY_DIR}/copy_onnxruntime_trt_rtx_ep.cmake") +file(GENERATE OUTPUT "${_onnxruntime_trt_rtx_ep_copy_script}" CONTENT [=[ +if(EP_MULTI_CONFIG) + set(_ep_candidates + "${EP_BINARY_DIR}/bin/${EP_CONFIGURATION}/${EP_LIBRARY_NAME}" + "${EP_BINARY_DIR}/${EP_CONFIGURATION}/${EP_LIBRARY_NAME}" + ) +else() + set(_ep_candidates + "${EP_BINARY_DIR}/bin/${EP_LIBRARY_NAME}" + "${EP_BINARY_DIR}/${EP_LIBRARY_NAME}" + ) +endif() + +foreach(_ep_candidate IN LISTS _ep_candidates) + if(EXISTS "${_ep_candidate}") + file(MAKE_DIRECTORY "${EP_RUNTIME_DIR}") + execute_process( + COMMAND "${CMAKE_COMMAND}" -E copy_if_different + "${_ep_candidate}" "${EP_RUNTIME_DIR}" + RESULT_VARIABLE _ep_copy_result + ERROR_VARIABLE _ep_copy_error + ) + if(NOT _ep_copy_result STREQUAL "0") + message(FATAL_ERROR "Failed to copy TensorRT RTX EP: ${_ep_copy_error}") + endif() + return() + endif() +endforeach() + +message(FATAL_ERROR "TensorRT RTX EP library was not found. Checked: ${_ep_candidates}") +]=]) + set(_onnxruntime_trt_rtx_ep_runtime_commands COMMAND "${CMAKE_COMMAND}" -E make_directory "${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" ) list(APPEND _onnxruntime_trt_rtx_ep_runtime_commands - COMMAND "${CMAKE_COMMAND}" -E copy_if_different - "${_onnxruntime_trt_rtx_ep_binary_dir}/$/${_onnxruntime_trt_rtx_ep_runtime_library_name}" - "${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" + COMMAND "${CMAKE_COMMAND}" + "-DEP_BINARY_DIR=${_onnxruntime_trt_rtx_ep_binary_dir}" + "-DEP_CONFIGURATION=$" + "-DEP_MULTI_CONFIG=$" + "-DEP_LIBRARY_NAME=${_onnxruntime_trt_rtx_ep_runtime_library_name}" + "-DEP_RUNTIME_DIR=${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" + -P "${_onnxruntime_trt_rtx_ep_copy_script}" ) foreach(_onnxruntime_trt_rtx_ep_runtime_dependency IN LISTS _onnxruntime_trt_rtx_ep_runtime_dependencies) diff --git a/cmake/random123.cmake b/cmake/random123.cmake new file mode 100644 index 0000000..31d4966 --- /dev/null +++ b/cmake/random123.cmake @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +include_guard(GLOBAL) +include(FetchContent) +set(FETCHCONTENT_UPDATES_DISCONNECTED ON) + +FetchContent_Declare(random123 + GIT_REPOSITORY https://github.com/DEShawResearch/random123.git + GIT_TAG v1.14.0 + GIT_SHALLOW TRUE +) +FetchContent_MakeAvailable(random123) + +# Random123 v1.14.0's MSVC feature header rejects ARM64 and otherwise selects +# x86 intrinsics for newer MSVC versions. Philox4x32 uses the scalar path. +if(MSVC AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(ARM64|arm64)$") + set(_din_r123_msvc_features "${random123_SOURCE_DIR}/include/Random123/features/msvcfeatures.h") + file(READ "${_din_r123_msvc_features}" _din_r123_features) + set(_din_r123_old_guard "#if !defined(_M_IX86) && !defined(_M_X64)") + set(_din_r123_new_guard "#if !defined(_M_IX86) && !defined(_M_X64) && !defined(_M_ARM64)") + if(NOT _din_r123_features MATCHES "_M_ARM64") + if(NOT _din_r123_features MATCHES "#if !defined\\(_M_IX86\\) && !defined\\(_M_X64\\)") + message(FATAL_ERROR "Random123 MSVC architecture guard changed; review the ARM64 patch") + endif() + string(REPLACE "${_din_r123_old_guard}" "${_din_r123_new_guard}" + _din_r123_features "${_din_r123_features}") + file(WRITE "${_din_r123_msvc_features}" "${_din_r123_features}") + endif() + unset(_din_r123_msvc_features) + unset(_din_r123_features) + unset(_din_r123_old_guard) + unset(_din_r123_new_guard) +endif() + +add_library(din_random123 INTERFACE) +add_library(Random123::Random123 ALIAS din_random123) +target_include_directories(din_random123 SYSTEM INTERFACE "${random123_SOURCE_DIR}/include") +if(MSVC AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(ARM64|arm64)$") + target_compile_definitions(din_random123 INTERFACE + R123_USE_SSE=0 + R123_USE_SSE4_1=0 + R123_USE_SSE4_2=0 + R123_USE_EMMINTRIN_H=0 + R123_USE_SMMINTRIN_H=0 + R123_USE_WMMINTRIN_H=0 + R123_USE_CPUID_MSVC=0 + ) +endif() diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index ebe3a1c..e7e2c3b 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -33,3 +33,20 @@ set_target_properties(din_common_ort PROPERTIES ) add_subdirectory(io) + +# CPU callers do not depend on CUDA headers or libraries. +include(random123) +add_library(din_common_rng STATIC rng.cpp) +target_include_directories(din_common_rng PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}") +target_compile_features(din_common_rng PUBLIC cxx_std_20) +target_link_libraries(din_common_rng PRIVATE Random123::Random123) +set_target_properties(din_common_rng PROPERTIES POSITION_INDEPENDENT_CODE ON) + +enable_language(CUDA) +add_library(din_common_rng_cuda STATIC rng_cuda.cu) +target_include_directories(din_common_rng_cuda PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}") +target_link_libraries(din_common_rng_cuda PUBLIC CUDA::cudart) +set_target_properties(din_common_rng_cuda PROPERTIES + CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}" + POSITION_INDEPENDENT_CODE ON +) diff --git a/common/ort_session.cpp b/common/ort_session.cpp index 81e6e66..9514ebd 100644 --- a/common/ort_session.cpp +++ b/common/ort_session.cpp @@ -916,17 +916,33 @@ Ort::ConstMemoryInfo OrtRunner::DeviceMemory() template TensorBuffer::TensorBuffer(OrtRunner& runner, const std::vector& shape, bool use_device_io, bool disable_uma) - : runner_(&runner) + : TensorBuffer(runner.env, runner.ep_device, runner.compute_stream, shape, use_device_io, disable_uma) +{ +} + +template +TensorBuffer::TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, + const std::vector& shape, bool use_device_io, bool disable_uma) + : env_(&env) + , compute_stream_(stream) , use_device_io_(use_device_io) - , use_unified_memory_(use_device_io && runner.UsesUnifiedMemory() && !disable_uma) + , use_unified_memory_(use_device_io && IsCudaUnifiedMemoryDevice(device) && !disable_uma) , shape_(shape) { if (use_device_io_) { - host_tensor_ = Ort::Value::CreateTensor(runner.PinnedAllocator(), shape_.data(), shape_.size()); + if (!device || stream == nullptr) + throw std::invalid_argument("Device tensors require a device and a persistent stream"); + const auto pinned = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_HOST_ACCESSIBLE)); + if (!pinned) + throw std::runtime_error("Missing shared pinned allocator"); + host_tensor_ = Ort::Value::CreateTensor(pinned, shape_.data(), shape_.size()); if (!use_unified_memory_) { - device_tensor_ = Ort::Value::CreateTensor(runner.DeviceAllocator(), shape_.data(), shape_.size()); + const auto allocator = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); + if (!allocator) + throw std::runtime_error("Missing shared device allocator"); + device_tensor_ = Ort::Value::CreateTensor(allocator, shape_.data(), shape_.size()); } } else @@ -971,7 +987,7 @@ void TensorBuffer::CopyAsyncToDevice() { if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&host_tensor_}, {&device_tensor_}, *runner_->compute_stream); + Copy(*env_, {&host_tensor_}, {&device_tensor_}, *compute_stream_); } } @@ -983,7 +999,7 @@ NotificationPtr TensorBuffer::CopyAsyncToDeviceWithNotification() { return NotificationPtr{nullptr}; } - auto notification = NotificationPtr(*runner_->compute_stream); + auto notification = NotificationPtr(*compute_stream_); // Even when the transfer is a unified-memory no-op, record the stream so // callers retain a usable ordering point for Sync(). notification.Record(); @@ -995,7 +1011,7 @@ void TensorBuffer::CopyAsyncToHost() { if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&device_tensor_}, {&host_tensor_}, *runner_->compute_stream); + Copy(*env_, {&device_tensor_}, {&host_tensor_}, *compute_stream_); } } @@ -1007,7 +1023,7 @@ NotificationPtr TensorBuffer::CopyAsyncToHostWithNotification() { return NotificationPtr{nullptr}; } - auto notification = NotificationPtr(*runner_->compute_stream); + auto notification = NotificationPtr(*compute_stream_); // A unified-memory buffer needs no copy, but the notification still waits // for device writes before its host data is consumed. notification.Record(); @@ -1023,7 +1039,7 @@ void TensorBuffer::CopyFrom(TensorBuffer& other) } if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&other.device_tensor_}, {&device_tensor_}, *runner_->compute_stream); + Copy(*env_, {&other.device_tensor_}, {&device_tensor_}, *compute_stream_); } else { diff --git a/common/ort_session.h b/common/ort_session.h index 3e0a09f..30c8572 100644 --- a/common/ort_session.h +++ b/common/ort_session.h @@ -199,6 +199,8 @@ class TensorBuffer { public: TensorBuffer(OrtRunner& runner, const std::vector& shape, bool use_device_io, bool disable_uma = false); + TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, const std::vector& shape, + bool use_device_io, bool disable_uma = false); Ort::Value& BindingValue(); @@ -228,7 +230,10 @@ class TensorBuffer [[nodiscard]] size_t ElementCount() const; - OrtRunner* runner_ = nullptr; + // Allocators belong to the environment. External streams must outlive this + // buffer, allowing a model session to be replaced without moving its tensors. + Ort::Env* env_ = nullptr; + Ort::SyncStream* compute_stream_ = nullptr; bool use_device_io_ = false; bool use_unified_memory_ = false; std::vector shape_; diff --git a/common/rng.cpp b/common/rng.cpp new file mode 100644 index 0000000..7ed5181 --- /dev/null +++ b/common/rng.cpp @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "rng.h" + +#include +#include +#include +#include +#include +#include + +#include + +namespace din::common +{ +namespace +{ +float round_bfloat16(float value) +{ + uint32_t bits = std::bit_cast(value); + if ((bits & 0x7fffffffU) > 0x7f800000U) + bits |= 0x00400000U; // Preserve NaNs as quiet NaNs. + else + bits += 0x7fffU + ((bits >> 16) & 1U); + return std::bit_cast(bits & 0xffff0000U); +} + +void box_muller(uint32_t x, uint32_t y, float* result) +{ + constexpr float a = 2.3283064e-10f; + constexpr float b = a * 6.2831855f; + // Match CUDA's fused multiply/add, including integer-to-FP32 conversion. + const float u = std::fma(static_cast(x), a, a / 2); + const float v = std::fma(static_cast(y), b, b / 2); + const float radius = std::sqrt(-2.0f * std::log(u)); + result[0] = radius * std::sin(v); + result[1] = radius * std::cos(v); +} +} // namespace + +void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const NormalOptions& options) +{ + if (rng != Rng::CpuStd && rng != Rng::CpuPhilox) + throw std::invalid_argument("CPU normal generation requires CpuStd or CpuPhilox"); + if (options.philox_stride == 0 || options.philox_stride > std::numeric_limits::max() / 4) + throw std::invalid_argument("Philox stride must be positive and allow four components without overflow"); + const size_t count = output.size(); + if (rng == Rng::CpuStd) + { + std::default_random_engine engine(static_cast(seed)); + std::normal_distribution normal(0.0f, 1.0f); + for (size_t i = 0; i < count; ++i) + { + const float sample = normal(engine); + output[i] = options.round_to_bfloat16 ? round_bfloat16(sample) : sample; + } + return; + } + + using Philox = r123::Philox4x32_R<10>; + const Philox::key_type key = {{static_cast(seed), static_cast(seed >> 32)}}; + const size_t step = 4 * options.philox_stride; + for (size_t thread = 0; thread < std::min(count, options.philox_stride); ++thread) + { + uint64_t iteration = 0; + for (size_t base = thread; base < count; ++iteration) + { + const Philox::ctr_type counter = {{static_cast(iteration), static_cast(iteration >> 32), + static_cast(thread), + static_cast(static_cast(thread) >> 32)}}; + const auto words = Philox{}(counter, key); + float samples[4]; + box_muller(words[0], words[1], samples); + box_muller(words[2], words[3], samples + 2); + for (size_t component = 0; component < 4; ++component) + { + const size_t delta = component * options.philox_stride; + if (delta >= count - base) + break; + const size_t index = base + delta; + output[index] = options.round_to_bfloat16 ? round_bfloat16(samples[component]) : samples[component]; + } + if (count - base <= step) + break; + base += step; + } + } +} +} // namespace din::common diff --git a/common/rng.h b/common/rng.h new file mode 100644 index 0000000..8ee2233 --- /dev/null +++ b/common/rng.h @@ -0,0 +1,35 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +namespace din::common +{ +enum class Rng +{ + CpuStd, + CudaPhilox, + CpuPhilox +}; + +// Reference launch geometry: H100 SXM5, 132 SMs * 2048 threads/SM. +inline constexpr size_t kH100PhiloxStride = 270336; + +struct NormalOptions +{ + bool round_to_bfloat16 = false; + // Normal4 components are separated by this many samples. Must be nonzero. + size_t philox_stride = kH100PhiloxStride; +}; + +// Fill a contiguous buffer with N(0, 1) samples, restarting at the seed each call. +// CpuStd preserves std::default_random_engine/normal_distribution behavior, +// including platform dependence and truncation of the seed to unsigned int. +// CpuPhilox uses Philox4x32-10 and cuRAND's Box-Muller convention and ordering. +// CPU math can differ from CUDA, including at BF16 rounding boundaries. +// CudaPhilox must be dispatched through GenerateNormalCuda in rng_cuda.h. +void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const NormalOptions& options = {}); +} // namespace din::common diff --git a/common/rng_cuda.cu b/common/rng_cuda.cu new file mode 100644 index 0000000..cba5954 --- /dev/null +++ b/common/rng_cuda.cu @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include + +#include +#include +#include + +#include "rng_cuda.h" +#include + +namespace din::common +{ +namespace +{ +constexpr unsigned int kThreadsPerBlock = 256; +void check_cuda(cudaError_t result) +{ + if (result != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(result)); +} +// Match ATen DistributionTemplates.h: one Philox subsequence per thread, +// normal4 components spaced by the entire grid, then grid-stride by 4. +__global__ __launch_bounds__(256, 4) void normal_philox_kernel(uint64_t seed, size_t count, bool round_to_bfloat16, + size_t channels, size_t spatial, float* output) +{ + const size_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t stride = static_cast(blockDim.x) * gridDim.x; + const size_t rounded = ((count - 1) / (stride * 4) + 1) * stride * 4; + curandStatePhilox4_32_10_t state; + curand_init(seed, thread, 0, &state); + for (size_t base = thread; base < rounded; base += stride * 4) + { + const float4 samples = curand_normal4(&state); + const float values[] = {samples.x, samples.y, samples.z, samples.w}; +#pragma unroll + for (int component = 0; component < 4; ++component) + { + const size_t index = base + stride * component; + if (index < count) + { + const size_t destination = + channels == 0 ? index + : ((index / (channels * spatial)) * spatial + index % spatial) * channels + + (index / spatial) % channels; + output[destination] = + round_to_bfloat16 ? __bfloat162float(__float2bfloat16_rn(values[component])) : values[component]; + } + } + } +} + +void generate_normal_cuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, bool round_to_bfloat16, + size_t channels, size_t spatial) +{ + if (count == 0) + return; + if (output == nullptr || count > std::numeric_limits::max() / 2) + throw std::invalid_argument("CUDA normal generation requires a valid buffer and supported count"); + int device = 0; + cudaDeviceProp properties{}; + check_cuda(cudaGetDevice(&device)); + check_cuda(cudaGetDeviceProperties(&properties, device)); + // ATen caps the grid by the device's thread capacity, not kernel occupancy. + const size_t capacity = static_cast(properties.multiProcessorCount) * + (properties.maxThreadsPerMultiProcessor / kThreadsPerBlock); + const unsigned int blocks = + static_cast(std::min(capacity, (count + kThreadsPerBlock - 1) / kThreadsPerBlock)); + normal_philox_kernel<<>>(seed, count, round_to_bfloat16, channels, spatial, + output); + check_cuda(cudaPeekAtLastError()); +} + +} // namespace + +void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, bool round_to_bfloat16) +{ + generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, 0, 0); +} + +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, size_t channels, + size_t spatial, bool round_to_bfloat16) +{ + if (channels == 0 || spatial == 0 || channels > std::numeric_limits::max() / spatial || + count % (channels * spatial) != 0) + throw std::invalid_argument("Packed CUDA normal dimensions must match the buffer"); + generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, channels, spatial); +} + +} // namespace din::common diff --git a/common/rng_cuda.h b/common/rng_cuda.h new file mode 100644 index 0000000..fdc8b65 --- /dev/null +++ b/common/rng_cuda.h @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +#include +#include + +namespace din::common +{ +// Asynchronously fill a contiguous device buffer with N(0, 1) samples. +// Fresh Philox seed, offset zero. Uses ATen's device-dependent launch geometry. +// An empty buffer is a no-op. The buffer and stream must remain valid until done. +void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, + bool round_to_bfloat16 = false); + +// Generate the same [batch, channels, spatial] samples directly into contiguous +// [batch, spatial, channels] storage. count must be a multiple of channels * spatial. +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, size_t channels, + size_t spatial, bool round_to_bfloat16 = false); +} // namespace din::common diff --git a/image_gen/flux2/CMakeLists.txt b/image_gen/flux2/CMakeLists.txt index 2308a15..72de3ca 100644 --- a/image_gen/flux2/CMakeLists.txt +++ b/image_gen/flux2/CMakeLists.txt @@ -4,9 +4,11 @@ enable_language(CUDA) include(slang) option(FLUX2_ENABLE_VULKAN_VALIDATION "Enable Vulkan validation layers for the Flux2 Vulkan target" ON) +option(DIN_BUILD_FLUX2_UI "Build the GLFW/OpenGL Flux2 desktop sample" ON) add_library(din_flux2_common STATIC utils.cpp + flux2_models.cpp ) target_include_directories(din_flux2_common @@ -18,6 +20,9 @@ target_include_directories(din_flux2_common target_link_libraries(din_flux2_common PUBLIC din_common_ort + din_common_rng + PRIVATE + nlohmann_json::nlohmann_json ) set_target_properties(din_flux2_common PROPERTIES @@ -68,7 +73,7 @@ configure_flux2_backend_target(din_flux2_cuda) set_target_properties(din_flux2_cuda PROPERTIES CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}" ) -target_link_libraries(din_flux2_cuda PRIVATE CUDA::cudart CUDA::nvtx3) +target_link_libraries(din_flux2_cuda PRIVATE din_common_rng_cuda CUDA::cudart CUDA::nvtx3) if(Vulkan_FOUND AND SLANGC) add_subdirectory(shaders) @@ -135,6 +140,11 @@ add_din_executable(din_flux2_cli flux2_cli.cpp ) target_link_libraries(din_flux2_cli PRIVATE argparse din_flux2_runtime) +if (DIN_BUILD_FLUX2_UI) + include(imgui) + add_din_ui_executable(din_flux2_ui flux2_ui.cpp) + target_link_libraries(din_flux2_ui PRIVATE argparse din_flux2_runtime CUDA::cudart) +endif () if(TARGET din_flux2_dx) target_link_libraries(din_flux2_cli PRIVATE din_flux2_dx) endif() diff --git a/image_gen/flux2/README.md b/image_gen/flux2/README.md index b0f66a0..4f68f99 100644 --- a/image_gen/flux2/README.md +++ b/image_gen/flux2/README.md @@ -1,58 +1,74 @@ # FLUX.2-klein -CLI image generation with ONNX Runtime. +Interactive image generation and a CLI for FLUX.2-klein-4B using ONNX Runtime and TensorRT RTX. + +The interactive studio was used to produce the results in +[The Weight Is Over – Interactive Diffusion on Consumer GPUs](https://arxiv.org/abs/2609.21849) +(Frieder Ganz and Maximilian Müller). See the paper for the small text-encoder translator, +interactive workflow, and speed/quality/memory trade-offs implemented here. ## Supported models -| Model | Hugging Face ID | Checkpoint | ONNX export | -| --- | --- | --- | --- | -| FLUX.2-klein-4B | `black-forest-labs/FLUX.2-klein-4B` | BF16 | FP32 I/O | -| FLUX.2-klein-4B FP8 | `black-forest-labs/FLUX.2-klein-4b-fp8` | FP8 | FP32 I/O | -| FLUX.2-klein-4B NVFP4 | `black-forest-labs/FLUX.2-klein-4b-nvfp4` | NVFP4 | FP32 I/O | +| Model | Hugging Face ID | Checkpoint | ONNX export | +|-----------------------|-------------------------------------------|------------|-------------| +| FLUX.2-klein-4B | `black-forest-labs/FLUX.2-klein-4B` | BF16 | FP32 I/O | +| FLUX.2-klein-4B FP8 | `black-forest-labs/FLUX.2-klein-4b-fp8` | FP8 | FP32 I/O | +| FLUX.2-klein-4B NVFP4 | `black-forest-labs/FLUX.2-klein-4b-nvfp4` | NVFP4 | FP32 I/O | -The CLI supports CPU or TensorRT RTX execution with CPU, CUDA, DirectX, DirectX CIG, and Vulkan processing where available. +The backend supports CPU or TensorRT RTX execution with CPU, CUDA, DirectX, DirectX CIG, and Vulkan processing where +available. -## Backends +## Export -`CiG` means compute-in-graphics: inference runs in the graphics GPU context so it can coexist with rendering workloads. +Run from the repository root: -| Mode | CLI | Inference and sampling | -| --- | --- | --- | -| Direct TensorRT RTX | `--provider trt-rtx --processing cuda` | TensorRT RTX inference with CUDA sampling. | -| Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | -| Vulkan interop | `--provider trt-rtx --processing vk` | TensorRT RTX inference with Vulkan sampling; shared buffers and timeline semaphores keep intermediate data on the GPU. | -| DirectX CiG | `--provider trt-rtx --processing dx-cig` | TensorRT RTX inference in a DirectX CiG context with DirectX sampling. | +```powershell +python image_gen/flux2/model_export/export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output models/flux_full --model all --transformer-precision bf16 +python image_gen/flux2/model_export/quantize_onnx.py --input models/flux_full/transformer_bf16 --quant fp8 --output models/flux_full/transformer_fp8 +``` -CiG is useful when inference must share the graphics GPU with rendering. It can reduce standalone inference performance because of graphics-concurrency limits. TensorRT RTX currently documents DirectX CiG support; Vulkan CiG support is planned. See [Simultaneous Compute and Graphics](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/inference-library/compute-graphics.html). +Use `--quant nvfp4` with an output directory of `transformer_nvfp4` for NVFP4. +All precisions share the text encoder, VAE, and tokenizer under one model root. -## Export +Optionally export Qwen3-0.6B + `gafr1/translator` as the small text encoder: ```powershell -cd model_export -python export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output D:\models\flux_full --model all --transformer-precision bf16 -python quantize_onnx.py --input D:\models\flux_full\transformer_bf16 --quant fp8 --output D:\models\flux_full\transformer_fp8 -python quantize_onnx.py --input D:\models\flux_full\transformer_bf16 --quant nvfp4 --output D:\models\flux_full\transformer_nvfp4 +python image_gen/flux2/model_export/export_flux2.py --model_name models/flux_full --output models/flux_full --model text_encoder_translator +python image_gen/flux2/model_export/verify_flux2.py --onnx_dir models/flux_full --encoder translator --encoder-only ``` -All precisions share one root directory. The export contains `text_encoder/model.onnx`, `vae_decoder/model.onnx`, and `tokenizer/tokenizer.json` once, plus one transformer per precision: `transformer_bf16/model.onnx`, `transformer_fp8/model.onnx`, and `transformer_nvfp4/model.onnx`. +Use `verify_flux2.py --help` for full-pipeline verification options. -## Verify +## Build and run ```powershell -python verify_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --onnx_dir D:\models\flux_full --precision fp8 --output_dir flux2-fp8 --provider trt-rtx +cmake --preset windows-x64 +cmake --build --preset windows-x64-release --target din_flux2_ui din_flux2_cli ``` -## Build +From the repo root you can then run: -```powershell -cmake --preset windows-x64 -cmake --build --preset windows-x64-debug --target din_flux2_cli ``` +out/build/windows-x64/bin/Release/din_flux2_ui.exe --model-dir /flux_full +``` + +The studio provides prompt editing, encoder/precision/backend selection, seed sweeps, +weight-streaming controls, and an image gallery with optional PNG output. It selects the +translator when exported; override with `--encoder 4b|translator`. +Runtime and compiled-model caching are always enabled; clear affected caches after re-exporting models in place. +The has configuration checks per stage and does partial reload of impacted models upon changes. +Changing RNG, denoising steps, or the runtime weight-streaming budget reuses all +loaded model sessions. For a full execution backend change all models are reloaded. + +The UI uses GLFW, Dear ImGui, and OpenGL. For Debian/Ubuntu, install +`libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev`. +Use `-DDIN_BUILD_FLUX2_UI=OFF` for headless builds. -## Run +For pure CLI you can use: ```powershell -out\build\windows-x64\bin\Debug\din_flux2_cli.exe --model-dir D:\models\flux_full --precision nvfp4 --provider trt-rtx --processing cuda --prompt "a red fox" --output out +out/build/windows-x64/bin/Release/din_flux2_cli.exe --model-dir models/flux_full --precision fp8 --provider trt-rtx --processing cuda --prompt "a red fox" --output out ``` -`--precision` selects `transformer_` and defaults to `bf16`. Only the transformer TensorRT RTX runtime-cache and EP-context paths are precision-qualified, so switching precisions never reuses an incompatible compiled transformer engine while the shared text encoder and VAE reuse their existing caches. +Use `--help` for available processing backends and options, including `--encoder`, +`--steps`, and startup-only `--ws`. CPU, CUDA, DirectX, and Vulkan processing are available where built. diff --git a/image_gen/flux2/cuda_kernels.h b/image_gen/flux2/cuda_kernels.h index 46a92d1..072b128 100644 --- a/image_gen/flux2/cuda_kernels.h +++ b/image_gen/flux2/cuda_kernels.h @@ -6,6 +6,7 @@ #include #include +#include void launch_flux_euler_kernel(cudaStream_t stream, float t_curr, float t_next, size_t total_elements, float* hidden_states, const float* transformer_output); diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index c8d7e08..4791c82 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -23,7 +24,8 @@ unsigned int parse_uint(const std::string& value, const char* flag_name) { size_t parsed_chars = 0; const unsigned long parsed = std::stoul(value, &parsed_chars, 10); - if (parsed_chars != value.size()) + if (value.empty() || value.front() == '-' || parsed_chars != value.size() || + parsed > std::numeric_limits::max()) { throw std::invalid_argument(std::string(flag_name) + " must be an unsigned integer"); } @@ -118,11 +120,11 @@ std::string parse_precision(std::string value) { c = static_cast(std::tolower(static_cast(c))); } - if (value == "bf16" || value == "fp8" || value == "nvfp4") + if (value == "bf16" || value == "fp16" || value == "fp8" || value == "nvfp4") { return value; } - throw std::invalid_argument("Precision must be one of: bf16, fp8, nvfp4"); + throw std::invalid_argument("Precision must be one of: bf16, fp16, fp8, nvfp4"); } void validate_config(const Flux2Config& config) @@ -136,7 +138,6 @@ void validate_config(const Flux2Config& config) Flux2Config parse_args(int argc, char* argv[]) { Flux2Config config; - config.model_dir = DEFAULT_MODEL_BASE_PATH; config.prompt = DEFAULT_PROMPT; argparse::ArgumentParser parser("din_flux2"); @@ -152,22 +153,22 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("cpu|trt-rtx") .help("Select the ONNX Runtime execution provider."); parser.add_argument("--model-dir") - .default_value(config.model_dir.string()) + .required() .nargs(1) .metavar("PATH") .help("Root directory with shared Flux2 ONNX artifacts and transformer_ directories."); parser.add_argument("--precision") - .default_value(config.precision) + .default_value(config.transformer.precision) .nargs(1) - .metavar("bf16|fp8|nvfp4") + .metavar("bf16|fp16|fp8|nvfp4") .help("Transformer precision to load from transformer_."); parser.add_argument("--ep-cache") - .default_value(config.ep_cache_dir.string()) + .default_value(config.text_encoder.ep_cache_dir.string()) .nargs(1) .metavar("PATH") .help("TensorRT RTX runtime cache directory."); parser.add_argument("--ep-context-dir") - .default_value(config.ep_context_dir.string()) + .default_value(config.text_encoder.ep_context_dir.string()) .nargs(1) .metavar("PATH") .help("Directory for ONNX Runtime EP-context models."); @@ -182,6 +183,16 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("TEXT") .help("Prompt text to encode with the Flux2 tokenizer."); parser.add_argument("--seed").default_value(std::to_string(config.seed)).nargs(1).metavar("N").help("Random seed."); + parser.add_argument("--rng") + .default_value(std::string{"cpu-std"}) + .nargs(1) + .metavar("cpu-std|cuda-philox|cpu-philox") + .help("Noise generator: cpu-std (default), cuda-philox (CUDA only), or cpu-philox (fixed H100 SXM5 ordering)."); + parser.add_argument("--encoder").default_value(std::string("4b")).help("4b or translator."); + parser.add_argument("--steps").default_value(std::string("4")).help("Denoise steps, 1 through 50."); + parser.add_argument("--ws") + .default_value(std::string("off")) + .help("Transformer resident budget: off or 0% through 100%."); parser.add_argument("--num-images") .default_value(std::to_string(config.num_images)) .nargs(1) @@ -200,18 +211,40 @@ Flux2Config parse_args(int argc, char* argv[]) config.processing = parse_processing_backend(parser.get("--processing")); config.provider = parse_execution_provider(parser.get("--provider")); - config.model_dir = parser.get("--model-dir"); - config.precision = parse_precision(parser.get("--precision")); - config.ep_cache_dir = parser.get("--ep-cache"); - config.ep_context_dir = parser.get("--ep-context-dir"); + const std::filesystem::path model_root = parser.get("--model-dir"); + config.transformer.precision = parse_precision(parser.get("--precision")); + config.text_encoder.ep_cache_dir = parser.get("--ep-cache"); + config.text_encoder.ep_context_dir = parser.get("--ep-context-dir"); config.prompt = parser.get("--prompt"); config.seed = parse_uint(parser.get("--seed"), "--seed"); + const auto rng = parser.get("--rng"); + if (rng == "cpu-std") + config.rng = din::common::Rng::CpuStd; + else if (rng == "cuda-philox") + config.rng = din::common::Rng::CudaPhilox; + else if (rng == "cpu-philox") + config.rng = din::common::Rng::CpuPhilox; + else + throw std::invalid_argument("--rng must be cpu-std, cuda-philox, or cpu-philox"); + const auto encoder = parser.get("--encoder"); + if (encoder != "4b" && encoder != "translator") + throw std::invalid_argument("--encoder must be 4b or translator"); + config.text_encoder.variant = encoder == "4b" ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + const auto steps = parse_uint(parser.get("--steps"), "--steps"); + if (steps < 1 || steps > 50) + throw std::invalid_argument("--steps must be 1 through 50"); + config.steps = static_cast(steps); + config.transformer.initial_weight_streaming_budget = parser.get("--ws"); + if (config.transformer.initial_weight_streaming_budget == "off") + config.transformer.initial_weight_streaming_budget.clear(); config.num_images = parse_uint(parser.get("--num-images"), "--num-images"); const std::filesystem::path output_dir = parser.get("--output"); std::filesystem::create_directories(output_dir); config.output_path = output_dir / "flux2.png"; + ConfigureFlux2Models(config, model_root, config.text_encoder.ep_cache_dir, config.text_encoder.ep_context_dir); validate_config(config); + ValidateFlux2Config(config); return config; } @@ -246,16 +279,21 @@ int main(int argc, char* argv[]) try { const Flux2Config config = parse_args(argc, argv); - std::cout << "Model dir: " << config.model_dir.string() << "\n" - << "Precision: " << config.precision << "\n" + std::cout << "Model dir: " << config.text_encoder.model_path.parent_path().parent_path().string() << "\n" + << "Precision: " << config.transformer.precision << "\n" << "Output: " << config.output_path.string() << "\n" << "Prompt: " << config.prompt << "\n" << "Seed: " << config.seed << "\n" + << "RNG: " + << (config.rng == din::common::Rng::CpuStd ? "cpu-std" + : config.rng == din::common::Rng::CudaPhilox ? "cuda-philox" + : "cpu-philox") + << "\n" << "Num images: " << config.num_images << "\n" << "Processing: " << to_string(config.processing) << "\n" << "Provider: " << to_string(config.provider) << "\n" - << "EP cache: " << config.ep_cache_dir.string() << "\n" - << "EP context dir: " << config.ep_context_dir.string() << "\n" + << "EP cache: " << config.text_encoder.ep_cache_dir.string() << "\n" + << "EP context dir: " << config.text_encoder.ep_context_dir.string() << "\n" << std::endl; auto pipeline = CreateFlux2Pipeline(config); @@ -268,6 +306,10 @@ int main(int argc, char* argv[]) << " (seed=" << current_seed << ") ==========" << std::endl; const auto image_start = std::chrono::steady_clock::now(); Flux2Image image = pipeline->GenerateImage(current_seed); + const auto& times = image.timings; + std::cout << "Timings (ms): encode=" << times.encode_ms << " rng=" << times.rng_ms + << " denoise=" << times.denoise_ms << " decode=" << times.decode_ms << " total=" << times.total_ms + << std::endl; const auto image_end = std::chrono::steady_clock::now(); save_image(make_output_path(config, image_index), image); const std::chrono::duration image_duration = image_end - image_start; diff --git a/image_gen/flux2/flux2_cli.h b/image_gen/flux2/flux2_cli.h index 75ceeeb..7e8f775 100644 --- a/image_gen/flux2/flux2_cli.h +++ b/image_gen/flux2/flux2_cli.h @@ -4,10 +4,15 @@ #pragma once #include +#include #include +#include +#include #include #include +#include "rng.h" + enum class Flux2ProcessingBackend { Cpu, @@ -24,25 +29,102 @@ enum class Flux2ExecutionProvider TrtRtx, }; +enum class Flux2TextEncoder +{ + Qwen3_4B, + Qwen3_06BTranslator +}; + +inline void ValidateFlux2Rng(din::common::Rng rng, Flux2ProcessingBackend processing) +{ + using din::common::Rng; + if (rng != Rng::CpuStd && rng != Rng::CudaPhilox && rng != Rng::CpuPhilox) + throw std::invalid_argument("Unknown random number generator"); + if (rng == Rng::CudaPhilox && processing != Flux2ProcessingBackend::Cuda) + throw std::invalid_argument("cuda-philox RNG requires --processing cuda"); +} + +// Called on the generation thread. Completed denoising steps are in [0, total]. +using Flux2Progress = std::function; + +struct Flux2Timings +{ + double encode_ms = 0; + double rng_ms = 0; + double denoise_ms = 0; + double decode_ms = 0; + double total_ms = 0; + double per_step_ms = 0; +}; + +struct Flux2ModelConfig +{ + std::filesystem::path model_path; + std::filesystem::path ep_cache_dir = "artifacts/flux2/trt_rtx_cache"; + std::filesystem::path ep_context_dir = "artifacts/flux2/ep_context"; + std::string cache_subpath; + bool operator==(const Flux2ModelConfig&) const = default; +}; + +struct Flux2TextEncoderConfig : Flux2ModelConfig +{ + Flux2TextEncoder variant = Flux2TextEncoder::Qwen3_4B; + std::filesystem::path tokenizer_dir; + bool operator==(const Flux2TextEncoderConfig&) const = default; +}; + +struct Flux2TransformerConfig : Flux2ModelConfig +{ + std::string precision = "bf16"; + // Empty keeps a non-streaming engine; -1 enables automatic streaming. + std::string initial_weight_streaming_budget; + bool operator==(const Flux2TransformerConfig&) const = default; +}; + +struct Flux2VaeConfig : Flux2ModelConfig +{ + std::filesystem::path latent_stats_path; + bool operator==(const Flux2VaeConfig&) const = default; +}; + +// Recoverable: the affected stage is unavailable; other sessions remain loaded. +class Flux2StageLoadError : public std::runtime_error +{ +public: + using std::runtime_error::runtime_error; +}; + struct Flux2Config { Flux2ProcessingBackend processing = Flux2ProcessingBackend::Cuda; Flux2ExecutionProvider provider = Flux2ExecutionProvider::TrtRtx; - std::filesystem::path model_dir; - std::string precision = "bf16"; - std::filesystem::path ep_cache_dir = "artifacts/flux2/trt_rtx_cache"; - std::filesystem::path ep_context_dir = "artifacts/flux2/ep_context"; + Flux2TextEncoderConfig text_encoder; + Flux2TransformerConfig transformer; + Flux2VaeConfig vae; + int steps = 4; std::filesystem::path output_path; std::string prompt; unsigned int seed = 42; + // CPU standard noise, CUDA Philox, or CPU Philox with fixed H100 SXM5 ordering. + din::common::Rng rng = din::common::Rng::CpuStd; unsigned int num_images = 5; }; +struct Flux2GenerationOptions +{ + // Omit to restore the startup budget. Applied via transformer session dynamic options only. + std::optional weight_streaming_budget; + // Omit to use the startup selections. Overrides apply to this generation only. + std::optional steps; + std::optional rng; +}; + struct Flux2Image { std::vector data; int height = 0; int width = 0; + Flux2Timings timings; }; class Flux2ProcessingPipeline @@ -50,8 +132,20 @@ class Flux2ProcessingPipeline public: virtual ~Flux2ProcessingPipeline() = default; virtual void Initialize() = 0; + virtual void UpdateModels(const Flux2Config& config) = 0; virtual void SetPrompt(std::string prompt) = 0; - virtual Flux2Image GenerateImage(unsigned int seed) = 0; + // Opaque identity for diagnostics; never dereference or retain after pipeline destruction. + virtual const void* TransformerSessionIdentity() const = 0; + virtual const void* TextEncoderSessionIdentity() const = 0; + virtual const void* VaeSessionIdentity() const = 0; + virtual Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress = {}, + const Flux2GenerationOptions& options = {}) = 0; }; std::unique_ptr CreateFlux2Pipeline(const Flux2Config& config); + +void ValidateFlux2Config(const Flux2Config& config); +// Resolve existing CLI/UI root controls into independent stage settings. +void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& model_root, + const std::filesystem::path& cache_root, const std::filesystem::path& context_root); +bool IsFlux2BackendAvailable(Flux2ProcessingBackend backend); diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index 420c629..1ac7d6e 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -26,8 +26,11 @@ #include "cu_helper.h" #include "cuda_kernels.h" #include "flux2.h" +#include "flux2_models.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" +#include "rng_cuda.h" #include "utils.h" #include #include @@ -155,19 +158,6 @@ void fill_position_ids(std::vector& img_ids, std::vector& txt_ } } -std::vector create_time_schedule() -{ - constexpr double MU = 2.291179894115571; - std::vector time_schedule(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - const double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << time_schedule[i] << std::endl; - } - return time_schedule; -} - class EulerStage { public: @@ -320,7 +310,7 @@ class PostprocessStage constexpr int HW = static_cast(LATENT_HEIGHT * LATENT_WIDTH); unpack_latents_with_ids(hidden, scratch, LATENT_HEIGHT, LATENT_WIDTH, LATENT_CHANNELS); - denormalize_latents(scratch, BN_MEAN, BN_STD, LATENT_CHANNELS, HW); + denormalize_latents(scratch, bn_mean_->HostData(), bn_std_->HostData(), LATENT_CHANNELS, HW); unpatchify_latents(scratch, decoder_input, 1, LATENT_CHANNELS / PATCH_SIZE / PATCH_SIZE, LATENT_HEIGHT, LATENT_WIDTH, PATCH_SIZE, PATCH_SIZE); if (device_is_cuda) @@ -389,14 +379,24 @@ struct CudaPipelineState std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; EulerStage euler; PostprocessStage postprocess; }; void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedule, SamplingBackend sampling_backend, - bool device_is_cuda, Ort::SyncStream* compute_stream, bool encode_prompt) + bool device_is_cuda, Ort::SyncStream* compute_stream, bool encode_prompt, Flux2Telemetry& telemetry) { nvtx3::scoped_range nvtx_pipeline("run_pipeline"); + Ort::RunOptions transformer_options; + if (compute_stream) + { + transformer_options.SetSyncStream(*compute_stream); + transformer_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); + } + Ort::RunOptions run_options; if (compute_stream != nullptr) { @@ -404,13 +404,16 @@ void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedul run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); } + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); nvtx3::scoped_range nvtx("text_encoder"); pipeline.text_encoder_session.Run(run_options, pipeline.text_encoder_io); } - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { nvtx3::scoped_range nvtx_step("diffusion_step_" + std::to_string(step)); const float t_curr = time_schedule[step]; @@ -425,14 +428,16 @@ void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedul { nvtx3::scoped_range nvtx_transformer("transformer"); - pipeline.transformer_session.Run(run_options, pipeline.transformer_io); + pipeline.transformer_session.Run(transformer_options, pipeline.transformer_io); } { nvtx3::scoped_range nvtx_euler("euler_step"); pipeline.euler.Run(compute_stream, sampling_backend, device_is_cuda, t_curr, t_next); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); { nvtx3::scoped_range nvtx_post("postprocess"); @@ -477,8 +482,6 @@ Ort::ConstEpDevice find_cpu_device(Ort::Env& env) void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const Flux2Config& config, ExecutionProviderMode provider_mode, SamplingBackend sampling_backend) { - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); - const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision); const bool trt_rtx_device = is_trt_rtx_device(ep_device); const bool cpu_device = is_cpu_device(ep_device); const bool has_separate_binding = !cpu_device; @@ -510,105 +513,55 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const Ort::SyncStream* compute_stream_ptr = state.compute_stream ? &*state.compute_stream : nullptr; - std::cout << "=== Loading ONNX Models ===" << std::endl; - std::string provider; - if (provider_mode == ExecutionProviderMode::Trt) - { - provider = "trt-rtx"; - } - else - { - provider = "cpu"; - } - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); - - auto make_profile = [](std::string cache_subpath) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.enable_cuda_graph = false; - profile.embed_ep_context = false; - profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); - return profile; - }; - - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), provider, cache_dir, ep_context, - make_profile(cache_paths.text_encoder), compute_stream_ptr); - std::cout << " Text encoder loaded" << std::endl; - state.transformer_runner = - std::make_unique(state.env, model_paths.transformer_model.string(), provider, cache_dir, - ep_context, make_profile(cache_paths.transformer), compute_stream_ptr); - std::cout << " Transformer loaded" << std::endl; - state.vae_decoder_runner = - std::make_unique(state.env, model_paths.vae_decoder_model.string(), provider, cache_dir, - ep_context, make_profile(cache_paths.vae_decoder), compute_stream_ptr); - std::cout << " VAE decoder loaded" << std::endl; + state.model_profile.enable_cuda_graph = false; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.token = std::make_unique>(*state.text_encoder_runner, token_shape, - has_separate_binding); + state.token = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + token_shape, has_separate_binding); std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.attention_mask = std::make_unique>(*state.text_encoder_runner, - attn_mask_shape, has_separate_binding); + state.attention_mask = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, attn_mask_shape, has_separate_binding); std::vector te_out_shape = {BATCH_SIZE, SEQUENCE_LENGTH, TEXT_ENCODER_EMBED_DIM}; - state.text_encoder_embeds = std::make_unique>(*state.text_encoder_runner, - te_out_shape, has_separate_binding); + state.text_encoder_embeds = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, te_out_shape, has_separate_binding); std::vector timestep_shape = {BATCH_SIZE}; - state.timestep = std::make_unique>(*state.transformer_runner, timestep_shape, - has_separate_binding); + state.timestep = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + timestep_shape, has_separate_binding); std::vector hidden_shape = {BATCH_SIZE, TRANSFORMER_HIDDEN_DIM, LATENT_CHANNELS}; - state.hidden_states = std::make_unique>(*state.transformer_runner, hidden_shape, - has_separate_binding); + state.hidden_states = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector img_ids_shape = {BATCH_SIZE, IMAGE_SEQUENCE, 4}; - state.img_ids = std::make_unique>(*state.transformer_runner, img_ids_shape, - has_separate_binding); + state.img_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + img_ids_shape, has_separate_binding); std::vector txt_ids_shape = {BATCH_SIZE, SEQUENCE_LENGTH, 4}; - state.txt_ids = std::make_unique>(*state.transformer_runner, txt_ids_shape, - has_separate_binding); + state.txt_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + txt_ids_shape, has_separate_binding); - state.transformer_output = std::make_unique>(*state.transformer_runner, - hidden_shape, has_separate_binding); + state.transformer_output = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector dec_latent_shape = {BATCH_SIZE, LATENT_CHANNELS / PATCH_SIZE / PATCH_SIZE, LATENT_HEIGHT * PATCH_SIZE, LATENT_WIDTH * PATCH_SIZE}; - state.decoder_latent = std::make_unique>(*state.vae_decoder_runner, - dec_latent_shape, has_separate_binding); + state.decoder_latent = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, dec_latent_shape, has_separate_binding); std::vector image_shape = {BATCH_SIZE, IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH}; - state.decoded_image = std::make_unique>(*state.vae_decoder_runner, image_shape, - has_separate_binding); + state.decoded_image = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, image_shape, has_separate_binding); std::vector bn_shape = {LATENT_CHANNELS}; - state.bn_mean = - std::make_unique>(*state.transformer_runner, bn_shape, has_separate_binding); - state.bn_std = - std::make_unique>(*state.transformer_runner, bn_shape, has_separate_binding); - - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token->BindingValue()); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask->BindingValue()); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds->BindingValue()); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states->BindingValue()); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds->BindingValue()); - state.transformer_io->BindInput("timestep", state.timestep->BindingValue()); - state.transformer_io->BindInput("img_ids", state.img_ids->BindingValue()); - state.transformer_io->BindInput("txt_ids", state.txt_ids->BindingValue()); - state.transformer_io->BindOutput("sample", state.transformer_output->BindingValue()); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_latent->BindingValue()); - state.vae_decoder_io->BindOutput("sample", state.decoded_image->BindingValue()); + state.bn_mean = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + bn_shape, has_separate_binding); + state.bn_std = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + bn_shape, has_separate_binding); std::vector img_ids_cpu(state.img_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); std::vector txt_ids_cpu(state.txt_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); @@ -616,29 +569,15 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const std::copy(img_ids_cpu.begin(), img_ids_cpu.end(), state.img_ids->HostData()); std::copy(txt_ids_cpu.begin(), txt_ids_cpu.end(), state.txt_ids->HostData()); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), - state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); - std::copy_n(BN_MEAN, LATENT_CHANNELS, state.bn_mean->HostData()); - std::copy_n(BN_STD, LATENT_CHANNELS, state.bn_std->HostData()); - if (has_separate_binding) { auto img_ids_ready = state.img_ids->CopyAsyncToDeviceWithNotification(); auto txt_ids_ready = state.txt_ids->CopyAsyncToDeviceWithNotification(); - auto token_ready = state.token->CopyAsyncToDeviceWithNotification(); - auto attention_mask_ready = state.attention_mask->CopyAsyncToDeviceWithNotification(); - auto bn_mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); - auto bn_std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); img_ids_ready.Sync(); txt_ids_ready.Sync(); - token_ready.Sync(); - attention_mask_ready.Sync(); - bn_mean_ready.Sync(); - bn_std_ready.Sync(); } - state.time_schedule = create_time_schedule(); + state.time_schedule = MakeFlux2Schedule(config.steps); state.euler.BindInput("hidden_states", *state.hidden_states); state.euler.BindInput("sample", *state.transformer_output); @@ -651,10 +590,94 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const state.postprocess.BindOutput("latent_sample", *state.decoder_latent); } -Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) +static void upload_cuda_prompt(CudaPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), + state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); + + if (state.has_separate_binding) + { + auto token_ready = state.token->CopyAsyncToDeviceWithNotification(); + auto attention_mask_ready = state.attention_mask->CopyAsyncToDeviceWithNotification(); + token_ready.Sync(); + attention_mask_ready.Sync(); + } +} + +static void update_cuda_models(CudaPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + if (state.compute_stream) + signal_and_wait_on_host(*state.compute_stream); + } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = state.compute_stream ? &*state.compute_stream : nullptr; + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, + {{"input_ids", &state.token->BindingValue()}, + {"attention_mask", &state.attention_mask->BindingValue()}}, + {{"prompt_embeds", &state.text_encoder_embeds->BindingValue()}}); + upload_cuda_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model( + session, + {{"hidden_states", &state.hidden_states->BindingValue()}, + {"encoder_hidden_states", &state.text_encoder_embeds->BindingValue()}, + {"timestep", &state.timestep->BindingValue()}, + {"img_ids", &state.img_ids->BindingValue()}, + {"txt_ids", &state.txt_ids->BindingValue()}}, + {{"sample", &state.transformer_output->BindingValue()}}); + })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = + BindFlux2Model(session, {{"latent_sample", &state.decoder_latent->BindingValue()}}, + {{"sample", &state.decoded_image->BindingValue()}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + std::copy_n(stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); + std::copy_n(stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); + if (state.has_separate_binding) + { + auto mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); + auto std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); + mean_ready.Sync(); + std_ready.Sync(); + } + return binding; + }); +} + +Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { Ort::SyncStream* compute_stream_ptr = state.compute_stream ? &*state.compute_stream : nullptr; + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + if (compute_stream_ptr) + signal_and_wait_on_host(*compute_stream_ptr); + }); FluxPipeline pipeline{state.text_encoder_runner->session, state.transformer_runner->session, state.vae_decoder_runner->session, @@ -669,8 +692,19 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) nvtx3::scoped_range nvtx("initialize_latents"); std::cout << "\n=== Initializing Latent (seed=" << seed << ") ===" << std::endl; const size_t hidden_count = state.hidden_states->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount(); - initialize_latent(state.hidden_states->HostData(), hidden_count, seed); - if (state.binding_is_cuda) + if (rng == din::common::Rng::CudaPhilox) + { + if (!state.binding_is_cuda || compute_stream_ptr == nullptr) + throw std::runtime_error("cuRAND RNG requires device-backed bindings and an ORT CUDA stream"); + const auto stream = reinterpret_cast(compute_stream_ptr->GetHandle()); + din::common::GenerateNormalCudaPacked( + stream, state.hidden_states->BindingValue().GetTensorMutableData(), hidden_count, seed, + LATENT_CHANNELS, LATENT_HEIGHT * LATENT_WIDTH, true); + } + else + initialize_cpu_latent(state.hidden_states->HostData(), hidden_count, seed, rng, BATCH_SIZE, LATENT_CHANNELS, + LATENT_HEIGHT, LATENT_WIDTH); + if (state.binding_is_cuda && rng != din::common::Rng::CudaPhilox) { auto hidden_uploaded = state.hidden_states->CopyAsyncToDeviceWithNotification(); hidden_uploaded.Sync(); @@ -678,7 +712,7 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) } run_pipeline(pipeline, state.time_schedule, state.sampling_backend, state.binding_is_cuda, compute_stream_ptr, - !state.prompt_embeds_valid); + !state.prompt_embeds_valid, telemetry); state.prompt_embeds_valid = true; Flux2Image image; @@ -692,6 +726,7 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) } const float* image_data = state.decoded_image->HostData(); std::copy_n(image_data, image.data.size(), image.data.data()); + image.timings = telemetry.Finish(); return image; } @@ -708,6 +743,7 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } @@ -721,16 +757,20 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline throw std::runtime_error("CUDA processing requires --provider trt-rtx"); } - std::cout << "Model dir: " << config_.model_dir.string() << "\n" - << "Execution provider: " << to_string(provider) << "\n" - << "Sampling backend: " << to_string(sampling_backend) << "\n" - << std::endl; - state_ = std::make_unique(runtime_.env); if (provider == ExecutionProviderMode::Cpu) { std::cout << "Using CPU execution provider" << std::endl; - initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); + try + { + initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); + } + catch (...) + { + state_.reset(); + throw; + } + UpdateModels(config_); return; } @@ -741,11 +781,35 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } std::cout << "Using TensorRT RTX execution provider" << std::endl; - initialize_ep(*state_, trt_device, config_, provider, sampling_backend); + try + { + initialize_ep(*state_, trt_device, config_, provider, sampling_backend); + } + catch (...) + { + state_.reset(); + throw; + } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_cuda_models(*state_, requested); } void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -753,27 +817,49 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state_->token->HostData(), - state_->attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); + upload_cuda_prompt(*state_, config_); + } - if (state_->has_separate_binding) - { - auto token_ready = state_->token->CopyAsyncToDeviceWithNotification(); - auto attention_mask_ready = state_->attention_mask->CopyAsyncToDeviceWithNotification(); - token_ready.Sync(); - attention_mask_ready.Sync(); - } + const void* TransformerSessionIdentity() const override + { + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; } - Flux2Image GenerateImage(unsigned int seed) override + const void* TextEncoderSessionIdentity() const override { - if (!state_) + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override + { + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 CUDA pipeline was not initialized"); } - return run_initialized_pipeline(*state_, seed); + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + state_->applied_budget = budget; + } + return run_initialized_pipeline(*state_, seed, rng, progress); } private: diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index 8cd2acb..b2670eb 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -18,7 +18,9 @@ #include "dx_helper.h" #include "flux2.h" +#include "flux2_models.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" #include @@ -289,13 +291,18 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr DxPostprocessPipeline& postprocess_pipeline, OrtD3D12Fence& sync_fence, Ort::SyncStream& ort_stream, DxBuffer& timestep_buf, DxBuffer& hidden_states, DxBuffer& transformer_output, DxBuffer& decoder_latent, DxBuffer& bn_mean, DxBuffer& bn_std, - bool encode_prompt) + bool encode_prompt, Flux2Telemetry& telemetry) { + Ort::RunOptions transformer_options; + transformer_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + Ort::RunOptions run_options; run_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); text_encoder_session.Run(run_options, text_encoder_io); } @@ -305,12 +312,13 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr const uint64_t initial_timestep_uploaded = sync_fence.signal_d3d12(dx.queue.Get()); sync_fence.wait_ort(ort_stream, initial_timestep_uploaded); - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { const float t_curr = time_schedule[step]; const float t_next = time_schedule[step + 1]; - transformer_session.Run(run_options, transformer_io); + transformer_session.Run(transformer_options, transformer_io); const uint64_t transformer_done = sync_fence.signal_ort(ort_stream); sync_fence.wait_d3d12(dx.queue.Get(), transformer_done); @@ -322,19 +330,21 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr dx.begin(); euler_pipeline.record(dx, hidden_states, transformer_output, euler); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { dx.record_upload(timestep_buf, &t_next, sizeof(float)); } dx.execute(); const uint64_t d3d_done = sync_fence.signal_d3d12(dx.queue.Get()); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { sync_fence.wait_ort(ort_stream, d3d_done); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); DxPostprocessConstants constants{}; constants.C = static_cast(LATENT_CHANNELS); constants.I = static_cast(LATENT_HEIGHT); @@ -440,6 +450,9 @@ struct DxPipelineState std::unique_ptr transformer_io; std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; }; static void initialize_dx_state(DxPipelineState& state, const Flux2Config& config, Ort::ConstEpDevice trt_device) @@ -449,13 +462,6 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi { throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); - const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx"); - - std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" - << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" - << std::endl; std::cout << "\n=== Initializing D3D12 ===" << std::endl; state.trt_device = trt_device; @@ -488,29 +494,9 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi graphics_ep_options.emplace_back("nv_length_aux_stream_array", "0"); } - std::cout << "=== Loading ONNX Models ===" << std::endl; - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); - auto make_profile = [&](std::string cache_subpath) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.enable_cuda_graph = state.use_cig ? false : true; - profile.embed_ep_context = false; - profile.extra_ep_options = graphics_ep_options; - return profile; - }; - - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.text_encoder), &*state.sync_stream); - state.transformer_runner = std::make_unique( - state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer), &*state.sync_stream); - state.vae_decoder_runner = std::make_unique( - state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.vae_decoder), &*state.sync_stream); + state.model_profile.enable_cuda_graph = !state.use_cig; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options = graphics_ep_options; std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; @@ -569,27 +555,8 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.decoded_image_tensor = state.tensor_importer->create_tensor(state.decoded_image_buf, image_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token_tensor); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask_tensor); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds_tensor); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states_tensor); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds_tensor); - state.transformer_io->BindInput("timestep", state.timestep_tensor); - state.transformer_io->BindInput("img_ids", state.img_ids_tensor); - state.transformer_io->BindInput("txt_ids", state.txt_ids_tensor); - state.transformer_io->BindOutput("sample", state.transformer_output_tensor); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_input_tensor); - state.vae_decoder_io->BindOutput("sample", state.decoded_image_tensor); - std::cout << "\n=== Uploading Static Inputs (D3D12) ===" << std::endl; state.dx.begin(); - state.dx.record_upload(state.bn_mean_buf, BN_MEAN, state.bn_mean_buf.size); - state.dx.record_upload(state.bn_std_buf, BN_STD, state.bn_std_buf.size); { std::vector img_ids_cpu(shape_numel(img_ids_shape)); @@ -627,36 +594,101 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.dx.record_upload(state.txt_ids_buf, txt_ids_cpu.data(), state.txt_ids_buf.size); } - constexpr double MU = 2.291179894115571; - state.time_schedule.resize(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - state.time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << state.time_schedule[i] << std::endl; - } - - std::cout << "\n=== Tokenizing Prompt ===" << std::endl; - { - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + state.time_schedule = MakeFlux2Schedule(config.steps); - std::vector tokens_cpu(shape_numel(token_shape)); - std::vector attn_mask_cpu(shape_numel(attn_mask_shape)); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - state.dx.record_upload(state.token_buf, tokens_cpu.data(), state.token_buf.size); - state.dx.record_upload(state.attn_mask_buf, attn_mask_cpu.data(), state.attn_mask_buf.size); - } state.dx.execute(); const uint64_t static_uploads_done = state.sync_fence->signal_d3d12(state.dx.queue.Get()); state.sync_fence->wait_ort(*state.sync_stream, static_uploads_done); } -static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) +static void upload_dx_prompt(DxPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), + BATCH_SIZE, SEQUENCE_LENGTH); + + state.dx.begin(); + state.dx.record_upload(state.token_buf, tokens_cpu.data(), state.token_buf.size); + state.dx.record_upload(state.attn_mask_buf, attn_mask_cpu.data(), state.attn_mask_buf.size); + state.dx.execute(); + const uint64_t prompt_uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, prompt_uploaded); +} + +static void update_dx_models(DxPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + const auto graphics_done = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, graphics_done); + signal_and_wait_on_host(*state.sync_stream); + } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = &*state.sync_stream; + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model( + session, + {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, + {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_dx_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model(session, + {{"hidden_states", &state.hidden_states_tensor}, + {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, + {"timestep", &state.timestep_tensor}, + {"img_ids", &state.img_ids_tensor}, + {"txt_ids", &state.txt_ids_tensor}}, + {{"sample", &state.transformer_output_tensor}}); + })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, + {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + state.dx.begin(); + state.dx.record_upload(state.bn_mean_buf, stats.mean.data(), state.bn_mean_buf.size); + state.dx.record_upload(state.bn_std_buf, stats.std.data(), state.bn_std_buf.size); + state.dx.execute(); + const auto uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, uploaded); + signal_and_wait_on_host(*state.sync_stream); + return binding; + }); +} + +static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + const auto completed = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, completed); + signal_and_wait_on_host(*state.sync_stream); + }); std::vector latent_cpu(state.hidden_count); - initialize_latent(latent_cpu.data(), latent_cpu.size(), seed); + initialize_cpu_latent(latent_cpu.data(), latent_cpu.size(), seed, rng, BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, + LATENT_WIDTH); state.dx.begin(); state.dx.record_upload(state.hidden_states, latent_cpu.data(), state.hidden_states.size); state.dx.execute(); @@ -668,7 +700,7 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) *state.vae_decoder_io, state.time_schedule, state.dx, *state.euler_pipeline, *state.postprocess_pipeline, *state.sync_fence, *state.sync_stream, state.timestep_buf, state.hidden_states, state.transformer_output, state.decoder_latent, state.bn_mean_buf, - state.bn_std_buf, !state.prompt_embeds_valid); + state.bn_std_buf, !state.prompt_embeds_valid, telemetry); state.prompt_embeds_valid = true; state.dx.readback(state.decoded_image_buf, state.readback_buffer, state.decoded_image_bytes); @@ -681,6 +713,7 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) std::copy_n(image_data, image.data.size(), image.data.data()); state.dx.unmap(state.readback_buffer); + image.timings = telemetry.Finish(); return image; } @@ -697,16 +730,41 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } std::cout << "Initializing Flux2 DirectX pipeline" << std::endl; const auto identity = din::common::ResolveOrtGraphicsDeviceIdentity(runtime_.trt_device); state_ = std::make_unique(runtime_.env, identity.luid); - initialize_dx_state(*state_, config_, runtime_.trt_device); + try + { + initialize_dx_state(*state_, config_, runtime_.trt_device); + } + catch (...) + { + state_.reset(); + throw; + } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_dx_models(*state_, requested); } void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -714,28 +772,49 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - state_->dx.begin(); - state_->dx.record_upload(state_->token_buf, tokens_cpu.data(), state_->token_buf.size); - state_->dx.record_upload(state_->attn_mask_buf, attn_mask_cpu.data(), state_->attn_mask_buf.size); - state_->dx.execute(); - const uint64_t prompt_uploaded = state_->sync_fence->signal_d3d12(state_->dx.queue.Get()); - state_->sync_fence->wait_ort(*state_->sync_stream, prompt_uploaded); + upload_dx_prompt(*state_, config_); } - Flux2Image GenerateImage(unsigned int seed) override + const void* TransformerSessionIdentity() const override { - if (!state_) + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; + } + + const void* TextEncoderSessionIdentity() const override + { + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override + { + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 DirectX pipeline was not initialized"); } - return run_dx_image(*state_, seed); + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + state_->applied_budget = budget; + } + return run_dx_image(*state_, seed, rng, progress); } private: diff --git a/image_gen/flux2/flux2_models.cpp b/image_gen/flux2/flux2_models.cpp new file mode 100644 index 0000000..749ed0f --- /dev/null +++ b/image_gen/flux2/flux2_models.cpp @@ -0,0 +1,93 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "flux2_models.h" + +void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& model_root, + const std::filesystem::path& cache_root, const std::filesystem::path& context_root) +{ + const auto root = model_root.empty() ? model_root : std::filesystem::weakly_canonical(model_root); + std::string prefix; + switch (config.processing) + { + case Flux2ProcessingBackend::Dx: + prefix = "dx"; + break; + case Flux2ProcessingBackend::DxCig: + prefix = "dx_cig"; + break; + case Flux2ProcessingBackend::Vk: + prefix = "vk"; + break; + case Flux2ProcessingBackend::VkCig: + prefix = "vk_cig"; + break; + default: + break; + } + if (!prefix.empty()) + prefix += '_'; + const std::string encoder = + config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator"; + const std::string transformer = "transformer_" + config.transformer.precision; + for (Flux2ModelConfig* model : + {static_cast(&config.text_encoder), static_cast(&config.transformer), + static_cast(&config.vae)}) + { + model->ep_cache_dir = std::filesystem::weakly_canonical(cache_root); + model->ep_context_dir = std::filesystem::weakly_canonical(context_root); + } + config.text_encoder.cache_subpath = prefix + encoder; + config.transformer.cache_subpath = + prefix + transformer + (config.transformer.initial_weight_streaming_budget.empty() ? "" : "_ws"); + config.vae.cache_subpath = prefix + "vae_decoder"; + if (root.empty()) + return; // The studio can start before a model root is selected. + config.text_encoder.model_path = root / encoder / "model.onnx"; + config.text_encoder.tokenizer_dir = root / "tokenizer"; + config.transformer.model_path = root / transformer / "model.onnx"; + config.vae.model_path = root / "vae_decoder/model.onnx"; + config.vae.latent_stats_path = root / "vae_decoder/latent_stats.json"; +} + +std::unique_ptr BindFlux2Model(Ort::Session& session, std::initializer_list inputs, + std::initializer_list outputs) +{ + auto binding = std::make_unique(session); + Ort::AllocatorWithDefaultOptions allocator; + auto bind = [&](std::initializer_list tensors, bool input) + { + const size_t count = input ? session.GetInputCount() : session.GetOutputCount(); + if (count != tensors.size()) + throw std::runtime_error("Unexpected model input/output count"); + for (const auto& tensor : tensors) + { + size_t index = 0; + for (; index < count; ++index) + { + const auto name = input ? session.GetInputNameAllocated(index, allocator) + : session.GetOutputNameAllocated(index, allocator); + if (std::string(name.get()) == tensor.name) + break; + } + if (index == count) + throw std::runtime_error(std::string("Missing tensor: ") + tensor.name); + const auto type = input ? session.GetInputTypeInfo(index) : session.GetOutputTypeInfo(index); + const auto info = type.GetTensorTypeAndShapeInfo(); + const auto actual = tensor.tensor->GetTensorTypeAndShapeInfo(); + const auto shape = info.GetShape(); + const auto expected = actual.GetShape(); + if (info.GetElementType() != actual.GetElementType() || shape.size() != expected.size()) + throw std::runtime_error(std::string("Incompatible tensor type/rank: ") + tensor.name); + for (size_t axis = 0; axis < shape.size(); ++axis) + if (shape[axis] >= 0 && shape[axis] != expected[axis]) + throw std::runtime_error(std::string("Incompatible tensor shape: ") + tensor.name); + if (input) + binding->BindInput(tensor.name, *tensor.tensor); + else + binding->BindOutput(tensor.name, *tensor.tensor); + } + }; + bind(inputs, true); + bind(outputs, false); + return binding; +} diff --git a/image_gen/flux2/flux2_models.h b/image_gen/flux2/flux2_models.h new file mode 100644 index 0000000..489f7b8 --- /dev/null +++ b/image_gen/flux2/flux2_models.h @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include + +#include "flux2_cli.h" +#include "ort_session.h" + +struct Flux2TensorBinding +{ + const char* name; + Ort::Value* tensor; +}; + +std::unique_ptr BindFlux2Model(Ort::Session& session, std::initializer_list inputs, + std::initializer_list outputs); + +struct Flux2LoadedModels +{ + std::optional text_encoder; + std::optional transformer; + std::optional vae; + bool Matches(const Flux2Config& config) const + { + return text_encoder == config.text_encoder && transformer == config.transformer && vae == config.vae; + } + bool Ready() const + { + return text_encoder && transformer && vae; + } +}; + +// The caller drains device work before entering this function. Tensors and the +// externally owned stream remain alive across session replacement. +template +bool ReloadFlux2Stage(const char* name, const Config& requested, std::optional& loaded, + std::unique_ptr& runner, std::unique_ptr& binding, + Ort::Env& env, const char* provider, Ort::SyncStream* stream, din::common::ModelProfile profile, + Bind&& bind) +{ + if (runner && binding && loaded == requested) + { + std::cout << "Reusing " << name << '\n'; + return false; + } + std::cout << "Loading " << name << ": " << requested.model_path.string() << '\n'; + binding.reset(); + runner.reset(); + loaded.reset(); + try + { + profile.cache_subpath = requested.cache_subpath; + runner = std::make_unique( + env, requested.model_path.string(), provider, requested.ep_cache_dir.string(), + din::common::EpContextOptions{requested.ep_context_dir.string()}, profile, stream); + binding = bind(runner->session); + loaded = requested; + } + catch (const std::exception& error) + { + binding.reset(); + runner.reset(); + loaded.reset(); + throw Flux2StageLoadError(std::string("Failed to load ") + name + " (" + requested.model_path.string() + + "): " + error.what()); + } + return true; +} diff --git a/image_gen/flux2/flux2_pipeline.cpp b/image_gen/flux2/flux2_pipeline.cpp index 3b69570..4443507 100644 --- a/image_gen/flux2/flux2_pipeline.cpp +++ b/image_gen/flux2/flux2_pipeline.cpp @@ -1,16 +1,19 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#include #include #include #include "flux2_runtime_context.h" +#include "rng.h" namespace { Flux2RuntimeContext& GetProcessRuntime(Flux2ExecutionProvider provider) { - static Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "Flux2"}; + static Ort::Env env{std::getenv("DIN_FLUX2_VERBOSE") ? ORT_LOGGING_LEVEL_VERBOSE : ORT_LOGGING_LEVEL_WARNING, + "Flux2"}; static Flux2RuntimeContext runtime{env}; static std::once_flag register_trt_once; @@ -28,6 +31,7 @@ Flux2RuntimeContext& GetProcessRuntime(Flux2ExecutionProvider provider) std::unique_ptr CreateFlux2Pipeline(const Flux2Config& config) { + ValidateFlux2Config(config); Flux2RuntimeContext& runtime = GetProcessRuntime(config.provider); if (config.processing == Flux2ProcessingBackend::Cpu || config.processing == Flux2ProcessingBackend::Cuda) @@ -55,3 +59,53 @@ std::unique_ptr CreateFlux2Pipeline(const Flux2Config& throw std::invalid_argument("Unknown Flux2 processing backend"); } + +bool IsFlux2BackendAvailable(Flux2ProcessingBackend backend) +{ + switch (backend) + { + case Flux2ProcessingBackend::Cpu: + case Flux2ProcessingBackend::Cuda: + return true; +#if defined(DIN_FLUX2_BUILD_DX) + case Flux2ProcessingBackend::Dx: + case Flux2ProcessingBackend::DxCig: + return true; +#endif +#if defined(DIN_FLUX2_BUILD_VK) + case Flux2ProcessingBackend::Vk: + case Flux2ProcessingBackend::VkCig: + return true; +#endif + default: + return false; + } +} + +void ValidateFlux2Config(const Flux2Config& config) +{ + if (!IsFlux2BackendAvailable(config.processing)) + throw std::invalid_argument("Selected processing backend was not built"); + if (config.provider == Flux2ExecutionProvider::Cpu && config.processing != Flux2ProcessingBackend::Cpu) + throw std::invalid_argument("CPU execution requires CPU processing"); + if (config.steps < 1 || config.steps > 50) + throw std::invalid_argument("Denoise steps must be in [1, 50]"); + if (config.num_images == 0) + throw std::invalid_argument("Image count must be positive"); + ValidateFlux2Rng(config.rng, config.processing); + if (config.transformer.precision != "bf16" && config.transformer.precision != "fp16" && + config.transformer.precision != "fp8" && config.transformer.precision != "nvfp4") + throw std::invalid_argument("Precision must be bf16, fp16, fp8, or nvfp4"); + if (!config.transformer.initial_weight_streaming_budget.empty()) + { + if (config.provider != Flux2ExecutionProvider::TrtRtx) + throw std::invalid_argument("Weight streaming requires TensorRT RTX"); + const auto& budget = config.transformer.initial_weight_streaming_budget; + if (budget == "-1") + return; + size_t parsed = 0; + const int value = std::stoi(budget, &parsed); + if (parsed + 1 != budget.size() || budget.back() != '%' || value < 0 || value > 100) + throw std::invalid_argument("Weight streaming budget must be 0% through 100%"); + } +} diff --git a/image_gen/flux2/flux2_telemetry.h b/image_gen/flux2/flux2_telemetry.h new file mode 100644 index 0000000..c985f4b --- /dev/null +++ b/image_gen/flux2/flux2_telemetry.h @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include + +#include "flux2_cli.h" + +// Stage timings are synchronized wall times, including stage transfers. +// Synchronization is supplied by each backend to respect its ORT/graphics ordering. +class Flux2Telemetry +{ +public: + Flux2Telemetry(int steps, Flux2Progress progress, std::function synchronize) + : steps_(steps) + , progress_(std::move(progress)) + , synchronize_(std::move(synchronize)) + , start_(Clock::now()) + , stage_start_(start_) + { + Report("rng", 0); + } + + void Begin(const char* stage, double Flux2Timings::* field) + { + synchronize_(); + const auto now = Clock::now(); + timings_.*field_ += Milliseconds(now - stage_start_); + stage_start_ = now; + field_ = field; + Report(stage, 0); + } + + void Step(int completed) + { + if (progress_) + { + synchronize_(); + Report("denoise", completed); + } + } + + Flux2Timings Finish() + { + synchronize_(); + const auto now = Clock::now(); + timings_.*field_ += Milliseconds(now - stage_start_); + timings_.total_ms = Milliseconds(now - start_); + timings_.per_step_ms = timings_.denoise_ms / steps_; + Report("done", steps_); + return timings_; + } + +private: + using Clock = std::chrono::steady_clock; + + static double Milliseconds(Clock::duration duration) + { + return std::chrono::duration(duration).count(); + } + + void Report(const char* stage, int completed) + { + if (progress_) + progress_(stage, completed, steps_); + } + + int steps_; + Flux2Progress progress_; + std::function synchronize_; + Clock::time_point start_, stage_start_; + Flux2Timings timings_; + double Flux2Timings::* field_ = &Flux2Timings::rng_ms; +}; diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp new file mode 100644 index 0000000..f6ee8d5 --- /dev/null +++ b/image_gen/flux2/flux2_ui.cpp @@ -0,0 +1,552 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "flux2.h" +#include "flux2_cli.h" +#include "io/image.h" +#include "utils.h" +#include +#include +#include +#include +#include + +namespace +{ +// Stable per-root namespace prevents a different model directory from reusing stale engines. +std::string CacheNamespace(const std::filesystem::path& root) +{ + const auto path = std::filesystem::weakly_canonical(root).generic_u8string(); + uint64_t hash = 14695981039346656037ull; + for (const auto c : path) + { + hash ^= static_cast(c); + hash *= 1099511628211ull; + } + return "root_" + std::to_string(hash); +} + +struct Request +{ + Flux2Config config; + Flux2GenerationOptions generation; + bool random_seeds = false; +}; + +struct Result +{ + Flux2Image image; + unsigned int seed = 0; + GLuint texture = 0; +}; + +// Only the worker touches the pipeline; all graphics resources belong to the UI thread. +class Generator +{ +public: + Generator() + : worker_( + [this] + { + Run(); + }) + { + } + + ~Generator() + { + { + std::lock_guard lock(mutex); + closing_ = true; + } + cv_.notify_one(); + worker_.join(); + } + + void Submit(Request request) + { + std::lock_guard lock(mutex); + if (busy) + return; + busy = true; + failed = false; + status = "Loading models (first compilation can take several minutes)"; + completed = 0; + total = request.config.steps * static_cast(request.config.num_images); + request_ = std::move(request); + cv_.notify_one(); + } + + std::mutex mutex; + bool busy = false, failed = false; + std::string status = "Ready"; + int completed = 0, total = 1; + size_t used_mb = 0, total_mb = 0; + std::vector ready; + +private: + static bool SameBackend(const Flux2Config& a, const Flux2Config& b) + { + return a.processing == b.processing && a.provider == b.provider; + } + + void UpdateMemory() + { + size_t free = 0, total_bytes = 0; + if (cudaMemGetInfo(&free, &total_bytes) == cudaSuccess) + { + std::lock_guard lock(mutex); + used_mb = (total_bytes - free) / (1024 * 1024); + total_mb = total_bytes / (1024 * 1024); + } + else + cudaGetLastError(); + } + + void Run() + { + std::unique_ptr pipeline; + std::optional loaded; + std::mt19937 random(std::random_device{}()); + for (;;) + { + std::unique_lock lock(mutex); + cv_.wait_for(lock, std::chrono::milliseconds(500), + [&] + { + return closing_ || request_.has_value(); + }); + if (closing_) + break; + if (!request_) + { + lock.unlock(); + if (loaded && loaded->provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + continue; + } + Request request = std::move(*request_); + request_.reset(); + lock.unlock(); + try + { + const auto& config = request.config; + ValidateFlux2Config(config); + if (!loaded || !SameBackend(*loaded, config)) + { + pipeline.reset(); // Release old engines before allocating the replacement. + loaded.reset(); + pipeline = CreateFlux2Pipeline(config); + loaded = config; + pipeline->Initialize(); + } + pipeline->UpdateModels(config); + loaded = config; + pipeline->SetPrompt(config.prompt); + for (unsigned int i = 0; i < config.num_images; ++i) + { + const unsigned int seed = request.random_seeds ? random() : config.seed + i; + auto image = pipeline->GenerateImage( + seed, + [&](const char* stage, int step, int) + { + { + std::lock_guard guard(mutex); + completed = static_cast(i) * config.steps + step; + status = "Image " + std::to_string(i + 1) + "/" + std::to_string(config.num_images) + + ": " + stage; + } + if (config.provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + }, + request.generation); + if (!config.output_path.empty()) + { + std::filesystem::create_directories(config.output_path); + const auto file = + config.output_path / ("flux2_" + std::to_string(i) + "_" + std::to_string(seed) + ".png"); + if (!din::io::SaveRgbFloatImage(file.string(), image.data.data(), image.height, image.width, + din::io::ImageValueRange::MinusOneToOne)) + throw std::runtime_error("Could not save " + file.string()); + } + std::lock_guard guard(mutex); + ready.push_back({std::move(image), seed, 0}); + } + std::lock_guard guard(mutex); + completed = total; + status = "Done"; + } + catch (const Flux2StageLoadError& error) + { + std::lock_guard guard(mutex); + status = std::string("Error: ") + error.what(); + failed = true; + } + catch (const std::exception& error) + { + pipeline.reset(); + loaded.reset(); + std::lock_guard guard(mutex); + status = std::string("Error: ") + error.what(); + failed = true; + } + std::lock_guard guard(mutex); + busy = false; + } + // Destroy sessions, streams and GPU allocations on their owning thread. + } + + std::condition_variable cv_; + bool closing_ = false; + std::optional request_; + std::thread worker_; +}; + +void Upload(Result& result) +{ + const auto& img = result.image; + const size_t pixels = static_cast(img.width) * img.height; + if (!pixels || img.data.size() != pixels * 3) + throw std::runtime_error("Invalid image returned by pipeline"); + std::vector rgb(pixels * 3); + for (size_t p = 0; p < pixels; ++p) + for (size_t c = 0; c < 3; ++c) + { + const float v = img.data[c * pixels + p]; + rgb[p * 3 + c] = + std::isfinite(v) + ? static_cast(std::lround(std::clamp(v * 127.5f + 127.5f, 0.0f, 255.0f))) + : 0; + } + glGenTextures(1, &result.texture); + glBindTexture(GL_TEXTURE_2D, result.texture); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR); + glPixelStorei(GL_UNPACK_ALIGNMENT, 1); + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, rgb.data()); +} + +int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit_after) +{ + Generator generator; + Flux2Config config = initial; + std::array model{}, output{}, cache{}; + bool cache_changed = false; + const auto cache_utf8 = initial.text_encoder.ep_cache_dir.parent_path().u8string(); + std::snprintf(cache.data(), cache.size(), "%s", reinterpret_cast(cache_utf8.c_str())); + std::array prompt{}; + const auto model_utf8 = config.text_encoder.model_path.parent_path().parent_path().u8string(); + std::snprintf(model.data(), model.size(), "%s", reinterpret_cast(model_utf8.c_str())); + std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); + std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); + int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. + int budget = 50; + int count = static_cast(config.num_images); + int encoder = config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B ? 0 : 1; + int precision = config.transformer.precision == "bf16" ? 0 + : config.transformer.precision == "fp16" ? 1 + : config.transformer.precision == "fp8" ? 2 + : 3; + const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; + const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; + int backend = static_cast(config.processing); + int provider = config.provider == Flux2ExecutionProvider::Cpu ? 0 : 1; + std::vector results; + int selected = -1, frames = 0; + bool launched = false; + int exit_code = 0; + GLFWwindow* window = glfwGetCurrentContext(); + while (!glfwWindowShouldClose(window)) + { + glfwPollEvents(); + ImGui_ImplOpenGL3_NewFrame(); + ImGui_ImplGlfw_NewFrame(); + ImGui::NewFrame(); + ImGui::SetNextWindowPos({0, 0}); + ImGui::SetNextWindowSize(ImGui::GetIO().DisplaySize); + ImGui::Begin("Flux2 Studio", nullptr, + ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize); + ImGui::TextUnformatted("DIN Deploy - FLUX.2-klein Interactive Studio"); + ImGui::Separator(); + bool busy, failed; + int completed, total; + std::string status; + size_t used_mb, total_mb; + { + std::lock_guard lock(generator.mutex); + busy = generator.busy; + failed = generator.failed; + status = generator.status; + completed = generator.completed; + total = generator.total; + used_mb = generator.used_mb; + total_mb = generator.total_mb; + for (auto& result : generator.ready) + { + results.push_back(std::move(result)); + selected = static_cast(results.size()) - 1; + } + generator.ready.clear(); + } + if (launched && !busy && exit_after) + { + exit_code = failed ? 1 : 0; + glfwSetWindowShouldClose(window, true); + } + ImGui::BeginChild("controls", ImVec2(430, 0), true); + ImGui::BeginDisabled(busy); + ImGui::InputText("Model root", model.data(), model.size()); + cache_changed |= ImGui::InputText("Cache root", cache.data(), cache.size()); + ImGui::InputText("Save directory", output.data(), output.size()); + ImGui::TextUnformatted("Prompt"); + ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); + ImGui::RadioButton("Qwen3-4B", &encoder, 0); + const bool translator_exists = + model[0] && std::filesystem::is_regular_file(std::filesystem::u8path(model.data()) / + "text_encoder_translator/model.onnx"); + ImGui::BeginDisabled(!translator_exists); + ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); + ImGui::EndDisabled(); + if (!translator_exists) + ImGui::TextWrapped("Translator ONNX is missing from this model root."); + ImGui::Combo("Precision", &precision, precisions, IM_ARRAYSIZE(precisions)); + ImGui::Combo("Provider", &provider, "CPU\0TensorRT RTX\0"); + if (provider == 0) + { + backend = 0; + streaming_mode = 0; + } + if (ImGui::BeginCombo("Processing", backends[backend])) + { + for (int i = 0; i < 6; ++i) + { + if (!IsFlux2BackendAvailable(static_cast(i)) || (provider == 0 && i != 0)) + continue; + if (ImGui::Selectable(backends[i], backend == i)) + backend = i; + } + ImGui::EndCombo(); + } + if (backend != static_cast(Flux2ProcessingBackend::Cuda) && config.rng == din::common::Rng::CudaPhilox) + config.rng = din::common::Rng::CpuPhilox; + const char* rng_names[] = {"cpu-std", "cuda-philox", "cpu-philox"}; + if (ImGui::BeginCombo("RNG", rng_names[static_cast(config.rng)])) + { + for (int i = 0; i < 3; ++i) + { + const bool disabled = i == 1 && backend != static_cast(Flux2ProcessingBackend::Cuda); + ImGui::BeginDisabled(disabled); + if (ImGui::Selectable(rng_names[i], static_cast(config.rng) == i)) + config.rng = static_cast(i); + ImGui::EndDisabled(); + } + ImGui::EndCombo(); + } + ImGui::BeginDisabled(provider == 0); + ImGui::Combo("Weight streaming", &streaming_mode, "Auto\0Manual\0Off\0"); + if (streaming_mode == 1) + ImGui::SliderInt("Resident weights %", &budget, 10, 100); + ImGui::EndDisabled(); + ImGui::InputScalar("Seed", ImGuiDataType_U32, &config.seed); + ImGui::SliderInt("Denoise steps", &config.steps, 1, 50); + ImGui::SliderInt("Images", &count, 1, 16); + const bool paths_valid = model[0] && (!cache_changed || cache[0]); + ImGui::BeginDisabled(!paths_valid); + bool generate = ImGui::Button("Generate", ImVec2(-1, 35)); + bool random = ImGui::Button("Random sweep", ImVec2(-1, 30)); + ImGui::EndDisabled(); + ImGui::EndDisabled(); + if (!busy && paths_valid && (generate || random || (auto_generate && !launched))) + { + for (auto& result : results) + if (result.texture) + glDeleteTextures(1, &result.texture); + results.clear(); + selected = -1; + const auto model_root = std::filesystem::u8path(model.data()); + config.output_path = output.data(); + config.prompt = prompt.data(); + const auto cache_namespace = CacheNamespace(model_root); + const auto cache_root = (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" + : initial.text_encoder.ep_cache_dir) / + cache_namespace; + const auto context_root = (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" + : initial.text_encoder.ep_context_dir) / + cache_namespace; + config.text_encoder.variant = + encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + config.transformer.precision = precisions[precision]; + config.num_images = static_cast(count); + config.processing = static_cast(backend); + config.provider = provider == 0 ? Flux2ExecutionProvider::Cpu : Flux2ExecutionProvider::TrtRtx; + config.transformer.initial_weight_streaming_budget = provider == 0 ? "" : "-1"; + ConfigureFlux2Models(config, model_root, cache_root, context_root); + Flux2GenerationOptions generation; + generation.steps = config.steps; + generation.rng = config.rng; + if (provider != 0 && streaming_mode != 0) + generation.weight_streaming_budget = streaming_mode == 1 ? std::to_string(budget) + "%" : "0"; + generator.Submit({config, generation, random}); + launched = true; + } + ImGui::ProgressBar(static_cast(completed) / std::max(total, 1), ImVec2(-1, 0)); + ImGui::TextWrapped("%s", status.c_str()); + if (provider != 0 && total_mb) + ImGui::Text("Device memory: %zu / %zu MiB", used_mb, total_mb); + if (selected >= 0) + { + const auto& t = results[selected].image.timings; + ImGui::Separator(); + ImGui::Text("Encode %.0f ms | RNG %.0f ms", t.encode_ms, t.rng_ms); + ImGui::Text("Denoise %.0f ms (%.0f ms/step)", t.denoise_ms, t.per_step_ms); + ImGui::Text("Decode %.0f ms | Total %.2f s", t.decode_ms, t.total_ms / 1000); + } + ImGui::EndChild(); + ImGui::SameLine(); + ImGui::BeginChild("images", ImVec2(0, 0), true); + for (auto& result : results) + if (!result.texture) + Upload(result); + ImGui::BeginChild("preview", ImVec2(0, -115), false); + if (selected >= 0) + { + const auto available = ImGui::GetContentRegionAvail(); + const float side = std::max(1.0f, std::min(available.x, available.y)); + ImGui::Image(static_cast(results[selected].texture), ImVec2(side, side)); + } + else + ImGui::TextDisabled("Generated images will appear here."); + ImGui::EndChild(); + ImGui::BeginChild("thumbnails", ImVec2(0, 110), false, ImGuiWindowFlags_HorizontalScrollbar); + for (int i = 0; i < static_cast(results.size()); ++i) + { + if (i) + ImGui::SameLine(); + ImGui::PushID(i); + ImGui::BeginGroup(); + if (ImGui::ImageButton("image", static_cast(results[i].texture), ImVec2(72, 72))) + selected = i; + ImGui::Text("%u", results[i].seed); + ImGui::EndGroup(); + ImGui::PopID(); + } + ImGui::EndChild(); + ImGui::EndChild(); + ImGui::End(); + ImGui::Render(); + int width, height; + glfwGetFramebufferSize(window, &width, &height); + glViewport(0, 0, width, height); + glClearColor(0.09f, 0.09f, 0.10f, 1); + glClear(GL_COLOR_BUFFER_BIT); + ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData()); + glfwSwapBuffers(window); + if (smoke && ++frames >= 3) + break; + } + for (auto& result : results) + if (result.texture) + glDeleteTextures(1, &result.texture); + return exit_code; +} +} // namespace + +int main(int argc, char** argv) +{ + GLFWwindow* window = nullptr; + bool imgui = false, glfw_backend = false, gl_backend = false; + int result = 1; + try + { + argparse::ArgumentParser args("din_flux2_ui"); +#ifdef _WIN32 + const wchar_t* local_app_data = _wgetenv(L"LOCALAPPDATA"); + if (!local_app_data || !*local_app_data) + throw std::runtime_error("LOCALAPPDATA is not set; cannot determine the default cache directory"); + const auto cache_root = std::filesystem::path(local_app_data) / "din_deploy"; +#else + const auto cache_root = std::filesystem::path("/tmp/din_deploy/cache"); +#endif + args.add_argument("--model-dir").default_value(std::string{}); + args.add_argument("--output").default_value(std::string{}); + const auto runtime_cache_utf8 = (cache_root / "runtime").u8string(); + const auto context_cache_utf8 = (cache_root / "ep_context").u8string(); + args.add_argument("--ep-cache") + .default_value(std::string(reinterpret_cast(runtime_cache_utf8.c_str()))); + args.add_argument("--ep-context-dir") + .default_value(std::string(reinterpret_cast(context_cache_utf8.c_str()))); + args.add_argument("--prompt").default_value(std::string(DEFAULT_PROMPT)); + args.add_argument("--encoder").default_value(std::string("auto")).choices("auto", "4b", "translator"); + args.add_argument("--precision").default_value(std::string("bf16")).choices("bf16", "fp16", "fp8", "nvfp4"); + args.add_argument("--smoke-test").default_value(false).implicit_value(true); + args.add_argument("--auto-generate").default_value(false).implicit_value(true); + args.add_argument("--exit-after-generation").default_value(false).implicit_value(true); + args.parse_args(argc, argv); + Flux2Config config; + const std::filesystem::path model_root = args.get("--model-dir"); + if (model_root.empty() && args.get("--auto-generate")) + throw std::runtime_error("--auto-generate requires --model-dir"); + config.output_path = args.get("--output"); + config.text_encoder.ep_cache_dir = std::filesystem::u8path(args.get("--ep-cache")); + config.text_encoder.ep_context_dir = std::filesystem::u8path(args.get("--ep-context-dir")); + std::filesystem::create_directories(config.text_encoder.ep_cache_dir); + std::filesystem::create_directories(config.text_encoder.ep_context_dir); + config.prompt = args.get("--prompt"); + config.transformer.precision = args.get("--precision"); + config.num_images = 1; + const auto encoder = args.get("--encoder"); + const bool translator = encoder == "translator" || + (encoder == "auto" && !model_root.empty() && + std::filesystem::is_regular_file(model_root / "text_encoder_translator/model.onnx")); + config.text_encoder.variant = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; + config.transformer.initial_weight_streaming_budget = "-1"; + ConfigureFlux2Models(config, model_root, config.text_encoder.ep_cache_dir, config.text_encoder.ep_context_dir); + if (!glfwInit()) + throw std::runtime_error("GLFW initialization failed"); + glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3); + glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 2); + window = glfwCreateWindow(1280, 850, "DIN Flux2 Studio", nullptr, nullptr); + if (!window) + throw std::runtime_error("OpenGL 3.2 window creation failed"); + glfwMakeContextCurrent(window); + glfwSwapInterval(1); + IMGUI_CHECKVERSION(); + ImGui::CreateContext(); + imgui = true; + ImGui::GetIO().IniFilename = nullptr; + ImGui::StyleColorsDark(); + glfw_backend = ImGui_ImplGlfw_InitForOpenGL(window, true); + gl_backend = ImGui_ImplOpenGL3_Init("#version 130"); + if (!glfw_backend || !gl_backend) + throw std::runtime_error("ImGui graphics initialization failed"); + result = Studio(config, args.get("--smoke-test"), args.get("--auto-generate"), + args.get("--exit-after-generation")); + } + catch (const std::exception& error) + { + std::cerr << error.what() << '\n'; + } + if (gl_backend) + ImGui_ImplOpenGL3_Shutdown(); + if (glfw_backend) + ImGui_ImplGlfw_Shutdown(); + if (imgui) + ImGui::DestroyContext(); + if (window) + glfwDestroyWindow(window); + glfwTerminate(); + return result; +} diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index 865f564..7eb3b97 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -26,7 +26,9 @@ #endif #include "flux2.h" +#include "flux2_models.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" #include "vk_helper.h" @@ -38,14 +40,6 @@ static const std::string POSTPROCESS_SHADER_PATH = std::string(SHADER_DIR) + "po // Upload latent noise for a new seed (CPU → staging → GPU via Vulkan copy) // ============================================================================ -void initialize_latents(unsigned int seed, size_t hidden_count, void* dest) -{ - nvtx3::scoped_range nvtx("initialize_latents"); - std::cout << "\n=== Initializing Latent (seed=" << seed << ") ===" << std::endl; - - initialize_latent(static_cast(dest), hidden_count, seed); -} - class OrtVulkanGraphicsInteropScope { public: @@ -438,17 +432,23 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor VkCommandBuffer cmd_postprocess, ComputePipelineResources& euler_shader, ComputePipelineResources& postprocess_shader, VulkanBuffer& hidden_states_buf, VulkanBuffer& transformer_output_buf, VulkanBuffer& decoder_latent_buf, StagingBuffer& staging, - bool encode_prompt) + bool encode_prompt, Flux2Telemetry& telemetry) { nvtx3::scoped_range nvtx_pipeline("run_pipeline"); + Ort::RunOptions transformer_options; + transformer_options.SetSyncStream(ort_stream); + transformer_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + Ort::RunOptions run_options; run_options.SetSyncStream(ort_stream); run_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); // --- Text Encoder --- + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); nvtx3::scoped_range nvtx("text_encoder"); text_encoder_session.Run(run_options, text_encoder_io); } @@ -482,7 +482,8 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor sync_semaphore.wait_ort(ort_stream, timestep_done); } - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { nvtx3::scoped_range nvtx_step("diffusion_step_" + std::to_string(step)); @@ -491,7 +492,7 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor { nvtx3::scoped_range nvtx_transformer("transformer"); - transformer_session.Run(run_options, transformer_io); + transformer_session.Run(transformer_options, transformer_io); } { @@ -527,7 +528,7 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { vkCmdUpdateBuffer(cmd_euler[step], timestep_buf.buffer, 0, sizeof(float), &t_next); } @@ -538,8 +539,10 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor const uint64_t euler_done = sync_semaphore.submit_vulkan(vk, cmd_euler[step], transformer_done); sync_semaphore.wait_ort(ort_stream, euler_done); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); // --- Post-process Latents (Vulkan compute) --- { nvtx3::scoped_range nvtx_post("postprocess"); @@ -685,8 +688,8 @@ struct VkPipelineState Ort::Value decoded_image_tensor{nullptr}; VkCommandPool command_pool = VK_NULL_HANDLE; - std::array cmd_bufs_euler{}; - std::array cmd_bufs_time{}; + std::vector cmd_bufs_euler{}; + std::vector cmd_bufs_time{}; VkCommandBuffer cmd_upload_lat = VK_NULL_HANDLE; VkCommandBuffer cmd_download = VK_NULL_HANDLE; VkCommandBuffer cmd_postprocess = VK_NULL_HANDLE; @@ -697,6 +700,9 @@ struct VkPipelineState std::unique_ptr transformer_io; std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; }; static void initialize_vk_state(VkPipelineState& state, const Flux2Config& config, Ort::ConstEpDevice trt_device) @@ -705,14 +711,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { throw std::runtime_error("Vulkan processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; - const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk"); - - std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" - << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" - << std::endl; // ----------------------------------------------------------------- // Init Vulkan @@ -742,11 +741,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Load ONNX models // ----------------------------------------------------------------- - std::cout << "=== Loading ONNX Models ===" << std::endl; - auto nvtx_scope_load = nvtx3::start_range("load_onnx_models"); - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); std::vector> graphics_ep_options; if (state.use_cig) { @@ -767,15 +761,9 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // graphics_ep_options.emplace_back("nv_use_sync_gpu_allocator", "1"); } - auto make_profile = [&](std::string cache_subpath) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.embed_ep_context = false; - profile.enable_cuda_graph = !state.use_cig; - profile.extra_ep_options = graphics_ep_options; - return profile; - }; + state.model_profile.enable_cuda_graph = !state.use_cig; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options = graphics_ep_options; // ----------------------------------------------------------------- // Allocate tensors from Vulkan buffers imported through ORT's external @@ -868,21 +856,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi state.bn_mean_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); state.bn_std_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); - { - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, BN_MEAN, LATENT_CHANNELS * sizeof(float)); - VkBufferCopy region = {0, 0, state.bn_mean_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_mean_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); - } - { - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, BN_STD, LATENT_CHANNELS * sizeof(float)); - VkBufferCopy region = {0, 0, state.bn_std_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_std_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); - } - state.postprocess_shader = loadComputeShader(POSTPROCESS_SHADER_PATH, *state.vk, sizeof(PostprocessPushConstants), {&state.hidden_states, &state.decoder_latent, &state.bn_mean_buf, &state.bn_std_buf}); @@ -901,31 +874,21 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // Sessions are intentionally last: their CIG context sees the fully initialized importer, // external queue, shared buffers, and timeline semaphore. - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.text_encoder), &*state.sync_stream); - std::cout << " Text encoder loaded" << std::endl; - - state.transformer_runner = std::make_unique( - state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer), &*state.sync_stream); - std::cout << " Transformer loaded" << std::endl; - - state.vae_decoder_runner = std::make_unique( - state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.vae_decoder), &*state.sync_stream); - std::cout << " VAE decoder loaded" << std::endl; - nvtx3::end_range(nvtx_scope_load); // ----------------------------------------------------------------- // Persistent command buffers (pool already created above) // ----------------------------------------------------------------- // Euler sampler - state.vk->device.allocateCommandBuffers(state.command_pool, FLOW_STEPS, state.cmd_bufs_euler.data()); + // Reserve all supported steps so changing the schedule needs no model reload + // or command-buffer allocation while generation is in progress. + constexpr unsigned int max_steps = 50; + state.cmd_bufs_euler.resize(max_steps); + state.cmd_bufs_time.resize(max_steps); + state.vk->device.allocateCommandBuffers(state.command_pool, max_steps, state.cmd_bufs_euler.data()); // Time scheduler - state.vk->device.allocateCommandBuffers(state.command_pool, FLOW_STEPS, state.cmd_bufs_time.data()); + state.vk->device.allocateCommandBuffers(state.command_pool, max_steps, state.cmd_bufs_time.data()); VkCommandBuffer cmd_bufs[3]; state.vk->device.allocateCommandBuffers(state.command_pool, 3, cmd_bufs); @@ -938,22 +901,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Create persistent IoBindings (all GPU tensors, reused every run) // ----------------------------------------------------------------- - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token_tensor); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask_tensor); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds_tensor); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states_tensor); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds_tensor); - state.transformer_io->BindInput("timestep", state.timestep_tensor); - state.transformer_io->BindInput("img_ids", state.img_ids_tensor); - state.transformer_io->BindInput("txt_ids", state.txt_ids_tensor); - state.transformer_io->BindOutput("sample", state.transformer_output_tensor); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_input_tensor); - state.vae_decoder_io->BindOutput("sample", state.decoded_image_tensor); // ----------------------------------------------------------------- // One-time upload: position IDs (via staging buffer) @@ -1001,55 +948,115 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi } // ----------------------------------------------------------------- - // Time schedule (computed once, constant across images) + // Initial time schedule; generation options can replace it per image. // ----------------------------------------------------------------- - constexpr double MU = 2.291179894115571; - state.time_schedule.resize(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - state.time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << state.time_schedule[i] << std::endl; - } + state.time_schedule = MakeFlux2Schedule(config.steps); // ----------------------------------------------------------------- // Load prompt tokens & upload to GPU (once for all images) // ----------------------------------------------------------------- - std::cout << "\n=== Tokenizing Prompt ===" << std::endl; - { - nvtx3::scoped_range nvtx("tokenize"); - - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); - - std::vector tokens_cpu(shape_numel(token_shape)); - std::vector attn_mask_cpu(shape_numel(attn_mask_shape)); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); +} - memcpy(state.staging.mapped, tokens_cpu.data(), state.token_buf.size); - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - VkBufferCopy region = {0, 0, state.token_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.token_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +static void upload_vk_prompt(VkPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), + BATCH_SIZE, SEQUENCE_LENGTH); + + memcpy(state.staging.mapped, tokens_cpu.data(), state.token_buf.size); + VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); + VkBufferCopy region = {0, 0, state.token_buf.size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, state.token_buf.buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); + + memcpy(state.staging.mapped, attn_mask_cpu.data(), state.attn_mask_buf.size); + cmd = state.vk->beginCommandBuffer(state.command_pool); + region = {0, 0, state.attn_mask_buf.size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, state.attn_mask_buf.buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +} - memcpy(state.staging.mapped, attn_mask_cpu.data(), state.attn_mask_buf.size); - cmd = state.vk->beginCommandBuffer(state.command_pool); - region = {0, 0, state.attn_mask_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.attn_mask_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +static void update_vk_models(VkPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + signal_and_wait_on_host(*state.sync_stream); + VK_CHECK(vkDeviceWaitIdle(state.vk->device.raw())); } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = &*state.sync_stream; + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model( + session, + {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, + {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_vk_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model(session, + {{"hidden_states", &state.hidden_states_tensor}, + {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, + {"timestep", &state.timestep_tensor}, + {"img_ids", &state.img_ids_tensor}, + {"txt_ids", &state.txt_ids_tensor}}, + {{"sample", &state.transformer_output_tensor}}); + })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, + {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + for (const auto& upload : {std::pair{&state.bn_mean_buf, stats.mean.data()}, + std::pair{&state.bn_std_buf, stats.std.data()}}) + { + memcpy(state.staging.mapped, upload.second, upload.first->size); + const auto cmd = state.vk->beginCommandBuffer(state.command_pool); + VkBufferCopy region = {0, 0, upload.first->size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, upload.first->buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); + } + return binding; + }); } -static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed) +static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { - initialize_latent(static_cast(state.staging.mapped), state.hidden_count, seed); + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + signal_and_wait_on_host(*state.sync_stream); + VK_CHECK(vkQueueWaitIdle(state.vk->queue)); + }); + initialize_cpu_latent(static_cast(state.staging.mapped), state.hidden_count, seed, rng, BATCH_SIZE, + LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); const uint64_t vae_done = run_pipeline( state.text_encoder_runner->session, state.transformer_runner->session, state.vae_decoder_runner->session, *state.text_encoder_io, *state.transformer_io, *state.vae_decoder_io, state.time_schedule, state.timestep_buf, *state.vk, *state.sync_semaphore, *state.sync_stream, state.cmd_upload_lat, state.cmd_bufs_time.data(), state.cmd_bufs_euler.data(), state.cmd_postprocess, state.euler_shader, state.postprocess_shader, - state.hidden_states, state.transformer_output, state.decoder_latent, state.staging, !state.prompt_embeds_valid); + state.hidden_states, state.transformer_output, state.decoder_latent, state.staging, !state.prompt_embeds_valid, + telemetry); state.prompt_embeds_valid = true; vkResetCommandBuffer(state.cmd_download, 0); @@ -1069,6 +1076,7 @@ static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed) image.width = static_cast(IMAGE_WIDTH); image.data.resize(state.decoded_image_bytes / sizeof(float)); std::copy_n(static_cast(state.staging.mapped), image.data.size(), image.data.data()); + image.timings = telemetry.Finish(); return image; } @@ -1085,15 +1093,40 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } std::cout << "Initializing Flux2 Vulkan pipeline" << std::endl; state_ = std::make_unique(runtime_.env); - initialize_vk_state(*state_, config_, runtime_.trt_device); + try + { + initialize_vk_state(*state_, config_, runtime_.trt_device); + } + catch (...) + { + state_.reset(); + throw; + } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_vk_models(*state_, requested); } void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -1101,33 +1134,49 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - memcpy(state_->staging.mapped, tokens_cpu.data(), state_->token_buf.size); - VkCommandBuffer cmd = state_->vk->beginCommandBuffer(state_->command_pool); - VkBufferCopy region = {0, 0, state_->token_buf.size}; - vkCmdCopyBuffer(cmd, state_->staging.buffer, state_->token_buf.buffer, 1, ®ion); - state_->vk->endAndSubmitCommandBuffer(cmd, state_->command_pool); - - memcpy(state_->staging.mapped, attn_mask_cpu.data(), state_->attn_mask_buf.size); - cmd = state_->vk->beginCommandBuffer(state_->command_pool); - region = {0, 0, state_->attn_mask_buf.size}; - vkCmdCopyBuffer(cmd, state_->staging.buffer, state_->attn_mask_buf.buffer, 1, ®ion); - state_->vk->endAndSubmitCommandBuffer(cmd, state_->command_pool); + upload_vk_prompt(*state_, config_); } - Flux2Image GenerateImage(unsigned int seed) override + const void* TransformerSessionIdentity() const override { - if (!state_) + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; + } + + const void* TextEncoderSessionIdentity() const override + { + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override + { + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 Vulkan pipeline was not initialized"); } - return run_vk_image(*state_, seed); + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + state_->applied_budget = budget; + } + return run_vk_image(*state_, seed, rng, progress); } private: diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 9ef75e9..d5f08bc 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -9,24 +9,20 @@ #include #include #include +#include #include -#include #include #include #include +#include "flux2_cli.h" +#include "rng.h" #include "tokenizer.h" // ============================================================================ // Configuration // ============================================================================ -#ifdef _WIN32 -static const std::filesystem::path DEFAULT_MODEL_BASE_PATH = "S:/din_deploy_artifacts/FLUX.2-klein-4B-onnx"; -#else -static const std::filesystem::path DEFAULT_MODEL_BASE_PATH = "/mnt/share/onnx/FLUX.2-klein-4B-onnx/"; -#endif - struct Flux2ModelPaths { std::filesystem::path base_dir; @@ -36,36 +32,10 @@ struct Flux2ModelPaths std::filesystem::path tokenizer_dir; }; -struct Flux2ModelCachePaths -{ - std::string text_encoder; - std::string transformer; - std::string vae_decoder; -}; - -inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, const std::string& precision) -{ - if (model_dir.empty()) - { - model_dir = DEFAULT_MODEL_BASE_PATH; - } - return { - model_dir, - model_dir / "text_encoder/model.onnx", - model_dir / ("transformer_" + precision) / "model.onnx", - model_dir / "vae_decoder/model.onnx", - model_dir / "tokenizer", - }; -} - -inline Flux2ModelCachePaths MakeFlux2ModelCachePaths(const std::string& precision, const std::string& prefix = {}) +inline Flux2ModelPaths MakeFlux2ModelPaths(const Flux2Config& config) { - const std::string prefix_separator = prefix.empty() ? "" : prefix + "_"; - return { - prefix_separator + "text_encoder", - prefix_separator + "transformer_" + precision, - prefix_separator + "vae_decoder", - }; + return {config.text_encoder.model_path.parent_path().parent_path(), config.text_encoder.model_path, + config.transformer.model_path, config.vae.model_path, config.text_encoder.tokenizer_dir}; } static const std::filesystem::path TEXT_ENCODER_MODEL_CTX = "text_encoder_model_ctx.onnx"; @@ -114,12 +84,36 @@ struct Flux2TextEncoderInputs int64_t pad_token_id = 0; }; +inline std::string FormatFlux2Prompt(const std::string& prompt) +{ + return "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; +} + +inline std::vector MakeFlux2Schedule(int steps) +{ + if (steps < 1 || steps > 50) + throw std::invalid_argument("Denoise steps must be in [1, 50]"); + constexpr double m200 = 0.00016927 * IMAGE_SEQUENCE + 0.45666666; + constexpr double m10 = 8.73809524e-05 * IMAGE_SEQUENCE + 1.89833333; + constexpr double a = (m200 - m10) / 190.0; + const double mu = IMAGE_SEQUENCE > 4300 ? m200 : a * steps + m200 - 200.0 * a; + const double e = std::exp(mu); + std::vector schedule(steps + 1); + for (int i = 0; i < steps; ++i) + { + const double t = 1.0 - static_cast(i) / steps; + schedule[i] = static_cast(e / (e + 1.0 / t - 1.0)); + } + schedule.back() = 0.0f; + return schedule; +} + inline Flux2TextEncoderInputs TokenizeFlux2Prompt(const Flux2ModelPaths& model_paths, const std::string& prompt) { din::io::Tokenizer tokenizer((model_paths.tokenizer_dir / "tokenizer.json").string(), din::io::TokenizerFormat::Json); - const std::string chat_text = "<|im_start|>user\n" + prompt + "<|im_end|>"; + const std::string chat_text = FormatFlux2Prompt(prompt); Flux2TextEncoderInputs inputs; inputs.token_ids = tokenizer.Encode(chat_text, false); inputs.pad_token_id = tokenizer.TokenId("<|endoftext|>"); @@ -155,12 +149,33 @@ inline void FillTextEncoderInputs(const std::vector& token_ids, int64_t // Random latent initialisation // ============================================================================ -inline void initialize_latent(float* latent, size_t num_elements, unsigned int seed = 42) +inline void initialize_cpu_latent(float* output, size_t count, uint64_t seed, din::common::Rng rng, unsigned int batch, + unsigned int channels, unsigned int height, unsigned int width) { - std::default_random_engine rng(seed); - std::normal_distribution dist(0.0f, 1.0f); - for (size_t i = 0; i < num_elements; ++i) - latent[i] = dist(rng); + size_t expected = 1; + for (unsigned int dimension : {batch, channels, height, width}) + { + if (dimension == 0 || expected > std::numeric_limits::max() / dimension) + throw std::invalid_argument("Latent dimensions must be nonempty and fit in size_t"); + expected *= dimension; + } + if (output == nullptr || count != expected) + throw std::invalid_argument("Noise buffer must match the latent dimensions"); + if (rng == din::common::Rng::CpuStd) + { + din::common::GenerateNormalCpu({output, count}, seed, rng); + return; + } + std::vector noise(count); + din::common::GenerateNormalCpu(noise, seed, rng, {.round_to_bfloat16 = true}); + const size_t spatial = static_cast(height) * width; + for (size_t index = 0; index < count; ++index) + { + const size_t b = index / (channels * spatial); + const size_t c = (index / spatial) % channels; + const size_t pixel = index % spatial; + output[(b * spatial + pixel) * channels + c] = noise[index]; + } } // ============================================================================ @@ -240,50 +255,3 @@ inline void stats(const std::vector& vec) { stats(vec.data(), vec.size()); } - -// ============================================================================ -// Batch-norm parameters (exported from vae.bn) -// ============================================================================ - -static const float BN_STD[128] = { - 1.8046875, 1.7734375, 1.7890625, 1.7890625, 1.7734375, 1.7578125, 1.7578125, 1.75, 1.734375, 1.734375, - 1.734375, 1.734375, 1.859375, 1.8515625, 1.859375, 1.859375, 1.7578125, 1.7578125, 1.7578125, 1.7578125, - 1.734375, 1.7421875, 1.734375, 1.7421875, 1.734375, 1.7265625, 1.7421875, 1.734375, 1.7578125, 1.75, - 1.7578125, 1.75, 1.84375, 1.84375, 1.8515625, 1.8515625, 1.828125, 1.78125, 1.7890625, 1.796875, - 1.8046875, 1.7890625, 1.765625, 1.765625, 1.7734375, 1.765625, 1.796875, 1.7890625, 1.7578125, 1.7421875, - 1.75, 1.734375, 1.796875, 1.796875, 1.796875, 1.796875, 1.7734375, 1.7734375, 1.7734375, 1.7734375, - 1.7578125, 1.7421875, 1.765625, 1.75, 1.7734375, 1.7734375, 1.78125, 1.78125, 1.8125, 1.8046875, - 1.8046875, 1.8046875, 1.7578125, 1.734375, 1.7578125, 1.75, 1.7578125, 1.75, 1.765625, 1.765625, - 1.765625, 1.7578125, 1.78125, 1.7734375, 1.7890625, 1.796875, 1.7890625, 1.7734375, 1.765625, 1.765625, - 1.7578125, 1.7578125, 1.7734375, 1.765625, 1.78125, 1.78125, 1.71875, 1.703125, 1.7265625, 1.7109375, - 1.78125, 1.7578125, 1.75, 1.7421875, 1.7578125, 1.75, 1.7578125, 1.734375, 1.765625, 1.7578125, - 1.7421875, 1.7421875, 1.75, 1.7578125, 1.765625, 1.765625, 1.7265625, 1.7109375, 1.734375, 1.7265625, - 1.765625, 1.765625, 1.7734375, 1.765625, 1.7421875, 1.7421875, 1.7578125, 1.7578125}; - -static const float BN_MEAN[128] = { - -0.0673828125f, -0.0712890625f, -0.0751953125f, -0.07470703125f, 0.0223388671875f, - 0.0179443359375f, 0.01422119140625f, 0.018310546875f, -6.29425048828125e-05f, -0.006256103515625f, - -0.000209808349609375f, -0.003143310546875f, -0.0272216796875f, -0.028076171875f, -0.027587890625f, - -0.029052734375f, -0.07666015625f, -0.0673828125f, -0.09033203125f, -0.08935546875f, - 0.016845703125f, 0.01519775390625f, 0.00787353515625f, 0.00860595703125f, 0.00836181640625f, - 0.0015411376953125f, 0.0002574920654296875f, -0.0042724609375f, -0.0439453125f, -0.0419921875f, - -0.043701171875f, -0.043212890625f, -0.01025390625f, -0.01318359375f, -0.006622314453125f, - -0.0047607421875f, -0.031005859375f, -0.030517578125f, -0.0279541015625f, -0.0179443359375f, - 0.003021240234375f, 0.00150299072265625f, 0.0125732421875f, 0.01446533203125f, 0.03466796875f, - 0.03369140625f, 0.03369140625f, 0.0283203125f, 0.001983642578125f, 0.004730224609375f, - 0.004669189453125f, 0.004974365234375f, 0.01226806640625f, 0.00811767578125f, 0.008056640625f, - 0.01458740234375f, 0.06787109375f, 0.06787109375f, 0.07666015625f, 0.0732421875f, - -0.046142578125f, -0.04736328125f, -0.039306640625f, -0.051025390625f, -0.052734375f, - -0.0478515625f, -0.047119140625f, -0.0517578125f, -0.03173828125f, -0.03173828125f, - -0.034423828125f, -0.0281982421875f, 0.051025390625f, 0.04443359375f, 0.057861328125f, - 0.0458984375f, -0.041259765625f, -0.0458984375f, -0.048828125f, -0.046630859375f, - -0.00885009765625f, -0.0106201171875f, -0.0087890625f, -0.004608154296875f, -0.03759765625f, - -0.043212890625f, -0.04345703125f, -0.0498046875f, 0.0118408203125f, 0.0166015625f, - 0.020263671875f, 0.0279541015625f, 0.01129150390625f, 0.01287841796875f, 0.001556396484375f, - 0.00714111328125f, -0.01177978515625f, -0.00183868408203125f, -0.01416015625f, -0.00537109375f, - -0.00909423828125f, -0.0137939453125f, -0.01446533203125f, -0.0186767578125f, 0.0322265625f, - 0.030517578125f, 0.02587890625f, 0.0299072265625f, 0.053955078125f, 0.0615234375f, - 0.049560546875f, 0.05908203125f, -0.051025390625f, -0.060302734375f, -0.0478515625f, - -0.052490234375f, -0.022705078125f, -0.0274658203125f, -0.015380859375f, -0.0255126953125f, - -0.05712890625f, -0.056396484375f, -0.0517578125f, -0.049560546875f, 0.0115966796875f, - 0.00543212890625f, 0.016357421875f, 0.0103759765625f}; diff --git a/image_gen/flux2/model_export/export_flux2.py b/image_gen/flux2/model_export/export_flux2.py index 3ec28af..c6a7e5a 100644 --- a/image_gen/flux2/model_export/export_flux2.py +++ b/image_gen/flux2/model_export/export_flux2.py @@ -33,13 +33,14 @@ from pathlib import Path import torch -from diffusers import Flux2KleinPipeline as FluxPipeline from huggingface_hub import snapshot_download +from translator import add_translator_arguments, load_translator_encoder EXTERNAL_DATA_NAME = "model.onnx_data" DEFAULT_MODEL_NAME = "black-forest-labs/FLUX.2-klein-4b" DEFAULT_MODEL_OPSETS = { "text_encoder": 25, + "text_encoder_translator": 25, "transformer": 25, "vae_encoder": 22, "vae_decoder": 22, @@ -302,6 +303,21 @@ def forward(self, latent_sample, return_dict=False): opset=opset, ) print(" vae_decoder exported.") + # Denormalization happens outside the decoder ONNX graph, in patchified channel order. + # Use learned BN buffers, not optional config.latents_mean/config.latents_std. + mean = vae.bn.running_mean + std = torch.sqrt(vae.bn.running_var + 1e-4) + expected_channels = vae_latent_channels * patch_size[0] * patch_size[1] + if mean.numel() != expected_channels or std.numel() != expected_channels: + raise ValueError("VAE latent statistics have an unexpected channel count") + if not torch.isfinite(mean).all() or not torch.isfinite(std).all() or not (std > 0).all(): + raise ValueError("VAE latent statistics must be finite, with positive standard deviations") + stats_path = output_path / "vae_decoder" / "latent_stats.json" + stats_path.write_text( + json.dumps({"bn_mean": mean.float().cpu().tolist(), "bn_std": std.float().cpu().tolist()}, allow_nan=False), + encoding="utf-8", + ) + print(f" VAE latent statistics exported to {stats_path}.") def resolve_model_snapshot(model_name: str, local_files_only: bool) -> Path: @@ -315,6 +331,8 @@ def copy_scheduler_and_tokenizer(model_snapshot: Path, output_path: Path): for name in ("scheduler", "tokenizer"): src = model_snapshot / name dst = output_path / name + if src.resolve() == dst.resolve(): + continue if src.exists(): if dst.exists(): shutil.rmtree(dst) @@ -428,7 +446,7 @@ def main(): nargs="+", default=None, metavar="NAME", - help="Model(s) to export: one or more of all, text_encoder, transformer, vae_encoder, vae_decoder (default: all)", + help="Models to export: all, text_encoder, text_encoder_translator, transformer, vae_encoder, vae_decoder. Translator is opt-in.", ) ap.add_argument( "--opset", @@ -486,6 +504,7 @@ def main(): default=[], help="Additional argument to pass to tensorrt_rtx. May be specified more than once.", ) + add_translator_arguments(ap) args = ap.parse_args() model_snapshot = resolve_model_snapshot(args.model_name, args.local_files_only) @@ -495,7 +514,7 @@ def main(): seq_len = 512 # static text sequence length for FLUX.2-klein # Parse --model: default all; otherwise one or more of all, text_encoder, transformer, … - model_choices = ["all"] + list(MODELS) + model_choices = ["all", "text_encoder_translator"] + list(MODELS) raw = args.model if args.model is not None else ["all"] to_export: list[str] = [] for m in raw: @@ -518,12 +537,26 @@ def main(): f"Loading pipeline from {args.model_name} ({model_snapshot}) " f"(dtype={dtype}, image_size={image_size}, io_precision={args.io_precision})..." ) - pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) + pipe = None + if any(name != "text_encoder_translator" for name in to_export): + from diffusers import Flux2KleinPipeline as FluxPipeline + pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) exported_model_paths: list[Path] = [] for name in to_export: model_opset = args.opset if args.opset is not None else DEFAULT_MODEL_OPSETS[name] print(f"Exporting {name}...") + if name == "text_encoder_translator": + encoder = load_translator_encoder(args, device) + ids = torch.zeros((1, seq_len), dtype=torch.int64, device=device) + mask = torch.ones_like(ids) + path = model_output_path(output_path, name, args.transformer_precision) + _onnx_export(encoder, (ids, mask), path, + ["input_ids", "attention_mask"], ["prompt_embeds"], model_opset) + exported_model_paths.append(path) + del encoder + torch.cuda.empty_cache() + continue fn = EXPORTERS[name] if name == "text_encoder": fn(pipe, output_path, device, model_opset, @@ -539,7 +572,7 @@ def main(): fn(pipe, output_path, device, model_opset, io_dtype=io_dtype) exported_model_paths.append(model_output_path(output_path, name, args.transformer_precision)) - if export_all: + if export_all or "text_encoder_translator" in to_export: copy_scheduler_and_tokenizer(model_snapshot, output_path) write_model_index(output_path) diff --git a/image_gen/flux2/model_export/translator.py b/image_gen/flux2/model_export/translator.py new file mode 100644 index 0000000..e228a1d --- /dev/null +++ b/image_gen/flux2/model_export/translator.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deployment architecture for gafr1/translator; shared by export and verification.""" +from pathlib import Path + +import torch +from torch import nn +from torch.nn import functional as F + +DEFAULT_TRANSLATOR_REPO = "gafr1/translator" +DEFAULT_TRANSLATOR_FILE = "mlp_d3072_b4_e2.pt" + + +class ResMLP(nn.Module): + def __init__(self, width, expansion): + super().__init__() + self.norm = nn.LayerNorm(width) + self.fc1 = nn.Linear(width, width * expansion) + self.fc2 = nn.Linear(width * expansion, width) + + def forward(self, x): + return x + self.fc2(F.gelu(self.fc1(self.norm(x)))) + + +class MLPTranslator(nn.Module): + def __init__(self, in_dim, width, blocks, expansion): + super().__init__() + self.in_proj = nn.Linear(in_dim, width) + self.blocks = nn.ModuleList([ResMLP(width, expansion) for _ in range(blocks)]) + self.norm = nn.LayerNorm(width) + self.out = nn.Linear(width, 7680) + + def forward(self, x): + x = self.in_proj(x) + for block in self.blocks: + x = block(x) + return self.out(self.norm(x)) + + +class TranslatorEncoder(nn.Module): + def __init__(self, qwen, translator, layers): + super().__init__() + self.qwen = qwen + self.translator = translator + self.layers = layers + + def forward(self, input_ids, attention_mask): + output = self.qwen(input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True, use_cache=False) + features = torch.cat([output.hidden_states[i] for i in self.layers], dim=-1) + return self.translator(features).float() + + +def add_translator_arguments(parser): + parser.add_argument("--qwen-model", default="Qwen/Qwen3-0.6B", help="HF repository or local Qwen3-0.6B snapshot") + parser.add_argument("--translator-checkpoint", type=Path, + help="Local checkpoint; takes precedence over HF download") + parser.add_argument("--translator-repo", default=DEFAULT_TRANSLATOR_REPO) + parser.add_argument("--translator-file", default=DEFAULT_TRANSLATOR_FILE) + parser.add_argument("--translator-revision", default="main", + help="HF revision; pin a commit for reproducible exports") + + +def load_translator_encoder(args, device): + from huggingface_hub import hf_hub_download + from transformers import AutoModel + + checkpoint = args.translator_checkpoint + if checkpoint is None: + checkpoint = Path(hf_hub_download( + args.translator_repo, args.translator_file, revision=args.translator_revision, + local_files_only=args.local_files_only)) + # Checkpoints contain only state tensors plus plain config/scalar dictionaries. + data = torch.load(checkpoint, map_location="cpu", weights_only=True) + cfg = data["cfg"] + if cfg["kind"] != "mlp" or cfg["layers"] != "q4" or data["in_dim"] != 4096: + raise ValueError("Expected a Qwen3-0.6B q4 MLP translator with input dimension 4096") + net = MLPTranslator(data["in_dim"], cfg["d"], cfg["nblocks"], cfg["expansion"]) + net.load_state_dict(data["model"], strict=True) + qwen = AutoModel.from_pretrained( + args.qwen_model, torch_dtype=torch.bfloat16, local_files_only=args.local_files_only, + attn_implementation="eager") + if qwen.config.hidden_size != 1024 or qwen.config.num_hidden_layers != 28: + raise ValueError("Translator requires Qwen3-0.6B (28 layers, width 1024)") + return TranslatorEncoder(qwen, net, (7, 14, 21, 28)).eval().to(device=device, dtype=torch.bfloat16) diff --git a/image_gen/flux2/model_export/verify_flux2.py b/image_gen/flux2/model_export/verify_flux2.py index 9dd004d..0607910 100644 --- a/image_gen/flux2/model_export/verify_flux2.py +++ b/image_gen/flux2/model_export/verify_flux2.py @@ -26,9 +26,9 @@ import onnx import torch import torch.nn as nn -from diffusers import Flux2KleinPipeline as FluxPipeline from huggingface_hub import snapshot_download from PIL import Image as PILImage +from translator import add_translator_arguments, load_translator_encoder EP_NAME = "nv_tensorrt_rtx" _REGISTERED_EP_NAME: str | None = None @@ -281,7 +281,7 @@ def transformer_model_path(onnx_dir: Path, precision: str) -> Path: def validate_onnx_dir(onnx_dir: Path, precision: str, provider: str, args) -> None: print(f"\n[Validate] {onnx_dir}") model_paths = [ - onnx_dir / "text_encoder" / "model.onnx", + onnx_dir / ("text_encoder_translator" if args.encoder == "translator" else "text_encoder") / "model.onnx", transformer_model_path(onnx_dir, precision), onnx_dir / "vae_encoder" / "model.onnx", onnx_dir / "vae_decoder" / "model.onnx", @@ -310,11 +310,13 @@ def validate_onnx_dir(onnx_dir: Path, precision: str, provider: str, args) -> No raise RuntimeError(f"No model.onnx files found under {onnx_dir}") -def run_pytorch_pipeline(pipe, prompt: str, seed: int, image_size: int, num_steps: int) -> PILImage.Image: +def run_pytorch_pipeline(pipe, prompt: str, seed: int, image_size: int, num_steps: int, + prompt_embeds=None) -> PILImage.Image: print("\n[PyTorch] Running pipeline...") generator = torch.Generator(device="cuda").manual_seed(seed) result = pipe( - prompt=prompt, + prompt=prompt if prompt_embeds is None else None, + prompt_embeds=prompt_embeds, height=image_size, width=image_size, num_inference_steps=num_steps, @@ -339,7 +341,7 @@ def run_ort_pipeline( prompt_embeds = None original_text_encoder = None - te_path = onnx_dir / "text_encoder" / "model.onnx" + te_path = onnx_dir / ("text_encoder_translator" if args.encoder == "translator" else "text_encoder") / "model.onnx" if te_path.exists(): original_text_encoder = pipe.text_encoder original_text_dtype = original_text_encoder.dtype @@ -439,12 +441,46 @@ def compare_images(ref: PILImage.Image, ort: PILImage.Image, output_dir: Path) - print(f"[Comparison] Side-by-side saved -> {cmp_path}") +def verify_translator_encoder(onnx_dir, args): + from transformers import AutoTokenizer + + if args.encoder != "translator": + raise ValueError("--encoder-only requires --encoder translator") + tokenizer = AutoTokenizer.from_pretrained(str(onnx_dir / "tokenizer")) + reference = load_translator_encoder(args, "cuda") + runner = OrtSessionRunner.from_path(onnx_dir / "text_encoder_translator/model.onnx", args.provider, args) + _assert_fp32_float_io(runner, onnx_dir / "text_encoder_translator/model.onnx") + if set(runner.input_names) != {"input_ids", "attention_mask"} or runner.output_names != ["prompt_embeds"]: + raise RuntimeError("Translator ONNX graph IO names do not match the C++ contract") + prompts = [args.prompt, "", "A café in Zürich, 日本の庭", "a detailed landscape " * 600] + for index, prompt in enumerate(prompts): + text = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=False) + expected_text = "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + if text != expected_text: + raise RuntimeError("Tokenizer chat template differs from the C++ runtime") + tokens = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to("cuda") + with torch.no_grad(): + expected = reference(tokens["input_ids"], tokens["attention_mask"]).float().cpu().numpy() + actual = runner.run({key: _to_numpy(tokens[key], np.int64) for key in ("input_ids", "attention_mask")})[0] + if actual.shape != (1, 512, 7680) or actual.dtype != np.float32 or not np.isfinite(actual).all(): + raise RuntimeError("Translator output must be finite fp32 [1,512,7680]") + cosine = float(np.mean(np.sum(expected * actual, axis=-1) / np.maximum( + np.linalg.norm(expected, axis=-1) * np.linalg.norm(actual, axis=-1), 1e-12))) + nrmse = float(np.sqrt(np.mean((expected - actual) ** 2)) / max(float(np.sqrt(np.mean(expected ** 2))), 1e-12)) + print(f"[Translator prompt {index}] cosine={cosine:.6f}, NRMSE={nrmse:.6f}", flush=True) + if not np.isfinite(cosine + nrmse) or cosine < args.embedding_cosine_min or nrmse > args.embedding_nrmse_max: + raise RuntimeError("Translator embedding parity failed") + + def main(): p = argparse.ArgumentParser(description="Verify FLUX.2-klein ONNX graphs with PyTorch and ONNX Runtime.") p.add_argument("--model_name", type=str, default=DEFAULT_MODEL_NAME, help="Hugging Face model repo ID, or a local snapshot directory") p.add_argument("--local_files_only", action="store_true", help="Resolve --model_name from the local Hugging Face cache only") p.add_argument("--onnx_dir", type=str, default="./flux2_klein_onnx") - p.add_argument("--precision", choices=["bf16", "fp8", "nvfp4"], default="bf16", help="Transformer precision to verify from transformer_.") + p.add_argument("--precision", choices=["bf16", "fp16", "fp8", "nvfp4"], default="bf16", + help="Transformer precision to verify from transformer_.") p.add_argument("--output_dir", type=str, default=None) p.add_argument("--prompt", type=str, default=DEFAULT_PROMPT) p.add_argument("--seed", type=int, default=42) @@ -456,6 +492,12 @@ def main(): p.add_argument("--ep_dll_dir", default=None, help="Dir with the plugin EP's bundled runtime DLLs") p.add_argument("--trt_bin", default=None, help="TensorRT-RTX bin dir, used as a DLL search path") p.add_argument("--validate_only", action="store_true", help="Only run ONNX Runtime validation; do not generate images") + p.add_argument("--encoder", choices=["4b", "translator"], default="4b") + p.add_argument("--encoder-only", action="store_true", + help="Compare fused translator PyTorch and ONNX embeddings without loading the diffusion pipeline") + p.add_argument("--embedding-cosine-min", type=float, default=0.99) + p.add_argument("--embedding-nrmse-max", type=float, default=0.10) + add_translator_arguments(p) args = p.parse_args() configure_stdio() @@ -463,6 +505,12 @@ def main(): output_dir = Path(args.output_dir).resolve() if args.output_dir else onnx_dir output_dir.mkdir(parents=True, exist_ok=True) + if args.encoder_only: + verify_translator_encoder(onnx_dir, args) + return + + if args.encoder == "translator" and not (onnx_dir / "text_encoder_translator/model.onnx").is_file(): + raise FileNotFoundError("Selected translator ONNX is missing") if args.validate_only: validate_onnx_dir(onnx_dir, args.precision, args.provider, args) return @@ -472,15 +520,55 @@ def main(): device = "cuda" model_snapshot = resolve_model_snapshot(args.model_name, args.local_files_only) + from diffusers import Flux2KleinPipeline as FluxPipeline print(f"Loading pipeline from {args.model_name} ({model_snapshot})...") - pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) + pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype) + if args.encoder == "translator": + # Avoid loading the unused 4B encoder on the GPU. + pipe.transformer.to(device) + pipe.vae.to(device) + else: + pipe.to(device) + reference_embeds = None + if args.encoder == "translator": + # Use the same fused student/translator for numerical parity, not the 4B teacher. + pipe.text_encoder.to("cpu") + torch.cuda.empty_cache() + reference = load_translator_encoder(args, device) + text = pipe.tokenizer.apply_chat_template( + [{"role": "user", "content": args.prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=False) + tokens = pipe.tokenizer(text, return_tensors="pt", padding="max_length", + truncation=True, max_length=512).to(device) + with torch.no_grad(): + reference_embeds = reference(tokens["input_ids"], tokens["attention_mask"]).detach() + del reference + torch.cuda.empty_cache() + path = onnx_dir / "text_encoder_translator/model.onnx" + runner = OrtSessionRunner.from_path(path, args.provider, args) + expected = reference_embeds.float().cpu().numpy() + actual = runner.run({k: _to_numpy(tokens[k], np.int64) + for k in ("input_ids", "attention_mask")})[0] + if actual.shape != (1, 512, 7680) or actual.dtype != np.float32 or not np.isfinite(actual).all(): + raise RuntimeError("Translator output must be finite fp32 [1,512,7680]") + cosine = float(np.mean(np.sum(expected * actual, axis=-1) / + np.maximum(np.linalg.norm(expected, axis=-1) * np.linalg.norm(actual, axis=-1), 1e-12))) + nrmse = float(np.sqrt(np.mean((expected - actual) ** 2)) / max(float(np.sqrt(np.mean(expected ** 2))), 1e-12)) + print(f"[Translator embeddings] cosine={cosine:.6f}, NRMSE={nrmse:.6f}") + if cosine < args.embedding_cosine_min or nrmse > args.embedding_nrmse_max: + raise RuntimeError("Translator embedding parity failed") + del runner + reference_embeds = reference_embeds.to(dtype=dtype) + # Precomputed embeddings bypass the teacher; keep pipeline device discovery on CUDA. + pipe.text_encoder = _DeviceAnchor(device, dtype) pytorch_image = run_pytorch_pipeline( pipe, prompt=args.prompt, seed=args.seed, image_size=args.image_size, num_steps=args.num_steps, + prompt_embeds=reference_embeds, ) pt_path = output_dir / "pytorch_output.png" pytorch_image.save(str(pt_path)) diff --git a/image_gen/flux2/utils.cpp b/image_gen/flux2/utils.cpp index 86b4343..a645da4 100644 --- a/image_gen/flux2/utils.cpp +++ b/image_gen/flux2/utils.cpp @@ -16,9 +16,47 @@ #include "utils.h" #include +#include #include #include #include +#include + +#include + +Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t channels) +{ + std::ifstream file(path); + if (!file) + throw std::runtime_error("Missing VAE latent statistics: " + path.string() + ". Re-export the VAE decoder."); + try + { + const auto document = nlohmann::json::parse(file); + const auto read = [&](const char* key, bool positive) + { + const auto& values = document.at(key); + if (!values.is_array() || values.size() != channels) + throw std::runtime_error(std::string(key) + " must contain " + std::to_string(channels) + " values"); + std::vector result; + result.reserve(channels); + for (const auto& value : values) + { + if (!value.is_number()) + throw std::runtime_error(std::string(key) + " must contain only numbers"); + const float number = value.get(); + if (!std::isfinite(number) || (positive && number <= 0.0f)) + throw std::runtime_error(std::string(key) + " contains an invalid value"); + result.push_back(number); + } + return result; + }; + return {read("bn_mean", false), read("bn_std", true)}; + } + catch (const std::exception& error) + { + throw std::runtime_error("Invalid VAE latent statistics in " + path.string() + ": " + error.what()); + } +} #ifdef _WIN32 #include diff --git a/image_gen/flux2/utils.h b/image_gen/flux2/utils.h index 2d3a174..d4aff19 100644 --- a/image_gen/flux2/utils.h +++ b/image_gen/flux2/utils.h @@ -42,6 +42,14 @@ OrtFileString toOrtFileString(const std::filesystem::path& path); std::filesystem::path get_executable_path(); std::filesystem::path get_executable_parent_path(); +struct Flux2LatentStats +{ + std::vector mean; + std::vector std; +}; + +Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t channels); + void flush_ort_stream(Ort::SyncStream& stream); void copy_tensors(Ort::Env& env, std::initializer_list src, std::initializer_list dst, Ort::SyncStream& stream);