syding9/HiPredict_Diffusion

HiPredict Diffusion

Bidirectional conditional flow-matching checkpoints for reconstructing HiRISE RED4 from IR10, or IR10 from RED4. Use these checkpoints with the HiPredict Python package.

Checkpoints

File	Model	Training steps	Gradient-loss weight
FM_Condition/latest.pt	BidirectionalFM, Condition mode	100,000	0.0
FM_Condition_gradient/latest.pt	BidirectionalFM, Condition mode	100,000	0.1
brightness/model.pt	Auxiliary RED3 boundary-residual brightness head	—	—

Choose either FM checkpoint as the main model. Both support ir2red and red2ir. The brightness head is optional: the default boundary-based correction does not need it.

Both main checkpoints include model weights, EMA weights, optimizer state, learning-rate scheduler state, and model/training settings. Each file is approximately 2.05 GB. The brightness head is approximately 32 KB.

Installation

Use a HiPredict checkout containing hipredict.diffusion and install its diffusion dependencies from the checkout root:

pip install ".[diffusion]"

For evaluation with HiPredict's shared metrics, install .[diffusion,postprocess] instead.

Download and load

from huggingface_hub import hf_hub_download
from hipredict.diffusion import load_checkpoint

checkpoint = hf_hub_download(
    repo_id="syding9/HiPredict_Diffusion",
    filename="FM_Condition/latest.pt",
)
model, model_config, training_config, step = load_checkpoint(
    checkpoint, device="cuda"
)

Select FM_Condition_gradient/latest.pt for the gradient-loss model.

HiPredict API

The HiPredict API adapter reconstructs RED4 (ir2red), on patches or full strips. Large images are predicted with overlapping tiles.

from hipredict.diffusion import build_hipredict_predictor

predictor, config, step = build_hipredict_predictor(checkpoint, device="cuda")
samples = predictor.load("files/data_record_evals.csv", mode="patches")
predictions = predictor.predict(samples)

For red2ir, use the diffusion inference command with the downloaded checkpoint path:

python -m hipredict.diffusion.inference.inference     --ckpt_path /path/to/latest.pt     --input_path observation_RED4.npy     --red3_path observation_RED3.npy     --direction red2ir     --output_dir outputs/inference/

Brightness correction

The default correction uses the RED3 boundary. To use the learned residual brightness head for dataset inference or evaluation, download brightness/model.pt, select --dc_correct head, and pass its local path as --head_ckpt.

Model configuration

Both FM checkpoints use a 128-channel base width, three residual blocks per level, Sobel edge conditioning, and noise-conditioned flow matching. Their normalization settings are stored in the checkpoint:

• Center: 0.103687465
• Scale: 0.04842728
• Normalized training values clipped to [-5, 5]

HiPredict reads these settings when loading the checkpoint. Evaluation uses original, unclipped physical targets. The checkpoint format is hipredict.diffusion, version 1.

File SHA-256 digests are recorded in checksums.json.