runnable quick-start example

example_sr_op.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""
+Genesis quick start — super-resolution + outpainting on a live Esri tile.
+
+Downloads the main JiT-B/16 checkpoints from https://huggingface.co/MVRL/genesis,
+fetches one real Esri World Imagery tile (Sydney Opera House / Circular Quay,
+zoom 16), then:
+
+  1. SR: the 256x256 parent tile (z16) -> 512x512 mosaic of its four z17
+     children in a single forward pass            -> genesis_sr_out.png
+  2. OP: the tile viewed as a 2x2 quadrant mosaic (each quadrant is the
+     footprint of one real z17 child); keep the upper-left quadrant real,
+     outpaint the other three                     -> genesis_op_out.png
+
+Run from the root of a https://github.com/subash-khanal/genesis checkout,
+with its environment active (`uv sync && source .venv/bin/activate`):
+
+    python example_sr_op.py
+
+Sampling uses the paper eval setting: 50 steps, cfg = 1.0 (the loader
+defaults). EMA weights are loaded automatically.
+
+Imagery: Esri World Imagery — Source: Esri, Maxar, Earthstar Geographics,
+and the GIS User Community.
+"""
+import os
+import sys
+
+# Repo root: this file at the genesis repo root, or run from a checkout.
+REPO = os.path.dirname(os.path.abspath(__file__))
+if not os.path.isdir(os.path.join(REPO, "src")):
+    REPO = os.getcwd()
+    assert os.path.isdir(os.path.join(REPO, "src")), \
+        "Run this script from the root of a genesis checkout (github.com/subash-khanal/genesis)."
+
+# Same sys.path recipe as demos/*.py — demos/ goes LAST (highest priority)
+# so its `utils/` package is not shadowed by src/*/utils.py.
+sys.path.insert(0, os.path.join(REPO, "src"))
+sys.path.insert(0, os.path.join(REPO, "demos"))
+
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from utils.genesis_common import (  # demos/utils/genesis_common.py
+    fetch_tile_at_zoom,
+    inference_device,
+    load_op_denoiser,
+    load_sr_denoiser,
+    outpaint_context_white_holes_preview,
+    run_outpainting,
+    run_superresolution,
+)
+
+HF_REPO = "MVRL/genesis"
+ARCH = "JiT-B/16"  # main B/16 checkpoints (H/16 SR alone is 23 GB)
+
+# Fixed scenic spot: Sydney Opera House / Circular Quay waterfront.
+LAT, LON, ZOOM = -33.8568, 151.2153, 16  # training zoom levels are 10..19
+
+
+def main() -> None:
+    device = inference_device()  # cuda if available, else cpu
+    print(f"Device: {device}")
+
+    # 1. Checkpoints from the HF Hub (cached under ~/.cache/huggingface).
+    #    The main SR model is DINOv3-conditioned; the main OP model is no-DINO.
+    sr_ckpt = hf_hub_download(HF_REPO, "main/superresolution/main_B16/sr-full-tile-stage3-step0800000.ckpt")
+    op_ckpt = hf_hub_download(HF_REPO, "main/outpainting/main_B16/op-new-stage3-step0800000.ckpt")
+    dino = hf_hub_download(HF_REPO, "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth")
+
+    # 2. One real 256x256 Esri World Imagery tile as the input.
+    tile = fetch_tile_at_zoom(LON, LAT, ZOOM)  # PIL RGB, 256x256
+    tile.save("genesis_sr_in.png")
+
+    # 3. Super-resolution: z16 parent -> 512x512 mosaic of its 4 z17 children.
+    #    target_zoom is the CHILD zoom level (feeds the model's GSD embedder).
+    sr_model, _ = load_sr_denoiser(sr_ckpt, ARCH, dino_weights=dino, device=device)
+    sr_out = run_superresolution(sr_model, tile, device, target_zoom=ZOOM + 1)
+    sr_out.save("genesis_sr_out.png")
+    print("Saved genesis_sr_in.png (256 parent) and genesis_sr_out.png (512 SR mosaic)")
+
+    del sr_model  # free memory before loading the OP model
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+
+    # 4. Outpainting: 2x2 quadrant mosaic — 1 real quadrant, 3 masked holes.
+    #    Mask convention: white (255) = hole to generate, black (0) = known.
+    mask = Image.new("L", (256, 256), 255)  # everything is a hole ...
+    mask.paste(0, (0, 0, 128, 128))         # ... except the upper-left quadrant
+
+    op_model, _ = load_op_denoiser(op_ckpt, ARCH, dino_weights="", device=device)
+    outpaint_context_white_holes_preview(tile, mask).save("genesis_op_in.png")
+    op_out = run_outpainting(op_model, tile, mask, device, tile_zoom=ZOOM)
+    op_out.save("genesis_op_out.png")
+    print("Saved genesis_op_in.png (masked input) and genesis_op_out.png (outpainted)")
+
+
+if __name__ == "__main__":
+    main()