diff --git a/R/dit_ltx23.R b/R/dit_ltx23.R index 725c545..f3265f8 100644 --- a/R/dit_ltx23.R +++ b/R/dit_ltx23.R @@ -203,7 +203,8 @@ ltx23_transformer <- torch::nn_module( spatio_temporal_guidance_blocks = NULL, perturbation_mask = NULL, use_cross_timestep = FALSE, - video_self_attention_mask = NULL + video_self_attention_mask = NULL, + cond_token_index = NULL ) { audio_timestep <- audio_timestep %||% timestep audio_sigma <- audio_sigma %||% sigma @@ -343,13 +344,25 @@ ltx23_transformer <- torch::nn_module( video_rotary_emb, audio_rotary_emb, video_ca_rotary_emb, audio_ca_rotary_emb, encoder_attention_mask, - audio_encoder_attention_mask + audio_encoder_attention_mask, + cond_token_index = cond_token_index ) hidden_states <- res[[1]] audio_hidden_states <- res[[2]] } block_gc <- isTRUE(getOption("diffuseR.block_gc")) + if (!is.null(cond_token_index)) { + # Expand the compact [B, V, ...] timestep variants per token + # (1-based index). The output layer always needs the expanded + # embedded_timestep; the eager blocks also need expanded temb + # (the compiled stack selects per token internally instead). + idx1 <- cond_token_index$to(dtype = torch::torch_long())$add(1L)$squeeze(1L) + embedded_timestep <- embedded_timestep$index_select(2L, idx1) + if (!use_jit) { + temb <- temb$index_select(2L, idx1) + } + } eager_blocks <- if (use_jit) integer(0) else seq_along(self$transformer_blocks) for (block_idx in eager_blocks) { is_stg_block <- (block_idx - 1L) %in% stg_blocks # reference indices are 0-based diff --git a/R/jit_ltx23.R b/R/jit_ltx23.R index 3d46e9e..7c1e96e 100644 --- a/R/jit_ltx23.R +++ b/R/jit_ltx23.R @@ -86,6 +86,17 @@ def mods(tbl: Tensor, temb: Tensor, num: int) -> Tensor: t = temb.size(1) return tbl.unsqueeze(0).unsqueeze(0) + temb.reshape([b, t, num, -1]) +def modtok(vada: Tensor, j: int, idx: Optional[Tensor]) -> Tensor: + # One modulation vector [B, T, D]. Without idx, T is 1 (global) or + # S (per-token) and broadcasts/aligns directly. With idx (prefix + # conditioning), T holds the distinct timestep variants and idx + # [S] selects per token -- materializing one [B, S, D] slice on + # demand instead of a [B, S, num, D] block. + v = vada.select(2, j) + if idx is None: + return v + return v.index_select(1, idx) + def attn_nf4(x: Tensor, ctx: Tensor, ws: List[Tensor], base: int, table: Tensor, heads: int, q_cos: Optional[Tensor], q_sin: Optional[Tensor], k_cos: Optional[Tensor], k_sin: Optional[Tensor], @@ -118,14 +129,15 @@ def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, v_cos: Tensor, v_sin: Tensor, a_cos: Tensor, a_sin: Tensor, cav_cos: Tensor, cav_sin: Tensor, caa_cos: Tensor, caa_sin: Tensor, enc_mask: Optional[Tensor], aenc_mask: Optional[Tensor], + cond_idx: Optional[Tensor], ws: List[Tensor], base: int, table: Tensor, heads: int, aheads: int) -> Tuple[Tensor, Tensor]: vada = mods(ws[base + 108], temb, 9) aada = mods(ws[base + 109], temb_a, 9) - nh = rmsn(h) * (vada.select(2, 1) + 1.0) + vada.select(2, 0) + nh = rmsn(h) * (modtok(vada, 1, cond_idx) + 1.0) + modtok(vada, 0, cond_idx) ax = attn_nf4(nh, nh, ws, base, table, heads, v_cos, v_sin, v_cos, v_sin, None) - h = h + ax * vada.select(2, 2) + h = h + ax * modtok(vada, 2, cond_idx) nah = rmsn(ah) * (aada.select(2, 1) + 1.0) + aada.select(2, 0) aax = attn_nf4(nah, nah, ws, base + 16, table, aheads, a_cos, a_sin, a_cos, a_sin, None) @@ -134,10 +146,10 @@ def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, pada = mods(ws[base + 110], tp, 2) apada = mods(ws[base + 111], tpa, 2) - nh = rmsn(h) * (vada.select(2, 7) + 1.0) + vada.select(2, 6) + nh = rmsn(h) * (modtok(vada, 7, cond_idx) + 1.0) + modtok(vada, 6, cond_idx) encm = enc * (pada.select(2, 1) + 1.0) + pada.select(2, 0) ax = attn_nf4(nh, encm, ws, base + 32, table, heads, None, None, None, None, enc_mask) - h = h + ax * vada.select(2, 8) + h = h + ax * modtok(vada, 8, cond_idx) nah = rmsn(ah) * (aada.select(2, 7) + 1.0) + aada.select(2, 6) aencm = aenc * (apada.select(2, 1) + 1.0) + apada.select(2, 0) @@ -161,8 +173,8 @@ def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, v2a = attn_nf4(mna, mnh, ws, base + 80, table, aheads, caa_cos, caa_sin, cav_cos, cav_sin, None) ah = ah + acg.select(2, 0) * v2a - nh = rmsn(h) * (vada.select(2, 4) + 1.0) + vada.select(2, 3) - h = h + ff_nf4(nh, ws, base + 96, table) * vada.select(2, 5) + nh = rmsn(h) * (modtok(vada, 4, cond_idx) + 1.0) + modtok(vada, 3, cond_idx) + h = h + ff_nf4(nh, ws, base + 96, table) * modtok(vada, 5, cond_idx) nah = rmsn(ah) * (aada.select(2, 4) + 1.0) + aada.select(2, 3) ah = ah + ff_nf4(nah, ws, base + 102, table) * aada.select(2, 5) @@ -176,6 +188,7 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, v_cos: Tensor, v_sin: Tensor, a_cos: Tensor, a_sin: Tensor, cav_cos: Tensor, cav_sin: Tensor, caa_cos: Tensor, caa_sin: Tensor, enc_mask: Optional[Tensor], aenc_mask: Optional[Tensor], + cond_idx: Optional[Tensor], ws: List[Tensor], table: Tensor, n_blocks: int, heads: int, aheads: int) -> Tuple[Tensor, Tensor]: i = 0 @@ -183,7 +196,7 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, h, ah = block_nf4(h, ah, enc, aenc, temb, temb_a, tcss, tcass, tcg, tcag, tp, tpa, v_cos, v_sin, a_cos, a_sin, cav_cos, cav_sin, caa_cos, caa_sin, - enc_mask, aenc_mask, ws, i * 114, table, heads, aheads) + enc_mask, aenc_mask, cond_idx, ws, i * 114, table, heads, aheads) i += 1 return (h, ah) " @@ -289,7 +302,8 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, audio_rotary_emb, ca_video_rotary_emb, ca_audio_rotary_emb, encoder_attention_mask = NULL, - audio_encoder_attention_mask = NULL) { + audio_encoder_attention_mask = NULL, + cond_token_index = NULL) { unit <- .ltx23_jit_unit() # unname: an nn_module_list yields named children, and a named R # list marshals to TorchScript as Dict[str, Tensor], not List[Tensor] @@ -318,7 +332,8 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, ca_video_rotary_emb[[1]], ca_video_rotary_emb[[2]], ca_audio_rotary_emb[[1]], ca_audio_rotary_emb[[2]], fix_mask(encoder_attention_mask), - fix_mask(audio_encoder_attention_mask), ws, table, + fix_mask(audio_encoder_attention_mask), + cond_token_index, ws, table, torch::jit_scalar(length(blocks)), torch::jit_scalar(as.integer(heads)), torch::jit_scalar(as.integer(aheads))) diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 9b66bff..b05e30d 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -99,6 +99,15 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { } else { NULL } + # Compact per-token timestep: the conditioned prefix sees exactly + # two values (t and 0), so the transformer gets a 2-entry timestep + # plus a per-token selector instead of a [B, S] vector -- the + # modulation tensors stay [B, 2, ...] until selected on demand + cond_idx <- if (!is.null(conditioning_mask)) { + conditioning_mask$squeeze(1L)$to(dtype = torch::torch_long()) + } else { + NULL + } video_coords <- transformer$rope$prepare_video_coords(latents$shape[1], latent_frames, latent_height, latent_width, device, fps = frame_rate) @@ -116,8 +125,9 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { t_video <- if (is.null(conditioning_mask)) { t } else { - # Per-token video timestep: conditioned tokens see zero - t * (1 - conditioning_mask) + # [t, 0]: index 0 = free tokens, index 1 = conditioned + torch::torch_tensor(c(sigma * scale_mult, 0), + device = device, dtype = f32) } out <- transformer( @@ -138,7 +148,8 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { audio_num_frames = audio_num_frames, video_coords = video_coords, audio_coords = audio_coords, - use_cross_timestep = TRUE + use_cross_timestep = TRUE, + cond_token_index = cond_idx ) # Euler velocity step in float32; dt is negative (sigma decreasing) diff --git a/inst/tinytest/test_jit_ltx23.R b/inst/tinytest/test_jit_ltx23.R index e68f56c..afdc364 100644 --- a/inst/tinytest/test_jit_ltx23.R +++ b/inst/tinytest/test_jit_ltx23.R @@ -212,3 +212,30 @@ torch::with_no_grad({ }) expect_true(max_abs_diff(out_pt[[1]], ref_pt[[1]]) < 1e-4) expect_true(max_abs_diff(out_pt[[2]], ref_pt[[2]]) < 1e-4) + +# Compact conditioned temb (2 variants + per-token index) must equal +# the eager block run with the explicitly expanded per-token temb +temb2 <- torch::torch_randn(B, 2L, 9L * dim) +cond_idx <- torch::torch_tensor(c(1L, 1L, 0L, 0L, 0L, 0L), + dtype = torch::torch_long()) # first two tokens conditioned +temb_full <- temb2$index_select(2L, cond_idx$add(1L)) +torch::with_no_grad({ + ref_ci <- blk1( + hidden_states = h, audio_hidden_states = ah, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + temb = temb_full, temb_audio = temb_a, + temb_ca_scale_shift = tcss, temb_ca_audio_scale_shift = tcass, + temb_ca_gate = tcg, temb_ca_audio_gate = tcag, + temb_prompt = tp, temb_prompt_audio = tpa, + video_rotary_emb = v_rope, audio_rotary_emb = a_rope, + ca_video_rotary_emb = cav_rope, ca_audio_rotary_emb = caa_rope + ) + out_ci <- diffuseR:::.ltx23_jit_run_stack( + list(blk1), h, ah, enc, aenc, + temb2, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v_rope, a_rope, cav_rope, caa_rope, + cond_token_index = cond_idx + ) +}) +expect_true(max_abs_diff(out_ci[[1]], ref_ci[[1]]) < 1e-4) +expect_true(max_abs_diff(out_ci[[2]], ref_ci[[2]]) < 1e-4)