SAT-JEPA-DIFF-Caption-Guided

Sat-JEPA-Diff (Caption-Guided Zero-RGB)

Caption-Guided Zero-RGB Satellite Image Forecasting via Self-Supervised Diffusion

Submitted to IEEE Geoscience and Remote Sensing Letters (GRSL)

[

Model Description

This is the Caption-Guided Zero-RGB version of Sat-JEPA-Diff. It advances the original model by completely eliminating the 32×32 coarse RGB input, replacing it with a Triple-Captioning system that textually drives temporal changes.

The model predicts future satellite images (t → t+1) through a dual-stream mechanism:

1. I-JEPA spatial anchor encodes the current image and predicts future latent embeddings — dictating where changes manifest
2. Triple-Captioning system provides hierarchical text conditioning — describing what changes:
   • Informative caption (t): dense pixel-level description
   • Geometric caption (t): spatial layout and topology
   • Semantic caption (t+1): forecasted future state via a learned Caption Forecaster
3. Multi-Stream Conditioning Adapter fuses all four signals (3 caption embeddings + IJEPA tokens) via cross-attention
4. Frozen Stable Diffusion 3.5 Medium + LoRA generates the final high-fidelity RGB imagery

This Zero-RGB paradigm bypasses physical sensor blindness — the model never sees past pixels during generation, relying entirely on semantic and latent representations.

What Changed from the Original Model

Key Results

Despite using zero optical inputs, the caption-guided model achieves 15% better perceptual realism (LPIPS) while maintaining topological integrity (GSSIM).

Architecture Details

How to Use

Note: This is a custom PyTorch checkpoint, not a standard transformers model. You need the source code from the GitHub repository to load and run inference.

1. Setup

git clone https://github.com/VU-AIML/SAT-JEPA-DIFF.git
cd SAT-JEPA-DIFF
git checkout caption-guided

conda create -n satjepa python=3.12
conda activate satjepa

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install diffusers transformers peft accelerate
pip install rasterio matplotlib pyyaml lpips

You need a Hugging Face token with access to Stable Diffusion 3.5 Medium:

export HF_TOKEN=hf_your_token_here

2. Download the Checkpoint

from huggingface_hub import hf_hub_download

checkpoint_path = hf_hub_download(
    repo_id="kursatkomurcu/SAT-JEPA-DIFF-Caption-Guided",
    filename="s2_future_jepa-best.pth.tar",
)

Or via CLI:

huggingface-cli download kursatkomurcu/SAT-JEPA-DIFF-Caption-Guided \
    s2_future_jepa-best.pth.tar --local-dir ./checkpoints

3. Run Inference

cd SAT-JEPA-DIFF/src

python inference.py \
    --checkpoint /path/to/s2_future_jepa-best.pth.tar \
    --output_dir ./results \
    --diffusion_steps 20 \
    --noise_strength 0.35

4. Programmatic Usage

import torch
from helper import init_model
from sd_models import load_sd_model, encode_caption_batch
from sd_joint_loss import diffusion_sample
from caption_forecaster import CaptionForecaster

device = torch.device("cuda")

# Load all components from checkpoint
encoder, predictor, sd_state, embed_dim = load_model(
    checkpoint_path="path/to/s2_future_jepa-best.pth.tar",
    device=device,
)

# Load caption forecaster
caption_forecaster = CaptionForecaster(
    ijepa_dim=768, text_dim=4096, hidden_dim=1024, num_layers=3
).to(device)

# Load a Sentinel-2 GeoTIFF
rgb_t = load_and_resize_tif("path/to/sentinel2_image.tif", target_size=128)

# Predict with caption conditioning
rgb_t1_pred = predict_next_frame(
    rgb_t=rgb_t,
    encoder=encoder,
    predictor=predictor,
    sd_state=sd_state,
    caption_forecaster=caption_forecaster,
    device=device,
    num_diffusion_steps=20,
    noise_strength=0.35,
)
# rgb_t1_pred: (3, 128, 128) tensor in [0, 1]

Checkpoint Contents

The .pth.tar file contains:

Caption Pipeline

The model requires precomputed caption embeddings. The full pipeline is documented in the GitHub repository:

1. Generate captions with EarthDial (generate_captions.py)
2. Precompute embeddings with frozen CLIP+T5 (precompute_caption_embeddings.py)
3. Train or run inference — text encoders are never loaded during training/inference

Training

See the GitHub repository for full training instructions. The model jointly optimizes three losses:

ℒ_total = ℒ_IJEPA + λ_sd · ℒ_diff + λ_cap · ℒ_caption

Training takes approximately 5 days on a single NVIDIA RTX 5090 (24GB).

Dataset

Sentinel-2 RGB imagery (10m GSD) paired with Alpha Earth Foundation Model embeddings and EarthDial captions across 100 global Regions of Interest (2017–2024). Available on Zenodo.

Requirements

• Python 3.12+
• PyTorch 2.0+ with CUDA
• ~24GB GPU VRAM (RTX 3090/4090/5090 or A100)
• Stable Diffusion 3.5 Medium access via HF token

Citation

@ARTICLE{11589351,
  author={Kömürcü, Kürşat and Petkevicius, Linas},
  journal={IEEE Geoscience and Remote Sensing Letters}, 
  title={Sat-JEPA-Diff: Caption-Guided Zero-RGB Satellite Image Forecasting via Self-Supervised Diffusion}, 
  year={2026},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/LGRS.2026.3708464}
}
@inproceedings{komurcu2026satjepadiff,
  title={Sat-{JEPA}-Diff: Bridging Self-Supervised Learning and Generative Diffusion for Remote Sensing},
  author={Kursat Komurcu and Linas Petkevicius},
  booktitle={4th ICLR Workshop on Machine Learning for Remote Sensing (Main Track)},
  year={2026},
  url={https://openreview.net/forum?id=WBHfQLbgZR}
}

Acknowledgments

This project was funded by the European Union (project No S-MIP-23-45) under the agreement with the Research Council of Lithuania (LMTLT). The I-JEPA implementation is based on Meta's I-JEPA. Diffusion backbone: SD 3.5 Medium. Captions: EarthDial.