ONNX Image Model
Run a PIX2PIX image to image model in ONNX format. Inference runs through ONNX Runtime Web on the GPU and does not block the interface: the node keeps showing its last output until the next result is ready.
The node reads the required input and output dimensions from the ONNX model metadata. Input images must match the model’s expected size, so put a Resize node in front of it.
The PIX2PIX tutorial walks through preparing a dataset in Figment, training with the script below, and running the result live.
To work well, it needs an input that is similar to what it has seen in training; e.g. if you trained it on a hand model, you need to feed it hand models with the same color and line thickness.
Parameters
- Model The .onnx model file.
- Optimize Convert the model for this GPU on load and use the converted copy. Off by default.
Optimizing a model
Two things make a model faster in Figment without retraining it: float16 weights and activations, which halve the file and speed up most GPUs, and a rewrite of stride-2 ConvTranspose layers, whose WebGPU kernel is several times slower than a regular convolution. Constant chains left behind by the exporter are folded as well.
With Optimize on, the node converts the model on first load and checks the result against the original on the same input: the rewritten graph has to reproduce the original to rounding, and the float16 version has to stay close to it. A conversion that fails either check is not used. The result is written next to the model as <name>.figment-optimized.onnx with a .json file that records what was done and why. Later loads use that copy as long as the model file, the converter and the GPU’s float16 support are unchanged; a retrained model or a machine without shader-f16 gets a fresh conversion. A rejected conversion leaves the original in use and a note in the console.
File > Optimize ONNX Model… runs the same conversion on a file you choose and shows the report: sizes, layers rewritten, and how closely the result matches the original. Use it to produce a file to hand to someone else, or to see what the option would do before turning it on.
Float16 models
A model with float16 weights and activations (for instance one converted with onnxconverter-common, keeping the input and output in float32) loads like any other and runs faster on most GPUs: about 1.5× for a pix2pix U-Net. Figment asks the GPU for the shader-f16 feature when it has one. On a GPU without it, a float16 model fails on its first frame with an error on the node that says so; keep the float32 model as the fallback for those machines.
During an export (File > Render, or --render) every frame waits for its own inference, so the rendered frames match their input. Live, the node keeps showing its last result while the next inference runs.
Checking that the model runs on the GPU
The node keeps the image on the GPU: the input texture is packed into a GPU buffer, ONNX Runtime runs the model on the WebGPU execution provider, and the output buffer is unpacked into a texture. That only holds when every operator in the model has a WebGPU kernel. ONNX Runtime silently runs any operator without one on the CPU and copies the activations to the CPU and back around it, on every frame. A single such operator can make the model several times slower.
The generator from the script below, and the ones from the figmentapp/pix2pix notebooks (U-Net with instance normalization, pix2pixHD with reflection padding), export to operators that all have WebGPU kernels: Conv, ConvTranspose, BatchNormalization or InstanceNormalization, LeakyRelu, Relu, Tanh, Concat and Pad. Operators to avoid when you change the architecture: PRelu, Selu and Softplus (so also Mish), and anything else ONNX Runtime has no WebGPU kernel for in the build Figment ships.
The reliable check is to load the model in Figment itself. From a checkout of the Figment repository, with npm install done:
node scripts/check-onnx-webgpu.mjs output/generator_epoch_14.onnxThe script opens Figment in a headless browser, creates a session for the model and reports whether every node landed on the WebGPU provider, or which operators fell back to the CPU and how many GPU to CPU copies that costs per run. Run it once on a freshly exported model, and again after every change to the network.
Training
Here’s a simple example of how to train a PIX2PIX model using PyTorch. You will need torch, torchvision, onnx and onnxruntime installed.
pip install torch==2.4.0 torchvision==0.19.0 onnx==1.16.1 onnxruntime==1.19.0Run the code with python train.py --input_dir datasets/trees --output_dir output.
import glob
import torch
import torch.nn as nn
import torch.optim as optim
import torch.onnx
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from PIL import Image
import os
import random
import argparse
from tqdm import tqdm
# Create the dataset class
class Pix2PixDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_files = [
f for f in os.listdir(root_dir) if f.endswith(".jpg") or f.endswith(".png")
]
def __len__(self):
return len(self.image_files)
def __getitem__(self, idx):
img_name = os.path.join(self.root_dir, self.image_files[idx])
image = Image.open(img_name)
# Split the image into input and target
w, h = image.size
target_image = image.crop((0, 0, w // 2, h))
input_image = image.crop((w // 2, 0, w, h))
if self.transform:
input_image = self.transform(input_image)
target_image = self.transform(target_image)
return input_image, target_image
# Implement the UNet architecture for the generator
class UNetBlock(nn.Module):
def __init__(self, in_channels, out_channels, down=True, bn=True, dropout=False):
super(UNetBlock, self).__init__()
self.conv = (
nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False)
if down
else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False)
)
self.bn = nn.BatchNorm2d(out_channels) if bn else None
self.dropout = nn.Dropout(0.5) if dropout else None
self.act = nn.LeakyReLU(0.2) if down else nn.ReLU()
def forward(self, x):
x = self.conv(x)
if self.bn:
x = self.bn(x)
if self.dropout:
x = self.dropout(x)
return self.act(x)
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.down1 = UNetBlock(3, 64, down=True, bn=False)
self.down2 = UNetBlock(64, 128)
self.down3 = UNetBlock(128, 256)
self.down4 = UNetBlock(256, 512)
self.down5 = UNetBlock(512, 512)
self.down6 = UNetBlock(512, 512)
self.down7 = UNetBlock(512, 512)
self.down8 = UNetBlock(512, 512, bn=False)
self.up1 = UNetBlock(512, 512, down=False, dropout=True)
self.up2 = UNetBlock(1024, 512, down=False, dropout=True)
self.up3 = UNetBlock(1024, 512, down=False, dropout=True)
self.up4 = UNetBlock(1024, 512, down=False)
self.up5 = UNetBlock(1024, 256, down=False)
self.up6 = UNetBlock(512, 128, down=False)
self.up7 = UNetBlock(256, 64, down=False)
self.final = nn.Sequential(nn.ConvTranspose2d(128, 3, 4, 2, 1), nn.Tanh())
def forward(self, x):
d1 = self.down1(x)
d2 = self.down2(d1)
d3 = self.down3(d2)
d4 = self.down4(d3)
d5 = self.down5(d4)
d6 = self.down6(d5)
d7 = self.down7(d6)
d8 = self.down8(d7)
u1 = self.up1(d8)
u2 = self.up2(torch.cat([u1, d7], 1))
u3 = self.up3(torch.cat([u2, d6], 1))
u4 = self.up4(torch.cat([u3, d5], 1))
u5 = self.up5(torch.cat([u4, d4], 1))
u6 = self.up6(torch.cat([u5, d3], 1))
u7 = self.up7(torch.cat([u6, d2], 1))
return self.final(torch.cat([u7, d1], 1))
# Implement the discriminator
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
UNetBlock(6, 64, bn=False),
UNetBlock(64, 128),
UNetBlock(128, 256),
UNetBlock(256, 512),
nn.Conv2d(512, 1, 4, 1, 1),
nn.Sigmoid(),
)
def forward(self, x, y):
return self.model(torch.cat([x, y], 1))
# Define the loss functions and optimizers
criterion_gan = nn.BCELoss()
criterion_pixel = nn.L1Loss()
# Load snapshot if available
def get_latest_snapshot(output_dir):
snapshots = glob.glob(os.path.join(output_dir, "snapshot_epoch_*.pth"))
if not snapshots:
return None
return max(snapshots, key=os.path.getctime)
def load_snapshot(generator, discriminator, g_optimizer, d_optimizer, snapshot_path):
checkpoint = torch.load(snapshot_path, map_location=device, weights_only=False)
generator.load_state_dict(checkpoint["generator"])
discriminator.load_state_dict(checkpoint["discriminator"])
g_optimizer.load_state_dict(checkpoint["g_optimizer"])
d_optimizer.load_state_dict(checkpoint["d_optimizer"])
start_epoch = int(os.path.basename(snapshot_path).split("_")[2].split(".")[0])
return start_epoch
# 6. Create the training loop
def train(generator, discriminator, dataloader, args):
g_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
# Get fixed input/output for visualization
fixed_set = next(iter(dataloader))
fixed_input = fixed_set[0][0].unsqueeze(0)
fixed_target = fixed_set[1][0].unsqueeze(0)
# fixed_input = next(iter(dataloader))[0][0].unsqueeze(0) # Get a fixed input for visualization
start_epoch = 0
if not args.restart:
latest_snapshot = get_latest_snapshot(args.output_dir)
if latest_snapshot:
start_epoch = load_snapshot(
generator, discriminator, g_optimizer, d_optimizer, latest_snapshot
)
print(f"Resuming training from epoch {start_epoch}")
else:
print("No snapshots found. Starting from scratch.")
else:
print("Restarting training from scratch.")
for epoch in range(args.epochs):
for i, (input_img, target_img) in enumerate(tqdm(dataloader)):
input_img = input_img.to(device)
target_img = target_img.to(device)
# Train Discriminator
d_optimizer.zero_grad()
fake_img = generator(input_img)
d_real = discriminator(input_img, target_img)
d_fake = discriminator(input_img, fake_img.detach())
d_loss_real = criterion_gan(d_real, torch.ones_like(d_real))
d_loss_fake = criterion_gan(d_fake, torch.zeros_like(d_fake))
d_loss = (d_loss_real + d_loss_fake) / 2
d_loss.backward()
d_optimizer.step()
# Train Generator
g_optimizer.zero_grad()
fake_img = generator(input_img)
d_fake = discriminator(input_img, fake_img)
g_loss_gan = criterion_gan(d_fake, torch.ones_like(d_fake))
g_loss_pixel = criterion_pixel(fake_img, target_img) * 100
g_loss = g_loss_gan + g_loss_pixel
g_loss.backward()
g_optimizer.step()
if i % args.sample_interval == 0:
with torch.no_grad():
fake_img = generator(fixed_input.to(device))
img_sample = torch.cat(
(fixed_input.cpu(), fake_img.cpu(), fixed_target.cpu()), -1
)
save_image(
img_sample,
f"{args.output_dir}/epoch_{epoch}_iter_{i}.jpg",
nrow=3,
normalize=True,
)
if (epoch + 1) % args.snapshot_interval == 0:
torch.save(
{
"generator": generator.state_dict(),
"discriminator": discriminator.state_dict(),
"g_optimizer": g_optimizer.state_dict(),
"d_optimizer": d_optimizer.state_dict(),
},
f"{args.output_dir}/snapshot_epoch_{epoch + 1}.pth",
)
# Save to ONNX format
onnx_path = f"{args.output_dir}/generator_epoch_{epoch + 1}.onnx"
generator.eval()
dummy_input = torch.randn(1, 3, 512, 512).to(device)
traced_script_module = torch.jit.trace(generator, dummy_input)
torch.onnx.export(
traced_script_module,
dummy_input,
onnx_path,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
)
print(f"ONNX model exported to {onnx_path}")
generator.train()
# 7. Implement the argument parser for configuration
def parse_args():
parser = argparse.ArgumentParser(
description="Conditional GAN with pix2pix architecture"
)
parser.add_argument(
"--input_dir",
type=str,
default="datasets/trees",
help="Input dataset directory",
)
parser.add_argument(
"--output_dir",
type=str,
default="output",
help="Output directory for generated images",
)
parser.add_argument(
"--sample_interval",
type=int,
default=100,
help="Interval for saving sample images",
)
parser.add_argument(
"--snapshot_interval",
type=int,
default=1,
help="Interval for saving model snapshots",
)
parser.add_argument(
"--epochs", type=int, default=200, help="Number of epochs to train"
)
parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
parser.add_argument(
"--restart", action="store_true", help="Restart training from scratch"
)
return parser.parse_args()
# 8. Set up the main function to run the training
def main():
args = parse_args()
os.makedirs(args.output_dir, exist_ok=True)
transform = transforms.Compose(
[
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
dataset = Pix2PixDataset(args.input_dir, transform=transform)
dataloader = DataLoader(
dataset, batch_size=args.batch_size, shuffle=True, num_workers=4
)
generator = Generator().to(device)
discriminator = Discriminator().to(device)
train(generator, discriminator, dataloader, args)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
main()