switch to pixelDiT repo and packaging

README.md

@@ -62,7 +62,7 @@ Use the model-specific README for the exporter command and CLI arguments.
 Create a Python environment and install the repository plus the extra required by the model you plan to export:
 
 ```bash
-python -m pip install -e ".[parakeet]"  # or .[nemotron] / .[flux]
+python -m pip install -e ".[parakeet]"  # or .[nemotron] / .[flux] / .[pixeldit]
 ```
 
 The base `pyproject.toml` contains common exporter dependencies; the optional extras add the model-family dependencies. Export the model to a directory, then pass that directory to the corresponding CLI with `--model-dir`.

image_gen/pixeldit/README.md

@@ -9,7 +9,7 @@ CLI image generation with ONNX Runtime.
 | 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.
+PixelDiT generates pixels directly, so there is no VAE. The DiT model code comes from the `pixeldit` pip package, installed from [NVlabs/PixelDiT](https://github.com/NVlabs/PixelDiT).
 
 ## Pipeline
 
@@ -37,11 +37,16 @@ DirectX and Vulkan processing are not implemented yet. `--provider cpu` is not s
 
 ## Export
 
+Install the export dependencies first, from the repository root. The `pixeldit` extra adds `pixeldit`, upstream's PixelDiT model package, to the common exporter dependencies (see [Requirements](../../README.md#requirements)). Export and verify run on the GPU, so install a CUDA build of PyTorch first (see [pytorch.org](https://pytorch.org/get-started/locally/)).
+
 ```bash
-cd model_export
+python -m pip install -e ".[pixeldit]"
+cd image_gen/pixeldit/model_export
 python export_pixeldit.py --output ./models/PixelDiT-1300M-1024px-onnx
 ```
 
+The DiT is upstream's `PixDiT_T2I` (imported as `pixdit_core`; the local `model_export/pixeldit/` folder holds din-deploy's export helpers). `export_pixeldit.py` patches it only while exporting: complex-valued RoPE has no ONNX lowering, and TensorRT RTX rejects the `Col2Im` op that `F.fold` exports to. See `model_export/pixeldit/upstream_patch.py` for the details. `test_upstream_patch.py` checks the patches on a small random model.
+
 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 <TensorRT-RTX>` to check that each model builds with the TensorRT RTX CLI.
@@ -52,7 +57,7 @@ Shapes are static. `--height` and `--width` (default 1024, multiples of 16) set
 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.
+Compares each ONNX model running on TensorRT RTX with the upstream PyTorch model from `pixeldit` (unpatched), 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.
 

image_gen/pixeldit/model_export/export_pixeldit.py

@@ -32,6 +32,7 @@
 import torch
 
 from pixeldit import pipeline as pixeldit_pipeline
+from pixeldit.upstream_patch import export_patches
 from pixeldit.weights import DEFAULT_MODEL_NAME, DEFAULT_TEXT_ENCODER_NAME, load_gemma_text_encoder, load_pixeldit
 
 EXTERNAL_DATA_NAME = "model.onnx_data"
@@ -114,15 +115,18 @@ def forward(self, input_ids, attention_mask, select_index):
 
 
 class PixelDiTExportWrapper(torch.nn.Module):
-    """Inputs are cast to bf16 inside PixelDiT (as upstream PixDiTTrainer does), output to io_dtype."""
+    """Inputs are cast to the model dtype , output to io_dtype."""
 
     def __init__(self, model, io_dtype: torch.dtype):
         super().__init__()
         self.model = model
         self.io_dtype = io_dtype
+        self.model_dtype = next(model.parameters()).dtype
 
     def forward(self, hidden_states, timestep, encoder_hidden_states):
-        return self.model(hidden_states, timestep, encoder_hidden_states).to(self.io_dtype)
+        dtype = self.model_dtype
+        velocity = self.model(hidden_states.to(dtype), timestep.to(dtype), encoder_hidden_states.to(dtype))
+        return velocity.to(self.io_dtype)
 
 
 @torch.no_grad()
@@ -384,8 +388,8 @@ def main():
         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)
+            with export_patches(model, args.height, args.width, model.txt_max_length):
+                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")

image_gen/pixeldit/model_export/pixeldit/__init__.py

@@ -1,3 +1,5 @@
 # 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."""
+"""PixelDiT T2I export helpers: upstream export patches, weights loader and sampling pipeline.
+
+The model itself is upstream's PixDiT_T2I from the `pixeldit` pip package."""