Janus-first2last

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.

Software

numpy>=1.24
torch>=2.1
transformers>=4.39
huggingface_hub>=0.20
safetensors>=0.4

Papers and Scientific Outputs

@article{Chu2026DiffUNet2,
  title = {DiffUNet^2: Bidirectional Prediction, Probabilistic Generation and Collaborative Visual Discovery for Scientific Data},
  author = {Chu, Mengdi and Yang, Jiaxin and Forbes, Angus G. and Debardeleben, Nathan and Lawrence, Earl and Biswas, Ayan and Shen, Han-Wei},
  year = {2026},
  doi = {https://doi.org/10.48550/arXiv.2606.03926}
}

Associated paper: https://doi.org/10.48550/arXiv.2606.03926

Model License

Other.

Contact Info and Model Card Authors

Mengdi Chu, chu.752@osu.edu

Intended Uses

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.

Install Instructions

pip install numpy torch transformers huggingface_hub safetensors

Training configuration

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.

• Source checkpoint: DDPM_first2last/first2last_latest.pt
• Published weights: model.safetensors
• 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]
• Conditional UNet input channels: 2
• Conditional UNet output channels: 1
• Base channels: 64
• Timestep embedding dimension: 256
• DDPM training timesteps: 1000
• Beta schedule: linear from 1e-4 to 0.02
• Batch size: 8
• Planned training epochs: 800
• Published checkpoint epoch: 533
• Optimizer: AdamW
• Learning rate: 1e-4
• Learning-rate scheduler: CosineAnnealingLR, T_max=800
• Gradient clipping: global norm 1.0
• Checkpointing: save every 2 epochs and whenever validation loss improves
• Training objective: MSE between predicted noise and sampled Gaussian noise
• Conditioning field: HEAT-PLI av_density
• Target field: HEAT-PLI av_density

Inference configuration

• Default sampling method: DDIM
• Default sampling steps: 50
• Default eta: 1.0
• Optional ensemble sampling: n_samples > 1
• Optional fixed random seed: seed
• Recommended loading mode: AutoModel.from_pretrained(repo_id, trust_remote_code=True)

Code snippets of how to use the model

Library usage from the Hub:

import numpy as np
from transformers import AutoModel

repo_id = "YOUR_USERNAME/Janus-first2last"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
model = model.to("cuda")  # optional

sample = np.load("examples/id01040/pli240420_id01040_pvi_idx00000.npz")
density = sample["av_density"]
prediction = model.predict(density, num_steps=50, eta=1.0, seed=42)
mean, std = model.predict_ensemble(density, n_samples=5, num_steps=50, eta=1.0)

Command-line inference after cloning or downloading the repository:

python api_ddpm.py \
  --model . \
  --input examples/id01040/pli240420_id01040_pvi_idx00000.npz \
  --output outputs/id01040_first2last_pred.npz \
  --field av_density \
  --num_steps 50 \
  --eta 1.0

Command-line inference directly from the Hub:

python api_ddpm.py \
  --model YOUR_USERNAME/Janus-first2last \
  --input examples/id01040/pli240420_id01040_pvi_idx00000.npz \
  --output outputs/id01040_first2last_pred.npz \
  --field av_density \
  --num_steps 50 \
  --eta 1.0

Limitations

Risks

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.
• Learning-rate scheduler: CosineAnnealingLR, T_max=800.
• Gradient clipping: global norm 1.0.
• Training loss: 0.000961887405287.
• Prompting templates/functions: not applicable.
• 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.

More Information (optional)

None.