| --- |
| 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`. |
|
|
