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())