Latent-space chaining seams for chunked video continuation

DESCRIPTION

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

NAMESPACE

@@ -153,6 +153,7 @@ export(ltx23_set_attn_chunk)
 export(ltx23_snake_beta)
 export(ltx23_split_keys)
 export(ltx23_stage2_distilled_sigmas)
+export(ltx23_tail_latents)
 export(ltx23_text_connectors)
 export(ltx23_tone_map_latents)
 export(ltx23_transformer)

NEWS.md

@@ -1,3 +1,16 @@
+# diffuseR 0.1.0.14 (development)
+
+* Latent-space chaining seams for chunked video continuation:
+  `txt2vid_ltx2()` gains `condition_latents=` (an already-encoded
+  prefix, bypassing the VAE), `resident=` (keep named components on
+  the compute device across back-to-back calls; onload is now
+  idempotent), and `trim_frames=` (deliver head-free pixels while the
+  returned latents keep the full sequence). Results carry
+  `latent_shape`, and the new `ltx23_tail_latents()` slices a
+  result's trailing latent frames into `condition_latents` form — so
+  a chunk chain can stay in latent space end to end: encode once,
+  denoise every chunk with the transformer resident, decode once.
+
 # diffuseR 0.1.0.13 (development)
 
 * `gemma3_quantize_nf4()` + `load_gemma3_nf4()` put the Gemma3 12B text

R/condition_ltx23.R

@@ -119,6 +119,53 @@ ltx23_encode_video_frames <- function(vae, frames) {
                             vae$latents_std$to(dtype = torch::torch_float32()))
 }
 
+#' Slice the trailing latent frames of a generation for chaining
+#'
+#' Cuts the last \code{k} latent frames out of a result's video
+#' latents, in the [1, 128, k, H', W'] layout that
+#' \code{txt2vid_ltx2(condition_latents = )} consumes, so one chunk
+#' can seed the next without leaving latent space: no decode, no
+#' re-encode, no video round-trip.
+#'
+#' Semantics caveat: a latent frame sliced from inside a sequence
+#' represents 8 pixel frames, while a fresh VAE encode of a k-frame
+#' tail represents 1 + 8(k - 1) pixel frames with its first latent in
+#' first-frame form. The frozen prefix the next generation sees is
+#' therefore not identical to the pixel-path prefix; compare both on
+#' real content before relying on latent-only joins.
+#'
+#' @param result A \code{\link{txt2vid_ltx2}} result list (uses its
+#'   \code{latents} and \code{latent_shape}), or the packed latents
+#'   tensor [1, S, 128] itself (then \code{latent_shape} is
+#'   required).
+#' @param k Integer. Trailing latent frames to keep (default 2 = the
+#'   standard 9-pixel-frame conditioning prefix).
+#' @param latent_shape Integer vector c(frames, height, width) of the
+#'   latent geometry; only needed when \code{result} is a raw tensor.
+#'
+#' @return Normalized latents [1, 128, k, H', W'] (float32), ready
+#'   for \code{txt2vid_ltx2(condition_latents = )}.
+#'
+#' @export
+ltx23_tail_latents <- function(result, k = 2L, latent_shape = NULL) {
+    if (is.list(result)) {
+        latents <- result$latents
+        latent_shape <- latent_shape %||% result$latent_shape
+    } else {
+        latents <- result
+    }
+    if (is.null(latent_shape)) {
+        stop("latent_shape is required when result is a raw latents tensor")
+    }
+    lf <- latent_shape[1]
+    if (k < 1L || k > lf) {
+        stop("k must be between 1 and the latent frame count")
+    }
+    video <- ltx23_unpack_video_latents(latents, lf, latent_shape[2],
+                                        latent_shape[3])
+    video$narrow(3L, lf - k + 1L, k)
+}
+
 #' Build conditioned initial latents and the conditioning mask
 #'
 #' i2v (\code{cond_latents} has one latent frame): the encoded frame is

R/txt2vid_ltx23.R

@@ -432,6 +432,23 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda",
 #'   \code{condition_video} (8k + 1, default 9 = 2 latent frames).
 #' @param cond_noise_scale Numeric in [0, 1]. Optional partial noising
 #'   of the conditioned tokens (0 = keep them exactly).
+#' @param condition_latents Optional continuation source already in
+#'   latent space: normalized video latents [1, 128, k, height/32,
+#'   width/32] (e.g. \code{\link{ltx23_tail_latents}} on a previous
+#'   result), used directly as the frozen prefix with no VAE encode.
+#'   Mutually exclusive with \code{image} and \code{condition_video}.
+#' @param resident Character vector of pipeline component names
+#'   ("transformer", "vae", "audio_vae", "connectors", "vocoder") to
+#'   keep on the compute device after their phase instead of
+#'   offloading, for callers running several generations back to back
+#'   (chained chunks). Components already on the device are not
+#'   re-copied on later calls.
+#' @param trim_frames Integer. Drop this many leading pixel frames
+#'   from the decoded video (and the saved file), e.g. the
+#'   conditioning-head overlap of a continuation. The returned
+#'   \code{latents} keep the full sequence (tail slicing for chaining
+#'   needs it). Audio is muxed exactly as supplied, so drop any head
+#'   padding from the conditioning audio when trimming.
 #' @param audio Optional conditioning audio for audio-driven generation
 #'   (lip sync): a file path (decoded via \code{av}) or a matrix
 #'   [2, samples] in [-1, 1] at 16 kHz. The audio is encoded into
@@ -446,7 +463,10 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda",
 #'
 #' @return Invisibly, a list with \code{video} (array
 #'   [frames, height, width, 3] in [0, 1]), \code{audio} (matrix
-#'   [2, samples] in [-1, 1]), \code{sample_rate}, and the raw latents.
+#'   [2, samples] in [-1, 1]), \code{sample_rate}, the raw
+#'   \code{latents} and \code{audio_latents}, and \code{latent_shape}
+#'   (c(frames, height, width) of the latent geometry, for
+#'   \code{\link{ltx23_tail_latents}}).
 #'
 #' @export
 txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
@@ -461,7 +481,8 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
                          tone_map_compression = 0, phase_offload = TRUE,
                          image = NULL, condition_video = NULL,
                          conditioning_frames = 9L, cond_noise_scale = 0,
-                         audio = NULL, verbose = TRUE) {
+                         condition_latents = NULL, resident = character(),
+                         trim_frames = 0L, audio = NULL, verbose = TRUE) {
     level <- .verbosity(verbose)
     verbose <- level == "steps"
     if (level != "silent") {
@@ -473,10 +494,22 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
     if (guidance_scale != 1) {
         stop("Only guidance_scale = 1 is supported (distilled checkpoints).")
     }
-    if (!is.null(image) && !is.null(condition_video)) {
-        stop("Provide either image (i2v) or condition_video (continuation), not both.")
+    n_sources <- sum(!is.null(image), !is.null(condition_video),
+                     !is.null(condition_latents))
+    if (n_sources > 1L) {
+        stop("Provide only one of image (i2v), condition_video, or condition_latents.")
+    }
+    conditioned <- n_sources > 0L
+    known_components <- c("transformer", "vae", "audio_vae", "connectors",
+                          "vocoder")
+    if (length(resident) && !all(resident %in% known_components)) {
+        stop("resident must name pipeline components: ",
+             paste(known_components, collapse = ", "))
+    }
+    trim_frames <- as.integer(trim_frames)
+    if (trim_frames < 0L || trim_frames >= num_frames) {
+        stop("trim_frames must be in [0, num_frames)")
     }
-    conditioned <- !is.null(image) || !is.null(condition_video)
     if (conditioned && two_stage) {
         stop("Prefix conditioning with two_stage = TRUE is not supported yet.")
     }
@@ -531,13 +564,21 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
     # the pageable path.
     phase_offload <- phase_offload && device != "cpu"
     staging <- pipeline$staging %||% list()
+    target_type <- torch::torch_device(device)$type
     onload <- function(what) {
         if (is.character(what)) {
             module <- pipeline[[what]]
         } else {
             module <- what
         }
         if (phase_offload) {
+            # Idempotent: a resident component is already in place on
+            # the second and later calls of a chained run
+            cur <- tryCatch(module$parameters[[1]]$device$type,
+                            error = function(e) NULL)
+            if (identical(cur, target_type)) {
+                return(module)
+            }
             if (is.character(what)) {
                 st <- staging[[what]]
             } else {
@@ -557,7 +598,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
         } else {
             module <- what
         }
-        if (phase_offload) {
+        if (phase_offload && !(is.character(what) && what %in% resident)) {
             # Decode traces capture weight tensors; drop them so the
             # module's GPU memory actually frees
             .ltx23_release_vae_traces()
@@ -619,7 +660,18 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
     # Prefix conditioning: encode the start image (i2v) or the tail of
     # a previous clip (continuation) into normalized latents
     cond_latents <- NULL
-    if (conditioned) {
+    if (!is.null(condition_latents)) {
+        # Already-encoded prefix (e.g. ltx23_tail_latents of a previous
+        # chunk): no VAE touch at all
+        cond_latents <- condition_latents$to(device = device, dtype = f32)
+        if (cond_latents$ndim != 5L || cond_latents$shape[2] != 128L ||
+            cond_latents$shape[4] != s1_height ||
+            cond_latents$shape[5] != s1_width) {
+            stop(sprintf(
+                         "condition_latents must be [1, 128, k, %d, %d] for %dx%d generation",
+                         s1_height, s1_width, width, height))
+        }
+    } else if (conditioned) {
         vae <- onload("vae")
         cond_frames <- if (!is.null(image)) {
             ltx23_preprocess_frames(image, height, width)
@@ -765,13 +817,17 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
     }
 
     # The transformer and its dequant buffers are not needed past
-    # denoising; free the VRAM before the decoders claim it
+    # denoising; free the VRAM before the decoders claim it. A
+    # resident transformer keeps its buffers for the next chunk.
     offload("transformer")
-    ltx23_release_dequant_buffers()
+    if (!("transformer" %in% resident)) {
+        ltx23_release_dequant_buffers()
+    }
 
     result <- list(
                    latents = latents,
                    audio_latents = audio_latents,
+                   latent_shape = c(latent_frames, latent_height, latent_width),
                    sample_rate = 48000L
     )
     if (audio_conditioned) {
@@ -815,6 +871,12 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
             message(sprintf("  video to R array: %.1fs",
                             as.numeric(difftime(Sys.time(), phase_t0, units = "secs"))))
         }
+        if (trim_frames > 0L) {
+            # Deliver head-free pixels (the conditioning overlap of a
+            # continuation); result$latents keeps the full sequence
+            result$video <- result$video[-seq_len(trim_frames),,,,
+                drop = FALSE]
+        }
     }
 
     if (decode_audio) {

inst/tinytest/test_txt2vid_ltx23.R

@@ -187,9 +187,107 @@ expect_error(
   txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
     image = start_img, condition_video = tail_arr,
     device = "cpu", dtype = "float32"),
-  pattern = "not both"
+  pattern = "only one"
 )
 
+# --- Latent-space chaining seams -----------------------------------------------------
+
+# latent_shape rides in the result: 17 frames @ 64x64 -> [3, 2, 2]
+expect_equal(res_cont$latent_shape, c(3L, 2L, 2L))
+
+# condition_latents reproduces the condition_video path exactly when
+# fed the same encoded tail (the pixel path is just one producer)
+enc <- ltx23_encode_video_frames(
+  vae, ltx23_preprocess_frames(tail_arr, 64L, 64L))
+res_lat <- txt2vid_ltx2(
+  prompt = "tiny continuation",
+  pipeline = pipe,
+  prompt_embeds = stub_embeds,
+  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
+  seed = 7L, device = "cpu", dtype = "float32",
+  condition_latents = enc,
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+expect_true(max(abs(res_lat$video - res_cont$video)) < 1e-6)
+
+# Tail slicing: last 2 latent frames, matching a manual unpack + narrow
+tail_lat <- ltx23_tail_latents(res_cont, k = 2L)
+expect_equal(as.integer(tail_lat$shape), c(1L, 128L, 2L, 2L, 2L))
+manual <- diffuseR:::ltx23_unpack_video_latents(res_cont$latents, 3L, 2L, 2L)$
+  narrow(3L, 2L, 2L)
+expect_equal(as.numeric((tail_lat - manual)$abs()$max()), 0)
+
+# Feeding the sliced tail back runs end to end (chunk N -> chunk N+1)
+res_next <- txt2vid_ltx2(
+  prompt = "tiny chunk 3",
+  pipeline = pipe,
+  prompt_embeds = stub_embeds,
+  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
+  seed = 8L, device = "cpu", dtype = "float32",
+  condition_latents = tail_lat,
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+expect_equal(dim(res_next$video), c(17L, 64L, 64L, 3L))
+expect_true(all(is.finite(res_next$video)))
+
+# trim_frames delivers head-free pixels; latents keep the full sequence
+res_trim <- txt2vid_ltx2(
+  prompt = "tiny continuation",
+  pipeline = pipe,
+  prompt_embeds = stub_embeds,
+  width = 64L, height = 64L, num_frames = 17L, frame_rate = 24,
+  seed = 7L, device = "cpu", dtype = "float32",
+  condition_video = tail_arr, trim_frames = 9L,
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+expect_equal(dim(res_trim$video), c(8L, 64L, 64L, 3L))
+expect_equal(max(abs(res_trim$video -
+                     res_cont$video[-(1:9), , , , drop = FALSE])), 0)
+expect_equal(res_trim$latent_shape, c(3L, 2L, 2L))
+expect_equal(as.integer(ltx23_tail_latents(res_trim)$shape[3]), 2L)
+
+# resident= is accepted (inert on CPU, where phase_offload is off)
+res_res <- txt2vid_ltx2(
+  prompt = "tiny resident",
+  pipeline = pipe,
+  prompt_embeds = stub_embeds,
+  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
+  seed = 7L, device = "cpu", dtype = "float32",
+  resident = "transformer",
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+expect_true(all(is.finite(res_res$video)))
+
+# Guardrails
+expect_error(
+  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
+    condition_latents = enc, image = start_img,
+    device = "cpu", dtype = "float32"),
+  pattern = "only one"
+)
+expect_error(
+  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
+    condition_latents = torch::torch_randn(1L, 128L, 2L, 3L, 3L),
+    device = "cpu", dtype = "float32"),
+  pattern = "condition_latents"
+)
+expect_error(
+  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, num_frames = 9L,
+    trim_frames = 9L, device = "cpu", dtype = "float32"),
+  pattern = "trim_frames"
+)
+expect_error(
+  txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds,
+    resident = "flux_capacitor", device = "cpu", dtype = "float32"),
+  pattern = "resident"
+)
+expect_error(ltx23_tail_latents(res_cont$latents), pattern = "latent_shape")
+expect_error(ltx23_tail_latents(res_cont, k = 5L), pattern = "k must be")
+
 
 # --- Audio-conditioned generation (lip-sync plumbing) --------------------------------
 

man/ltx23_tail_latents.Rd

@@ -0,0 +1,39 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{ltx23_tail_latents}
+\alias{ltx23_tail_latents}
+\title{Slice the trailing latent frames of a generation for chaining}
+\usage{
+ltx23_tail_latents(result, k = 2L, latent_shape = NULL)
+}
+\arguments{
+\item{result}{A \code{\link{txt2vid_ltx2}} result list (uses its
+\code{latents} and \code{latent_shape}), or the packed latents
+tensor [1, S, 128] itself (then \code{latent_shape} is
+required).}
+
+\item{k}{Integer. Trailing latent frames to keep (default 2 = the
+standard 9-pixel-frame conditioning prefix).}
+
+\item{latent_shape}{Integer vector c(frames, height, width) of the
+latent geometry; only needed when \code{result} is a raw tensor.}
+}
+\value{
+Normalized latents [1, 128, k, H', W'] (float32), ready
+  for \code{txt2vid_ltx2(condition_latents = )}.
+}
+\description{
+Cuts the last \code{k} latent frames out of a result's video
+latents, in the [1, 128, k, H', W'] layout that
+\code{txt2vid_ltx2(condition_latents = )} consumes, so one chunk
+can seed the next without leaving latent space: no decode, no
+re-encode, no video round-trip.
+}
+\details{
+Semantics caveat: a latent frame sliced from inside a sequence
+represents 8 pixel frames, while a fresh VAE encode of a k-frame
+tail represents 1 + 8(k - 1) pixel frames with its first latent in
+first-frame form. The frozen prefix the next generation sees is
+therefore not identical to the pixel-path prefix; compare both on
+real content before relying on latent-only joins.
+
+}