FLux2: Graphics interop fixes

.github/workflows/ci.yml

@@ -26,25 +26,21 @@ jobs:
       - name: Check out source
         uses: actions/checkout@v7
 
-      - name: Install clang-format
+      - name: Install released clang-format
         shell: bash
         run: |
-          set -euxo pipefail
-          sudo apt-get update
-          sudo apt-get install -y --no-install-recommends ca-certificates curl
-          llvm_install_script="${RUNNER_TEMP}/llvm.sh"
-          curl --fail --location --silent --show-error https://apt.llvm.org/llvm.sh --output "$llvm_install_script"
-          sudo bash "$llvm_install_script" 22
-          sudo apt-get install -y --no-install-recommends clang-format-22
+          set -euo pipefail
+          python3 -m pip install --user --disable-pip-version-check "clang-format==22.1.8"
+          echo "$HOME/.local/bin" >> "$GITHUB_PATH"
 
       - name: Check formatting
         shell: bash
         run: |
           set -euo pipefail
-          clang-format-22 --version
+          clang-format --version
           git ls-files -z -- \
             '*.c' '*.cc' '*.cpp' '*.cxx' '*.h' '*.hh' '*.hpp' '*.hxx' '*.cu' '*.cuh' \
-            | xargs -0 -r clang-format-22 --dry-run --Werror --
+            | xargs -0 -r clang-format --dry-run --Werror --
 
   linux-x64:
     name: Ubuntu 24.04 / Clang / x64

CMakeLists.txt

@@ -205,15 +205,12 @@ if(DIN_BUILD_TESTING)
                 dx-cig
         )
     endif ()
-    if (NOT WIN32)
-        # TODO this test should be cross platform but experiences a device lost error on windows
-        add_din_flux2_cli_test(
-                din_flux2_bf16
-                "${DIN_DEPLOY_ARTIFACTS_ROOT}/FLUX.2-klein-4B-onnx"
-                trt-rtx
-                vk
-        )
-    endif ()
+    add_din_flux2_cli_test(
+            din_flux2_bf16
+            "${DIN_DEPLOY_ARTIFACTS_ROOT}/FLUX.2-klein-4B-onnx"
+            trt-rtx
+            vk
+    )
 
     unset(_din_cli_test_audio)
 endif()

CMakePresets.json

@@ -8,7 +8,7 @@
       "binaryDir": "${sourceDir}/out/build/${presetName}",
       "cacheVariables": {
         "CMAKE_EXPORT_COMPILE_COMMANDS": "ON",
-        "ONNXRUNTIME_VERSION": "1.27.0",
+        "ONNXRUNTIME_VERSION": "1.29.0",
         "DIN_BUILD_TESTING": "ON",
         "DIN_ENABLE_NVTX": "ON"
       }

common/ort_session.cpp

@@ -724,6 +724,30 @@ int ChooseCudaDeviceOrdinal(Ort::ConstEpDevice ep_device, const char* override_e
     return 0;
 }
 
+OrtGraphicsDeviceIdentity ResolveOrtGraphicsDeviceIdentity(Ort::ConstEpDevice ep_device)
+{
+    OrtGraphicsDeviceIdentity identity{};
+    identity.vendor_id = ep_device.Device().VendorId();
+    identity.device_id = ep_device.Device().DeviceId();
+
+#ifdef _WIN32
+    const auto metadata = ep_device.Device().Metadata().GetKeyValuePairs();
+    const std::string* luid_text = FindMetadataValue(metadata, "LUID");
+    if (luid_text == nullptr)
+    {
+        throw std::runtime_error("ORT TRT-RTX device does not expose Windows LUID metadata");
+    }
+    const auto luid = ParseUint64(*luid_text);
+    if (!luid.has_value())
+    {
+        throw std::runtime_error("ORT TRT-RTX device exposes an invalid Windows LUID: " + *luid_text);
+    }
+    identity.luid = *luid;
+#endif
+
+    return identity;
+}
+
 CudaGraphicsInteropSharedMemoryInfo QueryCudaGraphicsInteropSharedMemoryInfo(int cuda_device_ordinal,
                                                                              bool high_priority_compute_queue)
 {

common/ort_session.h

@@ -57,11 +57,20 @@ struct CudaGraphicsInteropSharedMemoryInfo
     int max_shared_memory_bytes = 0;
 };
 
+struct OrtGraphicsDeviceIdentity
+{
+    uint32_t vendor_id = 0;
+    uint32_t device_id = 0;
+    uint64_t luid = 0;
+};
+
 [[nodiscard]] const char* ToString(CudaDeviceGeneration generation);
 
 [[nodiscard]] int ChooseCudaDeviceOrdinal(Ort::ConstEpDevice ep_device,
                                           const char* override_env_var = "FLUX_CUDA_DEVICE_ID");
 
+[[nodiscard]] OrtGraphicsDeviceIdentity ResolveOrtGraphicsDeviceIdentity(Ort::ConstEpDevice ep_device);
+
 [[nodiscard]] CudaGraphicsInteropSharedMemoryInfo
 QueryCudaGraphicsInteropSharedMemoryInfo(int cuda_device_ordinal, bool high_priority_compute_queue);
 

image_gen/flux2/README.md

@@ -27,20 +27,19 @@ CiG is useful when inference must share the graphics GPU with rendering. It can
 
 ## Export
 
-```bash
+```powershell
 cd model_export
-python export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output ./models/FLUX.2-klein-4B-onnx --model all
-python quantize_onnx.py --input ./models/FLUX.2-klein-4B-onnx/transformer --quant fp8 --output ./models/FLUX.2-klein-4B-fp8-onnx/transformer
-python quantize_onnx.py --input ./models/FLUX.2-klein-4B-onnx/transformer --quant nvfp4 --output ./models/FLUX.2-klein-4B-nvfp4-onnx/transformer
+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
 ```
 
-Ensure to copy other required models to the quantized ONNX model directory. 
-An export contains `text_encoder/model.onnx`, `transformer/model.onnx`, `vae_decoder/model.onnx`, and `tokenizer/tokenizer.json`.
+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`.
 
 ## Verify
 
-```bash
-python verify_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --onnx_dir ./models/FLUX.2-klein-4B-fp8-onnx --output_dir flux2-fp8 --provider trt-rtx
+```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
 ```
 
 ## Build
@@ -53,5 +52,7 @@ cmake --build --preset windows-x64-debug --target din_flux2_cli
 ## Run
 
 ```powershell
-out\build\windows-x64\bin\Debug\din_flux2_cli.exe --model-dir D:\models\FLUX.2-klein-4B-onnx --provider trt-rtx --processing cuda --prompt "a red fox" --output out
+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
 ```
+
+`--precision` selects `transformer_<precision>` 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.

image_gen/flux2/cuda_kernels.cu

@@ -10,7 +10,6 @@
 
 namespace
 {
-
 constexpr unsigned int kThreadsPerBlock = 256;
 
 __global__ void flux_euler_kernel(float t_curr, float t_next, size_t total_elements, float* hidden_states,
@@ -66,7 +65,6 @@ dim3 grid_for(size_t total_elements)
     }
     return dim3(static_cast<unsigned int>(blocks), 1, 1);
 }
-
 }  // namespace
 
 void launch_flux_euler_kernel(cudaStream_t stream, float t_curr, float t_next, size_t total_elements,

image_gen/flux2/flux2_cli.cpp

@@ -19,7 +19,6 @@
 
 namespace
 {
-
 unsigned int parse_uint(const std::string& value, const char* flag_name)
 {
     size_t parsed_chars = 0;
@@ -45,6 +44,8 @@ std::string to_string(Flux2ProcessingBackend backend)
         return "dx-cig";
     case Flux2ProcessingBackend::Vk:
         return "vk";
+    case Flux2ProcessingBackend::VkCig:
+        return "vk-cig";
     }
     return "unknown";
 }
@@ -87,7 +88,11 @@ Flux2ProcessingBackend parse_processing_backend(std::string value)
     {
         return Flux2ProcessingBackend::Vk;
     }
-    throw std::invalid_argument("Processing must be one of: cpu, cuda, dx, dx-cig, vk");
+    if (value == "vk-cig")
+    {
+        return Flux2ProcessingBackend::VkCig;
+    }
+    throw std::invalid_argument("Processing must be one of: cpu, cuda, dx, dx-cig, vk, vk-cig");
 }
 
 Flux2ExecutionProvider parse_execution_provider(std::string value)
@@ -107,6 +112,19 @@ Flux2ExecutionProvider parse_execution_provider(std::string value)
     throw std::invalid_argument("Provider must be one of: cpu, trt-rtx");
 }
 
+std::string parse_precision(std::string value)
+{
+    for (char& c : value)
+    {
+        c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
+    }
+    if (value == "bf16" || value == "fp8" || value == "nvfp4")
+    {
+        return value;
+    }
+    throw std::invalid_argument("Precision must be one of: bf16, fp8, nvfp4");
+}
+
 void validate_config(const Flux2Config& config)
 {
     if (config.provider == Flux2ExecutionProvider::Cpu && config.processing != Flux2ProcessingBackend::Cpu)
@@ -126,7 +144,7 @@ Flux2Config parse_args(int argc, char* argv[])
     parser.add_argument("--processing")
         .default_value(to_string(config.processing))
         .nargs(1)
-        .metavar("cpu|cuda|dx|dx-cig|vk")
+        .metavar("cpu|cuda|dx|dx-cig|vk|vk-cig")
         .help("Select the processing backend.");
     parser.add_argument("--provider")
         .default_value(to_string(config.provider))
@@ -137,7 +155,12 @@ Flux2Config parse_args(int argc, char* argv[])
         .default_value(config.model_dir.string())
         .nargs(1)
         .metavar("PATH")
-        .help("Directory with exported Flux2 ONNX artifacts.");
+        .help("Root directory with shared Flux2 ONNX artifacts and transformer_<precision> directories.");
+    parser.add_argument("--precision")
+        .default_value(config.precision)
+        .nargs(1)
+        .metavar("bf16|fp8|nvfp4")
+        .help("Transformer precision to load from transformer_<precision>.");
     parser.add_argument("--ep-cache")
         .default_value(config.ep_cache_dir.string())
         .nargs(1)
@@ -178,6 +201,7 @@ Flux2Config parse_args(int argc, char* argv[])
     config.processing = parse_processing_backend(parser.get<std::string>("--processing"));
     config.provider = parse_execution_provider(parser.get<std::string>("--provider"));
     config.model_dir = parser.get<std::string>("--model-dir");
+    config.precision = parse_precision(parser.get<std::string>("--precision"));
     config.ep_cache_dir = parser.get<std::string>("--ep-cache");
     config.ep_context_dir = parser.get<std::string>("--ep-context-dir");
     config.prompt = parser.get<std::string>("--prompt");
@@ -212,7 +236,6 @@ void save_image(const std::filesystem::path& output_path, const Flux2Image& imag
     }
     std::cout << "Image saved to " << output_path.string() << std::endl;
 }
-
 }  // namespace
 
 int main(int argc, char* argv[])
@@ -224,6 +247,7 @@ int main(int argc, char* argv[])
     {
         const Flux2Config config = parse_args(argc, argv);
         std::cout << "Model dir: " << config.model_dir.string() << "\n"
+                  << "Precision: " << config.precision << "\n"
                   << "Output: " << config.output_path.string() << "\n"
                   << "Prompt: " << config.prompt << "\n"
                   << "Seed:   " << config.seed << "\n"

image_gen/flux2/flux2_cli.h

@@ -15,6 +15,7 @@ enum class Flux2ProcessingBackend
     Dx,
     DxCig,
     Vk,
+    VkCig,
 };
 
 enum class Flux2ExecutionProvider
@@ -28,6 +29,7 @@ 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";
     std::filesystem::path output_path;

image_gen/flux2/flux2_cuda_backend.cpp

@@ -34,7 +34,6 @@
 
 namespace
 {
-
 enum class ExecutionProviderMode
 {
     Auto,
@@ -478,7 +477,8 @@ 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);
+    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;
@@ -536,15 +536,15 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const
 
     state.text_encoder_runner = std::make_unique<din::common::OrtRunner>(
         state.env, model_paths.text_encoder_model.string(), provider, cache_dir, ep_context,
-        make_profile("text_encoder"), compute_stream_ptr);
+        make_profile(cache_paths.text_encoder), compute_stream_ptr);
     std::cout << "  Text encoder loaded" << std::endl;
     state.transformer_runner =
         std::make_unique<din::common::OrtRunner>(state.env, model_paths.transformer_model.string(), provider, cache_dir,
-                                                 ep_context, make_profile("transformer"), compute_stream_ptr);
+                                                 ep_context, make_profile(cache_paths.transformer), compute_stream_ptr);
     std::cout << "  Transformer loaded" << std::endl;
     state.vae_decoder_runner =
         std::make_unique<din::common::OrtRunner>(state.env, model_paths.vae_decoder_model.string(), provider, cache_dir,
-                                                 ep_context, make_profile("vae_decoder"), compute_stream_ptr);
+                                                 ep_context, make_profile(cache_paths.vae_decoder), compute_stream_ptr);
     std::cout << "  VAE decoder loaded" << std::endl;
 
     std::vector<int64_t> token_shape = {BATCH_SIZE, SEQUENCE_LENGTH};
@@ -753,7 +753,7 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline
         }
         state_->prompt_embeds_valid = false;
 
-        const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir);
+        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);
@@ -781,7 +781,6 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline
     Flux2RuntimeContext& runtime_;
     std::unique_ptr<CudaPipelineState> state_;
 };
-
 }  // namespace
 
 std::unique_ptr<Flux2ProcessingPipeline> CreateFlux2CudaPipeline(const Flux2Config& config,

image_gen/flux2/flux2_dx_backend.cpp

@@ -355,8 +355,9 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr
 
 struct DxPipelineState
 {
-    explicit DxPipelineState(Ort::Env& environment)
-        : env(environment)
+    DxPipelineState(Ort::Env& environment, uint64_t ort_luid)
+        : dx(ort_luid)
+        , env(environment)
     {
     }
 
@@ -448,7 +449,9 @@ 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);
+    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"
@@ -501,13 +504,13 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi
 
     state.text_encoder_runner = std::make_unique<din::common::OrtRunner>(
         state.env, model_paths.text_encoder_model.string(), "trt-rtx", cache_dir, ep_context,
-        make_profile(state.use_cig ? "dx_cig_text_encoder" : "dx_text_encoder"), &*state.sync_stream);
+        make_profile(cache_paths.text_encoder), &*state.sync_stream);
     state.transformer_runner = std::make_unique<din::common::OrtRunner>(
         state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context,
-        make_profile(state.use_cig ? "dx_cig_transformer" : "dx_transformer"), &*state.sync_stream);
+        make_profile(cache_paths.transformer), &*state.sync_stream);
     state.vae_decoder_runner = std::make_unique<din::common::OrtRunner>(
         state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context,
-        make_profile(state.use_cig ? "dx_cig_vae_decoder" : "dx_vae_decoder"), &*state.sync_stream);
+        make_profile(cache_paths.vae_decoder), &*state.sync_stream);
 
     std::vector<int64_t> token_shape = {BATCH_SIZE, SEQUENCE_LENGTH};
     std::vector<int64_t> attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH};
@@ -697,7 +700,8 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline
             return;
         }
         std::cout << "Initializing Flux2 DirectX pipeline" << std::endl;
-        state_ = std::make_unique<DxPipelineState>(runtime_.env);
+        const auto identity = din::common::ResolveOrtGraphicsDeviceIdentity(runtime_.trt_device);
+        state_ = std::make_unique<DxPipelineState>(runtime_.env, identity.luid);
         initialize_dx_state(*state_, config_, runtime_.trt_device);
     }
 
@@ -710,7 +714,7 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline
         }
         state_->prompt_embeds_valid = false;
 
-        const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir);
+        const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision);
         const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt);
         std::vector<int64_t> tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH);
         std::vector<int64_t> attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH);