From 05ec646205d06ae4ef75420fafe3f867bed3c573 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Thu, 1 Oct 2026 10:21:20 -0700 Subject: [PATCH 01/12] add python export --- image_gen/pixeldit/README.md | 44 ++ .../pixeldit/model_export/export_pixeldit.py | 415 ++++++++++++ .../model_export/pixeldit/__init__.py | 3 + .../pixeldit/model_export/pixeldit/model.py | 441 +++++++++++++ .../model_export/pixeldit/pipeline.py | 205 ++++++ .../pixeldit/model_export/pixeldit/weights.py | 88 +++ .../pixeldit/model_export/verify_pixeldit.py | 593 ++++++++++++++++++ 7 files changed, 1789 insertions(+) create mode 100644 image_gen/pixeldit/README.md create mode 100644 image_gen/pixeldit/model_export/export_pixeldit.py create mode 100644 image_gen/pixeldit/model_export/pixeldit/__init__.py create mode 100644 image_gen/pixeldit/model_export/pixeldit/model.py create mode 100644 image_gen/pixeldit/model_export/pixeldit/pipeline.py create mode 100644 image_gen/pixeldit/model_export/pixeldit/weights.py create mode 100644 image_gen/pixeldit/model_export/verify_pixeldit.py diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md new file mode 100644 index 0000000..4f3f5be --- /dev/null +++ b/image_gen/pixeldit/README.md @@ -0,0 +1,44 @@ +# PixelDiT + +ONNX export and verification for PixelDiT text-to-image with ONNX Runtime. There is no C++ runtime yet. + +## 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 | 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`. + +## 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 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, or `--reference_repo ` to also compare the vendored model and sampler with upstream PixelDiT. + +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. 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..26ebd23 --- /dev/null +++ b/image_gen/pixeldit/model_export/export_pixeldit.py @@ -0,0 +1,415 @@ +# 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) + export_all = (args.model is None) or ("all" in raw) + + 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 export_all: + 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..8b1a239 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/model.py @@ -0,0 +1,441 @@ +# 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, with three +export-friendly changes: + * RoPE uses real cos/sin tables instead of complex tensors (torch.polar / view_as_complex). + * Patchify / unpatchify use reshape + permute instead of F.unfold / F.fold. + * Position tables are built once by set_image_size() instead of lazily filled Python dict + caches, so tracing never writes fake tensors into a cache and they export as graph + constants. 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..cb14cb3 --- /dev/null +++ b/image_gen/pixeldit/model_export/verify_pixeldit.py @@ -0,0 +1,593 @@ +# 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 against the vendored PyTorch model, and optionally verify the +vendored PyTorch model + sampler against the upstream NVlabs/PixelDiT repo. + +Checks (see .claude/plans/pixeldit-export-plan.md, "Validation plan"): + --validate_only V5: ONNX structure + dummy run of each graph, then exit + (default) V5, V6: per-model ONNX Runtime vs PyTorch, V7: end-to-end image comparison + --reference_repo V2-V4: vendored model / text path / sampler vs upstream PixelDiT + +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 + upstream_output.png - image from upstream PixelDiT (--reference_repo only) + +Pass/fail is decided by the per-model checks (V2, V3, V6). End-to-end images only fail 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, + resolve_checkpoint, +) + +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." +) +BANNED_OPS = {"Im2Col", "Col2Im", "DFT", "STFT"} + +# All comparisons are bf16 vs bf16. Thresholds are ~2x the bf16-vs-fp32 noise measured once on the +# 1024x1024 export (see .claude/plans/pixeldit-export-plan.md, "Findings"); bf16 rounding of the +# timestep makes the DiT far noisier near t*1000 = 1000 than elsewhere. +DIT_MAX_REL_L2 = {999.75: 0.55, 996.0: 0.025, 500.0: 0.025, 50.0: 0.025} # DiT timestep input -> max rel L2 +CHECK_TIMESTEPS = (999.75, 500.0, 50.0) +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 (seed 1 even changes the pose). The image checks only catch broken output; the +# per-model checks are the real gate. 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 reference (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(), + "bitwise": torch.equal(a, b), + } + + +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) + 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: + return np.ascontiguousarray(tensor.detach().to("cpu", torch.float32 if dtype == np.float32 else tensor.dtype).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 + + +# --------------------------------------------------------------------------- +# V5: 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]: + """V5 checks; returns the ORT sessions so later checks don't rebuild TensorRT engines.""" + print(f"\n[V5 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) + op_types = {node.op_type for node in model.graph.node} + for function in model.functions: + op_types.update(node.op_type for node in function.node) + 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})") + banned = sorted(op_types & BANNED_OPS) + _check(f"{name} no Im2Col/Col2Im/complex ops", not banned, f"found {banned}" if banned else "none") + 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 + + +# --------------------------------------------------------------------------- +# V6: per-model ONNX Runtime vs PyTorch +# --------------------------------------------------------------------------- + + +@torch.inference_mode() +def compare_text_encoder(tokenizer, text_encoder, runner, layout, prompt, negative_prompt, device): + print("\n[V6] 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})") + + +def dit_check(name: str, tv: float, out: torch.Tensor, ref: torch.Tensor) -> None: + s = _stats(out, ref) + limit = DIT_MAX_REL_L2[tv] + _check(f"{name} t={tv}", s["rel_l2"] <= limit, f"{_fmt(s)} (need rel_l2 <= {limit})") + + +@torch.inference_mode() +def compare_dit(dit, runner, y_neg, y_pos, height, width, device): + print("\n[V6] 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 in CHECK_TIMESTEPS: + t = torch.full((2,), tv, device=device) + dit_check("transformer velocity", tv, ort_dit(x, t, y), dit(x, t, y)) + + +# --------------------------------------------------------------------------- +# V7: 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 check_image_sanity(name: str, image: np.ndarray, ref: np.ndarray) -> None: + """Fail only on clearly broken output (black, noise, flat); otherwise report metrics for visual review.""" + psnr = _psnr(ref, image) + mae = float(np.mean(np.abs(ref.astype(np.float32) - image.astype(np.float32)))) + ratio = _detail(image) / max(_detail(ref), 1e-6) + lo, hi = DETAIL_RATIO_RANGE + _check(f"{name} not broken: PSNR", psnr >= MIN_E2E_PSNR, f"PSNR={psnr:.2f} dB (need >= {MIN_E2E_PSNR} dB)") + _check(f"{name} not broken: detail", lo <= ratio <= hi, f"detail ratio={ratio:.2f} (need {lo}-{hi})") + print(f" [INFO] {name}: 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.") + + +def compare_images(ref: np.ndarray, ort: np.ndarray, output_dir: Path) -> None: + check_image_sanity("end-to-end image", ort, ref) + 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 + + +# --------------------------------------------------------------------------- +# V2-V4: vendored code vs upstream PixelDiT (--reference_repo) +# --------------------------------------------------------------------------- + + +def run_reference_checks(args, tokenizer, text_encoder, layout, dit, height, width, device, pytorch_image, output_dir) -> None: + """pytorch_image() returns the vendored-pipeline image shared with V7 (generated once).""" + repo = Path(args.reference_repo).resolve() + sys.path[:0] = [str(repo), str(repo / "t2i")] + import yaml + from diffusion.model.dpm_solver import DPM_Solver, NoiseScheduleFlow + from diffusion.model.flow_dpm import DPMS + from diffusion.model.trainer import PixDiTTrainer + from pixdit_core import modules as upstream_modules + + from pixeldit.model import apply_rotary_emb, pixel_sincos_pos_embed, rope_2d_angles + + print(f"\n[Reference] Upstream PixelDiT at {repo}") + cfg = yaml.safe_load(open(repo / "t2i" / "configs" / "PixelDiT_1024px_pixel_diffusion_stage3.yaml", encoding="utf-8")) + _check("chi_prompt matches upstream config", "\n".join(cfg["text_encoder"]["chi_prompt"]) == layout.chi_prompt, "") + + # V2 building blocks + q, k = torch.randn(2, 4096, 24, 64), torch.randn(2, 4096, 24, 64) + uq, uk = upstream_modules.apply_rotary_emb(q, k, upstream_modules.precompute_freqs_cis_2d(64, 64, 64)) + ang = rope_2d_angles(64, 64, 64) + mq, mk = apply_rotary_emb(q, k, ang.cos(), ang.sin()) + err = max(_stats(mq, uq)["max_abs"], _stats(mk, uk)["max_abs"]) + _check("V2 real-valued RoPE vs complex", err < 1e-6, f"max_abs={err:.3e}") + x = torch.randn(2, 3, height, width) + hs, ws = height // 16, width // 16 + patches = x.view(2, 3, hs, 16, ws, 16).permute(0, 2, 4, 1, 3, 5).reshape(2, hs * ws, -1) + _check("V2 patchify vs F.unfold", torch.equal(patches, F.unfold(x, 16, stride=16).transpose(1, 2)), "bitwise") + xp = torch.randn(2, hs * ws, 256, 3) + folded = F.fold(xp.permute(0, 3, 2, 1).reshape(2, 768, hs * ws), (height, width), kernel_size=16, stride=16) + unpatched = xp.view(2, hs, ws, 16, 16, 3).permute(0, 5, 1, 3, 2, 4).reshape(2, 3, height, width) + _check("V2 unpatchify vs F.fold", torch.equal(unpatched, folded), "bitwise") + if height == width: + up_pos = upstream_modules.get_2d_sincos_pos_embed(16, height) + _check("V2 pixel pos embed", np.array_equal(up_pos, pixel_sincos_pos_embed(16, height, width)), "bitwise") + + # V2 full forward: upstream PixDiTTrainer vs vendored, both bf16 + checkpoint, _ = resolve_checkpoint(args.model_name, args.local_files_only) + full_state = torch.load(str(checkpoint), map_location="cpu", mmap=True, weights_only=True)["state_dict"] + upstream = PixDiTTrainer(image_size=1024, extra=cfg["model"]["extra"], caption_channels=2304, model_max_length=300) + upstream.load_state_dict(full_state, strict=True) + upstream = upstream.to(device).eval().requires_grad_(False).to(torch.bfloat16) + g = torch.Generator(device=device).manual_seed(1) + xb = torch.randn(2, 3, height, width, device=device, generator=g) + yb = torch.randn(2, 300, 2304, device=device, generator=g) + with torch.inference_mode(): + for tv in CHECK_TIMESTEPS + (996.0,): + t = torch.full((2,), tv, device=device) + dit_check("V2 DiT vendored vs upstream", tv, dit(xb, t, yb), upstream.forward_with_dpmsolver(xb, t, yb)) + + # V3 text path + prompt, negative_prompt = args.prompt, args.negative_prompt + with torch.inference_mode(): + sel = [0] + list(range(-300 + 1, 0)) + up_tok = tokenizer([layout.chi_prompt + prompt.strip()], max_length=layout.text_seq_len, padding="max_length", + truncation=True, return_tensors="pt").to(device) + up_pos = text_encoder(up_tok.input_ids, up_tok.attention_mask)[0][:, sel] + up_neg_tok = tokenizer(negative_prompt, max_length=300, padding="max_length", truncation=True, return_tensors="pt").to(device) + up_neg = text_encoder(up_neg_tok.input_ids, up_neg_tok.attention_mask)[0] + my_neg, my_pos = P.encode_prompts(tokenizer, torch_encode_fn(text_encoder), layout, prompt, negative_prompt, device) + s_pos, s_neg = _stats(my_pos, up_pos), _stats(my_neg, up_neg) + _check("V3 positive embeds vs upstream", s_pos["bitwise"], _fmt(s_pos)) + # The 506-token trick is exact in fp32 (rel L2 3.8e-6, measured once); in bf16 it differs only by Gemma's bf16 noise. + _check("V3 negative embeds 506 + select vs upstream 300", s_neg["cos"] >= TEXT_MIN_COS, + f"{_fmt(s_neg)} (need cos >= {TEXT_MIN_COS})") + + # V4 sampler + ns = NoiseScheduleFlow(schedule="discrete_flow") + for steps in (2, 20, args.steps): + up_t = DPM_Solver(lambda x_, t_: x_, ns).get_time_steps("time_uniform_flow", 1, 1.0 / 1000, steps, device, shift=args.flow_shift) + _check(f"V4 timesteps steps={steps}", torch.equal(P.flow_timesteps(steps, args.flow_shift, device), up_t), "bitwise") + + def fake_model(x_, t_, y_, **_): + return torch.sin(3 * x_) * (t_ / 1000).view(-1, 1, 1, 1) + y_.reshape(y_.shape[0], -1).mean(1).view(-1, 1, 1, 1) - 0.3 * x_ + + z = torch.randn(1, 3, 64, 64, device=device, generator=g, dtype=torch.float64) + yc = torch.randn(1, 1, 8, 4, device=device, dtype=torch.float64) + yu = torch.randn(1, 1, 8, 4, device=device, dtype=torch.float64) + up = DPMS(fake_model, condition=yc, uncondition=yu, cfg_scale=args.cfg_scale, model_type="flow", model_kwargs={}, + schedule="FLOW", interval_guidance=[0, 1]).sample( + z, steps=args.steps, order=2, skip_type="time_uniform_flow", method="multistep", flow_shift=args.flow_shift) + mine = P.FlowDPMSolver(fake_model, yu[:, 0], yc[:, 0], args.cfg_scale).sample(z, args.steps, args.flow_shift) + s = _stats(mine, up) + _check("V4 solver math (fake model, fp64)", s["max_abs"] < 1e-10 and torch.isfinite(mine).all().item(), _fmt(s)) + + # V4 end to end: upstream DPMS + PixDiTTrainer (bf16) vs vendored pipeline + print(" [V4] end-to-end upstream sample...") + with torch.inference_mode(): + z = P.initial_noise(height, width, args.seed, device) + hw = torch.tensor([[float(height), float(width)]], device=device) + kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": (hw[:, 0] / hw[:, 1]).unsqueeze(1)}, mask=up_tok.attention_mask[:, sel]) + up_sample = DPMS(upstream.forward_with_dpmsolver, condition=up_pos[:, None], uncondition=up_neg[:, None], + guidance_type="classifier-free", cfg_scale=args.cfg_scale, model_type="flow", model_kwargs=kwargs, + schedule="FLOW", interval_guidance=[0, 1]).sample( + z, steps=args.steps, order=2, skip_type="time_uniform_flow", method="multistep", flow_shift=args.flow_shift) + del upstream + torch.cuda.empty_cache() + upstream_image = P.to_uint8_image(up_sample) + PILImage.fromarray(upstream_image).save(str(output_dir / "upstream_output.png")) + print(f" [V4] Upstream image saved -> {output_dir / 'upstream_output.png'}") + check_image_sanity("V4 end-to-end vendored vs upstream", pytorch_image(), upstream_image) + + +# --------------------------------------------------------------------------- +# 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 (V7)") + p.add_argument("--reference_repo", type=str, default=None, help="Path to an NVlabs/PixelDiT clone for V2-V4 parity checks") + 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) + + sampler = P.SamplerConfig(args.steps, args.cfg_scale, args.flow_shift, args.negative_prompt) + encode_torch = torch_encode_fn(text_encoder) + cache = {} + + def pytorch_image() -> np.ndarray: + """The PyTorch-pipeline image, generated once and shared by V4 and V7.""" + if "image" not in cache: + print("[PyTorch] Running pipeline...") + cache["image"] = P.to_uint8_image(P.generate(tokenizer, encode_torch, dit, layout, args.prompt, height, width, + args.seed, sampler, device, _progress("PyTorch"))) + PILImage.fromarray(cache["image"]).save(str(output_dir / "pytorch_output.png")) + print(f"[PyTorch] Image saved -> {output_dir / 'pytorch_output.png'}") + return cache["image"] + + if args.reference_repo: + run_reference_checks(args, tokenizer, text_encoder, layout, dit, height, width, device, pytorch_image, output_dir) + + 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[V7] End to end") + pt_image = pytorch_image() + 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() From cb8fa2ba1af72ba06faf1babb5aaa6c9a261711f Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Thu, 1 Oct 2026 13:24:45 -0700 Subject: [PATCH 02/12] fix validation --- image_gen/pixeldit/README.md | 2 +- .../pixeldit/model_export/pixeldit/model.py | 18 +- .../pixeldit/model_export/verify_pixeldit.py | 232 ++++-------------- 3 files changed, 57 insertions(+), 195 deletions(-) diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md index 4f3f5be..a21799a 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -39,6 +39,6 @@ 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 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, or `--reference_repo ` to also compare the vendored model and sampler with upstream PixelDiT. +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. diff --git a/image_gen/pixeldit/model_export/pixeldit/model.py b/image_gen/pixeldit/model_export/pixeldit/model.py index 8b1a239..1aa1006 100644 --- a/image_gen/pixeldit/model_export/pixeldit/model.py +++ b/image_gen/pixeldit/model_export/pixeldit/model.py @@ -2,14 +2,18 @@ # 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, with three -export-friendly changes: - * RoPE uses real cos/sin tables instead of complex tensors (torch.polar / view_as_complex). - * Patchify / unpatchify use reshape + permute instead of F.unfold / F.fold. +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 fake tensors into a cache and they export as graph - constants. They are plain attributes (not buffers) so Module.to(dtype) leaves the fp32 - RoPE tables alone, as upstream's complex64 tables are. + 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. """ diff --git a/image_gen/pixeldit/model_export/verify_pixeldit.py b/image_gen/pixeldit/model_export/verify_pixeldit.py index cb14cb3..37c1d8a 100644 --- a/image_gen/pixeldit/model_export/verify_pixeldit.py +++ b/image_gen/pixeldit/model_export/verify_pixeldit.py @@ -3,23 +3,21 @@ #!/usr/bin/env python3 """ -Verify PixelDiT T2I ONNX graphs against the vendored PyTorch model, and optionally verify the -vendored PyTorch model + sampler against the upstream NVlabs/PixelDiT repo. +Verify PixelDiT T2I ONNX graphs: run the vendored PyTorch pipeline and the ONNX Runtime pipeline +(TensorRT RTX by default) and compare them. -Checks (see .claude/plans/pixeldit-export-plan.md, "Validation plan"): - --validate_only V5: ONNX structure + dummy run of each graph, then exit - (default) V5, V6: per-model ONNX Runtime vs PyTorch, V7: end-to-end image comparison - --reference_repo V2-V4: vendored model / text path / sampler vs upstream PixelDiT +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 - upstream_output.png - image from upstream PixelDiT (--reference_repo only) -Pass/fail is decided by the per-model checks (V2, V3, V6). End-to-end images only fail 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. +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 @@ -36,13 +34,7 @@ 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, - resolve_checkpoint, -) +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 @@ -51,19 +43,17 @@ "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." ) -BANNED_OPS = {"Im2Col", "Col2Im", "DFT", "STFT"} -# All comparisons are bf16 vs bf16. Thresholds are ~2x the bf16-vs-fp32 noise measured once on the -# 1024x1024 export (see .claude/plans/pixeldit-export-plan.md, "Findings"); bf16 rounding of the -# timestep makes the DiT far noisier near t*1000 = 1000 than elsewhere. -DIT_MAX_REL_L2 = {999.75: 0.55, 996.0: 0.025, 500.0: 0.025, 50.0: 0.025} # DiT timestep input -> max rel L2 -CHECK_TIMESTEPS = (999.75, 500.0, 50.0) +# 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 (seed 1 even changes the pose). The image checks only catch broken output; the -# per-model checks are the real gate. 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 reference (correct runs measured 0.72-1.01x). +# 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) @@ -89,7 +79,6 @@ def _stats(a: torch.Tensor, b: torch.Tensor) -> dict[str, float]: "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(), - "bitwise": torch.equal(a, b), } @@ -179,7 +168,8 @@ def run(self, feed: dict[str, np.ndarray]) -> list[np.ndarray]: def _to_numpy(tensor: torch.Tensor, dtype: np.dtype) -> np.ndarray: - return np.ascontiguousarray(tensor.detach().to("cpu", torch.float32 if dtype == np.float32 else tensor.dtype).numpy().astype(dtype, copy=False)) + 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): @@ -214,7 +204,7 @@ def encode(input_ids, attention_mask, select_index): # --------------------------------------------------------------------------- -# V5: graph structure + dummy run +# Graph structure + dummy run # --------------------------------------------------------------------------- EXPECTED_IO = { @@ -233,8 +223,8 @@ def _onnx_io_shapes(model: onnx.ModelProto) -> dict[str, tuple[int, ...]]: def validate_onnx_dir(onnx_dir: Path, provider: str, args, expected_shapes: dict) -> dict[str, OrtSessionRunner]: - """V5 checks; returns the ORT sessions so later checks don't rebuild TensorRT engines.""" - print(f"\n[V5 Validate] {onnx_dir}") + """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" @@ -244,15 +234,10 @@ def validate_onnx_dir(onnx_dir: Path, provider: str, args, expected_shapes: dict onnx.checker.check_model(str(onnx_path)) model = onnx.load(str(onnx_path), load_external_data=False) shapes = _onnx_io_shapes(model) - op_types = {node.op_type for node in model.graph.node} - for function in model.functions: - op_types.update(node.op_type for node in function.node) 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})") - banned = sorted(op_types & BANNED_OPS) - _check(f"{name} no Im2Col/Col2Im/complex ops", not banned, f"found {banned}" if banned else "none") 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 @@ -283,13 +268,13 @@ def validate_onnx_dir(onnx_dir: Path, provider: str, args, expected_shapes: dict # --------------------------------------------------------------------------- -# V6: per-model ONNX Runtime vs PyTorch +# Per-model ONNX Runtime vs PyTorch # --------------------------------------------------------------------------- @torch.inference_mode() def compare_text_encoder(tokenizer, text_encoder, runner, layout, prompt, negative_prompt, device): - print("\n[V6] text_encoder: ONNX Runtime vs PyTorch (bf16)") + 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)): @@ -299,26 +284,21 @@ def compare_text_encoder(tokenizer, text_encoder, runner, layout, prompt, negati f"{_fmt(s)} min_token_cos={per_token:.6f} (need cos >= {TEXT_MIN_COS})") -def dit_check(name: str, tv: float, out: torch.Tensor, ref: torch.Tensor) -> None: - s = _stats(out, ref) - limit = DIT_MAX_REL_L2[tv] - _check(f"{name} t={tv}", s["rel_l2"] <= limit, f"{_fmt(s)} (need rel_l2 <= {limit})") - - @torch.inference_mode() def compare_dit(dit, runner, y_neg, y_pos, height, width, device): - print("\n[V6] transformer: ONNX Runtime vs PyTorch (bf16)") + 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 in CHECK_TIMESTEPS: + for tv, limit in DIT_MAX_REL_L2.items(): t = torch.full((2,), tv, device=device) - dit_check("transformer velocity", tv, ort_dit(x, t, y), dit(x, t, y)) + 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})") # --------------------------------------------------------------------------- -# V7: end to end +# End to end # --------------------------------------------------------------------------- @@ -328,20 +308,16 @@ def _detail(image: np.ndarray) -> float: 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 check_image_sanity(name: str, image: np.ndarray, ref: np.ndarray) -> None: +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, image) - mae = float(np.mean(np.abs(ref.astype(np.float32) - image.astype(np.float32)))) - ratio = _detail(image) / max(_detail(ref), 1e-6) + 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(f"{name} not broken: PSNR", psnr >= MIN_E2E_PSNR, f"PSNR={psnr:.2f} dB (need >= {MIN_E2E_PSNR} dB)") - _check(f"{name} not broken: detail", lo <= ratio <= hi, f"detail ratio={ratio:.2f} (need {lo}-{hi})") - print(f" [INFO] {name}: PSNR={psnr:.2f} dB MAE={mae:.2f}. Fine details can differ between correct " + _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.") - - -def compare_images(ref: np.ndarray, ort: np.ndarray, output_dir: Path) -> None: - check_image_sanity("end-to-end image", ort, ref) 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)) @@ -358,113 +334,6 @@ def report(step): return report -# --------------------------------------------------------------------------- -# V2-V4: vendored code vs upstream PixelDiT (--reference_repo) -# --------------------------------------------------------------------------- - - -def run_reference_checks(args, tokenizer, text_encoder, layout, dit, height, width, device, pytorch_image, output_dir) -> None: - """pytorch_image() returns the vendored-pipeline image shared with V7 (generated once).""" - repo = Path(args.reference_repo).resolve() - sys.path[:0] = [str(repo), str(repo / "t2i")] - import yaml - from diffusion.model.dpm_solver import DPM_Solver, NoiseScheduleFlow - from diffusion.model.flow_dpm import DPMS - from diffusion.model.trainer import PixDiTTrainer - from pixdit_core import modules as upstream_modules - - from pixeldit.model import apply_rotary_emb, pixel_sincos_pos_embed, rope_2d_angles - - print(f"\n[Reference] Upstream PixelDiT at {repo}") - cfg = yaml.safe_load(open(repo / "t2i" / "configs" / "PixelDiT_1024px_pixel_diffusion_stage3.yaml", encoding="utf-8")) - _check("chi_prompt matches upstream config", "\n".join(cfg["text_encoder"]["chi_prompt"]) == layout.chi_prompt, "") - - # V2 building blocks - q, k = torch.randn(2, 4096, 24, 64), torch.randn(2, 4096, 24, 64) - uq, uk = upstream_modules.apply_rotary_emb(q, k, upstream_modules.precompute_freqs_cis_2d(64, 64, 64)) - ang = rope_2d_angles(64, 64, 64) - mq, mk = apply_rotary_emb(q, k, ang.cos(), ang.sin()) - err = max(_stats(mq, uq)["max_abs"], _stats(mk, uk)["max_abs"]) - _check("V2 real-valued RoPE vs complex", err < 1e-6, f"max_abs={err:.3e}") - x = torch.randn(2, 3, height, width) - hs, ws = height // 16, width // 16 - patches = x.view(2, 3, hs, 16, ws, 16).permute(0, 2, 4, 1, 3, 5).reshape(2, hs * ws, -1) - _check("V2 patchify vs F.unfold", torch.equal(patches, F.unfold(x, 16, stride=16).transpose(1, 2)), "bitwise") - xp = torch.randn(2, hs * ws, 256, 3) - folded = F.fold(xp.permute(0, 3, 2, 1).reshape(2, 768, hs * ws), (height, width), kernel_size=16, stride=16) - unpatched = xp.view(2, hs, ws, 16, 16, 3).permute(0, 5, 1, 3, 2, 4).reshape(2, 3, height, width) - _check("V2 unpatchify vs F.fold", torch.equal(unpatched, folded), "bitwise") - if height == width: - up_pos = upstream_modules.get_2d_sincos_pos_embed(16, height) - _check("V2 pixel pos embed", np.array_equal(up_pos, pixel_sincos_pos_embed(16, height, width)), "bitwise") - - # V2 full forward: upstream PixDiTTrainer vs vendored, both bf16 - checkpoint, _ = resolve_checkpoint(args.model_name, args.local_files_only) - full_state = torch.load(str(checkpoint), map_location="cpu", mmap=True, weights_only=True)["state_dict"] - upstream = PixDiTTrainer(image_size=1024, extra=cfg["model"]["extra"], caption_channels=2304, model_max_length=300) - upstream.load_state_dict(full_state, strict=True) - upstream = upstream.to(device).eval().requires_grad_(False).to(torch.bfloat16) - g = torch.Generator(device=device).manual_seed(1) - xb = torch.randn(2, 3, height, width, device=device, generator=g) - yb = torch.randn(2, 300, 2304, device=device, generator=g) - with torch.inference_mode(): - for tv in CHECK_TIMESTEPS + (996.0,): - t = torch.full((2,), tv, device=device) - dit_check("V2 DiT vendored vs upstream", tv, dit(xb, t, yb), upstream.forward_with_dpmsolver(xb, t, yb)) - - # V3 text path - prompt, negative_prompt = args.prompt, args.negative_prompt - with torch.inference_mode(): - sel = [0] + list(range(-300 + 1, 0)) - up_tok = tokenizer([layout.chi_prompt + prompt.strip()], max_length=layout.text_seq_len, padding="max_length", - truncation=True, return_tensors="pt").to(device) - up_pos = text_encoder(up_tok.input_ids, up_tok.attention_mask)[0][:, sel] - up_neg_tok = tokenizer(negative_prompt, max_length=300, padding="max_length", truncation=True, return_tensors="pt").to(device) - up_neg = text_encoder(up_neg_tok.input_ids, up_neg_tok.attention_mask)[0] - my_neg, my_pos = P.encode_prompts(tokenizer, torch_encode_fn(text_encoder), layout, prompt, negative_prompt, device) - s_pos, s_neg = _stats(my_pos, up_pos), _stats(my_neg, up_neg) - _check("V3 positive embeds vs upstream", s_pos["bitwise"], _fmt(s_pos)) - # The 506-token trick is exact in fp32 (rel L2 3.8e-6, measured once); in bf16 it differs only by Gemma's bf16 noise. - _check("V3 negative embeds 506 + select vs upstream 300", s_neg["cos"] >= TEXT_MIN_COS, - f"{_fmt(s_neg)} (need cos >= {TEXT_MIN_COS})") - - # V4 sampler - ns = NoiseScheduleFlow(schedule="discrete_flow") - for steps in (2, 20, args.steps): - up_t = DPM_Solver(lambda x_, t_: x_, ns).get_time_steps("time_uniform_flow", 1, 1.0 / 1000, steps, device, shift=args.flow_shift) - _check(f"V4 timesteps steps={steps}", torch.equal(P.flow_timesteps(steps, args.flow_shift, device), up_t), "bitwise") - - def fake_model(x_, t_, y_, **_): - return torch.sin(3 * x_) * (t_ / 1000).view(-1, 1, 1, 1) + y_.reshape(y_.shape[0], -1).mean(1).view(-1, 1, 1, 1) - 0.3 * x_ - - z = torch.randn(1, 3, 64, 64, device=device, generator=g, dtype=torch.float64) - yc = torch.randn(1, 1, 8, 4, device=device, dtype=torch.float64) - yu = torch.randn(1, 1, 8, 4, device=device, dtype=torch.float64) - up = DPMS(fake_model, condition=yc, uncondition=yu, cfg_scale=args.cfg_scale, model_type="flow", model_kwargs={}, - schedule="FLOW", interval_guidance=[0, 1]).sample( - z, steps=args.steps, order=2, skip_type="time_uniform_flow", method="multistep", flow_shift=args.flow_shift) - mine = P.FlowDPMSolver(fake_model, yu[:, 0], yc[:, 0], args.cfg_scale).sample(z, args.steps, args.flow_shift) - s = _stats(mine, up) - _check("V4 solver math (fake model, fp64)", s["max_abs"] < 1e-10 and torch.isfinite(mine).all().item(), _fmt(s)) - - # V4 end to end: upstream DPMS + PixDiTTrainer (bf16) vs vendored pipeline - print(" [V4] end-to-end upstream sample...") - with torch.inference_mode(): - z = P.initial_noise(height, width, args.seed, device) - hw = torch.tensor([[float(height), float(width)]], device=device) - kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": (hw[:, 0] / hw[:, 1]).unsqueeze(1)}, mask=up_tok.attention_mask[:, sel]) - up_sample = DPMS(upstream.forward_with_dpmsolver, condition=up_pos[:, None], uncondition=up_neg[:, None], - guidance_type="classifier-free", cfg_scale=args.cfg_scale, model_type="flow", model_kwargs=kwargs, - schedule="FLOW", interval_guidance=[0, 1]).sample( - z, steps=args.steps, order=2, skip_type="time_uniform_flow", method="multistep", flow_shift=args.flow_shift) - del upstream - torch.cuda.empty_cache() - upstream_image = P.to_uint8_image(up_sample) - PILImage.fromarray(upstream_image).save(str(output_dir / "upstream_output.png")) - print(f" [V4] Upstream image saved -> {output_dir / 'upstream_output.png'}") - check_image_sanity("V4 end-to-end vendored vs upstream", pytorch_image(), upstream_image) - - # --------------------------------------------------------------------------- # main # --------------------------------------------------------------------------- @@ -498,8 +367,7 @@ def main(): p.add_argument("--ep_dll_dir", default=None, help="Dir with the plugin EP's bundled runtime DLLs") p.add_argument("--trt_bin", default=None, help="TensorRT-RTX bin dir, used as a DLL search path") p.add_argument("--validate_only", action="store_true", help="Only run ONNX structure checks and a dummy run") - p.add_argument("--skip_e2e", action="store_true", help="Skip the end-to-end image comparison (V7)") - p.add_argument("--reference_repo", type=str, default=None, help="Path to an NVlabs/PixelDiT clone for V2-V4 parity checks") + p.add_argument("--skip_e2e", action="store_true", help="Skip the end-to-end image comparison") args = p.parse_args() configure_stdio() @@ -542,23 +410,7 @@ def main(): dit = load_pixeldit(args.model_name, device, torch.bfloat16, args.local_files_only) dit.set_image_size(height, width) - sampler = P.SamplerConfig(args.steps, args.cfg_scale, args.flow_shift, args.negative_prompt) encode_torch = torch_encode_fn(text_encoder) - cache = {} - - def pytorch_image() -> np.ndarray: - """The PyTorch-pipeline image, generated once and shared by V4 and V7.""" - if "image" not in cache: - print("[PyTorch] Running pipeline...") - cache["image"] = P.to_uint8_image(P.generate(tokenizer, encode_torch, dit, layout, args.prompt, height, width, - args.seed, sampler, device, _progress("PyTorch"))) - PILImage.fromarray(cache["image"]).save(str(output_dir / "pytorch_output.png")) - print(f"[PyTorch] Image saved -> {output_dir / 'pytorch_output.png'}") - return cache["image"] - - if args.reference_repo: - run_reference_checks(args, tokenizer, text_encoder, layout, dit, height, width, device, pytorch_image, output_dir) - 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) @@ -568,8 +420,14 @@ def pytorch_image() -> np.ndarray: compare_dit(dit, tr_runner, y_neg, y_pos, height, width, device) if not args.skip_e2e: - print("\n[V7] End to end") - pt_image = pytorch_image() + 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 From d5e8a428b8f42c2216f6559ad5f24bbdfe4e3211 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Fri, 2 Oct 2026 16:21:11 -0700 Subject: [PATCH 03/12] add trt cpu backend --- CMakeLists.txt | 11 + common/io/tokenizer.cpp | 300 +++++++- common/io/tokenizer.h | 31 + image_gen/pixeldit/CMakeLists.txt | 62 ++ image_gen/pixeldit/README.md | 27 +- image_gen/pixeldit/helpers/pixeldit.h | 343 +++++++++ image_gen/pixeldit/pixeldit_cli.cpp | 283 ++++++++ image_gen/pixeldit/pixeldit_cli.h | 65 ++ image_gen/pixeldit/pixeldit_config.cpp | 109 +++ image_gen/pixeldit/pixeldit_pipeline.cpp | 46 ++ image_gen/pixeldit/pixeldit_runtime_context.h | 26 + image_gen/pixeldit/pixeldit_trt_backend.cpp | 670 ++++++++++++++++++ 12 files changed, 1966 insertions(+), 7 deletions(-) create mode 100644 image_gen/pixeldit/CMakeLists.txt create mode 100644 image_gen/pixeldit/helpers/pixeldit.h create mode 100644 image_gen/pixeldit/pixeldit_cli.cpp create mode 100644 image_gen/pixeldit/pixeldit_cli.h create mode 100644 image_gen/pixeldit/pixeldit_config.cpp create mode 100644 image_gen/pixeldit/pixeldit_pipeline.cpp create mode 100644 image_gen/pixeldit/pixeldit_runtime_context.h create mode 100644 image_gen/pixeldit/pixeldit_trt_backend.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index c70acfc..391f954 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -97,6 +97,7 @@ endfunction() add_subdirectory(common) add_subdirectory(tests/base) add_subdirectory(image_gen/flux2) +add_subdirectory(image_gen/pixeldit) add_subdirectory(asr/rnnt) add_subdirectory(vision/sam2) add_subdirectory(asr/whisper) @@ -214,5 +215,15 @@ if(DIN_BUILD_TESTING) ) endif () + add_din_cli_test(din_pixeldit_bf16_trt-rtx_cpu din_pixeldit_cli + --model-dir "${DIN_DEPLOY_ARTIFACTS_ROOT}/PixelDiT-1300M-1024px-onnx" + --provider trt-rtx + --processing cpu + --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_cpu" + ) + unset(_din_cli_test_audio) endif() diff --git a/common/io/tokenizer.cpp b/common/io/tokenizer.cpp index 86dcc97..109a2a0 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,22 @@ 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)) + { + 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 +278,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 +452,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 +705,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 +881,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..0b5936d --- /dev/null +++ b/image_gen/pixeldit/CMakeLists.txt @@ -0,0 +1,62 @@ +find_package(CUDAToolkit REQUIRED) + +# The host sampling math mirrors the PyTorch reference operation by operation in fp32; keep the compiler from +# fusing multiply-adds so results stay comparable (MSVC /fp:precise does not contract by default). +set(PIXELDIT_FP_FLAGS "$<$,$>:-ffp-contract=off>") + +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 +) +set_target_properties(din_pixeldit_trt PROPERTIES + CXX_STANDARD 20 + CXX_EXTENSIONS OFF +) +target_compile_options(din_pixeldit_trt PRIVATE ${PIXELDIT_FP_FLAGS}) +target_link_libraries(din_pixeldit_trt + PUBLIC + din_pixeldit_common + PRIVATE + CUDA::cudart +) + +add_library(din_pixeldit_runtime STATIC + pixeldit_pipeline.cpp +) +set_target_properties(din_pixeldit_runtime PROPERTIES + CXX_STANDARD 20 + CXX_EXTENSIONS OFF +) +target_link_libraries(din_pixeldit_runtime PUBLIC din_pixeldit_trt) + +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_runtime) diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md index a21799a..3f0e176 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -1,6 +1,6 @@ # PixelDiT -ONNX export and verification for PixelDiT text-to-image with ONNX Runtime. There is no C++ runtime yet. +CLI image generation with ONNX Runtime. ## Supported models @@ -20,7 +20,15 @@ PixelDiT generates pixels directly, so there is no VAE. The model code is vendor | DiT | `transformer/model.onnx` | Batch of 2: negative and positive prompt for classifier-free guidance. Returns velocity. | | Sampling | Host | 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`. +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 | +| --- | --- | --- | +| Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | + +CUDA, DirectX, and Vulkan processing are not implemented yet. `--provider cpu` is not supported. ## Export @@ -42,3 +50,18 @@ python verify_pixeldit.py --onnx_dir ./models/PixelDiT-1300M-1024px-onnx --provi 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 cpu --prompt "a red fox" --output out +``` + +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/helpers/pixeldit.h b/image_gen/pixeldit/helpers/pixeldit.h new file mode 100644 index 0000000..78092ef --- /dev/null +++ b/image_gen/pixeldit/helpers/pixeldit.h @@ -0,0 +1,343 @@ +// 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 +// ============================================================================ + +#ifdef _WIN32 +static const std::filesystem::path DEFAULT_PIXELDIT_MODEL_DIR = "S:/din_deploy_artifacts/PixelDiT-1300M-1024px-onnx"; +#else +static const std::filesystem::path DEFAULT_PIXELDIT_MODEL_DIR = "/mnt/share/onnx/PixelDiT-1300M-1024px-onnx/"; +#endif + +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; + } + }); +} + +// DPM-Solver-1 (dpmsolver++) from s to t: x_t = (t / s) * x - (alpha_t * phi_1) * x0_s, phi_1 = expm1(-h). +// 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 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); + const float ratio = t / s; + const float coeff = alpha_t * phi_1; + ParallelFor(n, + [=](size_t begin, size_t end) + { + for (size_t i = begin; i < end; ++i) + { + const float a = ratio * x[i]; + const float b = coeff * model_s[i]; + x[i] = a - b; + } + }); +} + +// Multistep DPM-Solver-2 (dpmsolver++): uses x0 at t_prev_1 and t_prev_0. +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 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 inv_r0 = 1.0f / r0; + const float phi_1 = std::expm1(-h); + const float ratio = t / t_prev_0; + const float coeff = alpha_t * phi_1; + const float half_coeff = 0.5f * coeff; + 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 = 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 din::image_gen diff --git a/image_gen/pixeldit/pixeldit_cli.cpp b/image_gen/pixeldit/pixeldit_cli.cpp new file mode 100644 index 0000000..61e798b --- /dev/null +++ b/image_gen/pixeldit/pixeldit_cli.cpp @@ -0,0 +1,283 @@ +// 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"; + } + 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" || lower == "dx" || lower == "dx-cig" || lower == "vk") + { + throw std::invalid_argument("--processing " + lower + " is not implemented for PixelDiT yet; use cpu"); + } + throw std::invalid_argument("Processing must be: cpu"); +} + +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") + .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..45555b0 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_cli.h @@ -0,0 +1,65 @@ +// 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, +}; + +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_pipeline.cpp b/image_gen/pixeldit/pixeldit_pipeline.cpp new file mode 100644 index 0000000..7b6a5ad --- /dev/null +++ b/image_gen/pixeldit/pixeldit_pipeline.cpp @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include + +#include "pixeldit_runtime_context.h" + +namespace din::image_gen +{ +namespace +{ + +PixelDiTRuntimeContext& GetProcessRuntime() +{ + static Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "PixelDiT"}; + static PixelDiTRuntimeContext 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; +} + +} // namespace + +std::unique_ptr CreatePixelDiTPipeline(const PixelDiTConfig& config) +{ + if (config.provider != PixelDiTExecutionProvider::TrtRtx) + { + throw std::invalid_argument("PixelDiT supports only --provider trt-rtx"); + } + PixelDiTRuntimeContext& runtime = GetProcessRuntime(); + + if (config.processing == PixelDiTProcessingBackend::Cpu) + { + return CreatePixelDiTTrtPipeline(config, runtime); + } + + throw std::invalid_argument("Unknown PixelDiT processing backend"); +} + +} // namespace din::image_gen diff --git a/image_gen/pixeldit/pixeldit_runtime_context.h b/image_gen/pixeldit/pixeldit_runtime_context.h new file mode 100644 index 0000000..0fc55e1 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_runtime_context.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 "ort_session.h" +#include "pixeldit_cli.h" + +namespace din::image_gen +{ + +struct PixelDiTRuntimeContext +{ + explicit PixelDiTRuntimeContext(Ort::Env& environment) + : env(environment) + { + } + + Ort::Env& env; + Ort::ConstEpDevice trt_device{}; +}; + +std::unique_ptr CreatePixelDiTTrtPipeline(const PixelDiTConfig& config, + PixelDiTRuntimeContext& runtime); + +} // 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..1bbf085 --- /dev/null +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -0,0 +1,670 @@ +// 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. "cuda" (kernels on the same +// buffers, like image_gen/flux2/flux2_cuda_backend.cpp) can be added here later. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nvtx_helper.h" +#include "ort_session.h" +#include "pixeldit.h" +#include "pixeldit_runtime_context.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, +}; + +const char* to_string(SamplingBackend backend) +{ + switch (backend) + { + case SamplingBackend::Cpu: + return "cpu"; + } + 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(); +} + +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); + return type_info.GetTensorTypeAndShapeInfo().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(); +} + +void expect_shape(Ort::Session& session, const char* model, const char* name, bool input, + const std::vector& expected) +{ + const auto actual = session_shape(session, name, input); + if (actual != expected) + { + throw std::runtime_error(std::string(model) + " " + name + " has shape " + shape_string(actual) + + " but pipeline config expects " + shape_string(expected) + + "; check pipeline_config.json matches this export"); + } +} + +// 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}); + } + } + + void Run(SamplingBackend backend, bool device_is_cuda, float t, float cfg_scale, FloatBuffer& x0) + { + if (velocity_ == nullptr || sample_ == nullptr) + { + throw std::runtime_error("CfgDataPredictionStage bindings are incomplete"); + } + (void)backend; // only SamplingBackend::Cpu exists so far + if (device_is_cuda) + { + auto velocity_ready = velocity_->CopyAsyncToHostWithNotification(); + velocity_ready.Sync(); + } + const size_t n = element_count(*sample_); + 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(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"); + } + (void)backend; + const size_t n = element_count(*sample_); + 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 (host view; GPU work is inside `dit`). +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; + // Host-side sampling state: current sample and the two most recent x0 estimates. + std::unique_ptr sample; + std::unique_ptr model_a; + std::unique_ptr model_b; + + 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}; + expect_shape(te, "text_encoder", "input_ids", true, text_shape); + expect_shape(te, "text_encoder", "attention_mask", true, text_shape); + expect_shape(te, "text_encoder", "select_index", true, select_shape); + expect_shape(te, "text_encoder", "prompt_embeds", false, embeds_shape); + expect_shape(dit, "transformer", "hidden_states", true, image_shape); + expect_shape(dit, "transformer", "timestep", true, timestep_shape); + expect_shape(dit, "transformer", "encoder_hidden_states", true, dit_text_shape); + expect_shape(dit, "transformer", "velocity", false, image_shape); + + 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); + state.sample = std::make_unique(dit_runner, sample_shape, false); + state.model_a = std::make_unique(dit_runner, sample_shape, false); + state.model_b = std::make_unique(dit_runner, sample_shape, false); + + 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; +} + +// One model evaluation: DiT on [sample, sample] at time t, then CFG -> x0 estimate. +void data_prediction(PipelineState& state, float t, din::common::TensorBuffer& x0) +{ + 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(state.sampling_backend, false, 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()); + } + + // 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. + state.timing = {}; + const auto start = std::chrono::steady_clock::now(); + 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], *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.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, *curr); + } + if (step % 10 == 0 || step == steps) + { + std::cout << " step " << step << "/" << steps << std::endl; + } + } + 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; + 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; +} + +class TrtPixelDiTProcessingPipeline final : public PixelDiTProcessingPipeline +{ +public: + TrtPixelDiTProcessingPipeline(PixelDiTConfig config, PixelDiTRuntimeContext& 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); + 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_; + PixelDiTRuntimeContext& runtime_; + std::unique_ptr state_; +}; + +} // namespace + +std::unique_ptr CreatePixelDiTTrtPipeline(const PixelDiTConfig& config, + PixelDiTRuntimeContext& runtime) +{ + std::cout << std::unitbuf; + std::cerr << std::unitbuf; + return std::make_unique(config, runtime); +} + +} // namespace din::image_gen From 229bfff2f22db0a68111e809c616c66346637585 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Sat, 3 Oct 2026 12:28:27 -0700 Subject: [PATCH 04/12] add tokenizer test --- CMakeLists.txt | 1 + tests/tokenizer/CMakeLists.txt | 14 ++ tests/tokenizer/test_tokenizer.cpp | 306 +++++++++++++++++++++++++++++ 3 files changed, 321 insertions(+) create mode 100644 tests/tokenizer/CMakeLists.txt create mode 100644 tests/tokenizer/test_tokenizer.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 391f954..415ea61 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -96,6 +96,7 @@ 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) 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..752cbce --- /dev/null +++ b/tests/tokenizer/test_tokenizer.cpp @@ -0,0 +1,306 @@ +// 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. + +#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; +} + +} // 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); + std::filesystem::remove_all(dir); + return failures == 0 ? 0 : 1; + } + catch (const std::exception& e) + { + std::cerr << "Error: " << e.what() << std::endl; + return 1; + } +} From 06a944a5929fe0eafc3d41983dfebe5725315a91 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Thu, 8 Oct 2026 15:19:39 -0700 Subject: [PATCH 05/12] remove expect shape --- image_gen/pixeldit/pixeldit_trt_backend.cpp | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp index 1bbf085..08151fc 100644 --- a/image_gen/pixeldit/pixeldit_trt_backend.cpp +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -127,18 +127,6 @@ std::string shape_string(const std::vector& shape) return text.str(); } -void expect_shape(Ort::Session& session, const char* model, const char* name, bool input, - const std::vector& expected) -{ - const auto actual = session_shape(session, name, input); - if (actual != expected) - { - throw std::runtime_error(std::string(model) + " " + name + " has shape " + shape_string(actual) + - " but pipeline config expects " + shape_string(expected) + - "; check pipeline_config.json matches this export"); - } -} - // CFG + data prediction: velocity [2, n] ([negative, positive]) and the current sample -> x0 estimate. class CfgDataPredictionStage { @@ -381,15 +369,6 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const 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}; - expect_shape(te, "text_encoder", "input_ids", true, text_shape); - expect_shape(te, "text_encoder", "attention_mask", true, text_shape); - expect_shape(te, "text_encoder", "select_index", true, select_shape); - expect_shape(te, "text_encoder", "prompt_embeds", false, embeds_shape); - expect_shape(dit, "transformer", "hidden_states", true, image_shape); - expect_shape(dit, "transformer", "timestep", true, timestep_shape); - expect_shape(dit, "transformer", "encoder_hidden_states", true, dit_text_shape); - expect_shape(dit, "transformer", "velocity", false, image_shape); - auto& te_runner = *state.text_encoder_runner; auto& dit_runner = *state.transformer_runner; state.input_ids = std::make_unique(te_runner, text_shape, true); From 3653fd7db5d2da3c982cd0eee610dcb4851fc198 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Thu, 8 Oct 2026 16:40:57 -0700 Subject: [PATCH 06/12] add cuda --- CMakeLists.txt | 21 +-- image_gen/pixeldit/CMakeLists.txt | 13 +- image_gen/pixeldit/README.md | 9 +- image_gen/pixeldit/cuda_kernels.cu | 111 +++++++++++++ image_gen/pixeldit/cuda_kernels.h | 26 +++ image_gen/pixeldit/helpers/pixeldit.h | 79 ++++++--- image_gen/pixeldit/pixeldit_cli.cpp | 14 +- image_gen/pixeldit/pixeldit_cli.h | 1 + image_gen/pixeldit/pixeldit_pipeline.cpp | 2 +- image_gen/pixeldit/pixeldit_trt_backend.cpp | 168 ++++++++++++++++---- 10 files changed, 373 insertions(+), 71 deletions(-) create mode 100644 image_gen/pixeldit/cuda_kernels.cu create mode 100644 image_gen/pixeldit/cuda_kernels.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 415ea61..2cabf5e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -216,15 +216,18 @@ if(DIN_BUILD_TESTING) ) endif () - add_din_cli_test(din_pixeldit_bf16_trt-rtx_cpu din_pixeldit_cli - --model-dir "${DIN_DEPLOY_ARTIFACTS_ROOT}/PixelDiT-1300M-1024px-onnx" - --provider trt-rtx - --processing cpu - --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_cpu" - ) + 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/image_gen/pixeldit/CMakeLists.txt b/image_gen/pixeldit/CMakeLists.txt index 0b5936d..52b0579 100644 --- a/image_gen/pixeldit/CMakeLists.txt +++ b/image_gen/pixeldit/CMakeLists.txt @@ -1,8 +1,13 @@ find_package(CUDAToolkit REQUIRED) +enable_language(CUDA) -# The host sampling math mirrors the PyTorch reference operation by operation in fp32; keep the compiler from -# fusing multiply-adds so results stay comparable (MSVC /fp:precise does not contract by default). -set(PIXELDIT_FP_FLAGS "$<$,$>:-ffp-contract=off>") +# 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 @@ -30,10 +35,12 @@ set_target_properties(din_pixeldit_common PROPERTIES 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 diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md index 3f0e176..78577f6 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -18,7 +18,7 @@ PixelDiT generates pixels directly, so there is no VAE. The model code is vendor | 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 | CFG and flow DPM-Solver++ (2nd order), 50 steps, CFG scale 2.75, flow shift 4.0 by default. | +| 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. @@ -27,8 +27,11 @@ Sampling defaults and the token layout are written to `pipeline_config.json`. Th | Mode | CLI | Inference and sampling | | --- | --- | --- | | Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | +| TensorRT RTX with CUDA sampling | `--provider trt-rtx --processing cuda` | TensorRT RTX inference with CUDA sampling. | -CUDA, DirectX, and Vulkan processing are not implemented yet. `--provider cpu` is not supported. +CUDA sampling keeps the sample on the GPU, so the DiT calls run back to back with one host sync per image (about 48 ms instead of 54 ms per step at 1024x1024 on an RTX PRO 6000). Both modes produce identical images for the same seed. + +DirectX and Vulkan processing are not implemented yet. `--provider cpu` is not supported. ## Export @@ -61,7 +64,7 @@ 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 cpu --prompt "a red fox" --output out +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 ``` 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..ab5fee4 --- /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 index 78092ef..2db931e 100644 --- a/image_gen/pixeldit/helpers/pixeldit.h +++ b/image_gen/pixeldit/helpers/pixeldit.h @@ -282,7 +282,54 @@ inline void cfg_data_prediction(const float* velocity, const float* x, float 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) { @@ -291,49 +338,33 @@ inline void dpm_first_order_update(float* x, const float* model_s, float s, floa std::copy_n(model_s, n, x); return; } - 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); - const float ratio = t / s; - const float coeff = alpha_t * phi_1; + 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 = ratio * x[i]; - const float b = coeff * model_s[i]; + const float a = k.ratio * x[i]; + const float b = k.coeff * model_s[i]; x[i] = a - b; } }); } -// Multistep DPM-Solver-2 (dpmsolver++): uses x0 at t_prev_1 and t_prev_0. 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 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 inv_r0 = 1.0f / r0; - const float phi_1 = std::expm1(-h); - const float ratio = t / t_prev_0; - const float coeff = alpha_t * phi_1; - const float half_coeff = 0.5f * coeff; + 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 = 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 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; } diff --git a/image_gen/pixeldit/pixeldit_cli.cpp b/image_gen/pixeldit/pixeldit_cli.cpp index 61e798b..c3169da 100644 --- a/image_gen/pixeldit/pixeldit_cli.cpp +++ b/image_gen/pixeldit/pixeldit_cli.cpp @@ -58,6 +58,8 @@ std::string to_string(PixelDiTProcessingBackend backend) { case PixelDiTProcessingBackend::Cpu: return "cpu"; + case PixelDiTProcessingBackend::Cuda: + return "cuda"; } return "unknown"; } @@ -79,11 +81,15 @@ PixelDiTProcessingBackend parse_processing_backend(const std::string& value) { return PixelDiTProcessingBackend::Cpu; } - if (lower == "cuda" || lower == "dx" || lower == "dx-cig" || lower == "vk") + if (lower == "cuda") { - throw std::invalid_argument("--processing " + lower + " is not implemented for PixelDiT yet; use cpu"); + return PixelDiTProcessingBackend::Cuda; } - throw std::invalid_argument("Processing must be: cpu"); + 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) @@ -111,7 +117,7 @@ PixelDiTConfig parse_args(int argc, char* argv[]) parser.add_argument("--processing") .default_value(to_string(config.processing)) .nargs(1) - .metavar("cpu") + .metavar("cpu|cuda") .help("Select the processing backend."); parser.add_argument("--provider") .default_value(to_string(config.provider)) diff --git a/image_gen/pixeldit/pixeldit_cli.h b/image_gen/pixeldit/pixeldit_cli.h index 45555b0..e9323c6 100644 --- a/image_gen/pixeldit/pixeldit_cli.h +++ b/image_gen/pixeldit/pixeldit_cli.h @@ -15,6 +15,7 @@ namespace din::image_gen enum class PixelDiTProcessingBackend { Cpu, + Cuda, }; enum class PixelDiTExecutionProvider diff --git a/image_gen/pixeldit/pixeldit_pipeline.cpp b/image_gen/pixeldit/pixeldit_pipeline.cpp index 7b6a5ad..942f75d 100644 --- a/image_gen/pixeldit/pixeldit_pipeline.cpp +++ b/image_gen/pixeldit/pixeldit_pipeline.cpp @@ -35,7 +35,7 @@ std::unique_ptr CreatePixelDiTPipeline(const PixelDi } PixelDiTRuntimeContext& runtime = GetProcessRuntime(); - if (config.processing == PixelDiTProcessingBackend::Cpu) + if (config.processing == PixelDiTProcessingBackend::Cpu || config.processing == PixelDiTProcessingBackend::Cuda) { return CreatePixelDiTTrtPipeline(config, runtime); } diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp index 08151fc..6947ce4 100644 --- a/image_gen/pixeldit/pixeldit_trt_backend.cpp +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -2,8 +2,9 @@ // 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. "cuda" (kernels on the same -// buffers, like image_gen/flux2/flux2_cuda_backend.cpp) can be added here later. +// 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 @@ -22,6 +23,7 @@ #include #include +#include "cuda_kernels.h" #include "nvtx_helper.h" #include "ort_session.h" #include "pixeldit.h" @@ -47,6 +49,7 @@ namespace enum class SamplingBackend { Cpu, + Cuda, }; const char* to_string(SamplingBackend backend) @@ -55,6 +58,8 @@ const char* to_string(SamplingBackend backend) { case SamplingBackend::Cpu: return "cpu"; + case SamplingBackend::Cuda: + return "cuda"; } return "unknown"; } @@ -68,6 +73,18 @@ size_t element_count(din::common::TensorBuffer& buffer) return buffer.BindingValue().GetTensorTypeAndShapeInfo().GetElementCount(); } +// Pointer to the bound tensor of a device-backed buffer: device memory, or pinned host memory on unified-memory GPUs. +float* device_data(FloatBuffer& buffer) +{ + return buffer.BindingValue().GetTensorMutableData(); +} + +// cudaMemcpyDefault so the copy also works when the bound tensors are pinned host memory (unified memory). +void copy_device_async(float* dst, const float* src, size_t count, cudaStream_t stream) +{ + PIXELDIT_CUDA_CHECK(cudaMemcpyAsync(dst, src, count * sizeof(float), cudaMemcpyDefault, stream)); +} + template void write_raw(const std::filesystem::path& path, const T* data, size_t count) { @@ -147,19 +164,20 @@ class CfgDataPredictionStage } } - void Run(SamplingBackend backend, bool device_is_cuda, float t, float cfg_scale, FloatBuffer& x0) + // 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"); } - (void)backend; // only SamplingBackend::Cpu exists so far - if (device_is_cuda) + const size_t n = element_count(*sample_); + if (backend == SamplingBackend::Cuda) { - auto velocity_ready = velocity_->CopyAsyncToHostWithNotification(); - velocity_ready.Sync(); + launch_pixeldit_cfg_x0_kernel(stream, device_data(*velocity_), device_data(*sample_), t, cfg_scale, + device_data(x0), n); + return; } - const size_t n = element_count(*sample_); cfg_data_prediction(velocity_->HostData(), sample_->HostData(), t, cfg_scale, x0.HostData(), n); } @@ -181,15 +199,35 @@ class DpmSolverStage sample_ = &tensor; } - void Run(SamplingBackend backend, int order, FloatBuffer* model_prev_1, FloatBuffer& model_prev_0, float t_prev_1, - float t_prev_0, float t) + void Run(cudaStream_t 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"); } - (void)backend; const size_t n = element_count(*sample_); + if (backend == SamplingBackend::Cuda) + { + if (order == 1 && t == 0.0f) + { + // Final step: the update is exactly x0 (see dpm_first_order_update). + copy_device_async(device_data(*sample_), device_data(model_prev_0), n, 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); @@ -213,7 +251,8 @@ struct SamplerSettings std::string negative_prompt; }; -// Per-image wall time of each sampling stage (host view; GPU work is inside `dit`). +// 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 @@ -254,10 +293,15 @@ struct PipelineState std::unique_ptr dit_timestep; std::unique_ptr dit_encoder_hidden_states; std::unique_ptr dit_velocity; - // Host-side sampling state: current sample and the two most recent x0 estimates. + // 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; std::unique_ptr text_encoder_io; std::unique_ptr transformer_io; @@ -379,9 +423,15 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const 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); - state.sample = std::make_unique(dit_runner, sample_shape, false); - state.model_a = std::make_unique(dit_runner, sample_shape, false); - state.model_b = std::make_unique(dit_runner, sample_shape, false); + 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); + } state.text_encoder_io = std::make_unique(te); state.text_encoder_io->BindInput("input_ids", state.input_ids->BindingValue()); @@ -458,9 +508,44 @@ void encode_prompts(PipelineState& state, const PixelDiTConfig& config) 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) +{ + const PixelDiTPipelineConfig& pc = state.pipeline_config; + const size_t n = element_count(*state.sample); + const cudaStream_t stream = compute_stream(state); + float* hidden = device_data(*state.dit_hidden_states); + copy_device_async(hidden, device_data(*state.sample), n, stream); + copy_device_async(hidden + n, device_data(*state.sample), n, stream); + const size_t batch = static_cast(pc.dit_batch); + copy_device_async(device_data(*state.dit_timestep), + device_data(*state.timestep_table) + static_cast(eval_index) * batch, batch, stream); + { + 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, din::common::TensorBuffer& x0) +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(); @@ -487,7 +572,7 @@ void data_prediction(PipelineState& state, float t, din::common::TensorBufferHostData(); + 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], *curr); + data_prediction(state, timesteps[0], 0, *curr); int history = 1; for (int step = 1; step <= steps; ++step) { @@ -543,7 +644,8 @@ PixelDiTImage run_pipeline(PipelineState& state, const PixelDiTConfig& config, u throw std::logic_error("2nd-order update without two x0 estimates"); } const auto solver_start = std::chrono::steady_clock::now(); - state.solver.Run(state.sampling_backend, step_order, prev, *curr, t_prev_1, t_prev_0, t); + state.solver.Run(compute_stream(state), 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; @@ -551,20 +653,30 @@ PixelDiTImage run_pipeline(PipelineState& state, const PixelDiTConfig& config, u history = std::min(history + 1, 2); if (step < steps) { - data_prediction(state, t, *curr); + 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; - 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 (!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) { @@ -601,6 +713,8 @@ class TrtPixelDiTProcessingPipeline final : public PixelDiTProcessingPipeline 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; From 33f377596eab7bd370f33345eb7a17d1581bb2ba Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Fri, 9 Oct 2026 10:40:15 -0700 Subject: [PATCH 07/12] addess AI comments --- common/io/tokenizer.cpp | 5 ++++ .../pixeldit/model_export/export_pixeldit.py | 3 +-- .../pixeldit/model_export/verify_pixeldit.py | 8 ++++--- tests/tokenizer/test_tokenizer.cpp | 23 +++++++++++++++++++ 4 files changed, 34 insertions(+), 5 deletions(-) diff --git a/common/io/tokenizer.cpp b/common/io/tokenizer.cpp index 109a2a0..1c69c91 100644 --- a/common/io/tokenizer.cpp +++ b/common/io/tokenizer.cpp @@ -238,6 +238,11 @@ JsonTokenizerData LoadJsonTokenizerData(const std::string& path) } 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")) diff --git a/image_gen/pixeldit/model_export/export_pixeldit.py b/image_gen/pixeldit/model_export/export_pixeldit.py index 26ebd23..6e7f3e9 100644 --- a/image_gen/pixeldit/model_export/export_pixeldit.py +++ b/image_gen/pixeldit/model_export/export_pixeldit.py @@ -357,7 +357,6 @@ def main(): break if m not in to_export: to_export.append(m) - export_all = (args.model is None) or ("all" in raw) device = "cuda" io_dtype = IO_PRECISION_MAP[args.io_precision] @@ -391,7 +390,7 @@ def main(): torch.cuda.empty_cache() exported_model_paths.append(output_path / name / "model.onnx") - if export_all: + 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()) diff --git a/image_gen/pixeldit/model_export/verify_pixeldit.py b/image_gen/pixeldit/model_export/verify_pixeldit.py index 37c1d8a..06d215e 100644 --- a/image_gen/pixeldit/model_export/verify_pixeldit.py +++ b/image_gen/pixeldit/model_export/verify_pixeldit.py @@ -120,9 +120,11 @@ def _register_trt_rtx(args) -> str: 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) - 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)) + # 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) diff --git a/tests/tokenizer/test_tokenizer.cpp b/tests/tokenizer/test_tokenizer.cpp index 752cbce..d1c8b40 100644 --- a/tests/tokenizer/test_tokenizer.cpp +++ b/tests/tokenizer/test_tokenizer.cpp @@ -7,6 +7,7 @@ // * 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 @@ -283,6 +284,27 @@ int run_cases(const std::filesystem::path& dir, const char* name, const json& to 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() @@ -295,6 +317,7 @@ int main() 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; } From 22fa0a912f7fef2356d20d9ba9f3107aa0bcdd36 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Fri, 9 Oct 2026 10:41:29 -0700 Subject: [PATCH 08/12] fix readme --- image_gen/pixeldit/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md index 78577f6..cad0e6c 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -29,7 +29,7 @@ Sampling defaults and the token layout are written to `pipeline_config.json`. Th | Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | | TensorRT RTX with CUDA sampling | `--provider trt-rtx --processing cuda` | TensorRT RTX inference with CUDA sampling. | -CUDA sampling keeps the sample on the GPU, so the DiT calls run back to back with one host sync per image (about 48 ms instead of 54 ms per step at 1024x1024 on an RTX PRO 6000). Both modes produce identical images for the same seed. +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. From e234fb9f38c3cd44af3310a2044e50821a429cd2 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Tue, 13 Oct 2026 09:51:02 -0700 Subject: [PATCH 09/12] remove pipeline --- .gitignore | 3 ++ image_gen/pixeldit/CMakeLists.txt | 11 +--- image_gen/pixeldit/README.md | 8 ++- image_gen/pixeldit/helpers/pixeldit.h | 7 +-- image_gen/pixeldit/pixeldit_pipeline.cpp | 46 ----------------- image_gen/pixeldit/pixeldit_runtime_context.h | 26 ---------- image_gen/pixeldit/pixeldit_trt_backend.cpp | 51 +++++++++++++++---- 7 files changed, 54 insertions(+), 98 deletions(-) delete mode 100644 image_gen/pixeldit/pixeldit_pipeline.cpp delete mode 100644 image_gen/pixeldit/pixeldit_runtime_context.h diff --git a/.gitignore b/.gitignore index b8f79d9..5df9e42 100644 --- a/.gitignore +++ b/.gitignore @@ -94,3 +94,6 @@ audio/ TODO frames/ .run/ + +models/ +*.png \ No newline at end of file diff --git a/image_gen/pixeldit/CMakeLists.txt b/image_gen/pixeldit/CMakeLists.txt index 52b0579..5848117 100644 --- a/image_gen/pixeldit/CMakeLists.txt +++ b/image_gen/pixeldit/CMakeLists.txt @@ -50,15 +50,6 @@ target_link_libraries(din_pixeldit_trt CUDA::cudart ) -add_library(din_pixeldit_runtime STATIC - pixeldit_pipeline.cpp -) -set_target_properties(din_pixeldit_runtime PROPERTIES - CXX_STANDARD 20 - CXX_EXTENSIONS OFF -) -target_link_libraries(din_pixeldit_runtime PUBLIC din_pixeldit_trt) - add_din_executable(din_pixeldit_cli pixeldit_cli.cpp ) @@ -66,4 +57,4 @@ set_target_properties(din_pixeldit_cli PROPERTIES CXX_STANDARD 20 CXX_EXTENSIONS OFF ) -target_link_libraries(din_pixeldit_cli PRIVATE argparse din_pixeldit_runtime) +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 index cad0e6c..a62182d 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -26,8 +26,10 @@ Sampling defaults and the token layout are written to `pipeline_config.json`. Th | Mode | CLI | Inference and sampling | | --- | --- | --- | -| Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | -| TensorRT RTX with CUDA sampling | `--provider trt-rtx --processing cuda` | TensorRT RTX inference with CUDA 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. @@ -67,4 +69,6 @@ cmake --build --preset windows-x64-release --target din_pixeldit_cli 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/helpers/pixeldit.h b/image_gen/pixeldit/helpers/pixeldit.h index 2db931e..3bae932 100644 --- a/image_gen/pixeldit/helpers/pixeldit.h +++ b/image_gen/pixeldit/helpers/pixeldit.h @@ -22,11 +22,8 @@ namespace din::image_gen // Configuration // ============================================================================ -#ifdef _WIN32 -static const std::filesystem::path DEFAULT_PIXELDIT_MODEL_DIR = "S:/din_deploy_artifacts/PixelDiT-1300M-1024px-onnx"; -#else -static const std::filesystem::path DEFAULT_PIXELDIT_MODEL_DIR = "/mnt/share/onnx/PixelDiT-1300M-1024px-onnx/"; -#endif +// 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 { diff --git a/image_gen/pixeldit/pixeldit_pipeline.cpp b/image_gen/pixeldit/pixeldit_pipeline.cpp deleted file mode 100644 index 942f75d..0000000 --- a/image_gen/pixeldit/pixeldit_pipeline.cpp +++ /dev/null @@ -1,46 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -#include -#include - -#include "pixeldit_runtime_context.h" - -namespace din::image_gen -{ -namespace -{ - -PixelDiTRuntimeContext& GetProcessRuntime() -{ - static Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "PixelDiT"}; - static PixelDiTRuntimeContext 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; -} - -} // namespace - -std::unique_ptr CreatePixelDiTPipeline(const PixelDiTConfig& config) -{ - if (config.provider != PixelDiTExecutionProvider::TrtRtx) - { - throw std::invalid_argument("PixelDiT supports only --provider trt-rtx"); - } - PixelDiTRuntimeContext& runtime = GetProcessRuntime(); - - if (config.processing == PixelDiTProcessingBackend::Cpu || config.processing == PixelDiTProcessingBackend::Cuda) - { - return CreatePixelDiTTrtPipeline(config, runtime); - } - - throw std::invalid_argument("Unknown PixelDiT processing backend"); -} - -} // namespace din::image_gen diff --git a/image_gen/pixeldit/pixeldit_runtime_context.h b/image_gen/pixeldit/pixeldit_runtime_context.h deleted file mode 100644 index 0fc55e1..0000000 --- a/image_gen/pixeldit/pixeldit_runtime_context.h +++ /dev/null @@ -1,26 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once - -#include "ort_session.h" -#include "pixeldit_cli.h" - -namespace din::image_gen -{ - -struct PixelDiTRuntimeContext -{ - explicit PixelDiTRuntimeContext(Ort::Env& environment) - : env(environment) - { - } - - Ort::Env& env; - Ort::ConstEpDevice trt_device{}; -}; - -std::unique_ptr CreatePixelDiTTrtPipeline(const PixelDiTConfig& config, - PixelDiTRuntimeContext& runtime); - -} // namespace din::image_gen diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp index 6947ce4..698131b 100644 --- a/image_gen/pixeldit/pixeldit_trt_backend.cpp +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -27,7 +28,7 @@ #include "nvtx_helper.h" #include "ort_session.h" #include "pixeldit.h" -#include "pixeldit_runtime_context.h" +#include "pixeldit_cli.h" #include #include @@ -73,13 +74,15 @@ size_t element_count(din::common::TensorBuffer& buffer) return buffer.BindingValue().GetTensorTypeAndShapeInfo().GetElementCount(); } -// Pointer to the bound tensor of a device-backed buffer: device memory, or pinned host memory on unified-memory GPUs. +// 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(); } -// cudaMemcpyDefault so the copy also works when the bound tensors are pinned host memory (unified memory). +// cudaMemcpyDefault lets CUDA infer the copy direction from the pointers, so this works for either kind of bound +// tensor described above. void copy_device_async(float* dst, const float* src, size_t count, cudaStream_t stream) { PIXELDIT_CUDA_CHECK(cudaMemcpyAsync(dst, src, count * sizeof(float), cudaMemcpyDefault, stream)); @@ -126,7 +129,10 @@ std::vector session_shape(Ort::Session& session, const char* name, bool if (std::string_view{io_name.get()} == name) { const auto type_info = input ? session.GetInputTypeInfo(i) : session.GetOutputTypeInfo(i); - return type_info.GetTensorTypeAndShapeInfo().GetShape(); + // 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); @@ -689,10 +695,36 @@ PixelDiTImage run_pipeline(PipelineState& state, const PixelDiTConfig& config, u 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, PixelDiTRuntimeContext& runtime) + TrtPixelDiTProcessingPipeline(PixelDiTConfig config, ProcessRuntime& runtime) : config_(std::move(config)) , runtime_(runtime) { @@ -746,18 +778,19 @@ class TrtPixelDiTProcessingPipeline final : public PixelDiTProcessingPipeline private: PixelDiTConfig config_; - PixelDiTRuntimeContext& runtime_; + ProcessRuntime& runtime_; std::unique_ptr state_; }; } // namespace -std::unique_ptr CreatePixelDiTTrtPipeline(const PixelDiTConfig& config, - PixelDiTRuntimeContext& runtime) +// 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, runtime); + return std::make_unique(config, GetProcessRuntime()); } } // namespace din::image_gen From b0cd4d4750ba527f1d98eb8fec30e40ac42c13a2 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Tue, 13 Oct 2026 11:51:53 -0700 Subject: [PATCH 10/12] use copy tensor and clang format --- common/io/tokenizer.cpp | 10 +- image_gen/pixeldit/cuda_kernels.cu | 4 +- image_gen/pixeldit/helpers/pixeldit.h | 10 +- image_gen/pixeldit/pixeldit_trt_backend.cpp | 101 +++++++++++++------- tests/tokenizer/test_tokenizer.cpp | 46 +++++---- 5 files changed, 110 insertions(+), 61 deletions(-) diff --git a/common/io/tokenizer.cpp b/common/io/tokenizer.cpp index 1c69c91..2264482 100644 --- a/common/io/tokenizer.cpp +++ b/common/io/tokenizer.cpp @@ -147,8 +147,8 @@ bool CollectNormalizerSteps(const nlohmann::json& normalizer, std::vector(), - normalizer["content"].get()}); + steps.push_back( + {false, normalizer["pattern"]["String"].get(), normalizer["content"].get()}); return true; } return false; @@ -826,9 +826,9 @@ std::vector Tokenizer::EncodeSentencePieceBpe(const std::string& text, 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); + 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)) { diff --git a/image_gen/pixeldit/cuda_kernels.cu b/image_gen/pixeldit/cuda_kernels.cu index ab5fee4..2f486aa 100644 --- a/image_gen/pixeldit/cuda_kernels.cu +++ b/image_gen/pixeldit/cuda_kernels.cu @@ -103,8 +103,8 @@ void launch_pixeldit_dpm2_kernel(cudaStream_t stream, float* sample, const float 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); + pixeldit_dpm2_kernel<<>>(sample, model_prev_1, model_prev_0, ratio, coeff, + half_coeff, inv_r0, n); check_launch(); } diff --git a/image_gen/pixeldit/helpers/pixeldit.h b/image_gen/pixeldit/helpers/pixeldit.h index 3bae932..fdd4c03 100644 --- a/image_gen/pixeldit/helpers/pixeldit.h +++ b/image_gen/pixeldit/helpers/pixeldit.h @@ -194,8 +194,8 @@ 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}); + 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); @@ -208,7 +208,11 @@ void ParallelFor(size_t n, Fn&& fn) { const size_t begin = k * chunk; const size_t end = std::min(n, begin + chunk); - workers.emplace_back([&fn, begin, end] { fn(begin, end); }); + workers.emplace_back( + [&fn, begin, end] + { + fn(begin, end); + }); } fn(size_t{0}, std::min(n, chunk)); for (auto& worker : workers) diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp index 698131b..07b00c9 100644 --- a/image_gen/pixeldit/pixeldit_trt_backend.cpp +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -37,14 +37,14 @@ namespace din::image_gen namespace { -#define PIXELDIT_CUDA_CHECK(call) \ - do \ - { \ - const cudaError_t _err = (call); \ - if (_err != cudaSuccess) \ - { \ +#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 @@ -81,11 +81,28 @@ float* device_data(FloatBuffer& buffer) return buffer.BindingValue().GetTensorMutableData(); } -// cudaMemcpyDefault lets CUDA infer the copy direction from the pointers, so this works for either kind of bound -// tensor described above. -void copy_device_async(float* dst, const float* src, size_t count, cudaStream_t stream) +// ORT's CopyTensors only copies whole OrtValues, so a partial copy goes through a view: an OrtValue of `shape` over +// the bound tensor's memory starting `offset` elements in, with the same memory info as the bound tensor. +Ort::Value make_view(FloatBuffer& buffer, size_t offset, const std::vector& shape) { - PIXELDIT_CUDA_CHECK(cudaMemcpyAsync(dst, src, count * sizeof(float), cudaMemcpyDefault, stream)); + 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()); +} + +// Asynchronous copy through the execution provider's data transfer (OrtApi::CopyTensors), ordered on `stream`. +void copy_tensor(Ort::Env& env, Ort::SyncStream& stream, const Ort::Value& src, Ort::Value& dst) +{ + Ort::ThrowOnError(env.CopyTensor(src, dst, stream)); } template @@ -125,7 +142,8 @@ std::vector session_shape(Ort::Session& session, const char* name, bool Ort::AllocatorWithDefaultOptions allocator; for (size_t i = 0; i < count; ++i) { - const auto io_name = input ? session.GetInputNameAllocated(i, allocator) : session.GetOutputNameAllocated(i, allocator); + 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); @@ -205,7 +223,7 @@ class DpmSolverStage sample_ = &tensor; } - void Run(cudaStream_t stream, SamplingBackend backend, int order, FloatBuffer* model_prev_1, + 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) @@ -215,10 +233,11 @@ class DpmSolverStage 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) { // Final step: the update is exactly x0 (see dpm_first_order_update). - copy_device_async(device_data(*sample_), device_data(model_prev_0), n, stream); + copy_tensor(env, sync_stream, model_prev_0.BindingValue(), sample_->BindingValue()); } else if (order == 1) { @@ -308,6 +327,10 @@ struct PipelineState // 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; @@ -362,8 +385,8 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const << "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); + 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) { @@ -376,8 +399,8 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const 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; + 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; @@ -395,13 +418,13 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const 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); + 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); + 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; @@ -437,6 +460,17 @@ void initialize_state(PipelineState& state, Ort::ConstEpDevice trt_device, const { 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); @@ -523,15 +557,14 @@ cudaStream_t compute_stream(PipelineState& state) // `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) { - const PixelDiTPipelineConfig& pc = state.pipeline_config; - const size_t n = element_count(*state.sample); + Ort::SyncStream& sync_stream = *state.compute_stream; + for (Ort::Value& half : state.hidden_state_halves) + { + copy_tensor(state.env, sync_stream, state.sample->BindingValue(), half); + } + copy_tensor(state.env, sync_stream, state.timestep_rows[static_cast(eval_index)], + state.dit_timestep->BindingValue()); const cudaStream_t stream = compute_stream(state); - float* hidden = device_data(*state.dit_hidden_states); - copy_device_async(hidden, device_data(*state.sample), n, stream); - copy_device_async(hidden + n, device_data(*state.sample), n, stream); - const size_t batch = static_cast(pc.dit_batch); - copy_device_async(device_data(*state.dit_timestep), - device_data(*state.timestep_table) + static_cast(eval_index) * batch, batch, stream); { din::common::nvtx_scoped_range nvtx("transformer"); Ort::RunOptions run_options = make_run_options(state); @@ -650,8 +683,8 @@ PixelDiTImage run_pipeline(PipelineState& state, const PixelDiTConfig& config, u throw std::logic_error("2nd-order update without two x0 estimates"); } const auto solver_start = std::chrono::steady_clock::now(); - state.solver.Run(compute_stream(state), state.sampling_backend, step_order, prev, *curr, t_prev_1, t_prev_0, - t); + 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; diff --git a/tests/tokenizer/test_tokenizer.cpp b/tests/tokenizer/test_tokenizer.cpp index d1c8b40..c747f78 100644 --- a/tests/tokenizer/test_tokenizer.cpp +++ b/tests/tokenizer/test_tokenizer.cpp @@ -61,16 +61,17 @@ std::string utf8(uint32_t code) 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}}; + 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"}; + 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]; @@ -87,7 +88,8 @@ json sentencepiece_tokenizer(bool prepend) const json replace = {{"type", "Replace"}, {"pattern", {{"String", " "}}}, {"content", sp}}; const json normalizer = - prepend ? json{{"type", "Sequence"}, {"normalizers", {{{"type", "Prepend"}, {"prepend", sp}}, replace}}} : replace; + prepend ? json{{"type", "Sequence"}, {"normalizers", {{{"type", "Prepend"}, {"prepend", sp}}, replace}}} + : replace; return { {"version", "1.0"}, @@ -96,17 +98,24 @@ json sentencepiece_tokenizer(bool prepend) 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}}}, + {"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}}}}}}, + {"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}}}, + {{"type", "BPE"}, + {"unk_token", ""}, + {"fuse_unk", true}, + {"byte_fallback", true}, + {"ignore_merges", false}, + {"vocab", vocab}, + {"merges", merges}}}, }; } @@ -123,8 +132,8 @@ json byte_level_tokenizer() 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"}}; + {"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) @@ -144,7 +153,8 @@ json byte_level_tokenizer() {"version", "1.0"}, {"added_tokens", json::array()}, {"normalizer", nullptr}, - {"pre_tokenizer", {{"type", "ByteLevel"}, {"add_prefix_space", false}, {"trim_offsets", true}, {"use_regex", true}}}, + {"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}}}, }; @@ -289,7 +299,8 @@ int run_cases(const std::filesystem::path& dir, const char* name, const json& to 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"]}}}; + 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 @@ -301,7 +312,8 @@ int check_unsupported_normalizer_rejected(const std::filesystem::path& dir) 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; + std::cerr << "FAIL unsupported_normalizer.json: loaded a SentencePiece tokenizer with an NFKC normalizer" + << std::endl; return 1; } From 716b14b9f166907996afd7aaef639b09ae8c83ee Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Tue, 13 Oct 2026 11:57:55 -0700 Subject: [PATCH 11/12] remove copy wrapper --- image_gen/pixeldit/pixeldit_trt_backend.cpp | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/image_gen/pixeldit/pixeldit_trt_backend.cpp b/image_gen/pixeldit/pixeldit_trt_backend.cpp index 07b00c9..3976db2 100644 --- a/image_gen/pixeldit/pixeldit_trt_backend.cpp +++ b/image_gen/pixeldit/pixeldit_trt_backend.cpp @@ -81,8 +81,6 @@ float* device_data(FloatBuffer& buffer) return buffer.BindingValue().GetTensorMutableData(); } -// ORT's CopyTensors only copies whole OrtValues, so a partial copy goes through a view: an OrtValue of `shape` over -// the bound tensor's memory starting `offset` elements in, with the same memory info as the bound tensor. Ort::Value make_view(FloatBuffer& buffer, size_t offset, const std::vector& shape) { size_t count = 1; @@ -99,12 +97,6 @@ Ort::Value make_view(FloatBuffer& buffer, size_t offset, const std::vector void write_raw(const std::filesystem::path& path, const T* data, size_t count) { @@ -236,8 +228,7 @@ class DpmSolverStage const auto stream = reinterpret_cast(sync_stream.GetHandle()); if (order == 1 && t == 0.0f) { - // Final step: the update is exactly x0 (see dpm_first_order_update). - copy_tensor(env, sync_stream, model_prev_0.BindingValue(), sample_->BindingValue()); + Ort::ThrowOnError(env.CopyTensor(model_prev_0.BindingValue(), sample_->BindingValue(), sync_stream)); } else if (order == 1) { @@ -560,10 +551,10 @@ void data_prediction_cuda(PipelineState& state, float t, int eval_index, FloatBu Ort::SyncStream& sync_stream = *state.compute_stream; for (Ort::Value& half : state.hidden_state_halves) { - copy_tensor(state.env, sync_stream, state.sample->BindingValue(), half); + Ort::ThrowOnError(state.env.CopyTensor(state.sample->BindingValue(), half, sync_stream)); } - copy_tensor(state.env, sync_stream, state.timestep_rows[static_cast(eval_index)], - state.dit_timestep->BindingValue()); + 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"); From 5ded8998883271918afc57e393c039d7bc7b368d Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Tue, 13 Oct 2026 14:26:32 -0700 Subject: [PATCH 12/12] revert gitignore --- .gitignore | 3 --- 1 file changed, 3 deletions(-) diff --git a/.gitignore b/.gitignore index 5df9e42..b8f79d9 100644 --- a/.gitignore +++ b/.gitignore @@ -94,6 +94,3 @@ audio/ TODO frames/ .run/ - -models/ -*.png \ No newline at end of file