diff --git a/CMakeLists.txt b/CMakeLists.txt index 05e7abf..c788fdb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -96,7 +96,9 @@ endfunction() add_subdirectory(common) add_subdirectory(tests/base) +add_subdirectory(tests/tokenizer) add_subdirectory(image_gen/flux2) +add_subdirectory(image_gen/pixeldit) add_subdirectory(asr/rnnt) add_subdirectory(vision/sam2) add_subdirectory(asr/whisper) @@ -211,5 +213,18 @@ if(DIN_BUILD_TESTING) vk ) + foreach (_din_pixeldit_processing cpu cuda) + add_din_cli_test(din_pixeldit_bf16_trt-rtx_${_din_pixeldit_processing} din_pixeldit_cli + --model-dir "${DIN_DEPLOY_ARTIFACTS_ROOT}/PixelDiT-1300M-1024px-onnx" + --provider trt-rtx + --processing ${_din_pixeldit_processing} + --ep-cache "${DIN_CLI_TEST_EP_CACHE_DIR}" + --ep-context-dir "${DIN_CLI_TEST_EP_CONTEXT_DIR}" + --num-images 2 + --output "${PROJECT_SOURCE_DIR}/din_pixeldit_bf16_trt-rtx_${_din_pixeldit_processing}" + ) + endforeach () + unset(_din_pixeldit_processing) + unset(_din_cli_test_audio) endif() diff --git a/common/io/tokenizer.cpp b/common/io/tokenizer.cpp index 86dcc97..2264482 100644 --- a/common/io/tokenizer.cpp +++ b/common/io/tokenizer.cpp @@ -97,14 +97,86 @@ std::array BuildByteEncoder() return byte_to_code; } +constexpr std::string_view kSentencePieceSpace = "\xE2\x96\x81"; // U+2581 + +struct JsonNormalizerStep +{ + bool prepend = false; // otherwise Replace + std::string pattern; + std::string content; +}; + struct JsonTokenizerData { std::vector id_to_token; std::unordered_map token_to_id; std::unordered_map bpe_ranks; std::vector special_tokens; + std::vector added_tokens; + std::vector normalizer; + bool byte_fallback = false; + std::string unk_token; + std::vector prefix_special_tokens; + std::vector suffix_special_tokens; }; +// Collects Prepend / Replace(String) steps from a normalizer (or a Sequence of them). Other +// normalizer types are not supported by the SentencePiece encode path and are reported to the caller. +bool CollectNormalizerSteps(const nlohmann::json& normalizer, std::vector& steps) +{ + if (normalizer.is_null()) + { + return true; + } + const auto type = normalizer.value("type", std::string{}); + if (type == "Sequence") + { + for (const auto& step : normalizer["normalizers"]) + { + if (!CollectNormalizerSteps(step, steps)) + { + return false; + } + } + return true; + } + if (type == "Prepend") + { + steps.push_back({true, {}, normalizer["prepend"].get()}); + return true; + } + if (type == "Replace" && normalizer.contains("pattern") && normalizer["pattern"].contains("String")) + { + steps.push_back( + {false, normalizer["pattern"]["String"].get(), normalizer["content"].get()}); + return true; + } + return false; +} + +// TemplateProcessing "single": special tokens before and after the "A" sequence. +void CollectTemplateSpecialTokens(const nlohmann::json& post_processor, std::vector& prefix, + std::vector& suffix) +{ + if (post_processor.is_null() || post_processor.value("type", std::string{}) != "TemplateProcessing" || + !post_processor.contains("single")) + { + return; + } + bool after_sequence = false; + for (const auto& piece : post_processor["single"]) + { + if (piece.contains("Sequence")) + { + after_sequence = true; + } + else if (piece.contains("SpecialToken")) + { + (after_sequence ? suffix : prefix).push_back(piece["SpecialToken"]["id"].get()); + } + } +} + void AddToken(JsonTokenizerData& data, const std::string& token, int64_t id) { if (id < 0) @@ -150,6 +222,7 @@ JsonTokenizerData LoadJsonTokenizerData(const std::string& path) } const auto content = token["content"].get(); AddToken(tokenizer_data, content, token["id"].get()); + tokenizer_data.added_tokens.push_back(content); if (token.value("special", false)) { tokenizer_data.special_tokens.push_back(content); @@ -157,6 +230,27 @@ JsonTokenizerData LoadJsonTokenizerData(const std::string& path) } } + const auto& model = data["model"]; + tokenizer_data.byte_fallback = model.value("byte_fallback", false); + if (model.contains("unk_token") && model["unk_token"].is_string()) + { + tokenizer_data.unk_token = model["unk_token"].get(); + } + if (data.contains("normalizer") && !CollectNormalizerSteps(data["normalizer"], tokenizer_data.normalizer)) + { + if (tokenizer_data.byte_fallback) + { + throw std::runtime_error("tokenizer.json uses a normalizer that is not supported for SentencePiece BPE: " + + data["normalizer"].dump()); + } + tokenizer_data.normalizer.clear(); + } + if (data.contains("post_processor")) + { + CollectTemplateSpecialTokens(data["post_processor"], tokenizer_data.prefix_special_tokens, + tokenizer_data.suffix_special_tokens); + } + if (data["model"].contains("merges")) { int32_t rank = 0; @@ -189,14 +283,28 @@ JsonTokenizerData LoadJsonTokenizerData(const std::string& path) } } - std::sort(tokenizer_data.special_tokens.begin(), tokenizer_data.special_tokens.end(), - [](const auto& left, const auto& right) - { - return left.size() > right.size(); - }); + const auto longest_first = [](const auto& left, const auto& right) + { + return left.size() > right.size(); + }; + std::sort(tokenizer_data.special_tokens.begin(), tokenizer_data.special_tokens.end(), longest_first); + std::stable_sort(tokenizer_data.added_tokens.begin(), tokenizer_data.added_tokens.end(), longest_first); return tokenizer_data; } +bool IsSentencePieceBpe(const JsonTokenizerData& data) +{ + if (!data.byte_fallback) + { + return false; + } + return std::any_of(data.normalizer.begin(), data.normalizer.end(), + [](const JsonNormalizerStep& step) + { + return !step.prepend && step.pattern == " " && step.content == kSentencePieceSpace; + }); +} + std::vector LoadFlatVocab(const std::string& path) { std::ifstream stream(path); @@ -349,12 +457,33 @@ Tokenizer::Tokenizer(const std::string& path, TokenizerFormat format) case TokenizerFormat::Json: { auto tokenizer_data = LoadJsonTokenizerData(path); + const bool sentencepiece_bpe = IsSentencePieceBpe(tokenizer_data); id_to_token_ = std::move(tokenizer_data.id_to_token); token_to_id_ = std::move(tokenizer_data.token_to_id); bpe_ranks_ = std::move(tokenizer_data.bpe_ranks); special_tokens_ = std::move(tokenizer_data.special_tokens); byte_encoder_ = BuildByteEncoder(); decode_mode_ = DecodeMode::Pieces; + if (sentencepiece_bpe) + { + encode_mode_ = EncodeMode::SentencePieceBpe; + added_tokens_ = std::move(tokenizer_data.added_tokens); + for (const auto& step : tokenizer_data.normalizer) + { + normalizer_.push_back({step.prepend ? NormalizerStep::Kind::Prepend : NormalizerStep::Kind::Replace, + step.pattern, step.content}); + } + byte_fallback_ = tokenizer_data.byte_fallback; + unk_id_ = TokenId(tokenizer_data.unk_token); + for (const auto& token : tokenizer_data.prefix_special_tokens) + { + prefix_special_ids_.push_back(token_to_id_.at(token)); + } + for (const auto& token : tokenizer_data.suffix_special_tokens) + { + suffix_special_ids_.push_back(token_to_id_.at(token)); + } + } break; } case TokenizerFormat::Vocab: @@ -581,8 +710,170 @@ std::vector ApplyByteLevelBpe(std::string_view piece, const std::ar return tokens; } +std::string ReplaceAll(std::string text, std::string_view pattern, std::string_view content) +{ + if (pattern.empty()) + { + return text; + } + std::string out; + out.reserve(text.size()); + size_t pos = 0; + while (true) + { + const size_t found = text.find(pattern, pos); + if (found == std::string::npos) + { + out.append(text, pos, std::string::npos); + return out; + } + out.append(text, pos, found - pos); + out.append(content); + pos = found + pattern.size(); + } +} + +size_t Utf8CharLength(uint8_t lead) +{ + if (lead < 0x80) + return 1; + if ((lead & 0xE0) == 0xC0) + return 2; + if ((lead & 0xF0) == 0xE0) + return 3; + if ((lead & 0xF8) == 0xF0) + return 4; + return 1; // invalid lead byte: treat as a single byte +} + +std::string ByteFallbackToken(uint8_t byte) +{ + constexpr char hex[] = "0123456789ABCDEF"; + return std::string{"<0x"} + hex[byte >> 4] + hex[byte & 0x0F] + ">"; +} + +// BPE over UTF-8 characters, as HF tokenizers' BPE model: characters missing from the vocab become +// <0xXX> byte tokens (byte_fallback), then the lowest-rank pair is merged, leftmost first. +std::vector ApplySentencePieceBpe(std::string_view piece, + const std::unordered_map& token_to_id, + const std::unordered_map& bpe_ranks, + bool byte_fallback, const std::string& unk_token) +{ + std::vector symbols; + symbols.reserve(piece.size()); + for (size_t pos = 0; pos < piece.size();) + { + const size_t length = std::min(Utf8CharLength(static_cast(piece[pos])), piece.size() - pos); + std::string character(piece.substr(pos, length)); + pos += length; + if (token_to_id.count(character) != 0) + { + symbols.push_back(std::move(character)); + } + else if (byte_fallback) + { + for (const char byte : character) + { + symbols.push_back(ByteFallbackToken(static_cast(byte))); + } + } + else + { + symbols.push_back(unk_token); + } + } + + while (symbols.size() > 1) + { + auto best_rank = std::numeric_limits::max(); + size_t best_index = symbols.size(); + for (size_t i = 0; i + 1 < symbols.size(); ++i) + { + const auto rank = bpe_ranks.find(BpePairKey(symbols[i], symbols[i + 1])); + if (rank != bpe_ranks.end() && rank->second < best_rank) + { + best_rank = rank->second; + best_index = i; + } + } + if (best_index == symbols.size()) + { + break; + } + symbols[best_index] += symbols[best_index + 1]; + symbols.erase(symbols.begin() + static_cast(best_index + 1)); + } + return symbols; +} + } // namespace +std::vector Tokenizer::EncodeSentencePieceBpe(const std::string& text, bool add_special_tokens) const +{ + std::vector ids; + if (add_special_tokens) + { + ids.insert(ids.end(), prefix_special_ids_.begin(), prefix_special_ids_.end()); + } + + const std::string& unk_token = Token(unk_id_); + const auto encode_segment = [&](std::string_view segment) + { + if (segment.empty()) + { + return; + } + std::string normalized(segment); + for (const auto& step : normalizer_) + { + normalized = step.kind == NormalizerStep::Kind::Prepend + ? step.content + normalized + : ReplaceAll(std::move(normalized), step.pattern, step.content); + } + for (const auto& token : ApplySentencePieceBpe(normalized, token_to_id_, bpe_ranks_, byte_fallback_, unk_token)) + { + const auto id = token_to_id_.find(token); + if (id == token_to_id_.end()) + { + throw std::runtime_error("tokenizer vocab is missing BPE token: " + token); + } + ids.push_back(id->second); + } + }; + + // Added tokens (special or not) are matched on the raw text, leftmost-longest, before normalization. + size_t segment_start = 0; + size_t pos = 0; + while (pos < text.size()) + { + const std::string* matched = nullptr; + for (const auto& added : added_tokens_) + { + if (StartsWith(std::string_view(text).substr(pos), added)) + { + matched = &added; + break; + } + } + if (matched == nullptr) + { + ++pos; + continue; + } + encode_segment(std::string_view(text).substr(segment_start, pos - segment_start)); + ids.push_back(token_to_id_.at(*matched)); + pos += matched->size(); + segment_start = pos; + } + encode_segment(std::string_view(text).substr(segment_start)); + + if (add_special_tokens) + { + ids.insert(ids.end(), suffix_special_ids_.begin(), suffix_special_ids_.end()); + } + return ids; +} + int64_t Tokenizer::TokenId(const std::string& token) const { const auto it = token_to_id_.find(token); @@ -595,6 +886,10 @@ int64_t Tokenizer::TokenId(const std::string& token) const std::vector Tokenizer::Encode(const std::string& text, bool add_special_tokens) const { + if (encode_mode_ == EncodeMode::SentencePieceBpe) + { + return EncodeSentencePieceBpe(text, add_special_tokens); + } if (bpe_ranks_.empty()) { throw std::runtime_error("Tokenizer encoding is only available for byte-level BPE tokenizer.json format"); diff --git a/common/io/tokenizer.h b/common/io/tokenizer.h index eec5b97..b14aca5 100644 --- a/common/io/tokenizer.h +++ b/common/io/tokenizer.h @@ -51,6 +51,28 @@ class Tokenizer WhisperByteBpe, }; + // ByteLevelBpe: GPT-2 style (Qwen, Whisper). SentencePieceBpe: tokenizer.json BPE with a + // Replace(" " -> "▁") normalizer and byte_fallback (Gemma, Llama). + enum class EncodeMode + { + ByteLevelBpe, + SentencePieceBpe, + }; + + struct NormalizerStep + { + enum class Kind + { + Prepend, + Replace, + }; + Kind kind; + std::string pattern; + std::string content; + }; + + [[nodiscard]] std::vector EncodeSentencePieceBpe(const std::string& text, bool add_special_tokens) const; + std::vector id_to_token_; std::unordered_map token_to_id_; std::unordered_map bpe_ranks_; @@ -58,6 +80,15 @@ class Tokenizer std::vector special_tokens_; std::vector lang_codes_; DecodeMode decode_mode_ = DecodeMode::Pieces; + EncodeMode encode_mode_ = EncodeMode::ByteLevelBpe; + // SentencePieceBpe only: all added tokens (special or not), longest first, split before BPE. + std::vector added_tokens_; + std::vector normalizer_; + bool byte_fallback_ = false; + int64_t unk_id_ = -1; + // Post-processor TemplateProcessing "single" special tokens around the sequence. + std::vector prefix_special_ids_; + std::vector suffix_special_ids_; }; } // namespace din::io diff --git a/image_gen/pixeldit/CMakeLists.txt b/image_gen/pixeldit/CMakeLists.txt new file mode 100644 index 0000000..5848117 --- /dev/null +++ b/image_gen/pixeldit/CMakeLists.txt @@ -0,0 +1,60 @@ +find_package(CUDAToolkit REQUIRED) +enable_language(CUDA) + +# The sampling math mirrors the PyTorch reference operation by operation in fp32, on the host (--processing cpu) +# and in cuda_kernels.cu (--processing cuda); keep both compilers from fusing multiply-adds so the two paths give +# identical results (MSVC /fp:precise does not contract by default). +set(PIXELDIT_FP_FLAGS + "$<$,$,$>>:-ffp-contract=off>" + "$<$:--fmad=false>" +) + +add_library(din_pixeldit_common STATIC + pixeldit_config.cpp +) + +target_include_directories(din_pixeldit_common + PUBLIC + "${CMAKE_CURRENT_SOURCE_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/helpers" +) + +target_link_libraries(din_pixeldit_common + PUBLIC + din_common_ort + din_common_io + PRIVATE + nlohmann_json::nlohmann_json +) + +set_target_properties(din_pixeldit_common PROPERTIES + CXX_STANDARD 20 + CXX_EXTENSIONS OFF + POSITION_INDEPENDENT_CODE ON +) + +add_din_shared_library(din_pixeldit_trt + pixeldit_trt_backend.cpp + cuda_kernels.cu +) +set_target_properties(din_pixeldit_trt PROPERTIES + CXX_STANDARD 20 + CXX_EXTENSIONS OFF + CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}" +) +target_compile_options(din_pixeldit_trt PRIVATE ${PIXELDIT_FP_FLAGS}) +target_link_libraries(din_pixeldit_trt + PUBLIC + din_pixeldit_common + PRIVATE + CUDA::cudart +) + +add_din_executable(din_pixeldit_cli + pixeldit_cli.cpp +) +set_target_properties(din_pixeldit_cli PROPERTIES + CXX_STANDARD 20 + CXX_EXTENSIONS OFF +) +target_link_libraries(din_pixeldit_cli PRIVATE argparse din_pixeldit_trt) diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md new file mode 100644 index 0000000..a62182d --- /dev/null +++ b/image_gen/pixeldit/README.md @@ -0,0 +1,74 @@ +# PixelDiT + +CLI image generation with ONNX Runtime. + +## Supported models + +| Model | Hugging Face ID | Checkpoint | ONNX export | +| --- | --- | --- | --- | +| PixelDiT-1300M-1024px | `nvidia/PixelDiT-1300M-1024px` | FP32 (runs in BF16) | BF16 weights, FP32 I/O | +| Gemma-2-2B-IT (text encoder) | `Efficient-Large-Model/gemma-2-2b-it` | BF16 | BF16 weights, FP32 I/O | + +PixelDiT generates pixels directly, so there is no VAE. The model code is vendored in `model_export/pixeldit/`, so the [NVlabs/PixelDiT](https://github.com/NVlabs/PixelDiT) repo is not needed to export or verify. + +## Pipeline + +| Stage | Where | Details | +| --- | --- | --- | +| Tokenize | Host | Gemma tokenizer. The positive prompt is prefixed with a fixed instruction prompt and padded to 506 tokens; the negative prompt is padded to 506 too. | +| Text encoder | `text_encoder/model.onnx` | Gemma-2 decoder, then `select_index` picks the 300 rows the DiT uses. | +| DiT | `transformer/model.onnx` | Batch of 2: negative and positive prompt for classifier-free guidance. Returns velocity. | +| Sampling | Host or CUDA (`--processing`) | CFG and flow DPM-Solver++ (2nd order), 50 steps, CFG scale 2.75, flow shift 4.0 by default. | + +Sampling defaults and the token layout are written to `pipeline_config.json`. The C++ runtime reads it at startup and falls back to built-in defaults (matching the 1024x1024 export) for missing fields. + +## Backends + +| Mode | CLI | Inference and sampling | +| --- | --- | --- | +| TensorRT RTX + CPU | `--provider trt-rtx --processing cpu` | TensorRT RTX inference, sampling on the CPU. | +| TensorRT RTX + CUDA | `--provider trt-rtx --processing cuda` | TensorRT RTX inference, sampling in CUDA kernels. | + +Both modes run the models through ONNX Runtime with the TensorRT RTX execution provider. + +CUDA sampling keeps the sample on the GPU, so the DiT calls run back to back with one host sync per image. Both modes produce identical images for the same seed. + +DirectX and Vulkan processing are not implemented yet. `--provider cpu` is not supported. + +## Export + +```bash +cd model_export +python export_pixeldit.py --output ./models/PixelDiT-1300M-1024px-onnx +``` + +An export contains `text_encoder/model.onnx`, `transformer/model.onnx`, `tokenizer/tokenizer.json`, and `pipeline_config.json`. + +Shapes are static. `--height` and `--width` (default 1024, multiples of 16) set the image size; re-export to change it. Add `--compile_trt --trt_root ` to check that each model builds with the TensorRT RTX CLI. + +## Verify + +```bash +python verify_pixeldit.py --onnx_dir ./models/PixelDiT-1300M-1024px-onnx --provider trt-rtx +``` + +Compares each ONNX model running on TensorRT RTX with the vendored PyTorch model, then generates one image with each and saves `pytorch_output.png`, `ort_output.png`, and `comparison.png`. Pass `--validate_only` to only check the graphs. + +Per-model checks decide pass or fail. The image check only fails on broken output: 50-step BF16 sampling amplifies tiny rounding differences, so correct runs can differ in fine detail, or with some seeds in composition. Compare the images visually. + +## Build + +```powershell +cmake --preset windows-x64 +cmake --build --preset windows-x64-release --target din_pixeldit_cli +``` + +## Run + +```powershell +out\build\windows-x64\bin\Release\din_pixeldit_cli.exe --model-dir D:\models\PixelDiT-1300M-1024px-onnx --provider trt-rtx --processing cuda --prompt "a red fox" --output out +``` + +`--model-dir` defaults to `./models/PixelDiT-1300M-1024px-onnx`, relative to the working directory. + +Optional: `--negative-prompt`, `--seed`, `--num-images`, `--steps`, `--cfg-scale`, `--flow-shift` (defaults come from `pipeline_config.json`). The C++ noise generator differs from PyTorch's, so the same seed gives a different image than the Python scripts. diff --git a/image_gen/pixeldit/cuda_kernels.cu b/image_gen/pixeldit/cuda_kernels.cu new file mode 100644 index 0000000..2f486aa --- /dev/null +++ b/image_gen/pixeldit/cuda_kernels.cu @@ -0,0 +1,111 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include "cuda_kernels.h" + +namespace din::image_gen +{ +namespace +{ + +constexpr unsigned int kThreadsPerBlock = 256; + +void check_launch() +{ + const cudaError_t err = cudaPeekAtLastError(); + if (err != cudaSuccess) + { + throw std::runtime_error(std::string("CUDA kernel launch failed: ") + cudaGetErrorString(err)); + } +} + +dim3 grid_for(size_t total_elements) +{ + const size_t blocks = (total_elements + kThreadsPerBlock - 1) / kThreadsPerBlock; + if (blocks > std::numeric_limits::max()) + { + throw std::runtime_error("CUDA grid is too large"); + } + return dim3(static_cast(blocks), 1, 1); +} + +__global__ void pixeldit_cfg_x0_kernel(const float* velocity, const float* x, float t, float cfg_scale, float* x0, + size_t n) +{ + const size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) + { + return; + } + const float one_minus_t = 1.0f - t; + const float scaled_u = one_minus_t * velocity[i]; + const float noise_u = scaled_u + x[i]; + const float scaled_c = one_minus_t * velocity[n + i]; + const float noise_c = scaled_c + x[i]; + const float diff = noise_c - noise_u; + const float guided_delta = cfg_scale * diff; + const float noise = noise_u + guided_delta; + const float t_noise = t * noise; + const float numerator = x[i] - t_noise; + x0[i] = numerator / one_minus_t; +} + +__global__ void pixeldit_dpm1_kernel(float* x, const float* model_s, float ratio, float coeff, size_t n) +{ + const size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) + { + return; + } + const float a = ratio * x[i]; + const float b = coeff * model_s[i]; + x[i] = a - b; +} + +__global__ void pixeldit_dpm2_kernel(float* x, const float* model_prev_1, const float* model_prev_0, float ratio, + float coeff, float half_coeff, float inv_r0, size_t n) +{ + const size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) + { + return; + } + const float delta = model_prev_0[i] - model_prev_1[i]; + const float d1 = inv_r0 * delta; + const float a = ratio * x[i]; + const float b = coeff * model_prev_0[i]; + const float c = half_coeff * d1; + const float ab = a - b; + x[i] = ab - c; +} + +} // namespace + +void launch_pixeldit_cfg_x0_kernel(cudaStream_t stream, const float* velocity, const float* sample, float t, + float cfg_scale, float* x0, size_t n) +{ + pixeldit_cfg_x0_kernel<<>>(velocity, sample, t, cfg_scale, x0, n); + check_launch(); +} + +void launch_pixeldit_dpm1_kernel(cudaStream_t stream, float* sample, const float* model_s, float ratio, float coeff, + size_t n) +{ + pixeldit_dpm1_kernel<<>>(sample, model_s, ratio, coeff, n); + check_launch(); +} + +void launch_pixeldit_dpm2_kernel(cudaStream_t stream, float* sample, const float* model_prev_1, + const float* model_prev_0, float ratio, float coeff, float half_coeff, float inv_r0, + size_t n) +{ + pixeldit_dpm2_kernel<<>>(sample, model_prev_1, model_prev_0, ratio, coeff, + half_coeff, inv_r0, n); + check_launch(); +} + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/cuda_kernels.h b/image_gen/pixeldit/cuda_kernels.h new file mode 100644 index 0000000..f3ff4cf --- /dev/null +++ b/image_gen/pixeldit/cuda_kernels.h @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include + +namespace din::image_gen +{ + +// velocity: [2, n] ([negative, positive]); sample: [n] at time t -> x0: [n]. See cfg_data_prediction. +void launch_pixeldit_cfg_x0_kernel(cudaStream_t stream, const float* velocity, const float* sample, float t, + float cfg_scale, float* x0, size_t n); + +// sample = ratio * sample - coeff * model_s. See dpm_first_order_update (t > 0). +void launch_pixeldit_dpm1_kernel(cudaStream_t stream, float* sample, const float* model_s, float ratio, float coeff, + size_t n); + +// See dpm_second_order_update. +void launch_pixeldit_dpm2_kernel(cudaStream_t stream, float* sample, const float* model_prev_1, + const float* model_prev_0, float ratio, float coeff, float half_coeff, float inv_r0, + size_t n); + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/helpers/pixeldit.h b/image_gen/pixeldit/helpers/pixeldit.h new file mode 100644 index 0000000..fdd4c03 --- /dev/null +++ b/image_gen/pixeldit/helpers/pixeldit.h @@ -0,0 +1,375 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tokenizer.h" + +namespace din::image_gen +{ + +// ============================================================================ +// Configuration +// ============================================================================ + +// Relative to the working directory; pass --model-dir to use an export elsewhere. +static const std::filesystem::path DEFAULT_PIXELDIT_MODEL_DIR = "./models/PixelDiT-1300M-1024px-onnx"; + +struct PixelDiTModelPaths +{ + std::filesystem::path base_dir; + std::filesystem::path text_encoder_model; + std::filesystem::path transformer_model; + std::filesystem::path tokenizer_json; + std::filesystem::path pipeline_config; +}; + +inline PixelDiTModelPaths MakePixelDiTModelPaths(std::filesystem::path model_dir) +{ + if (model_dir.empty()) + { + model_dir = DEFAULT_PIXELDIT_MODEL_DIR; + } + return { + model_dir, + model_dir / "text_encoder/model.onnx", + model_dir / "transformer/model.onnx", + model_dir / "tokenizer/tokenizer.json", + model_dir / "pipeline_config.json", + }; +} + +inline constexpr const char* DEFAULT_PIXELDIT_PROMPT = + "A red fox sitting in a snowy forest at sunrise, soft golden light, photorealistic"; + +// Upstream t2i/configs/PixelDiT_1024px_pixel_diffusion_stage3.yaml text_encoder.chi_prompt, joined with "\n" +// (same as model_export/pixeldit/pipeline.py CHI_PROMPT). +inline constexpr const char* PIXELDIT_CHI_PROMPT = + "Given a user prompt, generate an \"Enhanced prompt\" that provides detailed visual descriptions suitable for " + "image generation. Evaluate the level of detail in the user prompt:\n" + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial " + "relationships to create vivid and concrete scenes.\n" + "- If the prompt is already detailed, refine and enhance the existing details slightly without " + "overcomplicating.\n" + "Here are examples of how to transform or refine prompts:\n" + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping " + "peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.\n" + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street " + "lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass " + "skyscrapers.\n" + "Please generate only the enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:\n" + "User Prompt: "; + +inline constexpr const char* PIXELDIT_DEFAULT_NEGATIVE_PROMPT = + "low quality, worst quality, over-saturated, blurry, deformed, watermark"; + +// pipeline_config.json written by export_pixeldit.py. Every field has a built-in default matching the +// 1024x1024 export; LoadPixelDiTPipelineConfig records which defaults were used. +struct PixelDiTPipelineConfig +{ + int64_t height = 1024; + int64_t width = 1024; + int64_t dit_batch = 2; // [negative, positive] for classifier-free guidance + float timestep_scale = 1000.0f; + + std::string chi_prompt = PIXELDIT_CHI_PROMPT; + int64_t text_seq_len = 506; // chi_prompt tokens (208) + txt_max_length (300) - 2 + int64_t txt_max_length = 300; + int64_t bos_token_id = 2; + int64_t pad_token_id = 0; + std::vector positive_select_index; // BOS + last (txt_max_length - 1) positions + std::vector negative_select_index; // first txt_max_length positions + + int steps = 50; + float cfg_scale = 2.75f; + float flow_shift = 4.0f; + std::string negative_prompt = PIXELDIT_DEFAULT_NEGATIVE_PROMPT; + float t_start = 1.0f; + float t_end = 0.001f; + int order = 2; + + bool file_found = false; + std::vector defaults_used; +}; + +PixelDiTPipelineConfig LoadPixelDiTPipelineConfig(const std::filesystem::path& path); + +inline std::vector DefaultPositiveSelectIndex(int64_t text_seq_len, int64_t txt_max_length) +{ + std::vector index{0}; + for (int64_t i = text_seq_len - txt_max_length + 1; i < text_seq_len; ++i) + { + index.push_back(i); + } + return index; +} + +inline std::vector DefaultNegativeSelectIndex(int64_t txt_max_length) +{ + std::vector index(static_cast(txt_max_length)); + for (int64_t i = 0; i < txt_max_length; ++i) + { + index[static_cast(i)] = i; + } + return index; +} + +// ============================================================================ +// Text encoder inputs +// ============================================================================ + +struct PixelDiTTextInputs +{ + std::vector input_ids; + std::vector attention_mask; + size_t num_tokens = 0; // before padding / truncation +}; + +// Python str.strip() for ASCII whitespace. +inline std::string StripWhitespace(const std::string& text) +{ + const char* whitespace = " \t\n\r\f\v"; + const auto first = text.find_first_not_of(whitespace); + if (first == std::string::npos) + { + return {}; + } + const auto last = text.find_last_not_of(whitespace); + return text.substr(first, last - first + 1); +} + +// HF tokenizer(text, max_length=seq_len, padding="max_length", truncation=True): + tokens, truncated, +// right-padded with pad_token_id. +inline PixelDiTTextInputs MakeTextEncoderInputs(const din::io::Tokenizer& tokenizer, const std::string& text, + int64_t seq_len, int64_t pad_token_id) +{ + const std::vector ids = tokenizer.Encode(text, true); + PixelDiTTextInputs inputs; + inputs.num_tokens = ids.size(); + inputs.input_ids.assign(static_cast(seq_len), pad_token_id); + inputs.attention_mask.assign(static_cast(seq_len), 0); + const size_t kept = std::min(ids.size(), static_cast(seq_len)); + for (size_t i = 0; i < kept; ++i) + { + inputs.input_ids[i] = ids[i]; + inputs.attention_mask[i] = 1; + } + return inputs; +} + +// ============================================================================ +// Random noise (same generator as Flux2's initialize_latent) +// ============================================================================ + +inline void initialize_noise(float* noise, size_t num_elements, unsigned int seed) +{ + std::default_random_engine rng(seed); + std::normal_distribution dist(0.0f, 1.0f); + for (size_t i = 0; i < num_elements; ++i) + noise[i] = dist(rng); +} + +// ============================================================================ +// Flow DPM-Solver++ (2M) sampling math, host side. +// +// Mirrors model_export/pixeldit/pipeline.py (and upstream dpm_solver.py) operation by operation in fp32, +// so results match the Python reference up to libm differences in the scalar log/exp terms. +// Build with floating-point contraction disabled (see CMakeLists.txt). +// ============================================================================ + +// Runs fn(begin, end) over [0, n) on several threads. The sampling math is elementwise, so results are +// identical to a single-threaded loop. +template +void ParallelFor(size_t n, Fn&& fn) +{ + constexpr size_t min_chunk = size_t{1} << 16; + const size_t threads = std::min( + {std::max(1u, std::thread::hardware_concurrency()), size_t{16}, (n + min_chunk - 1) / min_chunk}); + if (threads <= 1) + { + fn(size_t{0}, n); + return; + } + const size_t chunk = (n + threads - 1) / threads; + std::vector workers; + workers.reserve(threads - 1); + for (size_t k = 1; k < threads; ++k) + { + const size_t begin = k * chunk; + const size_t end = std::min(n, begin + chunk); + workers.emplace_back( + [&fn, begin, end] + { + fn(begin, end); + }); + } + fn(size_t{0}, std::min(n, chunk)); + for (auto& worker : workers) + { + worker.join(); + } +} + +// DPM_Solver.get_time_steps("time_uniform_flow"): torch.linspace(t_start, t_end, steps + 1) in fp32, then +// t = flip(shift * s / (1 + (shift - 1) * s)) with s = 1 - linspace. Returns steps + 1 values, ~1 down to 0. +inline std::vector flow_timesteps(int steps, float flow_shift, float t_start, float t_end) +{ + const int count = steps + 1; + const float step = (t_end - t_start) / static_cast(count - 1); + const int halfway = count / 2; + std::vector betas(static_cast(count)); + for (int i = 0; i < count; ++i) + { + betas[static_cast(i)] = + i < halfway ? t_start + step * static_cast(i) : t_end - step * static_cast(count - i - 1); + } + std::vector timesteps(static_cast(count)); + for (int i = 0; i < count; ++i) + { + const float s = 1.0f - betas[static_cast(i)]; + const float numerator = flow_shift * s; + const float scaled = (flow_shift - 1.0f) * s; + const float denominator = 1.0f + scaled; + timesteps[static_cast(count - 1 - i)] = numerator / denominator; + } + return timesteps; +} + +// Noise schedule (NoiseScheduleFlow): alpha_t = 1 - t, sigma_t = t, lambda_t = log(alpha_t) - log(sigma_t). +inline float flow_log_alpha(float t) +{ + return std::log(1.0f - t); +} + +inline float flow_lambda(float t) +{ + return flow_log_alpha(t) - std::log(t); +} + +// velocity: [2, n] DiT output for [negative, positive]; x: [n] current sample at time t. +// noise_b = (1 - t) * v_b + x; noise = noise_u + cfg * (noise_c - noise_u); x0 = (x - t * noise) / (1 - t). +inline void cfg_data_prediction(const float* velocity, const float* x, float t, float cfg_scale, float* x0, size_t n) +{ + const float one_minus_t = 1.0f - t; + const float* v_uncond = velocity; + const float* v_cond = velocity + n; + ParallelFor(n, + [=](size_t begin, size_t end) + { + for (size_t i = begin; i < end; ++i) + { + const float scaled_u = one_minus_t * v_uncond[i]; + const float noise_u = scaled_u + x[i]; + const float scaled_c = one_minus_t * v_cond[i]; + const float noise_c = scaled_c + x[i]; + const float diff = noise_c - noise_u; + const float guided_delta = cfg_scale * diff; + const float noise = noise_u + guided_delta; + const float t_noise = t * noise; + const float numerator = x[i] - t_noise; + x0[i] = numerator / one_minus_t; + } + }); +} + +// Per-step scalars of the solver updates, computed on the host for both the CPU loops below and the CUDA kernels +// (cuda_kernels.cu), so log/exp never run on the device and both paths see the same coefficients. +struct DpmFirstOrderCoeffs +{ + float ratio = 0.0f; + float coeff = 0.0f; +}; + +struct DpmSecondOrderCoeffs +{ + float ratio = 0.0f; + float coeff = 0.0f; + float half_coeff = 0.0f; + float inv_r0 = 0.0f; +}; + +// DPM-Solver-1 (dpmsolver++) from s to t: x_t = (t / s) * x - (alpha_t * phi_1) * x0_s, phi_1 = expm1(-h). +// Only for t > 0; at t == 0 the update is exactly x0_s (see dpm_first_order_update). +inline DpmFirstOrderCoeffs dpm_first_order_coeffs(float s, float t) +{ + const float h = flow_lambda(t) - flow_lambda(s); + const float alpha_t = std::exp(flow_log_alpha(t)); + const float phi_1 = std::expm1(-h); + DpmFirstOrderCoeffs coeffs; + coeffs.ratio = t / s; + coeffs.coeff = alpha_t * phi_1; + return coeffs; +} + +// Multistep DPM-Solver-2 (dpmsolver++) from t_prev_0 to t, using x0 at t_prev_1 and t_prev_0. +inline DpmSecondOrderCoeffs dpm_second_order_coeffs(float t_prev_1, float t_prev_0, float t) +{ + const float lambda_prev_1 = flow_lambda(t_prev_1); + const float lambda_prev_0 = flow_lambda(t_prev_0); + const float lambda_t = flow_lambda(t); + const float alpha_t = std::exp(flow_log_alpha(t)); + const float h_0 = lambda_prev_0 - lambda_prev_1; + const float h = lambda_t - lambda_prev_0; + const float r0 = h_0 / h; + const float phi_1 = std::expm1(-h); + DpmSecondOrderCoeffs coeffs; + coeffs.inv_r0 = 1.0f / r0; + coeffs.ratio = t / t_prev_0; + coeffs.coeff = alpha_t * phi_1; + coeffs.half_coeff = 0.5f * coeffs.coeff; + return coeffs; +} + +// At t == 0 (sigma_t = 0, alpha_t = 1, expm1(-inf) = -1) the update is exactly x0_s. +inline void dpm_first_order_update(float* x, const float* model_s, float s, float t, size_t n) +{ + if (t == 0.0f) + { + std::copy_n(model_s, n, x); + return; + } + const DpmFirstOrderCoeffs k = dpm_first_order_coeffs(s, t); + ParallelFor(n, + [=](size_t begin, size_t end) + { + for (size_t i = begin; i < end; ++i) + { + const float a = k.ratio * x[i]; + const float b = k.coeff * model_s[i]; + x[i] = a - b; + } + }); +} + +inline void dpm_second_order_update(float* x, const float* model_prev_1, const float* model_prev_0, float t_prev_1, + float t_prev_0, float t, size_t n) +{ + const DpmSecondOrderCoeffs k = dpm_second_order_coeffs(t_prev_1, t_prev_0, t); + ParallelFor(n, + [=](size_t begin, size_t end) + { + for (size_t i = begin; i < end; ++i) + { + const float delta = model_prev_0[i] - model_prev_1[i]; + const float d1 = k.inv_r0 * delta; + const float a = k.ratio * x[i]; + const float b = k.coeff * model_prev_0[i]; + const float c = k.half_coeff * d1; + const float ab = a - b; + x[i] = ab - c; + } + }); +} + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/model_export/export_pixeldit.py b/image_gen/pixeldit/model_export/export_pixeldit.py new file mode 100644 index 0000000..6e7f3e9 --- /dev/null +++ b/image_gen/pixeldit/model_export/export_pixeldit.py @@ -0,0 +1,414 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +#!/usr/bin/env python3 +""" +Export PixelDiT text-to-image (nvidia/PixelDiT-1300M-1024px) and its Gemma-2 text encoder to ONNX +for use with ONNX Runtime / TensorRT / C++ inference. + +Usage: + python export_pixeldit.py --output ./models/PixelDiT-1300M-1024px-onnx + python export_pixeldit.py --output ./out --model transformer --height 768 --width 1344 + +Output layout: + / + text_encoder/model.onnx + model.onnx_data + transformer/model.onnx + model.onnx_data + tokenizer/ (Gemma tokenizer files) + pipeline_config.json + +All input shapes are static. The DiT runs negative + positive prompts as one batch of 2 (CFG); +CFG and the flow DPM-Solver++ update stay outside the graphs (see pixeldit/pipeline.py). +""" + +import argparse +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import torch + +from pixeldit import pipeline as pixeldit_pipeline +from pixeldit.weights import DEFAULT_MODEL_NAME, DEFAULT_TEXT_ENCODER_NAME, load_gemma_text_encoder, load_pixeldit + +EXTERNAL_DATA_NAME = "model.onnx_data" +DEFAULT_MODEL_OPSETS = { + "text_encoder": 25, + "transformer": 25, +} +DIT_BATCH = 2 # [negative, positive] for classifier-free guidance + +# ONNX model IO is fp32 because the C++ pipeline owns fp32 buffers at graph boundaries. +IO_PRECISION_MAP = { + "fp32": torch.float32, +} + + +def _configure_stdio(): + """Keep PyTorch ONNX status output from failing on Windows cp1252 consoles.""" + for stream_name in ("stdout", "stderr"): + stream = getattr(sys, stream_name, None) + if hasattr(stream, "reconfigure"): + stream.reconfigure(encoding="utf-8", errors="replace") + + +def _onnx_export( + model, + model_args: tuple, + output_path: Path, + input_names: list, + output_names: list, + opset: int, + **kwargs +): + """Export model to ONNX and always save weights as external data in model.onnx_data.""" + import onnx + + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export( + model, + model_args, + str(output_path), + input_names=input_names, + output_names=output_names, + do_constant_folding=True, + dynamo=True, + opset_version=opset, + **kwargs + ) + # Reload and save with external data so we always use model.onnx_data + onnx_model = onnx.load(str(output_path)) + out_dir = output_path.parent + for f in out_dir.iterdir(): + if f.name != output_path.name: + f.unlink() + onnx.save_model( + onnx_model, + str(output_path), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=EXTERNAL_DATA_NAME, + convert_attribute=False, + ) + print(f" ONNX saved: {output_path} (opset={opset})") + + +class TextEncoderExportWrapper(torch.nn.Module): + """Gemma-2 decoder final hidden state, gathered to the 300 rows the DiT consumes. + + select_index picks BOS + the last 299 positions for the positive (chi_prompt-prefixed) + prompt, or the first 300 positions for the negative prompt. Output is cast to io_dtype. + """ + + def __init__(self, text_encoder, io_dtype: torch.dtype): + super().__init__() + self.text_encoder = text_encoder + self.io_dtype = io_dtype + + def forward(self, input_ids, attention_mask, select_index): + hidden = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)[0] + return hidden.index_select(1, select_index).to(self.io_dtype) + + +class PixelDiTExportWrapper(torch.nn.Module): + """Inputs are cast to bf16 inside PixelDiT (as upstream PixDiTTrainer does), output to io_dtype.""" + + def __init__(self, model, io_dtype: torch.dtype): + super().__init__() + self.model = model + self.io_dtype = io_dtype + + def forward(self, hidden_states, timestep, encoder_hidden_states): + return self.model(hidden_states, timestep, encoder_hidden_states).to(self.io_dtype) + + +@torch.no_grad() +def export_text_encoder(text_encoder, tokenizer, layout, output_path: Path, device: str, opset: int, io_dtype: torch.dtype): + """Export with static shape (1, text_seq_len) -> (1, txt_max_length, hidden).""" + wrapper = TextEncoderExportWrapper(text_encoder, io_dtype).eval() + dummy_ids, dummy_mask = pixeldit_pipeline.tokenize(tokenizer, layout.chi_prompt, layout.text_seq_len, device) + dummy_select = torch.tensor(layout.positive_select_index(), dtype=torch.long, device=device) + _onnx_export( + wrapper, + (dummy_ids, dummy_mask, dummy_select), + output_path / "text_encoder" / "model.onnx", + input_names=["input_ids", "attention_mask", "select_index"], + output_names=["prompt_embeds"], + opset=opset, + ) + print(" text_encoder exported.") + + +@torch.no_grad() +def export_transformer(model, output_path: Path, device: str, opset: int, height: int, width: int, io_dtype: torch.dtype): + """Export with static shapes (2, 3, height, width), (2,), (2, txt_max_length, txt_embed_dim).""" + wrapper = PixelDiTExportWrapper(model, io_dtype).eval() + txt_embed_dim = model.y_embedder.proj.in_features + dummy_x = torch.randn(DIT_BATCH, 3, height, width, device=device, dtype=io_dtype) + dummy_t = torch.full((DIT_BATCH,), 500.0, device=device, dtype=io_dtype) + dummy_y = torch.randn(DIT_BATCH, model.txt_max_length, txt_embed_dim, device=device, dtype=io_dtype) + _onnx_export( + wrapper, + (dummy_x, dummy_t, dummy_y), + output_path / "transformer" / "model.onnx", + input_names=["hidden_states", "timestep", "encoder_hidden_states"], + output_names=["velocity"], + opset=opset, + ) + print(" transformer exported.") + + +def write_pipeline_config(output_path: Path, args, tokenizer, layout, sampler) -> None: + config = { + "model_name": args.model_name, + "text_encoder_name": args.text_encoder_name, + "height": args.height, + "width": args.width, + "io_precision": args.io_precision, + "compute_dtype": "bf16", + "dit_batch": DIT_BATCH, + "dit_batch_order": ["negative", "positive"], + "timestep_scale": pixeldit_pipeline.NUM_TRAIN_TIMESTEPS, + "text": { + "chi_prompt": layout.chi_prompt, + "chi_prompt_tokens": layout.chi_prompt_tokens, + "text_seq_len": layout.text_seq_len, + "txt_max_length": layout.txt_max_length, + "padding_side": "right", + "bos_token_id": tokenizer.bos_token_id, + "pad_token_id": tokenizer.pad_token_id, + "positive_select_index": layout.positive_select_index(), + "negative_select_index": layout.negative_select_index(), + }, + "sampler": { + "algorithm": "flow_dpm_solver++_multistep", + "order": 2, + "lower_order_final": True, + "t_start": 1.0, + "t_end": 1.0 / pixeldit_pipeline.NUM_TRAIN_TIMESTEPS, + "steps": sampler.steps, + "cfg_scale": sampler.cfg_scale, + "flow_shift": sampler.flow_shift, + "negative_prompt": sampler.negative_prompt, + }, + } + with open(output_path / "pipeline_config.json", "w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + print(" pipeline_config.json written.") + + +def _find_trt_binary(trt_root: str | None, trt_bin: str | None) -> Path: + if trt_bin: + candidate = Path(trt_bin).expanduser().resolve() + if not candidate.is_file(): + raise FileNotFoundError(f"TensorRT RTX binary not found: {candidate}") + return candidate + + exe_name = "tensorrt_rtx.exe" if os.name == "nt" else "tensorrt_rtx" + if trt_root: + candidate = Path(trt_root).expanduser().resolve() / "bin" / exe_name + if not candidate.is_file(): + raise FileNotFoundError(f"TensorRT RTX binary not found: {candidate}") + return candidate + + found = shutil.which(exe_name) or shutil.which("tensorrt_rtx") + if found: + return Path(found).resolve() + + raise FileNotFoundError("TensorRT RTX binary not found; pass --trt_root or --trt_bin") + + +def _trt_subprocess_env(trt_root: str | None) -> dict[str, str]: + env = os.environ.copy() + if not trt_root: + return env + + root = Path(trt_root).expanduser().resolve() + if os.name == "nt": + env["PATH"] = str(root / "bin") + os.pathsep + env.get("PATH", "") + else: + env["LD_LIBRARY_PATH"] = str(root / "lib") + os.pathsep + env.get("LD_LIBRARY_PATH", "") + return env + + +def compile_with_trt_rtx( + onnx_path: Path, + trt_exe: Path, + trt_root: str | None, + save_engine: bool = False, + extra_args: list[str] | None = None, +): + cmd = [str(trt_exe), f"--onnx={onnx_path}", "--skipInference"] + if save_engine: + cmd.append(f"--saveEngine={onnx_path.with_suffix('.engine')}") + if extra_args: + cmd.extend(extra_args) + + print(f" TRT RTX compile: {onnx_path}") + subprocess.run(cmd, check=True, env=_trt_subprocess_env(trt_root)) + print(" TRT RTX compile passed.") + + +MODELS = ["text_encoder", "transformer"] + + +def main(): + _configure_stdio() + + ap = argparse.ArgumentParser( + description="Export PixelDiT T2I and its Gemma-2 text encoder to ONNX (each model in model.onnx + model.onnx_data)", + ) + ap.add_argument( + "--model_name", + type=str, + default=DEFAULT_MODEL_NAME, + help="Hugging Face model repo ID, a local snapshot directory, or a local pixeldit_t2i_v1.pth", + ) + ap.add_argument( + "--text_encoder_name", + type=str, + default=DEFAULT_TEXT_ENCODER_NAME, + help="Hugging Face repo ID or local directory of the Gemma-2 text encoder", + ) + ap.add_argument( + "--local_files_only", + action="store_true", + help="Resolve models from the local Hugging Face cache only.", + ) + ap.add_argument( + "--output", + type=str, + default="./pixeldit_onnx", + help="Output directory for ONNX models and configs", + ) + ap.add_argument( + "--model", + type=str, + nargs="+", + default=None, + metavar="NAME", + help="Model(s) to export: one or more of all, text_encoder, transformer (default: all)", + ) + ap.add_argument( + "--opset", + type=int, + default=None, + help=( + "Override ONNX opset for every selected model. By default the script " + f"uses per-model opsets: {DEFAULT_MODEL_OPSETS}." + ), + ) + ap.add_argument("--height", type=int, default=1024, help="Image height for static shapes (multiple of 16, default 1024)") + ap.add_argument("--width", type=int, default=1024, help="Image width for static shapes (multiple of 16, default 1024)") + ap.add_argument( + "--io_precision", + type=str, + choices=list(IO_PRECISION_MAP.keys()), + default="fp32", + help="IO precision of ONNX models; only fp32 is supported by the C++ pipeline", + ) + ap.add_argument( + "--compile_trt", + action="store_true", + help="After export, build each selected ONNX model with the TensorRT RTX CLI and --skipInference.", + ) + ap.add_argument( + "--trt_root", + type=str, + default=None, + help="TensorRT RTX install root containing bin/ and lib/ (for example TensorRT-RTX-1.5.0.114).", + ) + ap.add_argument( + "--trt_bin", + type=str, + default=None, + help="Explicit path to tensorrt_rtx executable. Overrides --trt_root.", + ) + ap.add_argument( + "--trt_save_engines", + action="store_true", + help="Save compiled engines next to each ONNX file as model.engine.", + ) + ap.add_argument( + "--trt_arg", + action="append", + default=[], + help="Additional argument to pass to tensorrt_rtx. May be specified more than once.", + ) + args = ap.parse_args() + + if args.height % 16 or args.width % 16: + ap.error("--height and --width must be multiples of 16 (the PixelDiT patch size)") + + output_path = Path(args.output).resolve() + output_path.mkdir(parents=True, exist_ok=True) + + model_choices = ["all"] + MODELS + raw = args.model if args.model is not None else ["all"] + to_export: list[str] = [] + for m in raw: + if m not in model_choices: + ap.error(f"--model must be one or more of: {', '.join(model_choices)}") + if m == "all": + to_export = list(MODELS) + break + if m not in to_export: + to_export.append(m) + + device = "cuda" + io_dtype = IO_PRECISION_MAP[args.io_precision] + + print(f"Loading tokenizer from {args.text_encoder_name}...") + tokenizer, text_encoder = None, None + if "text_encoder" in to_export: + tokenizer, text_encoder = load_gemma_text_encoder(args.text_encoder_name, device, args.local_files_only) + print(f" Gemma attention implementation: {text_encoder.config._attn_implementation}") + else: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, local_files_only=args.local_files_only) + layout = pixeldit_pipeline.TextLayout.from_tokenizer(tokenizer) + print(f" chi_prompt tokens={layout.chi_prompt_tokens}, text_seq_len={layout.text_seq_len}") + + 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": + export_text_encoder(text_encoder, tokenizer, layout, output_path, device, model_opset, io_dtype) + del text_encoder + torch.cuda.empty_cache() + elif name == "transformer": + print(f"Loading PixelDiT from {args.model_name} (dtype=bf16, {args.height}x{args.width})...") + model = load_pixeldit(args.model_name, device, torch.bfloat16, args.local_files_only) + model.set_image_size(args.height, args.width) + export_transformer(model, output_path, device, model_opset, args.height, args.width, io_dtype) + del model + torch.cuda.empty_cache() + exported_model_paths.append(output_path / name / "model.onnx") + + if "transformer" in to_export: + tokenizer.save_pretrained(str(output_path / "tokenizer")) + print(" Saved tokenizer/") + write_pipeline_config(output_path, args, tokenizer, layout, pixeldit_pipeline.SamplerConfig()) + + if args.compile_trt: + trt_exe = _find_trt_binary(args.trt_root, args.trt_bin) + print(f"Compiling exported ONNX models with TensorRT RTX: {trt_exe}") + for onnx_path in exported_model_paths: + compile_with_trt_rtx( + onnx_path, + trt_exe=trt_exe, + trt_root=args.trt_root, + save_engine=args.trt_save_engines, + extra_args=args.trt_arg, + ) + + print(f"\nONNX export done -> {output_path}") + + +if __name__ == "__main__": + main() diff --git a/image_gen/pixeldit/model_export/pixeldit/__init__.py b/image_gen/pixeldit/model_export/pixeldit/__init__.py new file mode 100644 index 0000000..dde4ab1 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone PixelDiT T2I model, weights loader and sampling pipeline for ONNX export.""" diff --git a/image_gen/pixeldit/model_export/pixeldit/model.py b/image_gen/pixeldit/model_export/pixeldit/model.py new file mode 100644 index 0000000..1aa1006 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/model.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PixelDiT text-to-image DiT, vendored from NVlabs/PixelDiT (pixdit_core/) for ONNX export. + +Numerically equivalent to upstream PixDiTTrainer + PixDiT_T2I at inference. Changes: + * RoPE uses real cos/sin tables instead of complex tensors. Required: upstream unsqueezes + complex tensors (precompute_freqs_cis_2d, and freqs_cis[None, :, None, :] in + apply_rotary_emb), and the ONNX exporter has no lowering for unsqueeze on complex inputs. + Other complex ops (polar, view_as_complex/real, multiply) would export. + * Patchify / unpatchify use reshape + permute instead of F.unfold / F.fold (bit-identical, + since kernel == stride). Required for TensorRT RTX: the model runs them in bf16, and + TensorRT RTX has no bf16 Col2Im (F.fold); the unfold+fold export fails to create a session. + * Position tables are built once by set_image_size() instead of lazily filled Python dict + caches, so tracing never writes trace-time tensors into a cache. They are plain + attributes (not buffers) so Module.to(dtype) leaves the fp32 RoPE tables alone, as + upstream's complex64 tables are. +Training-only pieces (REPA projector, attention masks, class embedder) are dropped. +""" + +import math + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Matches nvidia/PixelDiT-1300M-1024px config.json. +PIXELDIT_1300M_CONFIG = { + "in_channels": 3, + "patch_size": 16, + "num_groups": 24, + "hidden_size": 1536, + "pixel_hidden_size": 16, + "pixel_attn_hidden_size": 1152, + "pixel_num_groups": 16, + "patch_depth": 14, + "pixel_depth": 2, + "txt_embed_dim": 2304, + "txt_max_length": 300, + "text_rope_theta": 10000.0, +} + + +# --------------------------------------------------------------------------- +# Position tables (match pixdit_core/modules.py) +# --------------------------------------------------------------------------- + + +def _sincos_1d(embed_dim, pos): + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega + out = np.einsum("m,d->md", pos.reshape(-1), omega) + return np.concatenate([np.sin(out), np.cos(out)], axis=1) + + +def pixel_sincos_pos_embed(embed_dim: int, height: int, width: int) -> np.ndarray: + """[H*W, D] absolute pixel embedding (get_2d_sincos_pos_embed[_from_grid] upstream).""" + grid_h = np.arange(height, dtype=np.float32) + grid_w = np.arange(width, dtype=np.float32) + grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape(2, 1, height, width) + emb_h = _sincos_1d(embed_dim // 2, grid[0]) + emb_w = _sincos_1d(embed_dim // 2, grid[1]) + return np.concatenate([emb_h, emb_w], axis=1) + + +def rope_2d_angles(dim: int, height: int, width: int, theta: float = 10000.0, scale: float = 16.0) -> torch.Tensor: + """[H*W, dim/2] angles; upstream precompute_freqs_cis_2d is torch.polar(1, angles).""" + x_pos = torch.linspace(0, scale, width) + y_pos = torch.linspace(0, scale, height) + y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij") + freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) + x_freqs = torch.outer(x_pos.reshape(-1), freqs).float() + y_freqs = torch.outer(y_pos.reshape(-1), freqs).float() + return torch.stack([x_freqs, y_freqs], dim=-1).reshape(height * width, -1) + + +def rope_1d_angles(head_dim: int, length: int, theta: float, device=None) -> torch.Tensor: + """[L, head_dim/2] angles; upstream PixDiT_T2I.fetch_pos_text.""" + freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) + positions = torch.arange(0, length, device=device).float().unsqueeze(1) + return positions * freqs.unsqueeze(0) + + +def apply_rotary_emb(xq, xk, cos, sin): + """Real-valued equivalent of upstream complex apply_rotary_emb. + + xq/xk: [B, N, H, D] with (even, odd) pairs along D; cos/sin: [N, D/2]. + """ + cos = cos[None, :, None, :] + sin = sin[None, :, None, :] + + def rotate(x): + x = x.float().reshape(*x.shape[:-1], -1, 2) + x_even, x_odd = x[..., 0], x[..., 1] + out = torch.stack([x_even * cos - x_odd * sin, x_even * sin + x_odd * cos], dim=-1) + return out.flatten(3) + + return rotate(xq).type_as(xq), rotate(xk).type_as(xk) + + +# --------------------------------------------------------------------------- +# Building blocks (match pixdit_core/modules.py) +# --------------------------------------------------------------------------- + + +def apply_adaln(x, shift, scale): + return x * (1 + scale) + shift + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class TimestepConditioner(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10): + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[..., None].float() * freqs[None, ...] + return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + hidden_dim = int(2 * hidden_dim / 3) + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class MLP(nn.Module): + def __init__(self, dim: int, mlp_ratio: float = 4.0): + super().__init__() + hidden_dim = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden_dim) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_dim, dim) + + def forward(self, x): + return self.fc2(self.act(self.fc1(x))) + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm = RMSNorm(hidden_size, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x): + return self.linear(self.norm(x)) + + +class PatchTokenEmbedder(nn.Module): + def __init__(self, in_chans: int, embed_dim: int, norm_layer=None, bias: bool = True): + super().__init__() + self.proj = nn.Linear(in_chans, embed_dim, bias=bias) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + return self.norm(self.proj(x)) + + +class PixelTokenEmbedder(nn.Module): + """Per-pixel linear embedding + absolute sin/cos position, grouped into [B*L, P*P, D].""" + + def __init__(self, in_channels: int, hidden_size_output: int): + super().__init__() + self.hidden_size_output = int(hidden_size_output) + self.proj = nn.Linear(in_channels, self.hidden_size_output, bias=True) + + def forward(self, inputs, pos_full, patch_size: int): + B, _, H, W = inputs.shape + Hs, Ws = H // patch_size, W // patch_size + x = self.proj(inputs.permute(0, 2, 3, 1)) + x = x + pos_full.to(x.dtype).view(H, W, self.hidden_size_output).unsqueeze(0) + x = x.view(B, Hs, patch_size, Ws, patch_size, self.hidden_size_output) + x = x.permute(0, 1, 3, 2, 4, 5) + return x.reshape(B * Hs * Ws, patch_size * patch_size, self.hidden_size_output) + + +class RotaryAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = RMSNorm(dim // num_heads) + self.k_norm = RMSNorm(dim // num_heads) + self.proj = nn.Linear(dim, dim) + + def forward(self, x, cos, sin): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + q, k = apply_rotary_emb(self.q_norm(q), self.k_norm(k), cos, sin) + q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + return self.proj(x.transpose(1, 2).reshape(B, N, C)) + + +class PiTBlock(nn.Module): + """Pixel-level block: per-pixel adaLN from patch tokens + attention over compressed patches.""" + + def __init__(self, pixel_hidden_size, patch_hidden_size, patch_size, attn_hidden_size, attn_num_heads, mlp_ratio=4.0): + super().__init__() + self.pixel_dim = int(pixel_hidden_size) + self.attn_dim = int(attn_hidden_size) + p2 = patch_size * patch_size + self.compress_to_attn = nn.Linear(p2 * self.pixel_dim, self.attn_dim, bias=True) + self.expand_from_attn = nn.Linear(self.attn_dim, p2 * self.pixel_dim, bias=True) + self.norm1 = RMSNorm(self.pixel_dim, eps=1e-6) + self.attn = RotaryAttention(self.attn_dim, num_heads=attn_num_heads, qkv_bias=False) + self.norm2 = RMSNorm(self.pixel_dim, eps=1e-6) + self.mlp = MLP(self.pixel_dim, mlp_ratio=mlp_ratio) + self.adaLN_modulation = nn.Sequential(nn.Linear(patch_hidden_size, 6 * self.pixel_dim * p2, bias=True)) + + def forward(self, x, s_cond, batch, num_patches, cos, sin): + BL, P2, _ = x.shape + cond = self.adaLN_modulation(s_cond).view(BL, P2, 6 * self.pixel_dim) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond, 6, dim=-1) + x_norm = apply_adaln(self.norm1(x), shift_msa, scale_msa) + x_comp = self.compress_to_attn(x_norm.view(BL, P2 * self.pixel_dim)).view(batch, num_patches, self.attn_dim) + attn_out = self.attn(x_comp, cos, sin) + attn_exp = self.expand_from_attn(attn_out.view(BL, self.attn_dim)).view(BL, P2, self.pixel_dim) + x = x + gate_msa * attn_exp + return x + gate_mlp * self.mlp(apply_adaln(self.norm2(x), shift_mlp, scale_mlp)) + + +class MMDiTJointAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + assert dim % num_heads == 0, "dim should be divisible by num_heads" + self.num_heads = num_heads + head_dim = dim // num_heads + self.qkv_x = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.qkv_y = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm_x = RMSNorm(head_dim) + self.k_norm_x = RMSNorm(head_dim) + self.q_norm_y = RMSNorm(head_dim) + self.k_norm_y = RMSNorm(head_dim) + self.proj_x = nn.Linear(dim, dim) + self.proj_y = nn.Linear(dim, dim) + + def forward(self, x, y, img_cos, img_sin, txt_cos, txt_sin): + B, Nx, C = x.shape + Ny = y.shape[1] + qkv_x = self.qkv_x(x).reshape(B, Nx, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) + qx, kx, vx = self.q_norm_x(qkv_x[0]), self.k_norm_x(qkv_x[1]), qkv_x[2] + qkv_y = self.qkv_y(y).reshape(B, Ny, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) + qy, ky, vy = self.q_norm_y(qkv_y[0]), self.k_norm_y(qkv_y[1]), qkv_y[2] + + qx, kx = apply_rotary_emb(qx, kx, img_cos, img_sin) + qy, ky = apply_rotary_emb(qy, ky, txt_cos, txt_sin) + + q = torch.cat([qy.transpose(1, 2), qx.transpose(1, 2)], dim=2) + k = torch.cat([ky.transpose(1, 2), kx.transpose(1, 2)], dim=2) + v = torch.cat([vy.transpose(1, 2), vx.transpose(1, 2)], dim=2) + out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + + out_y = out[:, :, :Ny, :].transpose(1, 2).reshape(B, Ny, C) + out_x = out[:, :, Ny:, :].transpose(1, 2).reshape(B, Nx, C) + return self.proj_x(out_x), self.proj_y(out_y) + + +class MMDiTBlockT2I(nn.Module): + def __init__(self, hidden_size, groups, mlp_ratio=4.0): + super().__init__() + self.norm_x1 = RMSNorm(hidden_size, eps=1e-6) + self.norm_y1 = RMSNorm(hidden_size, eps=1e-6) + self.attn = MMDiTJointAttention(hidden_size, num_heads=groups, qkv_bias=False) + self.norm_x2 = RMSNorm(hidden_size, eps=1e-6) + self.norm_y2 = RMSNorm(hidden_size, eps=1e-6) + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.mlp_x = FeedForward(hidden_size, mlp_hidden_dim) + self.mlp_y = FeedForward(hidden_size, mlp_hidden_dim) + self.adaLN_modulation_img = nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) + self.adaLN_modulation_txt = nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) + + def forward(self, x, y, c, img_cos, img_sin, txt_cos, txt_sin): + shift_msa_x, scale_msa_x, gate_msa_x, shift_mlp_x, scale_mlp_x, gate_mlp_x = self.adaLN_modulation_img(c).chunk(6, dim=-1) + shift_msa_y, scale_msa_y, gate_msa_y, shift_mlp_y, scale_mlp_y, gate_mlp_y = self.adaLN_modulation_txt(c).chunk(6, dim=-1) + + x_norm = apply_adaln(self.norm_x1(x), shift_msa_x, scale_msa_x) + y_norm = apply_adaln(self.norm_y1(y), shift_msa_y, scale_msa_y) + attn_x, attn_y = self.attn(x_norm, y_norm, img_cos, img_sin, txt_cos, txt_sin) + x = x + gate_msa_x * attn_x + y = y + gate_msa_y * attn_y + + x = x + gate_mlp_x * self.mlp_x(apply_adaln(self.norm_x2(x), shift_mlp_x, scale_mlp_x)) + y = y + gate_mlp_y * self.mlp_y(apply_adaln(self.norm_y2(y), shift_mlp_y, scale_mlp_y)) + return x, y + + +# --------------------------------------------------------------------------- +# PixelDiT +# --------------------------------------------------------------------------- + + +class PixelDiT(nn.Module): + """PixDiT_T2I with upstream PixDiTTrainer input handling (inputs cast to the weight dtype). + + forward(x [B,3,H,W], timestep [B] (already scaled by 1000), y [B,L,txt_embed_dim]) -> [B,3,H,W] + velocity in the weight dtype. Call set_image_size(H, W) after moving/casting the model and + before forward/export. + """ + + def __init__( + self, + in_channels=3, + patch_size=16, + num_groups=24, + hidden_size=1536, + pixel_hidden_size=16, + pixel_attn_hidden_size=1152, + pixel_num_groups=16, + patch_depth=14, + pixel_depth=2, + txt_embed_dim=2304, + txt_max_length=300, + text_rope_theta=10000.0, + ): + super().__init__() + self.in_channels = int(in_channels) + self.patch_size = int(patch_size) + self.num_groups = int(num_groups) + self.hidden_size = int(hidden_size) + self.pixel_hidden_size = int(pixel_hidden_size) + self.pixel_attn_hidden_size = int(pixel_attn_hidden_size) + self.pixel_num_groups = int(pixel_num_groups) + self.txt_max_length = int(txt_max_length) + self.text_rope_theta = float(text_rope_theta) + + self.pixel_embedder = PixelTokenEmbedder(in_channels, self.pixel_hidden_size) + self.s_embedder = PatchTokenEmbedder(in_channels * patch_size**2, hidden_size, bias=True) + self.t_embedder = TimestepConditioner(hidden_size) + self.y_embedder = PatchTokenEmbedder(txt_embed_dim, hidden_size, bias=True, norm_layer=RMSNorm) + self.y_pos_embedding = nn.Parameter(torch.zeros(1, self.txt_max_length, hidden_size)) + self.patch_blocks = nn.ModuleList([MMDiTBlockT2I(hidden_size, num_groups) for _ in range(patch_depth)]) + self.pixel_blocks = nn.ModuleList( + [ + PiTBlock( + self.pixel_hidden_size, + hidden_size, + patch_size=self.patch_size, + attn_hidden_size=self.pixel_attn_hidden_size, + attn_num_heads=self.pixel_num_groups, + ) + for _ in range(pixel_depth) + ] + ) + self.final_layer = FinalLayer(self.pixel_hidden_size, self.in_channels) + self.image_size = None + + @property + def dtype(self) -> torch.dtype: + return self.final_layer.linear.weight.dtype + + def set_image_size(self, height: int, width: int) -> None: + """Build position tables for a fixed (height, width) on the model's device and dtype. + + Matches upstream: 2D RoPE angles on CPU, text RoPE on device, cos/sin in fp32, and the + float64 pixel sin/cos table cast straight to the weight dtype. + """ + if height % self.patch_size or width % self.patch_size: + raise ValueError(f"Image size {height}x{width} must be a multiple of patch_size={self.patch_size}") + device = self.final_layer.linear.weight.device + hs, ws = height // self.patch_size, width // self.patch_size + img = rope_2d_angles(self.hidden_size // self.num_groups, hs, ws) + pix = rope_2d_angles(self.pixel_attn_hidden_size // self.pixel_num_groups, hs, ws) + txt = rope_1d_angles(self.hidden_size // self.num_groups, self.txt_max_length, self.text_rope_theta, device) + pixel_pos = torch.from_numpy(pixel_sincos_pos_embed(self.pixel_hidden_size, height, width)) + self.img_cos, self.img_sin = img.cos().to(device), img.sin().to(device) + self.pix_cos, self.pix_sin = pix.cos().to(device), pix.sin().to(device) + self.txt_cos, self.txt_sin = txt.cos(), txt.sin() + self.pixel_pos = pixel_pos.to(device=device, dtype=self.dtype) + self.image_size = (height, width) + + def forward(self, x, timestep, y): + dtype = self.dtype + x = x.to(dtype) + timestep = timestep.to(dtype) + y = y.to(dtype) + + B, _, H, W = x.shape + if self.image_size != (H, W): + raise RuntimeError(f"Input is {H}x{W} but set_image_size() was {self.image_size}") + if self.pixel_pos.dtype != dtype or self.pixel_pos.device != x.device: + raise RuntimeError("Model was moved or cast after set_image_size(); call it again") + p = self.patch_size + Hs, Ws = H // p, W // p + L = Hs * Ws + + x_patches = x.view(B, self.in_channels, Hs, p, Ws, p).permute(0, 2, 4, 1, 3, 5).reshape(B, L, -1) + t_emb = self.t_embedder(timestep.view(-1)).view(B, -1, self.hidden_size) + + Ltxt = min(y.shape[1], self.txt_max_length) + y = y[:, :Ltxt, :] + y_emb = self.y_embedder(y).view(B, Ltxt, self.hidden_size) + y_emb = y_emb + self.y_pos_embedding[:, :Ltxt, :].to(y_emb.dtype) + + condition = F.silu(t_emb) + s = self.s_embedder(x_patches) + txt_cos, txt_sin = self.txt_cos[:Ltxt], self.txt_sin[:Ltxt] + for block in self.patch_blocks: + s, y_emb = block(s, y_emb, condition, self.img_cos, self.img_sin, txt_cos, txt_sin) + s = F.silu(t_emb + s) + + s_cond = s.view(B * L, self.hidden_size) + x_pixels = self.pixel_embedder(x, self.pixel_pos, p) + for block in self.pixel_blocks: + x_pixels = block(x_pixels, s_cond, B, L, self.pix_cos, self.pix_sin) + + x_pixels = self.final_layer(x_pixels) + x_pixels = x_pixels.view(B, Hs, Ws, p, p, self.in_channels).permute(0, 5, 1, 3, 2, 4) + return x_pixels.reshape(B, self.in_channels, H, W) diff --git a/image_gen/pixeldit/model_export/pixeldit/pipeline.py b/image_gen/pixeldit/model_export/pixeldit/pipeline.py new file mode 100644 index 0000000..5e24059 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/pipeline.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PixelDiT T2I sampling glue: prompt encoding, CFG, and flow DPM-Solver++ (2M). + +Reproduces upstream t2i/inference.py with sample_steps/cfg_scale/flow_shift and +DPMS(model_type="flow", schedule="FLOW").sample(order=2, skip_type="time_uniform_flow", +method="multistep"). The text encoder and DiT are callables so the same code drives +PyTorch modules and ONNX Runtime sessions: + + encode(input_ids [1,S] int64, attention_mask [1,S] int64, select_index [300] int64) -> [1,300,D] + dit(x [2,3,H,W] fp32, timestep [2] fp32 (t * 1000), y [2,300,D] fp32) -> velocity [2,3,H,W] + +Batch index 0 is the negative (unconditional) prompt, 1 is the positive prompt. +""" + +from dataclasses import dataclass, field +from typing import Callable + +import numpy as np +import torch + +# t2i/configs/PixelDiT_1024px_pixel_diffusion_stage3.yaml text_encoder.chi_prompt +CHI_PROMPT = "\n".join( + [ + 'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:', + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", + ] +) +DEFAULT_NEGATIVE_PROMPT = "low quality, worst quality, over-saturated, blurry, deformed, watermark" +TXT_MAX_LENGTH = 300 +NUM_TRAIN_TIMESTEPS = 1000 # NoiseScheduleFlow.total_N; DiT timestep = t * 1000, t_0 = 1 / 1000 + +EncodeFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] +DiTFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] + + +@dataclass +class SamplerConfig: + steps: int = 50 + cfg_scale: float = 2.75 + flow_shift: float = 4.0 + negative_prompt: str = DEFAULT_NEGATIVE_PROMPT + + +@dataclass +class TextLayout: + """Fixed text-encoder sequence layout; text_seq_len is 506 for the shipped chi_prompt.""" + + chi_prompt: str + chi_prompt_tokens: int + txt_max_length: int = TXT_MAX_LENGTH + text_seq_len: int = field(init=False) + + def __post_init__(self): + # inference.py: num_chi_prompt_tokens + model_max_length - 2 ("magic number 2: [bos], [_]") + self.text_seq_len = self.chi_prompt_tokens + self.txt_max_length - 2 + + @classmethod + def from_tokenizer(cls, tokenizer, chi_prompt: str = CHI_PROMPT, txt_max_length: int = TXT_MAX_LENGTH): + return cls(chi_prompt, len(tokenizer.encode(chi_prompt)), txt_max_length) + + def positive_select_index(self) -> list[int]: + # inference.py: [0] + list(range(-model_max_length + 1, 0)) + return [0] + list(range(self.text_seq_len - self.txt_max_length + 1, self.text_seq_len)) + + def negative_select_index(self) -> list[int]: + # Causal encoder + right padding: the first 300 positions of a 506-token encode equal a 300-token encode. + return list(range(self.txt_max_length)) + + +def tokenize(tokenizer, text: str, max_length: int, device) -> tuple[torch.Tensor, torch.Tensor]: + tok = tokenizer(text, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt") + return tok.input_ids.to(device), tok.attention_mask.to(device) + + +def encode_prompts(tokenizer, encode: EncodeFn, layout: TextLayout, prompt: str, negative_prompt: str, device): + """Return (y_negative, y_positive), each [1, 300, D], using the single 506-token layout.""" + select_neg = torch.tensor(layout.negative_select_index(), dtype=torch.long, device=device) + select_pos = torch.tensor(layout.positive_select_index(), dtype=torch.long, device=device) + y_neg = encode(*tokenize(tokenizer, negative_prompt, layout.text_seq_len, device), select_neg) + y_pos = encode(*tokenize(tokenizer, layout.chi_prompt + prompt.strip(), layout.text_seq_len, device), select_pos) + return y_neg, y_pos + + +def initial_noise(height: int, width: int, seed: int, device="cuda") -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn(1, 3, height, width, device=device, generator=generator) + + +def flow_timesteps(steps: int, flow_shift: float, device) -> torch.Tensor: + """DPM_Solver.get_time_steps(skip_type="time_uniform_flow"); t goes from ~1 down to 0.""" + t_T, t_0 = 1, 1.0 / NUM_TRAIN_TIMESTEPS + betas = torch.linspace(t_T, t_0, steps + 1).to(device) + sigmas = 1.0 - betas + return (flow_shift * sigmas / (1 + (flow_shift - 1) * sigmas)).flip(dims=[0]) + + +# Flow noise schedule (NoiseScheduleFlow): alpha_t = 1 - t, sigma_t = t. +def _log_alpha(t): + return torch.log(1 - t) + + +def _lambda(t): + return _log_alpha(t) - torch.log(t) + + +class FlowDPMSolver: + """DPM-Solver++ multistep (order 2, lower_order_final) with classifier-free guidance.""" + + def __init__(self, dit: DiTFn, y_neg: torch.Tensor, y_pos: torch.Tensor, cfg_scale: float): + self.dit = dit + self.y = torch.cat([y_neg, y_pos]) + self.cfg_scale = cfg_scale + + def data_prediction(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """x0 estimate at time t (model_wrapper classifier-free + DPM_Solver.data_prediction_fn).""" + t_in = t.expand(1) + t_in = torch.cat([t_in] * 2) + x_in = torch.cat([x] * 2) + out = self.dit(x_in, t_in * NUM_TRAIN_TIMESTEPS, self.y) + sigma = t_in.view(-1, 1, 1, 1).to(x) + noise_uncond, noise = ((1 - sigma) * out + x_in).chunk(2) + noise = noise_uncond + self.cfg_scale * (noise - noise_uncond) + return (x - t * noise) / (1 - t) + + @staticmethod + def first_order_update(x, s, t, model_s): + if t == 0: + return model_s # sigma_t = 0, alpha_t = 1, expm1(-inf) = -1: the update is exactly x0 + h = _lambda(t) - _lambda(s) + alpha_t = torch.exp(_log_alpha(t)) + phi_1 = torch.expm1(-h) + return t / s * x - alpha_t * phi_1 * model_s + + @staticmethod + def second_order_update(x, model_prev_list, t_prev_list, t): + model_prev_1, model_prev_0 = model_prev_list[-2], model_prev_list[-1] + t_prev_1, t_prev_0 = t_prev_list[-2], t_prev_list[-1] + lambda_prev_1, lambda_prev_0, lambda_t = _lambda(t_prev_1), _lambda(t_prev_0), _lambda(t) + alpha_t = torch.exp(_log_alpha(t)) + h_0 = lambda_prev_0 - lambda_prev_1 + h = lambda_t - lambda_prev_0 + r0 = h_0 / h + D1_0 = (1.0 / r0) * (model_prev_0 - model_prev_1) + phi_1 = torch.expm1(-h) + return (t / t_prev_0) * x - (alpha_t * phi_1) * model_prev_0 - 0.5 * (alpha_t * phi_1) * D1_0 + + def sample(self, x: torch.Tensor, steps: int, flow_shift: float, progress: Callable[[int], None] | None = None): + order = 2 + assert steps >= order + timesteps = flow_timesteps(steps, flow_shift, x.device) + t_prev_list = [timesteps[0]] + model_prev_list = [self.data_prediction(x, timesteps[0])] + for step in range(1, steps + 1): + t = timesteps[step] + step_order = 1 if step < order else min(order, steps + 1 - step) + if step_order == 1: + x = self.first_order_update(x, t_prev_list[-1], t, model_prev_list[-1]) + else: + x = self.second_order_update(x, model_prev_list, t_prev_list, t) + if len(t_prev_list) < order: + t_prev_list.append(t) + model_prev_list.append(None) + else: + t_prev_list = t_prev_list[1:] + [t] + model_prev_list = model_prev_list[1:] + [None] + if step < steps: # no model evaluation after the final step + model_prev_list[-1] = self.data_prediction(x, t) + if progress: + progress(step) + return x + + +def generate( + tokenizer, + encode: EncodeFn, + dit: DiTFn, + layout: TextLayout, + prompt: str, + height: int, + width: int, + seed: int, + sampler: SamplerConfig, + device="cuda", + progress: Callable[[int], None] | None = None, +) -> torch.Tensor: + """Full text-to-image run; returns the final sample [1,3,H,W] in [-1, 1] (fp32).""" + with torch.inference_mode(): + y_neg, y_pos = encode_prompts(tokenizer, encode, layout, prompt, sampler.negative_prompt, device) + z = initial_noise(height, width, seed, device) + solver = FlowDPMSolver(dit, y_neg, y_pos, sampler.cfg_scale) + return solver.sample(z, sampler.steps, sampler.flow_shift, progress) + + +def to_uint8_image(sample: torch.Tensor) -> np.ndarray: + """torchvision save_image(normalize=True, value_range=(-1, 1)) -> HxWx3 uint8.""" + img = sample[0].float().clamp(-1, 1).add(1).div(2) + img = img.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8) + return img.numpy() diff --git a/image_gen/pixeldit/model_export/pixeldit/weights.py b/image_gen/pixeldit/model_export/pixeldit/weights.py new file mode 100644 index 0000000..2481335 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/weights.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Load the PixelDiT T2I checkpoint and the Gemma-2 text encoder.""" + +import json +from pathlib import Path + +import torch +from huggingface_hub import hf_hub_download + +from .model import PIXELDIT_1300M_CONFIG, PixelDiT + +DEFAULT_MODEL_NAME = "nvidia/PixelDiT-1300M-1024px" +CHECKPOINT_FILE = "pixeldit_t2i_v1.pth" +DEFAULT_TEXT_ENCODER_NAME = "Efficient-Large-Model/gemma-2-2b-it" + +# Upstream PixDiTTrainer wraps PixDiT_T2I as `core`; _repa_projector is a training-only head. +CORE_PREFIX = "core." +TRAINING_ONLY_PREFIXES = ("_repa_projector.",) + + +def resolve_checkpoint(model_name: str, local_files_only: bool = False) -> tuple[Path, dict]: + """Return (checkpoint .pth path, model config) from a HF repo ID, local dir, or local .pth.""" + local = Path(model_name).expanduser() + if local.is_file(): + config_path = local.parent / "config.json" + checkpoint = local + elif local.is_dir(): + config_path = local / "config.json" + checkpoint = local / CHECKPOINT_FILE + else: + checkpoint = Path(hf_hub_download(model_name, CHECKPOINT_FILE, local_files_only=local_files_only)) + config_path = Path(hf_hub_download(model_name, "config.json", local_files_only=local_files_only)) + + config = dict(PIXELDIT_1300M_CONFIG) + if config_path.is_file(): + hf_config = json.loads(config_path.read_text(encoding="utf-8")) + config.update({k: hf_config[k] for k in PIXELDIT_1300M_CONFIG if k in hf_config}) + return checkpoint.resolve(), config + + +def core_state_dict(checkpoint: Path) -> dict[str, torch.Tensor]: + state = torch.load(str(checkpoint), map_location="cpu", mmap=True, weights_only=True) + state = state.get("state_dict", state) + core = {} + for key, value in state.items(): + if key.startswith(CORE_PREFIX): + core[key[len(CORE_PREFIX):]] = value + elif not key.startswith(TRAINING_ONLY_PREFIXES): + raise KeyError(f"Unexpected checkpoint key: {key}") + return core + + +def load_pixeldit( + model_name: str = DEFAULT_MODEL_NAME, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + local_files_only: bool = False, +) -> PixelDiT: + """Build PixelDiT, load weights strictly, cast to dtype (upstream inference uses bf16).""" + checkpoint, config = resolve_checkpoint(model_name, local_files_only) + model = PixelDiT(**config) + model.load_state_dict(core_state_dict(checkpoint), strict=True) + return model.eval().requires_grad_(False).to(device=device, dtype=dtype) + + +def load_gemma_text_encoder( + model_name: str = DEFAULT_TEXT_ENCODER_NAME, + device: str = "cuda", + local_files_only: bool = False, + attn_implementation: str | None = None, +): + """(tokenizer, Gemma-2 decoder stack) exactly as upstream builder.get_tokenizer_and_text_encoder.""" + from transformers import AutoModelForCausalLM, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=local_files_only) + tokenizer.padding_side = "right" + kwargs = {"attn_implementation": attn_implementation} if attn_implementation else {} + text_encoder = ( + AutoModelForCausalLM.from_pretrained( + model_name, dtype=torch.bfloat16, local_files_only=local_files_only, **kwargs + ) + .get_decoder() + .eval() + .requires_grad_(False) + .to(device) + ) + return tokenizer, text_encoder diff --git a/image_gen/pixeldit/model_export/verify_pixeldit.py b/image_gen/pixeldit/model_export/verify_pixeldit.py new file mode 100644 index 0000000..06d215e --- /dev/null +++ b/image_gen/pixeldit/model_export/verify_pixeldit.py @@ -0,0 +1,453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +#!/usr/bin/env python3 +""" +Verify PixelDiT T2I ONNX graphs: run the vendored PyTorch pipeline and the ONNX Runtime pipeline +(TensorRT RTX by default) and compare them. + +Checks: + --validate_only ONNX structure + dummy run of each graph, then exit + (default) also text encoder and DiT outputs, ONNX Runtime vs PyTorch, and an end-to-end image + +Outputs (written inside --output_dir if given, otherwise --onnx_dir): + pytorch_output.png - image from the PyTorch pipeline + ort_output.png - image from the ONNX Runtime-backed pipeline + comparison.png - PyTorch (left) | ONNX Runtime (right) side-by-side + +Pass/fail is decided by the per-model checks. The end-to-end image only fails if clearly broken: +50-step bf16 sampling amplifies tiny per-step differences, so correct runs can differ in fine detail +or even composition depending on the seed. Review comparison.png visually. +""" + +import argparse +import json +import math +import os +import sys +from pathlib import Path + +import numpy as np +import onnx +import torch +import torch.nn.functional as F +from PIL import Image as PILImage + +from pixeldit import pipeline as P +from pixeldit.weights import DEFAULT_MODEL_NAME, DEFAULT_TEXT_ENCODER_NAME, load_gemma_text_encoder, load_pixeldit + +EP_NAME = "nv_tensorrt_rtx" +_REGISTERED_EP_NAME: str | None = None +DEFAULT_PROMPT = ( + "Close-up portrait of a beautiful Baltic model wearing white flower-shaped earrings, emphasis on the earrings, " + "everything is in full focus, pores and skin imperfections are visible, neutral lighting from a large studio " + "softbox, realistic, high resolution, natural beauty, detailed texture, professional studio shooting." +) + +# All comparisons are bf16 vs bf16, with thresholds at ~2x the bf16-vs-fp32 difference measured on the +# 1024x1024 export. PixelDiT casts the timestep to bf16 (999.75 -> 1000), which makes the DiT far +# noisier near t*1000 = 1000 than elsewhere, hence the looser limit there. +DIT_MAX_REL_L2 = {999.75: 0.55, 500.0: 0.025, 50.0: 0.025} # DiT timestep input -> max relative L2 +TEXT_MIN_COS = 0.999 # Gemma-2 bf16 vs fp32 is itself 0.9997-0.9998 +# 50-step images diverge chaotically from tiny per-step bf16 differences: correct runs measured 18-40 dB +# depending on seed (one seed even changes the pose). The image checks only catch broken output. +# PSNR >= 12 dB rejects black (~5 dB) and noise (~9 dB) but not a flat image (~14 dB, these portraits +# have smooth backgrounds), so detail (mean |Laplacian|) must also stay within 0.5-2x of the PyTorch +# image (correct runs measured 0.72-1.01x). +MIN_E2E_PSNR = 12.0 +DETAIL_RATIO_RANGE = (0.5, 2.0) + +_failures: list[str] = [] + + +def configure_stdio() -> None: + if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(encoding="utf-8") + if hasattr(sys.stderr, "reconfigure"): + sys.stderr.reconfigure(encoding="utf-8") + + +def _check(name: str, ok: bool, detail: str) -> None: + print(f" [{'PASS' if ok else 'FAIL'}] {name}: {detail}") + if not ok: + _failures.append(name) + + +def _stats(a: torch.Tensor, b: torch.Tensor) -> dict[str, float]: + a, b = a.detach().double().flatten(), b.detach().double().flatten().to(a.device) + return { + "max_abs": (a - b).abs().max().item(), + "rel_l2": ((a - b).norm() / b.norm().clamp_min(1e-30)).item(), + "cos": F.cosine_similarity(a, b, dim=0).item(), + } + + +def _fmt(s: dict) -> str: + return f"cos={s['cos']:.6f} rel_l2={s['rel_l2']:.3e} max_abs={s['max_abs']:.3e}" + + +def _psnr(a: np.ndarray, b: np.ndarray) -> float: + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse < 1e-12 else 10.0 * math.log10(255.0**2 / mse) + + +# --------------------------------------------------------------------------- +# ONNX Runtime session helpers (same as verify_flux2.py) +# --------------------------------------------------------------------------- + + +def _np_dtype(type_name: str) -> np.dtype: + return { + "tensor(float)": np.float32, + "tensor(float16)": np.float16, + "tensor(int32)": np.int32, + "tensor(int64)": np.int64, + "tensor(bool)": np.bool_, + }[type_name] + + +def _register_trt_rtx(args) -> str: + global _REGISTERED_EP_NAME + if _REGISTERED_EP_NAME is not None: + return _REGISTERED_EP_NAME + + try: + import onnxruntime_ep_nv_tensorrt_rtx as ep # noqa: F401 + + name, lib = ep.get_ep_name(), ep.get_library_path() + print(f" [ORT] Using onnxruntime-ep-nv-tensorrt-rtx ({name})") + except ImportError: + if not args.ep_lib: + raise SystemExit("trt-rtx needs onnxruntime-ep-nv-tensorrt-rtx or --ep_lib") + name, lib = EP_NAME, str(args.ep_lib) + # Windows only; on Linux the loader uses LD_LIBRARY_PATH. + if hasattr(os, "add_dll_directory"): + for d in (args.ep_dll_dir, args.trt_bin): + if d and os.path.isdir(d): + os.add_dll_directory(os.path.abspath(d)) + import onnxruntime as ort + + ort.register_execution_provider_library(name, lib) + _REGISTERED_EP_NAME = name + return name + + +def _make_session_options(provider: str, args): + import onnxruntime as ort + + so = ort.SessionOptions() + if provider != "trt-rtx": + return so, ["CPUExecutionProvider"] + name = _register_trt_rtx(args) + devices = [d for d in ort.get_ep_devices() if d.ep_name == name] + if not devices: + raise SystemExit(f"EP '{name}' registered but no devices were discovered") + so.add_provider_for_devices(devices, {}) + return so, None + + +def _make_session(onnx_path: Path, provider: str, args): + import onnxruntime as ort + + so, providers = _make_session_options(provider, args) + return ort.InferenceSession(str(onnx_path), sess_options=so, providers=providers) + + +class OrtSessionRunner: + def __init__(self, session): + self.session = session + self.input_names = [i.name for i in session.get_inputs()] + self.output_names = [o.name for o in session.get_outputs()] + self.input_dtypes = {i.name: _np_dtype(i.type) for i in session.get_inputs()} + self.output_dtypes = {o.name: _np_dtype(o.type) for o in session.get_outputs()} + + @classmethod + def from_path(cls, onnx_path: Path, provider: str, args): + return cls(_make_session(onnx_path, provider, args)) + + def run(self, feed: dict[str, np.ndarray]) -> list[np.ndarray]: + return self.session.run(None, feed) + + +def _to_numpy(tensor: torch.Tensor, dtype: np.dtype) -> np.ndarray: + tensor = tensor.detach().to("cpu", torch.float32 if dtype == np.float32 else tensor.dtype) + return np.ascontiguousarray(tensor.numpy().astype(dtype, copy=False)) + + +def ort_encode_fn(runner: OrtSessionRunner, device): + def encode(input_ids, attention_mask, select_index): + feed = { + "input_ids": _to_numpy(input_ids, runner.input_dtypes["input_ids"]), + "attention_mask": _to_numpy(attention_mask, runner.input_dtypes["attention_mask"]), + "select_index": _to_numpy(select_index, runner.input_dtypes["select_index"]), + } + return torch.from_numpy(runner.run(feed)[0]).to(device) + + return encode + + +def ort_dit_fn(runner: OrtSessionRunner, device): + def dit(x, timestep, y): + feed = { + "hidden_states": _to_numpy(x, np.float32), + "timestep": _to_numpy(timestep, np.float32), + "encoder_hidden_states": _to_numpy(y, np.float32), + } + return torch.from_numpy(runner.run(feed)[0]).to(device) + + return dit + + +def torch_encode_fn(text_encoder): + def encode(input_ids, attention_mask, select_index): + return text_encoder(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)[0][:, select_index] + + return encode + + +# --------------------------------------------------------------------------- +# Graph structure + dummy run +# --------------------------------------------------------------------------- + +EXPECTED_IO = { + "text_encoder": (["input_ids", "attention_mask", "select_index"], ["prompt_embeds"]), + "transformer": (["hidden_states", "timestep", "encoder_hidden_states"], ["velocity"]), +} + + +def _onnx_io_shapes(model: onnx.ModelProto) -> dict[str, tuple[int, ...]]: + initializer_names = {initializer.name for initializer in model.graph.initializer} + shapes = {} + for value in list(model.graph.input) + list(model.graph.output): + if value.name not in initializer_names: + shapes[value.name] = tuple(int(d.dim_value) for d in value.type.tensor_type.shape.dim) + return shapes + + +def validate_onnx_dir(onnx_dir: Path, provider: str, args, expected_shapes: dict) -> dict[str, OrtSessionRunner]: + """Structure checks + dummy run; returns the ORT sessions so later checks don't rebuild TensorRT engines.""" + print(f"\n[Validate] {onnx_dir}") + runners = {} + for name, (inputs, outputs) in EXPECTED_IO.items(): + onnx_path = onnx_dir / name / "model.onnx" + if not onnx_path.exists(): + continue + print(f" [Validate] {name}") + onnx.checker.check_model(str(onnx_path)) + model = onnx.load(str(onnx_path), load_external_data=False) + shapes = _onnx_io_shapes(model) + io_names = ([i.name for i in model.graph.input], [o.name for o in model.graph.output]) + _check(f"{name} io names", io_names == (inputs, outputs), f"{io_names}") + for io_name, shape in expected_shapes.get(name, {}).items(): + _check(f"{name} {io_name} shape", shapes.get(io_name) == shape, f"{shapes.get(io_name)} (expected {shape})") + has_bf16 = any(t.data_type == onnx.TensorProto.BFLOAT16 for t in model.graph.initializer) + _check(f"{name} bf16 weights", has_bf16, "bf16 initializers present" if has_bf16 else "no bf16 initializers") + data_gb = (onnx_path.parent / "model.onnx_data").stat().st_size / 1e9 + print(f" model.onnx_data: {data_gb:.2f} GB") + + runner = OrtSessionRunner.from_path(onnx_path, provider, args) + float_io = [d for d in list(runner.input_dtypes.values()) + list(runner.output_dtypes.values()) if d in (np.float16, np.float32)] + _check(f"{name} fp32 float IO", all(d == np.float32 for d in float_io), str(float_io)) + feed = {} + for inp in runner.input_names: + dtype = runner.input_dtypes[inp] + shape = shapes[inp] + if dtype == np.float32: + feed[inp] = np.random.default_rng(0).standard_normal(shape).astype(dtype) + elif inp == "attention_mask": + feed[inp] = np.ones(shape, dtype) # an all-zero mask has fully masked attention rows + else: + feed[inp] = np.zeros(shape, dtype) + print(f" input {inp}: shape={shape}, dtype={dtype}") + for out_name, output in zip(runner.output_names, runner.run(feed), strict=True): + finite = bool(np.isfinite(output).all()) + print(f" output {out_name}: shape={tuple(output.shape)}, dtype={output.dtype}, finite={finite}") + _check(f"{name} dummy run output finite", finite, out_name) + runners[name] = runner + if not runners: + raise RuntimeError(f"No model.onnx files found under {onnx_dir}") + return runners + + +# --------------------------------------------------------------------------- +# Per-model ONNX Runtime vs PyTorch +# --------------------------------------------------------------------------- + + +@torch.inference_mode() +def compare_text_encoder(tokenizer, text_encoder, runner, layout, prompt, negative_prompt, device): + print("\n[text_encoder] ONNX Runtime vs PyTorch (bf16)") + torch_neg, torch_pos = P.encode_prompts(tokenizer, torch_encode_fn(text_encoder), layout, prompt, negative_prompt, device) + ort_neg, ort_pos = P.encode_prompts(tokenizer, ort_encode_fn(runner, device), layout, prompt, negative_prompt, device) + for label, ort, ref in (("positive", ort_pos, torch_pos), ("negative", ort_neg, torch_neg)): + s = _stats(ort, ref) + per_token = F.cosine_similarity(ort[0].double(), ref[0].double(), dim=-1).min().item() + _check(f"text_encoder {label} prompt_embeds", s["cos"] >= TEXT_MIN_COS, + f"{_fmt(s)} min_token_cos={per_token:.6f} (need cos >= {TEXT_MIN_COS})") + + +@torch.inference_mode() +def compare_dit(dit, runner, y_neg, y_pos, height, width, device): + print("\n[transformer] ONNX Runtime vs PyTorch (bf16)") + g = torch.Generator(device=device).manual_seed(1) + x = torch.randn(1, 3, height, width, device=device, generator=g).repeat(2, 1, 1, 1) + y = torch.cat([y_neg, y_pos]).float() + ort_dit = ort_dit_fn(runner, device) + for tv, limit in DIT_MAX_REL_L2.items(): + t = torch.full((2,), tv, device=device) + s = _stats(ort_dit(x, t, y), dit(x, t, y)) + _check(f"transformer velocity t={tv}", s["rel_l2"] <= limit, f"{_fmt(s)} (need rel_l2 <= {limit})") + + +# --------------------------------------------------------------------------- +# End to end +# --------------------------------------------------------------------------- + + +def _detail(image: np.ndarray) -> float: + """High-frequency energy: mean |Laplacian| of luminance.""" + g = image.astype(np.float64).mean(-1) + return float(np.abs(4 * g[1:-1, 1:-1] - g[:-2, 1:-1] - g[2:, 1:-1] - g[1:-1, :-2] - g[1:-1, 2:]).mean()) + + +def compare_images(ref: np.ndarray, ort: np.ndarray, output_dir: Path) -> None: + """Fail only on clearly broken output (black, noise, flat); otherwise report metrics for visual review.""" + psnr = _psnr(ref, ort) + mae = float(np.mean(np.abs(ref.astype(np.float32) - ort.astype(np.float32)))) + ratio = _detail(ort) / max(_detail(ref), 1e-6) + lo, hi = DETAIL_RATIO_RANGE + _check("end-to-end image not broken: PSNR", psnr >= MIN_E2E_PSNR, f"PSNR={psnr:.2f} dB (need >= {MIN_E2E_PSNR} dB)") + _check("end-to-end image not broken: detail", lo <= ratio <= hi, f"detail ratio={ratio:.2f} (need {lo}-{hi})") + print(f" [INFO] end-to-end image: PSNR={psnr:.2f} dB MAE={mae:.2f}. Fine details can differ between correct " + "bf16 runs (18-40 dB measured across seeds); compare the images visually.") + canvas = PILImage.new("RGB", (ref.shape[1] + ort.shape[1], max(ref.shape[0], ort.shape[0])), (255, 255, 255)) + canvas.paste(PILImage.fromarray(ref), (0, 0)) + canvas.paste(PILImage.fromarray(ort), (ref.shape[1], 0)) + cmp_path = output_dir / "comparison.png" + canvas.save(str(cmp_path)) + print(f"[Comparison] Side-by-side saved -> {cmp_path}") + + +def _progress(label): + def report(step): + if step % 10 == 0: + print(f" [{label}] step {step}") + + return report + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def _image_size(onnx_dir: Path) -> tuple[int, int]: + config_path = onnx_dir / "pipeline_config.json" + if config_path.exists(): + config = json.loads(config_path.read_text(encoding="utf-8")) + return int(config["height"]), int(config["width"]) + transformer = onnx_dir / "transformer" / "model.onnx" + shape = _onnx_io_shapes(onnx.load(str(transformer), load_external_data=False))["hidden_states"] + return shape[2], shape[3] + + +def main(): + p = argparse.ArgumentParser(description="Verify PixelDiT T2I ONNX graphs with PyTorch and ONNX Runtime.") + p.add_argument("--model_name", type=str, default=DEFAULT_MODEL_NAME, help="Hugging Face model repo ID, local dir, or .pth") + p.add_argument("--text_encoder_name", type=str, default=DEFAULT_TEXT_ENCODER_NAME) + p.add_argument("--local_files_only", action="store_true", help="Resolve models from the local Hugging Face cache only") + p.add_argument("--onnx_dir", type=str, default="./pixeldit_onnx") + p.add_argument("--output_dir", type=str, default=None) + p.add_argument("--prompt", type=str, default=DEFAULT_PROMPT) + p.add_argument("--negative_prompt", type=str, default=P.DEFAULT_NEGATIVE_PROMPT) + p.add_argument("--seed", type=int, default=2025) + p.add_argument("--steps", type=int, default=50) + p.add_argument("--cfg_scale", type=float, default=2.75) + p.add_argument("--flow_shift", type=float, default=4.0) + p.add_argument("--provider", choices=["cpu", "trt-rtx"], default="trt-rtx") + p.add_argument("--ep_lib", default=None, help="Path to onnxruntime_providers_nv_tensorrt_rtx.dll/.so") + 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 structure checks and a dummy run") + p.add_argument("--skip_e2e", action="store_true", help="Skip the end-to-end image comparison") + args = p.parse_args() + configure_stdio() + + onnx_dir = Path(args.onnx_dir).resolve() + output_dir = Path(args.output_dir).resolve() if args.output_dir else onnx_dir + output_dir.mkdir(parents=True, exist_ok=True) + device = "cuda" + height, width = _image_size(onnx_dir) + print(f"Image size from export: {height}x{width}") + + from transformers import AutoConfig, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, local_files_only=args.local_files_only) + hidden = AutoConfig.from_pretrained(args.text_encoder_name, local_files_only=args.local_files_only).hidden_size + layout = P.TextLayout.from_tokenizer(tokenizer) + expected_shapes = { + "text_encoder": { + "input_ids": (1, layout.text_seq_len), + "attention_mask": (1, layout.text_seq_len), + "select_index": (layout.txt_max_length,), + "prompt_embeds": (1, layout.txt_max_length, hidden), + }, + "transformer": { + "hidden_states": (2, 3, height, width), + "timestep": (2,), + "encoder_hidden_states": (2, layout.txt_max_length, hidden), + "velocity": (2, 3, height, width), + }, + } + + runners = validate_onnx_dir(onnx_dir, args.provider, args, expected_shapes) + if args.validate_only: + _finish() + return + + print(f"Loading tokenizer + Gemma from {args.text_encoder_name}...") + tokenizer, text_encoder = load_gemma_text_encoder(args.text_encoder_name, device, args.local_files_only) + + print(f"Loading PixelDiT from {args.model_name} (bf16)...") + dit = load_pixeldit(args.model_name, device, torch.bfloat16, args.local_files_only) + dit.set_image_size(height, width) + + encode_torch = torch_encode_fn(text_encoder) + te_runner, tr_runner = runners.get("text_encoder"), runners.get("transformer") + if te_runner is not None: + compare_text_encoder(tokenizer, text_encoder, te_runner, layout, args.prompt, args.negative_prompt, device) + if tr_runner is not None: + with torch.inference_mode(): + y_neg, y_pos = P.encode_prompts(tokenizer, encode_torch, layout, args.prompt, args.negative_prompt, device) + compare_dit(dit, tr_runner, y_neg, y_pos, height, width, device) + + if not args.skip_e2e: + print("\n[End to end]") + sampler = P.SamplerConfig(args.steps, args.cfg_scale, args.flow_shift, args.negative_prompt) + print("[PyTorch] Running pipeline...") + pt_image = P.to_uint8_image(P.generate(tokenizer, encode_torch, dit, layout, args.prompt, height, width, + args.seed, sampler, device, _progress("PyTorch"))) + PILImage.fromarray(pt_image).save(str(output_dir / "pytorch_output.png")) + print(f"[PyTorch] Image saved -> {output_dir / 'pytorch_output.png'}") + + print("[ONNX Runtime] Running pipeline...") + encode_ort = ort_encode_fn(te_runner, device) if te_runner is not None else encode_torch + dit_ort = ort_dit_fn(tr_runner, device) if tr_runner is not None else dit + ort_image = P.to_uint8_image(P.generate(tokenizer, encode_ort, dit_ort, layout, args.prompt, height, width, + args.seed, sampler, device, _progress("ORT"))) + PILImage.fromarray(ort_image).save(str(output_dir / "ort_output.png")) + print(f"[ONNX Runtime] Image saved -> {output_dir / 'ort_output.png'}") + compare_images(pt_image, ort_image, output_dir) + + _finish() + + +def _finish() -> None: + if _failures: + print(f"\nFAILED checks ({len(_failures)}): {', '.join(_failures)}") + sys.exit(1) + print("\nAll checks passed.") + + +if __name__ == "__main__": + main() diff --git a/image_gen/pixeldit/pixeldit_cli.cpp b/image_gen/pixeldit/pixeldit_cli.cpp new file mode 100644 index 0000000..c3169da --- /dev/null +++ b/image_gen/pixeldit/pixeldit_cli.cpp @@ -0,0 +1,289 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "pixeldit_cli.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "io/image.h" +#include "pixeldit.h" +#include + +namespace din::image_gen +{ +namespace +{ + +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()) + { + throw std::invalid_argument(std::string(flag_name) + " must be an unsigned integer"); + } + return static_cast(parsed); +} + +float parse_float(const std::string& value, const char* flag_name) +{ + size_t parsed_chars = 0; + const float parsed = std::stof(value, &parsed_chars); + if (parsed_chars != value.size()) + { + throw std::invalid_argument(std::string(flag_name) + " must be a number"); + } + return parsed; +} + +std::string to_lower(std::string value) +{ + for (char& c : value) + { + c = static_cast(std::tolower(static_cast(c))); + } + return value; +} + +std::string to_string(PixelDiTProcessingBackend backend) +{ + switch (backend) + { + case PixelDiTProcessingBackend::Cpu: + return "cpu"; + case PixelDiTProcessingBackend::Cuda: + return "cuda"; + } + return "unknown"; +} + +std::string to_string(PixelDiTExecutionProvider provider) +{ + switch (provider) + { + case PixelDiTExecutionProvider::TrtRtx: + return "trt-rtx"; + } + return "unknown"; +} + +PixelDiTProcessingBackend parse_processing_backend(const std::string& value) +{ + const std::string lower = to_lower(value); + if (lower == "cpu") + { + return PixelDiTProcessingBackend::Cpu; + } + if (lower == "cuda") + { + return PixelDiTProcessingBackend::Cuda; + } + if (lower == "dx" || lower == "dx-cig" || lower == "vk") + { + throw std::invalid_argument("--processing " + lower + " is not implemented for PixelDiT yet; use cpu or cuda"); + } + throw std::invalid_argument("Processing must be one of: cpu, cuda"); +} + +PixelDiTExecutionProvider parse_execution_provider(const std::string& value) +{ + const std::string lower = to_lower(value); + if (lower == "trt-rtx" || lower == "trt") + { + return PixelDiTExecutionProvider::TrtRtx; + } + if (lower == "cpu") + { + throw std::invalid_argument("--provider cpu is not supported for PixelDiT; use trt-rtx"); + } + throw std::invalid_argument("Provider must be: trt-rtx"); +} + +PixelDiTConfig parse_args(int argc, char* argv[]) +{ + PixelDiTConfig config; + config.model_dir = DEFAULT_PIXELDIT_MODEL_DIR; + config.prompt = DEFAULT_PIXELDIT_PROMPT; + + argparse::ArgumentParser parser("din_pixeldit"); + parser.add_description("Run PixelDiT text-to-image generation with a selectable processing backend."); + parser.add_argument("--processing") + .default_value(to_string(config.processing)) + .nargs(1) + .metavar("cpu|cuda") + .help("Select the processing backend."); + parser.add_argument("--provider") + .default_value(to_string(config.provider)) + .nargs(1) + .metavar("trt-rtx") + .help("Select the ONNX Runtime execution provider."); + parser.add_argument("--model-dir") + .default_value(config.model_dir.string()) + .nargs(1) + .metavar("PATH") + .help("Directory with exported PixelDiT ONNX artifacts (text_encoder, transformer, tokenizer, " + "pipeline_config.json)."); + parser.add_argument("--ep-cache") + .default_value(config.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()) + .nargs(1) + .metavar("PATH") + .help("Directory for ONNX Runtime EP-context models."); + parser.add_argument("--output") + .default_value(std::string{"."}) + .nargs(1) + .metavar("DIR") + .help("Directory where generated images are written."); + parser.add_argument("--prompt").default_value(config.prompt).nargs(1).metavar("TEXT").help("Prompt text."); + parser.add_argument("--negative-prompt") + .nargs(1) + .metavar("TEXT") + .help("Negative prompt for classifier-free guidance (default: from pipeline_config.json)."); + parser.add_argument("--seed").default_value(std::to_string(config.seed)).nargs(1).metavar("N").help("Random seed."); + parser.add_argument("--num-images") + .default_value(std::to_string(config.num_images)) + .nargs(1) + .metavar("N") + .help("Number of images to generate (seeds seed, seed+1, ...)."); + parser.add_argument("--steps").nargs(1).metavar("N").help("Sampling steps (default: from pipeline_config.json)."); + parser.add_argument("--cfg-scale").nargs(1).metavar("X").help("CFG scale (default: from pipeline_config.json)."); + parser.add_argument("--flow-shift").nargs(1).metavar("X").help("Flow shift (default: from pipeline_config.json)."); + parser.add_argument("--init-noise") + .nargs(1) + .metavar("FILE") + .help("Validation: raw fp32 [1, 3, H, W] initial noise used instead of the seeded generator."); + parser.add_argument("--dump-dir") + .nargs(1) + .metavar("DIR") + .help("Validation: write raw tokens, prompt embeddings, timesteps, noise and final sample here."); + + try + { + parser.parse_args(argc, argv); + } + catch (const std::exception& exception) + { + std::cerr << parser << '\n'; + throw std::runtime_error(exception.what()); + } + + config.processing = parse_processing_backend(parser.get("--processing")); + config.provider = parse_execution_provider(parser.get("--provider")); + config.model_dir = parser.get("--model-dir"); + config.ep_cache_dir = parser.get("--ep-cache"); + config.ep_context_dir = parser.get("--ep-context-dir"); + config.prompt = parser.get("--prompt"); + config.seed = parse_uint(parser.get("--seed"), "--seed"); + config.num_images = parse_uint(parser.get("--num-images"), "--num-images"); + if (auto value = parser.present("--negative-prompt")) + { + config.negative_prompt = *value; + } + if (auto value = parser.present("--steps")) + { + config.steps = static_cast(parse_uint(*value, "--steps")); + } + if (auto value = parser.present("--cfg-scale")) + { + config.cfg_scale = parse_float(*value, "--cfg-scale"); + } + if (auto value = parser.present("--flow-shift")) + { + config.flow_shift = parse_float(*value, "--flow-shift"); + } + if (auto value = parser.present("--init-noise")) + { + config.init_noise_path = *value; + } + if (auto value = parser.present("--dump-dir")) + { + config.dump_dir = *value; + } + + const std::filesystem::path output_dir = parser.get("--output"); + std::filesystem::create_directories(output_dir); + config.output_path = output_dir / "pixeldit.png"; + return config; +} + +std::filesystem::path make_output_path(const PixelDiTConfig& config, unsigned int image_index) +{ + const std::string stem = config.output_path.stem().string(); + const std::string extension = config.output_path.extension().string(); + return config.output_path.parent_path() / (stem + "_" + std::to_string(image_index) + extension); +} + +void save_image(const std::filesystem::path& output_path, const PixelDiTImage& image) +{ + if (image.data.empty()) + { + throw std::runtime_error("Backend returned an empty image"); + } + if (!din::io::SaveRgbFloatImage(output_path.string(), image.data.data(), image.height, image.width, + din::io::ImageValueRange::MinusOneToOne)) + { + throw std::runtime_error("Failed to save image to " + output_path.string()); + } + std::cout << "Image saved to " << output_path.string() << std::endl; +} + +} // namespace +} // namespace din::image_gen + +int main(int argc, char* argv[]) +{ + using namespace din::image_gen; + std::cout << std::unitbuf; + std::cerr << std::unitbuf; + + try + { + const PixelDiTConfig config = parse_args(argc, argv); + std::cout << "Model dir: " << config.model_dir.string() << "\n" + << "Output: " << config.output_path.string() << "\n" + << "Prompt: " << config.prompt << "\n" + << "Seed: " << config.seed << "\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" + << std::endl; + + auto pipeline = CreatePixelDiTPipeline(config); + pipeline->Initialize(); + + for (unsigned int image_index = 0; image_index < config.num_images; ++image_index) + { + const unsigned int current_seed = config.seed + image_index; + std::cout << "\n========== Image " << (image_index + 1) << "/" << config.num_images + << " (seed=" << current_seed << ") ==========" << std::endl; + const auto image_start = std::chrono::steady_clock::now(); + PixelDiTImage image = pipeline->GenerateImage(current_seed); + 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; + std::cout << "Image " << (image_index + 1) << " completed in " << image_duration.count() << " seconds" + << std::endl; + } + + return EXIT_SUCCESS; + } + catch (const std::exception& e) + { + std::cerr << "Error: " << e.what() << std::endl; + } + + return EXIT_FAILURE; +} diff --git a/image_gen/pixeldit/pixeldit_cli.h b/image_gen/pixeldit/pixeldit_cli.h new file mode 100644 index 0000000..e9323c6 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_cli.h @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include +#include +#include + +namespace din::image_gen +{ + +enum class PixelDiTProcessingBackend +{ + Cpu, + Cuda, +}; + +enum class PixelDiTExecutionProvider +{ + TrtRtx, +}; + +struct PixelDiTConfig +{ + PixelDiTProcessingBackend processing = PixelDiTProcessingBackend::Cpu; + PixelDiTExecutionProvider provider = PixelDiTExecutionProvider::TrtRtx; + std::filesystem::path model_dir; + std::filesystem::path ep_cache_dir = "artifacts/pixeldit/trt_rtx_cache"; + std::filesystem::path ep_context_dir = "artifacts/pixeldit/ep_context"; + std::filesystem::path output_path; + std::string prompt; + // Unset sampler options fall back to pipeline_config.json, then to built-in defaults. + std::optional negative_prompt; + std::optional steps; + std::optional cfg_scale; + std::optional flow_shift; + unsigned int seed = 42; + unsigned int num_images = 1; + // Debug / validation: raw fp32 [1, 3, H, W] initial noise instead of the seeded generator, and a + // directory for raw dumps of intermediate tensors. + std::filesystem::path init_noise_path; + std::filesystem::path dump_dir; +}; + +struct PixelDiTImage +{ + std::vector data; // [3, H, W] in [-1, 1] + int height = 0; + int width = 0; +}; + +class PixelDiTProcessingPipeline +{ +public: + virtual ~PixelDiTProcessingPipeline() = default; + virtual void Initialize() = 0; + virtual void SetPrompt(std::string prompt) = 0; + virtual PixelDiTImage GenerateImage(unsigned int seed) = 0; +}; + +std::unique_ptr CreatePixelDiTPipeline(const PixelDiTConfig& config); + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/pixeldit_config.cpp b/image_gen/pixeldit/pixeldit_config.cpp new file mode 100644 index 0000000..ed68196 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_config.cpp @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include "pixeldit.h" +#include + +namespace din::image_gen +{ +namespace +{ + +// Reads object[key] into value when present, otherwise records `label` as a default that was used. +template +void ReadField(const nlohmann::json& object, const char* key, T& value, const std::string& label, + PixelDiTPipelineConfig& config) +{ + if (object.is_object() && object.contains(key) && !object[key].is_null()) + { + value = object[key].get(); + } + else + { + config.defaults_used.push_back(label); + } +} + +} // namespace + +PixelDiTPipelineConfig LoadPixelDiTPipelineConfig(const std::filesystem::path& path) +{ + PixelDiTPipelineConfig config; + nlohmann::json root = nlohmann::json::object(); + std::ifstream stream(path); + if (stream) + { + root = nlohmann::json::parse(stream); + config.file_found = true; + } + + const nlohmann::json empty = nlohmann::json::object(); + const nlohmann::json& text = root.contains("text") ? root["text"] : empty; + const nlohmann::json& sampler = root.contains("sampler") ? root["sampler"] : empty; + + ReadField(root, "height", config.height, "height", config); + ReadField(root, "width", config.width, "width", config); + ReadField(root, "dit_batch", config.dit_batch, "dit_batch", config); + ReadField(root, "timestep_scale", config.timestep_scale, "timestep_scale", config); + + ReadField(text, "chi_prompt", config.chi_prompt, "text.chi_prompt", config); + ReadField(text, "text_seq_len", config.text_seq_len, "text.text_seq_len", config); + ReadField(text, "txt_max_length", config.txt_max_length, "text.txt_max_length", config); + ReadField(text, "bos_token_id", config.bos_token_id, "text.bos_token_id", config); + ReadField(text, "pad_token_id", config.pad_token_id, "text.pad_token_id", config); + ReadField(text, "positive_select_index", config.positive_select_index, "text.positive_select_index", config); + ReadField(text, "negative_select_index", config.negative_select_index, "text.negative_select_index", config); + + ReadField(sampler, "steps", config.steps, "sampler.steps", config); + ReadField(sampler, "cfg_scale", config.cfg_scale, "sampler.cfg_scale", config); + ReadField(sampler, "flow_shift", config.flow_shift, "sampler.flow_shift", config); + ReadField(sampler, "negative_prompt", config.negative_prompt, "sampler.negative_prompt", config); + ReadField(sampler, "t_start", config.t_start, "sampler.t_start", config); + ReadField(sampler, "t_end", config.t_end, "sampler.t_end", config); + ReadField(sampler, "order", config.order, "sampler.order", config); + + if (config.positive_select_index.empty()) + { + config.positive_select_index = DefaultPositiveSelectIndex(config.text_seq_len, config.txt_max_length); + } + if (config.negative_select_index.empty()) + { + config.negative_select_index = DefaultNegativeSelectIndex(config.txt_max_length); + } + + const auto check_select = [&](const std::vector& index, const char* name) + { + if (static_cast(index.size()) != config.txt_max_length) + { + throw std::runtime_error(std::string(name) + " has " + std::to_string(index.size()) + + " entries, expected txt_max_length=" + std::to_string(config.txt_max_length)); + } + for (const int64_t i : index) + { + if (i < 0 || i >= config.text_seq_len) + { + throw std::runtime_error(std::string(name) + " entry " + std::to_string(i) + " is outside [0, " + + std::to_string(config.text_seq_len) + ")"); + } + } + }; + check_select(config.positive_select_index, "positive_select_index"); + check_select(config.negative_select_index, "negative_select_index"); + if (config.dit_batch != 2) + { + throw std::runtime_error("Only dit_batch=2 ([negative, positive]) is supported, got " + + std::to_string(config.dit_batch)); + } + if (config.order != 2) + { + throw std::runtime_error("Only the 2nd-order DPM-Solver++ is supported, got order=" + + std::to_string(config.order)); + } + return config; +} + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp new file mode 100644 index 0000000..3976db2 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -0,0 +1,820 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// PixelDiT pipeline on TensorRT RTX (GPU). The text encoder and DiT always run in TensorRT RTX; --processing selects +// where the sampling math (CFG + flow DPM-Solver++ update) runs. "cpu": on the host, with the velocity downloaded and +// the sample uploaded around every DiT call. "cuda": kernels (cuda_kernels.cu) on the device buffers, all on the +// TensorRT RTX compute stream, with one host sync per image (like image_gen/flux2/flux2_cuda_backend.cpp). + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cuda_kernels.h" +#include "nvtx_helper.h" +#include "ort_session.h" +#include "pixeldit.h" +#include "pixeldit_cli.h" +#include +#include + +namespace din::image_gen +{ +namespace +{ + +#define PIXELDIT_CUDA_CHECK(call) \ + do \ + { \ + const cudaError_t _err = (call); \ + if (_err != cudaSuccess) \ + { \ + throw std::runtime_error(std::string("CUDA error: ") + cudaGetErrorString(_err) + " at " #call); \ + } \ + } while (0) + +enum class SamplingBackend +{ + Cpu, + Cuda, +}; + +const char* to_string(SamplingBackend backend) +{ + switch (backend) + { + case SamplingBackend::Cpu: + return "cpu"; + case SamplingBackend::Cuda: + return "cuda"; + } + return "unknown"; +} + +using FloatBuffer = din::common::TensorBuffer; +using Int64Buffer = din::common::TensorBuffer; + +template +size_t element_count(din::common::TensorBuffer& buffer) +{ + return buffer.BindingValue().GetTensorTypeAndShapeInfo().GetElementCount(); +} + +// Pointer to the tensor a device-backed buffer binds to the session. That is its device allocation, except when the +// runner uses unified memory (TensorBuffer then binds its host-accessible tensor and has no separate device copy). +float* device_data(FloatBuffer& buffer) +{ + return buffer.BindingValue().GetTensorMutableData(); +} + +Ort::Value make_view(FloatBuffer& buffer, size_t offset, const std::vector& shape) +{ + size_t count = 1; + for (const int64_t dim : shape) + { + count *= static_cast(dim); + } + if (offset + count > element_count(buffer)) + { + throw std::logic_error("view is outside the tensor"); + } + const Ort::ConstMemoryInfo memory_info = buffer.BindingValue().GetTensorMemoryInfo(); + return Ort::Value::CreateTensor(memory_info, device_data(buffer) + offset, count, shape.data(), + shape.size()); +} + +template +void write_raw(const std::filesystem::path& path, const T* data, size_t count) +{ + std::ofstream out(path, std::ios::binary); + if (!out) + { + throw std::runtime_error("Failed to open dump file " + path.string()); + } + out.write(reinterpret_cast(data), static_cast(count * sizeof(T))); +} + +std::vector read_raw_floats(const std::filesystem::path& path, size_t expected_count) +{ + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) + { + throw std::runtime_error("Failed to open --init-noise file " + path.string()); + } + const auto bytes = static_cast(in.tellg()); + if (bytes != expected_count * sizeof(float)) + { + throw std::runtime_error("--init-noise file " + path.string() + " has " + std::to_string(bytes) + + " bytes, expected " + std::to_string(expected_count * sizeof(float)) + + " (fp32 [1, 3, H, W])"); + } + std::vector values(expected_count); + in.seekg(0); + in.read(reinterpret_cast(values.data()), static_cast(bytes)); + return values; +} + +std::vector session_shape(Ort::Session& session, const char* name, bool input) +{ + const size_t count = input ? session.GetInputCount() : session.GetOutputCount(); + Ort::AllocatorWithDefaultOptions allocator; + for (size_t i = 0; i < count; ++i) + { + const auto io_name = + input ? session.GetInputNameAllocated(i, allocator) : session.GetOutputNameAllocated(i, allocator); + if (std::string_view{io_name.get()} == name) + { + const auto type_info = input ? session.GetInputTypeInfo(i) : session.GetOutputTypeInfo(i); + // Keep the shape info in a named variable: calling GetShape() on the temporary has lifetime issues on + // Linux (https://github.com/microsoft/onnxruntime/issues/24300). + const auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + return tensor_info.GetShape(); + } + } + throw std::runtime_error(std::string("ONNX model has no ") + (input ? "input " : "output ") + name); +} + +std::string shape_string(const std::vector& shape) +{ + std::ostringstream text; + text << "["; + for (size_t i = 0; i < shape.size(); ++i) + { + text << (i ? ", " : "") << shape[i]; + } + text << "]"; + return text.str(); +} + +// CFG + data prediction: velocity [2, n] ([negative, positive]) and the current sample -> x0 estimate. +class CfgDataPredictionStage +{ +public: + void BindInput(const char* name, FloatBuffer& tensor) + { + if (std::string_view{name} == "velocity") + { + velocity_ = &tensor; + } + else if (std::string_view{name} == "sample") + { + sample_ = &tensor; + } + else + { + throw std::runtime_error("Unsupported CfgDataPredictionStage input: " + std::string{name}); + } + } + + // cpu: the velocity must already be on the host. cuda: everything stays on the device, ordered on `stream`. + void Run(cudaStream_t stream, SamplingBackend backend, float t, float cfg_scale, FloatBuffer& x0) + { + if (velocity_ == nullptr || sample_ == nullptr) + { + throw std::runtime_error("CfgDataPredictionStage bindings are incomplete"); + } + const size_t n = element_count(*sample_); + if (backend == SamplingBackend::Cuda) + { + launch_pixeldit_cfg_x0_kernel(stream, device_data(*velocity_), device_data(*sample_), t, cfg_scale, + device_data(x0), n); + return; + } + cfg_data_prediction(velocity_->HostData(), sample_->HostData(), t, cfg_scale, x0.HostData(), n); + } + +private: + FloatBuffer* velocity_ = nullptr; + FloatBuffer* sample_ = nullptr; +}; + +// Flow DPM-Solver++ update of the sample from t_prev_0 to t (1st or 2nd order). +class DpmSolverStage +{ +public: + void BindOutput(const char* name, FloatBuffer& tensor) + { + if (std::string_view{name} != "sample") + { + throw std::runtime_error("Unsupported DpmSolverStage output: " + std::string{name}); + } + sample_ = &tensor; + } + + void Run(Ort::Env& env, Ort::SyncStream& sync_stream, SamplingBackend backend, int order, FloatBuffer* model_prev_1, + FloatBuffer& model_prev_0, float t_prev_1, float t_prev_0, float t) + { + if (sample_ == nullptr) + { + throw std::runtime_error("DpmSolverStage bindings are incomplete"); + } + const size_t n = element_count(*sample_); + if (backend == SamplingBackend::Cuda) + { + const auto stream = reinterpret_cast(sync_stream.GetHandle()); + if (order == 1 && t == 0.0f) + { + Ort::ThrowOnError(env.CopyTensor(model_prev_0.BindingValue(), sample_->BindingValue(), sync_stream)); + } + else if (order == 1) + { + const DpmFirstOrderCoeffs k = dpm_first_order_coeffs(t_prev_0, t); + launch_pixeldit_dpm1_kernel(stream, device_data(*sample_), device_data(model_prev_0), k.ratio, k.coeff, + n); + } + else + { + const DpmSecondOrderCoeffs k = dpm_second_order_coeffs(t_prev_1, t_prev_0, t); + launch_pixeldit_dpm2_kernel(stream, device_data(*sample_), device_data(*model_prev_1), + device_data(model_prev_0), k.ratio, k.coeff, k.half_coeff, k.inv_r0, n); + } + return; + } + if (order == 1) + { + dpm_first_order_update(sample_->HostData(), model_prev_0.HostData(), t_prev_0, t, n); + } + else + { + dpm_second_order_update(sample_->HostData(), model_prev_1->HostData(), model_prev_0.HostData(), t_prev_1, + t_prev_0, t, n); + } + } + +private: + FloatBuffer* sample_ = nullptr; +}; + +struct SamplerSettings +{ + int steps = 0; + float cfg_scale = 0.0f; + float flow_shift = 0.0f; + std::string negative_prompt; +}; + +// Per-image wall time of each sampling stage with --processing cpu (host view; GPU work is inside `dit`). With +// --processing cuda the host only enqueues work, so only the total sampling time is reported. +struct StageTiming +{ + std::chrono::duration upload{}; // host packing + host->device copies + std::chrono::duration dit{}; // DiT run + velocity device->host copy + std::chrono::duration cfg{}; // host CFG + x0 + std::chrono::duration solver{}; // host DPM-Solver++ update + int evaluations = 0; +}; + +struct PipelineState +{ + explicit PipelineState(Ort::Env& environment) + : env(environment) + { + } + + Ort::Env& env; + PixelDiTModelPaths model_paths; + PixelDiTPipelineConfig pipeline_config; + SamplerSettings sampler; + SamplingBackend sampling_backend = SamplingBackend::Cpu; + int64_t text_hidden = 0; + std::unique_ptr tokenizer; + bool prompt_embeds_valid = false; + std::string prompt; + + std::optional compute_stream; + std::unique_ptr text_encoder_runner; + std::unique_ptr transformer_runner; + + // Text encoder IO (device-backed with pinned host mirrors). + std::unique_ptr input_ids; + std::unique_ptr attention_mask; + std::unique_ptr select_index; + std::unique_ptr prompt_embeds; + // DiT IO. + std::unique_ptr dit_hidden_states; + std::unique_ptr dit_timestep; + std::unique_ptr dit_encoder_hidden_states; + std::unique_ptr dit_velocity; + // Sampling state: current sample and the two most recent x0 estimates. Host-only with --processing cpu, + // device-backed with --processing cuda. + std::unique_ptr sample; + std::unique_ptr model_a; + std::unique_ptr model_b; + // --processing cuda only: DiT timestep for every model evaluation, [steps, dit_batch], uploaded once per image + // and copied into dit_timestep on the device (rewriting the pinned dit_timestep mirror each step would race the + // still-queued copies of earlier steps). + std::unique_ptr timestep_table; + // --processing cuda only: views for the partial device copies, created once (the buffers never move). + // hidden_state_halves[b] is batch row b of dit_hidden_states; timestep_rows[i] is row i of timestep_table. + std::vector hidden_state_halves; + std::vector timestep_rows; + + std::unique_ptr text_encoder_io; + std::unique_ptr transformer_io; + + CfgDataPredictionStage cfg; + DpmSolverStage solver; + StageTiming timing; +}; + +Ort::RunOptions make_run_options(PipelineState& state) +{ + Ort::RunOptions run_options; + if (state.compute_stream) + { + run_options.SetSyncStream(*state.compute_stream); + run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); + } + return run_options; +} + +void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const PixelDiTConfig& config) +{ + state.model_paths = MakePixelDiTModelPaths(config.model_dir); + state.pipeline_config = LoadPixelDiTPipelineConfig(state.model_paths.pipeline_config); + const PixelDiTPipelineConfig& pc = state.pipeline_config; + if (!pc.file_found) + { + std::cout << "pipeline_config.json not found at " << state.model_paths.pipeline_config.string() + << "; using built-in defaults (1024x1024 export)" << std::endl; + } + else if (!pc.defaults_used.empty()) + { + std::cout << "pipeline_config.json is missing fields, using defaults for:"; + for (const auto& name : pc.defaults_used) + { + std::cout << " " << name; + } + std::cout << std::endl; + } + + state.sampler.steps = config.steps.value_or(pc.steps); + state.sampler.cfg_scale = config.cfg_scale.value_or(pc.cfg_scale); + state.sampler.flow_shift = config.flow_shift.value_or(pc.flow_shift); + state.sampler.negative_prompt = config.negative_prompt.value_or(pc.negative_prompt); + if (state.sampler.steps < pc.order) + { + throw std::invalid_argument("--steps must be at least " + std::to_string(pc.order)); + } + std::cout << "Image size: " << pc.height << "x" << pc.width << "\n" + << "Steps: " << state.sampler.steps << ", CFG scale: " << state.sampler.cfg_scale + << ", flow shift: " << state.sampler.flow_shift << "\n" + << "Negative prompt: " << state.sampler.negative_prompt << std::endl; + + std::cout << "Loading tokenizer " << state.model_paths.tokenizer_json.string() << std::endl; + state.tokenizer = + std::make_unique(state.model_paths.tokenizer_json.string(), din::io::TokenizerFormat::Json); + const size_t chi_tokens = state.tokenizer->Encode(pc.chi_prompt, true).size(); + if (static_cast(chi_tokens) + pc.txt_max_length - 2 != pc.text_seq_len) + { + throw std::runtime_error("chi_prompt tokenizes to " + std::to_string(chi_tokens) + + " tokens, inconsistent with text_seq_len=" + std::to_string(pc.text_seq_len) + + " and txt_max_length=" + std::to_string(pc.txt_max_length)); + } + + const int cuda_device_ordinal = din::common::ChooseCudaDeviceOrdinal(trt_device, "PIXELDIT_CUDA_DEVICE_ID"); + PIXELDIT_CUDA_CHECK(cudaSetDevice(cuda_device_ordinal)); + cudaDeviceProp cuda_device_prop{}; + PIXELDIT_CUDA_CHECK(cudaGetDeviceProperties(&cuda_device_prop, cuda_device_ordinal)); + std::cout << "TensorRT RTX GPU: " << cuda_device_prop.name << " (CUDA device ordinal " << cuda_device_ordinal << ")" + << std::endl; + state.compute_stream.emplace(din::common::CreateTensorRTRTXComputeStream(state.env)); + + 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(); + const 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; // both engines exceed the ~2 GB protobuf limit + // Leave TensorRT auxiliary streams at the EP default: forcing 0 (as Flux2 does) made the 1024x1024 DiT + // call take ~67 ms instead of ~48 ms on an RTX PRO 6000. + return profile; + }; + Ort::SyncStream* stream = &*state.compute_stream; + state.text_encoder_runner = + std::make_unique(state.env, state.model_paths.text_encoder_model.string(), "trt-rtx", + cache_dir, ep_context, make_profile("text_encoder"), stream); + std::cout << " Text encoder loaded" << std::endl; + state.transformer_runner = + std::make_unique(state.env, state.model_paths.transformer_model.string(), "trt-rtx", + cache_dir, ep_context, make_profile("transformer"), stream); + std::cout << " Transformer loaded" << std::endl; + + Ort::Session& te = state.text_encoder_runner->session; + Ort::Session& dit = state.transformer_runner->session; + const auto te_out = session_shape(te, "prompt_embeds", false); + if (te_out.size() != 3) + { + throw std::runtime_error("text_encoder prompt_embeds must be rank 3, got " + shape_string(te_out)); + } + state.text_hidden = te_out[2]; + const std::vector text_shape = {1, pc.text_seq_len}; + const std::vector select_shape = {pc.txt_max_length}; + const std::vector embeds_shape = {1, pc.txt_max_length, state.text_hidden}; + const std::vector image_shape = {pc.dit_batch, 3, pc.height, pc.width}; + const std::vector timestep_shape = {pc.dit_batch}; + const std::vector dit_text_shape = {pc.dit_batch, pc.txt_max_length, state.text_hidden}; + const std::vector sample_shape = {1, 3, pc.height, pc.width}; + auto& te_runner = *state.text_encoder_runner; + auto& dit_runner = *state.transformer_runner; + state.input_ids = std::make_unique(te_runner, text_shape, true); + state.attention_mask = std::make_unique(te_runner, text_shape, true); + state.select_index = std::make_unique(te_runner, select_shape, true); + state.prompt_embeds = std::make_unique(te_runner, embeds_shape, true); + state.dit_hidden_states = std::make_unique(dit_runner, image_shape, true); + state.dit_timestep = std::make_unique(dit_runner, timestep_shape, true); + state.dit_encoder_hidden_states = std::make_unique(dit_runner, dit_text_shape, true); + state.dit_velocity = std::make_unique(dit_runner, image_shape, true); + const bool sample_on_device = state.sampling_backend == SamplingBackend::Cuda; + state.sample = std::make_unique(dit_runner, sample_shape, sample_on_device); + state.model_a = std::make_unique(dit_runner, sample_shape, sample_on_device); + state.model_b = std::make_unique(dit_runner, sample_shape, sample_on_device); + if (sample_on_device) + { + const std::vector table_shape = {state.sampler.steps, pc.dit_batch}; + state.timestep_table = std::make_unique(dit_runner, table_shape, true); + const size_t n = element_count(*state.sample); + for (int64_t b = 0; b < pc.dit_batch; ++b) + { + state.hidden_state_halves.push_back( + make_view(*state.dit_hidden_states, static_cast(b) * n, sample_shape)); + } + for (int64_t i = 0; i < state.sampler.steps; ++i) + { + state.timestep_rows.push_back( + make_view(*state.timestep_table, static_cast(i * pc.dit_batch), timestep_shape)); + } + } + + state.text_encoder_io = std::make_unique(te); + state.text_encoder_io->BindInput("input_ids", state.input_ids->BindingValue()); + state.text_encoder_io->BindInput("attention_mask", state.attention_mask->BindingValue()); + state.text_encoder_io->BindInput("select_index", state.select_index->BindingValue()); + state.text_encoder_io->BindOutput("prompt_embeds", state.prompt_embeds->BindingValue()); + + state.transformer_io = std::make_unique(dit); + state.transformer_io->BindInput("hidden_states", state.dit_hidden_states->BindingValue()); + state.transformer_io->BindInput("timestep", state.dit_timestep->BindingValue()); + state.transformer_io->BindInput("encoder_hidden_states", state.dit_encoder_hidden_states->BindingValue()); + state.transformer_io->BindOutput("velocity", state.dit_velocity->BindingValue()); + + state.cfg.BindInput("velocity", *state.dit_velocity); + state.cfg.BindInput("sample", *state.sample); + state.solver.BindOutput("sample", *state.sample); +} + +// Encodes one prompt and writes it into row `row` of the DiT encoder_hidden_states host buffer. +void encode_prompt(PipelineState& state, const std::string& text, const std::vector& select, int64_t row, + const std::filesystem::path& dump_path) +{ + const PixelDiTPipelineConfig& pc = state.pipeline_config; + const PixelDiTTextInputs inputs = MakeTextEncoderInputs(*state.tokenizer, text, pc.text_seq_len, pc.pad_token_id); + if (inputs.num_tokens > static_cast(pc.text_seq_len)) + { + std::cout << "Prompt tokens truncated from " << inputs.num_tokens << " to " << pc.text_seq_len << std::endl; + } + std::copy(inputs.input_ids.begin(), inputs.input_ids.end(), state.input_ids->HostData()); + std::copy(inputs.attention_mask.begin(), inputs.attention_mask.end(), state.attention_mask->HostData()); + std::copy(select.begin(), select.end(), state.select_index->HostData()); + if (!dump_path.empty()) + { + write_raw(dump_path, inputs.input_ids.data(), inputs.input_ids.size()); + } + + auto ids_ready = state.input_ids->CopyAsyncToDeviceWithNotification(); + auto mask_ready = state.attention_mask->CopyAsyncToDeviceWithNotification(); + auto select_ready = state.select_index->CopyAsyncToDeviceWithNotification(); + ids_ready.Sync(); + mask_ready.Sync(); + select_ready.Sync(); + + Ort::RunOptions run_options = make_run_options(state); + state.text_encoder_runner->session.Run(run_options, *state.text_encoder_io); + auto embeds_ready = state.prompt_embeds->CopyAsyncToHostWithNotification(); + embeds_ready.Sync(); + + const size_t row_size = static_cast(pc.txt_max_length * state.text_hidden); + std::copy_n(state.prompt_embeds->HostData(), row_size, + state.dit_encoder_hidden_states->HostData() + static_cast(row) * row_size); +} + +void encode_prompts(PipelineState& state, const PixelDiTConfig& config) +{ + din::common::nvtx_scoped_range nvtx("encode_prompts"); + const PixelDiTPipelineConfig& pc = state.pipeline_config; + const bool dump = !config.dump_dir.empty(); + const auto start = std::chrono::steady_clock::now(); + // Batch row 0 is the negative prompt, row 1 the positive prompt (chi_prompt prefix + prompt). + encode_prompt(state, state.sampler.negative_prompt, pc.negative_select_index, 0, + dump ? config.dump_dir / "tokens_negative.bin" : std::filesystem::path{}); + encode_prompt(state, pc.chi_prompt + StripWhitespace(state.prompt), pc.positive_select_index, 1, + dump ? config.dump_dir / "tokens_positive.bin" : std::filesystem::path{}); + auto uploaded = state.dit_encoder_hidden_states->CopyAsyncToDeviceWithNotification(); + uploaded.Sync(); + if (dump) + { + write_raw(config.dump_dir / "prompt_embeds.bin", state.dit_encoder_hidden_states->HostData(), + element_count(*state.dit_encoder_hidden_states)); + } + state.prompt_embeds_valid = true; + const std::chrono::duration elapsed = std::chrono::steady_clock::now() - start; + std::cout << "Text encoding (negative + positive): " << elapsed.count() << " s" << std::endl; +} + +cudaStream_t compute_stream(PipelineState& state) +{ + return reinterpret_cast(state.compute_stream->GetHandle()); +} + +// --processing cuda: one model evaluation, enqueued on the compute stream without any host sync. +// `eval_index` selects the row of timestep_table (timestep t * timestep_scale). +void data_prediction_cuda(PipelineState& state, float t, int eval_index, FloatBuffer& x0) +{ + Ort::SyncStream& sync_stream = *state.compute_stream; + for (Ort::Value& half : state.hidden_state_halves) + { + Ort::ThrowOnError(state.env.CopyTensor(state.sample->BindingValue(), half, sync_stream)); + } + Ort::ThrowOnError(state.env.CopyTensor(state.timestep_rows[static_cast(eval_index)], + state.dit_timestep->BindingValue(), sync_stream)); + const cudaStream_t stream = compute_stream(state); + { + din::common::nvtx_scoped_range nvtx("transformer"); + Ort::RunOptions run_options = make_run_options(state); + state.transformer_runner->session.Run(run_options, *state.transformer_io); + } + { + din::common::nvtx_scoped_range nvtx("cfg_data_prediction"); + state.cfg.Run(stream, SamplingBackend::Cuda, t, state.sampler.cfg_scale, x0); + } + ++state.timing.evaluations; +} + +// One model evaluation: DiT on [sample, sample] at time t, then CFG -> x0 estimate. +void data_prediction(PipelineState& state, float t, int eval_index, din::common::TensorBuffer& x0) +{ + if (state.sampling_backend == SamplingBackend::Cuda) + { + data_prediction_cuda(state, t, eval_index, x0); + return; + } + const PixelDiTPipelineConfig& pc = state.pipeline_config; + const size_t n = element_count(*state.sample); + const auto start = std::chrono::steady_clock::now(); + float* hidden = state.dit_hidden_states->HostData(); + std::copy_n(state.sample->HostData(), n, hidden); + std::copy_n(state.sample->HostData(), n, hidden + n); + const float timestep = t * pc.timestep_scale; + std::fill_n(state.dit_timestep->HostData(), static_cast(pc.dit_batch), timestep); + auto hidden_ready = state.dit_hidden_states->CopyAsyncToDeviceWithNotification(); + auto timestep_ready = state.dit_timestep->CopyAsyncToDeviceWithNotification(); + hidden_ready.Sync(); + timestep_ready.Sync(); + const auto uploaded = std::chrono::steady_clock::now(); + + { + din::common::nvtx_scoped_range nvtx("transformer"); + Ort::RunOptions run_options = make_run_options(state); + state.transformer_runner->session.Run(run_options, *state.transformer_io); + auto velocity_ready = state.dit_velocity->CopyAsyncToHostWithNotification(); + velocity_ready.Sync(); + } + const auto dit_done = std::chrono::steady_clock::now(); + + { + din::common::nvtx_scoped_range nvtx("cfg_data_prediction"); + // Velocity is already on the host. + state.cfg.Run(compute_stream(state), SamplingBackend::Cpu, t, state.sampler.cfg_scale, x0); + } + const auto cfg_done = std::chrono::steady_clock::now(); + state.timing.upload += uploaded - start; + state.timing.dit += dit_done - uploaded; + state.timing.cfg += cfg_done - dit_done; + ++state.timing.evaluations; +} + +PixelDiTImage run_pipeline(PipelineState& state, const PixelDiTConfig& config, unsigned int seed) +{ + const PixelDiTPipelineConfig& pc = state.pipeline_config; + const bool dump = !config.dump_dir.empty(); + const size_t n = element_count(*state.sample); + if (!config.init_noise_path.empty()) + { + const std::vector noise = read_raw_floats(config.init_noise_path, n); + std::copy(noise.begin(), noise.end(), state.sample->HostData()); + std::cout << "Initial noise loaded from " << config.init_noise_path.string() << std::endl; + } + else + { + initialize_noise(state.sample->HostData(), n, seed); + } + if (dump) + { + write_raw(config.dump_dir / ("noise_seed" + std::to_string(seed) + ".bin"), state.sample->HostData(), n); + } + + const int steps = state.sampler.steps; + const std::vector timesteps = flow_timesteps(steps, state.sampler.flow_shift, pc.t_start, pc.t_end); + if (dump) + { + write_raw(config.dump_dir / "timesteps.bin", timesteps.data(), timesteps.size()); + } + + const bool cuda = state.sampling_backend == SamplingBackend::Cuda; + state.timing = {}; + const auto start = std::chrono::steady_clock::now(); + if (cuda) + { + // Model evaluations happen at timesteps[0 .. steps-1]. + float* table = state.timestep_table->HostData(); + const size_t batch = static_cast(pc.dit_batch); + for (size_t i = 0; i < static_cast(steps); ++i) + { + std::fill_n(table + i * batch, batch, timesteps[i] * pc.timestep_scale); + } + auto sample_uploaded = state.sample->CopyAsyncToDeviceWithNotification(); + auto table_uploaded = state.timestep_table->CopyAsyncToDeviceWithNotification(); + sample_uploaded.Sync(); + table_uploaded.Sync(); + } + + // FlowDPMSolver.sample (model_export/pixeldit/pipeline.py): 1st order first, 2nd order in the middle, + // 1st order last (lower_order_final), and no model evaluation after the final step. + din::common::TensorBuffer* prev = state.model_a.get(); // x0 at t_prev_1 + din::common::TensorBuffer* curr = state.model_b.get(); // x0 at t_prev_0 + float t_prev_1 = 0.0f; + float t_prev_0 = timesteps[0]; + data_prediction(state, timesteps[0], 0, *curr); + int history = 1; + for (int step = 1; step <= steps; ++step) + { + din::common::nvtx_scoped_range nvtx_step(("sampling_step_" + std::to_string(step)).c_str()); + const float t = timesteps[static_cast(step)]; + const int step_order = step < pc.order ? 1 : std::min(pc.order, steps + 1 - step); + if (step_order == 2 && history < 2) + { + throw std::logic_error("2nd-order update without two x0 estimates"); + } + const auto solver_start = std::chrono::steady_clock::now(); + state.solver.Run(state.env, *state.compute_stream, state.sampling_backend, step_order, prev, *curr, t_prev_1, + t_prev_0, t); + state.timing.solver += std::chrono::steady_clock::now() - solver_start; + std::swap(prev, curr); + t_prev_1 = t_prev_0; + t_prev_0 = t; + history = std::min(history + 1, 2); + if (step < steps) + { + data_prediction(state, t, step, *curr); + } + // With --processing cuda this reports enqueued steps; the GPU may still be behind. + if (step % 10 == 0 || step == steps) + { + std::cout << " step " << step << "/" << steps << std::endl; + } + } + if (cuda) + { + auto sample_ready = state.sample->CopyAsyncToHostWithNotification(); + sample_ready.Sync(); + } + const std::chrono::duration elapsed = std::chrono::steady_clock::now() - start; + std::cout << "Sampling: " << elapsed.count() << " s (" << elapsed.count() / steps * 1000.0 << " ms/step)" + << std::endl; + if (!cuda) + { + const double evals = std::max(state.timing.evaluations, 1); + std::cout << " per DiT evaluation: upload " << state.timing.upload.count() / evals * 1000.0 + << " ms, DiT + download " << state.timing.dit.count() / evals * 1000.0 << " ms, CFG " + << state.timing.cfg.count() / evals * 1000.0 << " ms; solver " + << state.timing.solver.count() / steps * 1000.0 << " ms/step" << std::endl; + } + + if (dump) + { + write_raw(config.dump_dir / ("final_seed" + std::to_string(seed) + ".bin"), state.sample->HostData(), n); + } + PixelDiTImage image; + image.height = static_cast(pc.height); + image.width = static_cast(pc.width); + image.data.assign(state.sample->HostData(), state.sample->HostData() + n); + return image; +} + +// One ORT environment per process, with the TensorRT RTX EP registered once. +struct ProcessRuntime +{ + explicit ProcessRuntime(Ort::Env& environment) + : env(environment) + { + } + + Ort::Env& env; + Ort::ConstEpDevice trt_device{}; +}; + +ProcessRuntime& GetProcessRuntime() +{ + static Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "PixelDiT"}; + static ProcessRuntime runtime{env}; + static std::once_flag register_trt_once; + + std::call_once(register_trt_once, + [] + { + runtime.trt_device = din::common::RegisterTensorRTRTXProvider(runtime.env); + }); + return runtime; +} + +class TrtPixelDiTProcessingPipeline final : public PixelDiTProcessingPipeline +{ +public: + TrtPixelDiTProcessingPipeline(PixelDiTConfig config, ProcessRuntime& runtime) + : config_(std::move(config)) + , runtime_(runtime) + { + } + + void Initialize() override + { + if (state_) + { + return; + } + if (std::strcmp(runtime_.trt_device.EpName(), kDinNvTensorRTRTXExecutionProvider) != 0) + { + throw std::runtime_error("TensorRT RTX execution provider requested but no device was found"); + } + if (!config_.dump_dir.empty()) + { + std::filesystem::create_directories(config_.dump_dir); + } + auto state = std::make_unique(runtime_.env); + state->sampling_backend = + config_.processing == PixelDiTProcessingBackend::Cuda ? SamplingBackend::Cuda : SamplingBackend::Cpu; + std::cout << "Sampling backend: " << to_string(state->sampling_backend) << std::endl; + initialize_state(*state, runtime_.trt_device, config_); + state->prompt = config_.prompt; + state_ = std::move(state); + } + + void SetPrompt(std::string prompt) override + { + config_.prompt = std::move(prompt); + if (state_) + { + state_->prompt = config_.prompt; + state_->prompt_embeds_valid = false; + } + } + + PixelDiTImage GenerateImage(unsigned int seed) override + { + if (!state_) + { + throw std::runtime_error("PixelDiT pipeline was not initialized"); + } + if (!state_->prompt_embeds_valid) + { + encode_prompts(*state_, config_); + } + return run_pipeline(*state_, config_, seed); + } + +private: + PixelDiTConfig config_; + ProcessRuntime& runtime_; + std::unique_ptr state_; +}; + +} // namespace + +// TensorRT RTX is the only execution provider for PixelDiT; config.processing (cpu or cuda) only selects where the +// sampling math runs, inside this pipeline. +std::unique_ptr CreatePixelDiTPipeline(const PixelDiTConfig& config) +{ + std::cout << std::unitbuf; + std::cerr << std::unitbuf; + return std::make_unique(config, GetProcessRuntime()); +} + +} // namespace din::image_gen diff --git a/tests/tokenizer/CMakeLists.txt b/tests/tokenizer/CMakeLists.txt new file mode 100644 index 0000000..b196a60 --- /dev/null +++ b/tests/tokenizer/CMakeLists.txt @@ -0,0 +1,14 @@ +# din::io::Tokenizer unit test. Needs no models, GPU or artifacts. + +# Plain add_executable: add_din_executable's post-build DLL copy fails when a target has no runtime DLLs +# (this one only links static libraries). +add_executable(din_tokenizer_test + test_tokenizer.cpp +) +set_target_properties(din_tokenizer_test PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${DIN_RUNTIME_DIR}") +target_link_libraries(din_tokenizer_test PRIVATE din_common_io nlohmann_json::nlohmann_json) + +if(DIN_BUILD_TESTING) + add_test(NAME din_tokenizer_test COMMAND din_tokenizer_test) + set_tests_properties(din_tokenizer_test PROPERTIES LABELS "unit") +endif() diff --git a/tests/tokenizer/test_tokenizer.cpp b/tests/tokenizer/test_tokenizer.cpp new file mode 100644 index 0000000..c747f78 --- /dev/null +++ b/tests/tokenizer/test_tokenizer.cpp @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Unit test for din::io::Tokenizer::Encode on three tiny tokenizer.json layouts built in code: +// * Gemma style (SentencePiece-style BPE): Replace(" " -> "▁"), byte_fallback, template, +// special and non-special added tokens, no dummy prefix. +// * Llama style: the same with a Prepend("▁") normalizer. +// * GPT-2 / Qwen style byte-level BPE (regression for the existing path). +// The expected ids were produced by Hugging Face `tokenizers` 0.23.2 from identical tokenizer.json content. +// It also checks that a Gemma-style tokenizer with an unsupported normalizer step fails to load. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tokenizer.h" +#include + +namespace +{ + +using nlohmann::json; + +constexpr const char* kSpace = "\xE2\x96\x81"; // U+2581, SentencePiece space marker + +struct Case +{ + std::string text; + bool add_special_tokens; + std::vector ids; +}; + +std::string utf8(uint32_t code) +{ + std::string out; + if (code < 0x80) + { + out.push_back(static_cast(code)); + } + else if (code < 0x800) + { + out.push_back(static_cast(0xC0 | (code >> 6))); + out.push_back(static_cast(0x80 | (code & 0x3F))); + } + else + { + out.push_back(static_cast(0xE0 | (code >> 12))); + out.push_back(static_cast(0x80 | ((code >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (code & 0x3F))); + } + return out; +} + +json added_token(int64_t id, const std::string& content, bool special) +{ + return {{"id", id}, {"content", content}, {"single_word", false}, {"lstrip", false}, + {"rstrip", false}, {"normalized", false}, {"special", special}}; +} + +// Vocab ids: \n \n\n a b c d e . ▁ ab bc abc ▁a ▁ab aa, then <0x00>..<0xFF> from 20. +json sentencepiece_tokenizer(bool prepend) +{ + const std::string sp = kSpace; + std::vector tokens = {"", "", "", "", "\n", "\n\n", "", + "a", "b", "c", "d", "e", ".", sp, + "ab", "bc", "abc", sp + "a", sp + "ab", "aa"}; + for (int b = 0; b < 256; ++b) + { + char name[8]; + std::snprintf(name, sizeof(name), "<0x%02X>", b); + tokens.emplace_back(name); + } + json vocab = json::object(); + for (size_t i = 0; i < tokens.size(); ++i) + { + vocab[tokens[i]] = static_cast(i); + } + // Rank order matters: "a b" before "b c"; "a a" last so ties resolve leftmost. + const json merges = json::array({{"a", "b"}, {"b", "c"}, {"ab", "c"}, {sp, "a"}, {sp + "a", "b"}, {"a", "a"}}); + + const json replace = {{"type", "Replace"}, {"pattern", {{"String", " "}}}, {"content", sp}}; + const json normalizer = + prepend ? json{{"type", "Sequence"}, {"normalizers", {{{"type", "Prepend"}, {"prepend", sp}}, replace}}} + : replace; + + return { + {"version", "1.0"}, + {"added_tokens", + {added_token(0, "", true), added_token(1, "", true), added_token(2, "", true), + added_token(3, "", true), added_token(4, "\n", false), added_token(5, "\n\n", false), + added_token(6, "", false)}}, + {"normalizer", normalizer}, + {"pre_tokenizer", prepend ? json(nullptr) + : json{{"type", "Split"}, + {"pattern", {{"String", " "}}}, + {"behavior", "MergedWithPrevious"}, + {"invert", false}}}, + {"post_processor", + {{"type", "TemplateProcessing"}, + {"single", + {{{"SpecialToken", {{"id", ""}, {"type_id", 0}}}}, {{"Sequence", {{"id", "A"}, {"type_id", 0}}}}}}, + {"special_tokens", {{"", {{"id", ""}, {"ids", {2}}, {"tokens", {""}}}}}}}}, + {"model", + {{"type", "BPE"}, + {"unk_token", ""}, + {"fuse_unk", true}, + {"byte_fallback", true}, + {"ignore_merges", false}, + {"vocab", vocab}, + {"merges", merges}}}, + }; +} + +// GPT-2 byte-level alphabet (256 code points) in code-point order, then merged pieces. +json byte_level_tokenizer() +{ + std::vector alphabet; + uint32_t extra = 0; + for (uint32_t b = 0; b < 256; ++b) + { + const bool printable = (b >= '!' && b <= '~') || (b >= 0xA1 && b <= 0xAC) || (b >= 0xAE && b <= 0xFF); + alphabet.push_back(printable ? b : 256 + extra++); + } + std::sort(alphabet.begin(), alphabet.end()); + const std::string g = utf8(0x120); // byte-level space + const std::vector> merge_list = { + {"h", "e"}, {"l", "l"}, {"he", "ll"}, {"hell", "o"}, {g, "w"}, {g + "w", "o"}, + {"a", "b"}, {"ab", "c"}, {g, "a"}, {g + "a", "b"}, {g + "ab", "c"}}; + json vocab = json::object(); + int64_t id = 0; + for (const uint32_t code : alphabet) + { + vocab[utf8(code)] = id++; + } + json merges = json::array(); + for (const auto& [left, right] : merge_list) + { + merges.push_back({left, right}); + if (!vocab.contains(left + right)) + { + vocab[left + right] = id++; + } + } + return { + {"version", "1.0"}, + {"added_tokens", json::array()}, + {"normalizer", nullptr}, + {"pre_tokenizer", + {{"type", "ByteLevel"}, {"add_prefix_space", false}, {"trim_offsets", true}, {"use_regex", true}}}, + {"post_processor", nullptr}, + {"model", {{"type", "BPE"}, {"byte_fallback", false}, {"vocab", vocab}, {"merges", merges}}}, + }; +} + +// clang-format off +const std::vector kGemmaStyleCases = { + {"", true, {2}}, + {"", false, {}}, + {"abc", true, {2, 16}}, + {"abc", false, {16}}, + {"aaa", true, {2, 19, 7}}, + {"aaa", false, {19, 7}}, + {"a b", true, {2, 7, 13, 8}}, + {"a b", false, {7, 13, 8}}, + {" ab", true, {2, 13, 13, 14}}, + {" ab", false, {13, 13, 14}}, + {"ab ", true, {2, 14, 13}}, + {"ab ", false, {14, 13}}, + {"a b", true, {2, 7, 13, 13, 8}}, + {"a b", false, {7, 13, 13, 8}}, + {"a\xC3\xA9", true, {2, 7, 215, 189}}, + {"a\xC3\xA9", false, {7, 215, 189}}, + {"a\xF0\x9F\xA6\x8A", true, {2, 7, 260, 179, 186, 158}}, + {"a\xF0\x9F\xA6\x8A", false, {7, 260, 179, 186, 158}}, + {"x", true, {2, 140}}, + {"x", false, {140}}, + {"a\nb", true, {2, 7, 4, 8}}, + {"a\nb", false, {7, 4, 8}}, + {"a\n\nb", true, {2, 7, 5, 8}}, + {"a\n\nb", false, {7, 5, 8}}, + {"a\n\n\nb", true, {2, 7, 5, 4, 8}}, + {"a\n\n\nb", false, {7, 5, 4, 8}}, + {"ab", true, {2, 6, 14, 80, 67, 125, 82}}, + {"ab", false, {6, 14, 80, 67, 125, 82}}, + {"ab", true, {2, 7, 1, 8}}, + {"ab", false, {7, 1, 8}}, + {"ab c.\n\na\xC3\xA9", true, {2, 14, 13, 9, 12, 5, 2, 7, 215, 189}}, + {"ab c.\n\na\xC3\xA9", false, {14, 13, 9, 12, 5, 2, 7, 215, 189}}, +}; + +const std::vector kLlamaStyleCases = { + {"", true, {2}}, + {"", false, {}}, + {"abc", true, {2, 13, 16}}, + {"abc", false, {13, 16}}, + {"aaa", true, {2, 17, 19}}, + {"aaa", false, {17, 19}}, + {"a b", true, {2, 17, 13, 8}}, + {"a b", false, {17, 13, 8}}, + {" ab", true, {2, 13, 13, 13, 14}}, + {" ab", false, {13, 13, 13, 14}}, + {"ab ", true, {2, 13, 14, 13}}, + {"ab ", false, {13, 14, 13}}, + {"a b", true, {2, 17, 13, 13, 8}}, + {"a b", false, {17, 13, 13, 8}}, + {"a\xC3\xA9", true, {2, 17, 215, 189}}, + {"a\xC3\xA9", false, {17, 215, 189}}, + {"a\xF0\x9F\xA6\x8A", true, {2, 17, 260, 179, 186, 158}}, + {"a\xF0\x9F\xA6\x8A", false, {17, 260, 179, 186, 158}}, + {"x", true, {2, 13, 140}}, + {"x", false, {13, 140}}, + {"a\nb", true, {2, 17, 4, 13, 8}}, + {"a\nb", false, {17, 4, 13, 8}}, + {"a\n\nb", true, {2, 17, 5, 13, 8}}, + {"a\n\nb", false, {17, 5, 13, 8}}, + {"a\n\n\nb", true, {2, 17, 5, 4, 13, 8}}, + {"a\n\n\nb", false, {17, 5, 4, 13, 8}}, + {"ab", true, {2, 6, 13, 14, 80, 67, 125, 82}}, + {"ab", false, {6, 13, 14, 80, 67, 125, 82}}, + {"ab", true, {2, 17, 1, 13, 8}}, + {"ab", false, {17, 1, 13, 8}}, + {"ab c.\n\na\xC3\xA9", true, {2, 13, 14, 13, 9, 12, 5, 2, 17, 215, 189}}, + {"ab c.\n\na\xC3\xA9", false, {13, 14, 13, 9, 12, 5, 2, 17, 215, 189}}, +}; + +// The byte-level path ignores add_special_tokens; this tokenizer has no special tokens. +const std::vector kByteLevelCases = { + {"", false, {}}, + {"abc abc", false, {263, 220, 263}}, + {"hello world", false, {259, 261, 81, 75, 67}}, + {"a,b. c", false, {64, 11, 65, 13, 220, 66}}, +}; +// clang-format on + +std::string escape(const std::string& text) +{ + std::ostringstream out; + for (const unsigned char c : text) + { + if (c == '\n') + out << "\\n"; + else if (c < 0x20 || c >= 0x7F) + { + char hex[8]; + std::snprintf(hex, sizeof(hex), "\\x%02X", c); + out << hex; + } + else + out << c; + } + return out.str(); +} + +std::string join(const std::vector& ids) +{ + std::ostringstream out; + out << "["; + for (size_t i = 0; i < ids.size(); ++i) + { + out << (i ? ", " : "") << ids[i]; + } + out << "]"; + return out.str(); +} + +// din::io::Tokenizer loads from a path, so each tokenizer.json is written to a temporary directory. +int run_cases(const std::filesystem::path& dir, const char* name, const json& tokenizer_json, + const std::vector& cases) +{ + const auto path = dir / name; + std::ofstream(path, std::ios::binary) << tokenizer_json.dump(); + const din::io::Tokenizer tokenizer(path.string(), din::io::TokenizerFormat::Json); + int failures = 0; + for (const auto& test_case : cases) + { + const auto got = tokenizer.Encode(test_case.text, test_case.add_special_tokens); + if (got != test_case.ids) + { + ++failures; + std::cerr << "FAIL " << name << " \"" << escape(test_case.text) + << "\" add_special_tokens=" << test_case.add_special_tokens << "\n expected " + << join(test_case.ids) << "\n got " << join(got) << std::endl; + } + } + std::cout << name << ": " << (cases.size() - failures) << "/" << cases.size() << " cases match" << std::endl; + return failures; +} + +// A SentencePiece-style tokenizer with an unsupported normalizer step must fail to load, not fall back to +// byte-level BPE. +int check_unsupported_normalizer_rejected(const std::filesystem::path& dir) +{ + json tokenizer_json = sentencepiece_tokenizer(false); + tokenizer_json["normalizer"] = {{"type", "Sequence"}, + {"normalizers", {{{"type", "NFKC"}}, tokenizer_json["normalizer"]}}}; + const auto path = dir / "unsupported_normalizer.json"; + std::ofstream(path, std::ios::binary) << tokenizer_json.dump(); + try + { + const din::io::Tokenizer tokenizer(path.string(), din::io::TokenizerFormat::Json); + } + catch (const std::runtime_error&) + { + std::cout << "unsupported_normalizer.json: rejected" << std::endl; + return 0; + } + std::cerr << "FAIL unsupported_normalizer.json: loaded a SentencePiece tokenizer with an NFKC normalizer" + << std::endl; + return 1; +} + +} // namespace + +int main() +{ + try + { + const auto dir = std::filesystem::temp_directory_path() / "din_tokenizer_test"; + std::filesystem::create_directories(dir); + int failures = 0; + failures += run_cases(dir, "gemma_style.json", sentencepiece_tokenizer(false), kGemmaStyleCases); + failures += run_cases(dir, "llama_style.json", sentencepiece_tokenizer(true), kLlamaStyleCases); + failures += run_cases(dir, "byte_level.json", byte_level_tokenizer(), kByteLevelCases); + failures += check_unsupported_normalizer_rejected(dir); + std::filesystem::remove_all(dir); + return failures == 0 ? 0 : 1; + } + catch (const std::exception& e) + { + std::cerr << "Error: " << e.what() << std::endl; + return 1; + } +}