Deploy ZeroGPU codebase

app.py

@@ -0,0 +1,461 @@
+import io
+import time
+import sys
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+from PIL import Image
+import spaces
+
+# ============================================================
+# PATHS
+# ============================================================
+
+ROOT = Path(__file__).resolve().parent
+
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+# ============================================================
+# ZERO-GPU MODEL LOADING (CPU initially)
+# ============================================================
+
+from models.hatsat.hatsat_inference import HATSATInference
+from models.esrgan.esrgan_inference import ESRGANInference
+
+print("Loading models to CPU for ZeroGPU deployment...")
+try:
+    hatsat_model = HATSATInference(device="cpu")
+    esrgan_model = ESRGANInference(
+        checkpoint_path=str(ROOT / "weights" / "esrgan" / "RRDB_ESRGAN_x4.pth"),
+        device="cpu"
+    )
+    print("Models loaded successfully on CPU.")
+except Exception as e:
+    print(f"Error loading models: {e}")
+    hatsat_model = None
+    esrgan_model = None
+
+# ============================================================
+# ZERO-GPU INFERENCE WRAPPERS
+# ============================================================
+
+@spaces.GPU
+def infer_hatsat(image):
+    if hatsat_model is None:
+        raise RuntimeError("HATSAT model is not loaded.")
+    return hatsat_model.predict(image)
+
+@spaces.GPU
+def infer_esrgan(image):
+    if esrgan_model is None:
+        raise RuntimeError("ESRGAN model is not loaded.")
+    return esrgan_model.predict(image)
+
+
+# ============================================================
+# CUSTOM CSS
+# ============================================================
+
+CSS = """
+/* ── Google Font ── */
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
+
+body, .gradio-container {
+    font-family: 'Inter', sans-serif !important;
+}
+
+/* ── Hero header ── */
+.hero-header {
+    background: linear-gradient(135deg, #0f172a 0%, #1e3a5f 50%, #0f172a 100%);
+    border-radius: 16px;
+    padding: 36px 40px 28px;
+    margin-bottom: 8px;
+    border: 1px solid rgba(56, 189, 248, 0.2);
+    box-shadow: 0 0 40px rgba(56, 189, 248, 0.08);
+    position: relative;
+    overflow: hidden;
+}
+
+.hero-header::before {
+    content: '';
+    position: absolute;
+    top: -60px; right: -60px;
+    width: 200px; height: 200px;
+    background: radial-gradient(circle, rgba(56,189,248,0.12) 0%, transparent 70%);
+    pointer-events: none;
+}
+
+.hero-title {
+    font-size: 2.1rem;
+    font-weight: 700;
+    color: #f0f9ff;
+    margin: 0 0 8px;
+    letter-spacing: -0.5px;
+}
+
+.hero-subtitle {
+    font-size: 0.95rem;
+    color: #94a3b8;
+    margin: 0;
+    line-height: 1.6;
+}
+
+.hero-badge {
+    display: inline-block;
+    background: rgba(56, 189, 248, 0.15);
+    color: #38bdf8;
+    border: 1px solid rgba(56, 189, 248, 0.3);
+    border-radius: 20px;
+    padding: 3px 12px;
+    font-size: 0.75rem;
+    font-weight: 600;
+    letter-spacing: 0.5px;
+    margin-bottom: 14px;
+}
+
+/* ── Section cards ── */
+.card {
+    background: #1e293b;
+    border: 1px solid #334155;
+    border-radius: 12px;
+    padding: 20px;
+    margin-bottom: 12px;
+}
+
+.section-label {
+    font-size: 0.7rem;
+    font-weight: 600;
+    letter-spacing: 1.2px;
+    text-transform: uppercase;
+    color: #64748b;
+    margin-bottom: 10px;
+}
+
+/* ── Generate button ── */
+#generate-btn {
+    background: linear-gradient(135deg, #0ea5e9 0%, #2563eb 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-weight: 600 !important;
+    font-size: 1rem !important;
+    padding: 14px !important;
+    border-radius: 10px !important;
+    letter-spacing: 0.3px;
+    transition: opacity 0.2s ease, transform 0.1s ease !important;
+    box-shadow: 0 4px 14px rgba(14, 165, 233, 0.35) !important;
+    width: 100%;
+}
+#generate-btn:hover {
+    opacity: 0.92 !important;
+    transform: translateY(-1px) !important;
+}
+#generate-btn:active {
+    transform: translateY(0px) !important;
+}
+
+/* ── Download button ── */
+#download-btn {
+    background: #0f172a !important;
+    border: 1px solid #38bdf8 !important;
+    color: #38bdf8 !important;
+    font-weight: 600 !important;
+    border-radius: 8px !important;
+    padding: 10px !important;
+    width: 100%;
+    transition: background 0.2s ease !important;
+}
+#download-btn:hover {
+    background: rgba(56,189,248,0.1) !important;
+}
+
+/* ── Status box ── */
+.status-box {
+    background: #0f172a;
+    border: 1px solid #1e3a5f;
+    border-radius: 10px;
+    padding: 16px 20px;
+    font-size: 0.88rem;
+    color: #94a3b8;
+    min-height: 120px;
+}
+
+/* ── Model info chips ── */
+.model-chip-hat {
+    display: inline-block;
+    background: rgba(16, 185, 129, 0.12);
+    color: #34d399;
+    border: 1px solid rgba(16,185,129,0.3);
+    border-radius: 6px;
+    padding: 2px 10px;
+    font-size: 0.72rem;
+    font-weight: 600;
+    margin-right: 6px;
+}
+
+.model-chip-esr {
+    display: inline-block;
+    background: rgba(168, 85, 247, 0.12);
+    color: #c084fc;
+    border: 1px solid rgba(168,85,247,0.3);
+    border-radius: 6px;
+    padding: 2px 10px;
+    font-size: 0.72rem;
+    font-weight: 600;
+}
+
+/* ── Footer ── */
+.footer-note {
+    background: rgba(249, 115, 22, 0.07);
+    border: 1px solid rgba(249, 115, 22, 0.25);
+    border-radius: 8px;
+    padding: 12px 16px;
+    font-size: 0.82rem;
+    color: #fb923c;
+    margin-top: 4px;
+}
+"""
+
+
+# ============================================================
+# INFERENCE FUNCTION
+# ============================================================
+
+def run_super_resolution(image, model_choice, progress=gr.Progress()):
+
+    # --------------------------------------------------------
+    # Validate input
+    # --------------------------------------------------------
+    if image is None:
+        raise gr.Error("⚠️ Please upload a satellite image before running.")
+
+    # --------------------------------------------------------
+    # Normalize model name
+    # --------------------------------------------------------
+    model_map = {"HATSAT — Recommended": "hatsat", "ESRGAN — Baseline": "esrgan"}
+    selected_model = model_map.get(model_choice)
+
+    if selected_model is None:
+        raise gr.Error("⚠️ Please select a valid model.")
+
+    # --------------------------------------------------------
+    # Ensure PIL image + RGB
+    # --------------------------------------------------------
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    image = image.convert("RGB")
+    input_w, input_h = image.size
+
+    # --------------------------------------------------------
+    # Basic sanity check
+    # --------------------------------------------------------
+    if input_w < 16 or input_h < 16:
+        raise gr.Error(
+            f"⚠️ Image too small ({input_w}×{input_h}). "
+            "Please upload an image at least 16×16 pixels."
+        )
+
+    if input_w > 2000 or input_h > 2000:
+        raise gr.Error(
+            f"⚠️ Image too large ({input_w}×{input_h}). "
+            "For local testing please use images up to 2000×2000 pixels."
+        )
+
+    # --------------------------------------------------------
+    # Run inference
+    # --------------------------------------------------------
+    model_label = "HATSAT" if selected_model == "hatsat" else "ESRGAN"
+    
+    start_time = time.time()
+
+    try:
+        progress(0.5, desc=f"Running {model_label} 4× super-resolution (ZeroGPU) …")
+        if selected_model == "hatsat":
+            result = infer_hatsat(image)
+        else:
+            result = infer_esrgan(image)
+    except Exception as e:
+        raise gr.Error(f"{model_label} inference failed: {e}")
+
+    elapsed = time.time() - start_time
+    progress(0.9, desc="Finalising output …")
+
+    output_w, output_h = result.size
+
+    # --------------------------------------------------------
+    # Save to a temp PNG for download
+    # --------------------------------------------------------
+    tmp = tempfile.NamedTemporaryFile(
+        suffix=f"_sr_{selected_model}_4x.png",
+        delete=False
+    )
+    result.save(tmp.name, format="PNG")
+    tmp.close()
+
+    # --------------------------------------------------------
+    # Info markdown
+    # --------------------------------------------------------
+    info = f"""
+<div class="status-box">
+
+**✅ Super-Resolution Complete**
+
+| Field | Value |
+|---|---|
+| **Model** | {model_label} |
+| **Scale** | 4× |
+| **Input size** | {input_w} × {input_h} px |
+| **Output size** | {output_w} × {output_h} px |
+| **Processing time** | {elapsed:.2f} s |
+| **Nominal output sampling** | ~2.5 m (from 10 m input) |
+
+</div>
+"""
+
+    progress(1.0, desc="Done.")
+    return result, info, tmp.name
+
+
+# ============================================================
+# CLEAR FUNCTION
+# ============================================================
+
+def clear_all():
+    return None, None, "", None
+
+
+# ============================================================
+# UI
+# ============================================================
+
+THEME = gr.themes.Base(
+    primary_hue="sky",
+    neutral_hue="slate",
+)
+
+with gr.Blocks(
+    title="Satellite Super-Resolution — SIH 26142 (ZeroGPU)",
+) as demo:
+
+    # ── Hero header ─────────────────────────────────────────
+    gr.HTML("""
+    <div class="hero-header">
+        <div class="hero-badge">SIH PROBLEM STATEMENT 26142 · NTRO · SPACE TECHNOLOGY</div>
+        <div class="hero-title">🛰️ Satellite Image Super-Resolution</div>
+        <div class="hero-subtitle">
+            Deep learning–based <strong style="color:#38bdf8">4× super-resolution</strong>
+            for medium-resolution satellite imagery (Sentinel-2 and similar sources).<br>
+            Upload an image, select a model, and generate a sharper super-resolved output.
+        </div>
+    </div>
+    """)
+
+    # ── Main layout ─────────────────────────────────────────
+    with gr.Row(equal_height=False):
+
+        # ── LEFT column: inputs ─────────────────────────────
+        with gr.Column(scale=5):
+
+            gr.HTML('<div class="section-label">📂 Input Image</div>')
+            input_image = gr.Image(
+                label="Upload Satellite Image",
+                type="pil",
+                sources=["upload", "clipboard"],
+                elem_id="input-img",
+            )
+
+            gr.HTML('<div class="section-label" style="margin-top:16px">🤖 Model Selection</div>')
+            model_selector = gr.Radio(
+                choices=["HATSAT — Recommended", "ESRGAN — Baseline"],
+                value="HATSAT — Recommended",
+                label="Super-Resolution Model",
+                info=(
+                    "HATSAT is fine-tuned for satellite imagery and is the primary model. "
+                    "ESRGAN is a general-purpose baseline included for comparison."
+                ),
+            )
+
+            gr.HTML("""
+            <div class="card" style="margin-top:4px; font-size:0.82rem; color:#94a3b8; line-height:1.7;">
+                <span class="model-chip-hat">HATSAT</span>
+                Satellite-oriented · HAT transformer · 4× upscaling · Recommended<br>
+                <span class="model-chip-esr" style="margin-top:6px; display:inline-block;">ESRGAN</span>
+                General-purpose · RRDB architecture · 4× upscaling · Baseline comparison
+            </div>
+            """)
+
+            with gr.Row():
+                generate_btn = gr.Button(
+                    "🚀 Run Super-Resolution",
+                    variant="primary",
+                    elem_id="generate-btn",
+                )
+                clear_btn = gr.Button(
+                    "✕ Clear",
+                    variant="secondary",
+                )
+
+        # ── RIGHT column: outputs ────────────────────────────
+        with gr.Column(scale=5):
+
+            gr.HTML('<div class="section-label">🖼️ Super-Resolved Output</div>')
+            output_image = gr.Image(
+                label="4× Super-Resolved Output",
+                type="pil",
+                interactive=False,
+                elem_id="output-img",
+            )
+
+            download_file = gr.File(
+                label="⬇️ Download Result (PNG)",
+                interactive=False,
+                visible=True,
+                elem_id="download-btn",
+            )
+
+            gr.HTML('<div class="section-label" style="margin-top:12px">📊 Processing Info</div>')
+            output_info = gr.Markdown(
+                value='<div class="status-box">Results will appear here after processing.</div>',
+            )
+
+    # ── Disclaimer footer ────────────────────────────────────
+    gr.HTML("""
+    <div class="footer-note">
+        ⚠️ <strong>Important:</strong>
+        4× super-resolution produces approximately <strong>2.5 m nominal output sampling</strong>
+        from a 10 m Sentinel-2 input. This does <em>not</em> mean the system physically recovers
+        true 2.5 m ground resolution. Results are generated outputs, not real sensor measurements.
+    </div>
+    """)
+
+    # ── Events ──────────────────────────────────────────────
+
+    generate_btn.click(
+        fn=run_super_resolution,
+        inputs=[input_image, model_selector],
+        outputs=[output_image, output_info, download_file],
+        show_progress="full",
+    )
+
+    clear_btn.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[input_image, output_image, output_info, download_file],
+    )
+
+
+# ============================================================
+# LAUNCH
+# ============================================================
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        theme=THEME,
+        css=CSS,
+    )

inference/geotiff.py

@@ -0,0 +1,318 @@
+
+from pathlib import Path
+
+import numpy as np
+import rasterio
+from rasterio.transform import Affine
+from PIL import Image
+
+
+# ==========================================================
+# Read Sentinel-2 RGB bands
+# ==========================================================
+
+def read_sentinel_rgb(
+    b02_path,
+    b03_path,
+    b04_path
+):
+    """
+    Read Sentinel-2 B02, B03 and B04 GeoTIFF bands.
+
+    Band mapping:
+
+        B04 -> Red
+        B03 -> Green
+        B02 -> Blue
+
+    Returns
+    -------
+    rgb : numpy.ndarray
+        RGB array with shape (H, W, 3)
+
+    profile : rasterio profile
+        Original geospatial metadata.
+    """
+
+    b02_path = Path(b02_path)
+    b03_path = Path(b03_path)
+    b04_path = Path(b04_path)
+
+    with rasterio.open(b02_path) as src_b02:
+        b02 = src_b02.read(1)
+        profile = src_b02.profile.copy()
+        transform = src_b02.transform
+        crs = src_b02.crs
+        width = src_b02.width
+        height = src_b02.height
+
+    with rasterio.open(b03_path) as src_b03:
+        b03 = src_b03.read(1)
+
+    with rasterio.open(b04_path) as src_b04:
+        b04 = src_b04.read(1)
+
+    # Make sure all bands have identical dimensions.
+    if b02.shape != b03.shape or b02.shape != b04.shape:
+        raise ValueError(
+            "B02, B03 and B04 dimensions do not match."
+        )
+
+    # ------------------------------------------------------
+    # Sentinel-2 RGB order:
+    # Red   = B04
+    # Green = B03
+    # Blue  = B02
+    # ------------------------------------------------------
+
+    rgb = np.stack(
+        [
+            b04,
+            b03,
+            b02
+        ],
+        axis=-1
+    )
+
+    return rgb, profile
+
+
+# ==========================================================
+# Convert raw Sentinel data to visualization RGB
+# ==========================================================
+
+def normalize_rgb(
+    rgb,
+    lower_percentile=2,
+    upper_percentile=98
+):
+    """
+    Convert Sentinel-2 reflectance/DN values to uint8 RGB.
+
+    Percentile normalization is performed independently
+    for each channel.
+    """
+
+    rgb = rgb.astype(np.float32)
+
+    result = np.zeros_like(
+        rgb,
+        dtype=np.uint8
+    )
+
+    for channel in range(3):
+
+        band = rgb[:, :, channel]
+
+        low = np.percentile(
+            band,
+            lower_percentile
+        )
+
+        high = np.percentile(
+            band,
+            upper_percentile
+        )
+
+        if high <= low:
+            result[:, :, channel] = np.clip(
+                band,
+                0,
+                255
+            ).astype(np.uint8)
+
+        else:
+
+            normalized = (
+                (band - low) /
+                (high - low)
+            )
+
+            normalized = np.clip(
+                normalized,
+                0,
+                1
+            )
+
+            result[:, :, channel] = (
+                normalized * 255
+            ).astype(np.uint8)
+
+    return result
+
+
+# ==========================================================
+# GeoTIFF writer
+# ==========================================================
+
+def save_sr_geotiff(
+    sr_image,
+    source_profile,
+    output_path,
+    scale=4
+):
+    """
+    Save a super-resolved RGB image as a GeoTIFF.
+
+    The geographic extent remains unchanged while
+    the pixel dimensions increase by `scale`.
+
+    Parameters
+    ----------
+    sr_image : PIL.Image.Image or numpy.ndarray
+
+    source_profile : dict
+        Rasterio profile from the original Sentinel image.
+
+    output_path : str or Path
+
+    scale : int
+        Super-resolution scale factor.
+    """
+
+    output_path = Path(output_path)
+    output_path.parent.mkdir(
+        parents=True,
+        exist_ok=True
+    )
+
+    # ------------------------------------------------------
+    # Convert image to numpy
+    # ------------------------------------------------------
+
+    if isinstance(
+        sr_image,
+        Image.Image
+    ):
+        rgb = np.asarray(
+            sr_image.convert("RGB")
+        )
+
+    else:
+        rgb = np.asarray(
+            sr_image
+        )
+
+    if rgb.ndim != 3 or rgb.shape[2] != 3:
+        raise ValueError(
+            "sr_image must have shape (H, W, 3)."
+        )
+
+    rgb = np.clip(
+        rgb,
+        0,
+        255
+    ).astype(np.uint8)
+
+    height, width, _ = rgb.shape
+
+    # ------------------------------------------------------
+    # Original transform
+    # ------------------------------------------------------
+
+    original_transform = source_profile[
+        "transform"
+    ]
+
+    # ------------------------------------------------------
+    # Scale pixel size by 1/scale
+    #
+    # Geographic bounds stay unchanged.
+    # ------------------------------------------------------
+
+    new_transform = (
+        original_transform *
+        Affine.scale(
+            1.0 / scale,
+            1.0 / scale
+        )
+    )
+
+    # ------------------------------------------------------
+    # Build output profile
+    # ------------------------------------------------------
+
+    profile = source_profile.copy()
+
+    profile.update(
+        {
+            "driver": "GTiff",
+            "height": height,
+            "width": width,
+            "count": 3,
+            "dtype": "uint8",
+            "transform": new_transform,
+            "compress": "deflate",
+            "photometric": "RGB",
+            "tiled": True,
+            "blockxsize": 256,
+            "blockysize": 256,
+        }
+    )
+
+    # Remove incompatible fields where necessary.
+    profile.pop(
+        "nodata",
+        None
+    )
+
+    # ------------------------------------------------------
+    # Write GeoTIFF
+    # ------------------------------------------------------
+    # Remove block settings when not using tiled GeoTIFF
+    profile.pop("blockxsize", None)
+    profile.pop("blockysize", None)
+    profile.pop("tiled", None)
+
+    with rasterio.open(
+        output_path,
+        "w",
+        **profile
+    ) as dst:
+
+        dst.write(
+            rgb[:, :, 0],
+            1
+        )
+
+        dst.write(
+            rgb[:, :, 1],
+            2
+        )
+
+        dst.write(
+            rgb[:, :, 2],
+            3
+        )
+
+    return output_path
+
+
+# ==========================================================
+# Convenience function
+# ==========================================================
+
+def sentinel_rgb_from_bands(
+    b02_path,
+    b03_path,
+    b04_path
+):
+    """
+    Read B02/B03/B04 and return a normalized PIL RGB image
+    together with the source GeoTIFF profile.
+    """
+
+    rgb, profile = read_sentinel_rgb(
+        b02_path,
+        b03_path,
+        b04_path
+    )
+
+    rgb_uint8 = normalize_rgb(rgb)
+
+    image = Image.fromarray(
+        rgb_uint8,
+        mode="RGB"
+    )
+
+    return image, profile

inference/pipeline.py

@@ -0,0 +1,504 @@
+
+from pathlib import Path
+from PIL import Image
+
+from .tiling import create_tiles
+from .stitching import stitch_tiles
+
+from models.hatsat.hatsat_inference import HATSATInference
+from models.esrgan.esrgan_inference import ESRGANInference
+
+
+# ==========================================================
+# Project paths
+# ==========================================================
+
+# satellite_sr_deploy/ is the parent of inference/
+ROOT = Path(__file__).resolve().parent.parent
+
+HATSAT_CHECKPOINT = (
+    ROOT /
+    "weights/hatsat/net_g_150000.pth"
+)
+
+ESRGAN_CHECKPOINT = (
+    ROOT /
+    "weights/esrgan/RRDB_ESRGAN_x4.pth"
+)
+
+
+# ==========================================================
+# Model cache
+# ==========================================================
+
+_MODEL_CACHE = {}
+
+
+def get_model(model_name, device=None):
+    """
+    Load a model once and reuse it.
+
+    Parameters
+    ----------
+    model_name : str
+        'hatsat' or 'esrgan'
+
+    device : optional
+        'cuda' or 'cpu'
+
+    Returns
+    -------
+    Loaded inference model.
+    """
+
+    model_name = model_name.lower().strip()
+
+    if model_name not in ("hatsat", "esrgan"):
+        raise ValueError(
+            "model_name must be 'hatsat' or 'esrgan'."
+        )
+
+    cache_key = (
+        model_name,
+        str(device) if device is not None else "auto"
+    )
+
+    if cache_key in _MODEL_CACHE:
+        return _MODEL_CACHE[cache_key]
+
+    # ------------------------------------------------------
+    # HATSAT
+    # ------------------------------------------------------
+
+    if model_name == "hatsat":
+
+        model = HATSATInference(
+            checkpoint_path=HATSAT_CHECKPOINT,
+            device=device
+        )
+
+    # ------------------------------------------------------
+    # ESRGAN
+    # ------------------------------------------------------
+
+    else:
+
+        model = ESRGANInference(
+            checkpoint_path=ESRGAN_CHECKPOINT,
+            device=device
+        )
+
+    _MODEL_CACHE[cache_key] = model
+
+    return model
+
+
+# ==========================================================
+# Main pipeline
+# ==========================================================
+
+def super_resolve(
+    image,
+    model_name="hatsat",
+    tile_size=128,
+    overlap=32,
+    device=None
+):
+    """
+    Complete satellite super-resolution pipeline.
+
+    Pipeline:
+
+        Input image
+            ↓
+        Tiling
+            ↓
+        Selected model
+            ↓
+        4x SR
+            ↓
+        Overlap stitching
+            ↓
+        Final SR image
+
+    Parameters
+    ----------
+    image : PIL.Image.Image
+
+    model_name : str
+        'hatsat' or 'esrgan'
+
+    tile_size : int
+        Input tile size.
+
+    overlap : int
+        Input overlap.
+
+    device : optional
+        'cuda' or 'cpu'
+
+    Returns
+    -------
+    PIL.Image.Image
+        Final 4x super-resolved image.
+    """
+
+    # ------------------------------------------------------
+    # Validate image
+    # ------------------------------------------------------
+
+    if not isinstance(image, Image.Image):
+        raise TypeError(
+            "image must be a PIL.Image.Image"
+        )
+
+    image = image.convert("RGB")
+
+    original_size = image.size
+
+    print(
+        f"Pipeline input: "
+        f"{original_size[0]} × {original_size[1]}"
+    )
+
+    # ------------------------------------------------------
+    # Get selected model
+    # ------------------------------------------------------
+
+    model = get_model(
+        model_name,
+        device=device
+    )
+
+    print(
+        f"Selected model: "
+        f"{model_name.upper()}"
+    )
+
+    # ------------------------------------------------------
+    # Create tiles
+    # ------------------------------------------------------
+
+    tiles = create_tiles(
+        image,
+        tile_size=tile_size,
+        overlap=overlap
+    )
+
+    print(
+        f"Number of tiles: {len(tiles)}"
+    )
+
+    # ------------------------------------------------------
+    # Run SR on each tile
+    # ------------------------------------------------------
+
+    sr_tiles = []
+
+    for tile_info in tiles:
+
+        print(
+            f"Processing tile "
+            f"{tile_info['index'] + 1}/"
+            f"{len(tiles)} "
+            f"(x={tile_info['x']}, "
+            f"y={tile_info['y']})"
+        )
+
+        sr_image = model.predict(
+            tile_info["image"]
+        )
+
+        sr_tiles.append(
+            {
+                "image": sr_image,
+                "x": tile_info["x"],
+                "y": tile_info["y"]
+            }
+        )
+
+    # ------------------------------------------------------
+    # Stitch SR tiles
+    # ------------------------------------------------------
+
+    result = stitch_tiles(
+        sr_tiles,
+        scale=4,
+        original_size=original_size,
+        overlap=overlap
+    )
+
+    print(
+        f"Pipeline output: "
+        f"{result.size[0]} × "
+        f"{result.size[1]}"
+    )
+
+    return result
+
+
+# ==========================================================
+# Convenience functions
+# ==========================================================
+
+def run_hatsat_pipeline(
+    image,
+    tile_size=128,
+    overlap=32,
+    device=None
+):
+    return super_resolve(
+        image=image,
+        model_name="hatsat",
+        tile_size=tile_size,
+        overlap=overlap,
+        device=device
+    )
+
+
+def run_esrgan_pipeline(
+    image,
+    tile_size=128,
+    overlap=32,
+    device=None
+):
+    return super_resolve(
+        image=image,
+        model_name="esrgan",
+        tile_size=tile_size,
+        overlap=overlap,
+        device=device
+    )
+
+
+# ==========================================================
+# COMPLETE SATELLITE PROCESSING
+# ==========================================================
+
+from .geotiff import (
+    sentinel_rgb_from_bands,
+    save_sr_geotiff
+)
+
+
+def process_satellite(
+    b02_path,
+    b03_path,
+    b04_path,
+    model_name="hatsat",
+    output_dir=None,
+    tile_size=128,
+    overlap=32,
+    device=None
+):
+    """
+    Complete Sentinel-2 satellite super-resolution workflow.
+
+    Input:
+        B02, B03 and B04 Sentinel-2 GeoTIFF files.
+
+    Processing:
+        B02/B03/B04
+            ↓
+        RGB conversion
+            ↓
+        Tiling
+            ↓
+        HATSAT / ESRGAN
+            ↓
+        Stitching
+            ↓
+        4x SR
+
+    Output:
+        PNG preview
+        Georeferenced RGB GeoTIFF
+
+    Returns
+    -------
+    dict
+        {
+            "model": model_name,
+            "input_size": (...),
+            "output_size": (...),
+            "png": Path,
+            "geotiff": Path
+        }
+    """
+
+    # ------------------------------------------------------
+    # Paths
+    # ------------------------------------------------------
+
+    b02_path = Path(b02_path)
+    b03_path = Path(b03_path)
+    b04_path = Path(b04_path)
+
+    if output_dir is None:
+        output_dir = (
+            ROOT /
+            "outputs" /
+            model_name.lower() /
+            "results"
+        )
+
+    output_dir = Path(output_dir)
+    output_dir.mkdir(
+        parents=True,
+        exist_ok=True
+    )
+
+    # ------------------------------------------------------
+    # Validate inputs
+    # ------------------------------------------------------
+
+    for path in [
+        b02_path,
+        b03_path,
+        b04_path
+    ]:
+
+        if not path.exists():
+            raise FileNotFoundError(
+                f"Input file not found: {path}"
+            )
+
+    # ------------------------------------------------------
+    # Read Sentinel-2 bands
+    # ------------------------------------------------------
+
+    print("\n========================================")
+    print("SATELLITE SUPER-RESOLUTION PIPELINE")
+    print("========================================")
+
+    print("\nReading Sentinel-2 bands...")
+
+    image, source_profile = sentinel_rgb_from_bands(
+        b02_path,
+        b03_path,
+        b04_path
+    )
+
+    print(
+        "Input RGB size:",
+        image.size
+    )
+
+    print(
+        "Source CRS:",
+        source_profile["crs"]
+    )
+
+    # ------------------------------------------------------
+    # Run existing SR pipeline
+    # ------------------------------------------------------
+
+    print(
+        "\nRunning model:",
+        model_name.upper()
+    )
+
+    sr_image = super_resolve(
+        image=image,
+        model_name=model_name,
+        tile_size=tile_size,
+        overlap=overlap,
+        device=device
+    )
+
+    print(
+        "\nSR image size:",
+        sr_image.size
+    )
+
+    # ------------------------------------------------------
+    # Output filenames
+    # ------------------------------------------------------
+
+    model_tag = model_name.lower()
+
+    png_path = (
+        output_dir /
+        f"sentinel2_{model_tag}_4x.png"
+    )
+
+    geotiff_path = (
+        output_dir /
+        f"sentinel2_{model_tag}_4x_georef.tif"
+    )
+
+    # ------------------------------------------------------
+    # Save PNG preview
+    # ------------------------------------------------------
+
+    sr_image.save(
+        png_path,
+        format="PNG"
+    )
+
+    print(
+        "\nPNG saved:",
+        png_path
+    )
+
+    # ------------------------------------------------------
+    # Save georeferenced GeoTIFF
+    # ------------------------------------------------------
+
+    save_sr_geotiff(
+        sr_image=sr_image,
+        source_profile=source_profile,
+        output_path=geotiff_path,
+        scale=4
+    )
+
+    print(
+        "GeoTIFF saved:",
+        geotiff_path
+    )
+
+    # ------------------------------------------------------
+    # Final summary
+    # ------------------------------------------------------
+
+    print("\n========================================")
+    print("PROCESSING COMPLETE")
+    print("========================================")
+
+    print(
+        "Model:",
+        model_name.upper()
+    )
+
+    print(
+        "Input:",
+        image.size
+    )
+
+    print(
+        "Output:",
+        sr_image.size
+    )
+
+    print(
+        "Scale: 4x"
+    )
+
+    print(
+        "PNG:",
+        png_path
+    )
+
+    print(
+        "GeoTIFF:",
+        geotiff_path
+    )
+
+    return {
+        "model": model_name.lower(),
+        "input_size": image.size,
+        "output_size": sr_image.size,
+        "png": png_path,
+        "geotiff": geotiff_path
+    }
+

inference/stitching.py

@@ -0,0 +1,216 @@
+
+import numpy as np
+from PIL import Image
+
+
+def create_blend_weight(
+    height,
+    width,
+    overlap
+):
+    """
+    Create a smooth 2D blending weight.
+
+    Pixels near the center receive higher weight.
+    Pixels near overlapping boundaries receive lower weight.
+    """
+
+    if overlap <= 0:
+        return np.ones(
+            (height, width),
+            dtype=np.float32
+        )
+
+    # Horizontal weights
+    wx = np.ones(width, dtype=np.float32)
+
+    transition = min(overlap, width // 2)
+
+    if transition > 0:
+
+        ramp = np.linspace(
+            0.01,
+            1.0,
+            transition,
+            dtype=np.float32
+        )
+
+        wx[:transition] = ramp
+
+        wx[-transition:] = ramp[::-1]
+
+    # Vertical weights
+    wy = np.ones(height, dtype=np.float32)
+
+    transition = min(overlap, height // 2)
+
+    if transition > 0:
+
+        ramp = np.linspace(
+            0.01,
+            1.0,
+            transition,
+            dtype=np.float32
+        )
+
+        wy[:transition] = ramp
+
+        wy[-transition:] = ramp[::-1]
+
+    return wy[:, None] * wx[None, :]
+
+
+def stitch_tiles(
+    sr_tiles,
+    scale=4,
+    original_size=None,
+    overlap=32
+):
+    """
+    Stitch overlapping super-resolution tiles.
+
+    Parameters
+    ----------
+    sr_tiles : list of dictionaries
+
+        Each dictionary must contain:
+
+            image
+            x
+            y
+
+        where x/y are coordinates in the ORIGINAL image.
+
+    scale : int
+        Super-resolution scale factor.
+
+    original_size : tuple
+        (width, height) of original image.
+
+    overlap : int
+        Overlap in ORIGINAL-image pixels.
+
+    Returns
+    -------
+    PIL.Image
+        Final stitched SR image.
+    """
+
+    if not sr_tiles:
+        raise ValueError("No SR tiles supplied.")
+
+    if original_size is None:
+        raise ValueError(
+            "original_size must be provided."
+        )
+
+    original_width, original_height = original_size
+
+    output_width = original_width * scale
+    output_height = original_height * scale
+
+    # Accumulate weighted RGB values
+    canvas = np.zeros(
+        (
+            output_height,
+            output_width,
+            3
+        ),
+        dtype=np.float32
+    )
+
+    # Accumulate weights
+    weights = np.zeros(
+        (
+            output_height,
+            output_width
+        ),
+        dtype=np.float32
+    )
+
+    sr_overlap = overlap * scale
+
+    for tile_info in sr_tiles:
+
+        image = tile_info["image"]
+
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        image = image.convert("RGB")
+
+        tile = np.asarray(
+            image,
+            dtype=np.float32
+        )
+
+        tile_height, tile_width = tile.shape[:2]
+
+        # Original-image coordinates → SR coordinates
+        x = int(tile_info["x"] * scale)
+        y = int(tile_info["y"] * scale)
+
+        # Do not allow the tile to exceed final canvas
+        valid_width = min(
+            tile_width,
+            output_width - x
+        )
+
+        valid_height = min(
+            tile_height,
+            output_height - y
+        )
+
+        if valid_width <= 0 or valid_height <= 0:
+            continue
+
+        tile = tile[
+            :valid_height,
+            :valid_width
+        ]
+
+        # --------------------------------------------------
+        # Build blending weight
+        # --------------------------------------------------
+
+        weight = create_blend_weight(
+            valid_height,
+            valid_width,
+            sr_overlap
+        )
+
+        # --------------------------------------------------
+        # Accumulate
+        # --------------------------------------------------
+
+        canvas[
+            y:y + valid_height,
+            x:x + valid_width
+        ] += tile * weight[..., None]
+
+        weights[
+            y:y + valid_height,
+            x:x + valid_width
+        ] += weight
+
+    # ------------------------------------------------------
+    # Normalize overlapping pixels
+    # ------------------------------------------------------
+
+    weights = np.maximum(
+        weights,
+        1e-8
+    )
+
+    canvas /= weights[..., None]
+
+    canvas = np.clip(
+        canvas,
+        0,
+        255
+    ).round().astype(np.uint8)
+
+    return Image.fromarray(
+        canvas,
+        mode="RGB"
+    )

inference/tiling.py

@@ -0,0 +1,161 @@
+
+from pathlib import Path
+from PIL import Image
+
+
+DEFAULT_TILE_SIZE = 128
+DEFAULT_OVERLAP = 32
+
+
+def calculate_positions(length, tile_size, overlap):
+    """
+    Calculate tile starting positions along one dimension.
+
+    The final tile always reaches the image boundary.
+    """
+
+    if tile_size <= 0:
+        raise ValueError("tile_size must be positive.")
+
+    if overlap < 0 or overlap >= tile_size:
+        raise ValueError(
+            "overlap must satisfy 0 <= overlap < tile_size."
+        )
+
+    stride = tile_size - overlap
+
+    if length <= tile_size:
+        return [0]
+
+    positions = list(range(0, length - tile_size + 1, stride))
+
+    # Ensure the final region is covered.
+    last_position = length - tile_size
+
+    if positions[-1] != last_position:
+        positions.append(last_position)
+
+    return positions
+
+
+def create_tiles(
+    image,
+    tile_size=DEFAULT_TILE_SIZE,
+    overlap=DEFAULT_OVERLAP
+):
+    """
+    Split a PIL image into overlapping tiles.
+
+    Returns:
+        list of dictionaries containing:
+            image
+            x
+            y
+            width
+            height
+            index
+    """
+
+    if not isinstance(image, Image.Image):
+        raise TypeError(
+            "image must be a PIL.Image.Image"
+        )
+
+    image = image.convert("RGB")
+
+    width, height = image.size
+
+    x_positions = calculate_positions(
+        width,
+        tile_size,
+        overlap
+    )
+
+    y_positions = calculate_positions(
+        height,
+        tile_size,
+        overlap
+    )
+
+    tiles = []
+
+    index = 0
+
+    for y in y_positions:
+
+        for x in x_positions:
+
+            right = min(
+                x + tile_size,
+                width
+            )
+
+            bottom = min(
+                y + tile_size,
+                height
+            )
+
+            tile = image.crop(
+                (x, y, right, bottom)
+            )
+
+            tiles.append(
+                {
+                    "image": tile,
+                    "x": x,
+                    "y": y,
+                    "width": right - x,
+                    "height": bottom - y,
+                    "index": index
+                }
+            )
+
+            index += 1
+
+    return tiles
+
+
+def save_tiles(
+    tiles,
+    output_dir,
+    prefix="tile"
+):
+    """
+    Save generated tiles to disk.
+    """
+
+    output_dir = Path(output_dir)
+    output_dir.mkdir(
+        parents=True,
+        exist_ok=True
+    )
+
+    metadata = []
+
+    for tile_info in tiles:
+
+        image = tile_info["image"]
+
+        filename = (
+            f"{prefix}_"
+            f"{tile_info['index']:03d}_"
+            f"x{tile_info['x']}_"
+            f"y{tile_info['y']}.png"
+        )
+
+        path = output_dir / filename
+
+        image.save(path)
+
+        metadata.append(
+            {
+                "index": tile_info["index"],
+                "x": tile_info["x"],
+                "y": tile_info["y"],
+                "width": tile_info["width"],
+                "height": tile_info["height"],
+                "filename": filename
+            }
+        )
+
+    return metadata

model_router.py

@@ -0,0 +1,179 @@
+
+import gc
+import torch
+from pathlib import Path
+
+
+class SatelliteSRRouter:
+
+    def __init__(self, device=None):
+
+        from models.hatsat.hatsat_inference import HATSATInference
+        from models.esrgan.esrgan_inference import ESRGANInference
+
+        self.HATSATInference = HATSATInference
+        self.ESRGANInference = ESRGANInference
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.device = device
+
+        # Lazy loading:
+        # Only the selected model is loaded.
+        self.hatsat = None
+        self.esrgan = None
+
+        print("=" * 55)
+        print("SATELLITE SR ROUTER")
+        print("=" * 55)
+        print("Device:", self.device)
+        print("Router initialized.")
+        print("Models will be loaded only when selected.")
+
+    def _clear_gpu(self):
+
+        gc.collect()
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    def _unload_hatsat(self):
+
+        if self.hatsat is not None:
+
+            print("Unloading HATSAT...")
+
+            try:
+                if hasattr(self.hatsat, "model"):
+                    self.hatsat.model.cpu()
+            except Exception:
+                pass
+
+            self.hatsat = None
+
+            self._clear_gpu()
+
+    def _unload_esrgan(self):
+
+        if self.esrgan is not None:
+
+            print("Unloading ESRGAN...")
+
+            try:
+                if hasattr(self.esrgan, "model"):
+                    self.esrgan.model.cpu()
+            except Exception:
+                pass
+
+            self.esrgan = None
+
+            self._clear_gpu()
+
+    def _load_hatsat(self):
+
+        if self.hatsat is None:
+
+            print("Loading HATSAT...")
+
+            self.hatsat = self.HATSATInference(
+                device=self.device
+            )
+
+            print("HATSAT loaded successfully.")
+
+    def _load_esrgan(self):
+
+        if self.esrgan is None:
+
+            print("Loading ESRGAN...")
+
+            checkpoint = (
+                Path(__file__).resolve().parent
+                / "weights"
+                / "esrgan"
+                / "RRDB_ESRGAN_x4.pth"
+            )
+
+            self.esrgan = self.ESRGANInference(
+                checkpoint_path=checkpoint,
+                device=self.device
+            )
+
+            print("ESRGAN loaded successfully.")
+
+    def available_models(self):
+
+        return {
+            "hatsat": {
+                "name": "HATSAT",
+                "description": "Satellite-oriented super-resolution model",
+                "scale": 4
+            },
+
+            "esrgan": {
+                "name": "ESRGAN",
+                "description": "General-purpose super-resolution baseline",
+                "scale": 4
+            }
+        }
+
+    def predict(self, image, model_name="hatsat"):
+
+        if image is None:
+            raise ValueError("Please upload an image.")
+
+        model_name = str(model_name).lower().strip()
+
+        print("=" * 55)
+        print("ROUTER INFERENCE")
+        print("Selected model:", model_name)
+        print("=" * 55)
+
+        # ==========================================
+        # HATSAT
+        # ==========================================
+
+        if model_name == "hatsat":
+
+            # Free GPU memory used by ESRGAN
+            self._unload_esrgan()
+
+            # Load HATSAT only when required
+            self._load_hatsat()
+
+            print("Running HATSAT...")
+
+            result = self.hatsat.predict(image)
+
+            print("HATSAT inference complete.")
+
+            return result
+
+        # ==========================================
+        # ESRGAN
+        # ==========================================
+
+        elif model_name == "esrgan":
+
+            # Free GPU memory used by HATSAT
+            self._unload_hatsat()
+
+            # Load ESRGAN only when required
+            self._load_esrgan()
+
+            print("Running ESRGAN...")
+
+            result = self.esrgan.predict(image)
+
+            print("ESRGAN inference complete.")
+
+            return result
+
+        else:
+
+            raise ValueError(
+                f"Unknown model '{model_name}'. "
+                f"Choose 'hatsat' or 'esrgan'."
+            )

models/esrgan/RRDBNet_arch.py

@@ -0,0 +1,78 @@
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def make_layer(block, n_layers):
+    layers = []
+    for _ in range(n_layers):
+        layers.append(block())
+    return nn.Sequential(*layers)
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    def __init__(self, nf=64, gc=32, bias=True):
+        super(ResidualDenseBlock_5C, self).__init__()
+        # gc: growth channel, i.e. intermediate channels
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        # mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    '''Residual in Residual Dense Block'''
+
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
+        self.RDB3 = ResidualDenseBlock_5C(nf, gc)
+
+    def forward(self, x):
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32):
+        super(RRDBNet, self).__init__()
+        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
+
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+        self.RRDB_trunk = make_layer(RRDB_block_f, nb)
+        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        #### upsampling
+        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        fea = self.conv_first(x)
+        trunk = self.trunk_conv(self.RRDB_trunk(fea))
+        fea = fea + trunk
+
+        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
+        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.HRconv(fea)))
+
+        return out

models/esrgan/esrgan_inference.py

@@ -0,0 +1,188 @@
+
+import torch
+import numpy as np
+from PIL import Image
+
+from .RRDBNet_arch import RRDBNet
+
+
+class ESRGANInference:
+
+    def __init__(self, checkpoint_path, device=None):
+
+        self.checkpoint_path = str(checkpoint_path)
+
+        # --------------------------------------------------
+        # Device
+        # --------------------------------------------------
+
+        if device is None:
+            self.device = torch.device(
+                "cuda" if torch.cuda.is_available() else "cpu"
+            )
+        else:
+            self.device = torch.device(device)
+
+        print(f"ESRGAN device: {self.device}")
+
+        # --------------------------------------------------
+        # Original ESRGAN architecture
+        # --------------------------------------------------
+
+        self.model = RRDBNet(
+            3,
+            3,
+            64,
+            23,
+            gc=32
+        )
+
+        # --------------------------------------------------
+        # Load pretrained checkpoint
+        # --------------------------------------------------
+
+        checkpoint = torch.load(
+            self.checkpoint_path,
+            map_location="cpu"
+        )
+
+        # Different ESRGAN checkpoints may store weights
+        # directly or under a nested key.
+        if isinstance(checkpoint, dict):
+
+            if "params_ema" in checkpoint:
+                state_dict = checkpoint["params_ema"]
+
+            elif "params" in checkpoint:
+                state_dict = checkpoint["params"]
+
+            elif "state_dict" in checkpoint:
+                state_dict = checkpoint["state_dict"]
+
+            else:
+                state_dict = checkpoint
+
+        else:
+            state_dict = checkpoint
+
+        # Remove possible "module." prefix
+        cleaned_state_dict = {}
+
+        for key, value in state_dict.items():
+
+            if key.startswith("module."):
+                key = key[7:]
+
+            cleaned_state_dict[key] = value
+
+        self.model.load_state_dict(
+            cleaned_state_dict,
+            strict=True
+        )
+
+        print("ESRGAN checkpoint loaded successfully.")
+
+        self.model.eval()
+        self.model.to(self.device)
+
+    # ------------------------------------------------------
+    # Preprocessing
+    # ------------------------------------------------------
+
+    @staticmethod
+    def prepare_image(image):
+
+        if isinstance(image, Image.Image):
+
+            image = image.convert("RGB")
+            image = np.array(image)
+
+        elif isinstance(image, np.ndarray):
+
+            if image.ndim == 2:
+                image = np.stack(
+                    [image] * 3,
+                    axis=-1
+                )
+
+            if image.shape[-1] != 3:
+                raise ValueError(
+                    "Input image must have 3 RGB channels."
+                )
+
+            image = image.astype(np.uint8)
+
+        else:
+
+            raise TypeError(
+                "Input must be a PIL image or NumPy array."
+            )
+
+        tensor = torch.from_numpy(
+            image.astype(np.float32) / 255.0
+        )
+
+        tensor = tensor.permute(
+            2, 0, 1
+        )
+
+        tensor = tensor.unsqueeze(0)
+
+        return tensor
+
+    # ------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(self, image):
+
+        input_tensor = self.prepare_image(image)
+
+        device = next(self.model.parameters()).device
+        input_tensor = input_tensor.to(device)
+
+        # --------------------------------------------------
+        # ESRGAN 4x inference
+        # --------------------------------------------------
+
+        output = self.model(input_tensor)
+
+        # --------------------------------------------------
+        # Convert output to RGB uint8
+        # --------------------------------------------------
+
+        output = output.clamp(0, 1)
+
+        output = output.squeeze(0)
+
+        output = output.permute(
+            1, 2, 0
+        )
+
+        output = (
+            output.cpu().numpy() * 255.0
+        ).round().astype(np.uint8)
+
+        return Image.fromarray(
+            output,
+            mode="RGB"
+        )
+
+
+# ----------------------------------------------------------
+# Convenience function
+# ----------------------------------------------------------
+
+def run_esrgan(
+    image,
+    checkpoint_path,
+    device=None
+):
+
+    model = ESRGANInference(
+        checkpoint_path=checkpoint_path,
+        device=device
+    )
+
+    return model.predict(image)

models/hatsat/components.py

@@ -0,0 +1,451 @@
+"""
+HAT model components and building blocks.
+"""
+
+import torch
+import torch.nn as nn
+import math
+from einops import rearrange
+
+
+def to_2tuple(x):
+    """Convert input to tuple of length 2."""
+    if isinstance(x, (tuple, list)):
+        return tuple(x)
+    return (x, x)
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    """Truncated normal initialization."""
+    def norm_cdf(x):
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    with torch.no_grad():
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+        tensor.erfinv_()
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0], ) + (1, ) * (x.ndim - 1)
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, num_feat, squeeze_factor=16):
+        super(ChannelAttention, self).__init__()
+        self.attention = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
+            nn.Sigmoid())
+
+    def forward(self, x):
+        y = self.attention(x)
+        return x * y
+
+
+class CAB(nn.Module):
+    def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):
+        super(CAB, self).__init__()
+        self.cab = nn.Sequential(
+            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
+            nn.GELU(),
+            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
+            ChannelAttention(num_feat, squeeze_factor)
+        )
+
+    def forward(self, x):
+        return self.cab(x)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    b, h, w, c = x.shape
+    x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c)
+    return windows
+
+
+def window_reverse(windows, window_size, h, w):
+    b = int(windows.shape[0] / (h * w / window_size / window_size))
+    x = windows.view(b, h // window_size, w // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, rpi, mask=None):
+        b_, n, c = x.shape
+        qkv = self.qkv(x).reshape(b_, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nw = mask.shape[0]
+            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, n, n)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(b_, n, c)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class HAB(nn.Module):
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 compress_ratio=3, squeeze_factor=30, conv_scale=0.01, mlp_ratio=4.,
+                 qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, 'shift_size must in 0-window_size'
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.conv_scale = conv_scale
+        self.conv_block = CAB(num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, x_size, rpi_sa, attn_mask):
+        h, w = x_size
+        b, _, c = x.shape
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(b, h, w, c)
+
+        # Conv_X
+        conv_x = self.conv_block(x.permute(0, 3, 1, 2))
+        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = attn_mask
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, c)
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask)
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)
+        shifted_x = window_reverse(attn_windows, self.window_size, h, w)
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            attn_x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            attn_x = shifted_x
+        attn_x = attn_x.view(b, h * w, c)
+
+        # FFN
+        x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class OCAB(nn.Module):
+    def __init__(self, dim, input_resolution, window_size, overlap_ratio, num_heads,
+                 qkv_bias=True, qk_scale=None, mlp_ratio=2, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.window_size = window_size
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+        self.overlap_win_size = int(window_size * overlap_ratio) + window_size
+
+        self.norm1 = norm_layer(dim)
+        self.qkv = nn.Linear(dim, dim * 3,  bias=qkv_bias)
+        self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size),
+                               stride=window_size, padding=(self.overlap_win_size-window_size)//2)
+
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((window_size + self.overlap_win_size - 1) * (window_size + self.overlap_win_size - 1), num_heads))
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+        self.proj = nn.Linear(dim,dim)
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)
+
+    def forward(self, x, x_size, rpi):
+        h, w = x_size
+        b, _, c = x.shape
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(b, h, w, c)
+
+        qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2)
+        q = qkv[0].permute(0, 2, 3, 1)
+        kv = torch.cat((qkv[1], qkv[2]), dim=1)
+
+        # partition windows
+        q_windows = window_partition(q, self.window_size)
+        q_windows = q_windows.view(-1, self.window_size * self.window_size, c)
+
+        kv_windows = self.unfold(kv)
+        kv_windows = rearrange(kv_windows, 'b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch',
+                              nc=2, ch=c, owh=self.overlap_win_size, oww=self.overlap_win_size).contiguous()
+        k_windows, v_windows = kv_windows[0], kv_windows[1]
+
+        b_, nq, _ = q_windows.shape
+        _, n, _ = k_windows.shape
+        d = self.dim // self.num_heads
+        q = q_windows.reshape(b_, nq, self.num_heads, d).permute(0, 2, 1, 3)
+        k = k_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3)
+        v = v_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
+            self.window_size * self.window_size, self.overlap_win_size * self.overlap_win_size, -1)
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        attn = self.softmax(attn)
+        attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim)
+        x = window_reverse(attn_windows, self.window_size, h, w)
+        x = x.view(b, h * w, self.dim)
+
+        x = self.proj(x) + shortcut
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AttenBlocks(nn.Module):
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size, compress_ratio,
+                 squeeze_factor, conv_scale, overlap_ratio, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop=0., attn_drop=0., drop_path=0., norm_layer=nn.LayerNorm, downsample=None,
+                 use_checkpoint=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            HAB(dim=dim, input_resolution=input_resolution, num_heads=num_heads, window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2, compress_ratio=compress_ratio,
+                squeeze_factor=squeeze_factor, conv_scale=conv_scale, mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop, attn_drop=attn_drop,
+                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                norm_layer=norm_layer) for i in range(depth)
+        ])
+
+        # OCAB
+        self.overlap_attn = OCAB(dim=dim, input_resolution=input_resolution, window_size=window_size,
+                                overlap_ratio=overlap_ratio, num_heads=num_heads, qkv_bias=qkv_bias,
+                                qk_scale=qk_scale, mlp_ratio=mlp_ratio, norm_layer=norm_layer)
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, x_size, params):
+        for blk in self.blocks:
+            x = blk(x, x_size, params['rpi_sa'], params['attn_mask'])
+
+        x = self.overlap_attn(x, x_size, params['rpi_oca'])
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+
+class RHAG(nn.Module):
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size, compress_ratio,
+                 squeeze_factor, conv_scale, overlap_ratio, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop=0., attn_drop=0., drop_path=0., norm_layer=nn.LayerNorm, downsample=None,
+                 use_checkpoint=False, img_size=224, patch_size=4, resi_connection='1conv'):
+        super(RHAG, self).__init__()
+
+        self.dim = dim
+        self.input_resolution = input_resolution
+
+        self.residual_group = AttenBlocks(
+            dim=dim, input_resolution=input_resolution, depth=depth, num_heads=num_heads,
+            window_size=window_size, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor,
+            conv_scale=conv_scale, overlap_ratio=overlap_ratio, mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop, attn_drop=attn_drop,
+            drop_path=drop_path, norm_layer=norm_layer, downsample=downsample,
+            use_checkpoint=use_checkpoint)
+
+        if resi_connection == '1conv':
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == 'identity':
+            self.conv = nn.Identity()
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)
+
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)
+
+    def forward(self, x, x_size, params):
+        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size, params), x_size))) + x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        x = x.flatten(2).transpose(1, 2)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+
+class PatchUnEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+    def forward(self, x, x_size):
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], self.embed_dim, x_size[0], x_size[1])
+        return x
+
+
+class Upsample(nn.Sequential):
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)

models/hatsat/hat_model.py

@@ -0,0 +1,221 @@
+"""
+HAT (Hybrid Attention Transformer) main model implementation.
+"""
+
+import torch
+import torch.nn as nn
+import math
+
+from .components import (
+    RHAG, PatchEmbed, PatchUnEmbed, Upsample,
+    trunc_normal_, window_partition, to_2tuple
+)
+
+
+class HAT(nn.Module):
+    def __init__(self, img_size=64, patch_size=1, in_chans=3, embed_dim=96, depths=(6, 6, 6, 6),
+                 num_heads=(6, 6, 6, 6), window_size=7, compress_ratio=3, squeeze_factor=30,
+                 conv_scale=0.01, overlap_ratio=0.5, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, norm_layer=nn.LayerNorm,
+                 ape=False, patch_norm=True, use_checkpoint=False, upscale=2, img_range=1.,
+                 upsampler='', resi_connection='1conv', **kwargs):
+        super(HAT, self).__init__()
+
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.overlap_ratio = overlap_ratio
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.4488, 0.4371, 0.4040)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = upscale
+        self.upsampler = upsampler
+
+        # relative position index
+        relative_position_index_SA = self.calculate_rpi_sa()
+        relative_position_index_OCA = self.calculate_rpi_oca()
+        self.register_buffer('relative_position_index_SA', relative_position_index_SA)
+        self.register_buffer('relative_position_index_OCA', relative_position_index_OCA)
+
+        # shallow feature extraction
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # deep feature extraction
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = embed_dim
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # merge non-overlapping patches into image
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        # build Residual Hybrid Attention Groups (RHAG)
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = RHAG(
+                dim=embed_dim,
+                input_resolution=(patches_resolution[0], patches_resolution[1]),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                compress_ratio=compress_ratio,
+                squeeze_factor=squeeze_factor,
+                conv_scale=conv_scale,
+                overlap_ratio=overlap_ratio,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                norm_layer=norm_layer,
+                downsample=None,
+                use_checkpoint=use_checkpoint,
+                img_size=img_size,
+                patch_size=patch_size,
+                resi_connection=resi_connection)
+            self.layers.append(layer)
+        self.norm = norm_layer(self.num_features)
+
+        # build the last conv layer in deep feature extraction
+        if resi_connection == '1conv':
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == 'identity':
+            self.conv_after_body = nn.Identity()
+
+        # high quality image reconstruction
+        if self.upsampler == 'pixelshuffle':
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True))
+            self.upsample = Upsample(upscale, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def calculate_rpi_sa(self):
+        coords_h = torch.arange(self.window_size)
+        coords_w = torch.arange(self.window_size)
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
+        coords_flatten = torch.flatten(coords, 1)
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += self.window_size - 1
+        relative_coords[:, :, 1] += self.window_size - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size - 1
+        relative_position_index = relative_coords.sum(-1)
+        return relative_position_index
+
+    def calculate_rpi_oca(self):
+        window_size_ori = self.window_size
+        window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size)
+
+        coords_h = torch.arange(window_size_ori)
+        coords_w = torch.arange(window_size_ori)
+        coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w]))
+        coords_ori_flatten = torch.flatten(coords_ori, 1)
+
+        coords_h = torch.arange(window_size_ext)
+        coords_w = torch.arange(window_size_ext)
+        coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w]))
+        coords_ext_flatten = torch.flatten(coords_ext, 1)
+
+        relative_coords = coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None]
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += window_size_ori - window_size_ext + 1
+        relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1
+        relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1
+        relative_position_index = relative_coords.sum(-1)
+        return relative_position_index
+
+    def calculate_mask(self, x_size):
+        h, w = x_size
+        img_mask = torch.zeros((1, h, w, 1))
+        h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        return attn_mask
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x):
+        x_size = (x.shape[2], x.shape[3])
+
+        attn_mask = self.calculate_mask(x_size).to(x.device)
+        params = {'attn_mask': attn_mask, 'rpi_sa': self.relative_position_index_SA, 'rpi_oca': self.relative_position_index_OCA}
+
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x, x_size, params)
+
+        x = self.norm(x)
+        x = self.patch_unembed(x, x_size)
+        return x
+
+    def forward(self, x):
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        if self.upsampler == 'pixelshuffle':
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+
+        x = x / self.img_range + self.mean
+        return x

models/hatsat/hatsat_inference.py

@@ -0,0 +1,196 @@
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+# ------------------------------------------------------
+# Paths
+# ------------------------------------------------------
+
+HERE = Path(__file__).resolve().parent
+
+WEIGHTS = (
+    HERE.parent.parent
+    / "weights"
+    / "hatsat"
+    / "net_g_150000.pth"
+)
+
+# ------------------------------------------------------
+# Import HAT
+# ------------------------------------------------------
+
+from .hat_model import HAT
+
+
+class HATSATInference:
+
+    def __init__(self, checkpoint_path=None, device=None):
+
+        self.device = device or (
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        self.device = torch.device(self.device)
+
+        if checkpoint_path is None:
+            checkpoint_path = WEIGHTS
+
+        checkpoint_path = Path(checkpoint_path)
+
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(
+                f"HATSAT checkpoint not found: {checkpoint_path}"
+            )
+
+        # --------------------------------------------------
+        # HATSAT configuration
+        # --------------------------------------------------
+
+        self.model = HAT(
+            upscale=4,
+            in_chans=3,
+            img_size=128,
+            window_size=16,
+            compress_ratio=3,
+            squeeze_factor=30,
+            conv_scale=0.01,
+            overlap_ratio=0.5,
+            img_range=1.0,
+            depths=[6, 6, 6, 6, 6, 6],
+            embed_dim=180,
+            num_heads=[6, 6, 6, 6, 6, 6],
+            mlp_ratio=2,
+            upsampler="pixelshuffle",
+            resi_connection="1conv"
+        )
+
+        checkpoint = torch.load(
+            checkpoint_path,
+            map_location="cpu"
+        )
+
+        if "params_ema" in checkpoint:
+            state_dict = checkpoint["params_ema"]
+        elif "params" in checkpoint:
+            state_dict = checkpoint["params"]
+        else:
+            state_dict = checkpoint
+
+        self.model.load_state_dict(
+            state_dict,
+            strict=True
+        )
+
+        self.model.eval()
+        self.model.to(self.device)
+
+        print("HATSAT checkpoint loaded successfully.")
+
+    # ------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(self, image):
+
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        image = image.convert("RGB")
+
+        original_w, original_h = image.size
+
+        img = np.array(image).astype(
+            np.float32
+        ) / 255.0
+
+        device = next(self.model.parameters()).device
+        tensor = torch.from_numpy(
+            img
+        ).permute(
+            2, 0, 1
+        ).unsqueeze(0).to(device)
+
+        # --------------------------------------------------
+        # HAT requires dimensions compatible with window=16
+        # --------------------------------------------------
+
+        _, _, h, w = tensor.shape
+
+        window = 16
+
+        pad_h = (
+            window - h % window
+        ) % window
+
+        pad_w = (
+            window - w % window
+        ) % window
+
+        if pad_h > 0 or pad_w > 0:
+
+            tensor = torch.nn.functional.pad(
+                tensor,
+                (0, pad_w, 0, pad_h),
+                mode="reflect"
+            )
+
+        # --------------------------------------------------
+        # Model
+        # --------------------------------------------------
+
+        output = self.model(tensor)
+
+        # --------------------------------------------------
+        # Crop padding from output
+        # --------------------------------------------------
+
+        output_h = original_h * 4
+        output_w = original_w * 4
+
+        output = output[
+            :,
+            :,
+            :output_h,
+            :output_w
+        ]
+
+        output = (
+            output
+            .clamp(0, 1)
+            .squeeze(0)
+            .permute(1, 2, 0)
+            .cpu()
+            .numpy()
+        )
+
+        output = (
+            output * 255.0
+        ).round().astype(np.uint8)
+
+        return Image.fromarray(
+            output,
+            mode="RGB"
+        )
+
+
+# ------------------------------------------------------
+# Convenience function
+# ------------------------------------------------------
+
+def run_hatsat(
+    image,
+    checkpoint_path=None,
+    device=None
+):
+
+    model = HATSATInference(
+        checkpoint_path=checkpoint_path,
+        device=device
+    )
+
+    return model.predict(image)

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0.0, !=2.0.1
+numpy<2.0.0
+Pillow>=10.0.0
+gradio>=4.0.0
+einops==0.8.1
+rasterio==1.4.3
+spaces

weights/RRDB_ESRGAN_x4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65fece06e1ccb48853242aa972bdf00ad07a7dd8938d2dcbdf4221b59f6372ce
+size 66929193

weights/esrgan/RRDB_ESRGAN_x4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65fece06e1ccb48853242aa972bdf00ad07a7dd8938d2dcbdf4221b59f6372ce
+size 66929193

weights/hatsat/net_g_150000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54b2eda431411bfecc50c542ec863267ed683128b7aa4163c109776c627bb721
+size 170285859

weights/net_g_150000.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54b2eda431411bfecc50c542ec863267ed683128b7aa4163c109776c627bb721
+size 170285859