app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import math
import time
import torch
import torch.nn.functional as F
import numpy as np
import yaml
from pathlib import Path
from PIL import Image
from einops import rearrange, repeat
from safetensors.torch import load_file as safe_load
from huggingface_hub import hf_hub_download

import gradio as gr
from diffusers import AutoencoderKL

# ---- Model code from the RFMSR repo (bundled) ----
from models.rfmsr import create_rfmsr
from models.dinov2_encoder import create_dinov2_encoder
from utils.color_fix import apply_color_fix


# ================================================================================
# Config
# ================================================================================

MODEL_REPO = "frozen2001/RFMSR"
VAE_SUBDIR = "ckpts/stable-diffusion-2-1-base"
RFMSR_CKPT = "ckpts/rfmsr_os.safetensors"
MODEL_CONFIG = "configs/rfmsr.yaml"

FLOW_SIGMA = 1.0
DEFAULT_STEPS = 15
DEFAULT_SCALE = 4.0
DEFAULT_SEED = 42


# ================================================================================
# Model loading (module scope, eager .to("cuda"))
# ================================================================================

def _load_models():
    """Load SD2.1 VAE + RFMSR + DINOv2 at module scope."""
    # Download VAE config + weights from the HF model repo
    vae_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=f"{VAE_SUBDIR}/vae/diffusion_pytorch_model.safetensors",
    )
    vae_config_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=f"{VAE_SUBDIR}/vae/config.json",
    )
    vae_dir = os.path.dirname(vae_config_path)
    print(f"Loading SD2.1 VAE from {vae_dir} -> cuda...")
    ae = AutoencoderKL.from_pretrained(vae_dir)
    ae = ae.to("cuda").eval()
    ae.requires_grad_(False)
    print(f"  VAE scaling_factor: {ae.config.scaling_factor}")

    # Download RFMSR checkpoint
    rfmsr_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=RFMSR_CKPT,
    )
    print(f"Loading RFMSR from {rfmsr_path} ...")
    config_path = os.path.join(os.path.dirname(__file__), MODEL_CONFIG)
    rfmsr = create_rfmsr(config_path)
    sd = safe_load(rfmsr_path)
    sd.pop("ema_scale", None)
    sd = {"dit." + k if not k.startswith("dit.") else k: v for k, v in sd.items()}
    missing, unexpected = rfmsr.load_state_dict(sd, strict=False)
    rfmsr = rfmsr.to("cuda", dtype=torch.float32)
    rfmsr.eval()
    rfmsr.dit.use_checkpoint = False
    n = sum(p.numel() for p in rfmsr.parameters()) / 1e6
    print(f"  Params: {n:.2f}M")
    if missing:
        print(f"  Missing keys: {len(missing)}")
    if unexpected:
        print(f"  Unexpected keys: {len(unexpected)}")

    # DINOv2 encoder
    print("Loading DINOv2 encoder ...")
    venc = create_dinov2_encoder(config_path, device="cuda")
    if venc is not None:
        print("  DINOv2: loaded")

    print("All models loaded.")
    return ae, rfmsr, venc


ae, rfmsr, venc = _load_models()


# ================================================================================
# Inference helpers
# ================================================================================

@torch.no_grad()
def vae_encode(img: torch.Tensor) -> torch.Tensor:
    """img [-1,1] -> SD2.1 latent [B,4,H,W] (scaled)."""
    return ae.encode(img.float()).latent_dist.sample() * ae.config.scaling_factor


@torch.no_grad()
def vae_decode(latent: torch.Tensor) -> torch.Tensor:
    """SD2.1 latent [B,4,H,W] (scaled) -> pixel [0,1]."""
    latent = latent / ae.config.scaling_factor
    img = ae.decode(latent).sample
    return torch.clamp((img + 1.0) / 2.0, min=0.0, max=1.0)


@torch.no_grad()
def reverse_flow(z_lr: torch.Tensor, steps: int = 1,
                 flow_sigma: float = 1.0, seed: int = 42,
                 lr_pixel=None) -> torch.Tensor:
    """RFMSR reverse flow integration: t=1 -> t=0."""
    device = z_lr.device
    B, C, H, W = z_lr.shape

    venc_fea = None
    if venc is not None and lr_pixel is not None:
        venc_fea = venc(lr_pixel.float())

    timesteps = torch.linspace(1.0, 0.0, steps + 1, device=device)
    generator = torch.Generator(device=device).manual_seed(seed)
    x = z_lr + flow_sigma * torch.randn(B, C, H, W, generator=generator, device=device)

    step_pairs = list(zip(timesteps[:-1], timesteps[1:]))
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        for t_curr, t_prev in step_pairs:
            t_batch = torch.full((B,), t_curr, device=device)
            dt = t_prev - t_curr
            v = rfmsr(x, t_batch, z_lr, venc_fea=venc_fea).float()
            x = x + dt * v

    return x


@torch.no_grad()
def reverse_flow_tiled(z_lr: torch.Tensor, steps: int = 1,
                       flow_sigma: float = 1.0, seed: int = 42,
                       lt_size: int = 64, lt_stride: int = 32,
                       lr_pixel=None) -> torch.Tensor:
    """Per-step tiled velocity prediction with Gaussian-weighted blending."""
    import cv2
    device = z_lr.device
    B, C, H, W = z_lr.shape
    AE_FACTOR = 8

    def _make_tile_grid(length, tile, stride):
        if length <= tile:
            return [(0, length)]
        positions = list(range(0, length - tile + 1, stride))
        if positions[-1] + tile < length:
            positions.append(length - tile)
        return [(p, p + tile) for p in sorted(set(positions))]

    def _gaussian_weights(tile_h, tile_w, channels, device):
        def _kernel_1d(ksize):
            sigma = 0.3 * ((ksize - 1) * 0.5 - 1) + 0.8
            if ksize % 2 == 0:
                kernel = cv2.getGaussianKernel(ksize=ksize + 1, sigma=sigma, ktype=cv2.CV_64F)
                kernel = kernel[1:,]
            else:
                kernel = cv2.getGaussianKernel(ksize=ksize, sigma=sigma, ktype=cv2.CV_64F)
            return kernel
        kernel_h = _kernel_1d(tile_h)
        kernel_w = _kernel_1d(tile_w)
        w = np.matmul(kernel_h, kernel_w.T)
        w = torch.from_numpy(w).float().unsqueeze(0).unsqueeze(0)
        return w.to(device).expand(1, channels, -1, -1)

    h_tiles = _make_tile_grid(H, lt_size, lt_stride)
    w_tiles = _make_tile_grid(W, lt_size, lt_stride)

    tile_venc = {}
    use_venc = venc is not None and lr_pixel is not None
    if use_venc:
        with torch.no_grad():
            for hs, he in h_tiles:
                for ws, we in w_tiles:
                    ph_s, pw_s = hs * AE_FACTOR, ws * AE_FACTOR
                    ph_e = min(he * AE_FACTOR, lr_pixel.shape[2])
                    pw_e = min(we * AE_FACTOR, lr_pixel.shape[3])
                    lq_crop = lr_pixel[:, :, ph_s:ph_e, pw_s:pw_e]
                    tile_venc[(hs, ws)] = venc(lq_crop)

    timesteps = torch.linspace(1.0, 0.0, steps + 1, device=device)
    generator = torch.Generator(device=device).manual_seed(seed)
    x = z_lr + flow_sigma * torch.randn(B, C, H, W, generator=generator, device=device)

    g_weight = _gaussian_weights(lt_size, lt_size, C, device)

    step_pairs = list(zip(timesteps[:-1], timesteps[1:]))
    for t_curr, t_prev in step_pairs:
        t_batch = torch.full((B,), t_curr, device=device)
        dt = t_prev - t_curr

        v_acc = torch.zeros(B, C, H, W, device=device)
        w_acc = torch.zeros(B, C, H, W, device=device)

        for hs, he in h_tiles:
            for ws, we in w_tiles:
                x_tile = x[:, :, hs:he, ws:we]
                z_lr_tile = z_lr[:, :, hs:he, ws:we]
                tile_fea = tile_venc.get((hs, ws), None) if use_venc else None

                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    v_tile = rfmsr(x_tile, t_batch, z_lr_tile, venc_fea=tile_fea).float()

                v_acc[:, :, hs:he, ws:we] += v_tile * g_weight
                w_acc[:, :, hs:he, ws:we] += g_weight

        v_total = v_acc / w_acc.clamp(min=1e-8)
        x = x + dt * v_total

    return x


# ================================================================================
# Main inference function
# ================================================================================

@spaces.GPU(duration=120)
def super_resolve(input_image: "Image.Image", scale: float = 4.0, steps: int = 15,
                  flow_sigma: float = 1.0, seed: int = 42,
                  color_correction: str = "wavelet",
                  use_tiling: bool = True,
                  tile_size: int = 512, tile_stride: int = 256,
                  progress=gr.Progress(track_tqdm=True)):
    """Super-resolve a low-quality image using RFMSR (Residual Flow Matching).

    Args:
        input_image: Low-quality input image to upscale.
        scale: Upscale factor (e.g. 4.0 for 4x super-resolution).
        steps: Number of reverse integration steps (1 = one-step fast mode, 8-25 = multi-step high quality; default 15).
        flow_sigma: Noise standard deviation for the residual flow.
        seed: Random seed for reproducibility.
        color_correction: Color correction method: 'wavelet', 'adain', 'ycbcr', or 'none'.
        use_tiling: Enable tiled inference for large images (prevents OOM).
        tile_size: Pixel-space tile size for tiled inference.
        tile_stride: Sliding window stride for tiled inference.
    """
    t0 = time.perf_counter()

    AE_FACTOR = 8
    PATCH_SIZE = 2
    MOD_PIXEL = 16

    src = input_image.convert("RGB")
    exact_w = int(src.size[0] * scale)
    exact_h = int(src.size[1] * scale)
    target = src.resize((exact_w, exact_h), Image.BICUBIC)

    im_np = np.array(target).astype(np.float32) / 255.0
    im_cond = torch.from_numpy(np.moveaxis(im_np, 2, 0)).unsqueeze(0)
    im_cond = im_cond.to(dtype=torch.bfloat16, device="cuda")
    ori_h, ori_w = im_cond.shape[-2:]

    # Align to multiple of 16
    h, w = im_cond.shape[-2:]
    pad_h = (math.ceil(h / MOD_PIXEL) * MOD_PIXEL) - h
    pad_w = (math.ceil(w / MOD_PIXEL) * MOD_PIXEL) - w
    if pad_h > 0 or pad_w > 0:
        im_cond = F.pad(im_cond, (0, pad_w, 0, pad_h), mode="reflect")

    # VAE encode
    image_tensor = im_cond * 2.0 - 1.0
    z_lr = vae_encode(image_tensor)
    lh, lw = z_lr.shape[2], z_lr.shape[3]

    # LR pixel for DINOv2
    lr_pixel = im_cond.float() if venc is not None else None

    # Tile params
    lt_size = max((tile_size // AE_FACTOR // PATCH_SIZE) * PATCH_SIZE, PATCH_SIZE)
    lt_stride = max((tile_stride // AE_FACTOR // PATCH_SIZE) * PATCH_SIZE, PATCH_SIZE)
    lt_size = min(lt_size, min(lh, lw))
    lt_stride = min(lt_stride, lt_size)

    do_tiling = use_tiling and (lh > lt_size or lw > lt_size)

    if not do_tiling:
        z_hr = reverse_flow(z_lr, steps=steps, flow_sigma=flow_sigma,
                           seed=seed, lr_pixel=lr_pixel)
    else:
        z_hr = reverse_flow_tiled(z_lr, steps=steps, flow_sigma=flow_sigma,
                                  seed=seed, lt_size=lt_size, lt_stride=lt_stride,
                                  lr_pixel=lr_pixel)

    # VAE decode
    res_sr = vae_decode(z_hr)
    res_sr = res_sr[:, :, 0:ori_h, 0:ori_w]

    img = torch.clamp(res_sr, 0.0, 1.0)[0]
    decoded = 255.0 * np.moveaxis(img.cpu().float().detach().numpy(), 0, 2)
    decoded = decoded.astype(np.uint8)
    sr_image = Image.fromarray(decoded)

    # Color correction
    if color_correction != "none":
        sr_image = apply_color_fix(sr_image, target, method=color_correction)

    elapsed = time.perf_counter() - t0
    print(f"RFMSR inference done in {elapsed:.2f}s (steps={steps}, scale={scale}, tiling={do_tiling})")

    return sr_image


# ================================================================================
# Gradio UI
# ================================================================================

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown("""
    # RFMSR: Residual Flow Matching for Image Super-Resolution

    Upload a low-quality image and get a 4× super-resolved result using [RFMSR](https://huggingface.co/papers/2607.12753).

    The model uses residual flow matching in the latent space of the SD2.1 VAE with DINOv2 semantic guidance.
    """)

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            input_img = gr.Image(label="Low-quality input", type="pil")
            run_btn = gr.Button("Super-Resolve", variant="primary")
        with gr.Column(scale=1):
            output_img = gr.Image(label="Super-resolved output", type="pil")

    with gr.Accordion("Advanced settings", open=False):
        with gr.Row():
            scale_slider = gr.Slider(label="Upscale factor", minimum=2.0, maximum=8.0, value=4.0, step=0.5)
            steps_slider = gr.Slider(label="Sampling steps (1=fast, 8-25=sharp/high quality)", minimum=1, maximum=30, value=15, step=1)
        with gr.Row():
            sigma_slider = gr.Slider(label="Flow sigma (noise std)", minimum=0.0, maximum=2.0, value=1.0, step=0.1)
            seed_input = gr.Number(label="Seed", value=42, precision=0)
        color_correction = gr.Dropdown(
            label="Color correction",
            choices=["wavelet", "adain", "ycbcr", "none"],
            value="wavelet",
        )
        with gr.Row():
            use_tiling = gr.Checkbox(label="Enable tiled inference (for large images)", value=True)
            tile_size = gr.Slider(label="Tile size (px)", minimum=128, maximum=1024, value=512, step=64)
            tile_stride = gr.Slider(label="Tile stride (px)", minimum=64, maximum=512, value=256, step=32)

    gr.Examples(
        examples=[
            ["example1.png"],
            ["example2.png"],
            ["example3.png"],
        ],
        inputs=[input_img],
        outputs=[output_img],
        fn=super_resolve,
        cache_examples=True,
        cache_mode="lazy",
    )

    gr.Markdown("""
    ### Links
    - [Paper (arXiv 2607.12753)](https://arxiv.org/abs/2607.12753)
    - [GitHub](https://github.com/Faze-Hsw/RFMSR)
    - [HF Model](https://huggingface.co/frozen2001/RFMSR)
    """)

    run_btn.click(
        fn=super_resolve,
        inputs=[input_img, scale_slider, steps_slider, sigma_slider, seed_input,
                color_correction, use_tiling, tile_size, tile_stride],
        outputs=[output_img],
        api_name="super_resolve",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus())