From 29dbb745c7a0f2e1a6852bc43054278294520743 Mon Sep 17 00:00:00 2001 From: Sarah Wu Date: Fri, 16 Oct 2026 14:55:36 -0700 Subject: [PATCH] switch to pixelDiT repo and packaging --- README.md | 2 +- image_gen/pixeldit/README.md | 11 +- .../pixeldit/model_export/export_pixeldit.py | 12 +- .../model_export/pixeldit/__init__.py | 4 +- .../pixeldit/model_export/pixeldit/model.py | 445 ------------------ .../model_export/pixeldit/upstream_patch.py | 207 ++++++++ .../pixeldit/model_export/pixeldit/weights.py | 33 +- .../model_export/test_upstream_patch.py | 107 +++++ .../pixeldit/model_export/verify_pixeldit.py | 17 +- pyproject.toml | 6 + 10 files changed, 374 insertions(+), 470 deletions(-) delete mode 100644 image_gen/pixeldit/model_export/pixeldit/model.py create mode 100644 image_gen/pixeldit/model_export/pixeldit/upstream_patch.py create mode 100644 image_gen/pixeldit/model_export/test_upstream_patch.py diff --git a/README.md b/README.md index 4957d58..9ad1c0b 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,7 @@ Use the model-specific README for the exporter command and CLI arguments. Create a Python environment and install the repository plus the extra required by the model you plan to export: ```bash -python -m pip install -e ".[parakeet]" # or .[nemotron] / .[flux] +python -m pip install -e ".[parakeet]" # or .[nemotron] / .[flux] / .[pixeldit] ``` The base `pyproject.toml` contains common exporter dependencies; the optional extras add the model-family dependencies. Export the model to a directory, then pass that directory to the corresponding CLI with `--model-dir`. diff --git a/image_gen/pixeldit/README.md b/image_gen/pixeldit/README.md index a62182d..67d0a15 100644 --- a/image_gen/pixeldit/README.md +++ b/image_gen/pixeldit/README.md @@ -9,7 +9,7 @@ CLI image generation with ONNX Runtime. | PixelDiT-1300M-1024px | `nvidia/PixelDiT-1300M-1024px` | FP32 (runs in BF16) | BF16 weights, FP32 I/O | | Gemma-2-2B-IT (text encoder) | `Efficient-Large-Model/gemma-2-2b-it` | BF16 | BF16 weights, FP32 I/O | -PixelDiT generates pixels directly, so there is no VAE. The model code is vendored in `model_export/pixeldit/`, so the [NVlabs/PixelDiT](https://github.com/NVlabs/PixelDiT) repo is not needed to export or verify. +PixelDiT generates pixels directly, so there is no VAE. The DiT model code comes from the `pixeldit` pip package, installed from [NVlabs/PixelDiT](https://github.com/NVlabs/PixelDiT). ## Pipeline @@ -37,11 +37,16 @@ DirectX and Vulkan processing are not implemented yet. `--provider cpu` is not s ## Export +Install the export dependencies first, from the repository root. The `pixeldit` extra adds `pixeldit`, upstream's PixelDiT model package, to the common exporter dependencies (see [Requirements](../../README.md#requirements)). Export and verify run on the GPU, so install a CUDA build of PyTorch first (see [pytorch.org](https://pytorch.org/get-started/locally/)). + ```bash -cd model_export +python -m pip install -e ".[pixeldit]" +cd image_gen/pixeldit/model_export python export_pixeldit.py --output ./models/PixelDiT-1300M-1024px-onnx ``` +The DiT is upstream's `PixDiT_T2I` (imported as `pixdit_core`; the local `model_export/pixeldit/` folder holds din-deploy's export helpers). `export_pixeldit.py` patches it only while exporting: complex-valued RoPE has no ONNX lowering, and TensorRT RTX rejects the `Col2Im` op that `F.fold` exports to. See `model_export/pixeldit/upstream_patch.py` for the details. `test_upstream_patch.py` checks the patches on a small random model. + An export contains `text_encoder/model.onnx`, `transformer/model.onnx`, `tokenizer/tokenizer.json`, and `pipeline_config.json`. Shapes are static. `--height` and `--width` (default 1024, multiples of 16) set the image size; re-export to change it. Add `--compile_trt --trt_root ` to check that each model builds with the TensorRT RTX CLI. @@ -52,7 +57,7 @@ Shapes are static. `--height` and `--width` (default 1024, multiples of 16) set python verify_pixeldit.py --onnx_dir ./models/PixelDiT-1300M-1024px-onnx --provider trt-rtx ``` -Compares each ONNX model running on TensorRT RTX with the vendored PyTorch model, then generates one image with each and saves `pytorch_output.png`, `ort_output.png`, and `comparison.png`. Pass `--validate_only` to only check the graphs. +Compares each ONNX model running on TensorRT RTX with the upstream PyTorch model from `pixeldit` (unpatched), then generates one image with each and saves `pytorch_output.png`, `ort_output.png`, and `comparison.png`. Pass `--validate_only` to only check the graphs. Per-model checks decide pass or fail. The image check only fails on broken output: 50-step BF16 sampling amplifies tiny rounding differences, so correct runs can differ in fine detail, or with some seeds in composition. Compare the images visually. diff --git a/image_gen/pixeldit/model_export/export_pixeldit.py b/image_gen/pixeldit/model_export/export_pixeldit.py index 6e7f3e9..a3d2dc6 100644 --- a/image_gen/pixeldit/model_export/export_pixeldit.py +++ b/image_gen/pixeldit/model_export/export_pixeldit.py @@ -32,6 +32,7 @@ import torch from pixeldit import pipeline as pixeldit_pipeline +from pixeldit.upstream_patch import export_patches from pixeldit.weights import DEFAULT_MODEL_NAME, DEFAULT_TEXT_ENCODER_NAME, load_gemma_text_encoder, load_pixeldit EXTERNAL_DATA_NAME = "model.onnx_data" @@ -114,15 +115,18 @@ def forward(self, input_ids, attention_mask, select_index): class PixelDiTExportWrapper(torch.nn.Module): - """Inputs are cast to bf16 inside PixelDiT (as upstream PixDiTTrainer does), output to io_dtype.""" + """Inputs are cast to the model dtype , output to io_dtype.""" def __init__(self, model, io_dtype: torch.dtype): super().__init__() self.model = model self.io_dtype = io_dtype + self.model_dtype = next(model.parameters()).dtype def forward(self, hidden_states, timestep, encoder_hidden_states): - return self.model(hidden_states, timestep, encoder_hidden_states).to(self.io_dtype) + dtype = self.model_dtype + velocity = self.model(hidden_states.to(dtype), timestep.to(dtype), encoder_hidden_states.to(dtype)) + return velocity.to(self.io_dtype) @torch.no_grad() @@ -384,8 +388,8 @@ def main(): elif name == "transformer": print(f"Loading PixelDiT from {args.model_name} (dtype=bf16, {args.height}x{args.width})...") model = load_pixeldit(args.model_name, device, torch.bfloat16, args.local_files_only) - model.set_image_size(args.height, args.width) - export_transformer(model, output_path, device, model_opset, args.height, args.width, io_dtype) + with export_patches(model, args.height, args.width, model.txt_max_length): + export_transformer(model, output_path, device, model_opset, args.height, args.width, io_dtype) del model torch.cuda.empty_cache() exported_model_paths.append(output_path / name / "model.onnx") diff --git a/image_gen/pixeldit/model_export/pixeldit/__init__.py b/image_gen/pixeldit/model_export/pixeldit/__init__.py index dde4ab1..d0fedbb 100644 --- a/image_gen/pixeldit/model_export/pixeldit/__init__.py +++ b/image_gen/pixeldit/model_export/pixeldit/__init__.py @@ -1,3 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Standalone PixelDiT T2I model, weights loader and sampling pipeline for ONNX export.""" +"""PixelDiT T2I export helpers: upstream export patches, weights loader and sampling pipeline. + +The model itself is upstream's PixDiT_T2I from the `pixeldit` pip package.""" diff --git a/image_gen/pixeldit/model_export/pixeldit/model.py b/image_gen/pixeldit/model_export/pixeldit/model.py deleted file mode 100644 index 1aa1006..0000000 --- a/image_gen/pixeldit/model_export/pixeldit/model.py +++ /dev/null @@ -1,445 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""PixelDiT text-to-image DiT, vendored from NVlabs/PixelDiT (pixdit_core/) for ONNX export. - -Numerically equivalent to upstream PixDiTTrainer + PixDiT_T2I at inference. Changes: - * RoPE uses real cos/sin tables instead of complex tensors. Required: upstream unsqueezes - complex tensors (precompute_freqs_cis_2d, and freqs_cis[None, :, None, :] in - apply_rotary_emb), and the ONNX exporter has no lowering for unsqueeze on complex inputs. - Other complex ops (polar, view_as_complex/real, multiply) would export. - * Patchify / unpatchify use reshape + permute instead of F.unfold / F.fold (bit-identical, - since kernel == stride). Required for TensorRT RTX: the model runs them in bf16, and - TensorRT RTX has no bf16 Col2Im (F.fold); the unfold+fold export fails to create a session. - * Position tables are built once by set_image_size() instead of lazily filled Python dict - caches, so tracing never writes trace-time tensors into a cache. They are plain - attributes (not buffers) so Module.to(dtype) leaves the fp32 RoPE tables alone, as - upstream's complex64 tables are. -Training-only pieces (REPA projector, attention masks, class embedder) are dropped. -""" - -import math - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - -# Matches nvidia/PixelDiT-1300M-1024px config.json. -PIXELDIT_1300M_CONFIG = { - "in_channels": 3, - "patch_size": 16, - "num_groups": 24, - "hidden_size": 1536, - "pixel_hidden_size": 16, - "pixel_attn_hidden_size": 1152, - "pixel_num_groups": 16, - "patch_depth": 14, - "pixel_depth": 2, - "txt_embed_dim": 2304, - "txt_max_length": 300, - "text_rope_theta": 10000.0, -} - - -# --------------------------------------------------------------------------- -# Position tables (match pixdit_core/modules.py) -# --------------------------------------------------------------------------- - - -def _sincos_1d(embed_dim, pos): - omega = np.arange(embed_dim // 2, dtype=np.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega - out = np.einsum("m,d->md", pos.reshape(-1), omega) - return np.concatenate([np.sin(out), np.cos(out)], axis=1) - - -def pixel_sincos_pos_embed(embed_dim: int, height: int, width: int) -> np.ndarray: - """[H*W, D] absolute pixel embedding (get_2d_sincos_pos_embed[_from_grid] upstream).""" - grid_h = np.arange(height, dtype=np.float32) - grid_w = np.arange(width, dtype=np.float32) - grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape(2, 1, height, width) - emb_h = _sincos_1d(embed_dim // 2, grid[0]) - emb_w = _sincos_1d(embed_dim // 2, grid[1]) - return np.concatenate([emb_h, emb_w], axis=1) - - -def rope_2d_angles(dim: int, height: int, width: int, theta: float = 10000.0, scale: float = 16.0) -> torch.Tensor: - """[H*W, dim/2] angles; upstream precompute_freqs_cis_2d is torch.polar(1, angles).""" - x_pos = torch.linspace(0, scale, width) - y_pos = torch.linspace(0, scale, height) - y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij") - freqs = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) - x_freqs = torch.outer(x_pos.reshape(-1), freqs).float() - y_freqs = torch.outer(y_pos.reshape(-1), freqs).float() - return torch.stack([x_freqs, y_freqs], dim=-1).reshape(height * width, -1) - - -def rope_1d_angles(head_dim: int, length: int, theta: float, device=None) -> torch.Tensor: - """[L, head_dim/2] angles; upstream PixDiT_T2I.fetch_pos_text.""" - freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) - positions = torch.arange(0, length, device=device).float().unsqueeze(1) - return positions * freqs.unsqueeze(0) - - -def apply_rotary_emb(xq, xk, cos, sin): - """Real-valued equivalent of upstream complex apply_rotary_emb. - - xq/xk: [B, N, H, D] with (even, odd) pairs along D; cos/sin: [N, D/2]. - """ - cos = cos[None, :, None, :] - sin = sin[None, :, None, :] - - def rotate(x): - x = x.float().reshape(*x.shape[:-1], -1, 2) - x_even, x_odd = x[..., 0], x[..., 1] - out = torch.stack([x_even * cos - x_odd * sin, x_even * sin + x_odd * cos], dim=-1) - return out.flatten(3) - - return rotate(xq).type_as(xq), rotate(xk).type_as(xk) - - -# --------------------------------------------------------------------------- -# Building blocks (match pixdit_core/modules.py) -# --------------------------------------------------------------------------- - - -def apply_adaln(x, shift, scale): - return x * (1 + scale) + shift - - -class RMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - -class TimestepConditioner(nn.Module): - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, hidden_size, bias=True), - nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True), - ) - self.frequency_embedding_size = frequency_embedding_size - - @staticmethod - def timestep_embedding(t, dim, max_period=10): - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half - ) - args = t[..., None].float() * freqs[None, ...] - return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) - - -class FeedForward(nn.Module): - def __init__(self, dim: int, hidden_dim: int): - super().__init__() - hidden_dim = int(2 * hidden_dim / 3) - self.w1 = nn.Linear(dim, hidden_dim, bias=False) - self.w3 = nn.Linear(dim, hidden_dim, bias=False) - self.w2 = nn.Linear(hidden_dim, dim, bias=False) - - def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - -class MLP(nn.Module): - def __init__(self, dim: int, mlp_ratio: float = 4.0): - super().__init__() - hidden_dim = int(dim * mlp_ratio) - self.fc1 = nn.Linear(dim, hidden_dim) - self.act = nn.GELU() - self.fc2 = nn.Linear(hidden_dim, dim) - - def forward(self, x): - return self.fc2(self.act(self.fc1(x))) - - -class FinalLayer(nn.Module): - def __init__(self, hidden_size, out_channels): - super().__init__() - self.norm = RMSNorm(hidden_size, eps=1e-6) - self.linear = nn.Linear(hidden_size, out_channels, bias=True) - - def forward(self, x): - return self.linear(self.norm(x)) - - -class PatchTokenEmbedder(nn.Module): - def __init__(self, in_chans: int, embed_dim: int, norm_layer=None, bias: bool = True): - super().__init__() - self.proj = nn.Linear(in_chans, embed_dim, bias=bias) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - return self.norm(self.proj(x)) - - -class PixelTokenEmbedder(nn.Module): - """Per-pixel linear embedding + absolute sin/cos position, grouped into [B*L, P*P, D].""" - - def __init__(self, in_channels: int, hidden_size_output: int): - super().__init__() - self.hidden_size_output = int(hidden_size_output) - self.proj = nn.Linear(in_channels, self.hidden_size_output, bias=True) - - def forward(self, inputs, pos_full, patch_size: int): - B, _, H, W = inputs.shape - Hs, Ws = H // patch_size, W // patch_size - x = self.proj(inputs.permute(0, 2, 3, 1)) - x = x + pos_full.to(x.dtype).view(H, W, self.hidden_size_output).unsqueeze(0) - x = x.view(B, Hs, patch_size, Ws, patch_size, self.hidden_size_output) - x = x.permute(0, 1, 3, 2, 4, 5) - return x.reshape(B * Hs * Ws, patch_size * patch_size, self.hidden_size_output) - - -class RotaryAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): - super().__init__() - assert dim % num_heads == 0, "dim should be divisible by num_heads" - self.num_heads = num_heads - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.q_norm = RMSNorm(dim // num_heads) - self.k_norm = RMSNorm(dim // num_heads) - self.proj = nn.Linear(dim, dim) - - def forward(self, x, cos, sin): - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) - q, k, v = qkv[0], qkv[1], qkv[2] - q, k = apply_rotary_emb(self.q_norm(q), self.k_norm(k), cos, sin) - q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) - return self.proj(x.transpose(1, 2).reshape(B, N, C)) - - -class PiTBlock(nn.Module): - """Pixel-level block: per-pixel adaLN from patch tokens + attention over compressed patches.""" - - def __init__(self, pixel_hidden_size, patch_hidden_size, patch_size, attn_hidden_size, attn_num_heads, mlp_ratio=4.0): - super().__init__() - self.pixel_dim = int(pixel_hidden_size) - self.attn_dim = int(attn_hidden_size) - p2 = patch_size * patch_size - self.compress_to_attn = nn.Linear(p2 * self.pixel_dim, self.attn_dim, bias=True) - self.expand_from_attn = nn.Linear(self.attn_dim, p2 * self.pixel_dim, bias=True) - self.norm1 = RMSNorm(self.pixel_dim, eps=1e-6) - self.attn = RotaryAttention(self.attn_dim, num_heads=attn_num_heads, qkv_bias=False) - self.norm2 = RMSNorm(self.pixel_dim, eps=1e-6) - self.mlp = MLP(self.pixel_dim, mlp_ratio=mlp_ratio) - self.adaLN_modulation = nn.Sequential(nn.Linear(patch_hidden_size, 6 * self.pixel_dim * p2, bias=True)) - - def forward(self, x, s_cond, batch, num_patches, cos, sin): - BL, P2, _ = x.shape - cond = self.adaLN_modulation(s_cond).view(BL, P2, 6 * self.pixel_dim) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(cond, 6, dim=-1) - x_norm = apply_adaln(self.norm1(x), shift_msa, scale_msa) - x_comp = self.compress_to_attn(x_norm.view(BL, P2 * self.pixel_dim)).view(batch, num_patches, self.attn_dim) - attn_out = self.attn(x_comp, cos, sin) - attn_exp = self.expand_from_attn(attn_out.view(BL, self.attn_dim)).view(BL, P2, self.pixel_dim) - x = x + gate_msa * attn_exp - return x + gate_mlp * self.mlp(apply_adaln(self.norm2(x), shift_mlp, scale_mlp)) - - -class MMDiTJointAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): - super().__init__() - assert dim % num_heads == 0, "dim should be divisible by num_heads" - self.num_heads = num_heads - head_dim = dim // num_heads - self.qkv_x = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.qkv_y = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.q_norm_x = RMSNorm(head_dim) - self.k_norm_x = RMSNorm(head_dim) - self.q_norm_y = RMSNorm(head_dim) - self.k_norm_y = RMSNorm(head_dim) - self.proj_x = nn.Linear(dim, dim) - self.proj_y = nn.Linear(dim, dim) - - def forward(self, x, y, img_cos, img_sin, txt_cos, txt_sin): - B, Nx, C = x.shape - Ny = y.shape[1] - qkv_x = self.qkv_x(x).reshape(B, Nx, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) - qx, kx, vx = self.q_norm_x(qkv_x[0]), self.k_norm_x(qkv_x[1]), qkv_x[2] - qkv_y = self.qkv_y(y).reshape(B, Ny, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) - qy, ky, vy = self.q_norm_y(qkv_y[0]), self.k_norm_y(qkv_y[1]), qkv_y[2] - - qx, kx = apply_rotary_emb(qx, kx, img_cos, img_sin) - qy, ky = apply_rotary_emb(qy, ky, txt_cos, txt_sin) - - q = torch.cat([qy.transpose(1, 2), qx.transpose(1, 2)], dim=2) - k = torch.cat([ky.transpose(1, 2), kx.transpose(1, 2)], dim=2) - v = torch.cat([vy.transpose(1, 2), vx.transpose(1, 2)], dim=2) - out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) - - out_y = out[:, :, :Ny, :].transpose(1, 2).reshape(B, Ny, C) - out_x = out[:, :, Ny:, :].transpose(1, 2).reshape(B, Nx, C) - return self.proj_x(out_x), self.proj_y(out_y) - - -class MMDiTBlockT2I(nn.Module): - def __init__(self, hidden_size, groups, mlp_ratio=4.0): - super().__init__() - self.norm_x1 = RMSNorm(hidden_size, eps=1e-6) - self.norm_y1 = RMSNorm(hidden_size, eps=1e-6) - self.attn = MMDiTJointAttention(hidden_size, num_heads=groups, qkv_bias=False) - self.norm_x2 = RMSNorm(hidden_size, eps=1e-6) - self.norm_y2 = RMSNorm(hidden_size, eps=1e-6) - mlp_hidden_dim = int(hidden_size * mlp_ratio) - self.mlp_x = FeedForward(hidden_size, mlp_hidden_dim) - self.mlp_y = FeedForward(hidden_size, mlp_hidden_dim) - self.adaLN_modulation_img = nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) - self.adaLN_modulation_txt = nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) - - def forward(self, x, y, c, img_cos, img_sin, txt_cos, txt_sin): - shift_msa_x, scale_msa_x, gate_msa_x, shift_mlp_x, scale_mlp_x, gate_mlp_x = self.adaLN_modulation_img(c).chunk(6, dim=-1) - shift_msa_y, scale_msa_y, gate_msa_y, shift_mlp_y, scale_mlp_y, gate_mlp_y = self.adaLN_modulation_txt(c).chunk(6, dim=-1) - - x_norm = apply_adaln(self.norm_x1(x), shift_msa_x, scale_msa_x) - y_norm = apply_adaln(self.norm_y1(y), shift_msa_y, scale_msa_y) - attn_x, attn_y = self.attn(x_norm, y_norm, img_cos, img_sin, txt_cos, txt_sin) - x = x + gate_msa_x * attn_x - y = y + gate_msa_y * attn_y - - x = x + gate_mlp_x * self.mlp_x(apply_adaln(self.norm_x2(x), shift_mlp_x, scale_mlp_x)) - y = y + gate_mlp_y * self.mlp_y(apply_adaln(self.norm_y2(y), shift_mlp_y, scale_mlp_y)) - return x, y - - -# --------------------------------------------------------------------------- -# PixelDiT -# --------------------------------------------------------------------------- - - -class PixelDiT(nn.Module): - """PixDiT_T2I with upstream PixDiTTrainer input handling (inputs cast to the weight dtype). - - forward(x [B,3,H,W], timestep [B] (already scaled by 1000), y [B,L,txt_embed_dim]) -> [B,3,H,W] - velocity in the weight dtype. Call set_image_size(H, W) after moving/casting the model and - before forward/export. - """ - - def __init__( - self, - in_channels=3, - patch_size=16, - num_groups=24, - hidden_size=1536, - pixel_hidden_size=16, - pixel_attn_hidden_size=1152, - pixel_num_groups=16, - patch_depth=14, - pixel_depth=2, - txt_embed_dim=2304, - txt_max_length=300, - text_rope_theta=10000.0, - ): - super().__init__() - self.in_channels = int(in_channels) - self.patch_size = int(patch_size) - self.num_groups = int(num_groups) - self.hidden_size = int(hidden_size) - self.pixel_hidden_size = int(pixel_hidden_size) - self.pixel_attn_hidden_size = int(pixel_attn_hidden_size) - self.pixel_num_groups = int(pixel_num_groups) - self.txt_max_length = int(txt_max_length) - self.text_rope_theta = float(text_rope_theta) - - self.pixel_embedder = PixelTokenEmbedder(in_channels, self.pixel_hidden_size) - self.s_embedder = PatchTokenEmbedder(in_channels * patch_size**2, hidden_size, bias=True) - self.t_embedder = TimestepConditioner(hidden_size) - self.y_embedder = PatchTokenEmbedder(txt_embed_dim, hidden_size, bias=True, norm_layer=RMSNorm) - self.y_pos_embedding = nn.Parameter(torch.zeros(1, self.txt_max_length, hidden_size)) - self.patch_blocks = nn.ModuleList([MMDiTBlockT2I(hidden_size, num_groups) for _ in range(patch_depth)]) - self.pixel_blocks = nn.ModuleList( - [ - PiTBlock( - self.pixel_hidden_size, - hidden_size, - patch_size=self.patch_size, - attn_hidden_size=self.pixel_attn_hidden_size, - attn_num_heads=self.pixel_num_groups, - ) - for _ in range(pixel_depth) - ] - ) - self.final_layer = FinalLayer(self.pixel_hidden_size, self.in_channels) - self.image_size = None - - @property - def dtype(self) -> torch.dtype: - return self.final_layer.linear.weight.dtype - - def set_image_size(self, height: int, width: int) -> None: - """Build position tables for a fixed (height, width) on the model's device and dtype. - - Matches upstream: 2D RoPE angles on CPU, text RoPE on device, cos/sin in fp32, and the - float64 pixel sin/cos table cast straight to the weight dtype. - """ - if height % self.patch_size or width % self.patch_size: - raise ValueError(f"Image size {height}x{width} must be a multiple of patch_size={self.patch_size}") - device = self.final_layer.linear.weight.device - hs, ws = height // self.patch_size, width // self.patch_size - img = rope_2d_angles(self.hidden_size // self.num_groups, hs, ws) - pix = rope_2d_angles(self.pixel_attn_hidden_size // self.pixel_num_groups, hs, ws) - txt = rope_1d_angles(self.hidden_size // self.num_groups, self.txt_max_length, self.text_rope_theta, device) - pixel_pos = torch.from_numpy(pixel_sincos_pos_embed(self.pixel_hidden_size, height, width)) - self.img_cos, self.img_sin = img.cos().to(device), img.sin().to(device) - self.pix_cos, self.pix_sin = pix.cos().to(device), pix.sin().to(device) - self.txt_cos, self.txt_sin = txt.cos(), txt.sin() - self.pixel_pos = pixel_pos.to(device=device, dtype=self.dtype) - self.image_size = (height, width) - - def forward(self, x, timestep, y): - dtype = self.dtype - x = x.to(dtype) - timestep = timestep.to(dtype) - y = y.to(dtype) - - B, _, H, W = x.shape - if self.image_size != (H, W): - raise RuntimeError(f"Input is {H}x{W} but set_image_size() was {self.image_size}") - if self.pixel_pos.dtype != dtype or self.pixel_pos.device != x.device: - raise RuntimeError("Model was moved or cast after set_image_size(); call it again") - p = self.patch_size - Hs, Ws = H // p, W // p - L = Hs * Ws - - x_patches = x.view(B, self.in_channels, Hs, p, Ws, p).permute(0, 2, 4, 1, 3, 5).reshape(B, L, -1) - t_emb = self.t_embedder(timestep.view(-1)).view(B, -1, self.hidden_size) - - Ltxt = min(y.shape[1], self.txt_max_length) - y = y[:, :Ltxt, :] - y_emb = self.y_embedder(y).view(B, Ltxt, self.hidden_size) - y_emb = y_emb + self.y_pos_embedding[:, :Ltxt, :].to(y_emb.dtype) - - condition = F.silu(t_emb) - s = self.s_embedder(x_patches) - txt_cos, txt_sin = self.txt_cos[:Ltxt], self.txt_sin[:Ltxt] - for block in self.patch_blocks: - s, y_emb = block(s, y_emb, condition, self.img_cos, self.img_sin, txt_cos, txt_sin) - s = F.silu(t_emb + s) - - s_cond = s.view(B * L, self.hidden_size) - x_pixels = self.pixel_embedder(x, self.pixel_pos, p) - for block in self.pixel_blocks: - x_pixels = block(x_pixels, s_cond, B, L, self.pix_cos, self.pix_sin) - - x_pixels = self.final_layer(x_pixels) - x_pixels = x_pixels.view(B, Hs, Ws, p, p, self.in_channels).permute(0, 5, 1, 3, 2, 4) - return x_pixels.reshape(B, self.in_channels, H, W) diff --git a/image_gen/pixeldit/model_export/pixeldit/upstream_patch.py b/image_gen/pixeldit/model_export/pixeldit/upstream_patch.py new file mode 100644 index 0000000..71b40a4 --- /dev/null +++ b/image_gen/pixeldit/model_export/pixeldit/upstream_patch.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Export-time patches for upstream PixelDiT T2I (pip package `pixeldit`, imported as `pixdit_core`). + +The DiT is upstream's `pixdit_core.pixeldit_t2i.PixDiT_T2I`, installed from NVlabs/PixelDiT at UPSTREAM_COMMIT +(the `pixeldit` extra in the repository's root pyproject.toml). Upstream runs as-is everywhere except inside +`export_patches(...)`, which makes three changes for the duration of the ONNX export and restores everything on exit: + + * RoPE uses real (cos, sin) pairs 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. Matches upstream up to float rounding (last-bit differences in the rotation tables and in the + rotation itself). + * F.unfold / F.fold (patchify / unpatchify, kernel == stride) become reshape + permute, which is bit-identical. + Required for TensorRT RTX: its ONNX parser rejects Col2Im (what F.fold exports to) in every dtype, and in bf16 + ONNX Runtime's CPU provider has no Col2Im to fall back to either. + * Position tables are computed eagerly for the export size and cached before tracing, so the exporter sees + constants and never writes traced tensors into upstream's lazily filled dict caches. + +The patches replace module-level functions and torch.nn.functional.fold/unfold process-wide while active, so the +export must not run concurrently with other code using them. +""" + +import contextlib +import importlib.metadata +import json +import warnings +from typing import Tuple + +import torch +import torch.nn.functional as F + +import pixdit_core.modules as upstream_modules +import pixdit_core.pixeldit_c2i as upstream_c2i +import pixdit_core.pixeldit_t2i as upstream_t2i + +# Commit of NVlabs/PixelDiT these patches were written and validated against (keep in sync with the `pixeldit` extra +# in the root pyproject.toml). +UPSTREAM_COMMIT = "f7d7c35f71893d81c6e71d2576ea80601a1e8529" + + +# --------------------------------------------------------------------------- +# Real-valued RoPE: rotations as [N, head_dim // 2, 2] (cos, sin), the layout of torch.view_as_real +# --------------------------------------------------------------------------- + + +def _rotations(angles: torch.Tensor) -> torch.Tensor: + return torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1) + + +def precompute_freqs_cis_2d(dim: int, height: int, width: int, theta: float = 10000.0, scale=16.0): + """Upstream modules.precompute_freqs_cis_2d, returning (cos, sin) pairs instead of complex numbers.""" + 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() + # [N, dim / 4, 2 (x, y), 2 (cos, sin)] -> [N, dim / 2, 2], alternating x and y like upstream's complex table. + rotations = torch.stack([_rotations(x_freqs), _rotations(y_freqs)], dim=-2) + return rotations.reshape(height * width, -1, 2) + + +def apply_rotary_emb( + xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Upstream modules.apply_rotary_emb with real arithmetic: (re + i im) * (cos + i sin).""" + freqs_cis = freqs_cis[None, :, None, :, :] + cos, sin = freqs_cis[..., 0], freqs_cis[..., 1] + + def rotate(x: torch.Tensor) -> torch.Tensor: + x_ = x.float().reshape(*x.shape[:-1], -1, 2) + real, imag = x_[..., 0], x_[..., 1] + return torch.stack([real * cos - imag * sin, real * sin + imag * cos], dim=-1).flatten(3).type_as(x) + + return rotate(xq), rotate(xk) + + +def fetch_pos_text(self, length, device): + """Upstream PixDiT_T2I.fetch_pos_text, building (cos, sin) pairs instead of torch.polar.""" + if length in self.precompute_pos_txt: + return self.precompute_pos_txt[length].to(device) + head_dim = self.hidden_size // self.num_groups + freqs = 1.0 / (self.text_rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) + angles = torch.arange(0, length, device=device).float().unsqueeze(1) * freqs.unsqueeze(0) + self.precompute_pos_txt[length] = _rotations(angles) + return self.precompute_pos_txt[length] + + +# --------------------------------------------------------------------------- +# F.unfold / F.fold for non-overlapping patches (kernel == stride) as reshape + permute +# --------------------------------------------------------------------------- + + +def _patch_size(kernel_size, stride, kwargs) -> int: + kernel = kernel_size if isinstance(kernel_size, int) else kernel_size[0] + step = stride if isinstance(stride, int) else stride[0] + square = isinstance(kernel_size, int) or kernel_size[0] == kernel_size[1] + if not square or kernel != step or kwargs.get("padding", 0) not in (0, (0, 0)) or kwargs.get( + "dilation", 1 + ) not in (1, (1, 1)): + raise NotImplementedError("export patch only supports unfold/fold with square kernel == stride, no padding") + return kernel + + +def unfold(x, kernel_size, stride=1, **kwargs): + """F.unfold for kernel == stride: [B, C, H, W] -> [B, C * p * p, L], same element order.""" + p = _patch_size(kernel_size, stride, kwargs) + B, C, H, W = x.shape + return x.reshape(B, C, H // p, p, W // p, p).permute(0, 1, 3, 5, 2, 4).reshape(B, C * p * p, (H // p) * (W // p)) + + +def fold(cols, output_size, kernel_size, stride=1, **kwargs): + """F.fold for kernel == stride: [B, C * p * p, L] -> [B, C, H, W], same result.""" + p = _patch_size(kernel_size, stride, kwargs) + H, W = output_size + B = cols.shape[0] + C = cols.shape[1] // (p * p) + return cols.reshape(B, C, p, p, H // p, W // p).permute(0, 1, 4, 2, 5, 3).reshape(B, C, H, W) + + +# --------------------------------------------------------------------------- +# Patch context +# --------------------------------------------------------------------------- + + +def installed_upstream_commit() -> str | None: + """Commit of the installed `pixeldit` package, if it was installed from git; None otherwise.""" + try: + direct_url = importlib.metadata.distribution("pixeldit").read_text("direct_url.json") + except importlib.metadata.PackageNotFoundError: + return None + if not direct_url: + return None + return json.loads(direct_url).get("vcs_info", {}).get("commit_id") + + +def _require(obj, name: str) -> None: + if not hasattr(obj, name): + owner = getattr(obj, "__name__", type(obj).__name__) + raise RuntimeError( + f"upstream PixelDiT changed: {owner}.{name} not found; update pixeldit/upstream_patch.py " + f"(written against NVlabs/PixelDiT@{UPSTREAM_COMMIT[:7]})" + ) + + +def _clear_caches(model) -> None: + model.precompute_pos.clear() + model.precompute_pos_txt.clear() + for block in model.pixel_blocks: + block._pos_cache.clear() + model.last_repa_tokens = None + + +@contextlib.contextmanager +def export_patches(model: upstream_t2i.PixDiT_T2I, height: int, width: int, txt_len: int): + """Patch upstream PixDiT_T2I for export at a fixed image size and text length; restores everything on exit.""" + commit = installed_upstream_commit() + if commit is not None and commit != UPSTREAM_COMMIT: + warnings.warn( + f"pixeldit is installed from NVlabs/PixelDiT@{commit[:7]}, but these export patches were validated " + f"against @{UPSTREAM_COMMIT[:7]}; re-run the PixelDiT export validation after changing the version" + ) + + changes = [ + (upstream_modules, "apply_rotary_emb", apply_rotary_emb), # RotaryAttention (PiT blocks) + (upstream_t2i, "apply_rotary_emb", apply_rotary_emb), # MMDiTJointAttention, imported by name + (upstream_modules, "precompute_freqs_cis_2d", precompute_freqs_cis_2d), + (upstream_t2i, "precompute_freqs_cis_2d", precompute_freqs_cis_2d), # PixDiT_T2I.fetch_pos + (upstream_c2i, "precompute_freqs_cis_2d", precompute_freqs_cis_2d), + (model, "fetch_pos_text", fetch_pos_text.__get__(model)), + (F, "unfold", unfold), # PixDiT_T2I.forward calls torch.nn.functional.unfold / fold inline + (F, "fold", fold), + ] + changes += [(block, "_rope_fn", precompute_freqs_cis_2d) for block in model.pixel_blocks] # bound at init + for obj, name, _ in changes: + _require(obj, name) + for name in ("precompute_pos", "precompute_pos_txt", "fetch_pos", "pixel_embedder", "last_repa_tokens"): + _require(model, name) + for block in model.pixel_blocks: + _require(block, "_pos_cache") + _require(block, "_fetch_pos") + _require(model.pixel_embedder, "_fetch_pixel_pos_image") + + # (obj, name, previous value, whether obj itself held the attribute, e.g. not a class method) + saved = [(obj, name, getattr(obj, name), name in vars(obj)) for obj, name, _ in changes] + try: + for obj, name, value in changes: + setattr(obj, name, value) + _clear_caches(model) # drop complex tables from earlier eager runs + # Fill the caches with real tensors so tracing reads constants instead of writing traced tensors. + parameter = next(model.parameters()) + p = model.patch_size + model.fetch_pos(height // p, width // p, parameter.device) + model.fetch_pos_text(min(txt_len, model.txt_max_length), parameter.device) + for block in model.pixel_blocks: + block._fetch_pos(height // p, width // p, parameter.device) + if model.use_pixel_abs_pos: + model.pixel_embedder._fetch_pixel_pos_image(height, width, parameter.device, parameter.dtype) + yield model + finally: + for obj, name, value, own in reversed(saved): + if own: + setattr(obj, name, value) + else: + delattr(obj, name) # e.g. the instance override of the PixDiT_T2I.fetch_pos_text method + _clear_caches(model) # upstream rebuilds its own complex tables on the next eager call diff --git a/image_gen/pixeldit/model_export/pixeldit/weights.py b/image_gen/pixeldit/model_export/pixeldit/weights.py index 2481335..50c3a17 100644 --- a/image_gen/pixeldit/model_export/pixeldit/weights.py +++ b/image_gen/pixeldit/model_export/pixeldit/weights.py @@ -2,13 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """Load the PixelDiT T2I checkpoint and the Gemma-2 text encoder.""" +import inspect import json from pathlib import Path import torch from huggingface_hub import hf_hub_download - -from .model import PIXELDIT_1300M_CONFIG, PixelDiT +from pixdit_core.pixeldit_t2i import PixDiT_T2I DEFAULT_MODEL_NAME = "nvidia/PixelDiT-1300M-1024px" CHECKPOINT_FILE = "pixeldit_t2i_v1.pth" @@ -20,7 +20,7 @@ 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.""" + """Return (checkpoint .pth path, PixDiT_T2I constructor arguments) 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" @@ -32,10 +32,11 @@ def resolve_checkpoint(model_name: str, local_files_only: bool = False) -> tuple 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}) + if not config_path.is_file(): + raise FileNotFoundError(f"PixelDiT config.json not found next to the checkpoint: {config_path}") + hf_config = json.loads(config_path.read_text(encoding="utf-8")) + parameters = inspect.signature(PixDiT_T2I.__init__).parameters + config = {k: v for k, v in hf_config.items() if k in parameters} return checkpoint.resolve(), config @@ -56,14 +57,26 @@ def load_pixeldit( 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).""" +) -> PixDiT_T2I: + """Build PixDiT_T2I, load weights strictly, cast to dtype.""" checkpoint, config = resolve_checkpoint(model_name, local_files_only) - model = PixelDiT(**config) + model = PixDiT_T2I(**config) model.load_state_dict(core_state_dict(checkpoint), strict=True) return model.eval().requires_grad_(False).to(device=device, dtype=dtype) +def dit_callable(model: PixDiT_T2I): + """dit(x, timestep, y) for pipeline.py: casts the inputs to the model dtype first, like upstream PixDiTTrainer. + + The bf16 cast of the timestep is part of the reference behaviour (999.75 rounds to 1000).""" + dtype = next(model.parameters()).dtype + + def dit(x, timestep, y): + return model(x.to(dtype), timestep.to(dtype), y.to(dtype)) + + return dit + + def load_gemma_text_encoder( model_name: str = DEFAULT_TEXT_ENCODER_NAME, device: str = "cuda", diff --git a/image_gen/pixeldit/model_export/test_upstream_patch.py b/image_gen/pixeldit/model_export/test_upstream_patch.py new file mode 100644 index 0000000..1edef84 --- /dev/null +++ b/image_gen/pixeldit/model_export/test_upstream_patch.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for pixeldit/upstream_patch.py on a small random upstream PixDiT_T2I (CPU, no checkpoint needed). + +Run from this directory: python -m pytest test_upstream_patch.py (or: python test_upstream_patch.py) +""" + +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import onnx +import onnxruntime as ort +import torch +from pixdit_core.pixeldit_t2i import PixDiT_T2I + +import pixdit_core.modules as upstream_modules +import torch.nn.functional as F +from export_pixeldit import PixelDiTExportWrapper +from pixeldit.upstream_patch import export_patches + +HEIGHT, WIDTH, TXT_LEN, TXT_DIM = 96, 128, 12, 64 + + +def small_model() -> PixDiT_T2I: + torch.manual_seed(0) + model = PixDiT_T2I(in_channels=3, num_groups=4, hidden_size=128, pixel_hidden_size=8, pixel_attn_hidden_size=96, + pixel_num_groups=4, patch_depth=2, pixel_depth=2, patch_size=16, txt_embed_dim=TXT_DIM, + txt_max_length=TXT_LEN, repa_encoder_index=1).eval().requires_grad_(False) + with torch.no_grad(): + for p in model.parameters(): # initialize_weights zeroes the final layer; make outputs non-trivial + p.normal_(0, 0.05) + return model + + +def inputs(): + g = torch.Generator().manual_seed(1) + x = torch.randn(2, 3, HEIGHT, WIDTH, generator=g) + t = torch.tensor([999.75, 500.0]) + y = torch.randn(2, TXT_LEN, TXT_DIM, generator=g) + return x, t, y + + +def rel_l2(a, b) -> float: + a, b = np.asarray(a, np.float64), np.asarray(b, np.float64) + return float(np.linalg.norm(a - b) / np.linalg.norm(a)) + + +class UpstreamPatchTest(unittest.TestCase): + def test_export_runs_and_matches(self): + model, (x, t, y) = small_model(), inputs() + with torch.no_grad(): + reference = model(x, t, y) # unpatched upstream + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "dit.onnx" + with export_patches(model, HEIGHT, WIDTH, TXT_LEN): + patched = model(x, t, y) + torch.onnx.export(PixelDiTExportWrapper(model, torch.float32).eval(), (x, t, y), str(path), + dynamo=True, opset_version=25, + input_names=["hidden_states", "timestep", "encoder_hidden_states"], + output_names=["velocity"]) + ops = {node.op_type for node in onnx.load(str(path)).graph.node} + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + exported = session.run(None, {"hidden_states": x.numpy(), "timestep": t.numpy(), + "encoder_hidden_states": y.numpy()})[0] + after = model(x, t, y) + + self.assertNotIn("Col2Im", ops) + self.assertLess(rel_l2(patched, exported), 1e-6) # ONNX Runtime vs patched PyTorch (fp32) + self.assertLess(rel_l2(reference, patched), 1e-6) # patched vs upstream: RoPE rounding only + self.assertTrue(torch.equal(reference, after)) # upstream restored exactly + + def test_everything_is_restored(self): + model, (x, t, y) = small_model(), inputs() + rope_fns = [block._rope_fn for block in model.pixel_blocks] + apply_rotary_emb, fold, unfold = upstream_modules.apply_rotary_emb, F.fold, F.unfold + with torch.no_grad(), export_patches(model, HEIGHT, WIDTH, TXT_LEN): + self.assertIsNot(F.fold, fold) + self.assertEqual(next(iter(model.precompute_pos.values())).dtype, torch.float32) # real tables + self.assertIs(upstream_modules.apply_rotary_emb, apply_rotary_emb) + self.assertIs(F.fold, fold) + self.assertIs(F.unfold, unfold) + self.assertEqual([block._rope_fn for block in model.pixel_blocks], rope_fns) + self.assertNotIn("fetch_pos_text", vars(model)) # class method again, no instance override + self.assertEqual(len(model.precompute_pos), 0) # upstream rebuilds its complex tables + with torch.no_grad(): + model(x, t, y) + self.assertTrue(next(iter(model.precompute_pos.values())).is_complex()) + + def test_guard_reports_upstream_changes(self): + model = small_model() + del model.pixel_blocks[0]._rope_fn # as if upstream renamed it + with self.assertRaisesRegex(RuntimeError, "upstream PixelDiT changed: .*_rope_fn"): + with export_patches(model, HEIGHT, WIDTH, TXT_LEN): + pass + self.assertEqual(upstream_modules.apply_rotary_emb.__module__, "pixdit_core.modules") # nothing was patched + + def test_unsupported_fold_arguments_are_rejected(self): + model = small_model() + with export_patches(model, HEIGHT, WIDTH, TXT_LEN): + with self.assertRaises(NotImplementedError): + F.unfold(torch.zeros(1, 3, 32, 32), kernel_size=4, stride=2) + + +if __name__ == "__main__": + unittest.main() diff --git a/image_gen/pixeldit/model_export/verify_pixeldit.py b/image_gen/pixeldit/model_export/verify_pixeldit.py index 06d215e..797fe93 100644 --- a/image_gen/pixeldit/model_export/verify_pixeldit.py +++ b/image_gen/pixeldit/model_export/verify_pixeldit.py @@ -3,8 +3,8 @@ #!/usr/bin/env python3 """ -Verify PixelDiT T2I ONNX graphs: run the vendored PyTorch pipeline and the ONNX Runtime pipeline -(TensorRT RTX by default) and compare them. +Verify PixelDiT T2I ONNX graphs: run the PyTorch pipeline (PixDiT_T2I from the `pixeldit` package, +unpatched) and the ONNX Runtime pipeline (TensorRT RTX by default) and compare them. Checks: --validate_only ONNX structure + dummy run of each graph, then exit @@ -34,7 +34,13 @@ 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 +from pixeldit.weights import ( + DEFAULT_MODEL_NAME, + DEFAULT_TEXT_ENCODER_NAME, + dit_callable, + load_gemma_text_encoder, + load_pixeldit, +) EP_NAME = "nv_tensorrt_rtx" _REGISTERED_EP_NAME: str | None = None @@ -408,9 +414,8 @@ def main(): 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) + print(f"Loading PixelDiT from {args.model_name} (bf16, upstream pixeldit package)...") + dit = dit_callable(load_pixeldit(args.model_name, device, torch.bfloat16, args.local_files_only)) encode_torch = torch_encode_fn(text_encoder) te_runner, tr_runner = runners.get("text_encoder"), runners.get("transformer") diff --git a/pyproject.toml b/pyproject.toml index 54f617f..957b36f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,12 @@ flux = [ "pillow", "transformers", ] +pixeldit = [ + # Upstream PixelDiT model definitions (imported as pixdit_core), pinned to the commit the export patches in + # image_gen/pixeldit/model_export/pixeldit/upstream_patch.py (UPSTREAM_COMMIT) were validated against. + "pixeldit @ git+https://github.com/NVlabs/PixelDiT@f7d7c35f71893d81c6e71d2576ea80601a1e8529", + "transformers", +] [project.scripts] parakeet-tdt = "rnnt.parakeet_tdt.cli:main"