luboslenco
841 lines
27 KiB
C
841 lines
27 KiB
C
/*
|
|
* iris_upscale.c - RealESRGAN_x4plus image upscaler (RRDBNet)
|
|
*
|
|
* Pure-C, CPU-only implementation of the RRDBNet super-resolution network
|
|
* used by RealESRGAN_x4plus. The architecture (from BasicSR) is:
|
|
*
|
|
* feat = conv_first(x) # 3 -> 64
|
|
* body_feat = conv_body(body(feat)) # 23x RRDB, then 64 -> 64
|
|
* feat = feat + body_feat
|
|
* feat = lrelu(conv_up1(nearest_2x(feat))) # upsample x2
|
|
* feat = lrelu(conv_up2(nearest_2x(feat))) # upsample x2 (total x4)
|
|
* out = conv_last(lrelu(conv_hr(feat))) # 64 -> 64 -> 3
|
|
*
|
|
* Each RRDB is three ResidualDenseBlocks with a 0.2 residual scale; each
|
|
* ResidualDenseBlock is five 3x3 convs over densely concatenated features
|
|
* (LeakyReLU 0.2 between them), again with a 0.2 residual scale.
|
|
*
|
|
* All convolutions are 3x3, stride 1, zero-padded. Weights are read directly
|
|
* from the memory-mapped safetensors file (all tensors are F32), so there is
|
|
* no separate weight copy.
|
|
*/
|
|
|
|
#include "iris_upscale.h"
|
|
#include "iris_safetensors.h"
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#if defined(USE_METAL)
|
|
#include "iris_metal.h"
|
|
#elif defined(USE_VULKAN)
|
|
#include "iris_vulkan.h"
|
|
#endif
|
|
|
|
#if defined(USE_METAL)
|
|
#define IRIS_GPU_AVAILABLE() iris_metal_available()
|
|
#elif defined(USE_VULKAN)
|
|
#define IRIS_GPU_AVAILABLE() iris_vulkan_available()
|
|
#endif
|
|
|
|
/* RRDBNet hyperparameters for RealESRGAN_x4plus */
|
|
#define RG_NUM_FEAT 64
|
|
#define RG_NUM_GROW_CH 32
|
|
#define RG_NUM_BLOCK 23
|
|
#define RG_SCALE 4
|
|
#define RG_LRELU_SLOPE 0.2f
|
|
#define RG_RES_SCALE 0.2f
|
|
|
|
/* Densely-concatenated width inside a ResidualDenseBlock:
|
|
* num_feat + 4*num_grow_ch = 64 + 128 = 192 channels. */
|
|
#define RG_CAT_CH (RG_NUM_FEAT + 4 * RG_NUM_GROW_CH)
|
|
|
|
struct iris_upscale {
|
|
safetensors_file_t *sf;
|
|
int tileable; /* make the upscaled image wrap seamlessly */
|
|
};
|
|
|
|
/* ========================================================================
|
|
* Weight access (zero-copy; all tensors are F32 in the mmap'd file)
|
|
* ======================================================================== */
|
|
|
|
static const float *rg_tensor(iris_upscale_t *m, const char *name) {
|
|
const safetensor_t *t = safetensors_find(m->sf, name);
|
|
if (!t) {
|
|
fprintf(stderr, "RealESRGAN: missing tensor '%s'\n", name);
|
|
return NULL;
|
|
}
|
|
if (t->dtype != DTYPE_F32) {
|
|
fprintf(stderr, "RealESRGAN: tensor '%s' is not F32\n", name);
|
|
return NULL;
|
|
}
|
|
return (const float *)safetensors_data(m->sf, t);
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Primitive ops (CPU)
|
|
* ======================================================================== */
|
|
|
|
/* 3x3 zero-padded, stride-1 convolution.
|
|
* in: [Cin, H, W]
|
|
* weight: [Cout, Cin, 3, 3]
|
|
* bias: [Cout]
|
|
* out: [Cout, H, W] (must not alias in)
|
|
* The inner loop runs contiguously over x for good auto-vectorization; the
|
|
* valid x-range per kernel tap is precomputed so there is no per-pixel branch.
|
|
*/
|
|
static void conv2d_3x3(const float *in, int Cin, int H, int W, const float *weight, const float *bias, int Cout, float *out) {
|
|
const size_t plane = (size_t)H * W;
|
|
|
|
#ifdef _OPENMP
|
|
#pragma omp parallel for schedule(static)
|
|
#endif
|
|
for (int oc = 0; oc < Cout; oc++) {
|
|
float *op = out + (size_t)oc * plane;
|
|
float b = bias[oc];
|
|
for (size_t i = 0; i < plane; i++)
|
|
op[i] = b;
|
|
|
|
for (int ic = 0; ic < Cin; ic++) {
|
|
const float *ip = in + (size_t)ic * plane;
|
|
const float *wk = weight + ((size_t)oc * Cin + ic) * 9;
|
|
for (int ky = 0; ky < 3; ky++) {
|
|
int dy = ky - 1;
|
|
for (int kx = 0; kx < 3; kx++) {
|
|
float wv = wk[ky * 3 + kx];
|
|
if (wv == 0.0f)
|
|
continue;
|
|
int dx = kx - 1;
|
|
int x0 = (dx < 0) ? 1 : 0;
|
|
int x1 = (dx > 0) ? W - 1 : W;
|
|
for (int y = 0; y < H; y++) {
|
|
int sy = y + dy;
|
|
if (sy < 0 || sy >= H)
|
|
continue;
|
|
float *orow = op + (size_t)y * W;
|
|
const float *irow = ip + (size_t)sy * W + dx;
|
|
for (int x = x0; x < x1; x++)
|
|
orow[x] += wv * irow[x];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/* In-place LeakyReLU over a [C, H, W] buffer. */
|
|
static void leaky_relu(float *buf, size_t n) {
|
|
for (size_t i = 0; i < n; i++)
|
|
if (buf[i] < 0.0f)
|
|
buf[i] *= RG_LRELU_SLOPE;
|
|
}
|
|
|
|
/* Nearest-neighbour 2x upsample: in [C,H,W] -> out [C,2H,2W]. */
|
|
static void upsample_nearest_2x(const float *in, int C, int H, int W, float *out) {
|
|
int OH = H * 2, OW = W * 2;
|
|
const size_t in_plane = (size_t)H * W;
|
|
const size_t out_plane = (size_t)OH * OW;
|
|
for (int c = 0; c < C; c++) {
|
|
const float *ip = in + (size_t)c * in_plane;
|
|
float *op = out + (size_t)c * out_plane;
|
|
for (int y = 0; y < OH; y++) {
|
|
const float *irow = ip + (size_t)(y / 2) * W;
|
|
float *orow = op + (size_t)y * OW;
|
|
for (int x = 0; x < OW; x++)
|
|
orow[x] = irow[x / 2];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ========================================================================
|
|
* RRDBNet blocks
|
|
* ======================================================================== */
|
|
|
|
/* ResidualDenseBlock: in/out are [num_feat, H, W]; out may equal in only if
|
|
* caller is fine with that (here we always use distinct buffers).
|
|
* `prefix` is e.g. "body.0.rdb1". `scratch_cat` is [RG_CAT_CH, H, W] and
|
|
* `scratch_x5` is [num_feat, H, W]; both are caller-owned work buffers. */
|
|
static int residual_dense_block(iris_upscale_t *m, const char *prefix, const float *in, int H, int W, float *scratch_cat, float *scratch_x5, float *out) {
|
|
const size_t plane = (size_t)H * W;
|
|
char name[256];
|
|
|
|
/* Concatenation buffer starts with the input features (channels 0..63). */
|
|
memcpy(scratch_cat, in, RG_NUM_FEAT * plane * sizeof(float));
|
|
|
|
/* conv1..conv4 each produce num_grow_ch channels appended to the cat
|
|
* buffer; conv reads cat[0:cur] and writes cat[cur:cur+grow] (disjoint). */
|
|
int cur = RG_NUM_FEAT;
|
|
for (int i = 1; i <= 4; i++) {
|
|
snprintf(name, sizeof(name), "%s.conv%d.weight", prefix, i);
|
|
const float *w = rg_tensor(m, name);
|
|
snprintf(name, sizeof(name), "%s.conv%d.bias", prefix, i);
|
|
const float *b = rg_tensor(m, name);
|
|
if (!w || !b)
|
|
return -1;
|
|
|
|
float *dst = scratch_cat + (size_t)cur * plane;
|
|
conv2d_3x3(scratch_cat, cur, H, W, w, b, RG_NUM_GROW_CH, dst);
|
|
leaky_relu(dst, (size_t)RG_NUM_GROW_CH * plane);
|
|
cur += RG_NUM_GROW_CH;
|
|
}
|
|
|
|
/* conv5: cat[0:192] -> num_feat channels (no activation). */
|
|
snprintf(name, sizeof(name), "%s.conv5.weight", prefix);
|
|
const float *w5 = rg_tensor(m, name);
|
|
snprintf(name, sizeof(name), "%s.conv5.bias", prefix);
|
|
const float *b5 = rg_tensor(m, name);
|
|
if (!w5 || !b5)
|
|
return -1;
|
|
conv2d_3x3(scratch_cat, cur, H, W, w5, b5, RG_NUM_FEAT, scratch_x5);
|
|
|
|
/* out = x5 * 0.2 + in */
|
|
size_t n = (size_t)RG_NUM_FEAT * plane;
|
|
for (size_t i = 0; i < n; i++)
|
|
out[i] = scratch_x5[i] * RG_RES_SCALE + in[i];
|
|
return 0;
|
|
}
|
|
|
|
/* RRDB: three ResidualDenseBlocks then a 0.2 residual.
|
|
* `tmp_a`/`tmp_b` are [num_feat,H,W] ping-pong buffers. */
|
|
static int rrdb(iris_upscale_t *m, int block_idx, const float *in, int H, int W, float *scratch_cat, float *scratch_x5, float *tmp_a, float *tmp_b,
|
|
float *out) {
|
|
char prefix[64];
|
|
const size_t plane = (size_t)H * W;
|
|
|
|
snprintf(prefix, sizeof(prefix), "body.%d.rdb1", block_idx);
|
|
if (residual_dense_block(m, prefix, in, H, W, scratch_cat, scratch_x5, tmp_a) != 0)
|
|
return -1;
|
|
snprintf(prefix, sizeof(prefix), "body.%d.rdb2", block_idx);
|
|
if (residual_dense_block(m, prefix, tmp_a, H, W, scratch_cat, scratch_x5, tmp_b) != 0)
|
|
return -1;
|
|
snprintf(prefix, sizeof(prefix), "body.%d.rdb3", block_idx);
|
|
if (residual_dense_block(m, prefix, tmp_b, H, W, scratch_cat, scratch_x5, tmp_a) != 0)
|
|
return -1;
|
|
|
|
/* out = rdb_out * 0.2 + in */
|
|
size_t n = (size_t)RG_NUM_FEAT * plane;
|
|
for (size_t i = 0; i < n; i++)
|
|
out[i] = tmp_a[i] * RG_RES_SCALE + in[i];
|
|
return 0;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Input / output conversion (shared by CPU and GPU paths)
|
|
* ======================================================================== */
|
|
|
|
/* RGB(A) uint8 image -> planar CHW float in [0,1], 3 channels.
|
|
* Returns a malloc'd [3*H*W] buffer (caller frees), or NULL on OOM. */
|
|
static float *rg_build_input(const iris_image *input) {
|
|
int H = input->height, W = input->width, ic = input->channels;
|
|
const size_t plane = (size_t)H * W;
|
|
float *x_in = malloc((size_t)3 * plane * sizeof(float));
|
|
if (!x_in)
|
|
return NULL;
|
|
for (int y = 0; y < H; y++) {
|
|
for (int xx = 0; xx < W; xx++) {
|
|
const uint8_t *px = input->data + ((size_t)y * W + xx) * ic;
|
|
size_t pi = (size_t)y * W + xx;
|
|
x_in[0 * plane + pi] = px[0] / 255.0f;
|
|
x_in[1 * plane + pi] = (ic > 1 ? px[1] : px[0]) / 255.0f;
|
|
x_in[2 * plane + pi] = (ic > 2 ? px[2] : px[0]) / 255.0f;
|
|
}
|
|
}
|
|
return x_in;
|
|
}
|
|
|
|
/* Planar CHW float [3, H4, W4] in [0,1] -> RGB uint8 image. */
|
|
static iris_image *rg_chw_to_image(const float *out, int W4, int H4) {
|
|
const size_t plane4 = (size_t)H4 * W4;
|
|
iris_image *result = iris_image_create(W4, H4, 3);
|
|
if (!result)
|
|
return NULL;
|
|
for (size_t pi = 0; pi < plane4; pi++) {
|
|
for (int c = 0; c < 3; c++) {
|
|
float v = out[(size_t)c * plane4 + pi];
|
|
if (v < 0.0f)
|
|
v = 0.0f;
|
|
if (v > 1.0f)
|
|
v = 1.0f;
|
|
result->data[pi * 3 + c] = (uint8_t)(v * 255.0f + 0.5f);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
/* ========================================================================
|
|
* Seamless tiling via periodic+smooth decomposition (Moisan 2011)
|
|
*
|
|
* A 4x upscale of a tileable texture is not itself tileable: the network has no
|
|
* notion of wrap-around, so opposite borders no longer match. Rather than
|
|
* blend edges per-row/column (which streaks any line whose two ends differ
|
|
* sharply), we subtract the single smoothest 2D field that makes the image
|
|
* periodic -- Moisan's "smooth component", the solution of a periodic Poisson
|
|
* equation whose source is the jump between opposite borders. A localized edge
|
|
* mismatch is diffused into a gentle bump instead of smeared along a scanline,
|
|
* so the result tiles exactly. Solved in the Fourier domain, so it requires
|
|
* power-of-two dimensions (always the case for textures); other sizes are left
|
|
* unchanged.
|
|
* ======================================================================== */
|
|
|
|
#ifndef IRIS_TWO_PI
|
|
#define IRIS_TWO_PI 6.28318530717958647692
|
|
#endif
|
|
|
|
static int rg_is_pow2(int n) {
|
|
return n > 0 && (n & (n - 1)) == 0;
|
|
}
|
|
|
|
/* In-place iterative radix-2 Cooley-Tukey FFT. inv=0 forward (e^-i), inv=1
|
|
* inverse (e^+i, scaled by 1/n). n must be a power of two. */
|
|
static void rg_fft1d(double *re, double *im, int n, int inv) {
|
|
for (int i = 1, j = 0; i < n; i++) {
|
|
int bit = n >> 1;
|
|
for (; j & bit; bit >>= 1)
|
|
j ^= bit;
|
|
j ^= bit;
|
|
if (i < j) {
|
|
double tr = re[i];
|
|
re[i] = re[j];
|
|
re[j] = tr;
|
|
double ti = im[i];
|
|
im[i] = im[j];
|
|
im[j] = ti;
|
|
}
|
|
}
|
|
for (int len = 2; len <= n; len <<= 1) {
|
|
double ang = IRIS_TWO_PI / len * (inv ? 1.0 : -1.0);
|
|
double wr = cos(ang), wi = sin(ang);
|
|
for (int i = 0; i < n; i += len) {
|
|
double cwr = 1.0, cwi = 0.0;
|
|
for (int k = 0; k < len / 2; k++) {
|
|
int a = i + k, b = i + k + len / 2;
|
|
double vr = re[b] * cwr - im[b] * cwi;
|
|
double vi = re[b] * cwi + im[b] * cwr;
|
|
re[b] = re[a] - vr;
|
|
im[b] = im[a] - vi;
|
|
re[a] += vr;
|
|
im[a] += vi;
|
|
double ncwr = cwr * wr - cwi * wi;
|
|
cwi = cwr * wi + cwi * wr;
|
|
cwr = ncwr;
|
|
}
|
|
}
|
|
}
|
|
if (inv) {
|
|
for (int i = 0; i < n; i++) {
|
|
re[i] /= n;
|
|
im[i] /= n;
|
|
}
|
|
}
|
|
}
|
|
|
|
/* 2D FFT: transform every row (length W), then every column (length H), using
|
|
* caller-provided column scratch of length >= H. */
|
|
static void rg_fft2d(double *re, double *im, int H, int W, int inv, double *cr, double *ci) {
|
|
for (int i = 0; i < H; i++)
|
|
rg_fft1d(re + (size_t)i * W, im + (size_t)i * W, W, inv);
|
|
for (int j = 0; j < W; j++) {
|
|
for (int i = 0; i < H; i++) {
|
|
cr[i] = re[(size_t)i * W + j];
|
|
ci[i] = im[(size_t)i * W + j];
|
|
}
|
|
rg_fft1d(cr, ci, H, inv);
|
|
for (int i = 0; i < H; i++) {
|
|
re[(size_t)i * W + j] = cr[i];
|
|
im[(size_t)i * W + j] = ci[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Subtract the smooth component so img becomes seamlessly tileable. Only acts on
|
|
* power-of-two dimensions (textures always qualify); other sizes are untouched. */
|
|
static void make_seamless_poisson(float *img, int H, int W) {
|
|
if (!rg_is_pow2(H) || !rg_is_pow2(W) || H < 2 || W < 2)
|
|
return;
|
|
|
|
size_t n = (size_t)H * W;
|
|
double *re = calloc(n, sizeof(double));
|
|
double *im = calloc(n, sizeof(double));
|
|
double *cr = malloc((size_t)H * sizeof(double));
|
|
double *ci = malloc((size_t)H * sizeof(double));
|
|
if (!re || !im || !cr || !ci) {
|
|
free(re);
|
|
free(im);
|
|
free(cr);
|
|
free(ci);
|
|
return;
|
|
}
|
|
|
|
/* Boundary jump field: the wrap-around difference across each border,
|
|
* accumulated (corners receive both a row and a column contribution). */
|
|
for (int i = 0; i < H; i++) {
|
|
double *row = re + (size_t)i * W;
|
|
double l = img[(size_t)i * W + 0], r = img[(size_t)i * W + (W - 1)];
|
|
row[0] += r - l;
|
|
row[W - 1] += l - r;
|
|
}
|
|
for (int j = 0; j < W; j++) {
|
|
double t = img[(size_t)0 * W + j], b = img[(size_t)(H - 1) * W + j];
|
|
re[(size_t)0 * W + j] += b - t;
|
|
re[(size_t)(H - 1) * W + j] += t - b;
|
|
}
|
|
|
|
rg_fft2d(re, im, H, W, 0, cr, ci);
|
|
|
|
/* Divide by the periodic-Laplacian eigenvalues to solve the Poisson eq;
|
|
* the DC term (constant offset) is undetermined, so pin it to zero. */
|
|
for (int q = 0; q < H; q++) {
|
|
double cq = 2.0 * cos(IRIS_TWO_PI * q / H);
|
|
for (int j = 0; j < W; j++) {
|
|
size_t idx = (size_t)q * W + j;
|
|
if (q == 0 && j == 0) {
|
|
re[idx] = 0.0;
|
|
im[idx] = 0.0;
|
|
continue;
|
|
}
|
|
double denom = cq + 2.0 * cos(IRIS_TWO_PI * j / W) - 4.0;
|
|
re[idx] /= denom;
|
|
im[idx] /= denom;
|
|
}
|
|
}
|
|
|
|
rg_fft2d(re, im, H, W, 1, cr, ci);
|
|
|
|
for (size_t i = 0; i < n; i++)
|
|
img[i] -= (float)re[i];
|
|
|
|
free(re);
|
|
free(im);
|
|
free(cr);
|
|
free(ci);
|
|
}
|
|
|
|
/* Make a planar CHW [3, H, W] float image tile seamlessly, per channel. */
|
|
static void rg_make_tileable(float *chw, int H, int W) {
|
|
const size_t plane = (size_t)H * W;
|
|
for (int c = 0; c < 3; c++)
|
|
make_seamless_poisson(chw + (size_t)c * plane, H, W);
|
|
}
|
|
|
|
/* ========================================================================
|
|
* GPU-resident forward path (Metal or Vulkan)
|
|
*
|
|
* Mirrors the CPU forward, but every convolution / activation / residual runs
|
|
* on the GPU and activations stay resident in VRAM between ops. Convolution is
|
|
* the entire cost of RRDBNet (~350 3x3 convs), so offloading it is the win.
|
|
* Both backends expose the same iris_gpu_* tensor surface, so this path is
|
|
* shared verbatim between them.
|
|
* ======================================================================== */
|
|
#if defined(USE_METAL) || defined(USE_VULKAN)
|
|
|
|
/* Run a named 3x3/pad-1/stride-1 conv on the GPU. Weights are F32 in the
|
|
* mmap'd file and cached in VRAM by pointer across the run. */
|
|
static iris_gpu_tensor_t rg_conv_gpu(iris_upscale_t *m, const char *prefix, iris_gpu_tensor_t x, int Cin, int H, int W, int Cout) {
|
|
char name[256];
|
|
snprintf(name, sizeof(name), "%s.weight", prefix);
|
|
const float *w = rg_tensor(m, name);
|
|
snprintf(name, sizeof(name), "%s.bias", prefix);
|
|
const float *b = rg_tensor(m, name);
|
|
if (!w || !b)
|
|
return NULL;
|
|
return iris_gpu_conv2d_f32(x, w, b, 1, Cin, Cout, H, W, 3, 3, 1, 1);
|
|
}
|
|
|
|
/* ResidualDenseBlock on GPU. `in` is [num_feat,H,W]; returns a new tensor. */
|
|
static iris_gpu_tensor_t rdb_gpu(iris_upscale_t *m, const char *prefix, iris_gpu_tensor_t in, int H, int W) {
|
|
const size_t plane = (size_t)H * W;
|
|
char name[64];
|
|
|
|
iris_gpu_tensor_t cat = iris_gpu_tensor_alloc((size_t)RG_CAT_CH * plane);
|
|
if (!cat)
|
|
return NULL;
|
|
iris_gpu_copy_region_f32(cat, 0, in, 0, (size_t)RG_NUM_FEAT * plane);
|
|
|
|
int cur = RG_NUM_FEAT;
|
|
for (int i = 1; i <= 4; i++) {
|
|
snprintf(name, sizeof(name), "%s.conv%d", prefix, i);
|
|
iris_gpu_tensor_t gch = rg_conv_gpu(m, name, cat, cur, H, W, RG_NUM_GROW_CH);
|
|
if (!gch) {
|
|
iris_gpu_tensor_free(cat);
|
|
return NULL;
|
|
}
|
|
iris_gpu_leaky_relu_f32(gch, gch, (int)(RG_NUM_GROW_CH * plane), RG_LRELU_SLOPE);
|
|
iris_gpu_copy_region_f32(cat, (size_t)cur * plane, gch, 0, (size_t)RG_NUM_GROW_CH * plane);
|
|
iris_gpu_tensor_free(gch);
|
|
cur += RG_NUM_GROW_CH;
|
|
}
|
|
|
|
snprintf(name, sizeof(name), "%s.conv5", prefix);
|
|
iris_gpu_tensor_t x5 = rg_conv_gpu(m, name, cat, cur, H, W, RG_NUM_FEAT);
|
|
iris_gpu_tensor_free(cat);
|
|
if (!x5)
|
|
return NULL;
|
|
|
|
iris_gpu_tensor_t out = iris_gpu_tensor_alloc((size_t)RG_NUM_FEAT * plane);
|
|
if (!out) {
|
|
iris_gpu_tensor_free(x5);
|
|
return NULL;
|
|
}
|
|
/* out = x5 * 0.2 + in */
|
|
iris_gpu_scale_add_f32(out, x5, in, RG_RES_SCALE, (int)(RG_NUM_FEAT * plane));
|
|
iris_gpu_tensor_free(x5);
|
|
return out;
|
|
}
|
|
|
|
/* RRDB on GPU: three RDBs then a 0.2 residual. Returns a new tensor. */
|
|
static iris_gpu_tensor_t rrdb_gpu(iris_upscale_t *m, int idx, iris_gpu_tensor_t in, int H, int W) {
|
|
const size_t plane = (size_t)H * W;
|
|
char p[64];
|
|
|
|
snprintf(p, sizeof(p), "body.%d.rdb1", idx);
|
|
iris_gpu_tensor_t a = rdb_gpu(m, p, in, H, W);
|
|
if (!a)
|
|
return NULL;
|
|
snprintf(p, sizeof(p), "body.%d.rdb2", idx);
|
|
iris_gpu_tensor_t b = rdb_gpu(m, p, a, H, W);
|
|
iris_gpu_tensor_free(a);
|
|
if (!b)
|
|
return NULL;
|
|
snprintf(p, sizeof(p), "body.%d.rdb3", idx);
|
|
iris_gpu_tensor_t c = rdb_gpu(m, p, b, H, W);
|
|
iris_gpu_tensor_free(b);
|
|
if (!c)
|
|
return NULL;
|
|
|
|
iris_gpu_tensor_t out = iris_gpu_tensor_alloc((size_t)RG_NUM_FEAT * plane);
|
|
if (!out) {
|
|
iris_gpu_tensor_free(c);
|
|
return NULL;
|
|
}
|
|
iris_gpu_scale_add_f32(out, c, in, RG_RES_SCALE, (int)(RG_NUM_FEAT * plane));
|
|
iris_gpu_tensor_free(c);
|
|
return out;
|
|
}
|
|
|
|
static iris_image *upscale_gpu(iris_upscale_t *m, const iris_image *input) {
|
|
int H = input->height, W = input->width;
|
|
const size_t plane = (size_t)H * W;
|
|
|
|
float *x_in = rg_build_input(input);
|
|
if (!x_in)
|
|
return NULL;
|
|
iris_gpu_tensor_t x = iris_gpu_tensor_create(x_in, 3 * plane);
|
|
free(x_in);
|
|
if (!x)
|
|
return NULL;
|
|
|
|
/* feat = conv_first(x) (3 -> 64); feat is kept for the body residual. */
|
|
iris_gpu_tensor_t feat = rg_conv_gpu(m, "conv_first", x, 3, H, W, RG_NUM_FEAT);
|
|
iris_gpu_tensor_free(x);
|
|
if (!feat)
|
|
return NULL;
|
|
|
|
/* body: 23 RRDB blocks. `cur` is the running activation (!= feat). */
|
|
iris_gpu_tensor_t cur = feat;
|
|
for (int i = 0; i < RG_NUM_BLOCK; i++) {
|
|
iris_gpu_tensor_t nxt = rrdb_gpu(m, i, cur, H, W);
|
|
if (cur != feat)
|
|
iris_gpu_tensor_free(cur);
|
|
if (!nxt) {
|
|
iris_gpu_tensor_free(feat);
|
|
return NULL;
|
|
}
|
|
cur = nxt;
|
|
}
|
|
|
|
/* body_feat = conv_body(body(feat)); feat = feat + body_feat */
|
|
iris_gpu_tensor_t body_feat = rg_conv_gpu(m, "conv_body", cur, RG_NUM_FEAT, H, W, RG_NUM_FEAT);
|
|
if (cur != feat)
|
|
iris_gpu_tensor_free(cur);
|
|
if (!body_feat) {
|
|
iris_gpu_tensor_free(feat);
|
|
return NULL;
|
|
}
|
|
iris_gpu_add_f32(feat, feat, body_feat, (int)(RG_NUM_FEAT * plane));
|
|
iris_gpu_tensor_free(body_feat);
|
|
|
|
/* Upsample stage 1: nearest 2x -> lrelu(conv_up1) */
|
|
int H2 = H * 2, W2 = W * 2;
|
|
iris_gpu_tensor_t up1 = iris_gpu_upsample_nearest_2x_f32(feat, RG_NUM_FEAT, H, W);
|
|
iris_gpu_tensor_free(feat);
|
|
if (!up1)
|
|
return NULL;
|
|
iris_gpu_tensor_t f2 = rg_conv_gpu(m, "conv_up1", up1, RG_NUM_FEAT, H2, W2, RG_NUM_FEAT);
|
|
iris_gpu_tensor_free(up1);
|
|
if (!f2)
|
|
return NULL;
|
|
iris_gpu_leaky_relu_f32(f2, f2, (int)(RG_NUM_FEAT * (size_t)H2 * W2), RG_LRELU_SLOPE);
|
|
|
|
/* Upsample stage 2: nearest 2x -> lrelu(conv_up2) */
|
|
int H4 = H * 4, W4 = W * 4;
|
|
iris_gpu_tensor_t up2 = iris_gpu_upsample_nearest_2x_f32(f2, RG_NUM_FEAT, H2, W2);
|
|
iris_gpu_tensor_free(f2);
|
|
if (!up2)
|
|
return NULL;
|
|
iris_gpu_tensor_t f3 = rg_conv_gpu(m, "conv_up2", up2, RG_NUM_FEAT, H4, W4, RG_NUM_FEAT);
|
|
iris_gpu_tensor_free(up2);
|
|
if (!f3)
|
|
return NULL;
|
|
const size_t plane4 = (size_t)H4 * W4;
|
|
iris_gpu_leaky_relu_f32(f3, f3, (int)(RG_NUM_FEAT * plane4), RG_LRELU_SLOPE);
|
|
|
|
/* conv_hr -> lrelu -> conv_last (64 -> 64 -> 3) */
|
|
iris_gpu_tensor_t hr = rg_conv_gpu(m, "conv_hr", f3, RG_NUM_FEAT, H4, W4, RG_NUM_FEAT);
|
|
iris_gpu_tensor_free(f3);
|
|
if (!hr)
|
|
return NULL;
|
|
iris_gpu_leaky_relu_f32(hr, hr, (int)(RG_NUM_FEAT * plane4), RG_LRELU_SLOPE);
|
|
|
|
iris_gpu_tensor_t out = rg_conv_gpu(m, "conv_last", hr, RG_NUM_FEAT, H4, W4, 3);
|
|
iris_gpu_tensor_free(hr);
|
|
if (!out)
|
|
return NULL;
|
|
|
|
float *outbuf = malloc((size_t)3 * plane4 * sizeof(float));
|
|
if (!outbuf) {
|
|
iris_gpu_tensor_free(out);
|
|
return NULL;
|
|
}
|
|
iris_gpu_tensor_read(out, outbuf);
|
|
iris_gpu_tensor_free(out);
|
|
|
|
if (m->tileable)
|
|
rg_make_tileable(outbuf, H4, W4);
|
|
|
|
iris_image *result = rg_chw_to_image(outbuf, W4, H4);
|
|
free(outbuf);
|
|
return result;
|
|
}
|
|
#endif /* USE_METAL || USE_VULKAN */
|
|
|
|
/* ========================================================================
|
|
* Public API
|
|
* ======================================================================== */
|
|
|
|
iris_upscale_t *iris_upscale_load(const char *path) {
|
|
safetensors_file_t *sf = safetensors_open(path);
|
|
if (!sf) {
|
|
fprintf(stderr, "RealESRGAN: failed to open '%s'\n", path);
|
|
return NULL;
|
|
}
|
|
/* Sanity check a couple of expected tensors. */
|
|
if (!safetensors_find(sf, "conv_first.weight") || !safetensors_find(sf, "conv_last.weight")) {
|
|
fprintf(stderr, "RealESRGAN: '%s' does not look like an RRDBNet model\n", path);
|
|
safetensors_close(sf);
|
|
return NULL;
|
|
}
|
|
iris_upscale_t *m = calloc(1, sizeof(iris_upscale_t));
|
|
if (!m) {
|
|
safetensors_close(sf);
|
|
return NULL;
|
|
}
|
|
m->sf = sf;
|
|
return m;
|
|
}
|
|
|
|
void iris_upscale_set_tileable(iris_upscale_t *m, int on) {
|
|
if (m)
|
|
m->tileable = on ? 1 : 0;
|
|
}
|
|
|
|
void iris_upscale_free(iris_upscale_t *m) {
|
|
if (!m)
|
|
return;
|
|
safetensors_close(m->sf);
|
|
free(m);
|
|
}
|
|
|
|
/* Helper: run a 3x3 conv by tensor prefix into `out`. Returns 0 on success. */
|
|
static int conv_by_name(iris_upscale_t *m, const char *prefix, const float *in, int Cin, int H, int W, int Cout, float *out) {
|
|
char name[256];
|
|
snprintf(name, sizeof(name), "%s.weight", prefix);
|
|
const float *w = rg_tensor(m, name);
|
|
snprintf(name, sizeof(name), "%s.bias", prefix);
|
|
const float *b = rg_tensor(m, name);
|
|
if (!w || !b)
|
|
return -1;
|
|
conv2d_3x3(in, Cin, H, W, w, b, Cout, out);
|
|
return 0;
|
|
}
|
|
|
|
iris_image *iris_upscale_run(iris_upscale_t *m, const iris_image *input) {
|
|
if (!m || !input || !input->data)
|
|
return NULL;
|
|
|
|
int H = input->height, W = input->width;
|
|
const size_t plane = (size_t)H * W;
|
|
|
|
#if defined(USE_METAL) || defined(USE_VULKAN)
|
|
/* GPU-resident path: convolution dominates RRDBNet, so offload it. Falls
|
|
* back to the CPU path below if the GPU forward fails. */
|
|
if (IRIS_GPU_AVAILABLE()) {
|
|
iris_image *r = upscale_gpu(m, input);
|
|
if (r)
|
|
return r;
|
|
fprintf(stderr, "RealESRGAN: GPU path failed, falling back to CPU\n");
|
|
}
|
|
#endif
|
|
|
|
/* ---- Input: RGB uint8 -> CHW float in [0,1] ---- */
|
|
float *x_in = rg_build_input(input);
|
|
if (!x_in)
|
|
return NULL;
|
|
|
|
/* ---- Work buffers ---- */
|
|
float *feat = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *feat_body = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *body_a = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *body_b = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *scratch_cat = malloc((size_t)RG_CAT_CH * plane * sizeof(float));
|
|
float *scratch_x5 = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *tmp_a = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
float *tmp_b = malloc((size_t)RG_NUM_FEAT * plane * sizeof(float));
|
|
|
|
if (!feat || !feat_body || !body_a || !body_b || !scratch_cat || !scratch_x5 || !tmp_a || !tmp_b) {
|
|
free(x_in);
|
|
free(feat);
|
|
free(feat_body);
|
|
free(body_a);
|
|
free(body_b);
|
|
free(scratch_cat);
|
|
free(scratch_x5);
|
|
free(tmp_a);
|
|
free(tmp_b);
|
|
fprintf(stderr, "RealESRGAN: out of memory for %dx%d input\n", W, H);
|
|
return NULL;
|
|
}
|
|
|
|
int ok = 1;
|
|
|
|
/* ---- conv_first: 3 -> 64 ---- */
|
|
if (ok && conv_by_name(m, "conv_first", x_in, 3, H, W, RG_NUM_FEAT, feat) != 0)
|
|
ok = 0;
|
|
free(x_in);
|
|
|
|
/* ---- body: 23 RRDB blocks, ping-pong between body_a / body_b ---- */
|
|
if (ok) {
|
|
const float *cur = feat; /* current block input */
|
|
float *dst = body_a; /* current block output */
|
|
for (int i = 0; i < RG_NUM_BLOCK; i++) {
|
|
if (rrdb(m, i, cur, H, W, scratch_cat, scratch_x5, tmp_a, tmp_b, dst) != 0) {
|
|
ok = 0;
|
|
break;
|
|
}
|
|
cur = dst;
|
|
dst = (dst == body_a) ? body_b : body_a;
|
|
}
|
|
/* conv_body: 64 -> 64, reading the last written body buffer (`cur`). */
|
|
if (ok && conv_by_name(m, "conv_body", cur, RG_NUM_FEAT, H, W, RG_NUM_FEAT, feat_body) != 0)
|
|
ok = 0;
|
|
}
|
|
|
|
/* ---- feat = feat + conv_body(body(feat)) ---- */
|
|
if (ok) {
|
|
size_t n = (size_t)RG_NUM_FEAT * plane;
|
|
for (size_t i = 0; i < n; i++)
|
|
feat[i] += feat_body[i];
|
|
}
|
|
|
|
free(body_a);
|
|
free(body_b);
|
|
free(feat_body);
|
|
|
|
if (!ok) {
|
|
free(feat);
|
|
free(scratch_cat);
|
|
free(scratch_x5);
|
|
free(tmp_a);
|
|
free(tmp_b);
|
|
return NULL;
|
|
}
|
|
|
|
/* ---- Upsample stage 1: nearest 2x then lrelu(conv_up1) ---- */
|
|
int H2 = H * 2, W2 = W * 2;
|
|
size_t plane2 = (size_t)H2 * W2;
|
|
float *up = malloc((size_t)RG_NUM_FEAT * plane2 * sizeof(float));
|
|
float *feat2 = malloc((size_t)RG_NUM_FEAT * plane2 * sizeof(float));
|
|
if (!up || !feat2) {
|
|
free(feat);
|
|
free(scratch_cat);
|
|
free(scratch_x5);
|
|
free(tmp_a);
|
|
free(tmp_b);
|
|
free(up);
|
|
free(feat2);
|
|
fprintf(stderr, "RealESRGAN: out of memory (upsample 1)\n");
|
|
return NULL;
|
|
}
|
|
upsample_nearest_2x(feat, RG_NUM_FEAT, H, W, up);
|
|
free(feat);
|
|
free(scratch_cat);
|
|
free(scratch_x5);
|
|
free(tmp_a);
|
|
free(tmp_b);
|
|
|
|
if (conv_by_name(m, "conv_up1", up, RG_NUM_FEAT, H2, W2, RG_NUM_FEAT, feat2) != 0) {
|
|
free(up);
|
|
free(feat2);
|
|
return NULL;
|
|
}
|
|
leaky_relu(feat2, (size_t)RG_NUM_FEAT * plane2);
|
|
free(up);
|
|
|
|
/* ---- Upsample stage 2: nearest 2x then lrelu(conv_up2) ---- */
|
|
int H4 = H * 4, W4 = W * 4;
|
|
size_t plane4 = (size_t)H4 * W4;
|
|
float *up2 = malloc((size_t)RG_NUM_FEAT * plane4 * sizeof(float));
|
|
float *feat3 = malloc((size_t)RG_NUM_FEAT * plane4 * sizeof(float));
|
|
if (!up2 || !feat3) {
|
|
free(feat2);
|
|
free(up2);
|
|
free(feat3);
|
|
fprintf(stderr, "RealESRGAN: out of memory (upsample 2)\n");
|
|
return NULL;
|
|
}
|
|
upsample_nearest_2x(feat2, RG_NUM_FEAT, H2, W2, up2);
|
|
free(feat2);
|
|
|
|
if (conv_by_name(m, "conv_up2", up2, RG_NUM_FEAT, H4, W4, RG_NUM_FEAT, feat3) != 0) {
|
|
free(up2);
|
|
free(feat3);
|
|
return NULL;
|
|
}
|
|
leaky_relu(feat3, (size_t)RG_NUM_FEAT * plane4);
|
|
free(up2);
|
|
|
|
/* ---- conv_hr -> lrelu -> conv_last: 64 -> 64 -> 3 ---- */
|
|
float *hr = malloc((size_t)RG_NUM_FEAT * plane4 * sizeof(float));
|
|
float *out = malloc((size_t)3 * plane4 * sizeof(float));
|
|
if (!hr || !out) {
|
|
free(feat3);
|
|
free(hr);
|
|
free(out);
|
|
fprintf(stderr, "RealESRGAN: out of memory (head)\n");
|
|
return NULL;
|
|
}
|
|
if (conv_by_name(m, "conv_hr", feat3, RG_NUM_FEAT, H4, W4, RG_NUM_FEAT, hr) != 0) {
|
|
free(feat3);
|
|
free(hr);
|
|
free(out);
|
|
return NULL;
|
|
}
|
|
free(feat3);
|
|
leaky_relu(hr, (size_t)RG_NUM_FEAT * plane4);
|
|
|
|
if (conv_by_name(m, "conv_last", hr, RG_NUM_FEAT, H4, W4, 3, out) != 0) {
|
|
free(hr);
|
|
free(out);
|
|
return NULL;
|
|
}
|
|
free(hr);
|
|
|
|
if (m->tileable)
|
|
rg_make_tileable(out, H4, W4);
|
|
|
|
/* ---- CHW float [0,1] -> RGB uint8 image ---- */
|
|
iris_image *result = rg_chw_to_image(out, W4, H4);
|
|
free(out);
|
|
return result;
|
|
}
|