Untiled LTX VAE decode when it fits, drop per-tile gc

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0.11
+Version: 0.1.0.12
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NEWS.md

@@ -1,3 +1,13 @@
+# diffuseR 0.1.0.12 (development)
+
+* LTX video decode runs untiled when the estimated activation cost
+  (~1 GB + ~360 B per output pixel-frame) fits the card - tiling
+  bounds VRAM, and at decode time the transformer is phase-offloaded,
+  so 768x512x49 decodes in one full-latent forward. The per-tile
+  explicit `gc()` is gone (storm-free with the allocator gates live).
+  In-render decode phase 12.3 s -> 1.7-1.8 s; warm render 54 -> ~44 s.
+  `options(diffuseR.vae_untiled = TRUE/FALSE)` forces either path.
+
 # diffuseR 0.1.0.11 (development)
 
 * `ltx23_tune_gc()` now takes effect: it used to set the allocator gate

R/vae_ltx23.R

@@ -341,10 +341,20 @@ ltx23_video_vae <- torch::nn_module(
     tile_lat_w <- self$tile_sample_min_width %/% 32L
     tile_lat_f <- self$tile_sample_min_num_frames %/% 8L
 
-    if (self$use_framewise_decoding && num_frames > tile_lat_f) {
+    wants_temporal <- self$use_framewise_decoding && num_frames > tile_lat_f
+    wants_spatial <- self$use_tiling &&
+    (width > tile_lat_w || height > tile_lat_h)
+    if ((wants_temporal || wants_spatial) && .ltx23_untiled_fits(z)) {
+        # Tiling exists to bound VRAM, not for correctness: when the
+        # whole sample fits, one full-latent forward skips the tile
+        # overlap recompute and seam blending (2.7 s vs 4.4 s at
+        # 768x512x49)
+        return(.ltx23_decode_tile(self, z, causal))
+    }
+    if (wants_temporal) {
         return(.ltx23_temporal_tiled_decode(self, z, causal = causal))
     }
-    if (self$use_tiling && (width > tile_lat_w || height > tile_lat_h)) {
+    if (wants_spatial) {
         return(.ltx23_tiled_decode(self, z, causal = causal))
     }
     .ltx23_decode_tile(self, z, causal)
@@ -398,6 +408,48 @@ ltx23_video_vae <- torch::nn_module(
     b
 }
 
+# Tiling exists to bound decoder activation VRAM. When the whole
+# sample's untiled decode fits in what is left of the card, one
+# full-latent forward beats the tiles. Activation cost is linear in
+# output pixel-frames: ~1 GB + ~350 B per pixel-frame measured on the
+# 2.3 decoder across 2.5M-19.3M pixel-frames; 360 B/pxf plus 15%
+# headroom below rounds up. options(diffuseR.vae_untiled = TRUE/FALSE)
+# forces a path; the "auto" default applies the fit test on CUDA and
+# keeps tiled behavior elsewhere.
+.ltx23_untiled_env <- new.env(parent = emptyenv())
+
+.ltx23_untiled_fits <- function(z) {
+    pref <- getOption("diffuseR.vae_untiled", "auto")
+    if (isTRUE(pref)) {
+        return(TRUE)
+    }
+    if (!identical(pref, "auto")) {
+        return(FALSE)
+    }
+    if (z$device$type != "cuda") {
+        return(FALSE)
+    }
+    total <- .ltx23_untiled_env$total_bytes
+    if (is.null(total)) {
+        gb <- tryCatch(.detect_vram(use_free = FALSE), error = function(e) NULL)
+        if (!isTRUE(gb > 0)) {
+            return(FALSE)
+        }
+        total <- gb * 1e9
+        .ltx23_untiled_env$total_bytes <- total
+    }
+    alloc <- tryCatch(
+                      as.numeric(torch::cuda_memory_stats()$allocated_bytes$all$current),
+                      error = function(e) NULL
+    )
+    if (is.null(alloc)) {
+        return(FALSE)
+    }
+    pxf <- ((z$shape[3] - 1) * 8 + 1) * z$shape[4] * 32 * z$shape[5] * 32
+    est <- 1e9 + 360 * pxf
+    est < 0.85 * (total - alloc)
+}
+
 # Spatially tiled decode: overlapping latent tiles, decoded separately,
 # crossfaded at the seams (diffusers AutoencoderKLLTX2Video.tiled_decode)
 .ltx23_tiled_decode <- function(vae, z, causal = NULL) {
@@ -425,12 +477,11 @@ ltx23_video_vae <- torch::nn_module(
                                        z$narrow(4L, i + 1L, h_len)$narrow(5L, j + 1L, w_len),
                                        causal
             )
+            # Dead tile handles are left to the allocator callback:
+            # with the gc gates live (ltx23_tune_gc) that no longer
+            # storms, and an explicit per-tile gc() here measured
+            # 1.9 s per decode
             row[[length(row) + 1L]] <- tile
-            if (!isTRUE(getOption("diffuseR.jit_vae", FALSE))) {
-                # Eager tiles leave GBs of dead handles; without this
-                # the allocator callback storms instead
-                gc(verbose = FALSE)
-            }
         }
         rows[[length(rows) + 1L]] <- row
     }
@@ -487,9 +538,6 @@ ltx23_video_vae <- torch::nn_module(
             decoded <- decoded$narrow(3L, 1L, decoded$shape[3] - 1L)
         }
         row[[length(row) + 1L]] <- decoded
-        if (!isTRUE(getOption("diffuseR.jit_vae", FALSE))) {
-            gc(verbose = FALSE)
-        }
     }
 
     result_row <- list()

inst/tinytest/test_vae_ltx23.R

@@ -211,3 +211,29 @@ torch::with_no_grad({
   b <- vae_t$decode(z_small)
 })
 expect_true(as.numeric((a - b)$abs()$max()) == 0)
+
+# --- Untiled-when-fits dispatch ----------------------------------------------------
+
+# Force-untiled: with tiling enabled and tiles smaller than the input,
+# the option must route to the single full-latent forward (identical
+# to the untiled reference, no seam blending)
+vae_t$enable_tiling(spatial = TRUE, temporal = TRUE)
+vae_t$tile_sample_min_height <- 64L
+vae_t$tile_sample_min_width <- 64L
+vae_t$tile_sample_stride_height <- 32L
+vae_t$tile_sample_stride_width <- 32L
+vae_t$tile_sample_min_num_frames <- 16L
+vae_t$tile_sample_stride_num_frames <- 8L
+options(diffuseR.vae_untiled = TRUE)
+torch::with_no_grad(forced <- vae_t$decode(z_big))
+expect_true(max_abs_diff(forced, ref_full) < 1e-6)
+
+# Force-tiled: FALSE must reproduce the tiled path exactly
+options(diffuseR.vae_untiled = FALSE)
+torch::with_no_grad(forced_tiled <- vae_t$decode(z_big))
+expect_true(max_abs_diff(forced_tiled, tiled_t) < 1e-6)
+
+# "auto" on CPU keeps the tiled path (no VRAM to test against)
+options(diffuseR.vae_untiled = NULL)
+torch::with_no_grad(auto_cpu <- vae_t$decode(z_big))
+expect_true(max_abs_diff(auto_cpu, tiled_t) < 1e-6)