Janus-first2last is a conditional DDPM/DDIM model for HEAT-PLI av_density field prediction. This checkpoint is the first-to-last prediction variant: condition on the first HEAT-PLI frame at simulation time 0.0 and predict the last HEAT-PLI frame at simulation time 25.0.
This Hugging Face release packages the model as custom Transformers code with AutoModel support. The original PyTorch .pt checkpoint was converted to model.safetensors so the public artifact contains model weights only, not optimizer state.
Last Updated: 2026-06-24
Model Changelog
2026-06-24 Hugging Face release package with AutoModel support and model.safetensors weights.
2026-05-22 Initial Janus-first2last model card for the DDPM checkpoint package.
Model short description
Conditional DDPM model with DDIM sampling for HEAT-PLI density-field first-to-last prediction.
Model description
Janus-first2last uses a conditional UNet diffusion model implemented in modeling_janus_ddpm.py. The model conditions on one HEAT-PLI av_density frame and generates the corresponding target frame by DDIM sampling.
The model operates directly on density fields rather than VAE latents. At inference time, the input field is resized to the model internal resolution (560, 200), min-max normalized per sample to [-1, 1], used as the condition channel, and decoded back to the input field resolution after sampling.
Finetuned from model (optional)
Not applicable. This checkpoint is not documented as a fine-tune, adapter, quantized model, or merge of an existing base model.
Model Type
Conditional DDPM/DDIM UNet for scientific image-to-image field prediction. The repository includes custom Transformers code and requires trust_remote_code=True when loaded through AutoModel.
Inputs and outputs
Input field: HEAT-PLI av_density, provided as a NumPy array of shape (H, W) or as an .npz file containing the av_density field.
Demo data shape in this package: (1120, 400) with float16 storage; three demo cases are included under examples/id01040, examples/id01045, and examples/id01050.
Public training data source: https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/.
Internal model resolution: (560, 200).
Output field: predicted av_density array with the same spatial resolution as the input.
Direction: condition on the first HEAT-PLI frame at simulation time 0.0; predict the last HEAT-PLI frame at simulation time 25.0.
Compute Infrastructure
The reported training and evaluation runs were performed on an 8x NVIDIA A100 GPU server.
Hardware
Use A100 GPUs with at least 30 GB of GPU memory for comparable training and evaluation runs. Inference can run on CPU or a single GPU, but GPU inference is recommended.
1@article{Chu2026DiffUNet2,
2 title = {DiffUNet^2: Bidirectional Prediction, Probabilistic Generation and Collaborative Visual Discovery for Scientific Data},
3 author = {Chu, Mengdi and Yang, Jiaxin and Forbes, Angus G. and Debardeleben, Nathan and Lawrence, Earl and Biswas, Ayan and Shen, Han-Wei},
4 year = {2026},
5 doi = {https://doi.org/10.48550/arXiv.2606.03926}
6}
This section describes intended research uses for the Janus-first2last checkpoint and documents use cases that are out of scope.
Intended Use
This model is intended for research use as a learned surrogate for HEAT-PLI density-field first-to-last prediction. It can be used for fast exploratory prediction, uncertainty sampling through repeated DDIM runs, and model-card/API demonstration on the packaged HEAT-PLI demo data.
Primary Intended Users
Researchers and engineers studying scientific machine learning, surrogate modeling, visual analytics, and uncertainty-aware prediction for physics simulation data.
Mission Relevance
This project is funded by ArtIMis.
Out-of-Scope Use Cases
This model is not intended for operational safety decisions, real-time hardware control, or extrapolation to HEAT configurations, fields, timesteps, or parameter ranges not represented in the training and evaluation data. Generated fields should be validated against domain constraints and high-fidelity simulations before scientific or engineering conclusions are drawn.
How to use
Use AutoModel.from_pretrained(..., trust_remote_code=True) or the packaged api_ddpm.py CLI.
Training configuration is reconstructed from the model implementation, source checkpoint metadata, and the private training notebook used to train this checkpoint. The notebook itself is not included in this Hugging Face release.
Public source dataset: https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/
Training field: av_density
Pairing strategy: first available frame in each PLI case folder as condition, last available frame as target
Data split: folder-level random split, 70% train, 15% validation, 15% test
Split seed: 42
Input preprocessing: load av_density, cast to float32, resize from (1120, 400) to (560, 200) with bilinear interpolation, per-sample min-max normalize to [-1, 1]
This is a domain-specific scientific surrogate model. The primary risks are over-trusting generated fields, applying the model outside the HEAT-PLI distribution, or using predictions without domain validation.
Limitations
The model predicts av_density only.
The API uses per-sample min-max normalization, so absolute calibration should be checked for downstream analysis.
The model resizes fields to (560, 200) internally and resizes predictions back to the input resolution.
Stochastic DDIM sampling can produce different outputs unless a fixed seed is provided.
Performance outside the represented HEAT-PLI first/last-frame prediction setting has not been established in this card.
Training details
In this section, training details are documented from the available checkpoint metadata and model implementation.
Training data
The model was trained on HEAT-PLI av_density endpoint pairs from the HEAT dataset. The public PLI fp16 dataset is available at https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/; the general HEAT dataset page is https://oceans11.lanl.gov/heat/, with DOI 10.25583/2571471.
The training code scanned case folders containing .npz files, kept folders with at least two frames, and used the first sorted frame as the conditioning frame and the last sorted frame as the target frame. The packaged demo data include three small case folders, examples/id01040, examples/id01045, and examples/id01050, each with two frames: idx00000 at simulation time 0.0 and idx00100 at simulation time 25.0. These examples are included only for smoke testing and inference demonstrations; use the public HEAT URL for the full training data.
Training Procedure
The source checkpoint metadata records:
epoch: 533
train loss: 0.000961887405287
validation loss: 0.00133784566235
best validation loss: 0.00101556984417
parameters: 10,722,177
model state tensors: 118
The architecture uses a conditional UNet with two input channels, one output channel, base channel width 64, sinusoidal timestep embeddings with dimension 256, three downsampling blocks, a middle attention block, and three upsampling blocks. The scheduler uses 1000 DDPM timesteps with beta range 1e-4 to 0.02; inference uses DDIM sampling, defaulting to 50 steps.
During training, each batch samples a random diffusion timestep, adds Gaussian noise to the normalized target frame, and trains the UNet to predict that noise conditioned on the normalized first frame. Validation uses the same denoising loss procedure without gradient updates.
Reproducibility Information (optional)
Random seed used for data split: 42
Machine/environment info: GPUs with at least 30 GB of GPU memory for comparable training runs.
Public training data URL: https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/
Link to evaluation or inference pipeline: api_ddpm.py
Published weights: model.safetensors
Demo data: examples/id01040, examples/id01045, and examples/id01050
Exact train/validation/test folder lists are not included in this release; reproduce the split by applying the documented seed and 70/15/15 folder-level split to the public PLI fp16 dataset.
Pre-training information
Hyperparameter tuning: not specified in checkpoint metadata.
Model initialization: random initialization from the PyTorch module defaults used by the training notebook.
Fine-tuning, chain-of-thought, and n-shot learning: not applicable.
Optimizer information: AdamW, learning rate 1e-4; optimizer state existed in the source .pt checkpoint but is intentionally not included in the Hugging Face release artifact.
Loss function: denoising diffusion training loss, reported as training and validation MSE-style loss in checkpoint metadata.
Stopping criterion: planned training length was 800 epochs; this card reports the published source checkpoint at epoch 533 and best validation loss stored in that checkpoint.
Number of training epochs: 533 checkpoint epoch from a planned 800 epoch run.
Number of training steps: not specified in checkpoint metadata.
Batch size: 8.
Training speed: not specified in checkpoint metadata.
Model optimization techniques: no distillation, quantization, pruning, or sparsity method is documented for this checkpoint.
Evaluation details
Evaluation details are documented from the packaged demo examples and checkpoint validation metadata.
Evaluation data
The packaged demo examples use HEAT-PLI folders id01040, id01045, and id01050, with av_density arrays of shape (1120, 400). The input example is examples/id01040/pli240420_id01040_pvi_idx00000.npz, and the corresponding target endpoint is examples/id01040/pli240420_id01040_pvi_idx00100.npz.
Evaluation Procedure
Run the model on the first frame and compare the generated result to the last frame. The training notebook used denoising validation loss during training and also included exploratory DDIM evaluations with PSNR and SSIM for sampled predictions. DDIM step counts explored in the notebook included 10, 25, 50, 70, 100, and 150; the published API defaults to 50 steps.
Uncertainty Quantification.
Uncertainty/variability can be estimated by running predict_ensemble or the CLI --n_samples option. The API returns an ensemble mean and standard deviation across stochastic DDIM samples. The default single-sample inference does not provide uncertainty estimates.
Evaluation results
Checkpoint validation MSE from source training metadata:
Model
Direction
Validation loss
Best validation loss
Janus-first2last
first-to-last prediction
0.00133784566235
0.00101556984417
Evaluation runtime depends on device, field size, --num_steps, and --n_samples. The reported runs used an A100 GPU.
