rsdiff 0.1.0__py3-none-any.whl

rsdiff/__init__.py

@@ -0,0 +1,3 @@
+"""RSDiff: text-to-satellite-image diffusion."""
+
+__version__ = "0.1.0"

rsdiff/cascade.py

@@ -0,0 +1,184 @@
+"""Build, load and sample the Imagen-style text-to-image cascade.
+
+The architecture comes from a :class:`~rsdiff.config.CascadeConfig`; weights come from an
+``ImagenTrainer`` checkpoint such as the released ``rsdiff1.5`` file. Sampling goes through
+``ImagenTrainer.sample`` so the EMA weights are used, exactly as in the published evaluation.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import random
+from collections.abc import Callable, Sequence
+from dataclasses import asdict
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+
+from rsdiff import __version__
+from rsdiff.config import CascadeConfig, RunConfig, UnetArch
+
+if TYPE_CHECKING:
+    from imagen_pytorch import Imagen, ImagenTrainer
+
+RELEASE_REPO = "asebaq/rsdiff-sr-cascade-ep650"
+RELEASE_REVISION = "0ac4d46c3aac32f4fcfb40831fcaf83d5e520fb0"
+RELEASE_FILENAME = "ckpt_sr_ep650_step89050.pt"
+RELEASE_SHA256 = "1c7721d5c32bd2fdebd938302435269b9a4f5ddb3691ec8f5226c11d110b09a1"
+
+SYNTHETIC_NOTICE = (
+    "Synthetic image generated by RSDiff from a text prompt. "
+    "It is not an observation or prediction of any real place."
+)
+
+# Callable with the ImagenTrainer.sample keywords we use: texts, cond_scale,
+# stop_at_unet_number, return_pil_images. Tests substitute a fake.
+Sampler = Callable[..., list[Image.Image]]
+
+
+def seed_everything(seed: int) -> None:
+    """Seed Python, NumPy and torch, matching the seeding used for the release."""
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def sha256_file(path: str | Path, chunk_size: int = 1 << 24) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as f:
+        while chunk := f.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify_checkpoint(path: str | Path, sha256: str = RELEASE_SHA256) -> None:
+    """Raise ``ValueError`` unless the file at ``path`` has the expected SHA-256."""
+    actual = sha256_file(path)
+    if actual != sha256:
+        raise ValueError(f"checksum mismatch for {path}: expected {sha256}, got {actual}")
+
+
+def download_release(cache_dir: str | Path | None = None) -> Path:
+    """Download the released ``rsdiff1.5`` checkpoint (1.9 GB) and verify its digest."""
+    from huggingface_hub import hf_hub_download
+
+    path = Path(
+        hf_hub_download(
+            repo_id=RELEASE_REPO,
+            filename=RELEASE_FILENAME,
+            revision=RELEASE_REVISION,
+            cache_dir=cache_dir,
+        )
+    )
+    verify_checkpoint(path)
+    return path
+
+
+def build_unet(arch: UnetArch):
+    """Construct one standalone ``imagen_pytorch.Unet`` stage from ``arch``."""
+    from imagen_pytorch import Unet
+
+    kwargs = {k: tuple(v) if isinstance(v, list) else v for k, v in asdict(arch).items()}
+    del kwargs["optimizer"], kwargs["params_m"]
+    return Unet(**kwargs)
+
+
+def build_imagen(cascade: CascadeConfig, text_embed_dim: int | None = None) -> Imagen:
+    """Construct the (untrained) cascade described by ``cascade``.
+
+    ``text_embed_dim`` skips looking up the text encoder's width on the Hugging Face Hub;
+    leave it ``None`` when sampling from text.
+    """
+    from imagen_pytorch import Imagen
+
+    unets = [build_unet(cascade.base)]
+    if cascade.sr is not None:
+        unets.append(build_unet(cascade.sr))
+    image_sizes = tuple(cascade.image_sizes[: len(unets)])
+    return Imagen(
+        text_encoder_name=cascade.text_encoder,
+        text_embed_dim=text_embed_dim,
+        unets=tuple(unets) if len(unets) > 1 else unets[0],
+        image_sizes=image_sizes if len(unets) > 1 else image_sizes[0],
+        timesteps=cascade.timesteps,
+        cond_drop_prob=cascade.cond_drop_prob,
+    )
+
+
+def load_cascade(cfg: RunConfig, checkpoint: str | Path, cpu: bool = False) -> ImagenTrainer:
+    """Build the cascade for ``cfg`` and load ``checkpoint`` into an ``ImagenTrainer``.
+
+    Runs on CUDA when available, otherwise on CPU. Apple MPS is not used: imagen-pytorch
+    2.1.0 casts to float64 while sampling, which MPS does not support.
+    """
+    from imagen_pytorch import ImagenTrainer
+
+    cpu = cpu or not torch.cuda.is_available()
+    trainer = ImagenTrainer(imagen=build_imagen(cfg.cascade), verbose=False, accelerate_cpu=cpu)
+    trainer.load(str(checkpoint))
+    return trainer
+
+
+def _slug(text: str, n: int = 40) -> str:
+    s = "".join(ch if ch.isalnum() else "_" for ch in text.strip().lower())
+    return "_".join(filter(None, s.split("_")))[:n]
+
+
+def output_paths(captions: Sequence[str], out_dir: Path) -> list[Path]:
+    """PNG path per caption; zero-padded so sorted file order equals caption order."""
+    pad = max(2, len(str(len(captions) - 1)))
+    return [out_dir / f"{i:0{pad}d}_{_slug(c)}.png" for i, c in enumerate(captions)]
+
+
+def generate(
+    sample: Sampler,
+    captions: Sequence[str],
+    out_dir: str | Path,
+    *,
+    cond_scale: float = 5.0,
+    batch_size: int = 2,
+    stop_at_unet: int | None = None,
+    on_batch: Callable[[int, int], None] | None = None,
+) -> list[Path]:
+    """Sample one image per caption into ``out_dir`` and return the PNG paths.
+
+    Writes ``captions.txt`` (one caption per line, in file order) before sampling. A batch
+    whose PNGs all exist is skipped, so an interrupted run resumes where it stopped. Every
+    PNG carries the prompt and a synthetic-content notice in its text chunks.
+    """
+    if not captions:
+        raise ValueError("no captions to sample")
+    if any(not c.strip() for c in captions):
+        raise ValueError("captions must be non-empty")
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "captions.txt").write_text("\n".join(captions) + "\n", encoding="utf-8")
+    paths = output_paths(captions, out_dir)
+
+    for start in range(0, len(captions), batch_size):
+        batch_caps = list(captions[start : start + batch_size])
+        batch_paths = paths[start : start + batch_size]
+        if all(p.exists() for p in batch_paths):
+            continue
+        kwargs: dict = {"texts": batch_caps, "cond_scale": cond_scale, "return_pil_images": True}
+        if stop_at_unet is not None:
+            kwargs["stop_at_unet_number"] = stop_at_unet
+        images = sample(**kwargs)
+        for path, caption, image in zip(batch_paths, batch_caps, images, strict=True):
+            info = PngInfo()
+            info.add_text("Software", f"rsdiff {__version__}")
+            info.add_text("Description", SYNTHETIC_NOTICE)
+            info.add_text("Prompt", caption)
+            image.save(path, pnginfo=info)
+        if on_batch is not None:
+            on_batch(min(start + batch_size, len(captions)), len(captions))
+    return paths

rsdiff/cli.py

@@ -0,0 +1,137 @@
+"""``rsdiff`` command line: sample the released cascade and score generation bundles."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+from rsdiff import __version__
+from rsdiff.config import DEFAULT_PRESET, PRESETS
+
+
+def cmd_sample(args: argparse.Namespace) -> int:
+    from rsdiff.cascade import download_release, generate, load_cascade, seed_everything
+    from rsdiff.config import resolve_config
+    from rsdiff.datasets.rsicd import sample_captions
+
+    if args.prompt:
+        captions = args.prompt * (args.n or 1)
+    else:
+        captions = sample_captions(args.rsicd, split=args.split, n=args.n, seed=args.seed)
+    cfg = resolve_config(args.config)
+
+    if args.checkpoint is None:
+        if args.config != DEFAULT_PRESET:
+            sys.exit("--checkpoint is required unless --config is the released rsdiff1.5")
+        print("fetching released checkpoint (1.9 GB on first use)", file=sys.stderr)
+        checkpoint = download_release()
+    else:
+        checkpoint = args.checkpoint
+
+    seed_everything(args.seed)
+    trainer = load_cascade(cfg, checkpoint, cpu=args.cpu)
+    print(f"sampling {len(captions)} image(s) on {trainer.device}", file=sys.stderr)
+    paths = generate(
+        trainer.sample,
+        captions,
+        args.out_dir,
+        cond_scale=args.cond_scale,
+        batch_size=args.batch_size,
+        stop_at_unet=1 if args.lr_only else None,
+        on_batch=lambda done, total: print(f"  {done}/{total}", file=sys.stderr),
+    )
+    print(f"wrote {len(paths)} PNGs to {args.out_dir}", file=sys.stderr)
+    return 0
+
+
+def cmd_eval_fid(args: argparse.Namespace) -> int:
+    from rsdiff.eval.fid import fid
+
+    res = fid(
+        gen_dir=args.gen_dir,
+        real_csv=args.rsicd / "dataset_rsicd.csv",
+        real_dir=args.rsicd / "RSICD_images",
+        split=args.split,
+        feature=args.feature,
+        image_size=args.size,
+        max_n=args.max_n,
+    )
+    print(json.dumps(asdict(res), indent=2))
+    return 0
+
+
+def cmd_eval_clip(args: argparse.Namespace) -> int:
+    from rsdiff.eval.clip import clip_score
+
+    print(json.dumps(asdict(clip_score(args.gen_dir, model_id=args.model)), indent=2))
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="rsdiff", description="Text-to-satellite-image diffusion (RSDiff)."
+    )
+    p.add_argument("--version", action="version", version=f"rsdiff {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    s = sub.add_parser(
+        "sample",
+        help="Generate images from text with the cascade.",
+        description="Generate synthetic 256x256 images. Downloads the released "
+        "rsdiff1.5 checkpoint unless --checkpoint is given.",
+    )
+    src = s.add_mutually_exclusive_group(required=True)
+    src.add_argument("--prompt", action="append", help="Caption to render (repeatable).")
+    src.add_argument("--rsicd", type=Path, help="RSICD root; sample captions from its CSV instead.")
+    s.add_argument("--split", default="test", help="RSICD split for --rsicd (default: test).")
+    s.add_argument(
+        "-n",
+        type=int,
+        default=None,
+        help="Images per --prompt (default 1), or captions drawn from --rsicd (default: all).",
+    )
+    s.add_argument("--out-dir", type=Path, default=Path("generated"))
+    s.add_argument(
+        "--config",
+        default=DEFAULT_PRESET,
+        help=f"Preset ({', '.join(PRESETS)}) or YAML path (default: {DEFAULT_PRESET}).",
+    )
+    s.add_argument("--checkpoint", type=Path, help="ImagenTrainer checkpoint to load.")
+    s.add_argument("--cond-scale", type=float, default=5.0, help="CFG scale (default: 5).")
+    s.add_argument("--seed", type=int, default=17)
+    s.add_argument("--batch-size", type=int, default=2)
+    s.add_argument("--lr-only", action="store_true", help="Stop at the 128x128 base UNet.")
+    s.add_argument("--cpu", action="store_true", help="Force CPU even if a GPU is available.")
+    s.set_defaults(fn=cmd_sample)
+
+    e = sub.add_parser("eval", help="Score a generation bundle.")
+    esub = e.add_subparsers(dest="metric", required=True)
+
+    f = esub.add_parser("fid", help="FID against the real RSICD split.")
+    f.add_argument("gen_dir", type=Path, help="Directory of generated PNGs.")
+    f.add_argument("--rsicd", type=Path, required=True, help="RSICD root (CSV + RSICD_images).")
+    f.add_argument("--split", default="test")
+    f.add_argument("--feature", type=int, choices=[64, 192, 768, 2048], default=2048)
+    f.add_argument("--size", type=int, default=256, help="Resize both sides to this first.")
+    f.add_argument("--max-n", type=int, default=None, help="Cap images per side (debug only).")
+    f.set_defaults(fn=cmd_eval_fid)
+
+    c = esub.add_parser("clip", help="CLIP score with a shuffled-caption baseline.")
+    c.add_argument("gen_dir", type=Path, help="Bundle written by `rsdiff sample`.")
+    c.add_argument(
+        "--model", default="openai/clip-vit-base-patch32", help="Hugging Face CLIP model id."
+    )
+    c.set_defaults(fn=cmd_eval_clip)
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    return args.fn(args)
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

rsdiff/config.py

@@ -0,0 +1,107 @@
+"""Cascade configuration and the packaged presets (``rsdiff1.5``, ``rsdiff1``)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from importlib import resources
+from pathlib import Path
+from typing import Any
+
+from omegaconf import OmegaConf
+
+PRESETS = ("rsdiff1.5", "rsdiff1")
+DEFAULT_PRESET = "rsdiff1.5"
+
+
+@dataclass
+class DataConfig:
+    name: str = "rsicd"
+    root: str = "data/RSICD_optimal"
+    image_size: int = 128
+    num_workers: int = 4
+
+
+@dataclass
+class UnetArch:
+    """imagen-pytorch ``Unet`` hyperparameters for one cascade stage.
+
+    ``dim`` through ``layer_cross_attns`` are passed to ``imagen_pytorch.Unet`` unchanged.
+    ``num_resnet_blocks`` is an int or a per-stage tuple. ``optimizer`` and ``params_m``
+    record the training recipe and the measured parameter count (millions).
+    """
+
+    dim: int = 128
+    cond_dim: int = 256
+    dim_mults: tuple[int, ...] = (1, 2, 2, 2)
+    num_resnet_blocks: Any = 0
+    layer_attns: tuple[bool, ...] = (False, True, True, True)
+    layer_cross_attns: tuple[bool, ...] = (False, True, True, True)
+    optimizer: str = "adam"  # adam | adafactor
+    params_m: float | None = None
+
+
+@dataclass
+class CascadeConfig:
+    """Imagen-style text-conditioned cascade: a base UNet plus an optional SR UNet."""
+
+    image_sizes: tuple[int, ...] = (128, 256)
+    timesteps: int = 1000
+    cond_drop_prob: float = 0.1
+    text_encoder: str = "t5-base"
+    base: UnetArch = field(default_factory=UnetArch)
+    sr: UnetArch | None = None
+
+
+@dataclass
+class OptimConfig:
+    lr: float = 1.0e-4
+    warmup_steps: int = 10000
+    weight_decay: float = 0.0
+
+
+@dataclass
+class TrainConfig:
+    seed: int = 42
+    output_dir: str = "outputs/run"
+    batch_size: int = 64
+    mixed_precision: str = "bf16"  # no | fp16 | bf16
+
+
+@dataclass
+class MetaConfig:
+    """Preset provenance: model line, parameter count, expected and measured FID."""
+
+    version: str = ""
+    reflects: str = ""
+    params_m: float | None = None
+    fid_expected: float | None = None
+    fid_measured: float | None = None
+    notes: str = ""
+
+
+@dataclass
+class RunConfig:
+    data: DataConfig = field(default_factory=DataConfig)
+    cascade: CascadeConfig = field(default_factory=CascadeConfig)
+    optim: OptimConfig = field(default_factory=OptimConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+    meta: MetaConfig = field(default_factory=MetaConfig)
+
+
+def load_config(path: str | Path) -> RunConfig:
+    """Load a YAML config, validating every key against :class:`RunConfig`."""
+    merged = OmegaConf.merge(OmegaConf.structured(RunConfig), OmegaConf.load(path))
+    return OmegaConf.to_object(merged)  # type: ignore[return-value]
+
+
+def load_preset(name: str = DEFAULT_PRESET) -> RunConfig:
+    """Load a packaged preset by name (see :data:`PRESETS`)."""
+    if name not in PRESETS:
+        raise KeyError(f"unknown preset {name!r}; known: {', '.join(PRESETS)}")
+    with resources.as_file(resources.files("rsdiff.configs") / f"{name}.yaml") as path:
+        return load_config(path)
+
+
+def resolve_config(spec: str | Path) -> RunConfig:
+    """Return the preset named ``spec``, or load ``spec`` as a YAML path."""
+    return load_preset(str(spec)) if str(spec) in PRESETS else load_config(spec)

rsdiff/configs/rsdiff1.5.yaml

@@ -0,0 +1,54 @@
+# rsdiff1.5: the released cascade. Lightweight base (the net that produced the thesis
+# FID 66.49) + a shrunk SR unet.
+#
+# base 27.2M + SR 92.7M = 119.9M pipeline (~6x lighter than rsdiff1's 723M).
+# SR shrink: dim_mults deepest 8x->4x (1,2,3,4) + num_resnet_blocks->(2,2,2,2).
+
+meta:
+  version: rsdiff1.5
+  reflects: "thesis actual base (27.2M, FID 66.49 lineage) + optimized SR"
+  params_m: 119.9
+  fid_expected: 65.70
+  fid_measured: 65.70
+  notes: "released merged SR ep650 checkpoint; FID-2048 on 1093 RSICD test images, cond_scale=5."
+
+data:
+  name: rsicd
+  root: data/RSICD_optimal
+  image_size: 128
+  num_workers: 4
+
+cascade:
+  image_sizes: [128, 256]
+  timesteps: 1000
+  cond_drop_prob: 0.1
+  text_encoder: t5-base
+  base:                       # unet 1 — 27.2M (frozen, seeded from LR run)
+    dim: 128
+    cond_dim: 256
+    dim_mults: [1, 2, 2, 2]
+    num_resnet_blocks: 0
+    layer_attns: [false, true, true, true]
+    layer_cross_attns: [false, true, true, true]
+    optimizer: adam
+    params_m: 27.2
+  sr:                         # unet 2 — 92.7M (shrunk Efficient-U-Net)
+    dim: 128
+    cond_dim: 512
+    dim_mults: [1, 2, 3, 4]
+    num_resnet_blocks: [2, 2, 2, 2]
+    layer_attns: [false, false, false, true]
+    layer_cross_attns: [false, false, false, true]
+    optimizer: adam
+    params_m: 92.7
+
+optim:
+  lr: 1.0e-4
+  warmup_steps: 10000
+  weight_decay: 0.0
+
+train:
+  seed: 42
+  output_dir: outputs/rsdiff1.5
+  batch_size: 64
+  mixed_precision: bf16

rsdiff/configs/rsdiff1.yaml

@@ -0,0 +1,61 @@
+# rsdiff1 — paper-faithful cascade (RSDiff, NCAA 2024, as described in §3.3).
+#
+# Two-stage Imagen-style cascade: 260.8M base (128²) + 462.4M SR (128→256²).
+# Pipeline = 723.2M ≈ the abstract's stated "0.75 billion" parameters.
+#
+# NOTE on the paper's internal inconsistency: §3.3 states 260M base + 260M SR
+# (= 520M), but the abstract says 0.75B. The committed code's SR unet is 462.4M,
+# and 260.8M + 462.4M = 723M ≈ 0.75B — so the abstract figure is the consistent
+# one and §3.3's "260M SR" is a misstatement. rsdiff1 encodes the runnable code
+# artifacts (big base + original SR), not the §3.3 prose figure.
+#
+# This is NOT the config that produced the reported FID 66.49 — see rsdiff1.5*.
+
+meta:
+  version: rsdiff1
+  reflects: "paper NCAA 2024 §3.3 / abstract 0.75B"
+  params_m: 723.2
+  fid_expected: null          # paper headline; provenance (stage/feature) unrecoverable
+  fid_measured: null
+  notes: "big base (nrb=3) + original Efficient-U-Net SR. Adafactor base / Adam SR per paper."
+
+data:
+  name: rsicd
+  root: data/RSICD_optimal
+  image_size: 128
+  num_workers: 4
+
+cascade:
+  image_sizes: [128, 256]
+  timesteps: 1000
+  cond_drop_prob: 0.1
+  text_encoder: t5-base
+  base:                       # unet 1 — 260.8M (Imagen_text_big_pytorch.py)
+    dim: 128
+    cond_dim: 512
+    dim_mults: [1, 2, 4, 8]
+    num_resnet_blocks: 3
+    layer_attns: [false, true, true, true]
+    layer_cross_attns: [false, true, true, true]
+    optimizer: adafactor
+    params_m: 260.8
+  sr:                         # unet 2 — 462.4M (Efficient-U-Net)
+    dim: 128
+    cond_dim: 512
+    dim_mults: [1, 2, 4, 8]
+    num_resnet_blocks: [2, 4, 8, 8]
+    layer_attns: [false, false, false, true]
+    layer_cross_attns: [false, false, false, true]
+    optimizer: adam
+    params_m: 462.4
+
+optim:
+  lr: 1.0e-4
+  warmup_steps: 10000         # paper: 10k linear warmup
+  weight_decay: 0.0
+
+train:
+  seed: 42
+  output_dir: outputs/rsdiff1
+  batch_size: 64
+  mixed_precision: bf16

rsdiff/datasets/__init__.py

@@ -0,0 +1,16 @@
+from rsdiff.datasets.base import RSDataset
+from rsdiff.datasets.rsicd import RSICD
+
+__all__ = ["RSDataset", "RSICD", "build_dataset"]
+
+
+_REGISTRY = {
+    "rsicd": RSICD,
+}
+
+
+def build_dataset(name: str, **kwargs) -> RSDataset:
+    key = name.lower()
+    if key not in _REGISTRY:
+        raise KeyError(f"Unknown dataset {name!r}. Known: {sorted(_REGISTRY)}")
+    return _REGISTRY[key](**kwargs)

rsdiff/datasets/base.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from torch.utils.data import Dataset
+
+
+@dataclass
+class Sample:
+    image: Any
+    caption: str | None = None
+    label: str | None = None
+    meta: dict[str, Any] | None = None
+
+
+class RSDataset(Dataset, ABC):
+    """Common interface for all remote-sensing datasets in rsdiff.
+
+    Implementations should return a :class:`Sample` from ``__getitem__``.
+    A torchvision-style ``transform`` may be applied to ``image`` before return.
+    """
+
+    name: str = "base"
+    splits: tuple[str, ...] = ("train", "val", "test")
+
+    def __init__(self, root: str | Path, split: str = "train", transform=None) -> None:
+        if split not in self.splits:
+            raise ValueError(f"split={split!r} not in {self.splits}")
+        self.root = Path(root)
+        self.split = split
+        self.transform = transform
+
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    @abstractmethod
+    def __getitem__(self, idx: int) -> Sample: ...
+
+    @property
+    def classes(self) -> list[str]:
+        return []