File size: 3,319 Bytes
04c42c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | ---
pipeline_tag: image-to-image
tags:
- pytorch
- hirise
- mars
- remote-sensing
- flow-matching
---
# 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](https://github.com/EricKing22/HiPredict).
## 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:
```bash
pip install ".[diffusion]"
```
For evaluation with HiPredict's shared metrics, install `.[diffusion,postprocess]` instead.
## Download and load
```python
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.
```python
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:
```bash
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`.
|