From 5a283b91f9fe044a525e551ec4cd5da1cbdfd9b9 Mon Sep 17 00:00:00 2001 From: luboslenco Date: Wed, 17 Jun 2026 23:27:42 +0200 Subject: [PATCH] tools: update iris.c --- base/tools/iris.c/iris_depth.c | 29 +++- base/tools/iris.c/iris_metal.h | 6 + base/tools/iris.c/iris_metal.m | 137 +++++++++++++++++-- base/tools/iris.c/iris_shaders.metal | 28 ++++ base/tools/iris.c/iris_upscale.c | 194 +++++++++++++++++++++++++-- base/tools/iris.c/iris_upscale.h | 7 + base/tools/iris.c/main.c | 24 ++-- 7 files changed, 387 insertions(+), 38 deletions(-) diff --git a/base/tools/iris.c/iris_depth.c b/base/tools/iris.c/iris_depth.c index a20a8c48..c7f07476 100644 --- a/base/tools/iris.c/iris_depth.c +++ b/base/tools/iris.c/iris_depth.c @@ -30,8 +30,9 @@ * a single-channel grayscale image (brighter = nearer). * * All weights are F32 in the memory-mapped safetensors file (zero-copy). The - * transformer linear/attention matrix multiplies are dispatched to the Vulkan - * GEMM backend when available; convolutions (small spatially) run on the CPU. + * transformer linear/attention matrix multiplies are dispatched to the Metal + * or Vulkan GEMM backend when available; convolutions (small spatially) run on + * the CPU. * * Note: the official image processor resizes with bicubic resampling; here we * use bilinear, which yields a visually equivalent relative-depth map. @@ -45,7 +46,9 @@ #include #include -#ifdef USE_VULKAN +#if defined(USE_METAL) +#include "iris_metal.h" +#elif defined(USE_VULKAN) #include "iris_vulkan.h" #endif @@ -119,7 +122,15 @@ static const float *da_tensorf(iris_depth_t *m, const char *fmt, ...) { /* C[M,N] = scale * A[M,K] @ B[N,K]^T (row-wise dot products). * cache_b: weight matrix is immutable -> cache it on the GPU by pointer. */ static void gemm_nt(int M, int N, int K, const float *A, int lda, const float *B, int ldb, float *C, int ldc, float scale, int cache_b) { -#ifdef USE_VULKAN +#if defined(USE_METAL) + if (iris_metal_available()) { + if (cache_b) + iris_metal_sgemm_cached(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc); + else + iris_metal_sgemm(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc); + return; + } +#elif defined(USE_VULKAN) if (iris_vulkan_available()) { if (cache_b) iris_vulkan_sgemm_cached(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc); @@ -127,9 +138,8 @@ static void gemm_nt(int M, int N, int K, const float *A, int lda, const float *B iris_vulkan_sgemm(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc); return; } -#else - (void)cache_b; #endif + (void)cache_b; for (int i = 0; i < M; i++) { const float *a = A + (size_t)i * lda; float *c = C + (size_t)i * ldc; @@ -145,7 +155,12 @@ static void gemm_nt(int M, int N, int K, const float *A, int lda, const float *B /* C[M,N] = A[M,K] @ B[K,N] (B not transposed). */ static void gemm_nn(int M, int N, int K, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { -#ifdef USE_VULKAN +#if defined(USE_METAL) + if (iris_metal_available()) { + iris_metal_sgemm(0, 0, M, N, K, 1.0f, A, lda, B, ldb, 0.0f, C, ldc); + return; + } +#elif defined(USE_VULKAN) if (iris_vulkan_available()) { iris_vulkan_sgemm(0, 0, M, N, K, 1.0f, A, lda, B, ldb, 0.0f, C, ldc); return; diff --git a/base/tools/iris.c/iris_metal.h b/base/tools/iris.c/iris_metal.h index d44eb4aa..f95b068e 100644 --- a/base/tools/iris.c/iris_metal.h +++ b/base/tools/iris.c/iris_metal.h @@ -461,6 +461,12 @@ void iris_gpu_swish_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, int n); /* Element-wise add on f32 GPU tensors: out = a + b, in-place safe */ void iris_gpu_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tensor_t b, int n); +/* In-place-safe LeakyReLU: out[i] = x[i] >= 0 ? x[i] : slope*x[i]. */ +void iris_gpu_leaky_relu_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, int n, float slope); + +/* Scaled residual add: out[i] = scale*a[i] + b[i], in-place safe. */ +void iris_gpu_scale_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tensor_t b, float scale, int n); + /* Nearest neighbor 2x upsample on f32 GPU tensor: [1, C, H, W] -> [1, C, 2H, 2W] (batch=1) */ iris_gpu_tensor_t iris_gpu_upsample_nearest_2x_f32(iris_gpu_tensor_t x, int channels, int H, int W); diff --git a/base/tools/iris.c/iris_metal.m b/base/tools/iris.c/iris_metal.m index a4ec0a5c..f94a1216 100644 --- a/base/tools/iris.c/iris_metal.m +++ b/base/tools/iris.c/iris_metal.m @@ -139,6 +139,9 @@ typedef struct { id buffer; float *cpu_ptr; size_t size; + int rows; /* M */ + int cols; /* N */ + int ldc; /* row stride of cpu_ptr, in elements */ } pending_output_t; static id g_batch_cmd = nil; @@ -146,6 +149,23 @@ static int g_in_batch = 0; static pending_output_t g_pending_outputs[MAX_BATCH_OUTPUTS]; static int g_pending_count = 0; +/* Copy an MPS result buffer back to a (possibly row-strided) CPU matrix. + * MPS writes only the N valid columns of each row (rowBytes = ldc), leaving the + * gap columns of a strided output (ldc > N) untouched. A single contiguous copy + * is therefore only correct when the matrix is packed (ldc == N); otherwise it + * would clobber the caller's gap columns and run past the end of C. When the + * output is strided we copy row by row, mirroring the Vulkan sgemm semantics + * (e.g. depth attention writes each head's head_dim-wide slice into a hidden- + * wide row). */ +static void sgemm_copy_result(float *C, const float *src, int M, int N, int ldc) { + if (ldc == N) { + memcpy(C, src, (size_t)M * ldc * sizeof(float)); + return; + } + for (int r = 0; r < M; r++) + memcpy(C + (size_t)r * ldc, src + (size_t)r * ldc, (size_t)N * sizeof(float)); +} + /* Input buffer cache - during batch mode, cache input buffers to avoid * redundant copies when the same tensor is used as input to multiple ops */ #define MAX_BATCH_INPUTS 64 @@ -226,6 +246,8 @@ static id g_group_norm_f32_pipeline; static id g_swish_f32_pipeline; static id g_add_f32_pipeline; static id g_upsample_nearest_2x_f32_pipeline; +static id g_leaky_relu_f32_pipeline; +static id g_scale_add_f32_pipeline; static int g_shaders_initialized; /* ======================================================================== @@ -621,9 +643,10 @@ void iris_metal_end_batch(void) { [g_batch_cmd commit]; [g_batch_cmd waitUntilCompleted]; - /* Copy all pending outputs back to CPU */ + /* Copy all pending outputs back to CPU (respecting row stride) */ for (int i = 0; i < g_pending_count; i++) { - memcpy(g_pending_outputs[i].cpu_ptr, [g_pending_outputs[i].buffer contents], g_pending_outputs[i].size); + sgemm_copy_result(g_pending_outputs[i].cpu_ptr, [g_pending_outputs[i].buffer contents], g_pending_outputs[i].rows, g_pending_outputs[i].cols, + g_pending_outputs[i].ldc); /* Don't nil the buffer - it's from the pool */ } @@ -757,6 +780,9 @@ static void iris_metal_sgemm_impl(int transpose_a, int transpose_b, int M, int N g_pending_outputs[g_pending_count].buffer = bufferC; g_pending_outputs[g_pending_count].cpu_ptr = C; g_pending_outputs[g_pending_count].size = sizeC; + g_pending_outputs[g_pending_count].rows = M; + g_pending_outputs[g_pending_count].cols = N; + g_pending_outputs[g_pending_count].ldc = ldc; g_pending_count++; /* Don't release bufferA if it came from batch input cache */ if (!bufferA_from_cache) { @@ -767,7 +793,7 @@ static void iris_metal_sgemm_impl(int transpose_a, int transpose_b, int M, int N /* Too many pending outputs - fall back to immediate sync */ [cmdBuffer commit]; [cmdBuffer waitUntilCompleted]; - memcpy(C, [bufferC contents], sizeC); + sgemm_copy_result(C, [bufferC contents], M, N, ldc); if (!bufferA_from_cache) { pool_release_buffer(bufferA); } @@ -778,7 +804,7 @@ static void iris_metal_sgemm_impl(int transpose_a, int transpose_b, int M, int N /* Not in batch mode: execute immediately */ [cmdBuffer commit]; [cmdBuffer waitUntilCompleted]; - memcpy(C, [bufferC contents], sizeC); + sgemm_copy_result(C, [bufferC contents], M, N, ldc); /* Release pooled buffers */ pool_release_buffer(bufferA); @@ -1345,16 +1371,28 @@ static conv2d_graph_cache_t *get_conv2d_graph_cache(int batch, int in_ch, int ou MPSGraphTensor *weight = [graph placeholderWithShape:weightShape dataType:MPSDataTypeFloat32 name:nil]; MPSGraphTensor *bias = [graph placeholderWithShape:biasShape dataType:MPSDataTypeFloat32 name:nil]; - /* Seamless tiling: wrap the input around a torus (circular padding) as a - * separate pad node, then run a VALID conv. MPSGraph's convolution only - * offers zero padding, so circular wrapping must precede it. */ + /* Seamless tiling: wrap the input around a torus (circular padding), then + * run a VALID conv. MPSGraph's convolution only offers zero padding, so the + * circular wrap must precede it. MPSGraphPaddingModePeriodic exists in the + * enum but is not implemented by the underlying MPS pad kernel (it asserts + * "Unsupported paddingMode"), so we build the wrap by hand from slices: take + * the trailing `padding` elements to the front and the leading ones to the + * back, concatenated around the original tensor. Layout is NCHW, so H is + * dimension 2 and W is dimension 3. */ MPSGraphTensor *convSource = input; NSUInteger convPad = (NSUInteger)padding; if (circular) { - NSArray *lowPad = @[ @0, @0, @(padding), @(padding) ]; - NSArray *highPad = @[ @0, @0, @(padding), @(padding) ]; - convSource = [graph padTensor:input withPaddingMode:MPSGraphPaddingModePeriodic leftPadding:lowPad rightPadding:highPad constantValue:0.0 name:nil]; - convPad = 0; /* padding already applied via the periodic pad node */ + /* Wrap width (dimension 3). */ + MPSGraphTensor *wLeft = [graph sliceTensor:input dimension:3 start:(W - padding) length:padding name:nil]; + MPSGraphTensor *wRight = [graph sliceTensor:input dimension:3 start:0 length:padding name:nil]; + MPSGraphTensor *wWrap = [graph concatTensors:@[ wLeft, input, wRight ] dimension:3 name:nil]; + + /* Wrap height (dimension 2) on the width-wrapped tensor. */ + MPSGraphTensor *hTop = [graph sliceTensor:wWrap dimension:2 start:(H - padding) length:padding name:nil]; + MPSGraphTensor *hBottom = [graph sliceTensor:wWrap dimension:2 start:0 length:padding name:nil]; + convSource = [graph concatTensors:@[ hTop, wWrap, hBottom ] dimension:2 name:nil]; + + convPad = 0; /* padding already applied via the circular wrap */ } MPSGraphConvolution2DOpDescriptor *desc = [MPSGraphConvolution2DOpDescriptor descriptorWithStrideInX:(NSUInteger)stride @@ -3476,6 +3514,14 @@ int iris_metal_init_shaders(void) { if (func) { g_upsample_nearest_2x_f32_pipeline = [g_device newComputePipelineStateWithFunction:func error:&error]; } + func = [g_shader_library newFunctionWithName:@"leaky_relu_f32"]; + if (func) { + g_leaky_relu_f32_pipeline = [g_device newComputePipelineStateWithFunction:func error:&error]; + } + func = [g_shader_library newFunctionWithName:@"scale_add_f32"]; + if (func) { + g_scale_add_f32_pipeline = [g_device newComputePipelineStateWithFunction:func error:&error]; + } g_shaders_initialized = 1; if (iris_verbose) @@ -6272,6 +6318,75 @@ void iris_gpu_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tenso } } +/* In-place-safe LeakyReLU on f32 GPU tensor: out = x >= 0 ? x : slope*x */ +void iris_gpu_leaky_relu_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t x, int n, float slope) { + if (!g_shaders_initialized || !g_leaky_relu_f32_pipeline) + return; + if (!out || !x || n <= 0) + return; + + @autoreleasepool { + id cmdBuffer = get_tensor_cmd(); + id encoder = [cmdBuffer computeCommandEncoder]; + + [encoder setComputePipelineState:g_leaky_relu_f32_pipeline]; + [encoder setBuffer:x->buffer offset:0 atIndex:0]; + [encoder setBuffer:out->buffer offset:0 atIndex:1]; + [encoder setBytes:&n length:sizeof(int) atIndex:2]; + [encoder setBytes:&slope length:sizeof(float) atIndex:3]; + + NSUInteger threads = 256; + NSUInteger groups = ((NSUInteger)n + threads - 1) / threads; + [encoder dispatchThreadgroups:MTLSizeMake(groups, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [encoder endEncoding]; + + out->has_pending_work = 1; + x->has_pending_work = 1; + if (!g_tensor_batch_mode) { + [cmdBuffer commit]; + [cmdBuffer waitUntilCompleted]; + out->has_pending_work = 0; + x->has_pending_work = 0; + } + } +} + +/* Scaled residual add on f32 GPU tensors: out = scale*a + b */ +void iris_gpu_scale_add_f32(iris_gpu_tensor_t out, iris_gpu_tensor_t a, iris_gpu_tensor_t b, float scale, int n) { + if (!g_shaders_initialized || !g_scale_add_f32_pipeline) + return; + if (!out || !a || !b || n <= 0) + return; + + @autoreleasepool { + id cmdBuffer = get_tensor_cmd(); + id encoder = [cmdBuffer computeCommandEncoder]; + + [encoder setComputePipelineState:g_scale_add_f32_pipeline]; + [encoder setBuffer:a->buffer offset:0 atIndex:0]; + [encoder setBuffer:b->buffer offset:0 atIndex:1]; + [encoder setBuffer:out->buffer offset:0 atIndex:2]; + [encoder setBytes:&n length:sizeof(int) atIndex:3]; + [encoder setBytes:&scale length:sizeof(float) atIndex:4]; + + NSUInteger threads = 256; + NSUInteger groups = ((NSUInteger)n + threads - 1) / threads; + [encoder dispatchThreadgroups:MTLSizeMake(groups, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + [encoder endEncoding]; + + out->has_pending_work = 1; + a->has_pending_work = 1; + b->has_pending_work = 1; + if (!g_tensor_batch_mode) { + [cmdBuffer commit]; + [cmdBuffer waitUntilCompleted]; + out->has_pending_work = 0; + a->has_pending_work = 0; + b->has_pending_work = 0; + } + } +} + /* Nearest neighbor 2x upsample on f32 GPU tensor */ iris_gpu_tensor_t iris_gpu_upsample_nearest_2x_f32(iris_gpu_tensor_t x, int channels, int H, int W) { if (!g_shaders_initialized || !g_upsample_nearest_2x_f32_pipeline) diff --git a/base/tools/iris.c/iris_shaders.metal b/base/tools/iris.c/iris_shaders.metal index 4583a6ce..68d38174 100644 --- a/base/tools/iris.c/iris_shaders.metal +++ b/base/tools/iris.c/iris_shaders.metal @@ -2134,3 +2134,31 @@ kernel void upsample_nearest_2x_f32( out[c * out_spatial + oy * out_w + ox] = x[c * in_h * in_w + iy * in_w + ix]; } + +/* LeakyReLU f32: out = x >= 0 ? x : slope*x, in-place safe (out can alias x) */ +kernel void leaky_relu_f32( + device const float *x [[buffer(0)]], + device float *out [[buffer(1)]], + constant int &n [[buffer(2)]], + constant float &slope [[buffer(3)]], + uint gid [[thread_position_in_grid]] +) { + if (gid < uint(n)) { + float v = x[gid]; + out[gid] = v >= 0.0f ? v : slope * v; + } +} + +/* Scaled residual add f32: out = scale*a + b, in-place safe (out can alias a or b) */ +kernel void scale_add_f32( + device const float *a [[buffer(0)]], + device const float *b [[buffer(1)]], + device float *out [[buffer(2)]], + constant int &n [[buffer(3)]], + constant float &scale [[buffer(4)]], + uint gid [[thread_position_in_grid]] +) { + if (gid < uint(n)) { + out[gid] = scale * a[gid] + b[gid]; + } +} diff --git a/base/tools/iris.c/iris_upscale.c b/base/tools/iris.c/iris_upscale.c index b9a96e07..e6ac13a3 100644 --- a/base/tools/iris.c/iris_upscale.c +++ b/base/tools/iris.c/iris_upscale.c @@ -22,14 +22,23 @@ #include "iris_upscale.h" #include "iris_safetensors.h" +#include #include #include #include -#ifdef USE_VULKAN +#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 @@ -44,6 +53,7 @@ struct iris_upscale { safetensors_file_t *sf; + int tileable; /* make the upscaled image wrap seamlessly */ }; /* ======================================================================== @@ -254,13 +264,170 @@ static iris_image *rg_chw_to_image(const float *out, int W4, int H4) { } /* ======================================================================== - * Vulkan GPU-resident forward path + * 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. * ======================================================================== */ -#ifdef USE_VULKAN +#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. */ @@ -346,7 +513,7 @@ static iris_gpu_tensor_t rrdb_gpu(iris_upscale_t *m, int idx, iris_gpu_tensor_t return out; } -static iris_image *upscale_vulkan(iris_upscale_t *m, const iris_image *input) { +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; @@ -433,11 +600,14 @@ static iris_image *upscale_vulkan(iris_upscale_t *m, const iris_image *input) { 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_VULKAN */ +#endif /* USE_METAL || USE_VULKAN */ /* ======================================================================== * Public API @@ -464,6 +634,11 @@ iris_upscale_t *iris_upscale_load(const char *path) { 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; @@ -491,11 +666,11 @@ iris_image *iris_upscale_run(iris_upscale_t *m, const iris_image *input) { int H = input->height, W = input->width; const size_t plane = (size_t)H * W; -#ifdef USE_VULKAN +#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_vulkan_available()) { - iris_image *r = upscale_vulkan(m, input); + 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"); @@ -655,6 +830,9 @@ iris_image *iris_upscale_run(iris_upscale_t *m, const iris_image *input) { } 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); diff --git a/base/tools/iris.c/iris_upscale.h b/base/tools/iris.c/iris_upscale.h index a1088e36..da7a4f13 100644 --- a/base/tools/iris.c/iris_upscale.h +++ b/base/tools/iris.c/iris_upscale.h @@ -24,6 +24,13 @@ typedef struct iris_upscale iris_upscale_t; */ iris_upscale_t *iris_upscale_load(const char *path); +/* + * Enable "tileable" mode: subtract the smoothest periodic field that makes the + * upscaled image wrap seamlessly, so a tileable input stays tileable after 4x + * upscaling. Off by default. + */ +void iris_upscale_set_tileable(iris_upscale_t *model, int on); + /* * Free the model and its resources. */ diff --git a/base/tools/iris.c/main.c b/base/tools/iris.c/main.c index 6ec3a02a..818ea015 100644 --- a/base/tools/iris.c/main.c +++ b/base/tools/iris.c/main.c @@ -228,8 +228,7 @@ static void print_usage(const char *prog) { fprintf(stderr, " --linear Use linear timestep schedule\n"); fprintf(stderr, " --power Use power curve timestep schedule (default alpha: 2.0)\n"); fprintf(stderr, " --power-alpha N Set power schedule exponent (default: 2.0)\n"); - fprintf(stderr, " --sigmoid Use Flux shifted sigmoid schedule\n"); - fprintf(stderr, " --circular Seamless/tileable output (circular conv padding)\n\n"); + fprintf(stderr, " --sigmoid Use Flux shifted sigmoid schedule\n\n"); fprintf(stderr, "Model options:\n"); fprintf(stderr, " --base Force base model mode (undistilled, CFG enabled)\n\n"); fprintf(stderr, "Reference images (img2img / multi-reference):\n"); @@ -240,9 +239,12 @@ static void print_usage(const char *prog) { fprintf(stderr, " (uses /RealESRGAN_x4plus.safetensors; no prompt needed)\n\n"); fprintf(stderr, "Depth estimation:\n"); fprintf(stderr, " --depth Estimate depth from input (-i) via Depth Anything 3, write to -o\n"); - fprintf(stderr, " (uses /da3-mono-large.safetensors; no prompt needed)\n"); - fprintf(stderr, " --tileable With --depth: flatten the perspective tilt so depth from a\n"); - fprintf(stderr, " seamless top-down texture also tiles seamlessly\n\n"); + fprintf(stderr, " (uses /da3-mono-large.safetensors; no prompt needed)\n\n"); + fprintf(stderr, "Seamless tiling:\n"); + fprintf(stderr, " --tileable Make the output seamlessly tileable.\n"); + fprintf(stderr, " Generation: circular conv padding.\n"); + fprintf(stderr, " With --upscale: keep a tileable input seamless after 4x.\n"); + fprintf(stderr, " With --depth: also flatten the perspective tilt.\n\n"); fprintf(stderr, "Output options:\n"); fprintf(stderr, " -q, --quiet Silent mode, no output\n"); fprintf(stderr, " -v, --verbose Detailed output\n\n"); @@ -291,7 +293,6 @@ int main(int argc, char *argv[]) { {"power", no_argument, 0, 256}, {"power-alpha", required_argument, 0, 257}, {"sigmoid", no_argument, 0, 260}, - {"circular", no_argument, 0, 259}, {"vae-tiling", no_argument, 0, 261}, {"upscale", no_argument, 0, 262}, {"depth", no_argument, 0, 263}, @@ -322,7 +323,7 @@ int main(int argc, char *argv[]) { int no_license_info = 0; int upscale_mode = 0; int depth_mode = 0; - int depth_tileable = 0; + int tileable = 0; int opt; while ((opt = getopt_long(argc, argv, "d:p:o:W:H:s:g:S:i:t:qvhVmMD", long_options, NULL)) != -1) { @@ -397,9 +398,6 @@ int main(int argc, char *argv[]) { case 260: params.schedule = IRIS_SCHEDULE_SIGMOID; break; - case 259: - params.circular = 1; - break; case 261: iris_vae_tiling = 1; break; @@ -410,7 +408,8 @@ int main(int argc, char *argv[]) { depth_mode = 1; break; case 264: - depth_tileable = 1; + tileable = 1; + params.circular = 1; break; case 258: no_license_info = 1; @@ -477,6 +476,7 @@ int main(int argc, char *argv[]) { fprintf(stderr, "\nError: Failed to load upscale model: %s\n", model_path); return 1; } + iris_upscale_set_tileable(up, tileable); LOG_NORMAL(" done (%.1fs)\n", timer_end()); iris_image *src = iris_image_load(input_paths[0]); @@ -539,7 +539,7 @@ int main(int argc, char *argv[]) { return 1; } LOG_NORMAL(" done (%.1fs)\n", timer_end()); - iris_depth_set_tileable(dm, depth_tileable); + iris_depth_set_tileable(dm, tileable); iris_image *src = iris_image_load(input_paths[0]); if (!src) {