diff --git a/.Rbuildignore b/.Rbuildignore index a83a93a..1d0835b 100644 --- a/.Rbuildignore +++ b/.Rbuildignore @@ -10,3 +10,5 @@ ^fyi\.md$ ^man-md$ ^inst/validation/\.venv$ +^ref$ +^tasks$ diff --git a/.gitignore b/.gitignore index 7397cae..464ad16 100644 --- a/.gitignore +++ b/.gitignore @@ -48,3 +48,8 @@ po/*~ # RStudio Connect folder rsconnect/ /devman-md/ + +# Upstream reference checkouts and working notes (never committed) +ref/ +tasks/ +.venv/ diff --git a/CLAUDE.md b/CLAUDE.md index e822882..f2683e9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,7 +51,7 @@ The `txt2img_*` and `img2img` functions default to `devices = "auto"`, which: ## GPU-Poor Execution Plan (TODO) -Profile-based memory optimization for constrained GPUs. Inspired by mmgp/Wan2GP approach. +Profile-based memory optimization for constrained GPUs. ### API @@ -188,7 +188,7 @@ vram_report <- function(label) { } ``` -### Future Optimizations (from mmgp) +### Future Optimizations Not implemented, but worth considering: @@ -320,66 +320,34 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt - [ ] Add SD3 model support - [ ] ControlNet integration -### LTX-2 Video Generation (In Progress) -- [x] FlowMatch scheduler (validated against Python) -- [x] RoPE positional embeddings (validated against Python) -- [x] LTX2 Video VAE (3D causal convolutions) - see learnings below -- [x] DiT transformer (audio-video) - see learnings below -- [x] Text encoder integration (connectors + flexible backends) -- [x] GPU-poor optimizations (wan2GP style memory profiles) -- [x] Pipeline integration (txt2vid_ltx2) -- [x] Video output utilities (save_video) -- [x] Weight loading from HuggingFace safetensors +### LTX-2.3 Video Generation (clean-room rewrite in progress) -#### LTX-2 Weight Loading +The original LTX-2.0 port was removed and is being replaced by a ground-up +LTX-2.3 implementation written from the Apache-2.0 diffusers reference only +(branch `feat/ltx23`; working checklist in `tasks/todo.md`). -Load LTX-2 model weights from HuggingFace safetensors: +**Reference checkouts** (in `ref/upstream/`, gitignored): +- `ref/upstream/diffusers` — the code reference (Apache-2.0) +- `ref/upstream/ltx2-official` — Lightricks monorepo (LTX-2 Community License; + used for numeric constants and file-format facts only, never as a code base) -```r -# Load VAE (2.44 GB) -vae <- load_ltx2_vae( - weights_path = "~/.cache/huggingface/hub/models--Lightricks--LTX-2/vae/diffusion_pytorch_model.safetensors", - config_path = "~/.cache/huggingface/hub/models--Lightricks--LTX-2/vae/config.json", - device = "cuda", - dtype = "float16" -) - -# Load transformer (37.8 GB, sharded across 8 files) -transformer <- load_ltx2_transformer( - weights_dir = "~/.cache/huggingface/hub/models--Lightricks--LTX-2/transformer", - device = "cpu", # Start on CPU, offload to GPU layer-by-layer - dtype = "float16" -) - -# Load text connectors (2.86 GB) -connectors <- load_ltx2_connectors( - weights_path = "~/.cache/huggingface/hub/models--Lightricks--LTX-2/connectors/diffusion_pytorch_model.safetensors", - config_path = "~/.cache/huggingface/hub/models--Lightricks--LTX-2/connectors/config.json" -) -``` - -**Model sizes:** -| Component | Size | Notes | -|-----------|------|-------| -| VAE | 2.44 GB | Single safetensors file | -| Transformer | 37.8 GB | Sharded across 8 files | -| Connectors | 2.86 GB | Single safetensors file | -| Total (19B) | 43.3 GB | Full precision | -| Total FP8 | 27.1 GB | Quantized | +**Rules:** +- Never derive code from Wan2GP (WanGP Community License 2.0 restricts + commercial embedding). Do not open Wan2GP source or conversions of it. +- Every technique must have a documented public source independent of + Wan2GP: record it in `inst/REFERENCES.md` when adding one. +- Model weights are downloaded by the user from `Lightricks/LTX-2.3` under the + LTX-2 Community License and never redistributed with the package. +- Kept support files with independent provenance: `gemma3_text_encoder.R` + (from HF transformers), `tokenizer_bpe.R` (original), `flowmatch_scheduler.R` + (from diffusers). -#### LTX2 VAE Implementation Learnings +#### Implementation facts carried over from the 2.0 port **Temporal dimension constraint for causal downsampling:** -For LTX2's causal 3D convolutions with stride S downsampling, the temporal dimension T must satisfy: -``` -(T + S - 1) % S == 0 -``` -This means T % S == 1, so **T must be odd for stride=2**. - -Example valid sequences for 2 spatiotemporal downsampling stages: -- T=5 → (5+1)/2=3 → (3+1)/2=2 ✓ -- T=9 → (9+1)/2=5 → (5+1)/2=3 ✓ -- T=4 → (4+1)/2=2.5 ✗ (fails unflatten) +For causal 3D convolutions with stride S downsampling, the temporal dimension T +must satisfy `(T + S - 1) %% S == 0`, i.e. T %% S == 1, so **T must be odd for +stride=2** (T=4 fails the unflatten; T=5 → 3 → 2 works). **R torch `unflatten` is 1-indexed:** ```r @@ -387,21 +355,9 @@ Example valid sequences for 2 spatiotemporal downsampling stages: # R: x$unflatten(3, c(-1, stride)) # dim 3 is 1-indexed ``` -**LTX2 decoder channel flow:** -In LTX2 up blocks, `in_channels == out_channels` always. The upsampler handles channel reduction via pixel shuffle. Test inputs must match this pattern. - -#### LTX2 DiT Transformer Learnings - -**cross_attention_dim must equal inner_dim:** -In the LTX2 transformer, `caption_projection` projects text embeddings from `caption_channels` to `inner_dim`. The transformer blocks then expect `encoder_hidden_states` to have dimension `cross_attention_dim`. These must be equal: -```r -# In model config: -inner_dim = num_attention_heads * attention_head_dim # e.g., 32 * 128 = 4096 -cross_attention_dim = 4096 # Must equal inner_dim! -``` - **R torch lacks nnf_scaled_dot_product_attention:** -Manual scaled dot-product attention is required: +Manual scaled dot-product attention is required (and materializes the full +[B, H, S, S] attention matrix — chunk queries at high resolution): ```r scale <- 1.0 / sqrt(head_dim) attn_weights <- torch::torch_matmul(query, key$transpose(-2L, -1L)) * scale @@ -410,32 +366,18 @@ attn_weights <- torch::nnf_softmax(attn_weights, dim = -1L) hidden_states <- torch::torch_matmul(attn_weights, value) ``` -**Avoid function name collisions across files:** -The `apply_interleaved_rotary_emb` function in `rope.R` expects `freqs$cos_freqs`, while a similar function in `dit_ltx2_modules.R` expects `freqs[[1]]`. Name collision caused segfaults - renamed to `apply_interleaved_rotary_emb_list` in dit module. - -#### LTX2 Text Encoder Learnings - -**Architecture: Gemma3 + Connectors:** -LTX-2 uses Gemma3 as the text encoder, with separate connector transformers for video and audio streams: -``` -Gemma3 → pack_text_embeds → text_proj_in → video_connector → video_embeds - → audio_connector → audio_embeds -``` - **Attention mask broadcasting:** -When using 2D attention masks [B, S], they must be expanded to [B, 1, 1, S] for broadcasting with attention weights [B, H, S, S]: +2D attention masks [B, S] must be expanded to [B, 1, 1, S] to broadcast with +attention weights [B, H, S, S]: ```r if (attention_mask$ndim == 2L) { attention_mask <- attention_mask$unsqueeze(2L)$unsqueeze(2L) } ``` -**Flexible text encoding backends:** -The `encode_text_ltx2()` function supports multiple backends: -- `"gemma3"`: Native R torch Gemma3 encoder (no Python dependency) -- `"precomputed"`: Load from file (cached embeddings) -- `"api"`: HTTP request to external service (Gemma container) -- `"random"`: Random embeddings for testing +**Avoid function name collisions across files:** +Two same-named helpers with different calling conventions in different files +once caused segfaults. One name, one signature, package-wide. #### Native Gemma3 Text Encoder diff --git a/DESCRIPTION b/DESCRIPTION index 081925d..7d344d6 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,8 +1,14 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.0.2 -Authors@R: - person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre")) +Version: 0.1.0 +Authors@R: c( + person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), + comment = c(ORCID = "0009-0005-4248-604X")), + person("cornball.ai", role = "cph"), + person("The HuggingFace Team", role = "cph", + comment = "portions ported from the diffusers library (Apache-2.0); see inst/COPYRIGHTS"), + person("Lightricks Ltd.", role = "cph", + comment = "LTX checkpoint layout and pipeline constants; see inst/COPYRIGHTS")) Description: A native R implementation of diffusion models providing a functional interface to state-of-the-art generative AI. Inspired by Hugging Face's Python 'diffusers' library, 'diffuseR' allows users to generate and manipulate images diff --git a/NAMESPACE b/NAMESPACE index 8133d0b..7f9d48e 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -1,22 +1,16 @@ # tinyrox says don't edit this manually, but it can't stop you! -export(apply_interleaved_rotary_emb) -export(apply_split_rotary_emb) export(auto_devices) export(bpe_tokenizer) export(clear_vram) export(CLIPTokenizer) -export(configure_vae_for_profile) export(ddim_scheduler_create) export(ddim_scheduler_step) export(decode_bpe) -export(dequantize_int4) -export(diagonal_gaussian_distribution) -export(dit_offloaded_forward) export(download_component) +export(download_ltx2) export(download_model) export(encode_bpe) -export(encode_text_ltx2) export(encode_with_gemma3) export(filename_from_prompt) export(flowmatch_calculate_shift) @@ -28,19 +22,10 @@ export(gemma3_config_ltx2) export(gemma3_text_model) export(gemma3_tokenizer) export(img2img) -export(int4_linear) -export(int4_linear_from_quantized) export(is_blackwell_gpu) export(latents_to_video) -export(linear_to_int4) export(load_decoder_weights) export(load_gemma3_text_encoder) -export(load_int4_into_model) -export(load_int4_weights) -export(load_int4_weights_into_model) -export(load_ltx2_connectors) -export(load_ltx2_transformer) -export(load_ltx2_vae) export(load_model_component) export(load_pipeline) export(load_text_encoder_weights) @@ -48,43 +33,97 @@ export(load_text_encoder2_weights) export(load_to_gpu) export(load_unet_sdxl_weights) export(load_unet_weights) -export(ltx_video_downsampler3d) -export(ltx_video_upsampler3d) -export(ltx2_memory_profile) -export(ltx2_text_connectors) -export(ltx2_video_causal_conv3d) -export(ltx2_video_decoder3d) -export(ltx2_video_down_block3d) -export(ltx2_video_encoder3d) -export(ltx2_video_mid_block3d) -export(ltx2_video_resnet_block3d) -export(ltx2_video_transformer_3d_model) -export(ltx2_video_up_block3d) -export(ltx2_video_vae) -export(make_linear) +export(ltx23_ada_layer_norm_single) +export(ltx23_adain_filter_latent) +export(ltx23_antialias_act1d) +export(ltx23_apply_interleaved_rotary_emb) +export(ltx23_apply_split_rotary_emb) +export(ltx23_attention) +export(ltx23_audio_causal_conv2d) +export(ltx23_audio_decoder) +export(ltx23_audio_downsample) +export(ltx23_audio_encoder) +export(ltx23_audio_mel_frontend) +export(ltx23_audio_resnet_block) +export(ltx23_audio_upsample) +export(ltx23_audio_vae) +export(ltx23_causal_conv3d) +export(ltx23_census) +export(ltx23_connector_transformer_1d) +export(ltx23_denormalize_latents) +export(ltx23_distilled_sigmas) +export(ltx23_downsample1d) +export(ltx23_encode_audio) +export(ltx23_encode_video_frames) +export(ltx23_feed_forward) +export(ltx23_fp8_linear) +export(ltx23_get_timestep_embedding) +export(ltx23_is_fp8_cast_key) +export(ltx23_kaiser_sinc_filter1d) +export(ltx23_latent_upsampler) +export(ltx23_load_group) +export(ltx23_load_pipeline) +export(ltx23_load_transformer_fp8) +export(ltx23_load_transformer_nf4) +export(ltx23_load_upsampler) +export(ltx23_map_audio_vae_key) +export(ltx23_map_connector_key) +export(ltx23_map_dit_key) +export(ltx23_map_vae_key) +export(ltx23_map_vocoder_key) +export(ltx23_mel_stft) +export(ltx23_memory_profile) +export(ltx23_nf4_dequantize) +export(ltx23_nf4_linear) +export(ltx23_nf4_quantize) +export(ltx23_normalize_latents) +export(ltx23_open_checkpoint) +export(ltx23_open_fp8_checkpoint) +export(ltx23_per_channel_rms_norm) +export(ltx23_per_token_rms_norm) +export(ltx23_prepare_conditioned_latents) +export(ltx23_preprocess_frames) +export(ltx23_quantize_fp8) +export(ltx23_quantize_nf4) +export(ltx23_read_audio) +export(ltx23_read_tail_frames) +export(ltx23_release_dequant_buffers) +export(ltx23_rms_norm) +export(ltx23_rotary_pos_embed) +export(ltx23_rotary_pos_embed_1d) +export(ltx23_set_attn_chunk) +export(ltx23_snake_beta) +export(ltx23_split_keys) +export(ltx23_stage2_distilled_sigmas) +export(ltx23_text_connectors) +export(ltx23_tone_map_latents) +export(ltx23_transformer) +export(ltx23_transformer_block) +export(ltx23_tune_gc) +export(ltx23_upsample1d) +export(ltx23_video_decoder3d) +export(ltx23_video_down_block3d) +export(ltx23_video_downsampler3d) +export(ltx23_video_encoder3d) +export(ltx23_video_mid_block3d) +export(ltx23_video_resnet_block3d) +export(ltx23_video_up_block3d) +export(ltx23_video_upsampler3d) +export(ltx23_video_vae) +export(ltx23_vocoder) +export(ltx23_vocoder_resblock) +export(ltx23_vocoder_with_bwe) export(models2devices) export(offload_to_cpu) -export(pack_text_embeds) -export(pack_video_latents) -export(per_channel_rms_norm) export(post_quant_conv) export(preprocess_image) export(quant_conv) -export(quantize_int4) -export(quantize_ltx2_transformer) -export(quantize_ltx2_vae) -export(quantize_model_int4) -export(quantize_safetensors_int4) -export(rope_embedder_create) -export(rope_forward) -export(rope_prepare_video_coords) export(save_frames) export(save_image) -export(save_int4_weights) export(save_video) +export(save_video_ltx23) export(scheduler_add_noise) export(sdxl_memory_profile) -export(sequential_cfg_forward) export(text_encoder_native) export(text_encoder2_native) export(tokenize_gemma3) @@ -96,10 +135,12 @@ export(unet_native) export(unet_native_from_torchscript) export(unet_sdxl_native) export(unet_sdxl_native_from_torchscript) -export(unpack_video_latents) export(vae_decoder_native) -export(validate_resolution) export(vocab_size) export(vram_report) +export(write_wav) S3method(print,bpe_tokenizer) +S3method(print,ltx23_checkpoint) + +importFrom(utils,head) diff --git a/R/CLIPTokenizer.R b/R/CLIPTokenizer.R index 47cdbad..a035043 100644 --- a/R/CLIPTokenizer.R +++ b/R/CLIPTokenizer.R @@ -1,33 +1,30 @@ - # Helper: get all adjacent symbol pairs token_get_pairs <- function(symbols) { - pairs <- character(0) - if (length(symbols) > 1) { - for (i in seq_len(length(symbols) - 1)) { - pairs <- c(pairs, paste(symbols[i], symbols[i + 1], sep = " ")) + pairs <- character(0) + if (length(symbols) > 1) { + for (i in seq_len(length(symbols) - 1)) { + pairs <- c(pairs, paste(symbols[i], symbols[i + 1], sep = " ")) + } } - } - unique(pairs) + unique(pairs) } # Helper: merge the first occurrence of a given bigram -token_merge_pair_once <- function( - symbols, - bigram -) { - parts <- strsplit(bigram, " ") [[1]] - out <- c() - i <- 1 - while (i <= length(symbols)) { - if (i < length(symbols) && symbols[i] == parts[1] && symbols[i + 1] == parts[2]) { - out <- c(out, paste0(parts[1], parts[2])) - i <- i + 2 - } else { - out <- c(out, symbols[i]) - i <- i + 1 +token_merge_pair_once <- function(symbols, bigram) { + parts <- strsplit(bigram, " ")[[1]] + out <- c() + i <- 1 + while (i <= length(symbols)) { + if (i < length(symbols) && symbols[i] == parts[1] && + symbols[i + 1] == parts[2]) { + out <- c(out, paste0(parts[1], parts[2])) + i <- i + 2 + } else { + out <- c(out, symbols[i]) + i <- i + 1 + } } - } - out + out } #' Tokenize a prompt @@ -36,81 +33,78 @@ token_merge_pair_once <- function( #' @param merges Path to the merges file (BPE merges). #' @param vocab_file Path to the vocabulary file (token->id mapping). #' @param pad_token The token ID used for padding (default is 0). -#`` +#' #' @return A 2D torch tensor of shape c(1, 77) containing the token IDs. #' @export -CLIPTokenizer <- function( - prompt, - merges = system.file("tokenizer/merges.txt", package = "diffuseR"), - vocab_file = system.file("tokenizer/vocab.json", package = "diffuseR"), - pad_token = 0L -) { - # 1. Load merges and build BPE rank map - merge_lines <- readLines(merges, encoding = "UTF-8") - merge_lines <- merge_lines[- 1]# drop header - merge_lines <- merge_lines[nzchar(merge_lines)]# drop blanks - merges_list <- strsplit(merge_lines, " ") - merge_keys <- sapply(merges_list, paste, collapse = " ") - bpe_ranks <- stats::setNames(seq_along(merge_keys), merge_keys) +CLIPTokenizer <- function(prompt, + merges = system.file("tokenizer/merges.txt", package = "diffuseR"), + vocab_file = system.file("tokenizer/vocab.json", package = "diffuseR"), + pad_token = 0L) { + # 1. Load merges and build BPE rank map + merge_lines <- readLines(merges, encoding = "UTF-8") + merge_lines <- merge_lines[-1] # drop header + merge_lines <- merge_lines[nzchar(merge_lines)] # drop blanks + merges_list <- strsplit(merge_lines, " ") + merge_keys <- sapply(merges_list, paste, collapse = " ") + bpe_ranks <- stats::setNames(seq_along(merge_keys), merge_keys) - # 2. Load vocabulary JSON (token->id) - vocab <- jsonlite::fromJSON(vocab_file) + # 2. Load vocabulary JSON (token->id) + vocab <- jsonlite::fromJSON(vocab_file) - # 3. Prepare text: lowercase and split into words - text <- tolower(prompt) - words <- strsplit(text, "\\s+") [[1]] + # 3. Prepare text: lowercase and split into words + text <- tolower(prompt) + words <- strsplit(text, "\\s+")[[1]] - # 4. BPE tokenization per word - tokens <- character(0) - for (word in words) { - chars <- strsplit(word, "") [[1]] - symbols <- if (length(chars) > 0) { - c( - if (length(chars) > 1) chars[- length(chars)] else character(0), - paste0(chars[length(chars)], "") - ) - } else { - character(0) - } - repeat { - pairs <- token_get_pairs(symbols) - valid <- intersect(pairs, names(bpe_ranks)) - if (length(valid) == 0) break - best <- valid[which.min(bpe_ranks[valid])] - symbols <- token_merge_pair_once(symbols, best) + # 4. BPE tokenization per word + tokens <- character(0) + for (word in words) { + chars <- strsplit(word, "")[[1]] + symbols <- if (length(chars) > 0) { + c(if (length(chars) > 1) chars[-length(chars)] else character(0), + paste0(chars[length(chars)], "")) + } else { + character(0) + } + repeat { + pairs <- token_get_pairs(symbols) + valid <- intersect(pairs, names(bpe_ranks)) + if (length(valid) == 0) { + break + } + best <- valid[which.min(bpe_ranks[valid])] + symbols <- token_merge_pair_once(symbols, best) + } + tokens <- c(tokens, symbols) } - tokens <- c(tokens, symbols) - } - # Truncate tokens to max length - max_tokens <- 77L - 2L# 2 for start and end tokens - if (length(tokens) > max_tokens) { - truncated_tokens <- tokens[(max_tokens + 1) :length(tokens)] - tokens <- tokens[1:max_tokens] - warning("Prompt was truncated. Consider shortening it.") - warning("Dropped prompt: ", paste(truncated_tokens, collapse = " ")) - } + # Truncate tokens to max length + max_tokens <- 77L - 2L # 2 for start and end tokens + if (length(tokens) > max_tokens) { + truncated_tokens <- tokens[(max_tokens + 1):length(tokens)] + tokens <- tokens[1:max_tokens] + warning("Prompt was truncated. Consider shortening it.") + warning("Dropped prompt: ", paste(truncated_tokens, collapse = " ")) + } - # 5. Map token strings to vocabulary IDs - ids <- vapply(tokens, function(tok) { - if (!tok %in% names(vocab)) stop(paste("Unknown token:", tok)) - vocab[[tok]] + # 5. Map token strings to vocabulary IDs + ids <- vapply(tokens, function(tok) { + if (!tok %in% names(vocab)) stop(paste("Unknown token:", tok)) + vocab[[tok]] }, integer(1)) - # 6. Add special tokens and pad/truncate - sot_id <- 49406L - eot_id <- 49407L - seq_ids <- c(sot_id, ids, eot_id) - max_len <- 77L - if (length(seq_ids) < max_len) { - seq_ids <- c(seq_ids, rep(pad_token, max_len - length(seq_ids))) - } else if (length(seq_ids) > max_len) { - seq_ids <- seq_ids[1:max_len] - } + # 6. Add special tokens and pad/truncate + sot_id <- 49406L + eot_id <- 49407L + seq_ids <- c(sot_id, ids, eot_id) + max_len <- 77L + if (length(seq_ids) < max_len) { + seq_ids <- c(seq_ids, rep(pad_token, max_len - length(seq_ids))) + } else if (length(seq_ids) > max_len) { + seq_ids <- seq_ids[1:max_len] + } - # 7. Return as 2D torch tensor [1, 77] - t <- torch::torch_tensor(seq_ids, dtype = torch::torch_long()) - t <- t$unsqueeze(1) # prepend batch dimension - t + # 7. Return as 2D torch tensor [1, 77] + t <- torch::torch_tensor(seq_ids, dtype = torch::torch_long()) + t <- t$unsqueeze(1) # prepend batch dimension + t } - diff --git a/R/audio_encode_ltx23.R b/R/audio_encode_ltx23.R new file mode 100644 index 0000000..f5e3925 --- /dev/null +++ b/R/audio_encode_ltx23.R @@ -0,0 +1,180 @@ +#' Audio Conditioning Frontend for LTX-2.3 +#' +#' Turns user audio into the normalized, packed audio latents the joint +#' denoiser conditions on (lip sync): decode to 16 kHz stereo PCM, +#' log-mel via a causal STFT (filter 1024, hop 160, 64 slaney-normed +#' mel bins to 8 kHz — the checkpoint's preprocessing spec), then the +#' audio VAE encoder in argmax mode. The STFT and mel-filterbank +#' constructors were verified against the checkpoint's stored vocoder +#' bases (identical up to bf16 rounding), so the convention matches +#' training. +#' +#' @name audio_encode_ltx23 +NULL + +# Periodic Hann window +.ltx23_hann <- function(n) { + 0.5 - 0.5 * cos(2 * pi * (0:(n - 1L)) / n) +} + +# Windowed Fourier basis in the layout ltx23_causal_stft expects: +# [n_freqs * 2, 1, filter_length], cos rows then -sin rows, each times +# the periodic Hann window (matches vocoder.mel_stft.stft_fn.forward_basis) +.ltx23_stft_basis <- function(filter_length) { + nf <- filter_length %/% 2L + 1L + n <- 0:(filter_length - 1L) + win <- .ltx23_hann(filter_length) + basis <- matrix(0, nf * 2L, filter_length) + for (k in 0:(nf - 1L)) { + basis[k + 1L,] <- cos(2 * pi * k * n / filter_length) * win + basis[nf + k + 1L,] <- -sin(2 * pi * k * n / filter_length) * win + } + torch::torch_tensor(basis, dtype = torch::torch_float32())$unsqueeze(2L) +} + +# Slaney-scale, slaney-normed mel filterbank (matches +# vocoder.mel_stft.mel_basis up to bf16 rounding) +.ltx23_mel_filterbank <- function(sample_rate, n_fft, n_mels, fmin, fmax) { + hz2mel <- function(f) { + ifelse(f < 1000, f * 3 / 200, 15 + log(f / 1000) / (log(6.4) / 27)) + } + mel2hz <- function(m) { + ifelse(m < 15, m * 200 / 3, 1000 * exp((m - 15) * log(6.4) / 27)) + } + nf <- n_fft %/% 2L + 1L + fft_freqs <- (0:(nf - 1L)) * sample_rate / n_fft + hz <- mel2hz(seq(hz2mel(fmin), hz2mel(fmax), length.out = n_mels + 2L)) + W <- matrix(0, n_mels, nf) + for (i in seq_len(n_mels)) { + lower <- (fft_freqs - hz[i]) / (hz[i + 1] - hz[i]) + upper <- (hz[i + 2] - fft_freqs) / (hz[i + 2] - hz[i + 1]) + W[i,] <- pmax(0, pmin(lower, upper)) * 2 / (hz[i + 2] - hz[i]) + } + torch::torch_tensor(W, dtype = torch::torch_float32()) +} + +#' Build the 16 kHz log-mel frontend for audio conditioning +#' +#' An \code{\link{ltx23_mel_stft}} whose STFT and mel bases are +#' constructed (not checkpoint-loaded) with the audio VAE's +#' preprocessing spec. +#' +#' @param filter_length,hop_length,n_mels,sample_rate,fmin,fmax The +#' checkpoint preprocessing parameters (defaults are LTX-2.3's). +#' +#' @return An \code{ltx23_mel_stft} module. +#' +#' @export +ltx23_audio_mel_frontend <- function(filter_length = 1024L, + hop_length = 160L, n_mels = 64L, + sample_rate = 16000L, fmin = 0, + fmax = 8000) { + frontend <- ltx23_mel_stft(filter_length = filter_length, + hop_length = hop_length, + window_length = filter_length, + num_mel_channels = n_mels) + torch::with_no_grad({ + frontend$stft_fn$forward_basis$copy_(.ltx23_stft_basis(filter_length)) + frontend$mel_basis$copy_(.ltx23_mel_filterbank(sample_rate, + filter_length, n_mels, fmin, fmax)) + }) + frontend$eval() + frontend +} + +#' Read an audio file as 16 kHz stereo PCM +#' +#' Decodes MP3/WAV/etc. via \code{av} to 16-bit PCM at the target rate +#' and parses the RIFF container in base R. +#' +#' @param path Audio file. +#' @param sample_rate Integer. +#' +#' @return Matrix [2, samples] in [-1, 1]. +#' +#' @export +ltx23_read_audio <- function(path, sample_rate = 16000L) { + if (!requireNamespace("av", quietly = TRUE)) { + stop("Reading audio requires the 'av' package.") + } + wav <- tempfile(fileext = ".wav") + on.exit(unlink(wav), add = TRUE) + av::av_audio_convert(path, wav, format = NULL, channels = 2L, + sample_rate = sample_rate, verbose = FALSE) + + con <- file(wav, "rb") + on.exit(close(con), add = TRUE) + riff <- readBin(con, "raw", 12L) + stopifnot(rawToChar(riff[1:4]) == "RIFF") + n_channels <- 2L + repeat { + hdr <- readBin(con, "raw", 8L) + if (length(hdr) < 8L) { + stop("No data chunk found in ", wav) + } + id <- rawToChar(hdr[1:4]) + size <- readBin(hdr[5:8], "integer", 1L, size = 4L, endian = "little") + if (id == "fmt ") { + fmt <- readBin(con, "raw", size) + n_channels <- readBin(fmt[3:4], "integer", 1L, size = 2L, + endian = "little") + } else if (id == "data") { + pcm <- readBin(con, "integer", size %/% 2L, size = 2L, + signed = TRUE, endian = "little") + break + } else { + invisible(readBin(con, "raw", size + size %% 2L)) + } + } + m <- matrix(pcm / 32768, nrow = n_channels) + if (n_channels == 1L) { + m <- rbind(m, m) + } + m +} + +#' Encode audio into normalized, packed conditioning latents +#' +#' Pads or trims the waveform so the latent length equals +#' \code{audio_num_frames} (mel frames \code{4L - 3}, mirroring the +#' decoder's \code{target_frames}), computes the log-mel, encodes in +#' argmax mode, packs, and normalizes with the checkpoint statistics. +#' +#' @param audio_vae An \code{ltx23_audio_vae} (with encoder weights). +#' @param wav Matrix [2, samples] in [-1, 1] at 16 kHz (see +#' \code{\link{ltx23_read_audio}}). +#' @param audio_num_frames Integer. Target latent length. +#' @param frontend Optional prebuilt \code{\link{ltx23_audio_mel_frontend}}. +#' +#' @return Packed normalized latents [1, audio_num_frames, 128] (float32). +#' +#' @export +ltx23_encode_audio <- function(audio_vae, wav, audio_num_frames, + frontend = NULL) { + frontend <- frontend %||% ltx23_audio_mel_frontend() + target_mel_frames <- 4L * audio_num_frames - 3L + + dev <- audio_vae$latents_mean$device + enc_dtype <- audio_vae$encoder$conv_in$conv$weight$dtype + frontend$to(device = dev) + + x <- torch::torch_tensor(wav, dtype = torch::torch_float32(), device = dev) + torch::with_no_grad({ + # Stereo channels through the STFT as a batch of mono signals + mel <- frontend(x$unsqueeze(2L)) # [2, 64, T] + # -> [1, 2, T, 64] + mel <- mel$permute(c(1L, 3L, 2L))$unsqueeze(1L) + t_cur <- mel$shape[3] + if (t_cur < target_mel_frames) { + mel <- torch::nnf_pad(mel, c(0L, 0L, 0L, target_mel_frames - t_cur)) + } else if (t_cur > target_mel_frames) { + mel <- mel$narrow(3L, 1L, target_mel_frames) + } + moments <- audio_vae$encode(mel$to(dtype = enc_dtype)) + latents <- moments$mean$to(dtype = torch::torch_float32()) + }) + packed <- ltx23_pack_audio_latents(latents) + mean <- audio_vae$latents_mean$to(dtype = torch::torch_float32()) + std <- audio_vae$latents_std$to(dtype = torch::torch_float32()) + (packed - mean) / std +} diff --git a/R/audio_vae_ltx23.R b/R/audio_vae_ltx23.R new file mode 100644 index 0000000..1bb2901 --- /dev/null +++ b/R/audio_vae_ltx23.R @@ -0,0 +1,454 @@ +#' LTX-2.3 Audio VAE +#' +#' Fresh R port of the LTX-2 audio autoencoder from the diffusers +#' reference (Apache-2.0, autoencoder_kl_ltx2_audio.py), configured per +#' the checkpoint: pixel norm, height-axis causality, base 128 channels +#' with multipliers (1, 2, 4), 8 latent channels, 64 mel bins, no +#' attention. The decoder produces mel for the vocoder; the encoder +#' turns user audio into conditioning latents (lip sync). +#' +#' @name audio_vae_ltx23 +NULL + +# Audio latents are 4x downsampled in time relative to mel frames +.ltx23_audio_latent_downsample_factor <- 4L + +#' Causal 2D convolution for audio spectrograms +#' +#' Pads asymmetrically along the causal axis ("height" = time frames for +#' LTX audio) before an unpadded Conv2d. +#' +#' @param in_channels,out_channels Integers. +#' @param kernel_size Integer or length-2 vector. +#' @param stride Integer. +#' @param causality_axis "height", "width", "width-compatibility", or "none". +#' +#' @export +ltx23_audio_causal_conv2d <- torch::nn_module( + "ltx23_audio_causal_conv2d", + initialize = function( + in_channels, + out_channels, + kernel_size = 3L, + stride = 1L, + causality_axis = "height" + ) { + if (length(kernel_size) == 1L) kernel_size <- rep(kernel_size, 2L) + pad_h <- kernel_size[1] - 1L + pad_w <- kernel_size[2] - 1L + + self$padding <- switch(causality_axis, + none = c(pad_w %/% 2L, pad_w - pad_w %/% 2L, pad_h %/% 2L, + pad_h - pad_h %/% 2L), + width = c(pad_w, 0L, pad_h %/% 2L, pad_h - pad_h %/% 2L), + `width-compatibility` = c(pad_w, 0L, pad_h %/% 2L, pad_h - pad_h %/% 2L), + height = c(pad_w %/% 2L, pad_w - pad_w %/% 2L, pad_h, 0L), + stop("Invalid causality_axis: ", causality_axis) + ) + self$conv <- torch::nn_conv2d( + in_channels, out_channels, kernel_size, + stride = stride, padding = 0L + ) +}, + forward = function(x) { + self$conv(torch::nnf_pad(x, self$padding)) +} +) + +#' LTX audio ResNet block +#' +#' PixelNorm -> SiLU -> causal conv, twice, with a 1x1 causal conv +#' shortcut (\code{nin_shortcut}) on channel change. +#' +#' @param in_channels,out_channels Integers. +#' @param causality_axis Character. +#' +#' @export +ltx23_audio_resnet_block <- torch::nn_module( + "ltx23_audio_resnet_block", + initialize = function( + in_channels, + out_channels = NULL, + causality_axis = "height" + ) { + out_channels <- out_channels %||% in_channels + self$changes_channels <- in_channels != out_channels + + self$norm1 <- ltx23_per_channel_rms_norm(eps = 1e-6) + self$conv1 <- ltx23_audio_causal_conv2d(in_channels, out_channels, + kernel_size = 3L, causality_axis = causality_axis) + self$norm2 <- ltx23_per_channel_rms_norm(eps = 1e-6) + self$conv2 <- ltx23_audio_causal_conv2d( + out_channels, out_channels, kernel_size = 3L, causality_axis = causality_axis + ) + if (self$changes_channels) { + self$nin_shortcut <- ltx23_audio_causal_conv2d( + in_channels, out_channels, kernel_size = 1L, causality_axis = causality_axis + ) + } +}, + forward = function(x) { + h <- self$norm1(x) + h <- torch::nnf_silu(h) + h <- self$conv1(h) + h <- self$norm2(h) + h <- torch::nnf_silu(h) + h <- self$conv2(h) + if (self$changes_channels) { + x <- self$nin_shortcut(x) + } + x + h +} +) + +#' LTX audio upsampler +#' +#' Nearest 2x interpolation, causal conv, then a crop of the first +#' element along the causal axis. +#' +#' @param in_channels Integer. +#' @param causality_axis Character. +#' +#' @export +ltx23_audio_upsample <- torch::nn_module( + "ltx23_audio_upsample", + initialize = function(in_channels, causality_axis = "height") { + self$causality_axis <- causality_axis + self$conv <- ltx23_audio_causal_conv2d(in_channels, in_channels, + kernel_size = 3L, causality_axis = causality_axis) +}, + forward = function(x) { + x <- torch::nnf_interpolate(x, scale_factor = 2, mode = "nearest") + x <- self$conv(x) + if (self$causality_axis == "height") { + # Drop the first (causally padded) frame; Python [:, :, 1:, :] + x <- x$narrow(3L, 2L, x$shape[3] - 1L) + } else if (self$causality_axis == "width") { + x <- x$narrow(4L, 2L, x$shape[4] - 1L) + } + x +} +) + +#' LTX audio downsampler +#' +#' Causal zero-pad followed by a plain stride-2 conv (reference +#' LTX2AudioDownsample; note the conv is unwrapped, so its checkpoint +#' key is \code{downsample.conv.*}). +#' +#' @param in_channels Integer. +#' @param causality_axis Character. +#' +#' @export +ltx23_audio_downsample <- torch::nn_module( + "ltx23_audio_downsample", + initialize = function(in_channels, causality_axis = "height") { + # Padding order: (left, right, top, bottom) + self$padding <- switch(causality_axis, none = c(0L, 1L, 0L, 1L), + width = c(2L, 0L, 0L, 1L), + height = c(0L, 1L, 2L, 0L), + `width-compatibility` = c(1L, 0L, 0L, 1L), + stop("Invalid causality_axis: ", causality_axis)) + self$conv <- torch::nn_conv2d(in_channels, in_channels, kernel_size = 3L, + stride = 2L, padding = 0L) +}, + forward = function(x) { + self$conv(torch::nnf_pad(x, self$padding)) +} +) + +#' LTX-2.3 audio VAE encoder +#' +#' Mel spectrogram [B, 2, T, 64] -> latent distribution moments +#' [B, 2 * latent_channels, ceil(T/4), 16]. Structure mirrors the +#' decoder: causal convs, parameterless pixel norms, ResNet stages with +#' stride-2 downsampling between levels (reference LTX2AudioEncoder). +#' +#' @param base_channels,num_res_blocks,latent_channels,ch_mult,causality_axis +#' See \code{\link{ltx23_audio_decoder}}. +#' @param in_channels Integer. Mel channels (2 = stereo). +#' +#' @export +ltx23_audio_encoder <- torch::nn_module( + "ltx23_audio_encoder", + initialize = function( + base_channels = 128L, + in_channels = 2L, + num_res_blocks = 2L, + latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), + causality_axis = "height" + ) { + num_levels <- length(ch_mult) + self$conv_in <- ltx23_audio_causal_conv2d(in_channels, base_channels, + kernel_size = 3L, causality_axis = causality_axis) + + block_in <- base_channels + stages <- list() + for (level in seq_len(num_levels)) { + block_out <- base_channels * ch_mult[level] + blocks <- list() + for (j in seq_len(num_res_blocks)) { + blocks[[j]] <- ltx23_audio_resnet_block( + if (j == 1L) block_in else block_out, block_out, + causality_axis = causality_axis + ) + } + block_in <- block_out + stage <- torch::nn_module( + "ltx23_audio_down_stage", + initialize = function(blocks, downsample) { + self$block <- torch::nn_module_list(blocks) + if (!is.null(downsample)) { + self$downsample <- downsample + } + }, + forward = function(x) { + for (i in seq_along(self$block)) { + x <- self$block[[i]](x) + } + if (!is.null(self$downsample)) { + x <- self$downsample(x) + } + x + } + ) + stages[[level]] <- stage( + blocks, + if (level != num_levels) { + ltx23_audio_downsample(block_in, causality_axis = causality_axis) + } else { + NULL + } + ) + } + self$down <- torch::nn_module_list(stages) + + self$mid <- torch::nn_module( + "ltx23_audio_mid", + initialize = function(channels, causality_axis) { + self$block_1 <- ltx23_audio_resnet_block(channels, + causality_axis = causality_axis) + self$block_2 <- ltx23_audio_resnet_block(channels, + causality_axis = causality_axis) + }, + forward = function(x) { + self$block_2(self$block_1(x)) + } + )(block_in, causality_axis) + + self$norm_out <- ltx23_per_channel_rms_norm(eps = 1e-6) + self$conv_out <- ltx23_audio_causal_conv2d( + block_in, 2L * latent_channels, kernel_size = 3L, + causality_axis = causality_axis + ) + self$latent_channels <- as.integer(latent_channels) +}, + forward = function(x) { + h <- self$conv_in(x) + for (level in seq_along(self$down)) { + h <- self$down[[level]](h) + } + h <- self$mid(h) + h <- self$norm_out(h) + h <- torch::nnf_silu(h) + self$conv_out(h) +} +) + +#' LTX-2.3 audio VAE decoder +#' +#' Latents [B, 8, L, 16] -> mel spectrogram [B, 2, 4L - 3, 64]. +#' +#' @param base_channels Integer. +#' @param output_channels Integer. Audio channels (2 = stereo). +#' @param num_res_blocks Integer. Per-level ResNet count (a stage runs +#' \code{num_res_blocks + 1} blocks). +#' @param latent_channels Integer. +#' @param ch_mult Integer vector. Channel multipliers per level. +#' @param causality_axis Character. +#' @param mel_bins Integer. Output mel bins (crop/pad target). +#' +#' @export +ltx23_audio_decoder <- torch::nn_module( + "ltx23_audio_decoder", + initialize = function( + base_channels = 128L, + output_channels = 2L, + num_res_blocks = 2L, + latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), + causality_axis = "height", + mel_bins = 64L + ) { + self$num_resolutions <- length(ch_mult) + self$num_res_blocks <- num_res_blocks + self$out_ch <- as.integer(output_channels) + self$mel_bins <- as.integer(mel_bins) + self$causal <- causality_axis != "none" + + base_block_channels <- base_channels * ch_mult[length(ch_mult)] + + self$conv_in <- ltx23_audio_causal_conv2d(latent_channels, + base_block_channels, kernel_size = 3L, + causality_axis = causality_axis) + + mid_container <- torch::nn_module( + "ltx23_audio_mid", + initialize = function(channels, causality_axis) { + self$block_1 <- ltx23_audio_resnet_block(channels, channels, + causality_axis = causality_axis) + self$block_2 <- ltx23_audio_resnet_block(channels, channels, + causality_axis = causality_axis) + }, + forward = function(x) { + self$block_2(self$block_1(x)) + } + ) + self$mid <- mid_container(base_block_channels, causality_axis) + + # Stages are stored at their level index; built from the top level + # down, tracking the running channel count like the reference + stages <- vector("list", self$num_resolutions) + block_in <- base_block_channels + for (level in rev(seq_len(self$num_resolutions))) { + block_out <- base_channels * ch_mult[level] + stage_blocks <- list() + for (j in seq_len(num_res_blocks + 1L)) { + stage_blocks[[j]] <- ltx23_audio_resnet_block( + block_in, block_out, causality_axis = causality_axis + ) + block_in <- block_out + } + stage <- torch::nn_module( + "ltx23_audio_up_stage", + initialize = function(blocks, upsample) { + self$block <- torch::nn_module_list(blocks) + if (!is.null(upsample)) self$upsample <- upsample + }, + forward = function(x) { + for (i in seq_along(self$block)) { + x <- self$block[[i]](x) + } + if (!is.null(self$upsample)) x <- self$upsample(x) + x + } + ) + upsample <- if (level != 1L) { + ltx23_audio_upsample(block_in, causality_axis = causality_axis) + } else { + NULL + } + stages[[level]] <- stage(stage_blocks, upsample) + } + self$up <- torch::nn_module_list(stages) + + self$norm_out <- ltx23_per_channel_rms_norm(eps = 1e-6) + self$conv_out <- ltx23_audio_causal_conv2d( + block_in, output_channels, kernel_size = 3L, causality_axis = causality_axis + ) +}, + forward = function(sample) { + frames <- sample$shape[3] + target_frames <- frames * .ltx23_audio_latent_downsample_factor + if (self$causal) { + target_frames <- max(target_frames - (.ltx23_audio_latent_downsample_factor - 1L), 1L) + } + + h <- self$conv_in(sample) + h <- self$mid(h) + for (level in rev(seq_len(self$num_resolutions))) { + h <- self$up[[level]](h) + } + + h <- self$norm_out(h) + h <- torch::nnf_silu(h) + h <- self$conv_out(h) + + # Crop/pad to the target (channels, time, mel bins) + cur_time <- h$shape[3] + cur_freq <- h$shape[4] + h <- h$narrow(2L, 1L, self$out_ch) + h <- h$narrow(3L, 1L, min(cur_time, target_frames)) + h <- h$narrow(4L, 1L, min(cur_freq, self$mel_bins)) + + time_pad <- target_frames - h$shape[3] + freq_pad <- self$mel_bins - h$shape[4] + if (time_pad > 0L || freq_pad > 0L) { + h <- torch::nnf_pad(h, c(0L, max(freq_pad, 0L), 0L, max(time_pad, 0L))) + } + h +} +) + +#' LTX-2.3 audio VAE +#' +#' Encoder + decoder plus the per-channel latent statistics loaded from +#' the checkpoint. Encoding is used for audio-conditioned generation +#' (lip sync); decoding for generated audio. +#' +#' @param base_channels,output_channels,num_res_blocks,latent_channels,ch_mult,causality_axis,mel_bins +#' See \code{\link{ltx23_audio_decoder}}. +#' @param in_channels Integer. Mel input channels (2 = stereo). +#' +#' @export +ltx23_audio_vae <- torch::nn_module( + "ltx23_audio_vae", + initialize = function( + base_channels = 128L, + output_channels = 2L, + num_res_blocks = 2L, + latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), + causality_axis = "height", + mel_bins = 64L, + in_channels = 2L + ) { + self$encoder <- ltx23_audio_encoder(base_channels = base_channels, + in_channels = in_channels, + num_res_blocks = num_res_blocks, + latent_channels = latent_channels, + ch_mult = ch_mult, + causality_axis = causality_axis) + self$decoder <- ltx23_audio_decoder(base_channels = base_channels, + output_channels = output_channels, + num_res_blocks = num_res_blocks, + latent_channels = latent_channels, + ch_mult = ch_mult, + causality_axis = causality_axis, + mel_bins = mel_bins) + # Statistics are stored at base_channels size in the checkpoint; + # only the first latent_channels entries are meaningful + self$latents_mean <- torch::nn_buffer(torch::torch_zeros(base_channels)) + self$latents_std <- torch::nn_buffer(torch::torch_ones(base_channels)) + self$latent_channels <- as.integer(latent_channels) +}, + encode = function(mel) { + moments <- self$encoder(mel) + n <- self$latent_channels + list( + mean = moments$narrow(2L, 1L, n), + logvar = moments$narrow(2L, n + 1L, n) + ) +}, + decode = function(z) { + .ltx23_traced_call(self$decoder, z) +}, + forward = function(z) { + self$decode(z) +} +) + +#' Map an official audio VAE checkpoint key to the R module name +#' +#' @param key Character. Checkpoint key. +#' +#' @return Character. +#' +#' @export +ltx23_map_audio_vae_key <- function(key) { + key <- sub("^audio_vae\\.", "", key) + key <- sub("^per_channel_statistics\\.mean-of-means$", "latents_mean", key) + key <- sub("^per_channel_statistics\\.std-of-means$", "latents_std", key) + key +} diff --git a/R/auto_devices.R b/R/auto_devices.R index c0879bd..6a22057 100644 --- a/R/auto_devices.R +++ b/R/auto_devices.R @@ -38,25 +38,22 @@ #' # Force CPU-only #' devices <- auto_devices("sdxl", strategy = "cpu_only") #' } -auto_devices <- function( - model = "sdxl", - strategy = "auto" -) { - # If gpuctl is available, use it +auto_devices <- function(model = "sdxl", strategy = "auto") { + # If gpuctl is available, use it - if (requireNamespace("gpu.ctl", quietly = TRUE)) { - return(gpu.ctl::recommended_devices(model = model, strategy = strategy)) - } + if (requireNamespace("gpu.ctl", quietly = TRUE)) { + return(gpu.ctl::recommended_devices(model = model, strategy = strategy)) + } - # Fallback when gpuctl not available - if (strategy == "auto") { - message("gpuctl not installed - using unet_gpu strategy as default") - message("Install gpuctl for auto-detection: install.packages('gpuctl')") - strategy <- "unet_gpu" - } + # Fallback when gpuctl not available + if (strategy == "auto") { + message("gpuctl not installed - using unet_gpu strategy as default") + message("Install gpuctl for auto-detection: install.packages('gpuctl')") + strategy <- "unet_gpu" + } - # Build device config manually - .build_fallback_devices(model, strategy) + # Build device config manually + .build_fallback_devices(model, strategy) } #' Build fallback device configuration @@ -65,35 +62,31 @@ auto_devices <- function( #' @param strategy Character. Memory strategy. #' @return Named list of device assignments. #' @keywords internal -.build_fallback_devices <- function( - model, - strategy -) { - # Components by model - components <- list( - sd21 = c("unet", "decoder", "text_encoder", "encoder"), - sdxl = c("unet", "decoder", "text_encoder", "text_encoder2", "encoder") - ) +.build_fallback_devices <- function(model, strategy) { + # Components by model + components <- list( + sd21 = c("unet", "decoder", "text_encoder", "encoder"), + sdxl = c("unet", "decoder", "text_encoder", "text_encoder2", "encoder") + ) - if (!model %in% names(components)) { - stop("Unsupported model: ", model) - } + if (!model %in% names(components)) { + stop("Unsupported model: ", model) + } - comp <- components[[model]] + comp <- components[[model]] - if (strategy == "full_gpu") { - devices <- as.list(rep("cuda", length(comp))) - names(devices) <- comp - } else if (strategy == "unet_gpu") { - devices <- as.list(rep("cpu", length(comp))) - names(devices) <- comp - devices$unet <- "cuda" - } else { - # cpu_only - devices <- as.list(rep("cpu", length(comp))) - names(devices) <- comp - } + if (strategy == "full_gpu") { + devices <- as.list(rep("cuda", length(comp))) + names(devices) <- comp + } else if (strategy == "unet_gpu") { + devices <- as.list(rep("cpu", length(comp))) + names(devices) <- comp + devices$unet <- "cuda" + } else { + # cpu_only + devices <- as.list(rep("cpu", length(comp))) + names(devices) <- comp + } - devices + devices } - diff --git a/R/checkpoint_ltx23.R b/R/checkpoint_ltx23.R new file mode 100644 index 0000000..59d2619 --- /dev/null +++ b/R/checkpoint_ltx23.R @@ -0,0 +1,244 @@ +#' LTX-2.3 Single-File Checkpoint Reader +#' +#' LTX 2.3 checkpoints ship as one safetensors file containing every +#' component (transformer, connectors, video VAE, audio VAE, vocoder), +#' with the model version and full component configuration embedded in +#' the safetensors metadata. These helpers open the file, validate the +#' version, split the key space by component, and stream tensors into +#' R torch modules one at a time so the 46 GB file is never fully +#' materialized in memory. +#' +#' @name checkpoint_ltx23 +NULL + +#' Open an LTX-2.3 checkpoint +#' +#' Opens a single-file LTX checkpoint lazily (header only), validates the +#' \code{model_version} metadata, and parses the embedded component +#' configuration. +#' +#' @param path Path to the checkpoint .safetensors file. +#' @param require_version Character. Required \code{model_version} prefix +#' (default "2.3"). Set to NULL to skip the check. +#' +#' @return An object of class \code{ltx23_checkpoint}: a list with +#' \code{handle} (safetensors reader), \code{keys}, \code{version}, +#' \code{config} (parsed component configs, or NULL), and \code{path}. +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' ckpt <- ltx23_open_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") +#' str(ltx23_split_keys(ckpt$keys), max.level = 1) +#' } +ltx23_open_checkpoint <- function(path, require_version = "2.3") { + if (!requireNamespace("safetensors", quietly = TRUE)) { + stop("The safetensors package is required to read LTX checkpoints.") + } + path <- path.expand(path) + if (!file.exists(path)) { + stop("Checkpoint not found: ", path) + } + + handle <- safetensors::safetensors$new(path, framework = "torch") + + meta <- handle$metadata[["__metadata__"]] + if (is.list(meta) || is.character(meta)) { + version <- meta[["model_version"]] + } else { + version <- NULL + } + + if (!is.null(require_version)) { + if (is.null(version)) { + stop("Checkpoint has no model_version metadata; expected an LTX ", + require_version, " single-file checkpoint: ", path) + } + if (!startsWith(version, require_version)) { + stop( + "Checkpoint model_version is '", version, "' but '", + require_version, "*' is required: ", path + ) + } + } + + config <- NULL + if (!is.null(meta[["config"]])) { + config <- tryCatch( + jsonlite::fromJSON(meta[["config"]], simplifyVector = TRUE), + error = function(e) NULL + ) + } + + keys <- setdiff(handle$keys(), "__metadata__") + + structure( + list( + handle = handle, + keys = keys, + version = version, + config = config, + path = path + ), + class = "ltx23_checkpoint" + ) +} + +#' @export +print.ltx23_checkpoint <- function(x, ...) { + cat("\n") + cat(" path: ", x$path, "\n") + cat(" version: ", x$version %||% "(none)", "\n") + cat(" tensors: ", length(x$keys), "\n") + groups <- ltx23_split_keys(x$keys) + for (g in names(groups)) { + cat(sprintf(" %-10s %d keys\n", g, length(groups[[g]]))) + } + invisible(x) +} + +#' Split checkpoint keys by component +#' +#' Splits the flat key space of an LTX single-file checkpoint into its +#' component groups. Connector tensors live under the +#' \code{model.diffusion_model.} prefix alongside the transformer, plus a +#' top-level \code{text_embedding_projection.} group; both are routed to +#' the \code{connectors} component. +#' +#' @param keys Character vector of checkpoint tensor names. +#' +#' @return Named list of character vectors: \code{dit}, +#' \code{connectors}, \code{vae}, \code{audio_vae}, \code{vocoder}, +#' and \code{other} (anything unrecognized; should be empty). +#' +#' @export +ltx23_split_keys <- function(keys) { + dm_prefix <- "model.diffusion_model." + connector_res <- c( + "^model\\.diffusion_model\\.video_embeddings_connector\\.", + "^model\\.diffusion_model\\.audio_embeddings_connector\\.", + "^text_embedding_projection\\." + ) + + is_connector <- Reduce(`|`, lapply(connector_res, grepl, x = keys)) + is_dit <- startsWith(keys, dm_prefix) & !is_connector + is_vae <- startsWith(keys, "vae.") + is_audio_vae <- startsWith(keys, "audio_vae.") + is_vocoder <- startsWith(keys, "vocoder.") + + claimed <- is_connector | is_dit | is_vae | is_audio_vae | is_vocoder + + list(dit = keys[is_dit], connectors = keys[is_connector], + vae = keys[is_vae], audio_vae = keys[is_audio_vae], + vocoder = keys[is_vocoder], other = keys[!claimed]) +} + +#' Stream a checkpoint key group into a module +#' +#' Reads tensors one at a time from an open checkpoint and copies them +#' into the matching parameters/buffers of \code{module}. Destination +#' names are derived by \code{map_key}; \code{$copy_()} handles any +#' dtype/device conversion, so the module may already live on its target +#' device in its target dtype. +#' +#' @param ckpt An \code{ltx23_checkpoint}. +#' @param keys Character vector of checkpoint keys to load (one group +#' from \code{\link{ltx23_split_keys}}). +#' @param module A torch nn_module to populate. +#' @param map_key Function mapping a checkpoint key to the module's +#' parameter/buffer name, or NA to skip the key deliberately. +#' @param verbose Logical. Report progress and coverage. +#' @param gc_every Integer. Run \code{gc()} after this many tensors. +#' +#' @return Invisibly, a list with \code{unmapped} (checkpoint keys that +#' found no destination), \code{skipped} (keys the mapper declined), +#' and \code{unfilled} (module parameters/buffers never written). +#' A perfectly loaded group has zero \code{unmapped} and zero +#' \code{unfilled}. +#' +#' @export +ltx23_load_group <- function(ckpt, keys, module, map_key = identity, + verbose = TRUE, gc_every = 50L) { + stopifnot(inherits(ckpt, "ltx23_checkpoint")) + + params <- module$named_parameters() + buffers <- module$named_buffers() + dests <- c(params, buffers) + + unmapped <- character(0) + skipped <- character(0) + filled <- character(0) + + n <- length(keys) + torch::with_no_grad({ + for (i in seq_along(keys)) { + key <- keys[[i]] + dest_name <- map_key(key) + + if (length(dest_name) != 1L || is.na(dest_name)) { + skipped <- c(skipped, key) + next + } + dest <- dests[[dest_name]] + if (is.null(dest)) { + unmapped <- c(unmapped, key) + next + } + + value <- ckpt$handle$get_tensor(key) + if (!identical(as.integer(dest$shape), as.integer(value$shape))) { + stop(sprintf( + "Shape mismatch for '%s' -> '%s': checkpoint [%s], module [%s]", + key, dest_name, + paste(value$shape, collapse = ","), + paste(dest$shape, collapse = ",") + )) + } + dest$copy_(value) + filled <- c(filled, dest_name) + rm(value) + + if (i %% gc_every == 0L) { + gc(verbose = FALSE) + if (verbose && n > 500L) { + message(sprintf(" loaded %d/%d tensors", i, n)) + } + } + } + }) + gc(verbose = FALSE) + + unfilled <- setdiff(names(dests), filled) + + if (verbose) { + message(sprintf( + "Loaded %d/%d tensors (%d skipped); unmapped: %d, unfilled: %d", + length(filled), n, length(skipped), + length(unmapped), length(unfilled) + )) + if (length(unmapped)) { + message(" unmapped e.g.: ", + paste(utils::head(unmapped, 3), collapse = ", ")) + } + if (length(unfilled)) { + message(" unfilled e.g.: ", paste(utils::head(unfilled, 3), collapse = ", ")) + } + } + + invisible(list(unmapped = unmapped, skipped = skipped, unfilled = unfilled)) +} + +#' Summarize checkpoint key coverage +#' +#' @param ckpt An \code{ltx23_checkpoint}. +#' +#' @return A data.frame with one row per component group and its key count. +#' +#' @export +ltx23_census <- function(ckpt) { + stopifnot(inherits(ckpt, "ltx23_checkpoint")) + groups <- ltx23_split_keys(ckpt$keys) + data.frame(group = names(groups), + keys = vapply(groups, length, integer(1)), row.names = NULL) +} diff --git a/R/condition_ltx23.R b/R/condition_ltx23.R new file mode 100644 index 0000000..4889d32 --- /dev/null +++ b/R/condition_ltx23.R @@ -0,0 +1,175 @@ +#' LTX-2.3 Prefix Conditioning (image-to-video, video continuation) +#' +#' Fresh R port of the frame-conditioning mechanics from the diffusers +#' reference (Apache-2.0, pipelines/ltx2/pipeline_ltx2_image2video.py +#' and pipeline_ltx2_condition.py), restricted to prefix conditioning +#' at latent index 0 with strength 1: a single start image (i2v) or the +#' leading pixel frames of a previous clip (continuation). Conditioned +#' latent tokens are initialized from the VAE-encoded pixels, see a +#' per-token timestep of zero, and are frozen through the Euler loop. +#' +#' @name condition_ltx23 +NULL + +#' Preprocess an image (or frame stack) for VAE encoding +#' +#' Mirrors the diffusers VideoProcessor: bilinear resize so the shorter +#' relative side matches, center-crop to the exact target, and scale to +#' [-1, 1]. +#' +#' @param x Path to a PNG/JPEG, or an array [H, W, 3] (values in +#' [0, 1]), or a [F, H, W, 3] array of frames. +#' @param height,width Integers. Target size (multiples of 32). +#' +#' @return Float32 tensor [1, 3, F, height, width] in [-1, 1]. +#' +#' @export +ltx23_preprocess_frames <- function(x, height, width) { + if (is.character(x)) { + x <- if (grepl("\\.png$", x, ignore.case = TRUE)) { + png::readPNG(x) + } else { + jpeg::readJPEG(x) + } + } + if (length(dim(x)) == 3L) { + dim(x) <- c(1L, dim(x)) + } + if (dim(x)[4] > 3L) { + # Drop alpha + x <- x[,,, 1:3, drop = FALSE] + } + frames <- torch::torch_tensor(x, dtype = torch::torch_float32()) + # [F, H, W, C] -> [F, C, H, W] + frames <- frames$permute(c(1L, 4L, 2L, 3L)) + + src_h <- frames$shape[3] + src_w <- frames$shape[4] + scale <- max(height / src_h, width / src_w) + new_h <- as.integer(round(src_h * scale)) + new_w <- as.integer(round(src_w * scale)) + frames <- torch::nnf_interpolate(frames, size = c(new_h, new_w), + mode = "bilinear", align_corners = FALSE) + top <- (new_h - height) %/% 2L + left <- (new_w - width) %/% 2L + frames <- frames$narrow(3L, top + 1L, height)$narrow(4L, left + 1L, width) + + # [F, C, H, W] -> [1, C, F, H, W], [0, 1] -> [-1, 1] + frames$permute(c(2L, 1L, 3L, 4L))$unsqueeze(1L)$mul(2)$sub(1) +} + +#' Read the trailing frames of a video file +#' +#' Extracts the last \code{n} frames of an MP4 (via \code{av}) for use +#' as continuation conditioning. +#' +#' @param path Video file. +#' @param n Integer. Trailing frame count. +#' +#' @return Array [n, H, W, 3] in [0, 1]. +#' +#' @export +ltx23_read_tail_frames <- function(path, n = 9L) { + if (!requireNamespace("av", quietly = TRUE)) { + stop("Reading video frames requires the 'av' package.") + } + info <- av::av_media_info(path) + total <- info$video$frames + if (is.null(total) || is.na(total)) { + # Some containers omit the frame count; derive from duration + total <- floor(info$duration * info$video$framerate) + } + out_dir <- tempfile("ltx23_tail_") + dir.create(out_dir) + on.exit(unlink(out_dir, recursive = TRUE), add = TRUE) + av::av_video_images(path, destdir = out_dir, format = "png") + files <- sort(list.files(out_dir, pattern = "\\.png$", full.names = TRUE)) + files <- utils::tail(files, n) + frames <- lapply(files, png::readPNG) + arr <- array(0, dim = c(length(frames), dim(frames[[1]])[1], + dim(frames[[1]])[2], 3L)) + for (i in seq_along(frames)) { + f <- frames[[i]] + arr[i,,,] <- f[,, 1:3] + } + arr +} + +#' Encode pixel frames to normalized video latents +#' +#' VAE encode in "argmax" mode (the distribution mean), then normalize +#' with the checkpoint's per-channel statistics — the exact inverse of +#' the decode path. +#' +#' @param vae An \code{ltx23_video_vae}. +#' @param frames Tensor [1, 3, F, H, W] in [-1, 1] (see +#' \code{\link{ltx23_preprocess_frames}}). +#' +#' @return Normalized latents [1, 128, F', H/32, W/32] (float32). +#' +#' @export +ltx23_encode_video_frames <- function(vae, frames) { + enc_dtype <- vae$encoder$conv_in$conv$weight$dtype + enc_device <- vae$encoder$conv_in$conv$weight$device + moments <- torch::with_no_grad( + vae$encode(frames$to(device = enc_device, dtype = enc_dtype)) + ) + latents <- moments$mean$to(dtype = torch::torch_float32()) + ltx23_normalize_latents(latents, vae$latents_mean$to(dtype = torch::torch_float32()), + vae$latents_std$to(dtype = torch::torch_float32())) +} + +#' Build conditioned initial latents and the conditioning mask +#' +#' i2v (\code{cond_latents} has one latent frame): the encoded frame is +#' repeated across all latent frames and only latent frame 0 is marked +#' conditioned. Continuation (k latent frames): the prefix tokens are +#' replaced and marked. Unconditioned positions start as pure noise. +#' +#' @param cond_latents Normalized condition latents +#' [1, 128, k, H', W'] from \code{\link{ltx23_encode_video_frames}}. +#' @param latent_frames,latent_height,latent_width Integers. Full +#' latent geometry of the generation. +#' @param noise Tensor [1, 128, F', H', W'] of standard noise (caller +#' provides so seeding stays in one place). +#' @param cond_noise_scale Numeric. Optional partial noising of the +#' conditioned tokens (diffusers \code{noise_scale}, default 0). +#' +#' @return list(latents [1, S, 128] float32 packed, +#' conditioning_mask [1, S] float32 packed). +#' +#' @export +ltx23_prepare_conditioned_latents <- function(cond_latents, latent_frames, + latent_height, latent_width, + noise, cond_noise_scale = 0) { + k <- cond_latents$shape[3] + stopifnot(k <= latent_frames, cond_latents$shape[4] == latent_height, + cond_latents$shape[5] == latent_width) + + init <- if (k == 1L) { + cond_latents$`repeat`(c(1L, 1L, latent_frames, 1L, 1L)) + } else { + torch::torch_cat(list( + cond_latents, + noise$narrow(3L, k + 1L, latent_frames - k) + ), dim = 3L) + } + + mask <- torch::torch_zeros(1L, 1L, latent_frames, latent_height, + latent_width, dtype = torch::torch_float32(), + device = noise$device) + mask$narrow(3L, 1L, k)$fill_(1) + + latents <- init$to(device = noise$device) * mask + noise * (1 - mask) + if (cond_noise_scale > 0) { + # Partially re-noise the conditioned tokens (reference + # _create_noised_state with per-token scale) + scale <- mask * cond_noise_scale + latents <- noise * scale + latents * (1 - scale) + } + + list( + latents = ltx23_pack_video_latents(latents), + conditioning_mask = ltx23_pack_video_latents(mask)$squeeze(3L) + ) +} diff --git a/R/connectors_ltx23.R b/R/connectors_ltx23.R new file mode 100644 index 0000000..b439c7c --- /dev/null +++ b/R/connectors_ltx23.R @@ -0,0 +1,392 @@ +#' LTX-2.3 Text Embedding Connectors +#' +#' Fresh R port of the LTX text connectors from the diffusers reference +#' (Apache-2.0, src/diffusers/pipelines/ltx2/connectors.py). The +#' connectors take raw stacked per-layer Gemma3 hidden states +#' [batch, seq, caption_channels, num_layers + 1], normalize and project +#' them per modality, replace padding with learnable registers, and run a +#' small 1D transformer per modality to produce the DiT text embeddings. +#' +#' @name connectors_ltx23 +NULL + +# Largest finite value for a float dtype; R torch's torch_finfo rejects +# bfloat16, so fall back to the known constants +.ltx23_finfo_max <- function(dtype) { + tryCatch( + torch::torch_finfo(dtype)$max, + error = function(e) { + switch(dtype$.type(), BFloat16 = 3.3895313892515355e38, Half = 65504, + 3.402823466e38) + } + ) +} + +#' Per-token RMS normalization over the channel axis +#' +#' @param x Tensor [B, S, C, L] of stacked per-layer hidden states. +#' @param eps Numeric. Stability epsilon. +#' +#' @return Tensor of the same shape. +#' +#' @export +ltx23_per_token_rms_norm <- function(x, eps = 1e-6) { + variance <- torch::torch_mean(x ^ 2, dim = 3L, keepdim = TRUE) + x * torch::torch_rsqrt(variance$add(eps)) +} + +#' 1D rotary embeddings for the text connectors +#' +#' @param dim Integer. Rotary dimension (connector inner dim). +#' @param base_seq_len Integer. Base sequence length for normalization. +#' @param theta Numeric. RoPE theta. +#' @param double_precision Logical. Compute base frequencies in float64. +#' @param rope_type "split" (LTX-2.3) or "interleaved". +#' @param num_attention_heads Integer. For the split per-head layout. +#' +#' @export +ltx23_rotary_pos_embed_1d <- torch::nn_module( + "ltx23_rotary_pos_embed_1d", + initialize = function( + dim, + base_seq_len = 4096L, + theta = 10000.0, + double_precision = TRUE, + rope_type = "split", + num_attention_heads = 32L + ) { + if (!rope_type %in% c("interleaved", "split")) { + stop("rope_type must be 'interleaved' or 'split', got: ", rope_type) + } + self$dim <- as.integer(dim) + self$base_seq_len <- base_seq_len + self$theta <- theta + self$double_precision <- double_precision + self$rope_type <- rope_type + self$num_attention_heads <- as.integer(num_attention_heads) +}, + forward = function(batch_size, pos, device) { + grid_1d <- torch::torch_arange(start = 0, end = pos - 1, + dtype = torch::torch_float32(), + device = device) / self$base_seq_len + grid <- grid_1d$unsqueeze(1L)$`repeat`(c(batch_size, 1L)) # [B, S] + + freqs_dtype <- if (self$double_precision) torch::torch_float64() else torch::torch_float32() + pow_indices <- torch::torch_pow( + self$theta, + torch::torch_linspace( + start = 0.0, end = 1.0, steps = self$dim %/% 2L, + dtype = freqs_dtype, device = device + ) + ) + freqs <- (pow_indices * pi / 2.0)$to(dtype = torch::torch_float32()) + freqs <- (grid$unsqueeze(-1L) * 2 - 1) * freqs # [B, S, dim/2] + + if (self$rope_type == "interleaved") { + cos_freqs <- freqs$cos()$repeat_interleave(2L, dim = -1L) + sin_freqs <- freqs$sin()$repeat_interleave(2L, dim = -1L) + } else { + cos_freqs <- freqs$cos() + sin_freqs <- freqs$sin() + b <- cos_freqs$shape[1] + t <- cos_freqs$shape[2] + cos_freqs <- cos_freqs$reshape(c(b, t, self$num_attention_heads, -1L))$transpose(2L, 3L) + sin_freqs <- sin_freqs$reshape(c(b, t, self$num_attention_heads, -1L))$transpose(2L, 3L) + } + list(cos_freqs, sin_freqs) +} +) + +# Pre-norm 1D transformer block: RMS norm -> attention -> RMS norm -> FF, +# both with plain residual connections. +ltx23_transformer_block_1d <- torch::nn_module( + "ltx23_transformer_block_1d", + initialize = function( + dim, + num_attention_heads, + attention_head_dim, + eps = 1e-6, + rope_type = "split", + apply_gated_attention = TRUE + ) { + self$norm1 <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = FALSE) + self$attn1 <- ltx23_attention(query_dim = dim, + heads = num_attention_heads, + dim_head = attention_head_dim, + rope_type = rope_type, + apply_gated_attention = apply_gated_attention) + self$norm2 <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = FALSE) + self$ff <- ltx23_feed_forward(dim) +}, + forward = function(hidden_states, attention_mask = NULL, rotary_emb = NULL) { + norm_hidden_states <- self$norm1(hidden_states) + hidden_states <- hidden_states + self$attn1( + norm_hidden_states, attention_mask = attention_mask, query_rotary_emb = rotary_emb + ) + norm_hidden_states <- self$norm2(hidden_states) + hidden_states + self$ff(norm_hidden_states) +} +) + +#' 1D connector transformer +#' +#' Replaces padded positions with learnable registers (valid tokens are +#' front-aligned in their original order; the tail is filled with +#' registers indexed by absolute position, after which the attention mask +#' is cleared), then runs 1D transformer blocks with rotary embeddings. +#' +#' @param num_attention_heads,attention_head_dim,num_layers Transformer shape. +#' @param num_learnable_registers Integer or NULL. Register count (the +#' sequence length must be divisible by it). +#' @param rope_base_seq_len,rope_theta,rope_double_precision,rope_type RoPE config. +#' @param eps Numeric. Norm epsilon. +#' @param gated_attention Logical. Per-head attention output gates. +#' +#' @export +ltx23_connector_transformer_1d <- torch::nn_module( + "ltx23_connector_transformer_1d", + initialize = function( + num_attention_heads = 32L, + attention_head_dim = 128L, + num_layers = 8L, + num_learnable_registers = 128L, + rope_base_seq_len = 4096L, + rope_theta = 10000.0, + rope_double_precision = TRUE, + eps = 1e-6, + rope_type = "split", + gated_attention = TRUE + ) { + self$num_attention_heads <- as.integer(num_attention_heads) + self$inner_dim <- as.integer(num_attention_heads * attention_head_dim) + self$num_learnable_registers <- num_learnable_registers + + if (!is.null(num_learnable_registers)) { + self$learnable_registers <- torch::nn_parameter( + torch::torch_rand(num_learnable_registers, self$inner_dim) * 2 - 1 + ) + } + + self$rope <- ltx23_rotary_pos_embed_1d(self$inner_dim, + base_seq_len = rope_base_seq_len, theta = rope_theta, + double_precision = rope_double_precision, rope_type = rope_type, + num_attention_heads = num_attention_heads) + + self$transformer_blocks <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { + ltx23_transformer_block_1d( + dim = self$inner_dim, + num_attention_heads = num_attention_heads, + attention_head_dim = attention_head_dim, + eps = eps, + rope_type = rope_type, + apply_gated_attention = gated_attention + ) + })) + + self$norm_out <- ltx23_rms_norm(self$inner_dim, eps = eps, elementwise_affine = FALSE) +}, + forward = function(hidden_states, attention_mask = NULL, + attn_mask_binarize_threshold = -9000.0) { + batch_size <- hidden_states$shape[1] + seq_len <- hidden_states$shape[2] + + if (!is.null(self$num_learnable_registers)) { + if (seq_len %% self$num_learnable_registers != 0L) { + stop( + "Sequence length ", seq_len, " must be divisible by the number of ", + "learnable registers ", self$num_learnable_registers + ) + } + num_repeats <- seq_len %/% self$num_learnable_registers + registers <- self$learnable_registers$unsqueeze(1L)$ + expand(c(num_repeats, -1L, -1L))$reshape(c(seq_len, -1L)) # [S, D] + + binary_attn_mask <- (attention_mask >= attn_mask_binarize_threshold)$to( + dtype = torch::torch_int() + ) + if (binary_attn_mask$ndim == 4L) { + binary_attn_mask <- binary_attn_mask$squeeze(2L)$squeeze(2L) # [B, 1, 1, S] -> [B, S] + } + + # Front-align valid tokens (stable sort keeps their order), fill the + # tail with registers indexed by absolute position + order <- torch::torch_sort(1L - binary_attn_mask, dim = 2L, stable = TRUE)[[2]] # [B, S] + front_aligned <- torch::torch_gather( + hidden_states, 2L, + order$unsqueeze(-1L)$expand(c(-1L, -1L, hidden_states$shape[3])) + ) + num_valid <- binary_attn_mask$sum(dim = 2L, keepdim = TRUE) # [B, 1] + positions <- torch::torch_arange( + start = 0, end = seq_len - 1, device = hidden_states$device + )$unsqueeze(1L) # [1, S] + front_mask <- (positions < num_valid)$unsqueeze(-1L) # [B, S, 1] + registers_expanded <- registers$unsqueeze(1L)$expand(c(batch_size, -1L, -1L)) + hidden_states <- torch::torch_where( + front_mask, front_aligned, registers_expanded$to(dtype = hidden_states$dtype) + ) + + # All positions are valid once registers replace the padding + attention_mask <- torch::torch_zeros_like(attention_mask) + } + + rotary_emb <- self$rope(batch_size, seq_len, device = hidden_states$device) + + for (i in seq_along(self$transformer_blocks)) { + hidden_states <- self$transformer_blocks[[i]]( + hidden_states, attention_mask = attention_mask, rotary_emb = rotary_emb + ) + } + hidden_states <- self$norm_out(hidden_states) + + list(hidden_states, attention_mask) +} +) + +#' LTX-2.3 text connectors +#' +#' Takes raw stacked per-layer text encoder hidden states and produces +#' the video and audio text embeddings for the DiT: per-token RMS norm, +#' per-modality sqrt(dim ratio) rescaling and projection, then a +#' per-modality 1D connector transformer. +#' +#' @param caption_channels Integer. Text encoder hidden size (3840 for +#' Gemma3-12B). +#' @param text_proj_in_factor Integer. Number of stacked hidden states +#' (num_layers + 1 = 49 for Gemma3-12B). +#' @param video_connector_num_attention_heads,video_connector_attention_head_dim,video_connector_num_layers +#' Video connector shape (LTX-2.3: 32 x 128, 8 layers). +#' @param video_connector_num_learnable_registers Integer or NULL. +#' @param video_gated_attn Logical. +#' @param audio_connector_num_attention_heads,audio_connector_attention_head_dim,audio_connector_num_layers +#' Audio connector shape (LTX-2.3: 32 x 64, 8 layers). +#' @param audio_connector_num_learnable_registers Integer or NULL. +#' @param audio_gated_attn Logical. +#' @param connector_rope_base_seq_len,rope_theta,rope_double_precision,rope_type RoPE config. +#' @param video_hidden_dim,audio_hidden_dim Integers. Projection targets +#' (DiT inner dims: 4096 / 2048). +#' @param proj_bias Logical. Projection bias (TRUE for LTX-2.3). +#' +#' @export +ltx23_text_connectors <- torch::nn_module( + "ltx23_text_connectors", + initialize = function( + caption_channels = 3840L, + text_proj_in_factor = 49L, + video_connector_num_attention_heads = 32L, + video_connector_attention_head_dim = 128L, + video_connector_num_layers = 8L, + video_connector_num_learnable_registers = 128L, + video_gated_attn = TRUE, + audio_connector_num_attention_heads = 32L, + audio_connector_attention_head_dim = 64L, + audio_connector_num_layers = 8L, + audio_connector_num_learnable_registers = 128L, + audio_gated_attn = TRUE, + connector_rope_base_seq_len = 4096L, + rope_theta = 10000.0, + rope_double_precision = TRUE, + rope_type = "split", + video_hidden_dim = 4096L, + audio_hidden_dim = 2048L, + proj_bias = TRUE + ) { + self$caption_channels <- as.integer(caption_channels) + self$video_hidden_dim <- as.integer(video_hidden_dim) + self$audio_hidden_dim <- as.integer(audio_hidden_dim) + + text_encoder_dim <- caption_channels * text_proj_in_factor + self$video_text_proj_in <- torch::nn_linear(text_encoder_dim, + video_hidden_dim, bias = proj_bias) + self$audio_text_proj_in <- torch::nn_linear(text_encoder_dim, audio_hidden_dim, bias = proj_bias) + + self$video_connector <- ltx23_connector_transformer_1d( + num_attention_heads = video_connector_num_attention_heads, + attention_head_dim = video_connector_attention_head_dim, + num_layers = video_connector_num_layers, + num_learnable_registers = video_connector_num_learnable_registers, + rope_base_seq_len = connector_rope_base_seq_len, + rope_theta = rope_theta, + rope_double_precision = rope_double_precision, + rope_type = rope_type, + gated_attention = video_gated_attn + ) + self$audio_connector <- ltx23_connector_transformer_1d( + num_attention_heads = audio_connector_num_attention_heads, + attention_head_dim = audio_connector_attention_head_dim, + num_layers = audio_connector_num_layers, + num_learnable_registers = audio_connector_num_learnable_registers, + rope_base_seq_len = connector_rope_base_seq_len, + rope_theta = rope_theta, + rope_double_precision = rope_double_precision, + rope_type = rope_type, + gated_attention = audio_gated_attn + ) +}, + forward = function(text_encoder_hidden_states, attention_mask) { + if (text_encoder_hidden_states$ndim == 3L) { + text_encoder_hidden_states <- text_encoder_hidden_states$unflatten( + 3L, c(self$caption_channels, -1L) + ) + } + + norm_states <- ltx23_per_token_rms_norm(text_encoder_hidden_states) + norm_states <- norm_states$flatten(start_dim = 3L, end_dim = 4L) + bool_mask <- attention_mask$to(dtype = torch::torch_bool())$unsqueeze(-1L) + norm_states <- torch::torch_where( + bool_mask, norm_states, torch::torch_zeros_like(norm_states) + ) + + # Rescale per modality by sqrt(target_dim / caption_channels) + video_norm <- norm_states$mul(sqrt(self$video_hidden_dim / self$caption_channels)) + audio_norm <- norm_states$mul(sqrt(self$audio_hidden_dim / self$caption_channels)) + + video_proj <- self$video_text_proj_in(video_norm) + audio_proj <- self$audio_text_proj_in(audio_norm) + + # Multiplicative [B, S] mask -> additive [B, 1, 1, S] with -finfo max + text_dtype <- video_proj$dtype + add_mask <- (attention_mask$to(dtype = torch::torch_int64()) - 1L)$to(dtype = text_dtype) + add_mask <- add_mask$reshape(c(add_mask$shape[1], 1L, -1L, add_mask$shape[2])) + add_mask <- add_mask$mul(.ltx23_finfo_max(text_dtype)) + + video_res <- self$video_connector(video_proj, add_mask) + video_text_embedding <- video_res[[1]] + video_attn_mask <- video_res[[2]] + + # Post-connector mask back to binary; mask the video embedding + binary_mask <- (video_attn_mask < 1e-6)$to(dtype = torch::torch_int64()) + binary_mask <- binary_mask$reshape(c( + video_text_embedding$shape[1], video_text_embedding$shape[2], 1L + )) + video_text_embedding <- video_text_embedding * binary_mask + + audio_res <- self$audio_connector(audio_proj, add_mask) + audio_text_embedding <- audio_res[[1]] + + list( + video_text_embedding = video_text_embedding, + audio_text_embedding = audio_text_embedding, + attention_mask = binary_mask$squeeze(-1L) + ) +} +) + +#' Map an official connectors checkpoint key to the R module name +#' +#' @param key Character. Checkpoint key. +#' +#' @return Character. Module parameter name. +#' +#' @export +ltx23_map_connector_key <- function(key) { + key <- sub("^model\\.diffusion_model\\.", "", key) + key <- gsub("video_embeddings_connector", "video_connector", key, + fixed = TRUE) + key <- gsub("audio_embeddings_connector", "audio_connector", key, fixed = TRUE) + key <- gsub("transformer_1d_blocks", "transformer_blocks", key, fixed = TRUE) + key <- gsub("text_embedding_projection.video_aggregate_embed", "video_text_proj_in", key, fixed = TRUE) + key <- gsub("text_embedding_projection.audio_aggregate_embed", "audio_text_proj_in", key, fixed = TRUE) + key <- gsub("q_norm", "norm_q", key, fixed = TRUE) + key <- gsub("k_norm", "norm_k", key, fixed = TRUE) + key +} diff --git a/R/ddim_scheduler.R b/R/ddim_scheduler.R index 386daf0..eda8a8c 100644 --- a/R/ddim_scheduler.R +++ b/R/ddim_scheduler.R @@ -54,46 +54,41 @@ #' ) #' } #' @export -ddim_scheduler_create <- function( - num_train_timesteps = 1000, - num_inference_steps = 50, - eta = 0, - beta_schedule = c("linear", "scaled_linear", "cosine"), - beta_start = 0.00085, - beta_end = 0.012, - rescale_betas_zero_snr = FALSE, - dtype = torch::torch_float32(), - device = c(torch::torch_device("cpu"), torch::torch_device("cuda")) -) { - betas <- switch(beta_schedule, - "linear" = seq(beta_start, beta_end, - length.out = num_train_timesteps), - "scaled_linear" = seq(sqrt(beta_start), sqrt(beta_end), - length.out = num_train_timesteps) ^ 2, - "cosine" = NULL) - if (rescale_betas_zero_snr) { - betas <- rescale_zero_terminal_snr(betas) - } - betas <- torch::torch_tensor(betas, dtype = dtype, device = device) - alphas <- 1 - betas - alphas <- torch::torch_tensor(alphas, dtype = dtype, device = device) - alphas_cumprod <- torch::torch_cumprod(alphas, dim = 1, dtype = dtype) - alphas_cumprod_prev <- torch::torch_cat(list(torch::torch_ones(1, device = device), - alphas_cumprod[1:(length(alphas_cumprod) - 1)])) +ddim_scheduler_create <- function(num_train_timesteps = 1000, + num_inference_steps = 50, eta = 0, + beta_schedule = c("linear", "scaled_linear", "cosine"), + beta_start = 0.00085, beta_end = 0.012, + rescale_betas_zero_snr = FALSE, + dtype = torch::torch_float32(), + device = c(torch::torch_device("cpu"), torch::torch_device("cuda"))) { + betas <- switch(beta_schedule, + "linear" = seq(beta_start, beta_end, length.out = num_train_timesteps), + "scaled_linear" = seq(sqrt(beta_start), sqrt(beta_end), + length.out = num_train_timesteps) ^ 2, + "cosine" = NULL) + if (rescale_betas_zero_snr) { + betas <- rescale_zero_terminal_snr(betas) + } + betas <- torch::torch_tensor(betas, dtype = dtype, device = device) + alphas <- 1 - betas + alphas <- torch::torch_tensor(alphas, dtype = dtype, device = device) + alphas_cumprod <- torch::torch_cumprod(alphas, dim = 1, dtype = dtype) + alphas_cumprod_prev <- torch::torch_cat(list(torch::torch_ones(1, device = device), + alphas_cumprod[1:(length(alphas_cumprod) - 1)])) - step_ratio <- num_train_timesteps %/% num_inference_steps - # Could add 2 here instead of 1 and remove it from timestep_index and prev_timestep_index below - timesteps <- as.integer(rev(round((0:(num_inference_steps - 1)) * step_ratio) + 1)) - # timesteps <- as.integer(seq(num_train_timesteps, 0, length.out = num_inference_steps + 1)[-1] + 1) + step_ratio <- num_train_timesteps %/% num_inference_steps + # Could add 2 here instead of 1 and remove it from timestep_index and prev_timestep_index below + timesteps <- as.integer(rev(round((0:(num_inference_steps - 1)) * step_ratio) + 1)) + # timesteps <- as.integer(seq(num_train_timesteps, 0, length.out = num_inference_steps + 1)[-1] + 1) - list( - betas = betas, - alphas = alphas, - alphas_cumprod = alphas_cumprod, - alphas_cumprod_prev = alphas_cumprod_prev, - timesteps = timesteps, - eta = eta - ) + list( + betas = betas, + alphas = alphas, + alphas_cumprod = alphas_cumprod, + alphas_cumprod_prev = alphas_cumprod_prev, + timesteps = timesteps, + eta = eta + ) } #' Perform a DDIM scheduler step @@ -171,131 +166,121 @@ ddim_scheduler_create <- function( #' prediction_type = "epsilon") #' } #' @export -ddim_scheduler_step <- function( - model_output, - timestep, - sample, - schedule, - eta = 0, - use_clipped_model_output = FALSE, - thresholding = FALSE, - generator = NULL, - variance_noise = NULL, - clip_sample = FALSE, - set_alpha_to_one = FALSE, - prediction_type = c("epsilon", "sample", "v_prediction"), - dtype = torch::torch_float32(), - device = "cpu" -) { - - # 1. get previous step value (= timestep + 1); i.e. python-indexing - timestep_index <- torch::torch_tensor(timestep + 1, - dtype = torch::torch_long(), - device = torch::torch_device(device)) - if (as.numeric(timestep_index) <= 2) { #schedule$timesteps[1]) { - prev_timestep <- 1#length(schedule$alphas) - prev_timestep_index <- torch::torch_tensor(prev_timestep, - dtype = torch::torch_long(), - device = torch::torch_device(device)) - } else { - prev_timestep <- schedule$timesteps[which(as.logical(schedule$timesteps == (timestep))) + 1] - prev_timestep_index <- torch::torch_tensor(prev_timestep + 1, - dtype = torch::torch_long(), - device = torch::torch_device(device)) - } +ddim_scheduler_step <- function(model_output, timestep, sample, schedule, + eta = 0, use_clipped_model_output = FALSE, + thresholding = FALSE, generator = NULL, + variance_noise = NULL, clip_sample = FALSE, + set_alpha_to_one = FALSE, + prediction_type = c("epsilon", "sample", "v_prediction"), + dtype = torch::torch_float32(), + device = "cpu") { + # 1. get previous step value (= timestep + 1); i.e. python-indexing + timestep_index <- torch::torch_tensor(timestep + 1, + dtype = torch::torch_long(), device = torch::torch_device(device)) + if (as.numeric(timestep_index) <= 2) { #schedule$timesteps[1]) { + prev_timestep <- 1 #length(schedule$alphas) + prev_timestep_index <- torch::torch_tensor(prev_timestep, + dtype = torch::torch_long(), + device = torch::torch_device(device)) + } else { + prev_timestep <- schedule$timesteps[which(as.logical(schedule$timesteps == (timestep))) + 1] + prev_timestep_index <- torch::torch_tensor(prev_timestep + 1, + dtype = torch::torch_long(), + device = torch::torch_device(device)) + } - # 2. compute alphas, betas - alpha_prod_t <- schedule$alphas_cumprod[timestep_index] - alpha_prod_t_prev <- schedule$alphas_cumprod[prev_timestep_index] - # alpha_prod_t <- alpha_prod_t + 1e-7 # Prevent division by zero - # alpha_prod_t_prev <- alpha_prod_t_prev + 1e-7 - if (set_alpha_to_one & prev_timestep == 1) { - alpha_prod_t_prev <- torch::torch_tensor(1.0, - dtype = dtype, - device = torch::torch_device(device)) - } else { + # 2. compute alphas, betas + alpha_prod_t <- schedule$alphas_cumprod[timestep_index] alpha_prod_t_prev <- schedule$alphas_cumprod[prev_timestep_index] - alpha_prod_t_prev <- alpha_prod_t_prev$to(dtype = dtype, device = torch::torch_device(device)) - } + # alpha_prod_t <- alpha_prod_t + 1e-7 # Prevent division by zero + # alpha_prod_t_prev <- alpha_prod_t_prev + 1e-7 + if (set_alpha_to_one & prev_timestep == 1) { + alpha_prod_t_prev <- torch::torch_tensor(1.0, + dtype = dtype, + device = torch::torch_device(device)) + } else { + alpha_prod_t_prev <- schedule$alphas_cumprod[prev_timestep_index] + alpha_prod_t_prev <- alpha_prod_t_prev$to(dtype = dtype, device = torch::torch_device(device)) + } - beta_prod_t <- 1 - alpha_prod_t - beta_prod_t <- beta_prod_t$to(dtype = dtype, device = torch::torch_device(device)) - beta_prod_t_prev = 1 - alpha_prod_t_prev - beta_prod_t_prev <- beta_prod_t_prev$to(dtype = dtype, device = torch::torch_device(device)) + beta_prod_t <- 1 - alpha_prod_t + beta_prod_t <- beta_prod_t$to(dtype = dtype, device = torch::torch_device(device)) + beta_prod_t_prev <- 1 - alpha_prod_t_prev + beta_prod_t_prev <- beta_prod_t_prev$to(dtype = dtype, device = torch::torch_device(device)) - # 3. Handle different prediction types (epsilon or v-prediction) - preds <- switch(prediction_type, - "epsilon" = list(pred_original_sample = (sample - beta_prod_t ^ (0.5) * model_output) / alpha_prod_t ^ (0.5), - pred_epsilon = model_output), - "sample" = list(pred_original_sample = model_output, - pred_epsilon = (sample - alpha_prod_t ^ (0.5) * model_output) / beta_prod_t ^ (0.5)), - "v_prediction" = list(pred_original_sample = (alpha_prod_t ^ 0.5) * sample - (beta_prod_t ^ 0.5) * model_output, - pred_epsilon = (alpha_prod_t ^ 0.5) * model_output + (beta_prod_t ^ 0.5) * sample)) - pred_original_sample <- preds$pred_original_sample$to(dtype = dtype, device = torch::torch_device(device)) - pred_epsilon <- preds$pred_epsilon$to(dtype = dtype, device = torch::torch_device(device)) - # if self.config.prediction_type == "epsilon": - # pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) - # pred_epsilon = model_output + # 3. Handle different prediction types (epsilon or v-prediction) + preds <- switch(prediction_type, + "epsilon" = list(pred_original_sample = (sample - beta_prod_t ^ (0.5) * model_output) / alpha_prod_t ^ (0.5), + pred_epsilon = model_output), + "sample" = list(pred_original_sample = model_output, + pred_epsilon = (sample - alpha_prod_t ^ (0.5) * model_output) / beta_prod_t ^ (0.5)), + "v_prediction" = list(pred_original_sample = (alpha_prod_t ^ 0.5) * sample - (beta_prod_t ^ 0.5) * model_output, + pred_epsilon = (alpha_prod_t ^ 0.5) * model_output + (beta_prod_t ^ 0.5) * sample)) + pred_original_sample <- preds$pred_original_sample$to(dtype = dtype, device = torch::torch_device(device)) + pred_epsilon <- preds$pred_epsilon$to(dtype = dtype, device = torch::torch_device(device)) + # if self.config.prediction_type == "epsilon": + # pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) + # pred_epsilon = model_output - # 4. Clip or threshold "predicted x_0" - # Check thresholding sample code - if (thresholding) { - flat <- pred_original_sample$view(c(pred_original_sample$size(1), - 1)) - abs_flat <- torch::torch_abs(flat, device = torch::torch_device(device)) - s <- torch::torch_quantile(abs_flat, - torch::torch_tensor(0.995, device = torch::torch_device(device)), - dim = 2, - device = torch::torch_device(device)) # dim=2 because batch is dim=1 in R torch - s <- torch::torch_max(s, device = torch::torch_device(device)) - s <- s$view(c(pred_original_sample$size(1), 1, 1, 1)) - pred_original_sample <- torch::torch_clamp(pred_original_sample, - min = - s, max = s, - device = torch::torch_device(device)) / s - } else { - if (clip_sample) { - # Clip to [-1, 1] range - pred_original_sample <- torch::torch_clamp(pred_original_sample, - min = - 1, max = 1, - device = torch::torch_device(device)) + # 4. Clip or threshold "predicted x_0" + # Check thresholding sample code + if (thresholding) { + flat <- pred_original_sample$view(c(pred_original_sample$size(1), -1)) + abs_flat <- torch::torch_abs(flat, device = torch::torch_device(device)) + s <- torch::torch_quantile(abs_flat, + torch::torch_tensor(0.995, device = torch::torch_device(device)), + dim = 2, + device = torch::torch_device(device)) # dim=2 because batch is dim=1 in R torch + s <- torch::torch_max(s, device = torch::torch_device(device)) + s <- s$view(c(pred_original_sample$size(1), 1, 1, 1)) + pred_original_sample <- torch::torch_clamp(pred_original_sample, + min = -s, max = s, + device = torch::torch_device(device)) / s + } else { + if (clip_sample) { + # Clip to [-1, 1] range + pred_original_sample <- torch::torch_clamp(pred_original_sample, + min = -1, max = 1, + device = torch::torch_device(device)) + } } - } - # 5. compute variance - variance <- (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) - std_dev_t = eta * sqrt(variance) + # 5. compute variance + variance <- (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) + std_dev_t <- eta * sqrt(variance) - # Re-compute model_output (noise prediction) from clipped pred_original_sample if needed - if (use_clipped_model_output) { - if (prediction_type == "v_prediction") { - # Recompute v-prediction from clipped pred_original_sample - model_output <- sqrt(alpha_prod_t) * pred_epsilon - sqrt(beta_prod_t) * pred_original_sample - } else { - # Recompute epsilon from clipped pred_original_sample - model_output <- (sample - sqrt(alpha_prod_t) * pred_original_sample) / sqrt(beta_prod_t) + # Re-compute model_output (noise prediction) from clipped pred_original_sample if needed + if (use_clipped_model_output) { + if (prediction_type == "v_prediction") { + # Recompute v-prediction from clipped pred_original_sample + model_output <- sqrt(alpha_prod_t) * pred_epsilon - sqrt(beta_prod_t) * pred_original_sample + } else { + # Recompute epsilon from clipped pred_original_sample + model_output <- (sample - sqrt(alpha_prod_t) * pred_original_sample) / sqrt(beta_prod_t) + } } - } - # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf - pred_sample_direction <- sqrt(1 - alpha_prod_t_prev - std_dev_t ^ 2) * pred_epsilon + # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + pred_sample_direction <- sqrt(1 - alpha_prod_t_prev - std_dev_t ^ 2) * pred_epsilon - # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf - prev_sample <- sqrt(alpha_prod_t_prev) * pred_original_sample + pred_sample_direction + # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + prev_sample <- sqrt(alpha_prod_t_prev) * pred_original_sample + pred_sample_direction - if (eta > 0) { - if (!is.null(variance_noise) && !is.null(generator)) { - stop("Provide either `generator` or `variance_noise`, not both.") + if (eta > 0) { + if (!is.null(variance_noise) && !is.null(generator)) { + stop("Provide either `generator` or `variance_noise`, not both.") + } + if (is.null(variance_noise)) { + variance_noise <- torch::torch_randn(model_output$shape, + generator = generator, + device = torch::torch_device(device), + dtype = dtype) + } + variance <- std_dev_t * eta * variance_noise + prev_sample <- prev_sample + variance } - if (is.null(variance_noise)) { - variance_noise <- torch::torch_randn(model_output$shape, - generator = generator, - device = torch::torch_device(device), - dtype = dtype) - } - variance <- std_dev_t * eta * variance_noise - prev_sample <- prev_sample + variance - } - prev_sample + prev_sample } #' Add noise to latents using DDIM scheduler @@ -338,56 +323,51 @@ ddim_scheduler_step <- function( #' } #' #' @export -scheduler_add_noise <- function( - original_latents, - noise, - timestep, - scheduler_obj -) { - # Get the alpha_cumprod value for this timestep - # In DDIM/DDPM schedulers, alphas_cumprod represents - # how much of the original signal remains at each timestep +scheduler_add_noise <- function(original_latents, noise, timestep, + scheduler_obj) { + # Get the alpha_cumprod value for this timestep + # In DDIM/DDPM schedulers, alphas_cumprod represents + # how much of the original signal remains at each timestep - timestep_index <- timestep + 1 - # Get alpha_cumprod for this timestep - alpha_cumprod <- scheduler_obj$alphas_cumprod[timestep_index] + timestep_index <- timestep + 1 + # Get alpha_cumprod for this timestep + alpha_cumprod <- scheduler_obj$alphas_cumprod[timestep_index] - # Calculate the noise scaling factors - sqrt_alpha_prod <- torch::torch_sqrt(torch::torch_tensor(alpha_cumprod, - device = original_latents$device)) - sqrt_one_minus_alpha_prod <- torch::torch_sqrt( - torch::torch_tensor(1 - alpha_cumprod, device = original_latents$device) - ) + # Calculate the noise scaling factors + sqrt_alpha_prod <- torch::torch_sqrt(torch::torch_tensor(alpha_cumprod, + device = original_latents$device)) + sqrt_one_minus_alpha_prod <- torch::torch_sqrt( + torch::torch_tensor(1 - alpha_cumprod, device = original_latents$device) + ) - # Add noise to the original latents according to the diffusion formula: - # noisy = sqrt(alpha) * original + sqrt(1-alpha) * noise - noised_latents <- sqrt_alpha_prod * original_latents + - sqrt_one_minus_alpha_prod * noise + # Add noise to the original latents according to the diffusion formula: + # noisy = sqrt(alpha) * original + sqrt(1-alpha) * noise + noised_latents <- sqrt_alpha_prod * original_latents + + sqrt_one_minus_alpha_prod * noise - return(noised_latents) + return(noised_latents) } rescale_zero_terminal_snr <- function(betas) { - alphas <- 1.0 - betas - alphas_cumprod <- cumprod(alphas) - alphas_bar_sqrt <- sqrt(alphas_cumprod) + alphas <- 1.0 - betas + alphas_cumprod <- cumprod(alphas) + alphas_bar_sqrt <- sqrt(alphas_cumprod) - # Store old values. - alphas_bar_sqrt_0 <- alphas_bar_sqrt[1] - alphas_bar_sqrt_T <- alphas_bar_sqrt[length(alphas_bar_sqrt)] + # Store old values. + alphas_bar_sqrt_0 <- alphas_bar_sqrt[1] + alphas_bar_sqrt_T <- alphas_bar_sqrt[length(alphas_bar_sqrt)] - # Shift so the last timestep is zero. - alphas_bar_sqrt <- alphas_bar_sqrt - alphas_bar_sqrt_T + # Shift so the last timestep is zero. + alphas_bar_sqrt <- alphas_bar_sqrt - alphas_bar_sqrt_T - # Scale so the first timestep is back to the old value. - alphas_bar_sqrt <- alphas_bar_sqrt * alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) + # Scale so the first timestep is back to the old value. + alphas_bar_sqrt <- alphas_bar_sqrt * alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - # Convert alphas_bar_sqrt to betas - alphas_bar <- alphas_bar_sqrt ^ 2# Revert sqrt - alphas = alphas_bar[- 1] / alphas_bar[- length(alphas_bar)]# Revert cumprod - alphas = c(alphas_bar[1], alphas) - betas = 1 - alphas + # Convert alphas_bar_sqrt to betas + alphas_bar <- alphas_bar_sqrt ^ 2 # Revert sqrt + alphas <- alphas_bar[-1] / alphas_bar[-length(alphas_bar)] # Revert cumprod + alphas <- c(alphas_bar[1], alphas) + betas <- 1 - alphas - return(betas) + return(betas) } - diff --git a/R/dit_ltx2.R b/R/dit_ltx2.R deleted file mode 100644 index 94288cb..0000000 --- a/R/dit_ltx2.R +++ /dev/null @@ -1,936 +0,0 @@ -# LTX2 DiT Transformer Model -# Full audio-video transformer matching HuggingFace diffusers implementation - -# ------------------------------------------------------------------------------ -# LTX2 Audio-Video Rotary Positional Embeddings -# ------------------------------------------------------------------------------ - -#' LTX2 Audio-Video Rotary Positional Embeddings -#' @keywords internal -ltx2_audio_video_rotary_pos_embed <- torch::nn_module( - "LTX2AudioVideoRotaryPosEmbed", - initialize = function( - dim, - patch_size = 1L, - patch_size_t = 1L, - base_num_frames = 20L, - base_height = 2048L, - base_width = 2048L, - sampling_rate = 16000L, - hop_length = 160L, - scale_factors = c(8L, 32L, 32L), - theta = 10000.0, - causal_offset = 1L, - modality = "video", - double_precision = TRUE, - rope_type = "interleaved", - num_attention_heads = 32L - ) { - - self$dim <- dim - self$patch_size <- patch_size - self$patch_size_t <- patch_size_t - - if (!rope_type %in% c("interleaved", "split")) { - stop("rope_type must be 'interleaved' or 'split'") - } - self$rope_type <- rope_type - - self$base_num_frames <- base_num_frames - self$num_attention_heads <- num_attention_heads - - # Video-specific - self$base_height <- base_height - self$base_width <- base_width - - # Audio-specific - self$sampling_rate <- sampling_rate - self$hop_length <- hop_length - self$audio_latents_per_second <- as.numeric(sampling_rate) / as.numeric(hop_length) / as.numeric(scale_factors[1]) - - self$scale_factors <- scale_factors - self$theta <- theta - self$causal_offset <- causal_offset - - self$modality <- modality - if (!modality %in% c("video", "audio")) { - stop("modality must be 'video' or 'audio'") - } - self$double_precision <- double_precision - }, - - prepare_video_coords = function( - batch_size, - num_frames, - height, - width, - device, - fps = 24.0 - ) { - # Generate grid coordinates for each spatiotemporal dimension - grid_f <- torch::torch_arange(start = 0, end = num_frames - 1L, step = self$patch_size_t, - dtype = torch::torch_float32(), device = device) - grid_h <- torch::torch_arange(start = 0, end = height - 1L, step = self$patch_size, - dtype = torch::torch_float32(), device = device) - grid_w <- torch::torch_arange(start = 0, end = width - 1L, step = self$patch_size, - dtype = torch::torch_float32(), device = device) - - # Create meshgrid (F, H, W order) - grids <- torch::torch_meshgrid(list(grid_f, grid_h, grid_w), indexing = "ij") - grid <- torch::torch_stack(grids, dim = 1L) # [3, N_F, N_H, N_W] - - # Get patch boundaries - patch_size_vec <- c(self$patch_size_t, self$patch_size, self$patch_size) - patch_size_delta <- torch::torch_tensor(patch_size_vec, dtype = grid$dtype, device = grid$device) - patch_ends <- grid + patch_size_delta$view(c(3L, 1L, 1L, 1L)) - - # Combine start and end coordinates - latent_coords <- torch::torch_stack(list(grid, patch_ends), dim = - 1L) # [3, N_F, N_H, N_W, 2] - - # Reshape to [batch_size, 3, num_patches, 2] - latent_coords <- latent_coords$flatten(start_dim = 2L, end_dim = 4L) - latent_coords <- latent_coords$unsqueeze(1L)$`repeat`(c(batch_size, 1L, 1L, 1L)) - - # Scale to pixel space - scale_tensor <- torch::torch_tensor(self$scale_factors, device = latent_coords$device) - pixel_coords <- latent_coords * scale_tensor$view(c(1L, 3L, 1L, 1L)) - - # Handle causal offset for temporal dimension - pixel_coords[, 1,,] <- (pixel_coords[, 1,,] + self$causal_offset - self$scale_factors[1])$clamp(min = 0) - - # Scale temporal by FPS - pixel_coords[, 1,,] <- pixel_coords[, 1,,] / fps - - pixel_coords - }, - - prepare_audio_coords = function( - batch_size, - num_frames, - device, - shift = 0L - ) { - # Generate coordinates in the frame (time) dimension - grid_f <- torch::torch_arange(start = shift, end = num_frames + shift - 1L, step = self$patch_size_t, - dtype = torch::torch_float32(), device = device) - - audio_scale_factor <- self$scale_factors[1] - - # Calculate start timestamps - grid_start_mel <- grid_f * audio_scale_factor - grid_start_mel <- (grid_start_mel + self$causal_offset - audio_scale_factor)$clamp(min = 0) - grid_start_s <- grid_start_mel * self$hop_length / self$sampling_rate - - # Calculate end timestamps - grid_end_mel <- (grid_f + self$patch_size_t) * audio_scale_factor - grid_end_mel <- (grid_end_mel + self$causal_offset - audio_scale_factor)$clamp(min = 0) - grid_end_s <- grid_end_mel * self$hop_length / self$sampling_rate - - audio_coords <- torch::torch_stack(list(grid_start_s, grid_end_s), dim = - 1L) # [num_patches, 2] - audio_coords <- audio_coords$unsqueeze(1L)$expand(c(batch_size, - 1L, - 1L)) # [batch_size, num_patches, 2] - audio_coords <- audio_coords$unsqueeze(2L) # [batch_size, 1, num_patches, 2] - - audio_coords - }, - - prepare_coords = function(...) { - if (self$modality == "video") { - self$prepare_video_coords(...) - } else { - self$prepare_audio_coords(...) - } - }, - - forward = function( - coords, - device = NULL, - dtype = NULL - ) { - if (is.null(device)) device <- coords$device - # Store target dtype for conversion at end (RoPE computed in float32 for precision) - - num_pos_dims <- coords$shape[2] - - # If coords are patch boundaries, use midpoint - if (coords$ndim == 4L) { - coords_chunks <- coords$chunk(2L, dim = - 1L) - coords_start <- coords_chunks[[1]] - coords_end <- coords_chunks[[2]] - coords <- (coords_start + coords_end) / 2.0 - coords <- coords$squeeze(- 1L) # [B, num_pos_dims, num_patches] - } - - # Get coordinates as fraction of base data shape - if (self$modality == "video") { - max_positions <- c(self$base_num_frames, self$base_height, self$base_width) - } else { - max_positions <- c(self$base_num_frames) - } - - # [B, num_pos_dims, num_patches] -> [B, num_patches, num_pos_dims] - grid_parts <- lapply(seq_len(num_pos_dims), function(i) { - coords[, i,] / max_positions[i] - }) - grid <- torch::torch_stack(grid_parts, dim = - 1L)$to(device = device) - - num_rope_elems <- num_pos_dims * 2L - - # Create frequency grid - if (self$double_precision) { - freqs_dtype <- torch::torch_float64() - } else { - freqs_dtype <- torch::torch_float32() - } - pow_indices <- torch::torch_pow( - self$theta, - torch::torch_linspace(start = 0.0, end = 1.0, steps = self$dim %/% num_rope_elems, - dtype = freqs_dtype, device = device) - ) - freqs <- (pow_indices * pi / 2.0)$to(dtype = torch::torch_float32()) - - # Outer product - freqs <- (grid$unsqueeze(- 1L) * 2 - 1) * freqs# [B, num_patches, num_pos_dims, dim/num_elems] - freqs <- freqs$transpose(- 1L, - 2L)$flatten(start_dim = 3L) # [B, num_patches, dim/2] - - # Get cos/sin frequencies - if (self$rope_type == "interleaved") { - cos_freqs <- freqs$cos()$repeat_interleave(2L, dim = - 1L) - sin_freqs <- freqs$sin()$repeat_interleave(2L, dim = - 1L) - - if (self$dim %% num_rope_elems != 0L) { - pad_size <- self$dim %% num_rope_elems - cos_padding <- torch::torch_ones_like(cos_freqs[,, 1:pad_size]) - sin_padding <- torch::torch_zeros_like(sin_freqs[,, 1:pad_size]) - cos_freqs <- torch::torch_cat(list(cos_padding, cos_freqs), dim = - 1L) - sin_freqs <- torch::torch_cat(list(sin_padding, sin_freqs), dim = - 1L) - } - } else { - # split type - expected_freqs <- self$dim %/% 2L - current_freqs <- freqs$shape[length(freqs$shape)] - pad_size <- expected_freqs - current_freqs - - cos_freq <- freqs$cos() - sin_freq <- freqs$sin() - - if (pad_size != 0L) { - cos_padding <- torch::torch_ones_like(cos_freq[,, 1:pad_size]) - sin_padding <- torch::torch_zeros_like(sin_freq[,, 1:pad_size]) - cos_freq <- torch::torch_cat(list(cos_padding, cos_freq), dim = - 1L) - sin_freq <- torch::torch_cat(list(sin_padding, sin_freq), dim = - 1L) - } - - # Reshape for multi-head attention - b <- cos_freq$shape[1] - t <- cos_freq$shape[2] - - cos_freq <- cos_freq$reshape(c(b, t, self$num_attention_heads, - 1L)) - sin_freq <- sin_freq$reshape(c(b, t, self$num_attention_heads, - 1L)) - - cos_freqs <- cos_freq$transpose(2L, 3L) # (B, H, T, D//2) - sin_freqs <- sin_freq$transpose(2L, 3L) - } - - # Convert to target dtype if specified (for mixed precision) - if (!is.null(dtype)) { - cos_freqs <- cos_freqs$to(dtype = dtype) - sin_freqs <- sin_freqs$to(dtype = dtype) - } - - list(cos_freqs, sin_freqs) - } -) - -# ------------------------------------------------------------------------------ -# LTX2 Video Transformer 3D Model -# ------------------------------------------------------------------------------ - -#' LTX2 Video Transformer 3D Model (Audio-Video) -#' -#' Full audio-video transformer matching HuggingFace diffusers implementation. -#' -#' @param in_channels Integer. Video input channels (default: 128). -#' @param out_channels Integer. Video output channels (default: 128). -#' @param patch_size Integer. Spatial patch size (default: 1). -#' @param patch_size_t Integer. Temporal patch size (default: 1). -#' @param num_attention_heads Integer. Video attention heads (default: 32). -#' @param attention_head_dim Integer. Video attention head dimension (default: 128). -#' @param cross_attention_dim Integer. Video cross-attention dimension (default: 4096). -#' @param vae_scale_factors Integer vector. VAE scale factors (default: c(8, 32, 32)). -#' @param pos_embed_max_pos Integer. Max position for RoPE (default: 20). -#' @param base_height Integer. Base height for RoPE (default: 2048). -#' @param base_width Integer. Base width for RoPE (default: 2048). -#' @param audio_in_channels Integer. Audio input channels (default: 128). -#' @param audio_out_channels Integer. Audio output channels (default: 128). -#' @param audio_patch_size Integer. Audio patch size (default: 1). -#' @param audio_patch_size_t Integer. Audio temporal patch size (default: 1). -#' @param audio_num_attention_heads Integer. Audio attention heads (default: 32). -#' @param audio_attention_head_dim Integer. Audio head dimension (default: 64). -#' @param audio_cross_attention_dim Integer. Audio cross-attention dim (default: 2048). -#' @param audio_scale_factor Integer. Audio scale factor (default: 4). -#' @param audio_pos_embed_max_pos Integer. Audio max position (default: 20). -#' @param audio_sampling_rate Integer. Audio sampling rate (default: 16000). -#' @param audio_hop_length Integer. Audio hop length (default: 160). -#' @param num_layers Integer. Number of transformer layers (default: 48). -#' @param activation_fn Character. Activation function (default: "gelu-approximate"). -#' @param qk_norm Character. QK normalization type (default: "rms_norm_across_heads"). -#' @param norm_elementwise_affine Logical. Use elementwise affine in norms (default: FALSE). -#' @param norm_eps Numeric. Epsilon for normalization (default: 1e-6). -#' @param caption_channels Integer. Caption embedding channels (default: 3840). -#' @param attention_bias Logical. Use bias in attention (default: TRUE). -#' @param attention_out_bias Logical. Use bias in attention output (default: TRUE). -#' @param rope_theta Numeric. Theta for RoPE (default: 10000). -#' @param rope_double_precision Logical. Use double precision for RoPE (default: TRUE). -#' @param causal_offset Integer. Causal offset for RoPE (default: 1). -#' @param timestep_scale_multiplier Numeric. Timestep scale (default: 1000). -#' @param cross_attn_timestep_scale_multiplier Numeric. Cross-attn timestep scale (default: 1000). -#' @param rope_type Character. RoPE type: "interleaved" or "split" (default: "interleaved"). -#' -#' @return An nn_module representing the LTX2 video transformer. -#' @export -ltx2_video_transformer_3d_model <- torch::nn_module( - "LTX2VideoTransformer3DModel", - initialize = function( - in_channels = 128L, - out_channels = 128L, - patch_size = 1L, - patch_size_t = 1L, - num_attention_heads = 32L, - attention_head_dim = 128L, - cross_attention_dim = 4096L, - vae_scale_factors = c(8L, 32L, 32L), - pos_embed_max_pos = 20L, - base_height = 2048L, - base_width = 2048L, - audio_in_channels = 128L, - audio_out_channels = 128L, - audio_patch_size = 1L, - audio_patch_size_t = 1L, - audio_num_attention_heads = 32L, - audio_attention_head_dim = 64L, - audio_cross_attention_dim = 2048L, - audio_scale_factor = 4L, - audio_pos_embed_max_pos = 20L, - audio_sampling_rate = 16000L, - audio_hop_length = 160L, - num_layers = 48L, - activation_fn = "gelu-approximate", - qk_norm = "rms_norm_across_heads", - norm_elementwise_affine = FALSE, - norm_eps = 1e-6, - caption_channels = 3840L, - attention_bias = TRUE, - attention_out_bias = TRUE, - rope_theta = 10000.0, - rope_double_precision = TRUE, - causal_offset = 1L, - timestep_scale_multiplier = 1000L, - cross_attn_timestep_scale_multiplier = 1000L, - rope_type = "interleaved" - ) { - - if (is.null(out_channels)) out_channels <- in_channels - if (is.null(audio_out_channels)) audio_out_channels <- audio_in_channels - - self$timestep_scale_multiplier <- timestep_scale_multiplier - self$cross_attn_timestep_scale_multiplier <- cross_attn_timestep_scale_multiplier - - inner_dim <- num_attention_heads * attention_head_dim - audio_inner_dim <- audio_num_attention_heads * audio_attention_head_dim - - # 1. Patchification input projections - self$proj_in <- make_linear(in_channels, inner_dim) - self$audio_proj_in <- make_linear(audio_in_channels, audio_inner_dim) - - # 2. Prompt embeddings - self$caption_projection <- pixart_alpha_text_projection( - in_features = caption_channels, hidden_size = inner_dim - ) - self$audio_caption_projection <- pixart_alpha_text_projection( - in_features = caption_channels, hidden_size = audio_inner_dim - ) - - # 3. Timestep modulation - # 3.1 Global timestep embedding and modulation - self$time_embed <- ltx2_ada_layer_norm_single(inner_dim, num_mod_params = 6L, use_additional_conditions = FALSE) - self$audio_time_embed <- ltx2_ada_layer_norm_single(audio_inner_dim, num_mod_params = 6L, use_additional_conditions = FALSE) - - # 3.2 Global cross-attention modulation - self$av_cross_attn_video_scale_shift <- ltx2_ada_layer_norm_single(inner_dim, num_mod_params = 4L, use_additional_conditions = FALSE) - self$av_cross_attn_audio_scale_shift <- ltx2_ada_layer_norm_single(audio_inner_dim, num_mod_params = 4L, use_additional_conditions = FALSE) - self$av_cross_attn_video_a2v_gate <- ltx2_ada_layer_norm_single(inner_dim, num_mod_params = 1L, use_additional_conditions = FALSE) - self$av_cross_attn_audio_v2a_gate <- ltx2_ada_layer_norm_single(audio_inner_dim, num_mod_params = 1L, use_additional_conditions = FALSE) - - # 3.3 Output layer modulation - self$scale_shift_table <- torch::nn_parameter(torch::torch_randn(2L, inner_dim) / sqrt(inner_dim)) - self$audio_scale_shift_table <- torch::nn_parameter(torch::torch_randn(2L, audio_inner_dim) / sqrt(audio_inner_dim)) - - # 4. Rotary Positional Embeddings - # Video self-attention RoPE - self$rope <- ltx2_audio_video_rotary_pos_embed( - dim = inner_dim, - patch_size = patch_size, - patch_size_t = patch_size_t, - base_num_frames = pos_embed_max_pos, - base_height = base_height, - base_width = base_width, - scale_factors = vae_scale_factors, - theta = rope_theta, - causal_offset = causal_offset, - modality = "video", - double_precision = rope_double_precision, - rope_type = rope_type, - num_attention_heads = num_attention_heads - ) - - # Audio self-attention RoPE - self$audio_rope <- ltx2_audio_video_rotary_pos_embed( - dim = audio_inner_dim, - patch_size = audio_patch_size, - patch_size_t = audio_patch_size_t, - base_num_frames = audio_pos_embed_max_pos, - sampling_rate = audio_sampling_rate, - hop_length = audio_hop_length, - scale_factors = c(audio_scale_factor), - theta = rope_theta, - causal_offset = causal_offset, - modality = "audio", - double_precision = rope_double_precision, - rope_type = rope_type, - num_attention_heads = audio_num_attention_heads - ) - - # Cross-attention RoPE - cross_attn_pos_embed_max_pos <- max(pos_embed_max_pos, audio_pos_embed_max_pos) - - self$cross_attn_rope <- ltx2_audio_video_rotary_pos_embed( - dim = audio_cross_attention_dim, - patch_size = patch_size, - patch_size_t = patch_size_t, - base_num_frames = cross_attn_pos_embed_max_pos, - base_height = base_height, - base_width = base_width, - theta = rope_theta, - causal_offset = causal_offset, - modality = "video", - double_precision = rope_double_precision, - rope_type = rope_type, - num_attention_heads = num_attention_heads - ) - - self$cross_attn_audio_rope <- ltx2_audio_video_rotary_pos_embed( - dim = audio_cross_attention_dim, - patch_size = audio_patch_size, - patch_size_t = audio_patch_size_t, - base_num_frames = cross_attn_pos_embed_max_pos, - sampling_rate = audio_sampling_rate, - hop_length = audio_hop_length, - theta = rope_theta, - causal_offset = causal_offset, - modality = "audio", - double_precision = rope_double_precision, - rope_type = rope_type, - num_attention_heads = audio_num_attention_heads - ) - - # 5. Transformer blocks - self$transformer_blocks <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { - ltx2_video_transformer_block( - dim = inner_dim, - num_attention_heads = num_attention_heads, - attention_head_dim = attention_head_dim, - cross_attention_dim = cross_attention_dim, - audio_dim = audio_inner_dim, - audio_num_attention_heads = audio_num_attention_heads, - audio_attention_head_dim = audio_attention_head_dim, - audio_cross_attention_dim = audio_cross_attention_dim, - qk_norm = qk_norm, - activation_fn = activation_fn, - attention_bias = attention_bias, - attention_out_bias = attention_out_bias, - eps = norm_eps, - elementwise_affine = norm_elementwise_affine, - rope_type = rope_type - ) - })) - - # 6. Output layers - self$norm_out <- torch::nn_layer_norm(inner_dim, eps = 1e-6, elementwise_affine = FALSE) - self$proj_out <- make_linear(inner_dim, out_channels) - - self$audio_norm_out <- torch::nn_layer_norm(audio_inner_dim, eps = 1e-6, elementwise_affine = FALSE) - self$audio_proj_out <- make_linear(audio_inner_dim, audio_out_channels) - - self$gradient_checkpointing <- FALSE - }, - - forward = function( - hidden_states, - audio_hidden_states, - encoder_hidden_states, - audio_encoder_hidden_states, - timestep, - audio_timestep = NULL, - encoder_attention_mask = NULL, - audio_encoder_attention_mask = NULL, - num_frames = NULL, - height = NULL, - width = NULL, - fps = 24.0, - audio_num_frames = NULL, - video_coords = NULL, - audio_coords = NULL - ) { - - if (is.null(audio_timestep)) audio_timestep <- timestep - - # Convert attention masks to bias (use tensor ops to preserve dtype) - # Formula: (1 - mask) * -10000.0 -> (mask - 1) * 10000.0 - if (!is.null(encoder_attention_mask) && encoder_attention_mask$ndim == 2L) { - mask_dtype <- hidden_states$dtype - encoder_attention_mask <- encoder_attention_mask$to(dtype = mask_dtype) - # Use tensor scalar to preserve dtype - scale <- torch::torch_tensor(- 10000.0, dtype = mask_dtype, device = encoder_attention_mask$device) - encoder_attention_mask <- encoder_attention_mask$sub(1)$neg()$mul(scale) - encoder_attention_mask <- encoder_attention_mask$unsqueeze(2L) - } - - if (!is.null(audio_encoder_attention_mask) && audio_encoder_attention_mask$ndim == 2L) { - mask_dtype <- audio_hidden_states$dtype - audio_encoder_attention_mask <- audio_encoder_attention_mask$to(dtype = mask_dtype) - scale <- torch::torch_tensor(- 10000.0, dtype = mask_dtype, device = audio_encoder_attention_mask$device) - audio_encoder_attention_mask <- audio_encoder_attention_mask$sub(1)$neg()$mul(scale) - audio_encoder_attention_mask <- audio_encoder_attention_mask$unsqueeze(2L) - } - - batch_size <- hidden_states$shape[1] - - # 1. Prepare RoPE positional embeddings - if (is.null(video_coords)) { - video_coords <- self$rope$prepare_video_coords(batch_size, num_frames, height, width, hidden_states$device, fps = fps) - } - if (is.null(audio_coords)) { - audio_coords <- self$audio_rope$prepare_audio_coords(batch_size, audio_num_frames, audio_hidden_states$device) - } - - video_rotary_emb <- self$rope(video_coords, device = hidden_states$device, dtype = hidden_states$dtype) - audio_rotary_emb <- self$audio_rope(audio_coords, device = audio_hidden_states$device, dtype = audio_hidden_states$dtype) - - video_cross_attn_rotary_emb <- self$cross_attn_rope(video_coords[, 1:1,,], device = hidden_states$device, dtype = hidden_states$dtype) - audio_cross_attn_rotary_emb <- self$cross_attn_audio_rope(audio_coords[, 1:1,,], device = audio_hidden_states$device, dtype = audio_hidden_states$dtype) - - # 2. Patchify input projections - hidden_states <- self$proj_in(hidden_states) - audio_hidden_states <- self$audio_proj_in(audio_hidden_states) - - # 3. Prepare timestep embeddings - timestep_cross_attn_gate_scale_factor <- self$cross_attn_timestep_scale_multiplier / self$timestep_scale_multiplier - - # 3.1 Global timestep embedding - temb_result <- self$time_embed(timestep$flatten(), batch_size = batch_size, hidden_dtype = hidden_states$dtype) - temb <- temb_result[[1]]$view(c(batch_size, - 1L, temb_result[[1]]$shape[length(temb_result[[1]]$shape)])) - embedded_timestep <- temb_result[[2]]$view(c(batch_size, - 1L, temb_result[[2]]$shape[length(temb_result[[2]]$shape)])) - - temb_audio_result <- self$audio_time_embed(audio_timestep$flatten(), batch_size = batch_size, hidden_dtype = audio_hidden_states$dtype) - temb_audio <- temb_audio_result[[1]]$view(c(batch_size, - 1L, temb_audio_result[[1]]$shape[length(temb_audio_result[[1]]$shape)])) - audio_embedded_timestep <- temb_audio_result[[2]]$view(c(batch_size, - 1L, temb_audio_result[[2]]$shape[length(temb_audio_result[[2]]$shape)])) - - # 3.2 Cross-attention modulation - video_cross_attn_scale_shift_result <- self$av_cross_attn_video_scale_shift(timestep$flatten(), batch_size = batch_size, hidden_dtype = hidden_states$dtype) - video_cross_attn_a2v_gate_result <- self$av_cross_attn_video_a2v_gate(timestep$flatten() * timestep_cross_attn_gate_scale_factor, batch_size = batch_size, hidden_dtype = hidden_states$dtype) - - video_cross_attn_scale_shift <- video_cross_attn_scale_shift_result[[1]]$view(c(batch_size, - 1L, video_cross_attn_scale_shift_result[[1]]$shape[length(video_cross_attn_scale_shift_result[[1]]$shape)])) - video_cross_attn_a2v_gate <- video_cross_attn_a2v_gate_result[[1]]$view(c(batch_size, - 1L, video_cross_attn_a2v_gate_result[[1]]$shape[length(video_cross_attn_a2v_gate_result[[1]]$shape)])) - - audio_cross_attn_scale_shift_result <- self$av_cross_attn_audio_scale_shift(audio_timestep$flatten(), batch_size = batch_size, hidden_dtype = audio_hidden_states$dtype) - audio_cross_attn_v2a_gate_result <- self$av_cross_attn_audio_v2a_gate(audio_timestep$flatten() * timestep_cross_attn_gate_scale_factor, batch_size = batch_size, hidden_dtype = audio_hidden_states$dtype) - - audio_cross_attn_scale_shift <- audio_cross_attn_scale_shift_result[[1]]$view(c(batch_size, - 1L, audio_cross_attn_scale_shift_result[[1]]$shape[length(audio_cross_attn_scale_shift_result[[1]]$shape)])) - audio_cross_attn_v2a_gate <- audio_cross_attn_v2a_gate_result[[1]]$view(c(batch_size, - 1L, audio_cross_attn_v2a_gate_result[[1]]$shape[length(audio_cross_attn_v2a_gate_result[[1]]$shape)])) - - # 4. Prepare prompt embeddings - encoder_hidden_states <- self$caption_projection(encoder_hidden_states) - encoder_hidden_states <- encoder_hidden_states$view(c(batch_size, - 1L, hidden_states$shape[length(hidden_states$shape)])) - - audio_encoder_hidden_states <- self$audio_caption_projection(audio_encoder_hidden_states) - audio_encoder_hidden_states <- audio_encoder_hidden_states$view(c(batch_size, - 1L, audio_hidden_states$shape[length(audio_hidden_states$shape)])) - - # 5. Run transformer blocks - for (i in seq_along(self$transformer_blocks)) { - block <- self$transformer_blocks[[i]] - result <- block( - hidden_states = hidden_states, - audio_hidden_states = audio_hidden_states, - encoder_hidden_states = encoder_hidden_states, - audio_encoder_hidden_states = audio_encoder_hidden_states, - temb = temb, - temb_audio = temb_audio, - temb_ca_scale_shift = video_cross_attn_scale_shift, - temb_ca_audio_scale_shift = audio_cross_attn_scale_shift, - temb_ca_gate = video_cross_attn_a2v_gate, - temb_ca_audio_gate = audio_cross_attn_v2a_gate, - video_rotary_emb = video_rotary_emb, - audio_rotary_emb = audio_rotary_emb, - ca_video_rotary_emb = video_cross_attn_rotary_emb, - ca_audio_rotary_emb = audio_cross_attn_rotary_emb, - encoder_attention_mask = encoder_attention_mask, - audio_encoder_attention_mask = audio_encoder_attention_mask - ) - hidden_states <- result[[1]] - audio_hidden_states <- result[[2]] - } - - # 6. Output layers - scale_shift_values <- self$scale_shift_table$unsqueeze(1)$unsqueeze(1)$to(dtype = hidden_states$dtype) + embedded_timestep$unsqueeze(3) - shift <- scale_shift_values[,, 1,] - scale <- scale_shift_values[,, 2,] - - hidden_states <- self$norm_out(hidden_states) - hidden_states <- hidden_states * scale$add(1) + shift - output <- self$proj_out(hidden_states) - - audio_scale_shift_values <- self$audio_scale_shift_table$unsqueeze(1)$unsqueeze(1)$to(dtype = audio_hidden_states$dtype) + audio_embedded_timestep$unsqueeze(3) - audio_shift <- audio_scale_shift_values[,, 1,] - audio_scale <- audio_scale_shift_values[,, 2,] - - audio_hidden_states <- self$audio_norm_out(audio_hidden_states) - audio_hidden_states <- audio_hidden_states * audio_scale$add(1) + audio_shift - audio_output <- self$audio_proj_out(audio_hidden_states) - - list(sample = output, audio_sample = audio_output) - } -) - -# ------------------------------------------------------------------------------ -# Weight Loading Functions -# ------------------------------------------------------------------------------ - -#' Load LTX2 DiT Transformer from safetensors -#' -#' Load pre-trained LTX2 transformer weights from HuggingFace safetensors files. -#' Supports both single file and sharded multi-file loading. -#' -#' @param weights_dir Character. Directory containing safetensors files. -#' @param config_path Character. Optional path to config.json. If NULL, uses default config. -#' @param device Character. Device to load weights to. Default: "cpu" -#' @param dtype Character. Data type ("float32", "float16", "bfloat16"). Default: "float16" -#' @param verbose Logical. Print loading progress. Default: TRUE -#' @return Initialized ltx2_video_transformer3d module -#' @export -load_ltx2_transformer <- function( - weights_dir, - config_path = NULL, - device = "cpu", - dtype = "float16", - verbose = TRUE -) { - - # Look for config - if (is.null(config_path)) { - config_path <- file.path(weights_dir, "config.json") - } - - config <- NULL - if (file.exists(config_path)) { - config <- jsonlite::fromJSON(config_path) - if (verbose) message("Loaded config from: ", config_path) - } - - # Create transformer with config or defaults - if (!is.null(config)) { - transformer <- ltx2_video_transformer_3d_model( - in_channels = config$in_channels %||% 128L, - out_channels = config$out_channels %||% 128L, - patch_size = config$patch_size %||% 1L, - patch_size_t = config$patch_size_t %||% 1L, - num_attention_heads = config$num_attention_heads %||% 32L, - attention_head_dim = config$attention_head_dim %||% 128L, - cross_attention_dim = config$cross_attention_dim %||% 4096L, - num_layers = config$num_layers %||% 48L, - caption_channels = config$caption_channels %||% 3840L, - qk_norm = config$qk_norm %||% "rms_norm_across_heads", - norm_elementwise_affine = config$norm_elementwise_affine %||% FALSE, - norm_eps = config$norm_eps %||% 1e-6, - timestep_scale_multiplier = config$timestep_scale_multiplier %||% 1000.0, - cross_attn_timestep_scale_multiplier = config$cross_attn_timestep_scale_multiplier %||% 1000.0, - base_height = config$base_height %||% 2048L, - base_width = config$base_width %||% 2048L, - pos_embed_max_pos = config$pos_embed_max_pos %||% 20L, - rope_theta = config$rope_theta %||% 10000.0, - rope_type = config$rope_type %||% "split", - causal_offset = config$causal_offset %||% 1L, - audio_in_channels = config$audio_in_channels %||% 128L, - audio_out_channels = config$audio_out_channels %||% 128L, - audio_num_attention_heads = config$audio_num_attention_heads %||% 32L, - audio_attention_head_dim = config$audio_attention_head_dim %||% 64L, - audio_cross_attention_dim = config$audio_cross_attention_dim %||% 2048L, - audio_sampling_rate = config$audio_sampling_rate %||% 16000L, - audio_hop_length = config$audio_hop_length %||% 160L, - audio_scale_factor = config$audio_scale_factor %||% 4L, - audio_pos_embed_max_pos = config$audio_pos_embed_max_pos %||% 20L - ) - } else { - transformer <- ltx2_video_transformer_3d_model() - } - - # Find weight files - index_path <- file.path(weights_dir, "diffusion_pytorch_model.safetensors.index.json") - single_path <- file.path(weights_dir, "diffusion_pytorch_model.safetensors") - - if (file.exists(index_path)) { - # Sharded weights - if (verbose) message("Loading sharded weights from: ", weights_dir) - load_ltx2_transformer_sharded(transformer, weights_dir, index_path, verbose = verbose) - } else if (file.exists(single_path)) { - # Single file - if (verbose) message("Loading weights from: ", single_path) - weights <- safetensors::safe_load_file(single_path, framework = "torch") - load_ltx2_transformer_weights(transformer, weights, verbose = verbose) - } else { - stop("No weights found in: ", weights_dir) - } - - # Convert dtype and move to device - torch_dtype <- switch(dtype, - "float32" = torch::torch_float32(), - "float16" = torch::torch_float16(), - "bfloat16" = torch::torch_bfloat16(), - torch::torch_float16() - ) - - transformer$to(device = device, dtype = torch_dtype) - - if (verbose) message("Transformer loaded successfully on device: ", device, " dtype: ", dtype) - transformer -} - -#' Load sharded transformer weights -#' -#' @param transformer LTX2 transformer module -#' @param weights_dir Directory containing sharded safetensors -#' @param index_path Path to index.json -#' @param verbose Print progress -#' @keywords internal -load_ltx2_transformer_sharded <- function( - transformer, - weights_dir, - index_path, - verbose = TRUE -) { - # Load index - index <- jsonlite::fromJSON(index_path) - weight_map <- index$weight_map - - # Get unique shard files - shard_files <- unique(unlist(weight_map)) - if (verbose) message(sprintf("Loading %d shards...", length(shard_files))) - - total_loaded <- 0L - total_skipped <- 0L - - for (i in seq_along(shard_files)) { - shard_file <- shard_files[i] - shard_path <- file.path(weights_dir, shard_file) - - if (!file.exists(shard_path)) { - warning("Shard not found: ", shard_path) - next - } - - if (verbose) message(sprintf("[%d/%d] Loading %s...", i, length(shard_files), shard_file)) - - weights <- safetensors::safe_load_file(shard_path, framework = "torch") - result <- load_ltx2_transformer_weights(transformer, weights, verbose = FALSE) - total_loaded <- total_loaded + result$loaded - total_skipped <- total_skipped + result$skipped - - # Free memory - rm(weights) - gc() - } - - if (verbose) { - message(sprintf("Total: %d loaded, %d skipped", total_loaded, total_skipped)) - } - - invisible(list(loaded = total_loaded, skipped = total_skipped)) -} - -#' Load weights into LTX2 transformer module -#' -#' @param transformer LTX2 transformer module -#' @param weights Named list of weight tensors -#' @param verbose Print progress -#' @keywords internal -load_ltx2_transformer_weights <- function( - transformer, - weights, - verbose = TRUE -) { - native_params <- names(transformer$parameters) - - # Remap HuggingFace names to R module names - remap_transformer_key <- function(key) { - # HuggingFace uses nn.ModuleList for FeedForward: - # ff.net.0.proj.weight -> ff.act_fn.proj.weight - # ff.net.2.weight -> ff.proj_out.weight - key <- gsub("\\.ff\\.net\\.0\\.", ".ff.act_fn.", key) - key <- gsub("\\.ff\\.net\\.2\\.", ".ff.proj_out.", key) - # Same for audio_ff - key <- gsub("\\.audio_ff\\.net\\.0\\.", ".audio_ff.act_fn.", key) - key <- gsub("\\.audio_ff\\.net\\.2\\.", ".audio_ff.proj_out.", key) - - # to_out.0 is used in both HF and our module (ModuleList) - # No remapping needed for to_out - key - } - - loaded <- 0L - skipped <- 0L - unmapped <- character(0) - - torch::with_no_grad({ - for (hf_name in names(weights)) { - native_name <- remap_transformer_key(hf_name) - - if (native_name %in% native_params) { - hf_tensor <- weights[[hf_name]] - native_tensor <- transformer$parameters[[native_name]] - - if (all(as.integer(hf_tensor$shape) == as.integer(native_tensor$shape))) { - native_tensor$copy_(hf_tensor) - loaded <- loaded + 1L - } else { - if (verbose) { - message("Shape mismatch: ", native_name, - " (HF: ", paste(as.integer(hf_tensor$shape), collapse = "x"), - " vs R: ", paste(as.integer(native_tensor$shape), collapse = "x"), ")") - } - skipped <- skipped + 1L - } - } else { - skipped <- skipped + 1L - unmapped <- c(unmapped, paste0(hf_name, " -> ", native_name)) - } - } - }) - - if (verbose) { - message(sprintf("Transformer weights: %d loaded, %d skipped", loaded, skipped)) - if (length(unmapped) > 0 && length(unmapped) <= 20) { - message("Unmapped parameters:") - for (u in unmapped[1:min(20, length(unmapped))]) { - message(" ", u) - } - } - if (length(unmapped) > 20) { - message(" ... and ", length(unmapped) - 20, " more") - } - } - - invisible(list(loaded = loaded, skipped = skipped, unmapped = unmapped)) -} - -# ------------------------------------------------------------------------------ -# Latent Packing/Unpacking for DiT -# ------------------------------------------------------------------------------ - -#' Pack Video Latents for DiT -#' -#' Transforms 5D video latents \[B, C, F, H, W\] to 3D sequence \[B, F*H*W, C\] -#' for input to the DiT transformer. -#' -#' @param latents Tensor of shape \[B, C, F, H, W\]. -#' @param patch_size Integer. Spatial patch size (default 1). -#' @param patch_size_t Integer. Temporal patch size (default 1). -#' -#' @return Packed tensor of shape \[B, num_patches, C*patch_size_t*patch_size^2\]. -#' -#' @export -pack_video_latents <- function( - latents, - patch_size = 1L, - patch_size_t = 1L -) { - batch_size <- latents$shape[1] - num_channels <- latents$shape[2] - num_frames <- latents$shape[3] - height <- latents$shape[4] - width <- latents$shape[5] - - post_patch_frames <- num_frames %/% patch_size_t - post_patch_height <- height %/% patch_size - post_patch_width <- width %/% patch_size - - # Reshape: [B, C, F, H, W] -> [B, C, F//pt, pt, H//p, p, W//p, p] - latents <- latents$reshape(c( - batch_size, num_channels, - post_patch_frames, patch_size_t, - post_patch_height, patch_size, - post_patch_width, patch_size - )) - - # Permute: [B, C, F//pt, pt, H//p, p, W//p, p] -> [B, F//pt, H//p, W//p, C, pt, p, p] - latents <- latents$permute(c(1L, 3L, 5L, 7L, 2L, 4L, 6L, 8L)) - - # Flatten patches: [B, F//pt, H//p, W//p, C*pt*p*p] - latents <- latents$flatten(start_dim = 5L, end_dim = 8L) - - # Flatten sequence: [B, F//pt * H//p * W//p, C*pt*p*p] - latents <- latents$flatten(start_dim = 2L, end_dim = 4L) - - latents -} - -#' Unpack Video Latents from DiT -#' -#' Transforms 3D sequence \[B, num_patches, D\] back to 5D video latents \[B, C, F, H, W\]. -#' -#' @param latents Tensor of shape \[B, num_patches, D\]. -#' @param num_frames Integer. Target number of latent frames. -#' @param height Integer. Target latent height. -#' @param width Integer. Target latent width. -#' @param patch_size Integer. Spatial patch size (default 1). -#' @param patch_size_t Integer. Temporal patch size (default 1). -#' -#' @return Unpacked tensor of shape \[B, C, F, H, W\]. -#' -#' @export -unpack_video_latents <- function( - latents, - num_frames, - height, - width, - patch_size = 1L, - patch_size_t = 1L -) { - batch_size <- latents$shape[1] - - post_patch_frames <- num_frames %/% patch_size_t - post_patch_height <- height %/% patch_size - post_patch_width <- width %/% patch_size - - # Unflatten sequence: [B, S, D] -> [B, F//pt, H//p, W//p, D] - latents <- latents$reshape(c(batch_size, post_patch_frames, post_patch_height, post_patch_width, - 1L)) - - # Calculate channel dimension - d <- latents$shape[5] - num_channels <- d %/% (patch_size_t * patch_size * patch_size) - - # Unflatten patches: [B, F//pt, H//p, W//p, C, pt, p, p] - latents <- latents$reshape(c( - batch_size, post_patch_frames, post_patch_height, post_patch_width, - num_channels, patch_size_t, patch_size, patch_size - )) - - # Permute back: [B, C, F//pt, pt, H//p, p, W//p, p] - latents <- latents$permute(c(1L, 5L, 2L, 6L, 3L, 7L, 4L, 8L)) - - # Flatten to video: [B, C, F, H, W] - latents <- latents$flatten(start_dim = 7L, end_dim = 8L) - latents <- latents$flatten(start_dim = 5L, end_dim = 6L) - latents <- latents$flatten(start_dim = 3L, end_dim = 4L) - - latents -} - diff --git a/R/dit_ltx23.R b/R/dit_ltx23.R new file mode 100644 index 0000000..725c545 --- /dev/null +++ b/R/dit_ltx23.R @@ -0,0 +1,449 @@ +#' LTX-2.3 Audio-Video Diffusion Transformer +#' +#' Fresh R port of the LTX-2 transformer from the diffusers reference +#' (Apache-2.0, src/diffusers/models/transformers/transformer_ltx2.py), +#' configured for LTX 2.3: gated attention, cross-attention modulation, +#' prompt AdaLN, split RoPE, and connector-projected text embeddings +#' (no in-model caption projection). +#' +#' @name dit_ltx23 +NULL + +#' LTX-2.3 video transformer model +#' +#' Dual-stream audio/video DiT. Text embeddings arrive already projected +#' to the video (\code{inner_dim}) and audio (\code{audio_inner_dim}) +#' dimensions by the connector modules. +#' +#' @param in_channels,out_channels Integers. Video latent channels. +#' @param patch_size,patch_size_t Integers. Video patch sizes. +#' @param num_attention_heads,attention_head_dim Video attention shape. +#' @param cross_attention_dim Integer. Video text embedding dimension. +#' @param vae_scale_factors Integer vector. VAE (time, height, width) scales. +#' @param pos_embed_max_pos,base_height,base_width RoPE base grid. +#' @param audio_in_channels,audio_out_channels Integers. Audio latent channels. +#' @param audio_patch_size,audio_patch_size_t Integers. Audio patch sizes. +#' @param audio_num_attention_heads,audio_attention_head_dim Audio attention shape. +#' @param audio_cross_attention_dim Integer. Audio text embedding dimension. +#' @param audio_scale_factor,audio_pos_embed_max_pos,audio_sampling_rate,audio_hop_length +#' Audio latent grid parameters. +#' @param num_layers Integer. Transformer block count. +#' @param norm_eps Numeric. Norm epsilon. +#' @param rope_theta,rope_double_precision,causal_offset RoPE parameters. +#' @param timestep_scale_multiplier,cross_attn_timestep_scale_multiplier +#' Timestep scaling (inputs arrive already scaled; the ratio modulates +#' the a2v/v2a gates). +#' @param rope_type "split" (LTX-2.3) or "interleaved". +#' @param gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn +#' LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints). +#' +#' @export +ltx23_transformer <- torch::nn_module( + "ltx23_transformer", + initialize = function( + in_channels = 128L, + out_channels = 128L, + patch_size = 1L, + patch_size_t = 1L, + num_attention_heads = 32L, + attention_head_dim = 128L, + cross_attention_dim = 4096L, + vae_scale_factors = c(8L, 32L, 32L), + pos_embed_max_pos = 20L, + base_height = 2048L, + base_width = 2048L, + gated_attn = TRUE, + cross_attn_mod = TRUE, + audio_in_channels = 128L, + audio_out_channels = 128L, + audio_patch_size = 1L, + audio_patch_size_t = 1L, + audio_num_attention_heads = 32L, + audio_attention_head_dim = 64L, + audio_cross_attention_dim = 2048L, + audio_scale_factor = 4L, + audio_pos_embed_max_pos = 20L, + audio_sampling_rate = 16000L, + audio_hop_length = 160L, + audio_gated_attn = TRUE, + audio_cross_attn_mod = TRUE, + num_layers = 48L, + norm_eps = 1e-6, + rope_theta = 10000.0, + rope_double_precision = TRUE, + causal_offset = 1L, + timestep_scale_multiplier = 1000, + cross_attn_timestep_scale_multiplier = 1000, + rope_type = "split", + perturbed_attn = TRUE + ) { + inner_dim <- num_attention_heads * attention_head_dim + audio_inner_dim <- audio_num_attention_heads * audio_attention_head_dim + self$inner_dim <- inner_dim + self$audio_inner_dim <- audio_inner_dim + self$timestep_scale_multiplier <- timestep_scale_multiplier + self$cross_attn_timestep_scale_multiplier <- cross_attn_timestep_scale_multiplier + self$prompt_modulation <- cross_attn_mod || audio_cross_attn_mod + + # 1. Patchified input projections + self$proj_in <- torch::nn_linear(in_channels, inner_dim) + self$audio_proj_in <- torch::nn_linear(audio_in_channels, audio_inner_dim) + + # 2. Global timestep modulation (9 params each with cross-attn mod) + video_mod_params <- if (cross_attn_mod) 9L else 6L + audio_mod_params <- if (audio_cross_attn_mod) 9L else 6L + self$time_embed <- ltx23_ada_layer_norm_single(inner_dim, video_mod_params) + self$audio_time_embed <- ltx23_ada_layer_norm_single(audio_inner_dim, + audio_mod_params) + + # Global a2v/v2a cross-attention modulation + self$av_cross_attn_video_scale_shift <- ltx23_ada_layer_norm_single(inner_dim, 4L) + self$av_cross_attn_audio_scale_shift <- ltx23_ada_layer_norm_single(audio_inner_dim, 4L) + self$av_cross_attn_video_a2v_gate <- ltx23_ada_layer_norm_single(inner_dim, 1L) + self$av_cross_attn_audio_v2a_gate <- ltx23_ada_layer_norm_single(audio_inner_dim, 1L) + + # Output layer modulation + self$scale_shift_table <- torch::nn_parameter( + torch::torch_randn(2L, inner_dim) / sqrt(inner_dim) + ) + self$audio_scale_shift_table <- torch::nn_parameter( + torch::torch_randn(2L, audio_inner_dim) / sqrt(audio_inner_dim) + ) + + # Prompt modulation (LTX-2.3) + if (self$prompt_modulation) { + self$prompt_adaln <- ltx23_ada_layer_norm_single(inner_dim, 2L) + self$audio_prompt_adaln <- ltx23_ada_layer_norm_single(audio_inner_dim, 2L) + } + + # 3. RoPE embedders (self-attention per modality + a2v/v2a cross) + self$rope <- ltx23_rotary_pos_embed( + dim = inner_dim, patch_size = patch_size, patch_size_t = patch_size_t, + base_num_frames = pos_embed_max_pos, base_height = base_height, + base_width = base_width, scale_factors = vae_scale_factors, + theta = rope_theta, causal_offset = causal_offset, modality = "video", + double_precision = rope_double_precision, rope_type = rope_type, + num_attention_heads = num_attention_heads + ) + self$audio_rope <- ltx23_rotary_pos_embed( + dim = audio_inner_dim, patch_size = audio_patch_size, + patch_size_t = audio_patch_size_t, base_num_frames = audio_pos_embed_max_pos, + sampling_rate = audio_sampling_rate, hop_length = audio_hop_length, + scale_factors = audio_scale_factor, + theta = rope_theta, causal_offset = causal_offset, modality = "audio", + double_precision = rope_double_precision, rope_type = rope_type, + num_attention_heads = audio_num_attention_heads + ) + ca_max_pos <- max(pos_embed_max_pos, audio_pos_embed_max_pos) + self$cross_attn_rope <- ltx23_rotary_pos_embed( + dim = audio_cross_attention_dim, patch_size = patch_size, + patch_size_t = patch_size_t, base_num_frames = ca_max_pos, + base_height = base_height, base_width = base_width, + theta = rope_theta, causal_offset = causal_offset, modality = "video", + double_precision = rope_double_precision, rope_type = rope_type, + num_attention_heads = num_attention_heads + ) + self$cross_attn_audio_rope <- ltx23_rotary_pos_embed( + dim = audio_cross_attention_dim, patch_size = audio_patch_size, + patch_size_t = audio_patch_size_t, base_num_frames = ca_max_pos, + sampling_rate = audio_sampling_rate, hop_length = audio_hop_length, + theta = rope_theta, causal_offset = causal_offset, modality = "audio", + double_precision = rope_double_precision, rope_type = rope_type, + num_attention_heads = audio_num_attention_heads + ) + + # 4. Transformer blocks + self$transformer_blocks <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { + ltx23_transformer_block( + dim = inner_dim, + num_attention_heads = num_attention_heads, + attention_head_dim = attention_head_dim, + cross_attention_dim = cross_attention_dim, + audio_dim = audio_inner_dim, + audio_num_attention_heads = audio_num_attention_heads, + audio_attention_head_dim = audio_attention_head_dim, + audio_cross_attention_dim = audio_cross_attention_dim, + video_gated_attn = gated_attn, + video_cross_attn_adaln = cross_attn_mod, + audio_gated_attn = audio_gated_attn, + audio_cross_attn_adaln = audio_cross_attn_mod, + eps = norm_eps, + elementwise_affine = FALSE, + rope_type = rope_type, + perturbed_attn = perturbed_attn + ) + })) + + # 5. Output layers + self$norm_out <- torch::nn_layer_norm(inner_dim, eps = 1e-6, elementwise_affine = FALSE) + self$proj_out <- torch::nn_linear(inner_dim, out_channels) + self$audio_norm_out <- torch::nn_layer_norm(audio_inner_dim, eps = 1e-6, + elementwise_affine = FALSE) + self$audio_proj_out <- torch::nn_linear(audio_inner_dim, audio_out_channels) +}, + forward = function( + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + timestep, + audio_timestep = NULL, + sigma = NULL, + audio_sigma = NULL, + encoder_attention_mask = NULL, + audio_encoder_attention_mask = NULL, + num_frames = NULL, + height = NULL, + width = NULL, + fps = 24.0, + audio_num_frames = NULL, + video_coords = NULL, + audio_coords = NULL, + isolate_modalities = FALSE, + spatio_temporal_guidance_blocks = NULL, + perturbation_mask = NULL, + use_cross_timestep = FALSE, + video_self_attention_mask = NULL + ) { + audio_timestep <- audio_timestep %||% timestep + audio_sigma <- audio_sigma %||% sigma + if (self$prompt_modulation && is.null(sigma)) { + stop("sigma is required for LTX-2.3 prompt modulation") + } + + # Multiplicative [B, S] masks -> additive bias [B, 1, S] + if (!is.null(encoder_attention_mask) && encoder_attention_mask$ndim == 2L) { + encoder_attention_mask <- encoder_attention_mask$to( + dtype = hidden_states$dtype + )$neg()$add(1)$mul(-10000.0)$unsqueeze(2L) + } + if (!is.null(audio_encoder_attention_mask) && + audio_encoder_attention_mask$ndim == 2L) { + audio_encoder_attention_mask <- audio_encoder_attention_mask$to( + dtype = audio_hidden_states$dtype + )$neg()$add(1)$mul(-10000.0)$unsqueeze(2L) + } + if (!is.null(video_self_attention_mask)) { + video_self_attention_mask <- video_self_attention_mask$to( + dtype = hidden_states$dtype + )$neg()$add(1)$mul(-10000.0) + } + + batch_size <- hidden_states$shape[1] + + # 1. RoPE frequencies + if (is.null(video_coords)) { + video_coords <- self$rope$prepare_video_coords( + batch_size, num_frames, height, width, hidden_states$device, fps = fps + ) + } + if (is.null(audio_coords)) { + audio_coords <- self$audio_rope$prepare_audio_coords( + batch_size, audio_num_frames, audio_hidden_states$device + ) + } + video_rotary_emb <- self$rope(video_coords, device = hidden_states$device) + audio_rotary_emb <- self$audio_rope(audio_coords, device = audio_hidden_states$device) + # Cross-modal attention uses the temporal axis only + video_ca_rotary_emb <- self$cross_attn_rope( + video_coords$narrow(2L, 1L, 1L), device = hidden_states$device + ) + audio_ca_rotary_emb <- self$cross_attn_audio_rope( + audio_coords$narrow(2L, 1L, 1L), device = audio_hidden_states$device + ) + + # 2. Input projections + hidden_states <- self$proj_in(hidden_states) + audio_hidden_states <- self$audio_proj_in(audio_hidden_states) + + # 3. Timestep embeddings and global modulation parameters + gate_scale_factor <- self$cross_attn_timestep_scale_multiplier / + self$timestep_scale_multiplier + + te <- self$time_embed(timestep$flatten(), hidden_dtype = hidden_states$dtype) + temb <- te[[1]]$view(c(batch_size, -1L, te[[1]]$shape[2])) + embedded_timestep <- te[[2]]$view(c(batch_size, -1L, te[[2]]$shape[2])) + + ate <- self$audio_time_embed(audio_timestep$flatten(), + hidden_dtype = audio_hidden_states$dtype) + temb_audio <- ate[[1]]$view(c(batch_size, -1L, ate[[1]]$shape[2])) + audio_embedded_timestep <- ate[[2]]$view(c(batch_size, -1L, ate[[2]]$shape[2])) + + if (self$prompt_modulation) { + pe <- self$prompt_adaln(sigma$flatten(), hidden_dtype = hidden_states$dtype) + temb_prompt <- pe[[1]]$view(c(batch_size, -1L, pe[[1]]$shape[2])) + ape <- self$audio_prompt_adaln(audio_sigma$flatten(), + hidden_dtype = audio_hidden_states$dtype) + temb_prompt_audio <- ape[[1]]$view(c(batch_size, -1L, ape[[1]]$shape[2])) + } else { + temb_prompt <- NULL + temb_prompt_audio <- NULL + } + + # a2v/v2a modulation; 2.3 modulates each modality by the *other* + # modality's sigma (use_cross_timestep) + video_ca_timestep <- if (use_cross_timestep) audio_sigma$flatten() else timestep$flatten() + vcss <- self$av_cross_attn_video_scale_shift( + video_ca_timestep, hidden_dtype = hidden_states$dtype + )[[1]] + vcg <- self$av_cross_attn_video_a2v_gate( + video_ca_timestep * gate_scale_factor, hidden_dtype = hidden_states$dtype + )[[1]] + video_ca_scale_shift <- vcss$view(c(batch_size, -1L, vcss$shape[2])) + video_ca_a2v_gate <- vcg$view(c(batch_size, -1L, vcg$shape[2])) + + audio_ca_timestep <- if (use_cross_timestep) sigma$flatten() else audio_timestep$flatten() + acss <- self$av_cross_attn_audio_scale_shift( + audio_ca_timestep, hidden_dtype = audio_hidden_states$dtype + )[[1]] + acg <- self$av_cross_attn_audio_v2a_gate( + audio_ca_timestep * gate_scale_factor, hidden_dtype = audio_hidden_states$dtype + )[[1]] + audio_ca_scale_shift <- acss$view(c(batch_size, -1L, acss$shape[2])) + audio_ca_v2a_gate <- acg$view(c(batch_size, -1L, acg$shape[2])) + + # 4. STG perturbation setup + stg_blocks <- spatio_temporal_guidance_blocks %||% integer(0) + if (length(stg_blocks) > 0L && is.null(perturbation_mask)) { + perturbation_mask <- torch::torch_zeros(batch_size, device = hidden_states$device) + } + if (!is.null(perturbation_mask) && perturbation_mask$ndim == 1L) { + perturbation_mask <- perturbation_mask$unsqueeze(2L)$unsqueeze(3L) + } + all_perturbed <- if (!is.null(perturbation_mask)) { + as.logical((perturbation_mask == 0)$all()$item()) + } else { + FALSE + } + + # 5. Transformer blocks + # Compiled path: the entire block stack in one TorchScript call (no + # per-op R dispatch, no R-side garbage, fused SDPA, no per-block + # gc). Requires the NF4-quantized 2.3 configuration and none of the + # STG/masking extras the script doesn't model. + use_jit <- isTRUE(getOption("diffuseR.jit_blocks", TRUE)) && + length(spatio_temporal_guidance_blocks %||% integer(0)) == 0L && + is.null(perturbation_mask) && + is.null(video_self_attention_mask) && + !isolate_modalities && + self$prompt_modulation && + !is.null(video_rotary_emb) && !is.null(audio_rotary_emb) && + !is.null(video_ca_rotary_emb) && !is.null(audio_ca_rotary_emb) && + .ltx23_jit_block_ok(self$transformer_blocks[[1]]) + + if (use_jit) { + res <- .ltx23_jit_run_stack( + self$transformer_blocks, + hidden_states, audio_hidden_states, + encoder_hidden_states, audio_encoder_hidden_states, + temb, temb_audio, + video_ca_scale_shift, audio_ca_scale_shift, + video_ca_a2v_gate, audio_ca_v2a_gate, + temb_prompt, temb_prompt_audio, + video_rotary_emb, audio_rotary_emb, + video_ca_rotary_emb, audio_ca_rotary_emb, + encoder_attention_mask, + audio_encoder_attention_mask + ) + hidden_states <- res[[1]] + audio_hidden_states <- res[[2]] + } + + block_gc <- isTRUE(getOption("diffuseR.block_gc")) + 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 + res <- self$transformer_blocks[[block_idx]]( + hidden_states = hidden_states, + audio_hidden_states = audio_hidden_states, + encoder_hidden_states = encoder_hidden_states, + audio_encoder_hidden_states = audio_encoder_hidden_states, + temb = temb, + temb_audio = temb_audio, + temb_ca_scale_shift = video_ca_scale_shift, + temb_ca_audio_scale_shift = audio_ca_scale_shift, + temb_ca_gate = video_ca_a2v_gate, + temb_ca_audio_gate = audio_ca_v2a_gate, + temb_prompt = temb_prompt, + temb_prompt_audio = temb_prompt_audio, + video_rotary_emb = video_rotary_emb, + audio_rotary_emb = audio_rotary_emb, + ca_video_rotary_emb = video_ca_rotary_emb, + ca_audio_rotary_emb = audio_ca_rotary_emb, + encoder_attention_mask = encoder_attention_mask, + audio_encoder_attention_mask = audio_encoder_attention_mask, + self_attention_mask = video_self_attention_mask, + use_a2v_cross_attention = !isolate_modalities, + use_v2a_cross_attention = !isolate_modalities, + perturbation_mask = if (is_stg_block) perturbation_mask else NULL, + all_perturbed = if (is_stg_block) all_perturbed else FALSE + ) + hidden_states <- res[[1]] + audio_hidden_states <- res[[2]] + if (isTRUE(getOption("diffuseR.debug"))) { + ms <- torch::cuda_memory_stats() + message(sprintf(" block %d: %.2f GB allocated, %.2f GB reserved", + block_idx, + ms$allocated_bytes$all$current / 1e9, + ms$reserved_bytes$all$current / 1e9)) + } + if (block_gc) { + # Only needed when the quantized linears allocate per call + # (buffered scratch makes per-block garbage negligible) + gc(verbose = FALSE) + } + } + + # 6. Output layers + scale_shift_values <- self$scale_shift_table$unsqueeze(1L)$unsqueeze(1L) + + embedded_timestep$unsqueeze(3L) + shift <- scale_shift_values[,, 1,] + scale <- scale_shift_values[,, 2,] + hidden_states <- self$norm_out(hidden_states) + hidden_states <- hidden_states * scale$add(1) + shift + output <- self$proj_out(hidden_states) + + audio_scale_shift_values <- self$audio_scale_shift_table$unsqueeze(1L)$unsqueeze(1L) + + audio_embedded_timestep$unsqueeze(3L) + audio_shift <- audio_scale_shift_values[,, 1,] + audio_scale <- audio_scale_shift_values[,, 2,] + audio_hidden_states <- self$audio_norm_out(audio_hidden_states) + audio_hidden_states <- audio_hidden_states * audio_scale$add(1) + audio_shift + audio_output <- self$audio_proj_out(audio_hidden_states) + + list(sample = output, audio_sample = audio_output) +} +) + +#' Map an official DiT checkpoint key to the R module name +#' +#' Applies the official-to-diffusers renames for the LTX-2.3 transformer +#' (cf. diffusers scripts/convert_ltx2_to_diffusers.py). Our module tree +#' matches the diffusers names, so this is the full mapping. +#' +#' @param key Character. Checkpoint key (with or without the +#' \code{model.diffusion_model.} prefix). +#' +#' @return Character. Module parameter/buffer name. +#' +#' @export +ltx23_map_dit_key <- function(key) { + key <- sub("^model\\.diffusion_model\\.", "", key) + + # Substring renames; *_adaln_single variants must precede the bare + # adaln_single prefix renames below + key <- gsub("av_ca_video_scale_shift_adaln_single", + "av_cross_attn_video_scale_shift", key, fixed = TRUE) + key <- gsub("av_ca_a2v_gate_adaln_single", "av_cross_attn_video_a2v_gate", key, fixed = TRUE) + key <- gsub("av_ca_audio_scale_shift_adaln_single", "av_cross_attn_audio_scale_shift", key, fixed = TRUE) + key <- gsub("av_ca_v2a_gate_adaln_single", "av_cross_attn_audio_v2a_gate", key, fixed = TRUE) + key <- gsub("scale_shift_table_a2v_ca_video", "video_a2v_cross_attn_scale_shift_table", key, fixed = TRUE) + key <- gsub("scale_shift_table_a2v_ca_audio", "audio_a2v_cross_attn_scale_shift_table", key, fixed = TRUE) + key <- gsub("prompt_adaln_single", "prompt_adaln", key, fixed = TRUE) + key <- sub("^audio_adaln_single\\.", "audio_time_embed.", key) + key <- sub("^adaln_single\\.", "time_embed.", key) + key <- gsub("patchify_proj", "proj_in", key, fixed = TRUE) + key <- gsub("q_norm", "norm_q", key, fixed = TRUE) + key <- gsub("k_norm", "norm_k", key, fixed = TRUE) + key +} diff --git a/R/dit_ltx23_modules.R b/R/dit_ltx23_modules.R new file mode 100644 index 0000000..25e71df --- /dev/null +++ b/R/dit_ltx23_modules.R @@ -0,0 +1,780 @@ +#' LTX-2.3 Transformer Building Blocks +#' +#' Fresh R port of the LTX-2 transformer components from the diffusers +#' reference (Apache-2.0, src/diffusers/models/transformers/ +#' transformer_ltx2.py and shared modules). Field names mirror the +#' diffusers module tree so checkpoint keys map 1:1. +#' +#' @name dit_ltx23_modules +NULL + +#' RMS normalization +#' +#' Variance is computed in float32; the result is cast back to the input +#' dtype (or the weight dtype when elementwise affine). +#' +#' @param dim Integer. Normalized dimension size. +#' @param eps Numeric. Stability epsilon. +#' @param elementwise_affine Logical. Learn a scale weight. +#' +#' @export +ltx23_rms_norm <- torch::nn_module( + "ltx23_rms_norm", + initialize = function(dim, eps = 1e-6, elementwise_affine = TRUE) { + self$eps <- eps + self$elementwise_affine <- elementwise_affine + if (elementwise_affine) { + self$weight <- torch::nn_parameter(torch::torch_ones(dim)) + } +}, + forward = function(x) { + input_dtype <- x$dtype + variance <- x$to(dtype = torch::torch_float32())$pow(2)$mean(dim = -1L, + keepdim = TRUE) + x <- x * torch::torch_rsqrt(variance + self$eps) + if (self$elementwise_affine) { + x <- x$to(dtype = self$weight$dtype) * self$weight + } else { + x <- x$to(dtype = input_dtype) + } + x +} +) + +#' Sinusoidal timestep embedding +#' +#' DDPM-style sinusoidal embedding. LTX uses \code{flip_sin_to_cos=TRUE} +#' (cos first) and \code{downscale_freq_shift=0}. +#' +#' @param timesteps 1D tensor of timestep values. +#' @param embedding_dim Integer. Output embedding size. +#' @param flip_sin_to_cos Logical. Put cos before sin. +#' @param downscale_freq_shift Numeric. Frequency delta control. +#' @param max_period Numeric. Maximum embedding frequency period. +#' +#' @return Tensor [N, embedding_dim]. +#' +#' @export +ltx23_get_timestep_embedding <- function(timesteps, embedding_dim, + flip_sin_to_cos = TRUE, + downscale_freq_shift = 0, + max_period = 10000) { + stopifnot(timesteps$ndim == 1L) + half_dim <- embedding_dim %/% 2L + exponent <- -log(max_period) * torch::torch_arange(start = 0, + end = half_dim - 1, dtype = torch::torch_float32(), + device = timesteps$device) + exponent <- exponent / (half_dim - downscale_freq_shift) + + emb <- torch::torch_exp(exponent) + emb <- timesteps$unsqueeze(2L)$to(dtype = torch::torch_float32()) * emb$unsqueeze(1L) + emb <- torch::torch_cat(list(torch::torch_sin(emb), torch::torch_cos(emb)), dim = -1L) + + if (flip_sin_to_cos) { + emb <- torch::torch_cat( + list(emb[, (half_dim + 1):(2 * half_dim)], emb[, 1:half_dim]), + dim = -1L + ) + } + if (embedding_dim %% 2L == 1L) { + emb <- torch::nnf_pad(emb, c(0L, 1L, 0L, 0L)) + } + emb +} + +# Two-layer timestep MLP (linear_1 -> silu -> linear_2), matching +# diffusers TimestepEmbedding state-dict names. +ltx23_timestep_embedding <- torch::nn_module( + "ltx23_timestep_embedding", + initialize = function(in_channels, time_embed_dim) { + self$linear_1 <- torch::nn_linear(in_channels, time_embed_dim) + self$linear_2 <- torch::nn_linear(time_embed_dim, time_embed_dim) +}, + forward = function(sample) { + self$linear_2(torch::nnf_silu(self$linear_1(sample))) +} +) + +# Sinusoidal projection (256 channels) + MLP, matching diffusers +# PixArtAlphaCombinedTimestepSizeEmbeddings without additional conditions. +ltx23_combined_timestep_embed <- torch::nn_module( + "ltx23_combined_timestep_embed", + initialize = function(embedding_dim) { + self$timestep_embedder <- ltx23_timestep_embedding(256L, embedding_dim) +}, + forward = function(timestep, hidden_dtype) { + proj <- ltx23_get_timestep_embedding(timestep, 256L, + flip_sin_to_cos = TRUE, downscale_freq_shift = 0) + self$timestep_embedder(proj$to(dtype = hidden_dtype)) +} +) + +#' Adaptive layer norm single (adaLN-single) +#' +#' Embeds a timestep/sigma and projects it to a configurable number of +#' modulation parameter vectors. +#' +#' @param embedding_dim Integer. Model dimension. +#' @param num_mod_params Integer. Number of modulation parameter vectors. +#' +#' @return Module whose forward returns +#' \code{list(mod_params [N, num_mod_params * dim], embedded_timestep [N, dim])}. +#' +#' @export +ltx23_ada_layer_norm_single <- torch::nn_module( + "ltx23_ada_layer_norm_single", + initialize = function(embedding_dim, num_mod_params = 6L) { + self$num_mod_params <- num_mod_params + self$emb <- ltx23_combined_timestep_embed(embedding_dim) + self$linear <- torch::nn_linear(embedding_dim, + num_mod_params * embedding_dim, + bias = TRUE) +}, + forward = function(timestep, hidden_dtype) { + embedded_timestep <- self$emb(timestep, hidden_dtype = hidden_dtype) + list(self$linear(torch::nnf_silu(embedded_timestep)), embedded_timestep) +} +) + +# torch_matmul_out / torch_softmax_out are unexported torch internals; +# resolve them once and fall back to allocating ops if they disappear +.ltx23_attn_out_fns <- local({ + fns <- NULL + function() { + if (is.null(fns)) { + fns <<- tryCatch( + list( + matmul = get("torch_matmul_out", + envir = asNamespace("torch")), + softmax = get("torch_softmax_out", envir = asNamespace("torch")) + ), + error = function(e) FALSE + ) + } + fns + } +}) + +# Persistent attention scratch (score matrix + context output), keyed by +# shape/dtype/device like the dequant buffers: attention temporaries are +# the dominant per-block garbage at high resolution +.ltx23_attn_scratch <- new.env(parent = emptyenv()) + +.ltx23_get_attn_buffer <- function(shape, dtype, device) { + key <- paste(paste(shape, collapse = "x"), dtype$.type(), + paste(device$type, device$index %||% 0L), sep = "|") + buf <- .ltx23_attn_scratch[[key]] + if (is.null(buf)) { + buf <- torch::torch_empty(shape, dtype = dtype, device = device) + .ltx23_attn_scratch[[key]] <- buf + } + buf +} + +# Scaled dot-product attention on [B, H, S, D] tensors. R torch has no +# fused SDPA, so the [B, H, Sq, Sk] score matrix materializes; queries +# are chunked adaptively against a memory budget and all large +# temporaries live in reusable scratch buffers. +.ltx23_sdpa <- function(query, key, value, attention_mask = NULL, + chunk_size = NULL) { + head_dim <- query$shape[length(query$shape)] + scale <- 1.0 / sqrt(head_dim) + key_t <- key$transpose(-2L, -1L) + fns <- .ltx23_attn_out_fns() + + b <- query$shape[1] + heads <- query$shape[2] + n_q <- query$shape[3] + n_k <- key$shape[3] + d_v <- value$shape[4] + + attend <- function(q, mask) { + attn <- torch::torch_matmul(q$mul(scale), key_t) + if (!is.null(mask)) { + attn <- attn + mask + } + attn <- torch::nnf_softmax(attn, dim = -1L) + torch::torch_matmul(attn, value) + } + + # Buffered variant: score matrix and context reuse persistent scratch + attend_into <- function(q, mask, attn_buf, out_buf) { + fns$matmul(attn_buf, q$mul(scale), key_t) + if (!is.null(mask)) { + attn_buf$add_(mask) + } + fns$softmax(attn_buf, attn_buf, dim = -1L, dtype = attn_buf$dtype) + fns$matmul(out_buf, attn_buf, value) + out_buf + } + + # Adaptive query chunking: bound the [B, H, chunk, S_k] score matrix + # to a memory budget regardless of sequence length. The default is + # sized so the self+cross scratch pair stays under ~400MB at + # 1280x704x121 (14080 video tokens), the validated ceiling of the + # 16GB NF4-resident profile; below ~2000 tokens it has no effect. + budget <- getOption("diffuseR.attn_budget", 1.5e8) + auto_chunk <- max(256L, as.integer(budget / (b * heads * n_k * 2))) + if (is.null(chunk_size)) { + chunk_size <- auto_chunk + } else { + chunk_size <- min(chunk_size, auto_chunk) + } + + if (isFALSE(fns)) { + # Fallback: allocating path + if (n_q <= chunk_size) { + return(attend(query, attention_mask)) + } + outs <- list() + start <- 1L + while (start <= n_q) { + len <- min(chunk_size, n_q - start + 1L) + q_chunk <- query$narrow(3L, start, len) + mask_chunk <- attention_mask + if (!is.null(mask_chunk) && mask_chunk$shape[3] > 1L) { + mask_chunk <- mask_chunk$narrow(3L, start, len) + } + outs[[length(outs) + 1L]] <- attend(q_chunk, mask_chunk) + start <- start + len + } + return(torch::torch_cat(outs, dim = 3L)) + } + + rows <- min(n_q, chunk_size) + attn_buf <- .ltx23_get_attn_buffer(c(b, heads, rows, n_k), query$dtype, + query$device) + out_buf <- .ltx23_get_attn_buffer(c(b, heads, n_q, d_v), query$dtype, + query$device) + + if (n_q <= chunk_size) { + return(attend_into(query, attention_mask, attn_buf, out_buf)) + } + + start <- 1L + while (start <= n_q) { + len <- min(chunk_size, n_q - start + 1L) + q_chunk <- query$narrow(3L, start, len) + mask_chunk <- attention_mask + if (!is.null(mask_chunk) && mask_chunk$shape[3] > 1L) { + mask_chunk <- mask_chunk$narrow(3L, start, len) + } + attend_into( + q_chunk, mask_chunk, + attn_buf$narrow(3L, 1L, len), + out_buf$narrow(3L, start, len) + ) + start <- start + len + } + out_buf +} + +#' Release the attention scratch buffers +#' +#' @return Invisibly, NULL. +#' +#' @keywords internal +.ltx23_release_attn_buffers <- function() { + rm(list = ls(.ltx23_attn_scratch), envir = .ltx23_attn_scratch) + invisible(NULL) +} + +#' LTX-2 attention layer +#' +#' Attention with RMS q/k norms across heads, optional per-head output +#' gating (LTX-2.3), separate query/key RoPE (for a2v/v2a cross +#' attention), and optional STG perturbation (skip attention, use the +#' value projection). +#' +#' @param query_dim Integer. Query feature dimension. +#' @param heads,kv_heads Integers. Attention head counts. +#' @param dim_head Integer. Per-head dimension. +#' @param bias,out_bias Logicals. Projection biases. +#' @param cross_attention_dim Integer or NULL. Key/value input dimension +#' (NULL for self-attention). +#' @param norm_eps Numeric. RMS norm epsilon. +#' @param norm_elementwise_affine Logical. RMS norms carry weights. +#' @param rope_type "split" or "interleaved". +#' @param apply_gated_attention Logical. Add per-head sigmoid output gates. +#' +#' @export +ltx23_attention <- torch::nn_module( + "ltx23_attention", + initialize = function( + query_dim, + heads = 8L, + kv_heads = NULL, + dim_head = 64L, + bias = TRUE, + cross_attention_dim = NULL, + out_bias = TRUE, + norm_eps = 1e-6, + norm_elementwise_affine = TRUE, + rope_type = "split", + apply_gated_attention = FALSE + ) { + kv_heads <- kv_heads %||% heads + self$heads <- as.integer(heads) + self$head_dim <- as.integer(dim_head) + inner_dim <- dim_head * heads + inner_kv_dim <- dim_head * kv_heads + cross_dim <- cross_attention_dim %||% query_dim + self$rope_type <- rope_type + self$attn_chunk <- NULL # set by memory profiles; NULL = unchunked + + self$norm_q <- ltx23_rms_norm(inner_dim, eps = norm_eps, + elementwise_affine = norm_elementwise_affine) + self$norm_k <- ltx23_rms_norm(inner_kv_dim, eps = norm_eps, + elementwise_affine = norm_elementwise_affine) + self$to_q <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$to_k <- torch::nn_linear(cross_dim, inner_kv_dim, bias = bias) + self$to_v <- torch::nn_linear(cross_dim, inner_kv_dim, bias = bias) + self$to_out <- torch::nn_module_list(list( + torch::nn_linear(inner_dim, query_dim, bias = out_bias) + )) + + self$apply_gated_attention <- apply_gated_attention + if (apply_gated_attention) { + # Per-head gate logits, computed on the raw block input (pre-Q) + self$to_gate_logits <- torch::nn_linear(query_dim, heads, bias = TRUE) + } +}, + forward = function( + hidden_states, + encoder_hidden_states = NULL, + attention_mask = NULL, + query_rotary_emb = NULL, + key_rotary_emb = NULL, + perturbation_mask = NULL, + all_perturbed = FALSE + ) { + if (is.null(encoder_hidden_states)) { + encoder_hidden_states <- hidden_states + } + if (!is.null(attention_mask)) { + # [B, 1, S] or [B, Tq, S] additive bias -> [B, 1, *, S] + if (attention_mask$ndim == 3L) { + attention_mask <- attention_mask$unsqueeze(2L) + } + } + + gate_logits <- NULL + if (self$apply_gated_attention) { + gate_logits <- self$to_gate_logits(hidden_states) # [B, T, H] + } + + value_flat <- self$to_v(encoder_hidden_states) + + if (isTRUE(all_perturbed)) { + # STG: skip attention entirely, use the value projection + out <- value_flat + } else { + query <- self$norm_q(self$to_q(hidden_states)) + key <- self$norm_k(self$to_k(encoder_hidden_states)) + + key_rope <- key_rotary_emb %||% query_rotary_emb + if (!is.null(query_rotary_emb)) { + if (self$rope_type == "interleaved") { + query <- ltx23_apply_interleaved_rotary_emb(query, query_rotary_emb) + key <- ltx23_apply_interleaved_rotary_emb(key, key_rope) + } else { + query <- ltx23_apply_split_rotary_emb(query, query_rotary_emb) + key <- ltx23_apply_split_rotary_emb(key, key_rope) + } + } + + # [B, S, H*D] -> [B, H, S, D] + query <- query$unflatten(3L, c(self$heads, -1L))$transpose(2L, 3L) + key <- key$unflatten(3L, c(self$heads, -1L))$transpose(2L, 3L) + value <- value_flat$unflatten(3L, c(self$heads, -1L))$transpose(2L, 3L) + + out <- .ltx23_sdpa(query, key, value, attention_mask, self$attn_chunk) + out <- out$transpose(2L, 3L)$flatten(start_dim = 3L) + out <- out$to(dtype = hidden_states$dtype) + + if (!is.null(perturbation_mask)) { + # Interpolate between the perturbed (value) and full attention paths + out <- torch::torch_lerp(value_flat, out, perturbation_mask) + } + } + + if (self$apply_gated_attention) { + out <- out$unflatten(3L, c(self$heads, -1L)) # [B, T, H, D] + # Factor 2 so zero-initialized gate logits give unit gates + gates <- gate_logits$sigmoid()$mul(2.0) # [B, T, H] + out <- out * gates$unsqueeze(-1L) + out <- out$flatten(start_dim = 3L) + } + + self$to_out[[1]](out) +} +) + +# In-place tanh GELU for large activations: nnf_gelu allocates a second +# copy of its input, which at high resolution doubles the feed-forward +# intermediate (the single largest transient in a block). Uses the +# internal in-place kernel when available, chunked so any internal +# temporaries stay small; falls back to the allocating path. +.ltx23_gelu_tanh_inplace <- function(x, chunk_elements = 2 ^ 24) { + fn <- get0("torch_gelu_", envir = asNamespace("torch")) + if (is.null(fn)) { + return(torch::nnf_gelu(x, approximate = "tanh")) + } + flat <- x$view(-1L) + n <- flat$numel() + start <- 1 + while (start <= n) { + len <- min(chunk_elements, n - start + 1) + fn(flat$narrow(1L, start, len), approximate = "tanh") + start <- start + len + } + x +} + +#' LTX feed-forward layer +#' +#' Linear -> GELU (tanh approximation) -> Linear with 4x hidden dim, +#' matching diffusers \code{FeedForward(activation_fn="gelu-approximate")} +#' state-dict names (net.0.proj, net.2). +#' +#' @param dim Integer. Input/output dimension. +#' @param mult Integer. Hidden dimension multiplier. +#' +#' @export +ltx23_feed_forward <- torch::nn_module( + "ltx23_feed_forward", + initialize = function(dim, mult = 4L) { + inner_dim <- as.integer(dim * mult) + gelu_proj <- torch::nn_module( + "ltx23_gelu", + initialize = function(dim_in, dim_out) { + self$proj <- torch::nn_linear(dim_in, dim_out) + }, + forward = function(x) { + h <- self$proj(x) + if (h$numel() > getOption("diffuseR.gelu_inplace_min", 1e8)) { + return(.ltx23_gelu_tanh_inplace(h)) + } + torch::nnf_gelu(h, approximate = "tanh") + } + ) + self$net <- torch::nn_module_list(list( + gelu_proj(dim, inner_dim), + torch::nn_identity(), # dropout slot in the reference; inference no-op + torch::nn_linear(inner_dim, dim) + )) +}, + forward = function(x) { + self$net[[3]](self$net[[1]](x)) +} +) + +# Combine a per-block scale/shift table with global modulation params: +# ada_values = table[None, None] + temb.reshape(B, temb_tokens, num, -1), +# returned as a list split along the parameter axis. +ltx23_get_mod_params <- function(scale_shift_table, temb, batch_size) { + num_params <- scale_shift_table$shape[1] + ada_values <- scale_shift_table$unsqueeze(1L)$unsqueeze(1L)$to(device = temb$device) + + temb$reshape(c(batch_size, temb$shape[2], num_params, -1L)) + torch::torch_unbind(ada_values, dim = 3L) +} + +#' LTX-2 transformer block +#' +#' Dual-stream (video + audio) block: modulated self-attention per +#' modality, text cross-attention per modality (with LTX-2.3 query and +#' key/value modulation), bidirectional audio-video cross-attention with +#' global+per-block modulation, and modulated feed-forward. +#' +#' @param dim,audio_dim Integers. Video/audio stream dimensions. +#' @param num_attention_heads,attention_head_dim Video attention shape. +#' @param cross_attention_dim Integer. Text embedding dim for video. +#' @param audio_num_attention_heads,audio_attention_head_dim Audio attention shape. +#' @param audio_cross_attention_dim Integer. Text embedding dim for audio. +#' @param video_gated_attn,audio_gated_attn Logicals. Per-head output gates. +#' @param video_cross_attn_adaln,audio_cross_attn_adaln Logicals. LTX-2.3 +#' text cross-attention modulation (9 mod params instead of 6). +#' @param eps Numeric. Norm epsilon. +#' @param elementwise_affine Logical. Block norms carry weights (FALSE for LTX). +#' @param rope_type "split" or "interleaved". +#' @param perturbed_attn Logical. Enable the STG perturbation arguments. +#' +#' @export +ltx23_transformer_block <- torch::nn_module( + "ltx23_transformer_block", + initialize = function( + dim, + num_attention_heads, + attention_head_dim, + cross_attention_dim, + audio_dim, + audio_num_attention_heads, + audio_attention_head_dim, + audio_cross_attention_dim, + video_gated_attn = TRUE, + video_cross_attn_adaln = TRUE, + audio_gated_attn = TRUE, + audio_cross_attn_adaln = TRUE, + eps = 1e-6, + elementwise_affine = FALSE, + rope_type = "split", + perturbed_attn = TRUE + ) { + self$perturbed_attn <- perturbed_attn + + # 1. Self-attention (video and audio) + self$norm1 <- ltx23_rms_norm(dim, eps = eps, + elementwise_affine = elementwise_affine) + self$attn1 <- ltx23_attention( + query_dim = dim, heads = num_attention_heads, dim_head = attention_head_dim, + rope_type = rope_type, apply_gated_attention = video_gated_attn + ) + self$audio_norm1 <- ltx23_rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) + self$audio_attn1 <- ltx23_attention( + query_dim = audio_dim, heads = audio_num_attention_heads, + dim_head = audio_attention_head_dim, + rope_type = rope_type, apply_gated_attention = audio_gated_attn + ) + + # 2. Prompt cross-attention + self$norm2 <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) + self$attn2 <- ltx23_attention( + query_dim = dim, cross_attention_dim = cross_attention_dim, + heads = num_attention_heads, dim_head = attention_head_dim, + rope_type = rope_type, apply_gated_attention = video_gated_attn + ) + self$audio_norm2 <- ltx23_rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) + self$audio_attn2 <- ltx23_attention( + query_dim = audio_dim, cross_attention_dim = audio_cross_attention_dim, + heads = audio_num_attention_heads, dim_head = audio_attention_head_dim, + rope_type = rope_type, apply_gated_attention = audio_gated_attn + ) + + # 3. Audio-to-video (Q: video) and video-to-audio (Q: audio) cross-attention + self$audio_to_video_norm <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) + self$audio_to_video_attn <- ltx23_attention( + query_dim = dim, cross_attention_dim = audio_dim, + heads = audio_num_attention_heads, dim_head = audio_attention_head_dim, + rope_type = rope_type, apply_gated_attention = video_gated_attn + ) + self$video_to_audio_norm <- ltx23_rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) + self$video_to_audio_attn <- ltx23_attention( + query_dim = audio_dim, cross_attention_dim = dim, + heads = audio_num_attention_heads, dim_head = audio_attention_head_dim, + rope_type = rope_type, apply_gated_attention = audio_gated_attn + ) + + # 4. Feed-forward + self$norm3 <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) + self$ff <- ltx23_feed_forward(dim) + self$audio_norm3 <- ltx23_rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) + self$audio_ff <- ltx23_feed_forward(audio_dim) + + # 5. Per-block modulation tables + self$video_cross_attn_adaln <- video_cross_attn_adaln + self$audio_cross_attn_adaln <- audio_cross_attn_adaln + self$cross_attn_adaln <- video_cross_attn_adaln || audio_cross_attn_adaln + video_mod_params <- if (video_cross_attn_adaln) 9L else 6L + audio_mod_params <- if (audio_cross_attn_adaln) 9L else 6L + self$scale_shift_table <- torch::nn_parameter( + torch::torch_randn(video_mod_params, dim) / sqrt(dim) + ) + self$audio_scale_shift_table <- torch::nn_parameter( + torch::torch_randn(audio_mod_params, audio_dim) / sqrt(audio_dim) + ) + if (self$cross_attn_adaln) { + self$prompt_scale_shift_table <- torch::nn_parameter(torch::torch_randn(2L, dim)) + self$audio_prompt_scale_shift_table <- torch::nn_parameter(torch::torch_randn(2L, audio_dim)) + } + self$video_a2v_cross_attn_scale_shift_table <- torch::nn_parameter(torch::torch_randn(5L, dim)) + self$audio_a2v_cross_attn_scale_shift_table <- torch::nn_parameter(torch::torch_randn(5L, audio_dim)) +}, + forward = function( + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + temb, + temb_audio, + temb_ca_scale_shift, + temb_ca_audio_scale_shift, + temb_ca_gate, + temb_ca_audio_gate, + temb_prompt = NULL, + temb_prompt_audio = NULL, + video_rotary_emb = NULL, + audio_rotary_emb = NULL, + ca_video_rotary_emb = NULL, + ca_audio_rotary_emb = NULL, + encoder_attention_mask = NULL, + audio_encoder_attention_mask = NULL, + self_attention_mask = NULL, + use_a2v_cross_attention = TRUE, + use_v2a_cross_attention = TRUE, + perturbation_mask = NULL, + all_perturbed = FALSE + ) { + batch_size <- hidden_states$shape[1] + + # 1.1 Video self-attention + video_ada <- ltx23_get_mod_params(self$scale_shift_table, temb, batch_size) + shift_msa <- video_ada[[1]]; scale_msa <- video_ada[[2]]; gate_msa <- video_ada[[3]] + shift_mlp <- video_ada[[4]]; scale_mlp <- video_ada[[5]]; gate_mlp <- video_ada[[6]] + + norm_hidden_states <- self$norm1(hidden_states) + norm_hidden_states <- norm_hidden_states * scale_msa$add(1) + shift_msa + + attn_hidden_states <- self$attn1( + norm_hidden_states, + query_rotary_emb = video_rotary_emb, + attention_mask = self_attention_mask, + perturbation_mask = if (self$perturbed_attn) { + perturbation_mask + } else { + NULL + }, + all_perturbed = if (self$perturbed_attn) all_perturbed else FALSE + ) + hidden_states <- hidden_states + attn_hidden_states * gate_msa + + # 1.2 Audio self-attention + audio_ada <- ltx23_get_mod_params(self$audio_scale_shift_table, temb_audio, batch_size) + audio_shift_msa <- audio_ada[[1]]; audio_scale_msa <- audio_ada[[2]] + audio_gate_msa <- audio_ada[[3]]; audio_shift_mlp <- audio_ada[[4]] + audio_scale_mlp <- audio_ada[[5]]; audio_gate_mlp <- audio_ada[[6]] + + norm_audio_hidden_states <- self$audio_norm1(audio_hidden_states) + norm_audio_hidden_states <- norm_audio_hidden_states * audio_scale_msa$add(1) + audio_shift_msa + + attn_audio_hidden_states <- self$audio_attn1( + norm_audio_hidden_states, + query_rotary_emb = audio_rotary_emb, + perturbation_mask = if (self$perturbed_attn) { + perturbation_mask + } else { + NULL + }, + all_perturbed = if (self$perturbed_attn) all_perturbed else FALSE + ) + audio_hidden_states <- audio_hidden_states + attn_audio_hidden_states * audio_gate_msa + + # 2. Text cross-attention modulation params (LTX-2.3) + if (self$cross_attn_adaln) { + prompt_ada <- ltx23_get_mod_params(self$prompt_scale_shift_table, temb_prompt, batch_size) + shift_text_kv <- prompt_ada[[1]]; scale_text_kv <- prompt_ada[[2]] + audio_prompt_ada <- ltx23_get_mod_params( + self$audio_prompt_scale_shift_table, temb_prompt_audio, batch_size + ) + audio_shift_text_kv <- audio_prompt_ada[[1]]; audio_scale_text_kv <- audio_prompt_ada[[2]] + } + + # 2.1 Video-text cross-attention + norm_hidden_states <- self$norm2(hidden_states) + if (self$video_cross_attn_adaln) { + shift_text_q <- video_ada[[7]]; scale_text_q <- video_ada[[8]]; gate_text_q <- video_ada[[9]] + norm_hidden_states <- norm_hidden_states * scale_text_q$add(1) + shift_text_q + } + enc_states <- encoder_hidden_states + if (self$cross_attn_adaln) { + enc_states <- enc_states * scale_text_kv$add(1) + shift_text_kv + } + attn_hidden_states <- self$attn2( + norm_hidden_states, + encoder_hidden_states = enc_states, + attention_mask = encoder_attention_mask + ) + if (self$video_cross_attn_adaln) { + attn_hidden_states <- attn_hidden_states * gate_text_q + } + hidden_states <- hidden_states + attn_hidden_states + + # 2.2 Audio-text cross-attention + norm_audio_hidden_states <- self$audio_norm2(audio_hidden_states) + if (self$audio_cross_attn_adaln) { + audio_shift_text_q <- audio_ada[[7]]; audio_scale_text_q <- audio_ada[[8]] + audio_gate_text_q <- audio_ada[[9]] + norm_audio_hidden_states <- norm_audio_hidden_states * audio_scale_text_q$add(1) + + audio_shift_text_q + } + audio_enc_states <- audio_encoder_hidden_states + if (self$cross_attn_adaln) { + audio_enc_states <- audio_enc_states * audio_scale_text_kv$add(1) + audio_shift_text_kv + } + attn_audio_hidden_states <- self$audio_attn2( + norm_audio_hidden_states, + encoder_hidden_states = audio_enc_states, + attention_mask = audio_encoder_attention_mask + ) + if (self$audio_cross_attn_adaln) { + attn_audio_hidden_states <- attn_audio_hidden_states * audio_gate_text_q + } + audio_hidden_states <- audio_hidden_states + attn_audio_hidden_states + + # 3. Audio-to-video and video-to-audio cross-attention + if (use_a2v_cross_attention || use_v2a_cross_attention) { + norm_hidden_states <- self$audio_to_video_norm(hidden_states) + norm_audio_hidden_states <- self$video_to_audio_norm(audio_hidden_states) + + video_ca_ada <- ltx23_get_mod_params( + self$video_a2v_cross_attn_scale_shift_table$narrow(1L, 1L, 4L), + temb_ca_scale_shift, batch_size + ) + video_ca_gate <- ltx23_get_mod_params( + self$video_a2v_cross_attn_scale_shift_table$narrow(1L, 5L, 1L), + temb_ca_gate, batch_size + ) + a2v_gate <- video_ca_gate[[1]]$squeeze(3L) + + audio_ca_ada <- ltx23_get_mod_params( + self$audio_a2v_cross_attn_scale_shift_table$narrow(1L, 1L, 4L), + temb_ca_audio_scale_shift, batch_size + ) + audio_ca_gate <- ltx23_get_mod_params( + self$audio_a2v_cross_attn_scale_shift_table$narrow(1L, 5L, 1L), + temb_ca_audio_gate, batch_size + ) + v2a_gate <- audio_ca_gate[[1]]$squeeze(3L) + + if (use_a2v_cross_attention) { + mod_norm_hidden <- norm_hidden_states * + video_ca_ada[[1]]$squeeze(3L)$add(1) + video_ca_ada[[2]]$squeeze(3L) + mod_norm_audio <- norm_audio_hidden_states * + audio_ca_ada[[1]]$squeeze(3L)$add(1) + audio_ca_ada[[2]]$squeeze(3L) + + a2v_attn <- self$audio_to_video_attn( + mod_norm_hidden, + encoder_hidden_states = mod_norm_audio, + query_rotary_emb = ca_video_rotary_emb, + key_rotary_emb = ca_audio_rotary_emb + ) + hidden_states <- hidden_states + a2v_gate * a2v_attn + } + + if (use_v2a_cross_attention) { + mod_norm_hidden <- norm_hidden_states * + video_ca_ada[[3]]$squeeze(3L)$add(1) + video_ca_ada[[4]]$squeeze(3L) + mod_norm_audio <- norm_audio_hidden_states * + audio_ca_ada[[3]]$squeeze(3L)$add(1) + audio_ca_ada[[4]]$squeeze(3L) + + v2a_attn <- self$video_to_audio_attn( + mod_norm_audio, + encoder_hidden_states = mod_norm_hidden, + query_rotary_emb = ca_audio_rotary_emb, + key_rotary_emb = ca_video_rotary_emb + ) + audio_hidden_states <- audio_hidden_states + v2a_gate * v2a_attn + } + } + + # 4. Feed-forward + norm_hidden_states <- self$norm3(hidden_states) * scale_mlp$add(1) + shift_mlp + hidden_states <- hidden_states + self$ff(norm_hidden_states) * gate_mlp + + norm_audio_hidden_states <- self$audio_norm3(audio_hidden_states) * + audio_scale_mlp$add(1) + audio_shift_mlp + audio_hidden_states <- audio_hidden_states + self$audio_ff(norm_audio_hidden_states) * + audio_gate_mlp + + list(hidden_states, audio_hidden_states) +} +) diff --git a/R/dit_ltx2_modules.R b/R/dit_ltx2_modules.R deleted file mode 100644 index 793da2f..0000000 --- a/R/dit_ltx2_modules.R +++ /dev/null @@ -1,802 +0,0 @@ -# LTX2 DiT Transformer Modules -# Audio-Video transformer matching HuggingFace diffusers implementation - -# ------------------------------------------------------------------------------ -# Helper modules: Timestep embeddings -# ------------------------------------------------------------------------------ - -#' Get timestep embedding (sinusoidal) -#' @keywords internal -get_timestep_embedding <- function( - timesteps, - embedding_dim, - flip_sin_to_cos = FALSE, - downscale_freq_shift = 1, - scale = 1, - max_period = 10000 -) { - # Ensure timesteps is 1D - if (timesteps$ndim > 1L) { - timesteps <- timesteps$flatten() - } - - # Store original dtype to convert back at end - orig_dtype <- timesteps$dtype - - half_dim <- embedding_dim %/% 2L - - # Compute in float32 for numerical precision - exponent <- - log(max_period) * torch::torch_arange( - start = 0, end = half_dim - 1L, dtype = torch::torch_float32(), device = timesteps$device - ) - exponent <- exponent / (half_dim - downscale_freq_shift) - - emb <- torch::torch_exp(exponent) - # timesteps: [N], emb: [half_dim] -> result: [N, half_dim] - emb <- timesteps$unsqueeze(- 1L)$to(dtype = torch::torch_float32()) * emb$unsqueeze(1L) - emb <- scale * emb - - # Concat sine and cosine -> [N, embedding_dim] - emb <- torch::torch_cat(list(torch::torch_sin(emb), torch::torch_cos(emb)), dim = - 1L) - - # Flip if needed (2D tensor: [N, dim]) - if (flip_sin_to_cos) { - emb <- torch::torch_cat(list( - emb[, (half_dim + 1L) :embedding_dim], - emb[, 1L:half_dim] - ), dim = - 1L) - } - - # Zero pad if odd - if (embedding_dim %% 2L == 1L) { - emb <- torch::nnf_pad(emb, c(0L, 1L, 0L, 0L)) - } - - # Convert back to original dtype (for mixed precision training) - emb <- emb$to(dtype = orig_dtype) - - emb -} - -#' Timesteps module -#' @keywords internal -timesteps_module <- torch::nn_module( - "Timesteps", - initialize = function( - num_channels, - flip_sin_to_cos = TRUE, - downscale_freq_shift = 0, - scale = 1L - ) { - self$num_channels <- num_channels - self$flip_sin_to_cos <- flip_sin_to_cos - self$downscale_freq_shift <- downscale_freq_shift - self$scale <- scale - }, - forward = function(timesteps) { - get_timestep_embedding( - timesteps, - self$num_channels, - flip_sin_to_cos = self$flip_sin_to_cos, - downscale_freq_shift = self$downscale_freq_shift, - scale = self$scale - ) - } -) - -#' Timestep embedding MLP -#' @keywords internal -timestep_embedding_module <- torch::nn_module( - "TimestepEmbedding", - initialize = function( - in_channels, - time_embed_dim, - act_fn = "silu", - out_dim = NULL - ) { - self$linear_1 <- make_linear(in_channels, time_embed_dim) - - if (act_fn == "silu") { - self$act <- torch::nn_silu() - } else if (act_fn == "gelu") { - self$act <- torch::nn_gelu() - } else { - self$act <- torch::nn_silu() - } - - if (!is.null(out_dim)) { - time_embed_dim_out <- out_dim - } else { - time_embed_dim_out <- time_embed_dim - } - self$linear_2 <- make_linear(time_embed_dim, time_embed_dim_out) - }, - forward = function(sample) { - sample <- self$linear_1(sample) - sample <- self$act(sample) - sample <- self$linear_2(sample) - sample - } -) - -#' PixArt Alpha Combined Timestep Size Embeddings -#' @keywords internal -pixart_alpha_combined_timestep_size_embeddings <- torch::nn_module( - "PixArtAlphaCombinedTimestepSizeEmbeddings", - initialize = function( - embedding_dim, - size_emb_dim, - use_additional_conditions = FALSE - ) { - self$outdim <- size_emb_dim - self$time_proj <- timesteps_module(num_channels = 256L, flip_sin_to_cos = TRUE, downscale_freq_shift = 0) - self$timestep_embedder <- timestep_embedding_module(in_channels = 256L, time_embed_dim = embedding_dim) - - self$use_additional_conditions <- use_additional_conditions - if (use_additional_conditions) { - self$additional_condition_proj <- timesteps_module(num_channels = 256L, flip_sin_to_cos = TRUE, downscale_freq_shift = 0) - self$resolution_embedder <- timestep_embedding_module(in_channels = 256L, time_embed_dim = size_emb_dim) - self$aspect_ratio_embedder <- timestep_embedding_module(in_channels = 256L, time_embed_dim = size_emb_dim) - } - }, - forward = function( - timestep, - resolution = NULL, - aspect_ratio = NULL, - batch_size = NULL, - hidden_dtype = NULL - ) { - timesteps_proj <- self$time_proj(timestep) - if (!is.null(hidden_dtype)) { - timesteps_proj <- timesteps_proj$to(dtype = hidden_dtype) - } - timesteps_emb <- self$timestep_embedder(timesteps_proj) - - if (self$use_additional_conditions && !is.null(resolution)) { - resolution_emb <- self$additional_condition_proj(resolution$flatten()) - if (!is.null(hidden_dtype)) { - resolution_emb <- resolution_emb$to(dtype = hidden_dtype) - } - resolution_emb <- self$resolution_embedder(resolution_emb)$reshape(c(batch_size, - 1L)) - - aspect_ratio_emb <- self$additional_condition_proj(aspect_ratio$flatten()) - if (!is.null(hidden_dtype)) { - aspect_ratio_emb <- aspect_ratio_emb$to(dtype = hidden_dtype) - } - aspect_ratio_emb <- self$aspect_ratio_embedder(aspect_ratio_emb)$reshape(c(batch_size, - 1L)) - - conditioning <- timesteps_emb + torch::torch_cat(list(resolution_emb, aspect_ratio_emb), dim = 2L) - } else { - conditioning <- timesteps_emb - } - - conditioning - } -) - -#' PixArt Alpha Text Projection -#' @keywords internal -pixart_alpha_text_projection <- torch::nn_module( - "PixArtAlphaTextProjection", - initialize = function( - in_features, - hidden_size, - out_features = NULL, - act_fn = "gelu_tanh" - ) { - if (is.null(out_features)) out_features <- hidden_size - - self$linear_1 <- make_linear(in_features, hidden_size) - - if (act_fn == "gelu_tanh") { - self$act_1 <- torch::nn_gelu(approximate = "tanh") - } else if (act_fn == "silu") { - self$act_1 <- torch::nn_silu() - } else { - self$act_1 <- torch::nn_gelu(approximate = "tanh") - } - - self$linear_2 <- make_linear(hidden_size, out_features) - }, - forward = function(caption) { - hidden_states <- self$linear_1(caption) - hidden_states <- self$act_1(hidden_states) - hidden_states <- self$linear_2(hidden_states) - hidden_states - } -) - -# ------------------------------------------------------------------------------ -# FeedForward module -# ------------------------------------------------------------------------------ - -#' GELU activation with optional approximation -#' @keywords internal -gelu_activation <- torch::nn_module( - "GELU", - initialize = function( - dim_in, - dim_out, - approximate = "none", - bias = TRUE - ) { - self$proj <- make_linear(dim_in, dim_out, bias = bias) - self$approximate <- approximate - }, - forward = function(hidden_states) { - hidden_states <- self$proj(hidden_states) - hidden_states <- torch::nnf_gelu(hidden_states, approximate = self$approximate) - hidden_states - } -) - -#' FeedForward module -#' @keywords internal -feed_forward <- torch::nn_module( - "FeedForward", - initialize = function( - dim, - dim_out = NULL, - mult = 4L, - dropout = 0.0, - activation_fn = "gelu-approximate", - inner_dim = NULL, - bias = TRUE - ) { - if (is.null(inner_dim)) inner_dim <- as.integer(dim * mult) - if (is.null(dim_out)) dim_out <- dim - - # Activation layer (projects in) - if (activation_fn == "gelu") { - self$act_fn <- gelu_activation(dim, inner_dim, approximate = "none", bias = bias) - } else if (activation_fn == "gelu-approximate") { - self$act_fn <- gelu_activation(dim, inner_dim, approximate = "tanh", bias = bias) - } else { - self$act_fn <- gelu_activation(dim, inner_dim, approximate = "tanh", bias = bias) - } - - self$dropout <- torch::nn_dropout(p = dropout) - self$proj_out <- make_linear(inner_dim, dim_out, bias = bias) - }, - forward = function(hidden_states) { - hidden_states <- self$act_fn(hidden_states) - hidden_states <- self$dropout(hidden_states) - hidden_states <- self$proj_out(hidden_states) - hidden_states - } -) - -# ------------------------------------------------------------------------------ -# RMSNorm -# ------------------------------------------------------------------------------ - -#' RMS Normalization -#' @keywords internal -rms_norm <- torch::nn_module( - "RMSNorm", - initialize = function( - dim, - eps = 1e-6, - elementwise_affine = TRUE - ) { - self$eps <- eps - self$elementwise_affine <- elementwise_affine - if (elementwise_affine) { - self$weight <- torch::nn_parameter(torch::torch_ones(dim)) - } - }, - forward = function(hidden_states) { - input_dtype <- hidden_states$dtype - hidden_states <- hidden_states$to(dtype = torch::torch_float32()) - variance <- hidden_states$pow(2)$mean(dim = - 1L, keepdim = TRUE) - hidden_states <- hidden_states * torch::torch_rsqrt(variance + self$eps) - if (self$elementwise_affine) { - hidden_states <- hidden_states * self$weight - } - hidden_states$to(dtype = input_dtype) - } -) - -# ------------------------------------------------------------------------------ -# LTX2 Adaptive Layer Norm Single -# ------------------------------------------------------------------------------ - -#' LTX2 AdaLayerNorm Single -#' @keywords internal -ltx2_ada_layer_norm_single <- torch::nn_module( - "LTX2AdaLayerNormSingle", - initialize = function( - embedding_dim, - num_mod_params = 6L, - use_additional_conditions = FALSE - ) { - self$num_mod_params <- num_mod_params - - self$emb <- pixart_alpha_combined_timestep_size_embeddings( - embedding_dim = embedding_dim, - size_emb_dim = embedding_dim %/% 3L, - use_additional_conditions = use_additional_conditions - ) - - self$silu <- torch::nn_silu() - self$linear <- make_linear(embedding_dim, num_mod_params * embedding_dim) - }, - forward = function( - timestep, - added_cond_kwargs = NULL, - batch_size = NULL, - hidden_dtype = NULL - ) { - if (is.null(added_cond_kwargs)) { - added_cond_kwargs <- list(resolution = NULL, aspect_ratio = NULL) - } - - embedded_timestep <- self$emb( - timestep, - resolution = added_cond_kwargs$resolution, - aspect_ratio = added_cond_kwargs$aspect_ratio, - batch_size = batch_size, - hidden_dtype = hidden_dtype - ) - - mod_params <- self$linear(self$silu(embedded_timestep)) - list(mod_params, embedded_timestep) - } -) - -# ------------------------------------------------------------------------------ -# LTX2 Attention -# ------------------------------------------------------------------------------ - -#' Apply interleaved rotary embedding -#' @keywords internal -apply_interleaved_rotary_emb_list <- function( - x, - freqs -) { - cos_freqs <- freqs[[1]] - sin_freqs <- freqs[[2]] - - # x: [B, S, C] - # Split into real and imaginary parts - x_shape <- x$shape - x_reshaped <- x$unflatten(3, c(- 1L, 2L)) # [B, S, C//2, 2] - x_real <- x_reshaped[,,, 1] - x_imag <- x_reshaped[,,, 2] - - # Rotate: [-x_imag, x_real] - x_rotated <- torch::torch_stack(list(- x_imag, x_real), dim = - 1L)$flatten(start_dim = 3L) - - # Apply rotation - out <- (x$to(dtype = torch::torch_float32()) * cos_freqs + - x_rotated$to(dtype = torch::torch_float32()) * sin_freqs)$to(dtype = x$dtype) - - out -} - -#' LTX2 Attention module -#' @keywords internal -ltx2_attention <- torch::nn_module( - "LTX2Attention", - initialize = function( - query_dim, - heads = 8L, - kv_heads = 8L, - dim_head = 64L, - dropout = 0.0, - bias = TRUE, - cross_attention_dim = NULL, - out_bias = TRUE, - qk_norm = "rms_norm_across_heads", - norm_eps = 1e-6, - norm_elementwise_affine = TRUE, - rope_type = "interleaved" - ) { - - self$head_dim <- dim_head - self$inner_dim <- dim_head * heads - if (is.null(kv_heads)) { - self$inner_kv_dim <- self$inner_dim - } else { - self$inner_kv_dim <- dim_head * kv_heads - } - self$query_dim <- query_dim - if (!is.null(cross_attention_dim)) { - self$cross_attention_dim <- cross_attention_dim - } else { - self$cross_attention_dim <- query_dim - } - self$use_bias <- bias - self$dropout_p <- dropout - self$out_dim <- query_dim - self$heads <- heads - self$rope_type <- rope_type - - # QK normalization - self$norm_q <- rms_norm(dim_head * heads, eps = norm_eps, elementwise_affine = norm_elementwise_affine) - if (is.null(kv_heads)) { - self$norm_k <- rms_norm(dim_head * heads, eps = norm_eps, elementwise_affine = norm_elementwise_affine) - } else { - self$norm_k <- rms_norm(dim_head * kv_heads, eps = norm_eps, elementwise_affine = norm_elementwise_affine) - } - - # Projections - self$to_q <- make_linear(query_dim, self$inner_dim, bias = bias) - self$to_k <- make_linear(self$cross_attention_dim, self$inner_kv_dim, bias = bias) - self$to_v <- make_linear(self$cross_attention_dim, self$inner_kv_dim, bias = bias) - - # Output projection - self$to_out <- torch::nn_sequential( - make_linear(self$inner_dim, self$out_dim, bias = out_bias), - torch::nn_dropout(p = dropout) - ) - }, - forward = function( - hidden_states, - encoder_hidden_states = NULL, - attention_mask = NULL, - query_rotary_emb = NULL, - key_rotary_emb = NULL - ) { - batch_size <- hidden_states$shape[1] - - if (is.null(encoder_hidden_states)) { - encoder_hidden_states <- hidden_states - } - - # Project Q, K, V - query <- self$to_q(hidden_states) - key <- self$to_k(encoder_hidden_states) - value <- self$to_v(encoder_hidden_states) - - # Normalize Q, K - query <- self$norm_q(query) - key <- self$norm_k(key) - - # Apply RoPE - if (!is.null(query_rotary_emb)) { - if (self$rope_type == "interleaved") { - query <- apply_interleaved_rotary_emb_list(query, query_rotary_emb) - if (!is.null(key_rotary_emb)) { - key_rope <- key_rotary_emb - } else { - key_rope <- query_rotary_emb - } - key <- apply_interleaved_rotary_emb_list(key, key_rope) - } - } - - # Reshape for multi-head attention [B, S, H, D] - query <- query$unflatten(3, c(self$heads, - 1L)) - key <- key$unflatten(3, c(self$heads, - 1L)) - value <- value$unflatten(3, c(self$heads, - 1L)) - - # Transpose to [B, H, S, D] - query <- query$transpose(2L, 3L) - key <- key$transpose(2L, 3L) - value <- value$transpose(2L, 3L) - - # Scaled dot-product attention (manual implementation) - scale <- 1.0 / sqrt(self$head_dim) - attn_weights <- torch::torch_matmul(query, key$transpose(- 2L, - 1L)) * scale - - if (!is.null(attention_mask)) { - # Expand attention mask to [B, 1, 1, S] for broadcasting with [B, H, S, S] - if (attention_mask$ndim == 2L) { - attention_mask <- attention_mask$unsqueeze(2L)$unsqueeze(2L) # [B, S] -> [B, 1, 1, S] - } else if (attention_mask$ndim == 3L) { - attention_mask <- attention_mask$unsqueeze(2L) # [B, 1, S] -> [B, 1, 1, S] - } - attn_weights <- attn_weights + attention_mask - } - - attn_weights <- torch::nnf_softmax(attn_weights, dim = - 1L) - - if (self$training && self$dropout_p > 0) { - attn_weights <- torch::nnf_dropout(attn_weights, p = self$dropout_p) - } - - hidden_states <- torch::torch_matmul(attn_weights, value) - - # Reshape back [B, H, S, D] -> [B, S, H*D] - hidden_states <- hidden_states$transpose(2L, 3L)$flatten(start_dim = 3L) - hidden_states <- hidden_states$to(dtype = query$dtype) - - # Output projection - hidden_states <- self$to_out(hidden_states) - - hidden_states - } -) - -# ------------------------------------------------------------------------------ -# LTX2 Video Transformer Block -# ------------------------------------------------------------------------------ - -#' LTX2 Video Transformer Block (Audio-Video) -#' @keywords internal -ltx2_video_transformer_block <- torch::nn_module( - "LTX2VideoTransformerBlock", - initialize = function( - dim, - num_attention_heads, - attention_head_dim, - cross_attention_dim, - audio_dim, - audio_num_attention_heads, - audio_attention_head_dim, - audio_cross_attention_dim, - qk_norm = "rms_norm_across_heads", - activation_fn = "gelu-approximate", - attention_bias = TRUE, - attention_out_bias = TRUE, - eps = 1e-6, - elementwise_affine = FALSE, - rope_type = "interleaved" - ) { - - # 1. Video Self-Attention - self$norm1 <- rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) - self$attn1 <- ltx2_attention( - query_dim = dim, - heads = num_attention_heads, - kv_heads = num_attention_heads, - dim_head = attention_head_dim, - bias = attention_bias, - cross_attention_dim = NULL, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # Audio Self-Attention - self$audio_norm1 <- rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) - self$audio_attn1 <- ltx2_attention( - query_dim = audio_dim, - heads = audio_num_attention_heads, - kv_heads = audio_num_attention_heads, - dim_head = audio_attention_head_dim, - bias = attention_bias, - cross_attention_dim = NULL, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # 2. Video Cross-Attention (with text) - self$norm2 <- rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) - self$attn2 <- ltx2_attention( - query_dim = dim, - cross_attention_dim = cross_attention_dim, - heads = num_attention_heads, - kv_heads = num_attention_heads, - dim_head = attention_head_dim, - bias = attention_bias, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # Audio Cross-Attention (with text) - self$audio_norm2 <- rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) - self$audio_attn2 <- ltx2_attention( - query_dim = audio_dim, - cross_attention_dim = audio_cross_attention_dim, - heads = audio_num_attention_heads, - kv_heads = audio_num_attention_heads, - dim_head = audio_attention_head_dim, - bias = attention_bias, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # 3. Audio-to-Video Cross-Attention (Q: Video, K/V: Audio) - self$audio_to_video_norm <- rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) - self$audio_to_video_attn <- ltx2_attention( - query_dim = dim, - cross_attention_dim = audio_dim, - heads = audio_num_attention_heads, - kv_heads = audio_num_attention_heads, - dim_head = audio_attention_head_dim, - bias = attention_bias, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # Video-to-Audio Cross-Attention (Q: Audio, K/V: Video) - self$video_to_audio_norm <- rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) - self$video_to_audio_attn <- ltx2_attention( - query_dim = audio_dim, - cross_attention_dim = dim, - heads = audio_num_attention_heads, - kv_heads = audio_num_attention_heads, - dim_head = audio_attention_head_dim, - bias = attention_bias, - out_bias = attention_out_bias, - qk_norm = qk_norm, - rope_type = rope_type - ) - - # 4. Feedforward layers - self$norm3 <- rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) - self$ff <- feed_forward(dim, activation_fn = activation_fn) - - self$audio_norm3 <- rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) - self$audio_ff <- feed_forward(audio_dim, activation_fn = activation_fn) - - # 5. Per-layer modulation parameters - self$scale_shift_table <- torch::nn_parameter( - torch::torch_randn(6L, dim) / sqrt(dim) - ) - self$audio_scale_shift_table <- torch::nn_parameter( - torch::torch_randn(6L, audio_dim) / sqrt(audio_dim) - ) - - # Cross-attention modulation parameters - self$video_a2v_cross_attn_scale_shift_table <- torch::nn_parameter(torch::torch_randn(5L, dim)) - self$audio_a2v_cross_attn_scale_shift_table <- torch::nn_parameter(torch::torch_randn(5L, audio_dim)) - }, - forward = function( - hidden_states, - audio_hidden_states, - encoder_hidden_states, - audio_encoder_hidden_states, - temb, - temb_audio, - temb_ca_scale_shift, - temb_ca_audio_scale_shift, - temb_ca_gate, - temb_ca_audio_gate, - video_rotary_emb = NULL, - audio_rotary_emb = NULL, - ca_video_rotary_emb = NULL, - ca_audio_rotary_emb = NULL, - encoder_attention_mask = NULL, - audio_encoder_attention_mask = NULL, - a2v_cross_attention_mask = NULL, - v2a_cross_attention_mask = NULL - ) { - - batch_size <- hidden_states$shape[1] - - # 1. Video and Audio Self-Attention - norm_hidden_states <- self$norm1(hidden_states) - - # Ada values for video - num_ada_params <- self$scale_shift_table$shape[1] - ada_values <- self$scale_shift_table$unsqueeze(1)$unsqueeze(1)$to(device = temb$device, dtype = temb$dtype) + - temb$reshape(c(batch_size, temb$shape[2], num_ada_params, - 1L)) - - shift_msa <- ada_values[,, 1,] - scale_msa <- ada_values[,, 2,] - gate_msa <- ada_values[,, 3,] - shift_mlp <- ada_values[,, 4,] - scale_mlp <- ada_values[,, 5,] - gate_mlp <- ada_values[,, 6,] - - norm_hidden_states <- norm_hidden_states * scale_msa$add(1) + shift_msa - - attn_hidden_states <- self$attn1( - hidden_states = norm_hidden_states, - encoder_hidden_states = NULL, - query_rotary_emb = video_rotary_emb - ) - hidden_states <- hidden_states + attn_hidden_states * gate_msa - - # Audio self-attention - norm_audio_hidden_states <- self$audio_norm1(audio_hidden_states) - - num_audio_ada_params <- self$audio_scale_shift_table$shape[1] - audio_ada_values <- self$audio_scale_shift_table$unsqueeze(1)$unsqueeze(1)$to(device = temb_audio$device, dtype = temb_audio$dtype) + - temb_audio$reshape(c(batch_size, temb_audio$shape[2], num_audio_ada_params, - 1L)) - - audio_shift_msa <- audio_ada_values[,, 1,] - audio_scale_msa <- audio_ada_values[,, 2,] - audio_gate_msa <- audio_ada_values[,, 3,] - audio_shift_mlp <- audio_ada_values[,, 4,] - audio_scale_mlp <- audio_ada_values[,, 5,] - audio_gate_mlp <- audio_ada_values[,, 6,] - - norm_audio_hidden_states <- norm_audio_hidden_states * audio_scale_msa$add(1) + audio_shift_msa - - attn_audio_hidden_states <- self$audio_attn1( - hidden_states = norm_audio_hidden_states, - encoder_hidden_states = NULL, - query_rotary_emb = audio_rotary_emb - ) - audio_hidden_states <- audio_hidden_states + attn_audio_hidden_states * audio_gate_msa - - # 2. Video and Audio Cross-Attention with text - norm_hidden_states <- self$norm2(hidden_states) - attn_hidden_states <- self$attn2( - norm_hidden_states, - encoder_hidden_states = encoder_hidden_states, - query_rotary_emb = NULL, - attention_mask = encoder_attention_mask - ) - hidden_states <- hidden_states + attn_hidden_states - - norm_audio_hidden_states <- self$audio_norm2(audio_hidden_states) - attn_audio_hidden_states <- self$audio_attn2( - norm_audio_hidden_states, - encoder_hidden_states = audio_encoder_hidden_states, - query_rotary_emb = NULL, - attention_mask = audio_encoder_attention_mask - ) - audio_hidden_states <- audio_hidden_states + attn_audio_hidden_states - - # 3. Audio-to-Video and Video-to-Audio Cross-Attention - norm_hidden_states <- self$audio_to_video_norm(hidden_states) - norm_audio_hidden_states <- self$video_to_audio_norm(audio_hidden_states) - - # Video cross-attention modulation - video_per_layer_ca_scale_shift <- self$video_a2v_cross_attn_scale_shift_table[1:4,] - video_per_layer_ca_gate <- self$video_a2v_cross_attn_scale_shift_table[5:5,] - - video_ca_scale_shift_table <- video_per_layer_ca_scale_shift$unsqueeze(1)$unsqueeze(1)$to(dtype = temb_ca_scale_shift$dtype) + - temb_ca_scale_shift$reshape(c(batch_size, temb_ca_scale_shift$shape[2], 4L, - 1L)) - video_ca_gate <- video_per_layer_ca_gate$unsqueeze(1)$unsqueeze(1)$to(dtype = temb_ca_gate$dtype) + - temb_ca_gate$reshape(c(batch_size, temb_ca_gate$shape[2], 1L, - 1L)) - - video_a2v_ca_scale <- video_ca_scale_shift_table[,, 1,] - video_a2v_ca_shift <- video_ca_scale_shift_table[,, 2,] - video_v2a_ca_scale <- video_ca_scale_shift_table[,, 3,] - video_v2a_ca_shift <- video_ca_scale_shift_table[,, 4,] - a2v_gate <- video_ca_gate[,, 1,] - - # Audio cross-attention modulation - audio_per_layer_ca_scale_shift <- self$audio_a2v_cross_attn_scale_shift_table[1:4,] - audio_per_layer_ca_gate <- self$audio_a2v_cross_attn_scale_shift_table[5:5,] - - audio_ca_scale_shift_table <- audio_per_layer_ca_scale_shift$unsqueeze(1)$unsqueeze(1)$to(dtype = temb_ca_audio_scale_shift$dtype) + - temb_ca_audio_scale_shift$reshape(c(batch_size, temb_ca_audio_scale_shift$shape[2], 4L, - 1L)) - audio_ca_gate <- audio_per_layer_ca_gate$unsqueeze(1)$unsqueeze(1)$to(dtype = temb_ca_audio_gate$dtype) + - temb_ca_audio_gate$reshape(c(batch_size, temb_ca_audio_gate$shape[2], 1L, - 1L)) - - audio_a2v_ca_scale <- audio_ca_scale_shift_table[,, 1,] - audio_a2v_ca_shift <- audio_ca_scale_shift_table[,, 2,] - audio_v2a_ca_scale <- audio_ca_scale_shift_table[,, 3,] - audio_v2a_ca_shift <- audio_ca_scale_shift_table[,, 4,] - v2a_gate <- audio_ca_gate[,, 1,] - - # Audio-to-Video Cross Attention - mod_norm_hidden_states <- norm_hidden_states * video_a2v_ca_scale$add(1) + video_a2v_ca_shift - mod_norm_audio_hidden_states <- norm_audio_hidden_states * audio_a2v_ca_scale$add(1) + audio_a2v_ca_shift - - a2v_attn_hidden_states <- self$audio_to_video_attn( - mod_norm_hidden_states, - encoder_hidden_states = mod_norm_audio_hidden_states, - query_rotary_emb = ca_video_rotary_emb, - key_rotary_emb = ca_audio_rotary_emb, - attention_mask = a2v_cross_attention_mask - ) - hidden_states <- hidden_states + a2v_gate * a2v_attn_hidden_states - - # Video-to-Audio Cross Attention - mod_norm_hidden_states <- norm_hidden_states * video_v2a_ca_scale$add(1) + video_v2a_ca_shift - mod_norm_audio_hidden_states <- norm_audio_hidden_states * audio_v2a_ca_scale$add(1) + audio_v2a_ca_shift - - v2a_attn_hidden_states <- self$video_to_audio_attn( - mod_norm_audio_hidden_states, - encoder_hidden_states = mod_norm_hidden_states, - query_rotary_emb = ca_audio_rotary_emb, - key_rotary_emb = ca_video_rotary_emb, - attention_mask = v2a_cross_attention_mask - ) - audio_hidden_states <- audio_hidden_states + v2a_gate * v2a_attn_hidden_states - - # 4. Feedforward - norm_hidden_states <- self$norm3(hidden_states) * scale_mlp$add(1) + shift_mlp - ff_output <- self$ff(norm_hidden_states) - hidden_states <- hidden_states + ff_output * gate_mlp - - norm_audio_hidden_states <- self$audio_norm3(audio_hidden_states) * audio_scale_mlp$add(1) + audio_shift_mlp - audio_ff_output <- self$audio_ff(norm_audio_hidden_states) - audio_hidden_states <- audio_hidden_states + audio_ff_output * audio_gate_mlp - - list(hidden_states, audio_hidden_states) - } -) - diff --git a/R/download_component.R b/R/download_component.R index b221a7b..142ff14 100644 --- a/R/download_component.R +++ b/R/download_component.R @@ -16,8 +16,9 @@ #' \dontrun{ #' path <- download_component("sd21", "text_encoder", "cpu") #' } -download_component <- function (model_name = "sd21", component, device = "cpu", - overwrite = FALSE, show_progress = TRUE) { +download_component <- function(model_name = "sd21", component, + device = "cpu", overwrite = FALSE, + show_progress = TRUE) { filename <- paste0(component, "-", device, ".pt") if (overwrite) { @@ -26,10 +27,9 @@ download_component <- function (model_name = "sd21", component, device = "cpu", stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") } repo_id <- paste0("cornball-ai/", model_name, "-R") - return(hfhub::hub_download(repo_id, filename, - repo_type = "dataset", force_download = TRUE)) + return(hfhub::hub_download(repo_id, filename, repo_type = "dataset", + force_download = TRUE)) } hf_download_pt(model_name, filename, download = TRUE) } - diff --git a/R/download_ltx23.R b/R/download_ltx23.R new file mode 100644 index 0000000..879c60e --- /dev/null +++ b/R/download_ltx23.R @@ -0,0 +1,149 @@ +#' Download and Prepare LTX-2.3 Model Weights +#' +#' Downloads the LTX-2.3 distilled checkpoint (46 GB, LTX-2 Community +#' License) and the Gemma3 text encoder from HuggingFace with an explicit +#' consent prompt, then quantizes the transformer to the local fp8 +#' artifact (~26 GB) used by the GPU-poor pipeline. +#' +#' @name download_ltx23 +NULL + +.ltx23_checkpoint_repo <- "Lightricks/LTX-2.3" +.ltx23_checkpoint_file <- "ltx-2.3-22b-distilled-1.1.safetensors" +.ltx23_text_encoder_repo <- "Lightricks/LTX-2" + +.ltx23_disk_free_gb <- function(path) { + out <- tryCatch( + system2("df", c("-Pk", shQuote(path)), stdout = TRUE, stderr = FALSE), + error = function(e) NULL + ) + if (is.null(out) || length(out) < 2L) { + return(NA_real_) + } + fields <- strsplit(trimws(out[[2]]), "\\s+")[[1]] + as.numeric(fields[4]) / 1024 ^ 2 +} + +.ltx23_consent <- function(what) { + if (isTRUE(getOption("diffuseR.consent"))) { + return(TRUE) + } + if (!interactive()) { + stop( + "Cannot download models in non-interactive mode without consent. ", + "Set options(diffuseR.consent = TRUE) to allow downloads.", + call. = FALSE + ) + } + isTRUE(utils::askYesNo(paste0("Download ", what, "?"))) +} + +#' Download the LTX-2.3 checkpoint and build the fp8 artifact +#' +#' Skips work that is already done: a valid fp8 manifest short-circuits +#' everything; a cached 46 GB source skips the download. The source file +#' may be deleted after quantization (it is never removed automatically). +#' +#' @param quantize Logical. Build the fp8 artifact after downloading. +#' @param output_dir Directory for the fp8 artifact. +#' @param text_encoder Logical. Also fetch the Gemma3 text encoder and +#' tokenizer (~25 GB, shared with LTX-2.0; from the Lightricks/LTX-2 +#' repo). +#' @param verbose Logical. +#' +#' @return Invisibly, a list with \code{checkpoint} (source path or NULL), +#' \code{fp8_dir}, and \code{text_encoder_dir}. +#' +#' @export +download_ltx2 <- function(quantize = TRUE, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), + text_encoder = TRUE, verbose = TRUE) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to download model weights.") + } + result <- list(checkpoint = NULL, fp8_dir = output_dir, + text_encoder_dir = NULL) + + manifest_path <- file.path(output_dir, "manifest.json") + have_fp8 <- file.exists(manifest_path) && { + m <- jsonlite::fromJSON(manifest_path) + all(file.exists(file.path(output_dir, m$shards))) + } + + if (!have_fp8 || !quantize) { + cached <- tryCatch( + hfhub::hub_download(.ltx23_checkpoint_repo, .ltx23_checkpoint_file, + local_files_only = TRUE), + error = function(e) NULL + ) + if (is.null(cached) && !have_fp8) { + free <- .ltx23_disk_free_gb(path.expand("~")) + if (!is.na(free) && free < 75) { + warning(sprintf( + "Only %.0f GB free; the download + fp8 artifact need ~75 GB.", free + )) + } + ok <- .ltx23_consent(paste0( + "the LTX-2.3 distilled checkpoint (46 GB) plus a ~26 GB local fp8 ", + "artifact from HuggingFace (weights under the LTX-2 Community License)" + )) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading ", .ltx23_checkpoint_file, " (46 GB)...") + } + cached <- hfhub::hub_download(.ltx23_checkpoint_repo, .ltx23_checkpoint_file) + } + result$checkpoint <- cached + + if (quantize && !have_fp8) { + if (verbose) { + message("Quantizing transformer linears to fp8 (one-time)...") + } + ltx23_quantize_fp8(cached, output_dir, verbose = verbose) + if (verbose) { + message( + "FP8 artifact ready: ", output_dir, "\n", + "The 46 GB source in the HuggingFace cache may be deleted if ", + "you do not need bf16 weights." + ) + } + } + } else if (verbose) { + message("FP8 artifact already present: ", output_dir) + } + + if (text_encoder) { + te_files <- c( + "text_encoder/config.json", + sprintf("text_encoder/model-%05d-of-00011.safetensors", 1:11), + "text_encoder/model.safetensors.index.json", + "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json", + "tokenizer/special_tokens_map.json" + ) + have_te <- !is.null(tryCatch( + hfhub::hub_download(.ltx23_text_encoder_repo, te_files[[2]], + local_files_only = TRUE), + error = function(e) NULL + )) + if (!have_te) { + ok <- .ltx23_consent( + "the Gemma3 text encoder and tokenizer (~25 GB, Lightricks/LTX-2)" + ) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + } + if (verbose && !have_te) { + message("Downloading Gemma3 text encoder...") + } + paths <- vapply(te_files, function(f) { + tryCatch(hfhub::hub_download(.ltx23_text_encoder_repo, f), + error = function(e) NA_character_) + }, character(1)) + result$text_encoder_dir <- dirname(paths[[2]]) + } + + invisible(result) +} diff --git a/R/download_model.R b/R/download_model.R index 1c61a95..e84d288 100644 --- a/R/download_model.R +++ b/R/download_model.R @@ -10,7 +10,7 @@ #' #' @return The local file path (character string). #' @keywords internal -hf_download_pt <- function (model_name, filename, download = TRUE) { +hf_download_pt <- function(model_name, filename, download = TRUE) { if (!requireNamespace("hfhub", quietly = TRUE)) { stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") } @@ -19,21 +19,25 @@ hf_download_pt <- function (model_name, filename, download = TRUE) { # 1. Try hfhub cache first (no network) path <- tryCatch( - hfhub::hub_download(repo_id, filename, - repo_type = "dataset", local_files_only = TRUE), - error = function (e) NULL + hfhub::hub_download(repo_id, filename, repo_type = "dataset", + local_files_only = TRUE), + error = function(e) NULL ) - if (!is.null(path) && file.exists(path)) return(path) + if (!is.null(path) && file.exists(path)) { + return(path) + } # 2. Check legacy R_user_dir location legacy_path <- file.path(tools::R_user_dir("diffuseR", "data"), - model_name, filename) - if (file.exists(legacy_path)) return(legacy_path) + model_name, filename) + if (file.exists(legacy_path)) { + return(legacy_path) + } # 3. Download via hfhub if (!download) { stop("Component '", filename, "' for model '", model_name, - "' not found locally. Set download=TRUE to download.") + "' not found locally. Set download=TRUE to download.") } hfhub::hub_download(repo_id, filename, repo_type = "dataset") @@ -62,30 +66,22 @@ hf_download_pt <- function (model_name, filename, download = TRUE) { #' paths <- download_model("sd21") #' } #' -download_model <- function( - model_name = "sd21", - devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), - unet_dtype_str = NULL, - overwrite = FALSE, - show_progress = TRUE, - download_models = FALSE -) { +download_model <- function(model_name = "sd21", + devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), + unet_dtype_str = NULL, overwrite = FALSE, + show_progress = TRUE, download_models = FALSE) { # Normalize 'devices' if (is.character(devices) && length(devices) == 1) { if (model_name == "sdxl") { - devices <- list( - unet = devices, - decoder = devices, - text_encoder = devices, - text_encoder2 = devices, - encoder = devices - ) + devices <- list(unet = devices, decoder = devices, + text_encoder = devices, text_encoder2 = devices, + encoder = devices) } else { devices <- list( - unet = devices, - decoder = devices, - text_encoder = devices, - encoder = devices + unet = devices, + decoder = devices, + text_encoder = devices, + encoder = devices ) } } else if (is.list(devices)) { @@ -128,7 +124,11 @@ download_model <- function( stop("Device not specified for component: ", name) } if (name == "unet" && device != "cpu") { - dtype <- if (is.null(unet_dtype_str)) "float16" else unet_dtype_str + if (is.null(unet_dtype_str)) { + dtype <- "float16" + } else { + dtype <- unet_dtype_str + } if (!dtype %in% c("float16", "float32")) { stop("Invalid unet_dtype_str: must be 'float16' or 'float32'") } @@ -140,17 +140,17 @@ download_model <- function( # Check which files are already available (hfhub cache or legacy) already_available <- vapply(filenames, function(f) { - # Check hfhub cache - path <- tryCatch( - hfhub::hub_download(paste0("cornball-ai/", model_name, "-R"), f, - repo_type = "dataset", local_files_only = TRUE), - error = function(e) NULL - ) - if (!is.null(path) && file.exists(path)) return(TRUE) - # Check legacy location - legacy <- file.path(tools::R_user_dir("diffuseR", "data"), model_name, f) - file.exists(legacy) - }, logical(1)) + # Check hfhub cache + path <- tryCatch( + hfhub::hub_download(paste0("cornball-ai/", model_name, "-R"), f, + repo_type = "dataset", local_files_only = TRUE), + error = function(e) NULL + ) + if (!is.null(path) && file.exists(path)) return(TRUE) + # Check legacy location + legacy <- file.path(tools::R_user_dir("diffuseR", "data"), model_name, f) + file.exists(legacy) + }, logical(1)) needs_download <- !already_available | overwrite @@ -158,8 +158,8 @@ download_model <- function( if (download_models && any(needs_download)) { if (interactive()) { ans <- utils::askYesNo( - paste0("Download '", model_name, "' model files from HuggingFace?"), - default = TRUE + paste0("Download '", model_name, "' model files from HuggingFace?"), + default = TRUE ) if (!isTRUE(ans)) { stop("Download cancelled.", call. = FALSE) @@ -174,16 +174,16 @@ download_model <- function( if (needs_download[i] && !download_models) { # Not available and not allowed to download — will error model_paths[i] <- tryCatch( - hf_download_pt(model_name, filenames[i], download = FALSE), - error = function(e) NA_character_ + hf_download_pt(model_name, filenames[i], download = FALSE), + error = function(e) NA_character_ ) } else { model_paths[i] <- tryCatch( - hf_download_pt(model_name, filenames[i], download = download_models), - error = function(e) { - warning("Download failed for ", filenames[i], ": ", e$message) - NA_character_ - } + hf_download_pt(model_name, filenames[i], download = download_models), + error = function(e) { + warning("Download failed for ", filenames[i], ": ", e$message) + NA_character_ + } ) } } @@ -195,4 +195,3 @@ download_model <- function( model_paths } - diff --git a/R/filename_from_prompt.R b/R/filename_from_prompt.R index 509ea06..c5cd6e6 100644 --- a/R/filename_from_prompt.R +++ b/R/filename_from_prompt.R @@ -1,7 +1,7 @@ -#` filename_from_prompt` function +#' filename_from_prompt #' @title Generate a filename from a prompt #' @description This function generates a filename from a prompt by removing all non-alphanumeric characters and replacing them with underscores. The filename is limited to 50 characters. If `datetime` is set to TRUE, the current date and time are prepended to the filename. -#`` +#' #' @param prompt A character string representing the prompt. #' @param datetime Logical indicating whether to prepend the current date and time to the filename. Default is TRUE. #' @return A character string representing the generated filename. @@ -10,22 +10,18 @@ #' filename_from_prompt("A beautiful sunset over the mountains") #' filename_from_prompt("A beautiful sunset over the mountains", datetime = FALSE) #' @export -filename_from_prompt <- function( - prompt, - datetime = TRUE -) { - # Remove all non-alphanumeric characters from the prompt - prompt_strs <- gsub("[^a-zA-Z0-9]", "_", prompt) - # limit prompt length to 50 characters - prompt_strs <- substr(prompt_strs, 1, 50) - if (datetime == FALSE) { - # Create a filename using the modified prompt - file_name <- paste0("prompt_", prompt_strs, ".png") - } else { - # date time stamp - datetime <- format(Sys.time(), "%Y%m%d_%H%M%S") - file_name <- paste0(datetime, "_", prompt_strs, ".png") - } - return(file_name) +filename_from_prompt <- function(prompt, datetime = TRUE) { + # Remove all non-alphanumeric characters from the prompt + prompt_strs <- gsub("[^a-zA-Z0-9]", "_", prompt) + # limit prompt length to 50 characters + prompt_strs <- substr(prompt_strs, 1, 50) + if (datetime == FALSE) { + # Create a filename using the modified prompt + file_name <- paste0("prompt_", prompt_strs, ".png") + } else { + # date time stamp + datetime <- format(Sys.time(), "%Y%m%d_%H%M%S") + file_name <- paste0(datetime, "_", prompt_strs, ".png") + } + return(file_name) } - diff --git a/R/flowmatch_scheduler.R b/R/flowmatch_scheduler.R index e91aac9..a62f807 100644 --- a/R/flowmatch_scheduler.R +++ b/R/flowmatch_scheduler.R @@ -55,53 +55,51 @@ #' scheduler <- flowmatch_set_timesteps(scheduler, num_inference_steps = 8) #' } #' @export -flowmatch_scheduler_create <- function( - num_train_timesteps = 1000L, - shift = 1.0, - use_dynamic_shifting = FALSE, - base_shift = 0.5, - max_shift = 1.15, - base_seq_len = 256L, - max_seq_len = 4096L, - invert_sigmas = FALSE, - shift_terminal = NULL, - time_shift_type = c("exponential", "linear") -) { - time_shift_type <- match.arg(time_shift_type) - - # Initial timesteps (reversed, from max to min) - timesteps <- seq(from = 1, to = num_train_timesteps, length.out = num_train_timesteps) - timesteps <- rev(timesteps) - - # Initial sigmas (normalized timesteps) - sigmas <- timesteps / num_train_timesteps - - # Apply shift if not using dynamic shifting - if (!use_dynamic_shifting) { - sigmas <- shift * sigmas / (1 + (shift - 1) * sigmas) - } - - timesteps <- sigmas * num_train_timesteps - - list( - sigmas = sigmas, - timesteps = timesteps, - num_train_timesteps = num_train_timesteps, - num_inference_steps = NULL, - step_index = NULL, - config = list( - num_train_timesteps = num_train_timesteps, - shift = shift, - use_dynamic_shifting = use_dynamic_shifting, - base_shift = base_shift, - max_shift = max_shift, - base_seq_len = base_seq_len, - max_seq_len = max_seq_len, - invert_sigmas = invert_sigmas, - shift_terminal = shift_terminal, - time_shift_type = time_shift_type +flowmatch_scheduler_create <- function(num_train_timesteps = 1000L, + shift = 1.0, + use_dynamic_shifting = FALSE, + base_shift = 0.5, max_shift = 1.15, + base_seq_len = 256L, + max_seq_len = 4096L, + invert_sigmas = FALSE, + shift_terminal = NULL, + time_shift_type = c("exponential", "linear")) { + time_shift_type <- match.arg(time_shift_type) + + # Initial timesteps (reversed, from max to min) + timesteps <- seq(from = 1, to = num_train_timesteps, + length.out = num_train_timesteps) + timesteps <- rev(timesteps) + + # Initial sigmas (normalized timesteps) + sigmas <- timesteps / num_train_timesteps + + # Apply shift if not using dynamic shifting + if (!use_dynamic_shifting) { + sigmas <- shift * sigmas / (1 + (shift - 1) * sigmas) + } + + timesteps <- sigmas * num_train_timesteps + + list( + sigmas = sigmas, + timesteps = timesteps, + num_train_timesteps = num_train_timesteps, + num_inference_steps = NULL, + step_index = NULL, + config = list( + num_train_timesteps = num_train_timesteps, + shift = shift, + use_dynamic_shifting = use_dynamic_shifting, + base_shift = base_shift, + max_shift = max_shift, + base_seq_len = base_seq_len, + max_seq_len = max_seq_len, + invert_sigmas = invert_sigmas, + shift_terminal = shift_terminal, + time_shift_type = time_shift_type + ) ) - ) } #' Calculate shift for dynamic shifting @@ -117,17 +115,13 @@ flowmatch_scheduler_create <- function( #' #' @return Numeric. The computed shift value (mu). #' @export -flowmatch_calculate_shift <- function( - seq_len, - base_seq_len = 256L, - max_seq_len = 4096L, - base_shift = 0.5, - max_shift = 1.15 -) { - m <- (max_shift - base_shift) / (max_seq_len - base_seq_len) - b <- base_shift - m * base_seq_len - mu <- seq_len * m + b - mu +flowmatch_calculate_shift <- function(seq_len, base_seq_len = 256L, + max_seq_len = 4096L, base_shift = 0.5, + max_shift = 1.15) { + m <- (max_shift - base_shift) / (max_seq_len - base_seq_len) + b <- base_shift - m * base_seq_len + mu <- seq_len * m + b + mu } #' Set timesteps for inference @@ -145,79 +139,72 @@ flowmatch_calculate_shift <- function( #' #' @return Updated scheduler with configured timesteps and sigmas. #' @export -flowmatch_set_timesteps <- function( - schedule, - num_inference_steps = 50L, - device = "cpu", - mu = NULL, - sigmas = NULL, - timesteps = NULL -) { - config <- schedule$config - - if (config$use_dynamic_shifting && is.null(mu)) { - stop("`mu` must be provided when use_dynamic_shifting is TRUE") - } - - schedule$num_inference_steps <- num_inference_steps - - # Get sigma_max and sigma_min - sigma_max <- max(schedule$sigmas) - sigma_min <- min(schedule$sigmas) - - # Create sigmas if not provided - if (is.null(sigmas)) { - if (is.null(timesteps)) { - # Linear spacing from sigma_max to sigma_min - timesteps <- seq( - from = sigma_max * config$num_train_timesteps, - to = sigma_min * config$num_train_timesteps, - length.out = num_inference_steps - ) +flowmatch_set_timesteps <- function(schedule, num_inference_steps = 50L, + device = "cpu", mu = NULL, sigmas = NULL, + timesteps = NULL) { + config <- schedule$config + + if (config$use_dynamic_shifting && is.null(mu)) { + stop("`mu` must be provided when use_dynamic_shifting is TRUE") } - sigmas <- timesteps / config$num_train_timesteps - } - - # Apply timestep shifting - if (config$use_dynamic_shifting) { - sigmas <- .flowmatch_time_shift( - mu = mu, - sigma = 1.0, - t = sigmas, - shift_type = config$time_shift_type - ) - } else { - shift <- config$shift - sigmas <- shift * sigmas / (1 + (shift - 1) * sigmas) - } - - # Stretch to terminal if configured - if (!is.null(config$shift_terminal)) { - sigmas <- .flowmatch_stretch_shift_to_terminal(sigmas, config$shift_terminal) - } - - # Convert to tensors - sigmas <- torch::torch_tensor(sigmas, dtype = torch::torch_float32()) - timesteps <- sigmas * config$num_train_timesteps - - # Append terminal sigma - if (config$invert_sigmas) { - sigmas <- 1.0 - sigmas + + schedule$num_inference_steps <- num_inference_steps + + # Get sigma_max and sigma_min + sigma_max <- max(schedule$sigmas) + sigma_min <- min(schedule$sigmas) + + # Create sigmas if not provided + if (is.null(sigmas)) { + if (is.null(timesteps)) { + # Linear spacing from sigma_max to sigma_min + timesteps <- seq(from = sigma_max * config$num_train_timesteps, + to = sigma_min * config$num_train_timesteps, + length.out = num_inference_steps) + } + sigmas <- timesteps / config$num_train_timesteps + } + + # Apply timestep shifting + if (config$use_dynamic_shifting) { + sigmas <- .flowmatch_time_shift( + mu = mu, + sigma = 1.0, + t = sigmas, + shift_type = config$time_shift_type + ) + } else { + shift <- config$shift + sigmas <- shift * sigmas / (1 + (shift - 1) * sigmas) + } + + # Stretch to terminal if configured + if (!is.null(config$shift_terminal)) { + sigmas <- .flowmatch_stretch_shift_to_terminal(sigmas, config$shift_terminal) + } + + # Convert to tensors + sigmas <- torch::torch_tensor(sigmas, dtype = torch::torch_float32()) timesteps <- sigmas * config$num_train_timesteps - sigmas <- torch::torch_cat(list(sigmas, torch::torch_ones(1))) - } else { - sigmas <- torch::torch_cat(list(sigmas, torch::torch_zeros(1))) - } - # Move to device - sigmas <- sigmas$to(device = device) - timesteps <- timesteps$to(device = device) + # Append terminal sigma + if (config$invert_sigmas) { + sigmas <- 1.0 - sigmas + timesteps <- sigmas * config$num_train_timesteps + sigmas <- torch::torch_cat(list(sigmas, torch::torch_ones(1))) + } else { + sigmas <- torch::torch_cat(list(sigmas, torch::torch_zeros(1))) + } - schedule$sigmas <- sigmas - schedule$timesteps <- timesteps - schedule$step_index <- NULL + # Move to device + sigmas <- sigmas$to(device = device) + timesteps <- timesteps$to(device = device) - schedule + schedule$sigmas <- sigmas + schedule$timesteps <- timesteps + schedule$step_index <- NULL + + schedule } #' Perform a FlowMatch scheduler step @@ -247,42 +234,34 @@ flowmatch_set_timesteps <- function( #' in continuous normalizing flows. #' #' @export -flowmatch_scheduler_step <- function( - model_output, - timestep, - sample, - schedule, - generator = NULL -) { - # Initialize step index if needed - if (is.null(schedule$step_index)) { - schedule$step_index <- .flowmatch_init_step_index(timestep, schedule) - } - - # Upcast sample for precision - sample <- sample$to(dtype = torch::torch_float32()) - - # Get current and next sigma - sigma_idx <- schedule$step_index - sigma <- schedule$sigmas[sigma_idx]$item() - sigma_next <- schedule$sigmas[sigma_idx + 1]$item() - - # Compute dt (step size) - dt <- sigma_next - sigma - - # Euler step: x_{t-1} = x_t + dt * v - prev_sample <- sample + dt * model_output - - # Cast back to model dtype - prev_sample <- prev_sample$to(dtype = model_output$dtype) - - # Increment step index - schedule$step_index <- schedule$step_index + 1L - - list( - prev_sample = prev_sample, - schedule = schedule - ) +flowmatch_scheduler_step <- function(model_output, timestep, sample, + schedule, generator = NULL) { + # Initialize step index if needed + if (is.null(schedule$step_index)) { + schedule$step_index <- .flowmatch_init_step_index(timestep, schedule) + } + + # Upcast sample for precision + sample <- sample$to(dtype = torch::torch_float32()) + + # Get current and next sigma + sigma_idx <- schedule$step_index + sigma <- schedule$sigmas[sigma_idx]$item() + sigma_next <- schedule$sigmas[sigma_idx + 1]$item() + + # Compute dt (step size) + dt <- sigma_next - sigma + + # Euler step: x_{t-1} = x_t + dt * v + prev_sample <- sample + dt * model_output + + # Cast back to model dtype + prev_sample <- prev_sample$to(dtype = model_output$dtype) + + # Increment step index + schedule$step_index <- schedule$step_index + 1L + + list(prev_sample = prev_sample, schedule = schedule) } #' Scale noise for flow matching forward process @@ -297,87 +276,67 @@ flowmatch_scheduler_step <- function( #' #' @return torch tensor. The noisy sample at timestep t. #' @export -flowmatch_scale_noise <- function( - sample, - timestep, - noise, - schedule -) { - # Get sigma for this timestep - sigmas <- schedule$sigmas$to(device = sample$device, dtype = sample$dtype) - - step_indices <- .flowmatch_index_for_timestep(timestep, schedule) - sigma <- sigmas[step_indices] - - # Expand sigma dimensions to match sample - while (length(sigma$shape) < length(sample$shape)) { - sigma <- sigma$unsqueeze(- 1) - } - - # Flow matching interpolation: x_t = (1 - sigma) * x_0 + sigma * noise - noisy_sample <- (1.0 - sigma) * sample + sigma * noise - - noisy_sample +flowmatch_scale_noise <- function(sample, timestep, noise, schedule) { + # Get sigma for this timestep + sigmas <- schedule$sigmas$to(device = sample$device, dtype = sample$dtype) + + step_indices <- .flowmatch_index_for_timestep(timestep, schedule) + sigma <- sigmas[step_indices] + + # Expand sigma dimensions to match sample + while (length(sigma$shape) < length(sample$shape)) { + sigma <- sigma$unsqueeze(-1) + } + + # Flow matching interpolation: x_t = (1 - sigma) * x_0 + sigma * noise + noisy_sample <- (1.0 - sigma) * sample + sigma * noise + + noisy_sample } # Internal helper functions -.flowmatch_time_shift <- function( - mu, - sigma, - t, - shift_type -) { - if (shift_type == "exponential") { - exp(mu) / (exp(mu) + (1 / t - 1) ^ sigma) - } else { - # linear - mu / (mu + (1 / t - 1) ^ sigma) - } +.flowmatch_time_shift <- function(mu, sigma, t, shift_type) { + if (shift_type == "exponential") { + exp(mu) / (exp(mu) + (1 / t - 1) ^ sigma) + } else { + # linear + mu / (mu + (1 / t - 1) ^ sigma) + } } -.flowmatch_stretch_shift_to_terminal <- function( - t, - shift_terminal -) { - one_minus_z <- 1 - t - scale_factor <- one_minus_z[length(one_minus_z)] / (1 - shift_terminal) - stretched_t <- 1 - (one_minus_z / scale_factor) - stretched_t +.flowmatch_stretch_shift_to_terminal <- function(t, shift_terminal) { + one_minus_z <- 1 - t + scale_factor <- one_minus_z[length(one_minus_z)] / (1 - shift_terminal) + stretched_t <- 1 - (one_minus_z / scale_factor) + stretched_t } -.flowmatch_init_step_index <- function( - timestep, - schedule -) { - # Find index for this timestep - idx <- .flowmatch_index_for_timestep(timestep, schedule) - idx +.flowmatch_init_step_index <- function(timestep, schedule) { + # Find index for this timestep + idx <- .flowmatch_index_for_timestep(timestep, schedule) + idx } -.flowmatch_index_for_timestep <- function( - timestep, - schedule -) { - timesteps <- schedule$timesteps - - if (inherits(timestep, "torch_tensor")) { - timestep <- timestep$item() - } - - # Find matching index - indices <- which(abs(as.numeric(timesteps) - timestep) < 1e-6) - - if (length(indices) == 0) { - # Find closest - indices <- which.min(abs(as.numeric(timesteps) - timestep)) - } - - # Return first match (or second if multiple, for numerical stability) - if (length(indices) > 1) { - indices[2] - } else { - indices[1] - } -} +.flowmatch_index_for_timestep <- function(timestep, schedule) { + timesteps <- schedule$timesteps + + if (inherits(timestep, "torch_tensor")) { + timestep <- timestep$item() + } + + # Find matching index + indices <- which(abs(as.numeric(timesteps) - timestep) < 1e-6) + if (length(indices) == 0) { + # Find closest + indices <- which.min(abs(as.numeric(timesteps) - timestep)) + } + + # Return first match (or second if multiple, for numerical stability) + if (length(indices) > 1) { + indices[2] + } else { + indices[1] + } +} diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R new file mode 100644 index 0000000..361be65 --- /dev/null +++ b/R/fp8_ltx23.R @@ -0,0 +1,372 @@ +#' FP8 Weight Storage for the LTX-2.3 Transformer +#' +#' GPU-poor weight handling: the large attention/FFN linears of the DiT +#' are stored as float8_e4m3fn with per-tensor scales (the official LTX +#' quantization policy), kept CPU-resident (optionally pinned), and +#' dequantized on the compute device inside each forward. Everything +#' else (norms, embeddings, modulation tables, biases) stays bfloat16. +#' Requires a safetensors build with F8 dtype support. +#' +#' @name fp8_ltx23 +NULL + +# Official LTX fp8 cast policy: attention and FFN linear weights inside +# the transformer blocks (the dotless ff suffixes also catch audio_ff) +.ltx23_fp8_cast_pattern <- paste0( + "^transformer_blocks\\.[0-9]+\\..*", + "(to_q|to_k|to_v|to_out\\.0|ff\\.net\\.0\\.proj|ff\\.net\\.2)\\.weight$" +) + +#' Test whether a mapped DiT key is in the official fp8 cast set +#' +#' @param mapped_key Character vector of mapped (diffusers-style) +#' parameter names. +#' +#' @return Logical vector. +#' +#' @export +ltx23_is_fp8_cast_key <- function(mapped_key) { + grepl(.ltx23_fp8_cast_pattern, mapped_key) +} + +#' FP8 linear layer +#' +#' Weight lives as float8_e4m3fn plus a float32 scale in plain module +#' fields (so \code{$to(device)} moves only the bias); the forward pass +#' ships 1 byte/param to the input's device, upcasts, rescales, and runs +#' \code{nnf_linear}. +#' +#' @param out_features,in_features Integers. +#' @param bias Logical. +#' +#' @export +ltx23_fp8_linear <- torch::nn_module( + "ltx23_fp8_linear", + initialize = function(out_features, in_features, bias = TRUE) { + self$out_features <- as.integer(out_features) + self$in_features <- as.integer(in_features) + self$weight_fp8 <- NULL + self$weight_scale <- NULL + if (bias) { + self$bias <- torch::nn_parameter(torch::torch_zeros(out_features)) + } +}, + set_fp8_weight = function(weight, scale, pin = FALSE) { + weight <- weight$to(device = "cpu") + if (pin && torch::cuda_is_available()) { + weight <- weight$pin_memory(device = torch::torch_device("cuda")) + } + self$weight_fp8 <- weight + self$weight_scale <- scale$to(device = "cpu", + dtype = torch::torch_float32()) + invisible(self) +}, + forward = function(x) { + # Transfer fp8 bytes first, cast on the compute device second + w <- self$weight_fp8$to(device = x$device, non_blocking = TRUE) + w <- w$to(dtype = x$dtype) * self$weight_scale$to(device = x$device, dtype = x$dtype) + out <- torch::nnf_linear(x, w, self$bias) + rm(w) + out +} +) + +# Navigate a dotted parameter path to a submodule. Numeric segments +# index module lists (0-based names -> 1-based R); named segments go +# through the nn_module `$` accessor. +.ltx23_walk_module <- function(module, segments) { + cur <- module + for (seg in segments) { + cur <- if (grepl("^[0-9]+$", seg)) { + cur[[as.integer(seg) + 1L]] + } else { + do.call(`$`, list(cur, seg)) + } + if (is.null(cur)) { + return(NULL) + } + } + cur +} + +#' Quantize an LTX-2.3 checkpoint to FP8 shards +#' +#' Streams the single-file bf16 checkpoint tensor by tensor. DiT +#' attention/FFN linear weights are stored as float8_e4m3fn with a +#' float32 absmax/448 per-tensor scale (\code{_scale} sibling); +#' everything else is copied through unchanged. Output shards carry the +#' original key names plus a manifest for skip-if-exists. +#' +#' @param checkpoint_path Source .safetensors (46 GB bf16 single file). +#' @param output_dir Output directory for shards + manifest. +#' @param shard_bytes Numeric. Approximate shard size (default 4 GB). +#' @param force Logical. Re-quantize even if a valid manifest exists. +#' @param verbose Logical. +#' +#' @return Invisibly, the manifest list. +#' +#' @export +ltx23_quantize_fp8 <- function(checkpoint_path, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), + shard_bytes = 4e9, force = FALSE, + verbose = TRUE) { + manifest_path <- file.path(output_dir, "manifest.json") + if (!force && file.exists(manifest_path)) { + manifest <- jsonlite::fromJSON(manifest_path) + if (all(file.exists(file.path(output_dir, manifest$shards)))) { + if (verbose) { + message("FP8 artifact already present: ", output_dir) + } + return(invisible(manifest)) + } + } + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + + ckpt <- ltx23_open_checkpoint(checkpoint_path) + fp8 <- torch::torch_float8_e4m3fn() + + shard <- list() + shard_size <- 0 + shard_files <- character(0) + n_cast <- 0L + + flush_shard <- function() { + if (!length(shard)) { + return() + } + fname <- sprintf("ltx2.3-fp8-%05d.safetensors", + length(shard_files) + 1L) + safetensors::safe_save_file(shard, file.path(output_dir, fname)) + shard_files[[length(shard_files) + 1L]] <<- fname + if (verbose) { + message(sprintf(" wrote %s (%.2f GB, %d tensors)", + fname, shard_size / 1e9, length(shard))) + } + shard <<- list() + shard_size <<- 0 + gc(verbose = FALSE) + } + + keys <- ckpt$keys + for (i in seq_along(keys)) { + key <- keys[[i]] + tensor <- ckpt$handle$get_tensor(key) + + mapped <- ltx23_map_dit_key(key) + if (startsWith(key, "model.diffusion_model.") && + ltx23_is_fp8_cast_key(mapped)) { + torch::with_no_grad({ + scale <- tensor$abs()$max()$to(dtype = torch::torch_float32())$ + clamp(min = 1e-12) / 448 + shard[[key]] <- (tensor$to(dtype = torch::torch_float32()) / scale)$to(dtype = fp8) + shard[[paste0(key, "_scale")]] <- scale + }) + shard_size <- shard_size + prod(tensor$shape) + n_cast <- n_cast + 1L + } else { + shard[[key]] <- tensor + shard_size <- shard_size + prod(tensor$shape) * 2 + } + rm(tensor) + + if (shard_size >= shard_bytes) { + flush_shard() + } + if (i %% 200L == 0L) { + gc(verbose = FALSE) + if (verbose) { + message(sprintf(" quantizing %d/%d tensors", i, length(keys))) + } + } + } + flush_shard() + + manifest <- list( + source = basename(checkpoint_path), + model_version = ckpt$version, + shards = shard_files, + tensors = length(keys), + fp8_cast = n_cast, + config = ckpt$config + ) + jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE, pretty = TRUE) + if (verbose) { + message(sprintf("Quantized %d/%d tensors to fp8 across %d shards: %s", + n_cast, length(keys), length(shard_files), output_dir)) + } + invisible(manifest) +} + +#' Open an FP8 shard directory as a checkpoint +#' +#' Presents the sharded fp8 artifact through the same interface as +#' \code{\link{ltx23_open_checkpoint}} so the group loaders work +#' unchanged. +#' +#' @param dir The fp8 artifact directory (with manifest.json). +#' +#' @return An \code{ltx23_checkpoint}. +#' +#' @export +ltx23_open_fp8_checkpoint <- function(dir) { + manifest_path <- file.path(dir, "manifest.json") + if (!file.exists(manifest_path)) { + stop("No manifest.json in ", dir, "; run ltx23_quantize_fp8() first.") + } + manifest <- jsonlite::fromJSON(manifest_path, simplifyVector = TRUE) + + handles <- lapply(file.path(dir, manifest$shards), function(p) { + safetensors::safetensors$new(p, framework = "torch") + }) + key_to_handle <- list() + for (h in handles) { + for (k in setdiff(h$keys(), "__metadata__")) { + key_to_handle[[k]] <- h + } + } + + handle <- list( + get_tensor = function(key) { + h <- key_to_handle[[key]] + if (is.null(h)) stop("Key not found in fp8 shards: ", key) + h$get_tensor(key) + } + ) + + structure( + list(handle = handle, keys = names(key_to_handle), + version = manifest$model_version, config = manifest$config, + format = manifest$format %||% "fp8", path = dir), + class = "ltx23_checkpoint" + ) +} + +#' Load the LTX-2.3 transformer with FP8 weights +#' +#' Builds the transformer, swaps the official cast-set linears for +#' \code{\link{ltx23_fp8_linear}}, loads fp8 weights CPU-side (optionally +#' pinned) and everything else as bfloat16 on \code{device}. Sets +#' \code{options(diffuseR.block_gc = TRUE)} so the transformer runs +#' per-block garbage collection over the dequantized temporaries. +#' +#' @param ckpt An fp8 \code{ltx23_checkpoint} +#' (\code{\link{ltx23_open_fp8_checkpoint}}). +#' @param device Character. Device for the resident (non-fp8) weights. +#' @param pin Logical. Pin the fp8 host memory for faster transfers. +#' @param verbose Logical. +#' @param ... Passed to \code{\link{ltx23_transformer}} (tiny test configs). +#' +#' @return The loaded \code{ltx23_transformer}. +#' +#' @export +ltx23_load_transformer_fp8 <- function(ckpt, device = "cuda", pin = TRUE, + verbose = TRUE, ...) { + stopifnot(inherits(ckpt, "ltx23_checkpoint")) + model <- ltx23_transformer(...) + model$to(dtype = torch::torch_bfloat16()) + + groups <- ltx23_split_keys(ckpt$keys) + dit_keys <- groups$dit + scale_keys <- dit_keys[endsWith(dit_keys, ".weight_scale")] + main_keys <- setdiff(dit_keys, scale_keys) + + dests <- c(model$named_parameters(), model$named_buffers()) + filled <- character(0) + unmapped <- character(0) + + torch::with_no_grad({ + for (i in seq_along(main_keys)) { + key <- main_keys[[i]] + mapped <- ltx23_map_dit_key(key) + + if (ltx23_is_fp8_cast_key(mapped) && + paste0(key, "_scale") %in% scale_keys) { + segments <- strsplit(mapped, ".", fixed = TRUE)[[1]] + parent <- .ltx23_walk_module(model, utils::head(segments, -2L)) + leaf <- segments[length(segments) - 1L] + old <- .ltx23_walk_module(parent, leaf) + if (is.null(old)) { + unmapped <- c(unmapped, key) + next + } + weight <- ckpt$handle$get_tensor(key) + scale <- ckpt$handle$get_tensor(paste0(key, "_scale")) + fp8_mod <- ltx23_fp8_linear(weight$shape[1], weight$shape[2], + bias = !is.null(old$bias)) + if (!is.null(old$bias)) { + # Adopt the original bias parameter so the separate bias key, + # which copies through the pre-swap destination map, lands here + fp8_mod$bias <- old$bias + } + fp8_mod$set_fp8_weight(weight, scale, pin = pin) + do.call(`$<-`, list(parent, leaf, fp8_mod)) + filled <- c(filled, mapped) + rm(weight, scale) + } else { + dest <- dests[[mapped]] + if (is.null(dest)) { + unmapped <- c(unmapped, key) + next + } + dest$copy_(ckpt$handle$get_tensor(key)) + filled <- c(filled, mapped) + } + + if (i %% 100L == 0L) { + gc(verbose = FALSE) + if (verbose && i %% 500L == 0L) { + message(sprintf(" loaded %d/%d transformer tensors", i, length(main_keys))) + } + } + } + }) + gc(verbose = FALSE) + + if (length(unmapped)) { + stop("FP8 transformer load: ", length(unmapped), " unmapped keys, e.g. ", + paste(utils::head(unmapped, 3), collapse = ", ")) + } + + # Weight params replaced by fp8 modules won't be "filled"; account for them + expected_missing <- grepl(.ltx23_fp8_cast_pattern, names(dests)) + unfilled <- setdiff(names(dests)[!expected_missing], filled) + if (length(unfilled)) { + stop("FP8 transformer load: ", length(unfilled), " unfilled params, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + + # Residents (norms, embeddings, tables, biases) to the compute device; + # fp8 fields stay on the CPU because they are plain fields + model$to(device = device) + model$eval() + # fp8 dequant allocates a full weight per linear; per-block gc + # keeps those temporaries from accumulating + options(diffuseR.block_gc = TRUE) + if (verbose) { + message("Transformer ready: fp8 weights CPU-resident, rest on ", device) + } + model +} + +#' Set the attention query-chunk size across a transformer +#' +#' R torch has no fused attention, so the [B, H, S, S] matrix +#' materializes; chunking queries bounds the peak. NULL disables +#' chunking. +#' +#' @param transformer An \code{ltx23_transformer}. +#' @param chunk Integer or NULL. +#' +#' @return Invisibly, the transformer. +#' +#' @export +ltx23_set_attn_chunk <- function(transformer, chunk) { + for (i in seq_along(transformer$transformer_blocks)) { + block <- transformer$transformer_blocks[[i]] + for (name in c("attn1", "audio_attn1", "attn2", "audio_attn2", + "audio_to_video_attn", "video_to_audio_attn")) { + block[[name]]$attn_chunk <- chunk + } + } + invisible(transformer) +} diff --git a/R/gemma3_text_encoder.R b/R/gemma3_text_encoder.R index 5d111f3..3c5fd90 100644 --- a/R/gemma3_text_encoder.R +++ b/R/gemma3_text_encoder.R @@ -22,22 +22,22 @@ #' @param eps Numeric. Epsilon for numerical stability. #' @keywords internal gemma3_rms_norm <- torch::nn_module( - "Gemma3RMSNorm", - initialize = function( - dim, - eps = 1e-6 - ) { + "Gemma3RMSNorm", + initialize = function( + dim, + eps = 1e-6 + ) { self$eps <- eps self$weight <- torch::nn_parameter(torch::torch_zeros(dim)) - }, - forward = function(x) { +}, + forward = function(x) { input_dtype <- x$dtype x <- x$to(dtype = torch::torch_float32()) - variance <- x$pow(2)$mean(dim = - 1L, keepdim = TRUE) + variance <- x$pow(2)$mean(dim = -1L, keepdim = TRUE) x <- x * torch::torch_rsqrt(variance + self$eps) # Gemma adds 1 to the weight: (1 + weight) * x x$to(dtype = input_dtype) * (self$weight$add(1)) - } +} ) # ----------------------------------------------------------------------------- @@ -54,13 +54,13 @@ gemma3_rms_norm <- torch::nn_module( #' @param scaling_factor Numeric. Optional scaling factor for extended context. #' @keywords internal gemma3_rotary_embedding <- torch::nn_module( - "Gemma3RotaryEmbedding", - initialize = function( - dim, - max_position_embeddings = 8192L, - base = 10000.0, - scaling_factor = 1.0 - ) { + "Gemma3RotaryEmbedding", + initialize = function( + dim, + max_position_embeddings = 8192L, + base = 10000.0, + scaling_factor = 1.0 + ) { self$dim <- dim self$max_position_embeddings <- max_position_embeddings self$base <- base @@ -71,12 +71,12 @@ gemma3_rotary_embedding <- torch::nn_module( inv_freq <- 1.0 / (base ^ (torch::torch_arange(0, dim - 1L, 2L)$to(dtype = torch::torch_float32()) / dim)) inv_freq <- inv_freq / scaling_factor self$inv_freq <- torch::nn_buffer(inv_freq, persistent = FALSE) - }, +}, - forward = function( - x, - position_ids - ) { + forward = function( + x, + position_ids + ) { # x: [batch, seq_len, ...] # position_ids: [batch, seq_len] or [1, seq_len] @@ -85,17 +85,18 @@ gemma3_rotary_embedding <- torch::nn_module( # HuggingFace uses matmul: inv_freq[None,:,None] @ position_ids[:,None,:] # We use einsum which is equivalent - freqs <- torch::torch_einsum("bs,d->bsd", list(position_ids_float, self$inv_freq$to(device = x$device))) + freqs <- torch::torch_einsum("bs,d->bsd", + list(position_ids_float, self$inv_freq$to(device = x$device))) # Concatenate for full dimension: [batch, seq_len, dim] # This duplicates freqs so that first half and second half are identical - emb <- torch::torch_cat(list(freqs, freqs), dim = - 1L) + emb <- torch::torch_cat(list(freqs, freqs), dim = -1L) cos_emb <- torch::torch_cos(emb)$to(dtype = x$dtype) sin_emb <- torch::torch_sin(emb)$to(dtype = x$dtype) list(cos_emb, sin_emb) - } +} ) #' Apply rotary position embeddings @@ -105,33 +106,28 @@ gemma3_rotary_embedding <- torch::nn_module( #' @param cos Cosine embeddings [batch, seq, head_dim] #' @param sin Sine embeddings [batch, seq, head_dim] #' @keywords internal -apply_rotary_pos_emb <- function( - q, - k, - cos, - sin -) { - # Reshape cos/sin for broadcasting: [batch, 1, seq, head_dim] - cos <- cos$unsqueeze(2L) - sin <- sin$unsqueeze(2L) - - # Apply rotation - q_embed <- (q * cos) + (rotate_half(q) * sin) - k_embed <- (k * cos) + (rotate_half(k) * sin) - - list(q_embed, k_embed) +apply_rotary_pos_emb <- function(q, k, cos, sin) { + # Reshape cos/sin for broadcasting: [batch, 1, seq, head_dim] + cos <- cos$unsqueeze(2L) + sin <- sin$unsqueeze(2L) + + # Apply rotation + q_embed <- (q * cos) + (rotate_half(q) * sin) + k_embed <- (k * cos) + (rotate_half(k) * sin) + + list(q_embed, k_embed) } #' Rotate half of the hidden dims #' @keywords internal rotate_half <- function(x) { - # x: [..., dim] - # Split into two halves and rotate - dim <- x$shape[length(x$shape)] - half_dim <- dim %/% 2L - x1 <- x[.., 1:half_dim] - x2 <- x[.., (half_dim + 1L) :dim] - torch::torch_cat(list(- x2, x1), dim = - 1L) + # x: [..., dim] + # Split into two halves and rotate + dim <- x$shape[length(x$shape)] + half_dim <- dim %/% 2L + x1 <- x[.., 1:half_dim] + x2 <- x[.., (half_dim + 1L):dim] + torch::torch_cat(list(-x2, x1), dim = -1L) } #' Repeat KV heads for GQA (Grouped Query Attention) @@ -141,22 +137,20 @@ rotate_half <- function(x) { #' @param n_rep Number of repetitions per KV head #' @return Tensor of shape [batch, num_kv_heads * n_rep, seq_len, head_dim] #' @keywords internal -repeat_kv <- function( - hidden_states, - n_rep -) { - if (n_rep == 1L) { - return(hidden_states) - } - batch <- hidden_states$shape[1] - num_kv_heads <- hidden_states$shape[2] - seq_len <- hidden_states$shape[3] - head_dim <- hidden_states$shape[4] - - # Add dimension and expand: [batch, kv_heads, 1, seq, dim] -> [batch, kv_heads, n_rep, seq, dim] - hidden_states <- hidden_states$unsqueeze(3L)$expand(c(batch, num_kv_heads, n_rep, seq_len, head_dim)) - # Reshape to interleave: [batch, kv_heads * n_rep, seq, dim] - hidden_states$reshape(c(batch, num_kv_heads * n_rep, seq_len, head_dim)) +repeat_kv <- function(hidden_states, n_rep) { + if (n_rep == 1L) { + return(hidden_states) + } + batch <- hidden_states$shape[1] + num_kv_heads <- hidden_states$shape[2] + seq_len <- hidden_states$shape[3] + head_dim <- hidden_states$shape[4] + + # Add dimension and expand: [batch, kv_heads, 1, seq, dim] -> [batch, kv_heads, n_rep, seq, dim] + hidden_states <- hidden_states$unsqueeze(3L)$expand(c(batch, num_kv_heads, + n_rep, seq_len, head_dim)) + # Reshape to interleave: [batch, kv_heads * n_rep, seq, dim] + hidden_states$reshape(c(batch, num_kv_heads * n_rep, seq_len, head_dim)) } # ----------------------------------------------------------------------------- @@ -170,22 +164,23 @@ repeat_kv <- function( #' @param config List with hidden_size and intermediate_size. #' @keywords internal gemma3_mlp <- torch::nn_module( - "Gemma3MLP", - initialize = function(config) { + "Gemma3MLP", + initialize = function(config) { self$hidden_size <- config$hidden_size self$intermediate_size <- config$intermediate_size - self$gate_proj <- torch::nn_linear(self$hidden_size, self$intermediate_size, bias = FALSE) + self$gate_proj <- torch::nn_linear(self$hidden_size, + self$intermediate_size, bias = FALSE) self$up_proj <- torch::nn_linear(self$hidden_size, self$intermediate_size, bias = FALSE) self$down_proj <- torch::nn_linear(self$intermediate_size, self$hidden_size, bias = FALSE) # Gemma uses approximate GELU self$act_fn <- function(x) torch::nnf_gelu(x, approximate = "tanh") - }, - forward = function(x) { +}, + forward = function(x) { # Gated Linear Unit: down(act(gate(x)) * up(x)) self$down_proj(self$act_fn(self$gate_proj(x)) * self$up_proj(x)) - } +} ) # ----------------------------------------------------------------------------- @@ -201,11 +196,11 @@ gemma3_mlp <- torch::nn_module( #' @param layer_idx Integer. Layer index for layer-specific settings. #' @keywords internal gemma3_attention <- torch::nn_module( - "Gemma3Attention", - initialize = function( - config, - layer_idx = 0L - ) { + "Gemma3Attention", + initialize = function( + config, + layer_idx = 0L + ) { self$config <- config self$layer_idx <- layer_idx @@ -223,7 +218,9 @@ gemma3_attention <- torch::nn_module( self$attn_logit_softcapping <- config$attn_logit_softcapping %||% NULL # Projections - self$q_proj <- torch::nn_linear(self$hidden_size, self$num_heads * self$head_dim, bias = FALSE) + self$q_proj <- torch::nn_linear(self$hidden_size, + self$num_heads * self$head_dim, + bias = FALSE) self$k_proj <- torch::nn_linear(self$hidden_size, self$num_key_value_heads * self$head_dim, bias = FALSE) self$v_proj <- torch::nn_linear(self$hidden_size, self$num_key_value_heads * self$head_dim, bias = FALSE) self$o_proj <- torch::nn_linear(self$num_heads * self$head_dim, self$hidden_size, bias = FALSE) @@ -235,21 +232,21 @@ gemma3_attention <- torch::nn_module( # Scaling factor self$scaling <- config$query_pre_attn_scalar %||% self$head_dim self$scaling <- 1.0 / sqrt(self$scaling) - }, +}, - get_is_sliding = function(layer_idx) { + get_is_sliding = function(layer_idx) { # Gemma3 uses alternating sliding window pattern # Default: every 6th layer is global attention sliding_window_pattern <- self$config$sliding_window_pattern %||% 6L return((layer_idx + 1L) %% sliding_window_pattern != 0L) - }, - - forward = function( - hidden_states, - attention_mask = NULL, - position_embeddings_global = NULL, - position_embeddings_local = NULL - ) { +}, + + forward = function( + hidden_states, + attention_mask = NULL, + position_embeddings_global = NULL, + position_embeddings_local = NULL + ) { batch_size <- hidden_states$shape[1] seq_len <- hidden_states$shape[2] @@ -269,42 +266,42 @@ gemma3_attention <- torch::nn_module( # Apply rotary embeddings - use local for sliding window layers, global otherwise if (self$is_sliding) { - position_embeddings <- position_embeddings_local + position_embeddings <- position_embeddings_local } else { - position_embeddings <- position_embeddings_global + position_embeddings <- position_embeddings_global } if (!is.null(position_embeddings)) { - cos <- position_embeddings[[1]] - sin <- position_embeddings[[2]] - rope_result <- apply_rotary_pos_emb(query_states, key_states, cos, sin) - query_states <- rope_result[[1]] - key_states <- rope_result[[2]] + cos <- position_embeddings[[1]] + sin <- position_embeddings[[2]] + rope_result <- apply_rotary_pos_emb(query_states, key_states, cos, sin) + query_states <- rope_result[[1]] + key_states <- rope_result[[2]] } # Repeat K/V for GQA (interleaved: [k0,k0,k1,k1,...] not [k0,k1,...,k0,k1,...]) if (self$num_key_value_groups > 1L) { - key_states <- repeat_kv(key_states, self$num_key_value_groups) - value_states <- repeat_kv(value_states, self$num_key_value_groups) + key_states <- repeat_kv(key_states, self$num_key_value_groups) + value_states <- repeat_kv(value_states, self$num_key_value_groups) } # Compute attention scores - attn_weights <- torch::torch_matmul(query_states, key_states$transpose(- 2L, - 1L)) * self$scaling + attn_weights <- torch::torch_matmul(query_states, key_states$transpose(-2L, -1L)) * self$scaling # Apply softcapping if configured if (!is.null(self$attn_logit_softcapping)) { - attn_weights <- attn_weights / self$attn_logit_softcapping - attn_weights <- torch::torch_tanh(attn_weights) - attn_weights <- attn_weights * self$attn_logit_softcapping + attn_weights <- attn_weights / self$attn_logit_softcapping + attn_weights <- torch::torch_tanh(attn_weights) + attn_weights <- attn_weights * self$attn_logit_softcapping } # Apply attention mask (includes causal mask from layer) # Note: sliding window constraint is handled at the model level, not here if (!is.null(attention_mask)) { - attn_weights <- attn_weights + attention_mask + attn_weights <- attn_weights + attention_mask } # Softmax - attn_weights <- torch::nnf_softmax(attn_weights, dim = - 1L) + attn_weights <- torch::nnf_softmax(attn_weights, dim = -1L) attn_weights <- attn_weights$to(dtype = value_states$dtype) # Apply attention to values @@ -318,31 +315,27 @@ gemma3_attention <- torch::nn_module( attn_output <- self$o_proj(attn_output) attn_output - } +} ) #' Create sliding window causal attention mask #' @keywords internal -create_sliding_window_mask <- function( - seq_len, - window_size, - device = "cpu" -) { - # Create causal mask - mask <- torch::torch_ones(c(seq_len, seq_len), device = device) - mask <- torch::torch_triu(mask, diagonal = 1L) - - # Create sliding window mask (tokens beyond window are masked) - row_idx <- torch::torch_arange(0, seq_len - 1L, device = device)$unsqueeze(2L) - col_idx <- torch::torch_arange(0, seq_len - 1L, device = device)$unsqueeze(1L) - window_mask <- (row_idx - col_idx) > window_size - - # Combine: mask future tokens and tokens outside window - mask <- mask$logical_or(window_mask) - - # Convert to additive mask - mask <- mask$to(dtype = torch::torch_float32()) * (- 1e9) - mask +create_sliding_window_mask <- function(seq_len, window_size, device = "cpu") { + # Create causal mask + mask <- torch::torch_ones(c(seq_len, seq_len), device = device) + mask <- torch::torch_triu(mask, diagonal = 1L) + + # Create sliding window mask (tokens beyond window are masked) + row_idx <- torch::torch_arange(0, seq_len - 1L, device = device)$unsqueeze(2L) + col_idx <- torch::torch_arange(0, seq_len - 1L, device = device)$unsqueeze(1L) + window_mask <- (row_idx - col_idx) > window_size + + # Combine: mask future tokens and tokens outside window + mask <- mask$logical_or(window_mask) + + # Convert to additive mask + mask <- mask$to(dtype = torch::torch_float32()) * (-1e9) + mask } # ----------------------------------------------------------------------------- @@ -357,38 +350,39 @@ create_sliding_window_mask <- function( #' @param layer_idx Integer. Layer index. #' @keywords internal gemma3_decoder_layer <- torch::nn_module( - "Gemma3DecoderLayer", - initialize = function( - config, - layer_idx = 0L - ) { + "Gemma3DecoderLayer", + initialize = function( + config, + layer_idx = 0L + ) { self$hidden_size <- config$hidden_size self$self_attn <- gemma3_attention(config, layer_idx = layer_idx) self$mlp <- gemma3_mlp(config) - self$input_layernorm <- gemma3_rms_norm(config$hidden_size, eps = config$rms_norm_eps %||% 1e-6) + self$input_layernorm <- gemma3_rms_norm(config$hidden_size, + eps = config$rms_norm_eps %||% 1e-6) self$post_attention_layernorm <- gemma3_rms_norm(config$hidden_size, eps = config$rms_norm_eps %||% 1e-6) # Gemma3 has additional pre-feedforward norm self$pre_feedforward_layernorm <- gemma3_rms_norm(config$hidden_size, eps = config$rms_norm_eps %||% 1e-6) self$post_feedforward_layernorm <- gemma3_rms_norm(config$hidden_size, eps = config$rms_norm_eps %||% 1e-6) - }, - - forward = function( - hidden_states, - attention_mask = NULL, - position_embeddings_global = NULL, - position_embeddings_local = NULL - ) { +}, + + forward = function( + hidden_states, + attention_mask = NULL, + position_embeddings_global = NULL, + position_embeddings_local = NULL + ) { residual <- hidden_states # Pre-norm self-attention hidden_states <- self$input_layernorm(hidden_states) hidden_states <- self$self_attn(hidden_states, - attention_mask = attention_mask, - position_embeddings_global = position_embeddings_global, - position_embeddings_local = position_embeddings_local) + attention_mask = attention_mask, + position_embeddings_global = position_embeddings_global, + position_embeddings_local = position_embeddings_local) hidden_states <- self$post_attention_layernorm(hidden_states) hidden_states <- residual + hidden_states @@ -400,7 +394,7 @@ gemma3_decoder_layer <- torch::nn_module( hidden_states <- residual + hidden_states hidden_states - } +} ) # ----------------------------------------------------------------------------- @@ -414,48 +408,49 @@ gemma3_decoder_layer <- torch::nn_module( #' @param config Model configuration list. #' @export gemma3_text_model <- torch::nn_module( - "Gemma3TextModel", - initialize = function(config) { + "Gemma3TextModel", + initialize = function(config) { self$config <- config self$hidden_size <- config$hidden_size self$num_layers <- config$num_hidden_layers # Token embeddings (scaled by sqrt(hidden_size)) - self$embed_tokens <- torch::nn_embedding(config$vocab_size, config$hidden_size) + self$embed_tokens <- torch::nn_embedding(config$vocab_size, + config$hidden_size) self$embed_scale <- sqrt(config$hidden_size) # Rotary embeddings - Gemma3 uses different frequencies for global vs local layers head_dim <- config$head_dim %||% (config$hidden_size %/% config$num_attention_heads) # Global rotary embedding (for every 6th layer - layers 5, 11, 17, etc.) self$rotary_emb <- gemma3_rotary_embedding( - dim = head_dim, - max_position_embeddings = config$max_position_embeddings %||% 8192L, - base = config$rope_theta %||% 1000000.0, # Global default - scaling_factor = config$rope_scaling$factor %||% 1.0 + dim = head_dim, + max_position_embeddings = config$max_position_embeddings %||% 8192L, + base = config$rope_theta %||% 1000000.0, # Global default + scaling_factor = config$rope_scaling$factor %||% 1.0 ) # Local rotary embedding (for sliding window layers) self$rotary_emb_local <- gemma3_rotary_embedding( - dim = head_dim, - max_position_embeddings = config$max_position_embeddings %||% 8192L, - base = config$rope_local_base_freq %||% 10000.0, # Local default - scaling_factor = 1.0# No scaling for local + dim = head_dim, + max_position_embeddings = config$max_position_embeddings %||% 8192L, + base = config$rope_local_base_freq %||% 10000.0, # Local default + scaling_factor = 1.0 # No scaling for local ) # Decoder layers self$layers <- torch::nn_module_list(lapply(seq_len(self$num_layers), function(i) { - gemma3_decoder_layer(config, layer_idx = i - 1L) # 0-indexed - })) + gemma3_decoder_layer(config, layer_idx = i - 1L) # 0-indexed + })) # Final norm self$norm <- gemma3_rms_norm(config$hidden_size, eps = config$rms_norm_eps %||% 1e-6) - }, - - forward = function( - input_ids, - attention_mask = NULL, - position_ids = NULL, - output_hidden_states = FALSE - ) { +}, + + forward = function( + input_ids, + attention_mask = NULL, + position_ids = NULL, + output_hidden_states = FALSE + ) { batch_size <- input_ids$shape[1] seq_len <- input_ids$shape[2] @@ -465,8 +460,8 @@ gemma3_text_model <- torch::nn_module( # Position IDs if (is.null(position_ids)) { - position_ids <- torch::torch_arange(0, seq_len - 1L, device = input_ids$device) - position_ids <- position_ids$unsqueeze(1L)$expand(c(batch_size, - 1L)) + position_ids <- torch::torch_arange(0, seq_len - 1L, device = input_ids$device) + position_ids <- position_ids$unsqueeze(1L)$expand(c(batch_size, -1L)) } # Compute rotary embeddings (both global and local) @@ -475,56 +470,51 @@ gemma3_text_model <- torch::nn_module( # Prepare causal attention mask if (is.null(attention_mask)) { - # Create causal mask - causal_mask <- torch::torch_ones(c(seq_len, seq_len), device = hidden_states$device) - causal_mask <- torch::torch_triu(causal_mask, diagonal = 1L) - causal_mask <- causal_mask$to(dtype = hidden_states$dtype) * (- 1e9) - causal_mask <- causal_mask$unsqueeze(1L)$unsqueeze(1L) + # Create causal mask + causal_mask <- torch::torch_ones(c(seq_len, seq_len), device = hidden_states$device) + causal_mask <- torch::torch_triu(causal_mask, diagonal = 1L) + causal_mask <- causal_mask$to(dtype = hidden_states$dtype) * (-1e9) + causal_mask <- causal_mask$unsqueeze(1L)$unsqueeze(1L) } else { - # Convert attention mask to additive mask - # attention_mask: [batch, seq_len] with 1 for valid, 0 for padding - causal_mask <- torch::torch_ones(c(seq_len, seq_len), device = hidden_states$device) - causal_mask <- torch::torch_triu(causal_mask, diagonal = 1L) + # Convert attention mask to additive mask + # attention_mask: [batch, seq_len] with 1 for valid, 0 for padding + causal_mask <- torch::torch_ones(c(seq_len, seq_len), device = hidden_states$device) + causal_mask <- torch::torch_triu(causal_mask, diagonal = 1L) - # Expand padding mask: [batch, 1, 1, seq_len] - padding_mask <- (1 - attention_mask)$unsqueeze(2L)$unsqueeze(2L) + # Expand padding mask: [batch, 1, 1, seq_len] + padding_mask <- (1 - attention_mask)$unsqueeze(2L)$unsqueeze(2L) - # Combine masks - causal_mask <- (causal_mask$unsqueeze(1L)$unsqueeze(1L) + padding_mask)$to(dtype = hidden_states$dtype) * (- 1e9) + # Combine masks + causal_mask <- (causal_mask$unsqueeze(1L)$unsqueeze(1L) + padding_mask)$to(dtype = hidden_states$dtype) * (-1e9) } - # Collect hidden states if requested - if (output_hidden_states) { - all_hidden_states <- list(hidden_states) - } else { - all_hidden_states <- NULL - } + # HF transformers hidden-state semantics: record the state BEFORE + # each layer, then the post-final-norm output — num_layers + 1 + # entries total. (The un-normed last-layer output never appears; + # the LTX connectors expect exactly this 49-state stack.) + all_hidden_states <- NULL # Apply decoder layers for (i in seq_along(self$layers)) { - layer <- self$layers[[i]] - hidden_states <- layer(hidden_states, - attention_mask = causal_mask, - position_embeddings_global = position_embeddings_global, - position_embeddings_local = position_embeddings_local) - - if (output_hidden_states) { - all_hidden_states <- c(all_hidden_states, list(hidden_states)) - } + if (output_hidden_states) { + all_hidden_states <- c(all_hidden_states, list(hidden_states)) + } + layer <- self$layers[[i]] + hidden_states <- layer(hidden_states, + attention_mask = causal_mask, + position_embeddings_global = position_embeddings_global, + position_embeddings_local = position_embeddings_local) } # Final normalization hidden_states <- self$norm(hidden_states) if (output_hidden_states) { - all_hidden_states <- c(all_hidden_states, list(hidden_states)) + all_hidden_states <- c(all_hidden_states, list(hidden_states)) } - list( - last_hidden_state = hidden_states, - hidden_states = all_hidden_states - ) - } + list(last_hidden_state = hidden_states, hidden_states = all_hidden_states) +} ) # ----------------------------------------------------------------------------- @@ -538,23 +528,23 @@ gemma3_text_model <- torch::nn_module( #' @return List with model configuration parameters. #' @export gemma3_config_ltx2 <- function() { - list( - vocab_size = 262208L, - hidden_size = 3840L, - intermediate_size = 15360L, - num_hidden_layers = 48L, - num_attention_heads = 16L, - num_key_value_heads = 8L, - head_dim = 256L, - max_position_embeddings = 131072L, - rms_norm_eps = 1e-6, - rope_theta = 1000000.0, # Gemma3 uses 1M, not 10K - rope_scaling = list(factor = 8.0, type = "linear"), - sliding_window = 1024L, - sliding_window_pattern = 6L, - attn_logit_softcapping = NULL, # Gemma3 doesn't use softcapping by default - query_pre_attn_scalar = 256L - ) + list( + vocab_size = 262208L, + hidden_size = 3840L, + intermediate_size = 15360L, + num_hidden_layers = 48L, + num_attention_heads = 16L, + num_key_value_heads = 8L, + head_dim = 256L, + max_position_embeddings = 131072L, + rms_norm_eps = 1e-6, + rope_theta = 1000000.0, # Gemma3 uses 1M, not 10K + rope_scaling = list(factor = 8.0, type = "linear"), + sliding_window = 1024L, + sliding_window_pattern = 6L, + attn_logit_softcapping = NULL, # Gemma3 doesn't use softcapping by default + query_pre_attn_scalar = 256L + ) } # ----------------------------------------------------------------------------- @@ -571,162 +561,159 @@ gemma3_config_ltx2 <- function() { #' @param verbose Logical. Print loading progress. #' @return Initialized gemma3_text_model with loaded weights. #' @export -load_gemma3_text_encoder <- function( - model_path, - device = "cpu", - dtype = "float16", - verbose = TRUE -) { - # Load config - config_path <- file.path(model_path, "config.json") - if (!file.exists(config_path)) { - stop("Config file not found: ", config_path) - } - - config_raw <- jsonlite::fromJSON(config_path) - - # Extract text config if this is a multimodal model - if (!is.null(config_raw$text_config)) { - config_raw <- config_raw$text_config - } - - config <- list( - vocab_size = config_raw$vocab_size %||% 262208L, - hidden_size = config_raw$hidden_size %||% 3840L, - intermediate_size = config_raw$intermediate_size %||% 15360L, - num_hidden_layers = config_raw$num_hidden_layers %||% 48L, - num_attention_heads = config_raw$num_attention_heads %||% 16L, - num_key_value_heads = config_raw$num_key_value_heads %||% 8L, - head_dim = config_raw$head_dim %||% 256L, - max_position_embeddings = config_raw$max_position_embeddings %||% 131072L, - rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6, - rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M - rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0), - sliding_window = config_raw$sliding_window %||% 1024L, - sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L, - attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default) - query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L - ) - - if (verbose) { - message(sprintf("Creating Gemma3 model: %d layers, hidden_size=%d", - config$num_hidden_layers, config$hidden_size)) - } - - # Create model - model <- gemma3_text_model(config) - - # Find safetensor files - safetensor_files <- list.files(model_path, pattern = "\\.safetensors$", full.names = TRUE) - # Prefer model-* files over diffusion_pytorch_model-* - model_files <- grep("^model-", basename(safetensor_files), value = TRUE) - if (length(model_files) == 0) { - model_files <- grep("diffusion_pytorch_model", safetensor_files, value = TRUE) - } else { - model_files <- file.path(model_path, model_files) - safetensor_files <- model_files - } - - if (length(safetensor_files) == 0) { - stop("No safetensor files found in: ", model_path) - } - - if (verbose) { - message(sprintf("Loading weights from %d safetensor files...", length(safetensor_files))) - } - - # Load and apply weights - total_loaded <- 0L - total_skipped <- 0L - - for (sf_path in safetensor_files) { - if (verbose) message(" Loading: ", basename(sf_path)) - weights <- safetensors::safe_load_file(sf_path, framework = "torch") - - result <- load_gemma3_weights(model, weights, verbose = FALSE) - total_loaded <- total_loaded + result$loaded - total_skipped <- total_skipped + result$skipped - } - - if (verbose) { - message(sprintf("Loaded %d parameters, skipped %d", total_loaded, total_skipped)) - } - - # Move to device with dtype - torch_dtype <- switch(dtype, - "float32" = torch::torch_float32(), - "float16" = torch::torch_float16(), - "bfloat16" = torch::torch_bfloat16(), - torch::torch_float32() - ) - - model$to(device = device, dtype = torch_dtype) - - if (verbose) message("Gemma3 text encoder loaded on device: ", device) - - model +load_gemma3_text_encoder <- function(model_path, device = "cpu", + dtype = "float16", verbose = TRUE) { + # Load config + config_path <- file.path(model_path, "config.json") + if (!file.exists(config_path)) { + stop("Config file not found: ", config_path) + } + + config_raw <- jsonlite::fromJSON(config_path) + + # Extract text config if this is a multimodal model + if (!is.null(config_raw$text_config)) { + config_raw <- config_raw$text_config + } + + config <- list( + vocab_size = config_raw$vocab_size %||% 262208L, + hidden_size = config_raw$hidden_size %||% 3840L, + intermediate_size = config_raw$intermediate_size %||% 15360L, + num_hidden_layers = config_raw$num_hidden_layers %||% 48L, + num_attention_heads = config_raw$num_attention_heads %||% 16L, + num_key_value_heads = config_raw$num_key_value_heads %||% 8L, + head_dim = config_raw$head_dim %||% 256L, + max_position_embeddings = config_raw$max_position_embeddings %||% 131072L, + rms_norm_eps = config_raw$rms_norm_eps %||% 1e-6, + rope_theta = config_raw$rope_theta %||% 1000000.0, # Gemma3 uses 1M + rope_scaling = list(factor = config_raw$rope_scaling$factor %||% 8.0), + sliding_window = config_raw$sliding_window %||% 1024L, + sliding_window_pattern = config_raw$sliding_window_pattern %||% 6L, + attn_logit_softcapping = config_raw$attn_logit_softcapping, # NULL = no softcapping (HF default) + query_pre_attn_scalar = config_raw$query_pre_attn_scalar %||% 256L + ) + + if (verbose) { + message(sprintf("Creating Gemma3 model: %d layers, hidden_size=%d", + config$num_hidden_layers, config$hidden_size)) + } + + # Create model + model <- gemma3_text_model(config) + + # Find safetensor files + safetensor_files <- list.files(model_path, pattern = "\\.safetensors$", full.names = TRUE) + # Prefer model-* files over diffusion_pytorch_model-* + model_files <- grep("^model-", basename(safetensor_files), value = TRUE) + if (length(model_files) == 0) { + model_files <- grep("diffusion_pytorch_model", safetensor_files, value = TRUE) + } else { + model_files <- file.path(model_path, model_files) + safetensor_files <- model_files + } + + if (length(safetensor_files) == 0) { + stop("No safetensor files found in: ", model_path) + } + + if (verbose) { + message(sprintf("Loading weights from %d safetensor files...", length(safetensor_files))) + } + + # Load and apply weights + total_loaded <- 0L + total_skipped <- 0L + + for (sf_path in safetensor_files) { + if (verbose) { + message(" Loading: ", basename(sf_path)) + } + weights <- safetensors::safe_load_file(sf_path, framework = "torch") + + result <- load_gemma3_weights(model, weights, verbose = FALSE) + total_loaded <- total_loaded + result$loaded + total_skipped <- total_skipped + result$skipped + } + + if (verbose) { + message(sprintf("Loaded %d parameters, skipped %d", total_loaded, total_skipped)) + } + + # Move to device with dtype + torch_dtype <- switch(dtype, + "float32" = torch::torch_float32(), + "float16" = torch::torch_float16(), + "bfloat16" = torch::torch_bfloat16(), + torch::torch_float32() + ) + + model$to(device = device, dtype = torch_dtype) + + if (verbose) { + message("Gemma3 text encoder loaded on device: ", device) + } + + model } #' Load weights into Gemma3 model #' @keywords internal -load_gemma3_weights <- function( - model, - weights, - verbose = TRUE -) { - native_params <- names(model$parameters) - - # Remap HuggingFace keys to our module structure - remap_gemma3_key <- function(key) { - # Remove 'language_model.' prefix if present (for multimodal models) - key <- sub("^language_model\\.", "", key) - - # Map HuggingFace structure to our structure: - # model.embed_tokens -> embed_tokens - # model.layers.0.self_attn.q_proj -> layers.1.self_attn.q_proj (1-indexed) - # model.norm -> norm - - key <- sub("^model\\.", "", key) - - # Note: R torch model$parameters uses 0-indexed layer names even though - # R list indexing is 1-based. So we keep the layer numbers as-is. - - key - } - - loaded <- 0L - skipped <- 0L - - torch::with_no_grad({ - for (hf_name in names(weights)) { - native_name <- remap_gemma3_key(hf_name) - - if (native_name %in% native_params) { - hf_tensor <- weights[[hf_name]] - native_tensor <- model$parameters[[native_name]] - - if (all(as.integer(hf_tensor$shape) == as.integer(native_tensor$shape))) { - native_tensor$copy_(hf_tensor) - loaded <- loaded + 1L - } else { - if (verbose) { - message("Shape mismatch: ", native_name, - " (HF: ", paste(as.integer(hf_tensor$shape), collapse = "x"), - " vs R: ", paste(as.integer(native_tensor$shape), collapse = "x"), ")") +load_gemma3_weights <- function(model, weights, verbose = TRUE) { + native_params <- names(model$parameters) + + # Remap HuggingFace keys to our module structure + remap_gemma3_key <- function(key) { + # Remove 'language_model.' prefix if present (for multimodal models) + key <- sub("^language_model\\.", "", key) + + # Map HuggingFace structure to our structure: + # model.embed_tokens -> embed_tokens + # model.layers.0.self_attn.q_proj -> layers.1.self_attn.q_proj (1-indexed) + # model.norm -> norm + + key <- sub("^model\\.", "", key) + + # Note: R torch model$parameters uses 0-indexed layer names even though + # R list indexing is 1-based. So we keep the layer numbers as-is. + + key + } + + loaded <- 0L + skipped <- 0L + + torch::with_no_grad({ + for (hf_name in names(weights)) { + native_name <- remap_gemma3_key(hf_name) + + if (native_name %in% native_params) { + hf_tensor <- weights[[hf_name]] + native_tensor <- model$parameters[[native_name]] + + if (all(as.integer(hf_tensor$shape) == as.integer(native_tensor$shape))) { + native_tensor$copy_(hf_tensor) + loaded <- loaded + 1L + } else { + if (verbose) { + message("Shape mismatch: ", native_name, + " (HF: ", paste(as.integer(hf_tensor$shape), + collapse = "x"), + " vs R: ", paste(as.integer(native_tensor$shape), collapse = "x"), ")") + } + skipped <- skipped + 1L + } + } else { + skipped <- skipped + 1L } - skipped <- skipped + 1L - } - } else { - skipped <- skipped + 1L } - } }) - if (verbose) { - message(sprintf("Gemma3 weights: %d loaded, %d skipped", loaded, skipped)) - } + if (verbose) { + message(sprintf("Gemma3 weights: %d loaded, %d skipped", loaded, skipped)) + } - invisible(list(loaded = loaded, skipped = skipped)) + invisible(list(loaded = loaded, skipped = skipped)) } # ----------------------------------------------------------------------------- @@ -742,17 +729,17 @@ load_gemma3_weights <- function( #' @return A gemma3_tokenizer object (extends bpe_tokenizer). #' @export gemma3_tokenizer <- function(tokenizer_path) { - # Load base BPE tokenizer - tokenizer <- bpe_tokenizer(tokenizer_path) + # Load base BPE tokenizer + tokenizer <- bpe_tokenizer(tokenizer_path) - # Add Gemma-specific configuration - tokenizer$padding_side <- "left"# Gemma uses left padding + # Add Gemma-specific configuration + tokenizer$padding_side <- "left" # Gemma uses left padding - # Update class + # Update class - class(tokenizer) <- c("gemma3_tokenizer", class(tokenizer)) + class(tokenizer) <- c("gemma3_tokenizer", class(tokenizer)) - tokenizer + tokenizer } #' Tokenize text for Gemma3 @@ -764,29 +751,20 @@ gemma3_tokenizer <- function(tokenizer_path) { #' @param return_tensors Character. Return type ("list" or "pt" for torch tensors). #' @return List with input_ids and attention_mask. #' @export -tokenize_gemma3 <- function( - tokenizer, - text, - max_length = 1024L, - padding = "max_length", - return_tensors = "pt" -) { - if (!inherits(tokenizer, "gemma3_tokenizer") && !inherits(tokenizer, "bpe_tokenizer")) { - stop("tokenizer must be a gemma3_tokenizer or bpe_tokenizer object") - } - - # Use native BPE encoding - result <- encode_bpe( - tokenizer = tokenizer, - text = text, - add_special_tokens = TRUE, - max_length = max_length, - padding = padding, - truncation = TRUE, - return_tensors = return_tensors - ) - - result +tokenize_gemma3 <- function(tokenizer, text, max_length = 1024L, + padding = "max_length", return_tensors = "pt") { + if (!inherits(tokenizer, "gemma3_tokenizer") && + !inherits(tokenizer, "bpe_tokenizer")) { + stop("tokenizer must be a gemma3_tokenizer or bpe_tokenizer object") + } + + # Use native BPE encoding + result <- encode_bpe(tokenizer = tokenizer, text = text, + add_special_tokens = TRUE, max_length = max_length, + padding = padding, truncation = TRUE, + return_tensors = return_tensors) + + result } # ----------------------------------------------------------------------------- @@ -796,94 +774,81 @@ tokenize_gemma3 <- function( #' Encode text with Gemma3 for LTX-2 #' #' Full pipeline for encoding text prompts using Gemma3 text encoder. -#' Returns packed embeddings ready for LTX-2 connectors. +#' Returns the raw stacked per-layer hidden states (embedding layer plus all +#' transformer layers) for downstream connector modules, which handle +#' normalization and projection themselves. #' #' @param prompts Character vector of prompts. #' @param model Gemma3 text model (or path to load from). #' @param tokenizer Gemma3 tokenizer (or path to load from). #' @param max_sequence_length Integer. Maximum sequence length. -#' @param scale_factor Numeric. Scale factor for packing (default 8). #' @param device Character. Device for computation. #' @param dtype Character. Data type. #' @param verbose Logical. Print progress. -#' @return List with prompt_embeds and prompt_attention_mask. +#' @return List with prompt_embeds (raw stacked hidden states, +#' shape \code{[batch, seq_len, hidden_size, num_layers + 1]}) and +#' prompt_attention_mask. #' @export -encode_with_gemma3 <- function( - prompts, - model = NULL, - tokenizer = NULL, - max_sequence_length = 1024L, - scale_factor = 8, - device = "cuda", - dtype = "float16", - verbose = TRUE -) { - # Load model if path provided - if (is.character(model)) { - model <- load_gemma3_text_encoder(model, device = device, dtype = dtype, verbose = verbose) - } - - # Load tokenizer if path provided - if (is.character(tokenizer)) { - tokenizer <- gemma3_tokenizer(tokenizer) - } - - if (is.null(model) || is.null(tokenizer)) { - stop("Both model and tokenizer are required") - } - - # Ensure prompts is a list - if (is.character(prompts)) { - prompts <- as.list(prompts) - } - - # Tokenize (left-padding to max_length for Gemma3) - if (verbose) message("Tokenizing prompts...") - tokens <- tokenize_gemma3(tokenizer, unlist(prompts), - max_length = max_sequence_length, - padding = "max_length") - - input_ids <- tokens$input_ids$to(device = device) - attention_mask <- tokens$attention_mask$to(device = device) - - # Run through model - if (verbose) message("Encoding with Gemma3...") - torch::with_no_grad({ - output <- model(input_ids, attention_mask = attention_mask, output_hidden_states = TRUE) +encode_with_gemma3 <- function(prompts, model = NULL, tokenizer = NULL, + max_sequence_length = 1024L, device = "cuda", + dtype = "float16", verbose = TRUE) { + # Load model if path provided + if (is.character(model)) { + model <- load_gemma3_text_encoder(model, device = device, + dtype = dtype, verbose = verbose) + } + + # Load tokenizer if path provided + if (is.character(tokenizer)) { + tokenizer <- gemma3_tokenizer(tokenizer) + } + + if (is.null(model) || is.null(tokenizer)) { + stop("Both model and tokenizer are required") + } + + # Ensure prompts is a list + if (is.character(prompts)) { + prompts <- as.list(prompts) + } + + # Tokenize (left-padding to max_length for Gemma3) + if (verbose) { + message("Tokenizing prompts...") + } + tokens <- tokenize_gemma3(tokenizer, unlist(prompts), + max_length = max_sequence_length, + padding = "max_length") + + input_ids <- tokens$input_ids$to(device = device) + attention_mask <- tokens$attention_mask$to(device = device) + + # Run through model + if (verbose) { + message("Encoding with Gemma3...") + } + torch::with_no_grad({ + output <- model(input_ids, attention_mask = attention_mask, output_hidden_states = TRUE) }) - # Stack hidden states from transformer layers (skip embedding layer at index 1) - hidden_states_list <- output$hidden_states - # hidden_states_list[[1]] is embedding layer, [[2]] onwards are transformer layers - # LTX-2 connectors expect 49 layers (text_proj_in_factor=49) - transformer_hidden_states <- hidden_states_list[2:length(hidden_states_list)] - # Stack: [batch, seq_len, hidden_size, num_layers] - hidden_states_stacked <- torch::torch_stack(transformer_hidden_states, dim = - 1L) - - # Compute sequence lengths from attention mask - sequence_lengths <- as.integer(attention_mask$sum(dim = 2L)$cpu()) - - # Pack embeddings - if (verbose) message("Packing embeddings...") - prompt_embeds <- pack_text_embeds( - hidden_states_stacked, - sequence_lengths = sequence_lengths, - padding_side = "left", - scale_factor = scale_factor, - device = device - ) - - list( - prompt_embeds = prompt_embeds, - prompt_attention_mask = attention_mask - ) + # Stack ALL hidden states: embedding layer + every transformer layer. + # The LTX connectors consume the full stack (their projection expects + # hidden_size * (num_layers + 1) inputs) and normalize it themselves. + hidden_states_list <- output$hidden_states + # Stack: [batch, seq_len, hidden_size, num_layers + 1] + hidden_states_stacked <- torch::torch_stack(hidden_states_list, dim = -1L) + + list( + prompt_embeds = hidden_states_stacked, + prompt_attention_mask = attention_mask + ) } # Null-coalescing operator if (!exists("%||%", mode = "function")) { - `%||%` <- function( - x, - y - ) if (is.null(x)) y else x + `%||%` <- function(x, y) if (is.null(x)) { + y + } else { + x + } } - diff --git a/R/gpu_poor.R b/R/gpu_poor.R deleted file mode 100644 index 39d7db0..0000000 --- a/R/gpu_poor.R +++ /dev/null @@ -1,1785 +0,0 @@ -#' GPU-Poor Memory Management for LTX-2 -#' -#' wan2GP-style memory optimizations for running LTX-2 video generation -#' on limited VRAM (6-16GB). -#' -#' @name gpu_poor -NULL - -#' Get LTX-2 Memory Profile -#' -#' Determines optimal memory configuration based on available VRAM. -#' -#' @param vram_gb Numeric. Available VRAM in GB, or NULL for auto-detection. -#' @param model Character. Model variant: "ltx2-19b-fp4" (default), "ltx2-19b-fp8", -#' or "ltx2-19b-distilled". -#' -#' @return A list with memory profile settings. -#' -#' @details -#' LTX-2 is a 19B parameter model. Even with FP4 quantization (~10GB weights), -#' it requires careful memory management. The GPU-poor approach: -#' -#' 1. Text encoding runs on CPU (cached) -#' 2. DiT loaded in chunks, processed layer-by-layer, unloaded -#' 3. VAE loaded after DiT unload, decode with tiling, unload -#' -#' Memory profiles: -#' \describe{ -#' \item{high}{16GB+ - FP4 DiT with chunk loading, VAE on GPU} -#' \item{medium}{12GB - FP4 DiT chunk loading, VAE tiled} -#' \item{low}{8GB - FP4 DiT layer-by-layer, VAE tiled small} -#' \item{very_low}{6GB - FP4 layer-by-layer, VAE on CPU} -#' \item{cpu_only}{All on CPU} -#' } -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Auto-detect profile -#' profile <- ltx2_memory_profile() -#' -#' # Specific VRAM -#' profile <- ltx2_memory_profile(vram_gb = 8) -#' } -ltx2_memory_profile <- function( - vram_gb = NULL, - model = "ltx2-19b-fp4" -) { - # Auto-detect free VRAM if not provided - if (is.null(vram_gb)) { - vram_gb <- .detect_vram(use_free = TRUE) - message(sprintf("Detected %.1f GB free VRAM", vram_gb)) - } - - # Determine profile level - if (vram_gb >= 16) { - profile <- "high" - } else if (vram_gb >= 12) { - profile <- "medium" - } else if (vram_gb >= 8) { - profile <- "low" - } else if (vram_gb >= 6) { - profile <- "very_low" - } else { - profile <- "cpu_only" - } - - # Build profile config - # Note: LTX-2 19B has 48 transformer layers - # At FP4, ~10GB total model weights - # Layer chunk size determines how many layers loaded at once - - profiles <- list( - high = list( - name = "high", - # Stage 1: Text encoding (always CPU) - text_device = "cpu", - text_backend = "native", # Native Gemma3 encoder - # Stage 2: DiT denoising - dit_device = "cuda", - dit_offload = "chunk", # Load layers in chunks - dit_chunk_size = 12L, # 12 layers at a time (~2.5GB) - # Stage 3: VAE decode - vae_device = "cuda", - vae_tiling = FALSE, - vae_tile_size = c(512L, 512L), - vae_tile_frames = 16L, - # General settings - dtype = "float16", - model_precision = "fp4", # Preferred quantization - max_resolution = c(720L, 1280L), # height, width - max_frames = 121L, - cfg_mode = "batched"# Distilled uses CFG=1, so this is moot - ), - medium = list( - name = "medium", - text_device = "cpu", - text_backend = "native", - dit_device = "cuda", - dit_offload = "chunk", - dit_chunk_size = 8L, # 8 layers at a time (~1.7GB) - vae_device = "cuda", - vae_tiling = TRUE, - vae_tile_size = c(512L, 512L), - vae_tile_frames = 16L, - dtype = "float16", - model_precision = "fp4", - max_resolution = c(720L, 1280L), - max_frames = 121L, - cfg_mode = "batched" - ), - low = list( - name = "low", - text_device = "cpu", - text_backend = "native", - dit_device = "cuda", - dit_offload = "layer", # One layer at a time - dit_chunk_size = 1L, - vae_device = "cuda", - vae_tiling = TRUE, - vae_tile_size = c(256L, 256L), - vae_tile_frames = 8L, - dtype = "float16", - model_precision = "fp4", - max_resolution = c(480L, 854L), - max_frames = 61L, - cfg_mode = "sequential" - ), - very_low = list( - name = "very_low", - text_device = "cpu", - text_backend = "native", - dit_device = "cuda", - dit_offload = "layer", - dit_chunk_size = 1L, - vae_device = "cpu", # VAE on CPU - vae_tiling = TRUE, - vae_tile_size = c(128L, 128L), - vae_tile_frames = 4L, - dtype = "float16", - model_precision = "fp4", - max_resolution = c(480L, 640L), - max_frames = 33L, - cfg_mode = "sequential" - ), - cpu_only = list( - name = "cpu_only", - text_device = "cpu", - text_backend = "native", - dit_device = "cpu", - dit_offload = "none", - dit_chunk_size = 48L, # All layers (CPU has more RAM) - vae_device = "cpu", - vae_tiling = TRUE, - vae_tile_size = c(256L, 256L), - vae_tile_frames = 8L, - dtype = "float32", # CPU often faster with float32 - model_precision = "fp4", - max_resolution = c(480L, 640L), - max_frames = 33L, - cfg_mode = "sequential" - ) - ) - - profiles[[profile]] -} - -#' Get SDXL Memory Profile -#' -#' Determines optimal memory configuration for SDXL image generation -#' based on available VRAM. -#' -#' @param vram_gb Numeric. Available VRAM in GB, or NULL for auto-detection. -#' -#' @return A list with memory profile settings. -#' -#' @details -#' Memory profiles for SDXL: -#' \describe{ -#' \item{full_gpu}{16GB+ - All components on CUDA} -#' \item{balanced}{10-12GB - UNet + decoder on CUDA, text encoders on CPU} -#' \item{unet_gpu}{6-10GB - Only UNet on CUDA, everything else CPU} -#' \item{cpu_only}{<6GB - All on CPU} -#' } -#' -#' Each profile also specifies: -#' - cfg_mode: "batched" or "sequential" (sequential halves peak memory) -#' - cleanup: "none", "phase", or "step" (when to clear VRAM) -#' - dtype: "float16" or "float32" -#' - max_resolution: maximum image dimension -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Auto-detect profile -#' profile <- sdxl_memory_profile() -#' -#' # Specific VRAM -#' profile <- sdxl_memory_profile(vram_gb = 8) -#' } -sdxl_memory_profile <- function(vram_gb = NULL) { - # Auto-detect free VRAM if not provided - if (is.null(vram_gb)) { - vram_gb <- .detect_vram(use_free = TRUE) - message(sprintf("Detected %.1f GB free VRAM", vram_gb)) - } - - # Determine profile level - if (vram_gb >= 16) { - profile <- "full_gpu" - } else if (vram_gb >= 10) { - profile <- "balanced" - } else if (vram_gb >= 6) { - profile <- "unet_gpu" - } else { - profile <- "cpu_only" - } - - # Build profile config - profiles <- list( - full_gpu = list( - name = "full_gpu", - devices = list( - unet = "cuda", - decoder = "cuda", - text_encoder = "cuda", - text_encoder2 = "cuda", - encoder = "cuda" - ), - dtype = "float16", - cfg_mode = "batched", - cleanup = "none", - max_resolution = 1536L, - step_cleanup_interval = 0L# No step cleanup - ), - balanced = list( - name = "balanced", - devices = list( - unet = "cuda", - decoder = "cuda", - text_encoder = "cpu", - text_encoder2 = "cpu", - encoder = "cpu" - ), - dtype = "float16", - cfg_mode = "batched", - cleanup = "phase", # Cleanup between text encoding and denoising - max_resolution = 1024L, - step_cleanup_interval = 0L - ), - unet_gpu = list( - name = "unet_gpu", - devices = list( - unet = "cuda", - decoder = "cpu", - text_encoder = "cpu", - text_encoder2 = "cpu", - encoder = "cpu" - ), - dtype = "float16", - cfg_mode = "sequential", # Sequential CFG halves peak memory - cleanup = "phase", - max_resolution = 1024L, - step_cleanup_interval = 10L# Cleanup every 10 steps - ), - cpu_only = list( - name = "cpu_only", - devices = list( - unet = "cpu", - decoder = "cpu", - text_encoder = "cpu", - text_encoder2 = "cpu", - encoder = "cpu" - ), - dtype = "float32", # CPU often faster with float32 - cfg_mode = "sequential", - cleanup = "none", # No GPU to clean - max_resolution = 768L, - step_cleanup_interval = 0L - ) - ) - - profiles[[profile]] -} - -#' Check if GPU is Blackwell Architecture -#' -#' Blackwell GPUs (RTX 50xx) may need special handling. -#' -#' @return Logical. TRUE if Blackwell GPU detected. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' if (is_blackwell_gpu()) { -#' message("Using Blackwell-compatible settings") -#' } -#' } -is_blackwell_gpu <- function() { - # Use gpuctl if available - if (requireNamespace("gpu.ctl", quietly = TRUE)) { - return(gpu.ctl::gpu_is_blackwell()) - } - - # Fallback: check compute capability via torch - if (torch::cuda_is_available()) { - props <- tryCatch( - torch::cuda_get_device_properties(0L), - error = function(e) NULL - ) - if (!is.null(props)) { - # Blackwell is compute 12.x - major <- props$major - return(major >= 12) - } - } - - FALSE -} - -#' Detect Available VRAM -#' -#' Uses gpuctl if available. -#' -#' @param use_free Logical. If TRUE, return free VRAM. If FALSE, return total. -#' -#' @return Numeric. VRAM in GB, or 0 if no GPU detected. -#' @keywords internal -.detect_vram <- function(use_free = FALSE) { - # Try gpuctl (preferred - uses nvidia-smi) - if (requireNamespace("gpu.ctl", quietly = TRUE)) { - info <- gpu.ctl::gpu_detect() - if (!is.null(info)) { - if (use_free && !is.null(info$vram_free_gb)) { - return(info$vram_free_gb) - } - if (!is.null(info$vram_total_gb)) { - return(info$vram_total_gb) - } - } - } - - # Fallback: check if CUDA available but can't determine VRAM - if (torch::cuda_is_available()) { - # Conservative estimate - assume 8GB if we can't detect - message("Could not detect VRAM. Install gpuctl for accurate detection.") - return(8) - } - - # No GPU detected - 0 -} - -#' Offload Module to CPU -#' -#' Moves a torch module and all its parameters to CPU. -#' -#' @param module A torch nn_module. -#' @param gc Logical. Run garbage collection after offload. -#' -#' @return The module (modified in place). -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' model$to(device = "cuda") -#' output <- model(x) -#' offload_to_cpu(model) -#' } -offload_to_cpu <- function( - module, - gc = TRUE -) { - module$to(device = "cpu") - if (gc && torch::cuda_is_available()) { - gc() - torch::cuda_empty_cache() - } - invisible(module) -} - -#' Load Module to GPU -#' -#' Moves a torch module and all its parameters to CUDA. -#' -#' @param module A torch nn_module. -#' @param device Character. Target device (default "cuda"). -#' -#' @return The module (modified in place). -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' load_to_gpu(model) -#' output <- model(x) -#' offload_to_cpu(model) -#' } -load_to_gpu <- function( - module, - device = "cuda" -) { - module$to(device = device) - invisible(module) -} - -#' Report VRAM Usage -#' -#' Prints current VRAM usage using gpuctl. -#' -#' @param label Character. Label for the report. -#' -#' @return Invisibly returns a list with used and free VRAM in GB. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' vram_report("After model load") -#' } -vram_report <- function(label = "") { - if (!torch::cuda_is_available()) { - message("[", label, "] No CUDA available") - return(invisible(list(used = 0, free = 0))) - } - - # Use gpuctl for accurate reporting - if (requireNamespace("gpu.ctl", quietly = TRUE)) { - info <- gpu.ctl::gpu_detect() - if (!is.null(info)) { - used <- info$vram_used_gb - free <- info$vram_free_gb - message(sprintf("[%s] VRAM: %.2f GB used, %.2f GB free", - label, used, free)) - return(invisible(list(used = used, free = free))) - } - } - - message("[", label, "] VRAM: (install gpuctl for detailed stats)") - invisible(list(used = NA, free = NA)) -} - -#' Clear VRAM Cache -#' -#' Forces garbage collection and clears CUDA memory cache. -#' -#' @param verbose Logical. Print memory status before/after. -#' -#' @return Invisibly returns NULL. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' clear_vram() -#' } -clear_vram <- function(verbose = FALSE) { - if (!torch::cuda_is_available()) { - return(invisible(NULL)) - } - - if (verbose) { - vram_report("Before clear") - } - - gc() - tryCatch( - torch::cuda_empty_cache(), - error = function(e) NULL - ) - - if (verbose) { - vram_report("After clear") - } - - invisible(NULL) -} - -#' DiT Chunk-based Forward Pass -#' -#' Runs transformer layers in chunks, moving each chunk to GPU before -#' computation and back to CPU after. Balances memory usage with speed. -#' -#' @param hidden_states Input tensor. -#' @param layers List of transformer layers (on CPU). -#' @param chunk_size Integer. Number of layers to load at once (default 1). -#' @param device Target device for computation. -#' @param verbose Logical. Print progress. -#' @param ... Additional arguments passed to each layer. -#' -#' @return Output tensor (on CPU). -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Layer-by-layer for 8GB VRAM -#' output <- dit_offloaded_forward( -#' hidden_states, -#' model$transformer_blocks, -#' chunk_size = 1, -#' device = "cuda" -#' ) -#' -#' # Chunk-based for 16GB VRAM -#' output <- dit_offloaded_forward( -#' hidden_states, -#' model$transformer_blocks, -#' chunk_size = 12, # 12 layers at a time -#' device = "cuda" -#' ) -#' } -dit_offloaded_forward <- function( - hidden_states, - layers, - chunk_size = 1L, - device = "cuda", - verbose = FALSE, - ... -) { - n_layers <- length(layers) - chunk_size <- as.integer(chunk_size) - - # Move input to target device - x <- hidden_states$to(device = device) - - # Process in chunks - chunk_start <- 1L - while (chunk_start <= n_layers) { - chunk_end <- min(chunk_start + chunk_size - 1L, n_layers) - - if (verbose) { - message(sprintf(" Processing layers %d-%d of %d", chunk_start, chunk_end, n_layers)) - } - - # Load chunk to GPU - for (i in chunk_start:chunk_end) { - layers[[i]]$to(device = device) - } - - # Forward pass through chunk - for (i in chunk_start:chunk_end) { - x <- layers[[i]](x, ...) - } - - # Offload chunk back to CPU - for (i in chunk_start:chunk_end) { - layers[[i]]$to(device = "cpu") - } - - # Clear cache after each chunk - if (device != "cpu") { - torch::cuda_empty_cache() - } - - chunk_start <- chunk_end + 1L - } - - # Return result on CPU - x$to(device = "cpu") -} - -#' Sequential CFG Forward Pass -#' -#' Runs unconditional and conditional forward passes separately to halve -#' peak activation memory. For GPU-poor scenarios. -#' -#' @param model The DiT model. -#' @param latents Current latent tensor. -#' @param timestep Current timestep tensor. -#' @param prompt_embeds Conditional prompt embeddings. -#' @param negative_prompt_embeds Unconditional prompt embeddings. -#' @param guidance_scale CFG scale. -#' @param ... Additional arguments to model forward pass. -#' -#' @return The CFG-combined noise prediction. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' noise_pred <- sequential_cfg_forward( -#' model, latents, timestep, -#' prompt_embeds, negative_prompt_embeds, -#' guidance_scale = 4.0 -#' ) -#' } -sequential_cfg_forward <- function( - model, - latents, - timestep, - prompt_embeds, - negative_prompt_embeds, - guidance_scale, - ... -) { - torch::with_no_grad({ - # Unconditional pass - noise_pred_uncond <- model( - hidden_states = latents, - encoder_hidden_states = negative_prompt_embeds, - timestep = timestep, - ... - )$sample - - # Conditional pass - noise_pred_cond <- model( - hidden_states = latents, - encoder_hidden_states = prompt_embeds, - timestep = timestep, - ... - )$sample - - # CFG combination - noise_pred <- noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) - - # Clean up intermediate tensors - rm(noise_pred_uncond, noise_pred_cond) - }) - - noise_pred -} - -#' Validate Resolution Against Profile -#' -#' Checks if requested resolution fits within memory profile limits. -#' -#' @param height Integer. Requested height. -#' @param width Integer. Requested width. -#' @param num_frames Integer. Requested number of frames. -#' @param profile Memory profile from `ltx2_memory_profile()`. -#' -#' @return List with adjusted height, width, num_frames and warning if adjusted. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' profile <- ltx2_memory_profile(vram_gb = 8) -#' validated <- validate_resolution(720, 1280, 60, profile) -#' } -validate_resolution <- function( - height, - width, - num_frames, - profile -) { - adjusted <- FALSE - warnings <- character(0) - - max_h <- profile$max_resolution[1] - max_w <- profile$max_resolution[2] - max_f <- profile$max_frames - - if (height > max_h) { - warnings <- c(warnings, sprintf("Height %d exceeds profile max %d", height, max_h)) - height <- max_h - adjusted <- TRUE - } - - if (width > max_w) { - warnings <- c(warnings, sprintf("Width %d exceeds profile max %d", width, max_w)) - width <- max_w - adjusted <- TRUE - } - - if (num_frames > max_f) { - warnings <- c(warnings, sprintf("Frames %d exceeds profile max %d", num_frames, max_f)) - num_frames <- max_f - adjusted <- TRUE - } - - if (adjusted && length(warnings) > 0) { - warning("Resolution adjusted for memory profile '", profile$name, "':\n ", - paste(warnings, collapse = "\n ")) - } - - list( - height = height, - width = width, - num_frames = num_frames, - adjusted = adjusted - ) -} - -#' Configure VAE for Memory Profile -#' -#' Sets VAE tiling parameters based on memory profile. -#' -#' @param vae The LTX2 VAE module. -#' @param profile Memory profile from `ltx2_memory_profile()`. -#' -#' @return The VAE (modified in place). -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' profile <- ltx2_memory_profile(vram_gb = 8) -#' vae <- load_ltx2_vae(...) -#' configure_vae_for_profile(vae, profile) -#' } -configure_vae_for_profile <- function( - vae, - profile -) { - if (profile$vae_tiling) { - vae$enable_tiling( - tile_sample_min_height = profile$vae_tile_size[1], - tile_sample_min_width = profile$vae_tile_size[2], - tile_sample_min_num_frames = profile$vae_tile_frames - ) - } else { - vae$disable_tiling() - } - - invisible(vae) -} - -#' Quantize Tensor to INT4 -#' -#' Quantizes a float tensor to 4-bit integers with block-wise scaling. -#' Two INT4 values are packed per byte for 7-8x compression. -#' -#' @param x Tensor. Input float tensor. -#' @param block_size Integer. Number of values per scale factor (default 64). -#' -#' @return A list with: -#' - `packed`: uint8 tensor with packed INT4 values -#' - `scales`: float tensor with per-block scale factors -#' - `orig_shape`: original tensor shape -#' - `orig_numel`: original number of elements -#' - `block_size`: block size used -#' -#' @details -#' INT4 range is -8 to 7. Values are scaled per block, quantized, shifted to -#' unsigned (0-15), and packed two per byte. Compression is ~7x vs float32, -#' ~3.5x vs float16. Typical reconstruction error is 10-12% of std. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' x <- torch_randn(c(4096, 4096)) * 0.02 -#' q <- quantize_int4(x) -#' x_back <- dequantize_int4(q) -#' } -quantize_int4 <- function( - x, - block_size = 64L -) { - orig_shape <- x$shape - orig_dtype <- x$dtype - x_flat <- x$to(dtype = torch::torch_float32())$flatten() - n <- x_flat$shape[1] - - # Pad to multiple of block_size * 2 (2 values per byte) - pad_to <- ceiling(n / (block_size * 2)) * block_size * 2 - if (pad_to > n) { - x_flat <- torch::torch_cat(list( - x_flat, - torch::torch_zeros(pad_to - n, dtype = torch::torch_float32(), device = x$device) - )) - } - - # Reshape into blocks - n_blocks <- as.integer(pad_to / block_size) - x_blocks <- x_flat$reshape(c(n_blocks, block_size)) - - # Compute scale per block (absmax / 7) - scales <- x_blocks$abs()$max(dim = 2) [[1]] / 7.0 - scales <- scales$clamp(min = 1e-10) - - # Quantize: scale, round, clamp to -8..7, shift to 0..15 - x_scaled <- x_blocks / scales$unsqueeze(2) - x_int <- x_scaled$round()$clamp(- 8, 7) + 8L - x_uint <- x_int$to(dtype = torch::torch_uint8()) - - # Pack pairs into bytes (high nibble * 16 + low nibble) - x_uint <- x_uint$reshape(c(n_blocks, block_size %/% 2L, 2L)) - high <- x_uint[,, 1]$to(torch::torch_int32()) * 16L - low <- x_uint[,, 2]$to(torch::torch_int32()) - packed <- (high + low)$to(torch::torch_uint8()) - - list( - packed = packed$flatten(), - scales = scales, - orig_shape = orig_shape, - orig_numel = n, - block_size = block_size - ) -} - -#' Dequantize INT4 Tensor -#' -#' Reconstructs a float tensor from INT4-quantized data. -#' -#' @param q List. Quantized data from `quantize_int4()`. -#' @param dtype Torch dtype. Output dtype (default float16). -#' @param device Character. Target device. -#' -#' @return Tensor with original shape and specified dtype. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' q <- quantize_int4(weights) -#' weights_approx <- dequantize_int4(q, dtype = torch_float16(), device = "cuda") -#' } -dequantize_int4 <- function( - q, - dtype = torch::torch_float16(), - device = "cpu" -) { - packed <- q$packed$to(dtype = torch::torch_int32(), device = device) - - # Unpack bytes: high nibble = floor(x/16), low nibble = x mod 16 - high <- torch::torch_floor(packed$to(torch::torch_float32()) / 16)$to(torch::torch_int32()) - low <- packed - high * 16L - - # Interleave high and low nibbles - x_uint <- torch::torch_stack(list(high, low), dim = 2L)$flatten() - - # Shift back to signed (-8 to 7) - x_int <- x_uint - 8L - - # Apply per-block scales - block_size <- q$block_size - x_blocks <- x_int$reshape(c(- 1L, block_size))$to(dtype = dtype) - scales_dev <- q$scales$to(dtype = dtype, device = device) - x_scaled <- x_blocks * scales_dev$unsqueeze(2) - - # Flatten and trim to original size - x_flat <- x_scaled$flatten()[1:q$orig_numel] - x_flat$reshape(q$orig_shape) -} - -#' Create Linear Layer (Standard or INT4) -#' -#' Factory function that creates either a standard nn_linear or INT4 linear layer -#' based on package options. Use this instead of torch::nn_linear() in model code. -#' -#' @param in_features Integer. Input dimension. -#' @param out_features Integer. Output dimension. -#' @param bias Logical. Include bias term. -#' -#' @return nn_linear or int4_linear module. -#' -#' @details -#' Behavior controlled by options: -#' - `diffuseR.use_int4`: If TRUE, create INT4 layer (default FALSE) -#' - `diffuseR.int4_device`: Device for INT4 layers (default "cuda") -#' - `diffuseR.int4_dtype`: Dtype for INT4 operations (default torch_float16()) -#' -#' @export -make_linear <- function( - in_features, - out_features, - bias = TRUE -) { - - if (getOption("diffuseR.use_int4", FALSE)) { - device <- getOption("diffuseR.int4_device", "cuda") - dtype <- getOption("diffuseR.int4_dtype", torch::torch_float16()) - int4_linear(in_features, out_features, bias = bias, device = device, dtype = dtype) - } else { - torch::nn_linear(in_features, out_features, bias = bias) - } -} - -#' INT4 Linear Layer -#' -#' A linear layer that stores weights in INT4 format and dequantizes on-the-fly -#' during forward pass. This enables running large models on limited VRAM by -#' keeping weights compressed on GPU. -#' -#' @param in_features Integer. Size of each input sample. -#' @param out_features Integer. Size of each output sample. -#' @param bias Logical. If TRUE, adds a learnable bias (stored in float16). -#' @param device Character. Device for the layer. -#' @param dtype torch_dtype. Data type for dequantized operations. -#' -#' @return nn_module with INT4 weight storage. -#' -#' @details -#' The layer stores: -#' - `weight_packed`: uint8 tensor with packed INT4 values -#' - `weight_scales`: float32 tensor with per-block scales -#' - `weight_shape`: original weight shape -#' - `bias`: optional float16 bias -#' -#' During forward(), weights are dequantized to the target dtype, the matmul -#' is performed, and the dequantized tensor is freed. This allows ~40GB models -#' to run with ~10GB of VRAM for weights. -#' -#' @export -int4_linear <- torch::nn_module( - "INT4Linear", - initialize = function( - in_features, - out_features, - bias = TRUE, - device = "cuda", - dtype = torch::torch_float16() - ) { - self$in_features <- in_features - self$out_features <- out_features - self$dtype <- dtype - self$device <- device - self$block_size <- 64L - - # Placeholder - will be set by load_int4_weight() - self$weight_packed <- NULL - self$weight_scales <- NULL - self$weight_shape <- c(out_features, in_features) - self$weight_numel <- out_features * in_features - - if (bias) { - self$bias <- torch::nn_parameter(torch::torch_zeros(out_features, - dtype = dtype, device = device)) - } else { - self$bias <- NULL - } - }, - -#' Load INT4 quantized weight into this layer -#' @param q List with packed, scales, orig_shape from quantize_int4() - load_int4_weight = function(q) { - # Store INT4 data as buffers (not parameters) - self$weight_packed <- q$packed$to(device = self$device) - self$weight_scales <- q$scales$to(device = self$device) - self$weight_shape <- q$orig_shape - self$weight_numel <- q$orig_numel - invisible(self) - }, - - forward = function(x) { - if (is.null(self$weight_packed)) { - stop("INT4 weight not loaded. Call load_int4_weight() first.") - } - - # Dequantize weight on-the-fly - q <- list( - packed = self$weight_packed, - scales = self$weight_scales, - orig_shape = self$weight_shape, - orig_numel = self$weight_numel, - block_size = self$block_size - ) - weight <- dequantize_int4(q, dtype = self$dtype, device = self$device) - - # Linear operation - out <- torch::nnf_linear(x, weight, self$bias) - - # Weight tensor goes out of scope and will be freed - out - } -) - -#' Create INT4 Linear from Standard Linear -#' -#' Converts a standard nn_linear layer to an INT4 linear layer. -#' -#' @param linear_module nn_linear module to convert. -#' @param device Character. Target device for INT4 weights. -#' @param dtype torch_dtype. Target dtype for dequantized operations. -#' -#' @return int4_linear module with quantized weights. -#' -#' @export -linear_to_int4 <- function( - linear_module, - device = "cuda", - dtype = torch::torch_float16() -) { - in_features <- linear_module$in_features - out_features <- linear_module$out_features - has_bias <- !is.null(linear_module$bias) - - # Create INT4 layer - int4_layer <- int4_linear(in_features, out_features, bias = has_bias, - device = device, dtype = dtype) - - # Quantize and load weight - q <- quantize_int4(linear_module$weight) - int4_layer$load_int4_weight(q) - - # Copy bias if present (use with_no_grad to avoid in-place error on parameter) - if (has_bias) { - torch::with_no_grad({ - int4_layer$bias$copy_(linear_module$bias$to(dtype = dtype, device = device)) - }) - } - - int4_layer -} - -#' Create INT4 Linear from Pre-quantized Weights -#' -#' Creates an INT4 linear layer from pre-quantized weight data. -#' -#' @param q_weight List with packed, scales, orig_shape from load_int4_weights(). -#' @param q_bias Optional. Quantized bias (or NULL for no bias). -#' @param bias_tensor Optional. Float tensor for bias (if not quantized). -#' @param device Character. Target device. -#' @param dtype torch_dtype. Target dtype for operations. -#' -#' @return int4_linear module with loaded weights. -#' -#' @export -int4_linear_from_quantized <- function( - q_weight, - q_bias = NULL, - bias_tensor = NULL, - device = "cuda", - dtype = torch::torch_float16() -) { - out_features <- q_weight$orig_shape[1] - in_features <- q_weight$orig_shape[2] - has_bias <- !is.null(q_bias) || !is.null(bias_tensor) - - # Create INT4 layer - int4_layer <- int4_linear(in_features, out_features, bias = has_bias, - device = device, dtype = dtype) - - # Load quantized weight - int4_layer$load_int4_weight(q_weight) - - # Load bias if present (use with_no_grad to avoid in-place error on parameter) - if (!is.null(bias_tensor)) { - torch::with_no_grad({ - int4_layer$bias$copy_(bias_tensor$to(dtype = dtype, device = device)) - }) - } else if (!is.null(q_bias)) { - # Dequantize bias - bias_dequant <- dequantize_int4(q_bias, dtype = dtype, device = device) - torch::with_no_grad({ - int4_layer$bias$copy_(bias_dequant) - }) - } - - int4_layer -} - -#' Load INT4 Weights into Model -#' -#' Replaces linear layers in a model with INT4 versions and loads quantized weights. -#' This is the main entry point for running large models with INT4 quantization. -#' -#' @param model nn_module. The model to convert. -#' @param int4_weights List of quantized weights from `load_int4_weights()`. -#' @param device Character. Target device for INT4 weights. -#' @param dtype torch_dtype. Target dtype for dequantized operations. -#' @param verbose Logical. Print progress. -#' -#' @return The model with linear layers replaced by INT4 versions. -#' -#' @details -#' This function: -#' 1. Identifies linear layers by matching parameter names ending in ".weight" -#' 2. Creates INT4Linear layers with matching dimensions -#' 3. Loads quantized weights and biases -#' 4. Replaces the original layers in the model -#' -#' The INT4 weights stay compressed on GPU (~10GB for a 40GB model). -#' During forward(), each layer dequantizes on-the-fly, keeping memory usage low. -#' -#' @export -load_int4_into_model <- function( - model, - int4_weights, - device = "cuda", - dtype = torch::torch_float16(), - verbose = TRUE -) { - # Get all module names that have .weight in the quantized weights - weight_names <- grep("\\.weight$", names(int4_weights), value = TRUE) - - if (verbose) { - message(sprintf("Loading %d INT4 weights into model...", length(weight_names))) - } - - loaded <- 0 - skipped <- 0 - - for (weight_name in weight_names) { - # Extract module path (e.g., "transformer_blocks.0.attn1.to_q") - module_path <- sub("\\.weight$", "", weight_name) - bias_name <- paste0(module_path, ".bias") - - q_weight <- int4_weights[[weight_name]] - if (bias_name %in% names(int4_weights)) { - q_bias <- int4_weights[[bias_name]] - } else { - q_bias <- NULL - } - - # Check dimensions - only process 2D weights (linear layers) - if (length(q_weight$orig_shape) != 2) { - skipped <- skipped + 1 - next - } - - # Create INT4 layer - out_features <- q_weight$orig_shape[1] - in_features <- q_weight$orig_shape[2] - has_bias <- !is.null(q_bias) - - int4_layer <- int4_linear(in_features, out_features, bias = has_bias, - device = device, dtype = dtype) - int4_layer$load_int4_weight(q_weight) - - # Load bias if present - if (has_bias) { - bias_dequant <- dequantize_int4(q_bias, dtype = dtype, device = device) - torch::with_no_grad({ - int4_layer$bias$copy_(bias_dequant) - }) - } - - # Store the INT4 layer for later assignment - # Note: Direct module replacement in R torch is complex - # For now, store in a separate list that can be used during forward - if (!exists("int4_layers", where = model)) { - model$int4_layers <- list() - } - model$int4_layers[[module_path]] <- int4_layer - - loaded <- loaded + 1 - } - - if (verbose) { - message(sprintf("Loaded %d INT4 layers, skipped %d non-linear", loaded, skipped)) - } - - invisible(model) -} - -#' Load INT4 Weights into INT4 Model -#' -#' Loads pre-quantized INT4 weights into a model created with `make_linear()` -#' when `diffuseR.use_int4 = TRUE`. -#' -#' @param model nn_module created with INT4 layers. -#' @param int4_weights List from `load_int4_weights()`. -#' @param verbose Logical. Print progress. -#' -#' @return Model with INT4 weights loaded (invisibly). -#' -#' @export -load_int4_weights_into_model <- function( - model, - int4_weights, - verbose = TRUE -) { - # Get model's named modules (flattened) - params <- model$parameters - param_names <- names(params) - - loaded <- 0 - skipped <- 0 - - # Name mapping from HuggingFace to R model structure - # FFN layers have different naming: - # HF: ff.net.0.proj, ff.net.2 - # R: ff.act_fn.proj, ff.proj_out - map_hf_to_r_name <- function(hf_name) { - r_name <- hf_name - # Map FFN layer names - r_name <- gsub("\\.ff\\.net\\.0\\.proj\\.", ".ff.act_fn.proj.", r_name) - r_name <- gsub("\\.ff\\.net\\.2\\.", ".ff.proj_out.", r_name) - r_name <- gsub("\\.audio_ff\\.net\\.0\\.proj\\.", ".audio_ff.act_fn.proj.", r_name) - r_name <- gsub("\\.audio_ff\\.net\\.2\\.", ".audio_ff.proj_out.", r_name) - # Handle end-of-string cases - r_name <- gsub("\\.ff\\.net\\.0\\.proj$", ".ff.act_fn.proj", r_name) - r_name <- gsub("\\.ff\\.net\\.2$", ".ff.proj_out", r_name) - r_name <- gsub("\\.audio_ff\\.net\\.0\\.proj$", ".audio_ff.act_fn.proj", r_name) - r_name <- gsub("\\.audio_ff\\.net\\.2$", ".audio_ff.proj_out", r_name) - r_name - } - - for (int4_name in names(int4_weights)) { - # Map HuggingFace name to R model name - r_name <- map_hf_to_r_name(int4_name) - # Check if this weight (with mapped name) exists in model - if (r_name %in% param_names) { - # This is a regular parameter (bias, norm weights, etc.) - # Dequantize and copy - q <- int4_weights[[int4_name]] - if (length(q$orig_shape) == 1) { - # 1D tensor (bias, norm) - dequantize to model's device/dtype - param <- params[[r_name]] - dequant <- dequantize_int4(q, dtype = param$dtype, device = as.character(param$device)) - torch::with_no_grad({ - param$copy_(dequant) - }) - loaded <- loaded + 1 - } - } else if (grepl("\\.weight$", r_name)) { - # This might be an INT4 linear weight - find the layer - # Weight name: "module.path.weight" -> layer path: "module.path" - layer_path <- sub("\\.weight$", "", r_name) - - # Try to find corresponding INT4 layer by navigating module tree - tryCatch({ - # Navigate to the layer using R model path - parts <- strsplit(layer_path, "\\.") [[1]] - current <- model - - for (part in parts) { - if (grepl("^[0-9]+$", part)) { - # Numeric index (0-based in Python, 1-based in R) - idx <- as.integer(part) + 1L - current <- current[[idx]] - } else { - current <- current[[part]] - } - } - - # Check if this is an INT4Linear layer (has load_int4_weight method) - if (!is.null(current$load_int4_weight)) { - # Load INT4 weight directly (using original name for data access) - q <- int4_weights[[int4_name]] - current$load_int4_weight(q) - loaded <- loaded + 1 - } else { - skipped <- skipped + 1 - } - }, error = function(e) { - skipped <<- skipped + 1 - }) - } else { - skipped <- skipped + 1 - } - } - - if (verbose) { - message(sprintf("Loaded %d weights, skipped %d", loaded, skipped)) - } - - invisible(model) -} - -#' Quantize Model Weights to INT4 -#' -#' Quantizes all parameters in a torch module to INT4 format. -#' -#' @param module nn_module. The model to quantize. -#' @param block_size Integer. Block size for quantization. -#' @param verbose Logical. Print progress. -#' -#' @return List of quantized parameters (does not modify original module). -#' -#' @export -quantize_model_int4 <- function( - module, - block_size = 64L, - verbose = TRUE -) { - params <- module$parameters - quantized <- list() - - total_orig <- 0 - total_quant <- 0 - - for (name in names(params)) { - p <- params[[name]] - orig_bytes <- prod(p$shape) * 2# Assume float16 - - q <- quantize_int4(p, block_size = block_size) - quant_bytes <- q$packed$shape[1] + prod(q$scales$shape) * 4 - - quantized[[name]] <- q - total_orig <- total_orig + orig_bytes - total_quant <- total_quant + quant_bytes - - if (verbose) { - message(sprintf(" %s: %.2f MB -> %.2f MB", - name, orig_bytes / 1e6, quant_bytes / 1e6)) - } - } - - if (verbose) { - message(sprintf("Total: %.2f MB -> %.2f MB (%.1fx compression)", - total_orig / 1e6, total_quant / 1e6, total_orig / total_quant)) - } - - quantized -} - -#' Save INT4 Quantized Weights -#' -#' Saves INT4 quantized weights to disk as sharded safetensors files. -#' -#' @param quantized_params List of quantized parameters from `quantize_model_int4()`. -#' @param path Character. Base path for safetensors files. If multiple shards needed, -#' files will be named `path-00001-of-NNNNN.safetensors`. -#' @param max_shard_size Numeric. Maximum bytes per shard (default 2GB to avoid R -#' integer overflow issues). -#' @param verbose Logical. Print progress. -#' -#' @return Invisible character vector of saved file paths. -#' -#' @details -#' Weights are saved in safetensors format with the following structure: -#' \itemize{ -#' \item `{name}::packed` - uint8 tensor with packed INT4 values -#' \item `{name}::scales` - float32 tensor with per-block scales -#' \item `{name}::shape` - int64 tensor with original shape -#' } -#' -#' Large models are automatically sharded to avoid R's 2GB vector limit. -#' The block size is fixed at 64 (standard for INT4 quantization). -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' q <- quantize_model_int4(model) -#' save_int4_weights(q, "model_int4.safetensors") -#' } -save_int4_weights <- function( - quantized_params, - path, - max_shard_size = 2e9, - verbose = TRUE -) { - if (verbose) message(sprintf("Preparing %d parameters for safetensors...", length(quantized_params))) - - # Calculate total size and estimate number of shards - total_bytes <- 0 - param_sizes <- list() - for (name in names(quantized_params)) { - q <- quantized_params[[name]] - size <- prod(q$packed$shape) + prod(q$scales$shape) * 4 + length(q$orig_shape) * 8 - param_sizes[[name]] <- size - total_bytes <- total_bytes + size - } - - n_shards <- max(1L, ceiling(total_bytes / max_shard_size)) - - if (n_shards == 1) { - # Single file - use original path - tensors <- list() - for (name in names(quantized_params)) { - q <- quantized_params[[name]] - tensors[[paste0(name, "::packed")]] <- q$packed$cpu() - tensors[[paste0(name, "::scales")]] <- q$scales$cpu() - tensors[[paste0(name, "::shape")]] <- torch::torch_tensor(q$orig_shape, dtype = torch::torch_int64()) - } - if (verbose) message(sprintf("Saving %d tensors to %s...", length(tensors), path)) - safetensors::safe_save_file(tensors, path) - file_size <- file.info(path)$size / 1e6 - if (verbose) message(sprintf("Saved %.2f MB", file_size)) - return(invisible(path)) - } - - # Multiple shards - split params across files - if (verbose) message(sprintf("Sharding into %d files (max %.1f GB each)...", n_shards, max_shard_size / 1e9)) - - # Remove extension for base path - base_path <- sub("\\.safetensors$", "", path) - param_names <- names(quantized_params) - params_per_shard <- ceiling(length(param_names) / n_shards) - saved_paths <- character(0) - - for (shard_idx in seq_len(n_shards)) { - start_idx <- (shard_idx - 1) * params_per_shard + 1 - end_idx <- min(shard_idx * params_per_shard, length(param_names)) - - if (start_idx > length(param_names)) break - - shard_names <- param_names[start_idx:end_idx] - tensors <- list() - - for (name in shard_names) { - q <- quantized_params[[name]] - tensors[[paste0(name, "::packed")]] <- q$packed$cpu() - tensors[[paste0(name, "::scales")]] <- q$scales$cpu() - tensors[[paste0(name, "::shape")]] <- torch::torch_tensor(q$orig_shape, dtype = torch::torch_int64()) - } - - shard_path <- sprintf("%s-%05d-of-%05d.safetensors", base_path, shard_idx, n_shards) - if (verbose) message(sprintf(" [%d/%d] Saving %d params to %s...", - shard_idx, n_shards, length(shard_names), basename(shard_path))) - safetensors::safe_save_file(tensors, shard_path) - saved_paths <- c(saved_paths, shard_path) - } - - total_size <- sum(file.info(saved_paths)$size) / 1e6 - if (verbose) message(sprintf("Total: %.2f MB across %d shards", total_size, length(saved_paths))) - - invisible(saved_paths) -} - -#' Load INT4 Quantized Weights -#' -#' Loads INT4 quantized weights from safetensors file(s). -#' -#' @param path Character. Path to safetensors file or base path for sharded files. -#' For sharded files, pass the base path (e.g., "model_int4.safetensors") and -#' the function will find all shards matching the pattern. -#' @param verbose Logical. Print progress. -#' -#' @return List of quantized parameter structures ready for `dequantize_int4()`. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' q <- load_int4_weights("model_int4.safetensors") -#' # Dequantize specific parameter on GPU -#' weight <- dequantize_int4(q[["linear.weight"]], device = "cuda") -#' } -load_int4_weights <- function( - path, - verbose = TRUE -) { - path <- path.expand(path) - - # Check for sharded files - base_path <- sub("\\.safetensors$", "", path) - shard_pattern <- sprintf("%s-[0-9]+-of-[0-9]+\\.safetensors$", basename(base_path)) - shard_dir <- dirname(path) - shard_files <- list.files(shard_dir, pattern = shard_pattern, full.names = TRUE) - - if (length(shard_files) > 0) { - # Load sharded files - shard_files <- sort(shard_files) - if (verbose) { - total_size <- sum(file.info(shard_files)$size) / 1e6 - message(sprintf("Loading INT4 weights from %d shards (%.2f MB total)...", - length(shard_files), total_size)) - } - paths <- shard_files - } else if (file.exists(path)) { - # Single file - if (verbose) { - size_mb <- file.info(path)$size / 1e6 - message(sprintf("Loading INT4 weights from %s (%.2f MB)...", path, size_mb)) - } - paths <- path - } else { - stop("File not found: ", path) - } - - # Load all files - quantized <- list() - for (i in seq_along(paths)) { - p <- paths[i] - if (verbose && length(paths) > 1) { - message(sprintf(" [%d/%d] Loading %s...", i, length(paths), basename(p))) - } - - tensors <- safetensors::safe_load_file(p, framework = "torch") - - # Parse tensor names to reconstruct parameter structures - packed_names <- grep("::packed$", names(tensors), value = TRUE) - param_names <- sub("::packed$", "", packed_names) - - for (name in param_names) { - packed <- tensors[[paste0(name, "::packed")]] - scales <- tensors[[paste0(name, "::scales")]] - shape_tensor <- tensors[[paste0(name, "::shape")]] - orig_shape <- as.integer(as.array(shape_tensor)) - orig_numel <- prod(orig_shape) - - quantized[[name]] <- list( - packed = packed, - scales = scales, - orig_shape = orig_shape, - orig_numel = orig_numel, - block_size = 64L# Standard block size - ) - } - } - - if (verbose) message(sprintf("Done. Loaded %d parameters.", length(quantized))) - quantized -} - -#' Quantize Safetensor Weights to INT4 -#' -#' Loads weights from safetensors file(s) and quantizes to INT4. -#' Useful for quantizing large models without loading the full module. -#' -#' @param paths Character vector. Paths to safetensor files. -#' @param output_path Character. Path to save INT4 weights (.safetensors). -#' @param block_size Integer. Block size for quantization (default 64). -#' @param verbose Logical. Print progress. -#' -#' @return Invisible NULL. Writes quantized weights to output_path. -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Quantize sharded transformer weights -#' paths <- list.files("~/.cache/huggingface/.../transformer", -#' pattern = "safetensors$", full.names = TRUE) -#' quantize_safetensors_int4(paths, "dit_int4.safetensors") -#' } -quantize_safetensors_int4 <- function( - paths, - output_path, - block_size = 64L, - verbose = TRUE -) { - all_quantized <- list() - total_orig <- 0 - total_quant <- 0 - - for (path in paths) { - if (verbose) message(sprintf("Loading %s...", basename(path))) - - weights <- safetensors::safe_load_file(path, framework = "torch") - - for (name in names(weights)) { - w <- weights[[name]] - orig_bytes <- prod(w$shape) * 2# Assume float16 - - q <- quantize_int4(w, block_size = block_size) - quant_bytes <- length(as.array(q$packed)) + prod(q$scales$shape) * 4 - - all_quantized[[name]] <- q - total_orig <- total_orig + orig_bytes - total_quant <- total_quant + quant_bytes - - if (verbose && prod(w$shape) > 1e6) { - message(sprintf(" %s: %.2f MB -> %.2f MB", - name, orig_bytes / 1e6, quant_bytes / 1e6)) - } - } - - # Clear memory between shards - rm(weights) - gc() - } - - if (verbose) { - message(sprintf("\nTotal: %.2f GB -> %.2f GB (%.1fx compression)", - total_orig / 1e9, total_quant / 1e9, total_orig / total_quant)) - } - - save_int4_weights(all_quantized, output_path, verbose = verbose) -} - -#' Quantize LTX-2 Transformer to INT4 -#' -#' Downloads (if needed) and quantizes the LTX-2 19B transformer to INT4 format. -#' The quantized weights are cached for future use. -#' -#' @param model_id Character. HuggingFace model ID (default "Lightricks/LTX-2"). -#' @param output_dir Character. Directory to save quantized weights. -#' Default uses `tools::R_user_dir("diffuseR", "cache")`. -#' @param block_size Integer. Block size for INT4 quantization (default 64). -#' @param force Logical. Re-quantize even if cached file exists. -#' @param download Logical. If TRUE, download model from HuggingFace if not cached. -#' @param verbose Logical. Print progress. -#' -#' @return Character. Path to the quantized weights file. -#' -#' @details -#' LTX-2 is a 19B parameter model (~40GB in BF16). INT4 quantization reduces -#' this to ~5.7GB, fitting in 16GB VRAM with room for activations. -#' -#' The function: -#' 1. Uses hfhub to locate/download the model from HuggingFace -#' 2. Loads each safetensor shard -#' 3. Quantizes all weights to INT4 (block-wise, 64 values per scale) -#' 4. Saves as safetensors file -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Quantize and cache (first run takes ~10-20 minutes) -#' weights_path <- quantize_ltx2_transformer() -#' -#' # Load quantized weights -#' q <- load_int4_weights(weights_path) -#' -#' # Dequantize specific layer on GPU -#' layer_weight <- dequantize_int4(q[["transformer_blocks.0.attn1.to_q.weight"]], -#' device = "cuda", dtype = torch_float16()) -#' } -quantize_ltx2_transformer <- function( - model_id = "Lightricks/LTX-2", - output_dir = NULL, - block_size = 64L, - force = FALSE, - download = FALSE, - verbose = TRUE -) { - if (!requireNamespace("hfhub", quietly = TRUE)) { - stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") - } - - # Use R_user_dir for CRAN-compliant cache - - if (is.null(output_dir)) { - output_dir <- tools::R_user_dir("diffuseR", "cache") - } - - output_file <- file.path(output_dir, "ltx2_transformer_int4.safetensors") - - # Check if already cached - if (file.exists(output_file) && !force) { - if (verbose) { - size_gb <- file.info(output_file)$size / 1e9 - message(sprintf("Using cached INT4 weights: %s (%.2f GB)", output_file, size_gb)) - } - return(output_file) - } - - # Check if model is available locally via transformer/config.json - transformer_dir <- tryCatch({ - config_path <- hfhub::hub_download(model_id, "transformer/config.json", - local_files_only = TRUE) - dirname(config_path) - }, error = function(e) NULL) - - if (is.null(transformer_dir)) { - if (!download) { - stop("Model '", model_id, "' transformer not found in HuggingFace cache.\n", - "Run with download = TRUE to download, or use:\n", - " huggingface-cli download ", model_id) - } - - # Interactive consent before downloading - if (interactive()) { - ans <- utils::askYesNo( - paste0("Download '", model_id, "' transformer (~40GB) from HuggingFace?"), - default = TRUE - ) - if (!isTRUE(ans)) { - stop("Download cancelled.", call. = FALSE) - } - } - - if (verbose) message("Downloading transformer weights from HuggingFace...") - model_path <- hfhub::hub_snapshot(model_id, - allow_patterns = "transformer/*") - transformer_dir <- file.path(model_path, "transformer") - } - if (!dir.exists(transformer_dir)) { - stop("Transformer directory not found: ", transformer_dir) - } - - safetensor_files <- list.files(transformer_dir, pattern = "\\.safetensors$", - full.names = TRUE) - if (length(safetensor_files) == 0) { - stop("No safetensor files found in: ", transformer_dir) - } - - if (verbose) { - message(sprintf("Found %d safetensor files in: %s", length(safetensor_files), transformer_dir)) - total_size <- sum(file.info(safetensor_files)$size) / 1e9 - message(sprintf("Total size: %.2f GB (will compress to ~%.2f GB)", - total_size, total_size / 7)) - } - - # Create output directory only when actually writing (CRAN policy) - dir.create(output_dir, showWarnings = FALSE, recursive = TRUE) - - # Quantize - quantize_safetensors_int4( - paths = safetensor_files, - output_path = output_file, - block_size = block_size, - verbose = verbose - ) - - output_file -} - -#' Quantize LTX-2 VAE to INT4 -#' -#' Quantizes the LTX-2 video VAE to INT4 format. -#' -#' @param model_id Character. HuggingFace model ID (default "Lightricks/LTX-2"). -#' @param output_dir Character. Directory to save quantized weights. -#' Default uses `tools::R_user_dir("diffuseR", "cache")`. -#' @param block_size Integer. Block size for INT4 quantization. -#' @param force Logical. Re-quantize even if cached file exists. -#' @param download Logical. If TRUE, download model from HuggingFace if not cached. -#' @param verbose Logical. Print progress. -#' -#' @return Character. Path to the quantized weights file. -#' -#' @export -quantize_ltx2_vae <- function( - model_id = "Lightricks/LTX-2", - output_dir = NULL, - block_size = 64L, - force = FALSE, - download = FALSE, - verbose = TRUE -) { - if (!requireNamespace("hfhub", quietly = TRUE)) { - stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") - } - - # Use R_user_dir for CRAN-compliant cache - - if (is.null(output_dir)) { - output_dir <- tools::R_user_dir("diffuseR", "cache") - } - - output_file <- file.path(output_dir, "ltx2_vae_int4.safetensors") - - # Check if already cached - if (file.exists(output_file) && !force) { - if (verbose) { - size_mb <- file.info(output_file)$size / 1e6 - message(sprintf("Using cached INT4 VAE weights: %s (%.2f MB)", output_file, size_mb)) - } - return(output_file) - } - - # Check if model is available locally via vae/config.json - vae_dir <- tryCatch({ - config_path <- hfhub::hub_download(model_id, "vae/config.json", - local_files_only = TRUE) - dirname(config_path) - }, error = function(e) NULL) - - if (is.null(vae_dir)) { - if (!download) { - stop("Model '", model_id, "' VAE not found in HuggingFace cache.\n", - "Run with download = TRUE to download, or use:\n", - " huggingface-cli download ", model_id) - } - - # Interactive consent before downloading - if (interactive()) { - ans <- utils::askYesNo( - paste0("Download '", model_id, "' VAE from HuggingFace?"), - default = TRUE - ) - if (!isTRUE(ans)) { - stop("Download cancelled.", call. = FALSE) - } - } - - if (verbose) message("Downloading VAE weights from HuggingFace...") - model_path <- hfhub::hub_snapshot(model_id, - allow_patterns = "vae/*") - vae_dir <- file.path(model_path, "vae") - } - if (!dir.exists(vae_dir)) { - stop("VAE directory not found: ", vae_dir) - } - - safetensor_files <- list.files(vae_dir, pattern = "\\.safetensors$", - full.names = TRUE) - if (length(safetensor_files) == 0) { - stop("No safetensor files found in: ", vae_dir) - } - - if (verbose) { - total_size <- sum(file.info(safetensor_files)$size) / 1e6 - message(sprintf("Found %d VAE safetensor files (%.2f MB)", - length(safetensor_files), total_size)) - } - - # Create output directory only when actually writing (CRAN policy) - dir.create(output_dir, showWarnings = FALSE, recursive = TRUE) - - # Quantize - quantize_safetensors_int4( - paths = safetensor_files, - output_path = output_file, - block_size = block_size, - verbose = verbose - ) - - output_file -} - diff --git a/R/img2img.R b/R/img2img.R index 6fbcdd5..4bf3402 100644 --- a/R/img2img.R +++ b/R/img2img.R @@ -29,16 +29,16 @@ #' @return An image array and metadata #' @export -img2img <- function (input_image, prompt, negative_prompt = NULL, - img_dim = 512, model_name = c("sd21", "sdxl"), - pipeline = NULL, devices = "auto", - unet_dtype_str = "float16", download_models = FALSE, - scheduler = "ddim", num_inference_steps = 50, - strength = 0.8, guidance_scale = 7.5, seed = NULL, - save_file = TRUE, filename = NULL, metadata_path = NULL, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, use_native_unet = FALSE, - ...) { +img2img <- function(input_image, prompt, negative_prompt = NULL, + img_dim = 512, model_name = c("sd21", "sdxl"), + pipeline = NULL, devices = "auto", + unet_dtype_str = "float16", download_models = FALSE, + scheduler = "ddim", num_inference_steps = 50, + strength = 0.8, guidance_scale = 7.5, seed = NULL, + save_file = TRUE, filename = NULL, metadata_path = NULL, + use_native_decoder = FALSE, + use_native_text_encoder = FALSE, use_native_unet = FALSE, + ...) { if (model_name %in% c("sd21", "sdxl")) { num_train_timesteps <- 1000 } else { @@ -51,8 +51,9 @@ img2img <- function (input_image, prompt, negative_prompt = NULL, } # 1. Get models - m2d <- models2devices(model_name, devices = devices, unet_dtype_str = NULL, - download_models = download_models) + m2d <- models2devices(model_name, devices = devices, + unet_dtype_str = NULL, + download_models = download_models) devices <- m2d$devices unet_dtype <- m2d$unet_dtype device_cpu <- m2d$device_cpu @@ -60,26 +61,26 @@ img2img <- function (input_image, prompt, negative_prompt = NULL, if (is.null(pipeline)) { pipeline <- load_pipeline(model_name = model_name, m2d = m2d, - i2i = TRUE, - unet_dtype_str = unet_dtype_str, - use_native_decoder = use_native_decoder, - use_native_text_encoder = use_native_text_encoder, - use_native_unet = use_native_unet) + i2i = TRUE, + unet_dtype_str = unet_dtype_str, + use_native_decoder = use_native_decoder, + use_native_text_encoder = use_native_text_encoder, + use_native_unet = use_native_unet) } # 2. Encode input image to latents image_tensor <- preprocess_image(input_image, width = img_dim, height = img_dim, - device = torch::torch_device(devices$encoder)) # Resize & normalize + device = torch::torch_device(devices$encoder)) # Resize & normalize message("Encoding image...") encoded <- pipeline$encoder(image_tensor) # message("Loading quant_conv...") conv_latents <- quant_conv(encoded, dtype = unet_dtype, - device = devices$unet) + device = devices$unet) - latents_mean <- conv_latents[, 1:4,,]# First 4 channels - latents_log_var <- conv_latents[, 5:8,,]# Last 4 channels + latents_mean <- conv_latents[, 1:4,,] # First 4 channels + latents_log_var <- conv_latents[, 5:8,,] # Last 4 channels init_latents <- latents_mean$to(dtype = unet_dtype, - device = torch::torch_device(devices$unet)) * 0.18215 + device = torch::torch_device(devices$unet)) * 0.18215 # Need to FIX # Sample from the distribution (reparameterization trick) # if(eps > 0){ @@ -91,9 +92,9 @@ img2img <- function (input_image, prompt, negative_prompt = NULL, # 3. Compute noise timestep from strength t_strength <- as.integer(strength * num_train_timesteps) schedule <- ddim_scheduler_create(num_train_timesteps = 1000, - num_inference_steps = num_inference_steps, - beta_schedule = "scaled_linear", - device = torch::torch_device(devices$unet)) + num_inference_steps = num_inference_steps, + beta_schedule = "scaled_linear", + device = torch::torch_device(devices$unet)) all_inference_timesteps <- schedule$timesteps timestep_idx <- which.min(abs(all_inference_timesteps - t_strength)) @@ -102,31 +103,32 @@ img2img <- function (input_image, prompt, negative_prompt = NULL, # 4. Add noise to latents message("Adding noise to latent image...") - if (!is.null(seed)) set.seed(seed) + if (!is.null(seed)) { + set.seed(seed) + } noised_latents <- scheduler_add_noise(original_latents = init_latents, noise = torch::torch_randn_like(init_latents), timestep = timestep_start, scheduler_obj = schedule) noised_latents <- noised_latents$to(dtype = unet_dtype, - device = torch::torch_device(devices$unet)) + device = torch::torch_device(devices$unet)) txt2img( - prompt = prompt, - negative_prompt = negative_prompt, - img_dim = img_dim, - model_name = model_name, - pipeline = pipeline, - devices = devices, - unet_dtype_str = unet_dtype_str, - scheduler = "ddim", - timesteps = timesteps, - initial_latents = noised_latents, - num_inference_steps = num_inference_steps, - guidance_scale = guidance_scale, - seed = seed, - save_file = save_file, - filename = filename, - metadata_path = metadata_path + prompt = prompt, + negative_prompt = negative_prompt, + img_dim = img_dim, + model_name = model_name, + pipeline = pipeline, + devices = devices, + unet_dtype_str = unet_dtype_str, + scheduler = "ddim", + timesteps = timesteps, + initial_latents = noised_latents, + num_inference_steps = num_inference_steps, + guidance_scale = guidance_scale, + seed = seed, + save_file = save_file, + filename = filename, + metadata_path = metadata_path ) } - diff --git a/R/jit_ltx23.R b/R/jit_ltx23.R new file mode 100644 index 0000000..3d46e9e --- /dev/null +++ b/R/jit_ltx23.R @@ -0,0 +1,326 @@ +#' LTX-2.3 JIT Block Stack +#' +#' TorchScript compilation of the 48-block NF4 transformer step (cf. +#' the torch skill's JIT-decode pattern proven in whisper and +#' chatterbox). Eager execution crosses R -> lantern per op (~190 us +#' each) and leaves every intermediate as an R tensor handle that only +#' dies at gc(); at high resolution that forces a per-block gc() +#' costing the vast majority of step time. Compiled, the whole block +#' stack is one crossing: intermediates are freed eagerly by libtorch, +#' no R garbage accumulates, no per-block gc is needed, and attention +#' runs through the fused \code{scaled_dot_product_attention} kernel +#' instead of a materialized score matrix. +#' +#' Weights are passed per call as a flat \code{List[Tensor]} (borrowed +#' by reference, no copies) with a fixed per-block layout; the packer +#' and the TorchScript indices must stay in lockstep (parity-tested). +#' +#' @name jit_ltx23 +NULL + +# Per-attention tensor layout (16 slots): +# 0 gate_w, 1 gate_b, +# 2 q_packed, 3 q_absmax, 4 q_bias, +# 5 k_packed, 6 k_absmax, 7 k_bias, +# 8 v_packed, 9 v_absmax, 10 v_bias, +# 11 out_packed, 12 out_absmax, 13 out_bias, +# 14 norm_q_w, 15 norm_k_w +# Per-block layout (114 slots): attn1 @0, audio_attn1 @16, attn2 @32, +# audio_attn2 @48, audio_to_video_attn @64, video_to_audio_attn @80, +# ff @96 (proj_packed, proj_absmax, proj_bias, net2_packed, +# net2_absmax, net2_bias), audio_ff @102, then the modulation tables: +# 108 scale_shift_table, 109 audio_scale_shift_table, +# 110 prompt_scale_shift_table, 111 audio_prompt_scale_shift_table, +# 112 video_a2v_cross_attn_scale_shift_table, +# 113 audio_a2v_cross_attn_scale_shift_table. +.ltx23_jit_slots_per_block <- 114L + +.ltx23_jit_source <- function() { + " +def nf4_lin(x: Tensor, packed: Tensor, absmax: Tensor, bias: Tensor, table: Tensor) -> Tensor: + rows = bias.size(0) + n = packed.size(0) + step = 4194304 + outs: List[Tensor] = [] + i = 0 + while i < n: + j = min(i + step, n) + chunk = packed.narrow(0, i, j - i) + hi = torch.bitwise_right_shift(chunk, 4).long() + lo = torch.bitwise_and(chunk, 15).long() + idx = torch.stack([hi, lo], -1).flatten() + vals = torch.index_select(table, 0, idx) + sc = absmax.narrow(0, i * 2 // 64, (j - i) * 2 // 64) + outs.append((vals.reshape(-1, 64) * sc.unsqueeze(1)).flatten().type_as(x)) + i = j + w = torch.cat(outs, 0).reshape([rows, n * 2 // rows]) + return torch.linear(x, w, bias) + +def rmsn(x: Tensor) -> Tensor: + v = x.float().pow(2).mean(-1, keepdim=True) + return (x * torch.rsqrt(v + 1e-6)).type_as(x) + +def rmsn_w(x: Tensor, w: Tensor) -> Tensor: + v = x.float().pow(2).mean(-1, keepdim=True) + return (x * torch.rsqrt(v + 1e-6)).type_as(w) * w + +def rope(x: Tensor, cs: Tensor, sn: Tensor) -> Tensor: + if cs.dim() == 4: + # Per-head split layout [B, H, T, r]: x [B, T, H*D] -> [B, H, T, D] + b = cs.size(0) + hh = cs.size(1) + t = cs.size(2) + xh = x.reshape([b, t, hh, -1]).transpose(1, 2) + r = xh.size(-1) // 2 + xf = xh.narrow(-1, 0, r).float() + xs = xh.narrow(-1, r, r).float() + o = torch.cat([xf * cs - xs * sn, xs * cs + xf * sn], -1) + return o.transpose(1, 2).reshape([b, t, -1]).type_as(x) + r2 = x.size(-1) // 2 + xf2 = x.narrow(-1, 0, r2).float() + xs2 = x.narrow(-1, r2, r2).float() + return torch.cat([xf2 * cs - xs2 * sn, xs2 * cs + xf2 * sn], -1).type_as(x) + +def mods(tbl: Tensor, temb: Tensor, num: int) -> Tensor: + b = temb.size(0) + t = temb.size(1) + return tbl.unsqueeze(0).unsqueeze(0) + temb.reshape([b, t, num, -1]) + +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], + mask: Optional[Tensor]) -> Tensor: + gl = torch.linear(x, ws[base], ws[base + 1]) + q = rmsn_w(nf4_lin(x, ws[base + 2], ws[base + 3], ws[base + 4], table), ws[base + 14]) + k = rmsn_w(nf4_lin(ctx, ws[base + 5], ws[base + 6], ws[base + 7], table), ws[base + 15]) + v = nf4_lin(ctx, ws[base + 8], ws[base + 9], ws[base + 10], table) + if q_cos is not None and q_sin is not None: + q = rope(q, q_cos, q_sin) + if k_cos is not None and k_sin is not None: + k = rope(k, k_cos, k_sin) + qh = q.unflatten(-1, [heads, -1]).transpose(1, 2) + kh = k.unflatten(-1, [heads, -1]).transpose(1, 2) + vh = v.unflatten(-1, [heads, -1]).transpose(1, 2) + o = torch.scaled_dot_product_attention(qh, kh, vh, attn_mask=mask) + o = o.transpose(1, 2).flatten(2).type_as(x) + gates = torch.sigmoid(gl) * 2.0 + o = (o.unflatten(-1, [heads, -1]) * gates.unsqueeze(-1)).flatten(2) + return nf4_lin(o, ws[base + 11], ws[base + 12], ws[base + 13], table) + +def ff_nf4(x: Tensor, ws: List[Tensor], base: int, table: Tensor) -> Tensor: + h = torch.gelu(nf4_lin(x, ws[base], ws[base + 1], ws[base + 2], table), approximate=\"tanh\") + return nf4_lin(h, ws[base + 3], ws[base + 4], ws[base + 5], table) + +def block_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, + temb: Tensor, temb_a: Tensor, + tcss: Tensor, tcass: Tensor, tcg: Tensor, tcag: Tensor, + tp: Tensor, tpa: 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], + 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) + 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) + + 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) + ah = ah + aax * aada.select(2, 2) + + 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) + 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) + + 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) + aax = attn_nf4(nah, aencm, ws, base + 48, table, aheads, None, None, None, None, aenc_mask) + ah = ah + aax * aada.select(2, 8) + + nh = rmsn(h) + nah = rmsn(ah) + vca = mods(ws[base + 112].narrow(0, 0, 4), tcss, 4) + vcg = mods(ws[base + 112].narrow(0, 4, 1), tcg, 1) + aca = mods(ws[base + 113].narrow(0, 0, 4), tcass, 4) + acg = mods(ws[base + 113].narrow(0, 4, 1), tcag, 1) + + mnh = nh * (vca.select(2, 0) + 1.0) + vca.select(2, 1) + mna = nah * (aca.select(2, 0) + 1.0) + aca.select(2, 1) + a2v = attn_nf4(mnh, mna, ws, base + 64, table, aheads, cav_cos, cav_sin, caa_cos, caa_sin, None) + h = h + vcg.select(2, 0) * a2v + + mnh = nh * (vca.select(2, 2) + 1.0) + vca.select(2, 3) + mna = nah * (aca.select(2, 2) + 1.0) + aca.select(2, 3) + 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) + + 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) + + return (h, ah) + +def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, + temb: Tensor, temb_a: Tensor, + tcss: Tensor, tcass: Tensor, tcg: Tensor, tcag: Tensor, + tp: Tensor, tpa: 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], + ws: List[Tensor], table: Tensor, + n_blocks: int, heads: int, aheads: int) -> Tuple[Tensor, Tensor]: + i = 0 + while i < n_blocks: + 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) + i += 1 + return (h, ah) +" +} + +# Compile once per session +.ltx23_jit_env <- new.env(parent = emptyenv()) + +.ltx23_jit_unit <- function() { + unit <- .ltx23_jit_env$unit + if (is.null(unit)) { + unit <- torch::jit_compile(.ltx23_jit_source()) + .ltx23_jit_env$unit <- unit + } + unit +} + +# Flat List[Tensor] for one attention module (16 slots) +.ltx23_jit_pack_attn <- function(attn) { + list(attn$to_gate_logits$weight, attn$to_gate_logits$bias, + attn$to_q$weight_nf4, attn$to_q$weight_absmax, attn$to_q$bias, + attn$to_k$weight_nf4, attn$to_k$weight_absmax, attn$to_k$bias, + attn$to_v$weight_nf4, attn$to_v$weight_absmax, attn$to_v$bias, + attn$to_out[[1]]$weight_nf4, attn$to_out[[1]]$weight_absmax, + attn$to_out[[1]]$bias, attn$norm_q$weight, attn$norm_k$weight) +} + +.ltx23_jit_pack_ff <- function(ff) { + list(ff$net[[1]]$proj$weight_nf4, ff$net[[1]]$proj$weight_absmax, + ff$net[[1]]$proj$bias, ff$net[[3]]$weight_nf4, + ff$net[[3]]$weight_absmax, ff$net[[3]]$bias) +} + +#' Pack a transformer block's weights for the JIT stack +#' +#' Returns the block's tensors in the fixed 114-slot layout consumed by +#' the compiled \code{stack_nf4}/\code{block_nf4} TorchScript functions. +#' Tensor handles are borrowed, not copied. +#' +#' @param block An NF4-quantized \code{ltx23_transformer_block}. +#' +#' @return List of 114 tensors. +#' +#' @keywords internal +.ltx23_jit_pack_block <- function(block) { + c( + .ltx23_jit_pack_attn(block$attn1), + .ltx23_jit_pack_attn(block$audio_attn1), + .ltx23_jit_pack_attn(block$attn2), + .ltx23_jit_pack_attn(block$audio_attn2), + .ltx23_jit_pack_attn(block$audio_to_video_attn), + .ltx23_jit_pack_attn(block$video_to_audio_attn), + .ltx23_jit_pack_ff(block$ff), + .ltx23_jit_pack_ff(block$audio_ff), + list(block$scale_shift_table, block$audio_scale_shift_table, + block$prompt_scale_shift_table, + block$audio_prompt_scale_shift_table, + block$video_a2v_cross_attn_scale_shift_table, + block$audio_a2v_cross_attn_scale_shift_table) + ) +} + +# A block is JIT-eligible when every cast linear is NF4 (loader swap +# applied) and the gated/adaln 2.3 features the script bakes in are on +.ltx23_jit_block_ok <- function(block) { + inherits(block$attn1$to_q, "ltx23_nf4_linear") && + inherits(block$ff$net[[3]], "ltx23_nf4_linear") && + isTRUE(block$attn1$apply_gated_attention) && + isTRUE(block$video_cross_attn_adaln) && + isTRUE(block$audio_cross_attn_adaln) +} + +# NF4 level table on the right device (cached per device) +.ltx23_jit_table <- function(device) { + key <- paste(device$type, device$index %||% 0L, sep = "|") + tbl <- .ltx23_jit_env[[key]] + if (is.null(tbl)) { + tbl <- torch::torch_tensor(.ltx23_nf4_table, + dtype = torch::torch_float32(), + device = device) + .ltx23_jit_env[[key]] <- tbl + } + tbl +} + +#' Run the block stack through the compiled TorchScript path +#' +#' One R-to-libtorch crossing for all blocks: no per-op dispatch, no R +#' tensor garbage, fused SDPA. Masks must already be additive +#' \code{[B, 1, 1, S]} (or NULL); rope tensors are the \code{[.., r]} +#' cos/sin pairs used by the eager path. +#' +#' @return list(hidden_states, audio_hidden_states) +#' +#' @keywords internal +.ltx23_jit_run_stack <- function(blocks, hidden_states, audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, temb, + temb_audio, temb_ca_scale_shift, + temb_ca_audio_scale_shift, temb_ca_gate, + temb_ca_audio_gate, temb_prompt, + temb_prompt_audio, video_rotary_emb, + audio_rotary_emb, ca_video_rotary_emb, + ca_audio_rotary_emb, + encoder_attention_mask = NULL, + audio_encoder_attention_mask = 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] + ws <- unname(do.call(c, lapply(blocks, .ltx23_jit_pack_block))) + table <- .ltx23_jit_table(hidden_states$device) + heads <- blocks[[1]]$attn1$heads + aheads <- blocks[[1]]$audio_attn1$heads + + # Eager attention takes [B, 1, S] additive masks and unsqueezes to + # [B, 1, 1, S]; SDPA wants them pre-broadcast + fix_mask <- function(m) { + if (!is.null(m) && m$ndim == 3L) { + m$unsqueeze(2L) + } else { + m + } + } + + res <- unit$stack_nf4(hidden_states, audio_hidden_states, + encoder_hidden_states, audio_encoder_hidden_states, + temb, temb_audio, temb_ca_scale_shift, + temb_ca_audio_scale_shift, temb_ca_gate, + temb_ca_audio_gate, temb_prompt, temb_prompt_audio, + video_rotary_emb[[1]], video_rotary_emb[[2]], + audio_rotary_emb[[1]], audio_rotary_emb[[2]], + 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, + torch::jit_scalar(length(blocks)), + torch::jit_scalar(as.integer(heads)), + torch::jit_scalar(as.integer(aheads))) + list(res[[1]], res[[2]]) +} diff --git a/R/jit_vae_ltx23.R b/R/jit_vae_ltx23.R new file mode 100644 index 0000000..e79799a --- /dev/null +++ b/R/jit_vae_ltx23.R @@ -0,0 +1,118 @@ +#' JIT-Traced Decode for the LTX-2.3 VAEs and Vocoder +#' +#' The video/audio decoders and the vocoder are static feed-forward +#' graphs, so \code{torch::jit_trace} converts them wholesale: one +#' R-to-libtorch crossing per forward, intermediates freed eagerly by +#' libtorch instead of accumulating as R handles until gc. Traces are +#' shape-specialized (runtime sizes bake into the graph as constants; +#' a mismatched input errors), so they are cached per instance, input +#' shape, dtype, device, and call tag, and re-traced on a miss. +#' +#' A trace captures the module's weight tensors, which would pin them +#' on the GPU across phase offloads; the pipeline releases all traces +#' whenever a component offloads (\code{.ltx23_release_vae_traces}). +#' +#' Tracing hazard on this lantern build: if the allocator callback +#' runs R's gc \emph{during} trace recording (memory pressure), the +#' recorded graph can capture garbage argument values (observed as +#' corrupted \code{narrow} starts on the full-size decoder; verified +#' 5/5 clean once gc cannot fire mid-trace). Defenses, in order: a +#' gc + cache flush right before each trace so pressure starts near +#' zero, \code{tryCatch} around trace and replay, and a one-time +#' validation of every fresh trace against the eager output — any +#' mismatch permanently blacklists that shape and runs eager. With +#' those in place the traced path cannot corrupt output — but per +#' render it measured slower than eager (traces are released on phase +#' offload, so every render re-pays trace + validation), so it stays +#' opt-in: \code{options(diffuseR.jit_vae = TRUE)}. +#' +#' @name jit_vae_ltx23 +NULL + +.ltx23_vae_traces <- new.env(parent = emptyenv()) + +#' Run a module forward through a shape-specialized trace +#' +#' @param module The nn_module (identity for the cache key; also the +#' default callable). +#' @param x Input tensor. +#' @param forward Optional closure wrapping the call (for extra fixed +#' arguments like \code{causal}); must be pure in \code{x}. +#' @param tag Character. Distinguishes call variants of one module. +#' +#' @return The forward result. +#' +#' @keywords internal +.ltx23_traced_call <- function(module, x, forward = NULL, tag = "") { + run <- forward %||% module + if (!isTRUE(getOption("diffuseR.jit_vae", FALSE))) { + return(run(x)) + } + key <- paste(format(environment(module)), tag, + paste(as.integer(x$shape), collapse = "x"), x$dtype$.type(), + x$device$type, x$device$index %||% 0L, sep = "|") + tr <- .ltx23_vae_traces[[key]] + if (isFALSE(tr)) { + # This shape failed tracing or validation before: eager only + return(run(x)) + } + if (is.null(tr)) { + # Start the trace from minimal memory pressure so the + # allocator callback cannot run R gc mid-recording (which + # corrupts captured argument values on this lantern build) + gc(verbose = FALSE) + if (x$device$type == "cuda") { + torch::cuda_empty_cache() + } + # Tracing captures parameters as constants, which requires + # grad-free tensors; these modules are inference-only. Always + # trace through a plain closure: jit_trace treats a bare + # nn_module via a separate (and here broken) code path. + for (p in module$parameters) { + p$requires_grad_(FALSE) + } + out <- tryCatch({ + tr <- torch::jit_trace(function(z) run(z), x) + replay <- tr(x) + ref <- run(x) + d <- (replay$to(dtype = torch::torch_float32()) - + ref$to(dtype = torch::torch_float32()))$abs()$max() + if (as.numeric(d) > 0) { + warning("diffuseR: traced decode failed validation for ", + "shape [", paste(as.integer(x$shape), collapse = ", "), + "]; using the eager module for it", call. = FALSE) + .ltx23_vae_traces[[key]] <- FALSE + ref + } else { + .ltx23_vae_traces[[key]] <- tr + replay + } + }, error = function(e) { + .ltx23_vae_traces[[key]] <- FALSE + run(x) + }) + return(out) + } + tr(x) +} + +#' Release all cached decode traces +#' +#' Traces hold references to the weight tensors they captured; drop +#' them when a component leaves the GPU so its memory actually frees. +#' +#' @return Invisibly, NULL. +#' +#' @keywords internal +.ltx23_release_vae_traces <- function() { + rm(list = ls(.ltx23_vae_traces), envir = .ltx23_vae_traces) + invisible(NULL) +} + +# One choke point for the three video-decoder call sites (direct, +# spatially tiled, temporally tiled) +.ltx23_decode_tile <- function(vae, x, causal) { + dec <- vae$decoder + .ltx23_traced_call(dec, x, forward = function(z) dec(z, causal = causal), + tag = paste0("causal:", format(causal))) +} diff --git a/R/load_model_component.R b/R/load_model_component.R index 76380cc..b5b16a5 100644 --- a/R/load_model_component.R +++ b/R/load_model_component.R @@ -19,19 +19,20 @@ #' \dontrun{ #' unet <- load_model_component("unet", "sd21", "cpu") #' } -load_model_component <- function (component, model_name = "sd21", - device = "cpu", unet_dtype_str = NULL, - download = TRUE, use_native = FALSE) { +load_model_component <- function(component, model_name = "sd21", + device = "cpu", unet_dtype_str = NULL, + download = TRUE, use_native = FALSE) { # Set valid components based on model if (model_name == "sdxl") { - valid_components <- c("unet", "decoder", "text_encoder", "text_encoder2", "encoder") + valid_components <- c("unet", "decoder", "text_encoder", + "text_encoder2", "encoder") } else { valid_components <- c("unet", "decoder", "text_encoder", "encoder") } if (!component %in% valid_components) { stop("Invalid component name for model '", model_name, "'. Must be one of: ", - paste(valid_components, collapse = ", ")) + paste(valid_components, collapse = ", ")) } # Determine filename for this component @@ -56,7 +57,8 @@ load_model_component <- function (component, model_name = "sd21", } else { model <- unet_native_from_torchscript(file_path, verbose = FALSE) } - if (device == "cuda" && (is.null(unet_dtype_str) || unet_dtype_str == "float16")) { + if (device == "cuda" && + (is.null(unet_dtype_str) || unet_dtype_str == "float16")) { model$to(device = torch::torch_device(device), dtype = torch::torch_float16()) } else { model$to(device = torch::torch_device(device)) @@ -68,25 +70,25 @@ load_model_component <- function (component, model_name = "sd21", } else if (use_native && component == "text_encoder") { arch <- detect_text_encoder_architecture(file_path) model <- text_encoder_native( - vocab_size = arch$vocab_size, - context_length = arch$context_length, - embed_dim = arch$embed_dim, - num_layers = arch$num_layers, - num_heads = arch$num_heads, - mlp_dim = arch$mlp_dim, - apply_final_ln = arch$apply_final_ln + vocab_size = arch$vocab_size, + context_length = arch$context_length, + embed_dim = arch$embed_dim, + num_layers = arch$num_layers, + num_heads = arch$num_heads, + mlp_dim = arch$mlp_dim, + apply_final_ln = arch$apply_final_ln ) load_text_encoder_weights(model, file_path, verbose = FALSE) model$to(device = torch::torch_device(device)) } else if (use_native && component == "text_encoder2") { arch <- detect_text_encoder_architecture(file_path) model <- text_encoder2_native( - vocab_size = arch$vocab_size, - context_length = arch$context_length, - embed_dim = arch$embed_dim, - num_layers = arch$num_layers, - num_heads = arch$num_heads, - mlp_dim = arch$mlp_dim + vocab_size = arch$vocab_size, + context_length = arch$context_length, + embed_dim = arch$embed_dim, + num_layers = arch$num_layers, + num_heads = arch$num_heads, + mlp_dim = arch$mlp_dim ) load_text_encoder2_weights(model, file_path, verbose = FALSE) model$to(device = torch::torch_device(device)) @@ -98,7 +100,7 @@ load_model_component <- function (component, model_name = "sd21", } # Build the expected filename for a component -component_filename <- function (component, device, unet_dtype_str = NULL) { +component_filename <- function(component, device, unet_dtype_str = NULL) { if (component == "unet" && device != "cpu") { dtype <- if (is.null(unet_dtype_str) || unet_dtype_str == "float16") { "float16" @@ -114,20 +116,17 @@ component_filename <- function (component, device, unet_dtype_str = NULL) { } # Convenience function to load both text encoders for SDXL -load_text_encoders <- function (model_name = "sdxl", device = "cpu", - download = TRUE) { +load_text_encoders <- function(model_name = "sdxl", device = "cpu", + download = TRUE) { if (model_name != "sdxl") { return(list( - text_encoder = load_model_component("text_encoder", model_name, device, download = download) + text_encoder = load_model_component("text_encoder", + model_name, device, download = download) )) } text_encoder <- load_model_component("text_encoder", model_name, device, download = download) text_encoder2 <- load_model_component("text_encoder2", model_name, device, download = download) - list( - text_encoder = text_encoder, - text_encoder2 = text_encoder2 - ) + list(text_encoder = text_encoder, text_encoder2 = text_encoder2) } - diff --git a/R/load_pipeline.R b/R/load_pipeline.R index e6df28f..c3f08ce 100644 --- a/R/load_pipeline.R +++ b/R/load_pipeline.R @@ -24,59 +24,51 @@ #' pipeline <- load_pipeline("my_model", device = "cuda") #' } #' -load_pipeline <- function( - model_name, - m2d, - i2i = FALSE, - unet_dtype_str, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, - ... -) { - # Create an environment to store the pipeline components - # pipeline <- new.env(parent = emptyenv()) - devices <- m2d$devices - unet_dtype <- m2d$unet_dtype - device_cpu <- m2d$device_cpu - device_cuda <- m2d$device_cuda +load_pipeline <- function(model_name, m2d, i2i = FALSE, unet_dtype_str, + use_native_decoder = FALSE, + use_native_text_encoder = FALSE, + use_native_unet = FALSE, ...) { + # Create an environment to store the pipeline components + # pipeline <- new.env(parent = emptyenv()) + devices <- m2d$devices + unet_dtype <- m2d$unet_dtype + device_cpu <- m2d$device_cpu + device_cuda <- m2d$device_cuda - pipeline <- list() - # Load models into the environment - if (i2i) { - message("Loading image encoder...") - pipeline$encoder <- load_model_component(component = "encoder", - model_name, - devices$encoder) - } - message("Loading text_encoder...") - pipeline$text_encoder <- load_model_component("text_encoder", model_name, - devices$text_encoder, - use_native = use_native_text_encoder) - if (model_name == "sdxl") { - message("Loading text_encoder2...") - pipeline$text_encoder2 <- load_model_component(component = "text_encoder2", - model_name, - devices$text_encoder2, - use_native = use_native_text_encoder) - } - message("Loading unet...") - pipeline$unet <- load_model_component("unet", model_name, - device = devices$unet, - unet_dtype_str = unet_dtype_str, - use_native = use_native_unet) + pipeline <- list() + # Load models into the environment + if (i2i) { + message("Loading image encoder...") + pipeline$encoder <- load_model_component(component = "encoder", + model_name, devices$encoder) + } + message("Loading text_encoder...") + pipeline$text_encoder <- load_model_component("text_encoder", model_name, + devices$text_encoder, + use_native = use_native_text_encoder) + if (model_name == "sdxl") { + message("Loading text_encoder2...") + pipeline$text_encoder2 <- load_model_component(component = "text_encoder2", + model_name, + devices$text_encoder2, + use_native = use_native_text_encoder) + } + message("Loading unet...") + pipeline$unet <- load_model_component("unet", model_name, + device = devices$unet, + unet_dtype_str = unet_dtype_str, + use_native = use_native_unet) - message("Loading image decoder...") - pipeline$decoder <- load_model_component("decoder", model_name, - devices$decoder, - use_native = use_native_decoder) + message("Loading image decoder...") + pipeline$decoder <- load_model_component("decoder", model_name, + devices$decoder, + use_native = use_native_decoder) - # Store configuration - pipeline$devices <- devices + # Store configuration + pipeline$devices <- devices - # Add a class for S3 method dispatch if needed - # class(pipeline) <- c("diffusion_pipeline", "environment") + # Add a class for S3 method dispatch if needed + # class(pipeline) <- c("diffusion_pipeline", "environment") - return(pipeline) + return(pipeline) } - diff --git a/R/memory_ltx23.R b/R/memory_ltx23.R new file mode 100644 index 0000000..5f92a48 --- /dev/null +++ b/R/memory_ltx23.R @@ -0,0 +1,151 @@ +#' LTX-2.3 Memory Profiles and CUDA GC Tuning +#' +#' Memory management for running the 22B LTX-2.3 transformer on limited +#' VRAM, built on the patterns proven in the whisper and chatterbox +#' packages: torch allocator GC tuning before the first CUDA op, fp8 +#' CPU-resident streaming weights, query-chunked attention, and +#' phase-sequential component placement. +#' +#' @name memory_ltx23 +NULL + +#' Get an LTX-2.3 memory profile +#' +#' Selects transformer precision, component placement, and attention +#' chunking for the available VRAM. Measured on an RTX 5060 Ti (16 GB): +#' fp8 streaming peaks ~11.6 GB (without phase offloading) at +#' 512x320x49; NF4 keeps the whole 22B transformer resident (~12.5 GB) +#' and removes the ~21 GB/step PCIe weight streaming. The NF4 profile +#' renders 1280x704x121 with audio in ~23 min at a 15.7 GB peak +#' (tiled VAE decode, in-place feed-forward GELU, and the default +#' \code{diffuseR.attn_budget} of 1.5e8 all required at that size). +#' +#' \describe{ +#' \item{precision "nf4"}{Weights resident on the GPU; fastest steps; +#' about 9 percent weight round-trip error.} +#' \item{precision "fp8"}{Weights CPU-resident, streamed per linear; +#' near-bf16 quality; each step pays the PCIe transfer.} +#' } +#' +#' @param vram_gb Numeric or NULL (auto-detect free VRAM). +#' +#' @return Named list with device/dtype placement, \code{attn_chunk}, +#' \code{pin_weights}, and resolution caps. +#' +#' @export +ltx23_memory_profile <- function(vram_gb = NULL) { + if (is.null(vram_gb)) { + vram_gb <- .detect_vram(use_free = TRUE) + message(sprintf("Detected %.1f GB free VRAM", vram_gb)) + } + + profile <- if (vram_gb >= 14) { + "high" + } else if (vram_gb >= 10) { + "medium" + } else if (vram_gb >= 7) { + "low" + } else { + "cpu_only" + } + + profiles <- list( + high = list(name = "high", device = "cuda", dtype = "bfloat16", + precision = "nf4", phase_offload = TRUE, + pin_weights = FALSE, attn_chunk = NULL, + text_device = "cpu", max_resolution = c(704L, 1280L), + max_frames = 121L), + medium = list( + name = "medium", + device = "cuda", + dtype = "bfloat16", + precision = "fp8", + phase_offload = TRUE, + pin_weights = TRUE, + attn_chunk = 4096L, + text_device = "cpu", + max_resolution = c(576L, 1024L), + max_frames = 121L + ), + low = list( + name = "low", + device = "cuda", + dtype = "bfloat16", + precision = "fp8", + phase_offload = TRUE, + pin_weights = TRUE, + attn_chunk = 2048L, + text_device = "cpu", + max_resolution = c(512L, 768L), + max_frames = 65L + ), + cpu_only = list( + name = "cpu_only", + device = "cpu", + dtype = "float32", + precision = "fp8", + phase_offload = FALSE, + pin_weights = FALSE, + attn_chunk = NULL, + text_device = "cpu", + max_resolution = c(384L, 640L), + max_frames = 33L + ) + ) + + profiles[[profile]] +} + +#' Tune the torch CUDA allocator for large-resident inference +#' +#' Stops the allocator GC storm (cf. ~/skills/torch +#' torch-jit-gc-performance.md): lantern proactively calls R's gc() +#' whenever reserved memory exceeds \code{torch.cuda_allocator_reserved_rate} +#' (default 0.20) of the card. With ~75\% of VRAM occupied by resident +#' weights that fires on nearly every allocation. Raising the rate to the +#' actual footprint is safe here because the LTX hot loops compute into +#' persistent scratch buffers (near-zero per-step garbage). Also raises +#' the host-allocation GC threshold and defaults +#' \code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. Must run before +#' the first CUDA op; user-set options win. +#' +#' @param footprint_gb Numeric. Expected resident GPU footprint in GB +#' (NF4 transformer: ~12). +#' @param total_gb Numeric or NULL (auto-detect total VRAM). +#' +#' @return Invisibly, the applied reserved rate (NULL if skipped). +#' +#' @export +ltx23_tune_gc <- function(footprint_gb = 12, total_gb = NULL) { + if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { + Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True") + } + if (!torch::cuda_is_available()) { + return(invisible(NULL)) + } + if (is.null(total_gb)) { + total_gb <- .detect_vram(use_free = FALSE) + if (!isTRUE(total_gb > 0)) { + return(invisible(NULL)) + } + } + if (is.null(getOption("torch.threshold_call_gc"))) { + options(torch.threshold_call_gc = 16000) + } + # The other two allocator-callback gates (defaults 0.8): measured + # 32.7s -> 21.5s on the tiled decode under expandable segments, and + # R-gc share 50% -> 32% on the native backend, with no wall-time + # downside in any condition + if (is.null(getOption("torch.cuda_allocator_allocated_rate"))) { + options(torch.cuda_allocator_allocated_rate = 0.95) + } + if (is.null(getOption("torch.cuda_allocator_allocated_reserved_rate"))) { + options(torch.cuda_allocator_allocated_reserved_rate = 0.95) + } + rate <- NULL + if (is.null(getOption("torch.cuda_allocator_reserved_rate"))) { + rate <- min(0.92, max(0.20, footprint_gb / total_gb)) + options(torch.cuda_allocator_reserved_rate = rate) + } + invisible(rate) +} diff --git a/R/memory_sdxl.R b/R/memory_sdxl.R new file mode 100644 index 0000000..4b206c4 --- /dev/null +++ b/R/memory_sdxl.R @@ -0,0 +1,114 @@ +#' Get SDXL Memory Profile +#' +#' Determines optimal memory configuration for SDXL image generation +#' based on available VRAM. +#' +#' @param vram_gb Numeric. Available VRAM in GB, or NULL for auto-detection. +#' +#' @return A list with memory profile settings. +#' +#' @details +#' Memory profiles for SDXL: +#' \describe{ +#' \item{full_gpu}{16GB+ - All components on CUDA} +#' \item{balanced}{10-12GB - UNet + decoder on CUDA, text encoders on CPU} +#' \item{unet_gpu}{6-10GB - Only UNet on CUDA, everything else CPU} +#' \item{cpu_only}{<6GB - All on CPU} +#' } +#' +#' Each profile also specifies: +#' - cfg_mode: "batched" or "sequential" (sequential halves peak memory) +#' - cleanup: "none", "phase", or "step" (when to clear VRAM) +#' - dtype: "float16" or "float32" +#' - max_resolution: maximum image dimension +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' # Auto-detect profile +#' profile <- sdxl_memory_profile() +#' +#' # Specific VRAM +#' profile <- sdxl_memory_profile(vram_gb = 8) +#' } +sdxl_memory_profile <- function(vram_gb = NULL) { + # Auto-detect free VRAM if not provided + if (is.null(vram_gb)) { + vram_gb <- .detect_vram(use_free = TRUE) + message(sprintf("Detected %.1f GB free VRAM", vram_gb)) + } + + # Determine profile level + if (vram_gb >= 16) { + profile <- "full_gpu" + } else if (vram_gb >= 10) { + profile <- "balanced" + } else if (vram_gb >= 6) { + profile <- "unet_gpu" + } else { + profile <- "cpu_only" + } + + # Build profile config + profiles <- list( + full_gpu = list( + name = "full_gpu", + devices = list(unet = "cuda", decoder = "cuda", + text_encoder = "cuda", text_encoder2 = "cuda", + encoder = "cuda"), + dtype = "float16", + cfg_mode = "batched", + cleanup = "none", + max_resolution = 1536L, + step_cleanup_interval = 0L # No step cleanup + ), + balanced = list( + name = "balanced", + devices = list( + unet = "cuda", + decoder = "cuda", + text_encoder = "cpu", + text_encoder2 = "cpu", + encoder = "cpu" + ), + dtype = "float16", + cfg_mode = "batched", + cleanup = "phase", # Cleanup between text encoding and denoising + max_resolution = 1024L, + step_cleanup_interval = 0L + ), + unet_gpu = list( + name = "unet_gpu", + devices = list( + unet = "cuda", + decoder = "cpu", + text_encoder = "cpu", + text_encoder2 = "cpu", + encoder = "cpu" + ), + dtype = "float16", + cfg_mode = "sequential", # Sequential CFG halves peak memory + cleanup = "phase", + max_resolution = 1024L, + step_cleanup_interval = 10L # Cleanup every 10 steps + ), + cpu_only = list( + name = "cpu_only", + devices = list( + unet = "cpu", + decoder = "cpu", + text_encoder = "cpu", + text_encoder2 = "cpu", + encoder = "cpu" + ), + dtype = "float32", # CPU often faster with float32 + cfg_mode = "sequential", + cleanup = "none", # No GPU to clean + max_resolution = 768L, + step_cleanup_interval = 0L + ) + ) + + profiles[[profile]] +} diff --git a/R/models2devices.R b/R/models2devices.R index 11317e8..5af8188 100644 --- a/R/models2devices.R +++ b/R/models2devices.R @@ -9,8 +9,8 @@ #' #' @return A list containing the device configuration, UNet data type, and CPU/CUDA devices. #' @export -models2devices <- function (model_name, devices = "cpu", unet_dtype_str = NULL, - download_models = FALSE) { +models2devices <- function(model_name, devices = "cpu", + unet_dtype_str = NULL, download_models = FALSE) { # Validation (same as before) if (is.null(model_name) || !is.character(model_name)) { stop("Invalid model name") @@ -38,13 +38,13 @@ models2devices <- function (model_name, devices = "cpu", unet_dtype_str = NULL, # Verify model files are available (downloads if allowed) download_model(model_name, devices, unet_dtype_str, - download_models = download_models) + download_models = download_models) return(list( - devices = devices, - unet_dtype = unet_dtype, - device_cpu = device_cpu, - device_cuda = device_cuda + devices = devices, + unet_dtype = unet_dtype, + device_cpu = device_cpu, + device_cuda = device_cuda )) } @@ -52,13 +52,14 @@ models2devices <- function (model_name, devices = "cpu", unet_dtype_str = NULL, #' @description This function returns a list of required components for each supported model type. #' @param model_name A character string representing the name of the model. #' @return A character vector of required components for the specified model. -get_required_components <- function (model_name) { +get_required_components <- function(model_name) { components <- list( - # "sd15" = c("unet", "decoder", "text_encoder", "encoder"), - "sd21" = c("unet", "decoder", "text_encoder", "encoder"), - "sdxl" = c("unet", "decoder", "text_encoder", "text_encoder2", "encoder") - # "sd3" = c("transformer", "decoder", "text_encoder", "text_encoder2", "text_encoder3", "encoder"), - # "cascade" = c("prior", "decoder", "text_encoder", "vqgan") + # "sd15" = c("unet", "decoder", "text_encoder", "encoder"), + "sd21" = c("unet", "decoder", "text_encoder", "encoder"), + "sdxl" = c("unet", "decoder", "text_encoder", "text_encoder2", + "encoder") + # "sd3" = c("transformer", "decoder", "text_encoder", "text_encoder2", "text_encoder3", "encoder"), + # "cascade" = c("prior", "decoder", "text_encoder", "vqgan") ) if (!model_name %in% names(components)) { @@ -74,7 +75,7 @@ get_required_components <- function (model_name) { #' @param devices A character string or a named list specifying the devices for model components. #' @param required_components A character vector of required components for the model. #' @return A named list of devices for each required component. -standardize_devices <- function (devices, required_components) { +standardize_devices <- function(devices, required_components) { if (is.character(devices) && length(devices) == 1) { # Single device string - apply to all components device_list <- as.list(rep(devices, length(required_components))) @@ -91,12 +92,15 @@ standardize_devices <- function (devices, required_components) { for (component in missing_components) { if (component == "encoder" && "decoder" %in% names(devices)) { devices$encoder <- devices$decoder - } else if (component == "text_encoder2" && "text_encoder" %in% names(devices)) { + } else if (component == "text_encoder2" && + "text_encoder" %in% names(devices)) { devices$text_encoder2 <- devices$text_encoder message("text_encoder2 is set to text_encoder") - } else if (component == "text_encoder3" && "text_encoder" %in% names(devices)) { + } else if (component == "text_encoder3" && + "text_encoder" %in% names(devices)) { devices$text_encoder3 <- devices$text_encoder - } else if (component == "transformer" && "unet" %in% names(devices)) { + } else if (component == "transformer" && + "unet" %in% names(devices)) { devices$transformer <- devices$unet } else { stop("Missing required component: ", component) @@ -113,7 +117,7 @@ standardize_devices <- function (devices, required_components) { #' @param devices A character string or a named list specifying the devices for model components. #' @param unet_dtype_str A character string specifying the data type for the UNet model. #' @return A torch dtype object based on the main computation device. -setup_dtype <- function (devices, unet_dtype_str) { +setup_dtype <- function(devices, unet_dtype_str) { # Find the main computation device (unet or transformer) main_device <- if ("unet" %in% names(devices)) { devices$unet @@ -139,4 +143,3 @@ setup_dtype <- function (devices, unet_dtype_str) { stop("Invalid device: ", main_device) } } - diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R new file mode 100644 index 0000000..4c454a9 --- /dev/null +++ b/R/nf4_ltx23.R @@ -0,0 +1,475 @@ +#' NF4 Weight Storage for the LTX-2.3 Transformer +#' +#' 4-bit NormalFloat quantization (the QLoRA scheme: per-block absmax +#' normalization against a 16-level quantile code, two indices packed +#' per byte). At ~4.5 bits/parameter the 22B transformer fits in about +#' 12.5 GB, small enough to stay resident on a 16 GB GPU: no per-step +#' PCIe weight streaming, at a small quality cost relative to fp8. +#' Quantization and dequantization are pure torch ops (bucketize, +#' index_select) - no custom kernels. +#' +#' @name nf4_ltx23 +NULL + +# The 16 NF4 quantile levels (QLoRA, Dettmers et al. 2023) +.ltx23_nf4_table <- c(-1.0, -0.6961928009986877, -0.5250730514526367, + -0.39491748809814453, -0.28444138169288635, + -0.18477343022823334, -0.09105003625154495, 0.0, + 0.07958029955625534, 0.16093020141124725, + 0.24611230194568634, 0.33791524171829224, + 0.44070982933044434, 0.5626170039176941, + 0.7229568362236023, 1.0) + +.ltx23_nf4_block_size <- 64L + +#' Quantize a tensor to NF4 +#' +#' @param x Float tensor (any shape; total elements must be a multiple +#' of 128, i.e. two 64-element blocks - always true for the LTX +#' linears). +#' +#' @return List with \code{packed} (uint8, two indices per byte) and +#' \code{absmax} (float32, one per 64-element block). +#' +#' @export +ltx23_nf4_quantize <- function(x) { + n <- prod(x$shape) + block <- .ltx23_nf4_block_size + if (n %% block != 0L) { + stop("Tensor length must be a multiple of ", block) + } + + table <- torch::torch_tensor(.ltx23_nf4_table, + dtype = torch::torch_float32(), + device = x$device) + # Decision boundaries at the midpoints between adjacent levels + midpoints <- (table$narrow(1L, 1L, 15L) + table$narrow(1L, 2L, 15L)) / 2 + + blocks <- x$detach()$to(dtype = torch::torch_float32())$reshape(c(-1L, block)) + absmax <- blocks$abs()$amax(dim = 2L)$clamp(min = 1e-12) + normalized <- blocks / absmax$unsqueeze(2L) + + idx <- torch::torch_bucketize(normalized$flatten(), midpoints) # 0..15 + idx <- idx$to(dtype = torch::torch_uint8()) + + # Pack pairs: first index in the high nibble + pairs <- idx$reshape(c(-1L, 2L)) + packed <- pairs$narrow(2L, 1L, 1L)$squeeze(2L) * 16L + + pairs$narrow(2L, 2L, 1L)$squeeze(2L) + + list(packed = packed, absmax = absmax) +} + +#' Dequantize NF4 data to a float tensor +#' +#' @param packed uint8 tensor of packed index pairs. +#' @param absmax float32 tensor of per-block scales. +#' @param shape Integer vector. Original tensor shape. +#' @param dtype Target torch dtype. +#' @param chunk_elements Integer. Elements dequantized per slice (bounds +#' the int64 index temporary). +#' @param out Optional preallocated tensor of \code{shape} to write into +#' (avoids allocating a fresh weight tensor per call). +#' +#' @return Tensor of \code{shape} in \code{dtype} on the input's device. +#' +#' @export +ltx23_nf4_dequantize <- function(packed, absmax, shape, + dtype = torch::torch_bfloat16(), + chunk_elements = 8388608L, out = NULL) { + block <- .ltx23_nf4_block_size + + if (is.null(out)) { + out <- torch::torch_empty(shape, dtype = dtype, device = packed$device) + } + out_flat <- out$view(-1L) + + n_bytes <- packed$shape[1] + bytes_per_chunk <- max(chunk_elements %/% 2L, block) + scratch <- .ltx23_get_dequant_scratch(min(bytes_per_chunk, n_bytes), + packed$device) + + start <- 1L + torch::with_no_grad({ + while (start <= n_bytes) { + len <- min(bytes_per_chunk, n_bytes - start + 1L) + chunk <- packed$narrow(1L, start, len) + + # Fully in-place nibble unpack into persistent scratch: + # hi = byte %/% 16, lo = byte - 16 * hi + hi <- scratch$hi$narrow(1L, 1L, len) + lo <- scratch$lo$narrow(1L, 1L, len) + hi$copy_(chunk)$div_(16L, rounding_mode = "floor") + lo$copy_(hi)$mul_(-16L)$add_(chunk) + + # Interleave into the (1-based) int64 index scratch + idx <- scratch$idx$narrow(1L, 1L, len * 2L) + idx_pairs <- idx$view(c(-1L, 2L)) + idx_pairs$narrow(2L, 1L, 1L)$squeeze(2L)$copy_(hi) + idx_pairs$narrow(2L, 2L, 1L)$squeeze(2L)$copy_(lo) + idx$add_(1L) + + vals <- scratch$vals$narrow(1L, 1L, len * 2L) + .ltx23_index_select_into(vals, scratch$table, idx) + + block_start <- ((start - 1L) * 2L) %/% block + 1L + n_blocks <- (len * 2L) %/% block + scales <- absmax$narrow(1L, block_start, n_blocks) + vals$view(c(-1L, block))$mul_(scales$unsqueeze(2L)) + + out_flat$narrow(1L, (start - 1L) * 2L + 1L, len * 2L)$copy_(vals) + start <- start + len + } + }) + out +} + +# torch_index_select_out is not exported from torch; fall back to an +# allocating index_select if it ever disappears +.ltx23_index_select_fn <- local({ + fn <- NULL + function() { + if (is.null(fn)) { + fn <<- tryCatch( + get("torch_index_select_out", envir = asNamespace("torch")), + error = function(e) FALSE + ) + } + fn + } +}) + +.ltx23_index_select_into <- function(out, table, idx) { + fn <- .ltx23_index_select_fn() + if (isFALSE(fn)) { + out$copy_(torch::torch_index_select(table, 1L, idx)) + } else { + fn(out, table, 1L, idx) + } + invisible(out) +} + +# Persistent per-device dequantization scratch (nibbles, indices, values, +# and the level table), sized to the chunk length +.ltx23_dequant_scratch <- new.env(parent = emptyenv()) + +.ltx23_get_dequant_scratch <- function(n_bytes, device) { + key <- paste(device$type, device$index %||% 0L, sep = "|") + scratch <- .ltx23_dequant_scratch[[key]] + if (is.null(scratch) || scratch$n_bytes < n_bytes) { + scratch <- list( + n_bytes = n_bytes, + hi = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(), + device = device), + lo = torch::torch_empty(n_bytes, dtype = torch::torch_uint8(), + device = device), + idx = torch::torch_empty(n_bytes * 2L, + dtype = torch::torch_long(), + device = device), + vals = torch::torch_empty(n_bytes * 2L, + dtype = torch::torch_float32(), + device = device), + table = torch::torch_tensor(.ltx23_nf4_table, + dtype = torch::torch_float32(), + device = device) + ) + .ltx23_dequant_scratch[[key]] <- scratch + } + scratch +} + +# Reusable dequantization buffers, keyed by shape/dtype/device: each +# distinct weight shape gets one long-lived buffer, so per-step +# dequantization allocates nothing (the buffer is overwritten in place +# by the next linear of the same shape) +.ltx23_dequant_buffers <- new.env(parent = emptyenv()) + +.ltx23_get_dequant_buffer <- function(shape, dtype, device) { + key <- paste(paste(shape, collapse = "x"), dtype$.type(), + paste(device$type, device$index %||% 0L), sep = "|") + buf <- .ltx23_dequant_buffers[[key]] + if (is.null(buf)) { + buf <- torch::torch_empty(shape, dtype = dtype, device = device) + .ltx23_dequant_buffers[[key]] <- buf + } + buf +} + +#' Release the NF4 dequantization buffers +#' +#' Frees the cached per-shape weight buffers (e.g. before decoding at +#' high resolution). +#' +#' @return Invisibly, NULL. +#' +#' @export +ltx23_release_dequant_buffers <- function() { + rm(list = ls(.ltx23_dequant_buffers), envir = .ltx23_dequant_buffers) + rm(list = ls(.ltx23_dequant_scratch), envir = .ltx23_dequant_scratch) + .ltx23_release_attn_buffers() + gc(verbose = FALSE) + if (torch::cuda_is_available()) { + tryCatch(torch::cuda_empty_cache(), error = function(e) NULL) + } + invisible(NULL) +} + +#' NF4 linear layer +#' +#' Packed weights and per-block scales are registered as buffers, so +#' they move with the module (uint8 packs are untouched by dtype +#' conversions). The forward pass dequantizes on the weight's device. +#' +#' @param out_features,in_features Integers. +#' @param bias Logical. +#' +#' @export +ltx23_nf4_linear <- torch::nn_module( + "ltx23_nf4_linear", + initialize = function(out_features, in_features, bias = TRUE) { + self$out_features <- as.integer(out_features) + self$in_features <- as.integer(in_features) + n <- self$out_features * self$in_features + self$weight_nf4 <- torch::nn_buffer( + torch::torch_zeros(n %/% 2L, dtype = torch::torch_uint8()) + ) + self$weight_absmax <- torch::nn_buffer( + torch::torch_ones(n %/% .ltx23_nf4_block_size, + dtype = torch::torch_float32()) + ) + if (bias) { + self$bias <- torch::nn_parameter(torch::torch_zeros(out_features)) + } +}, + set_nf4_weight = function(packed, absmax) { + torch::with_no_grad({ + self$weight_nf4$copy_(packed) + self$weight_absmax$copy_(absmax) + }) + invisible(self) +}, + forward = function(x) { + w <- .ltx23_get_dequant_buffer( + c(self$out_features, self$in_features), x$dtype, + self$weight_nf4$device + ) + ltx23_nf4_dequantize( + self$weight_nf4, self$weight_absmax, + c(self$out_features, self$in_features), + dtype = x$dtype, out = w + ) + torch::nnf_linear(x, w, self$bias) +} +) + +#' Quantize an LTX-2.3 checkpoint to NF4 shards +#' +#' Same streaming layout and cast policy as +#' \code{\link{ltx23_quantize_fp8}}, but cast-set weights are stored as +#' NF4 (\code{} packed uint8 + \code{_absmax} float32 blocks + +#' the original shape recovered from the model config at load time). +#' Non-cast tensors are copied through unchanged. The manifest carries +#' \code{format = "nf4"}. +#' +#' @param checkpoint_path Source .safetensors (bf16 single file). +#' @param output_dir Output directory for shards + manifest. +#' @param shard_bytes Numeric. Approximate shard size. +#' @param force Logical. Re-quantize even if a valid manifest exists. +#' @param verbose Logical. +#' +#' @return Invisibly, the manifest list. +#' +#' @export +ltx23_quantize_nf4 <- function(checkpoint_path, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-nf4"), + shard_bytes = 4e9, force = FALSE, + verbose = TRUE) { + manifest_path <- file.path(output_dir, "manifest.json") + if (!force && file.exists(manifest_path)) { + manifest <- jsonlite::fromJSON(manifest_path) + if (all(file.exists(file.path(output_dir, manifest$shards)))) { + if (verbose) { + message("NF4 artifact already present: ", output_dir) + } + return(invisible(manifest)) + } + } + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + + ckpt <- ltx23_open_checkpoint(checkpoint_path) + + shard <- list() + shard_size <- 0 + shard_files <- character(0) + n_cast <- 0L + + flush_shard <- function() { + if (!length(shard)) { + return() + } + fname <- sprintf("ltx2.3-nf4-%05d.safetensors", + length(shard_files) + 1L) + safetensors::safe_save_file(shard, file.path(output_dir, fname)) + shard_files[[length(shard_files) + 1L]] <<- fname + if (verbose) { + message(sprintf(" wrote %s (%.2f GB, %d tensors)", fname, + shard_size / 1e9, length(shard))) + } + shard <<- list() + shard_size <<- 0 + gc(verbose = FALSE) + } + + keys <- ckpt$keys + for (i in seq_along(keys)) { + key <- keys[[i]] + tensor <- ckpt$handle$get_tensor(key) + + mapped <- ltx23_map_dit_key(key) + if (startsWith(key, "model.diffusion_model.") && + ltx23_is_fp8_cast_key(mapped)) { + torch::with_no_grad({ + q <- ltx23_nf4_quantize(tensor) + shard[[key]] <- q$packed + shard[[paste0(key, "_absmax")]] <- q$absmax + }) + shard_size <- shard_size + prod(tensor$shape) * 0.5625 + n_cast <- n_cast + 1L + } else { + shard[[key]] <- tensor + shard_size <- shard_size + prod(tensor$shape) * 2 + } + rm(tensor) + + if (shard_size >= shard_bytes) { + flush_shard() + } + if (i %% 200L == 0L) { + gc(verbose = FALSE) + if (verbose) { + message(sprintf(" quantizing %d/%d tensors", i, length(keys))) + } + } + } + flush_shard() + + manifest <- list( + source = basename(checkpoint_path), + model_version = ckpt$version, + format = "nf4", + shards = shard_files, + tensors = length(keys), + nf4_cast = n_cast, + config = ckpt$config + ) + jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE, pretty = TRUE) + if (verbose) { + message(sprintf("Quantized %d/%d tensors to nf4 across %d shards: %s", + n_cast, length(keys), length(shard_files), output_dir)) + } + invisible(manifest) +} + +#' Load the LTX-2.3 transformer with resident NF4 weights +#' +#' Builds the transformer, swaps the cast-set linears for +#' \code{\link{ltx23_nf4_linear}}, and loads everything onto +#' \code{device}: at ~4.5 bits/parameter the whole 22B transformer stays +#' GPU-resident, avoiding per-step weight transfers. +#' +#' @param ckpt An NF4 \code{ltx23_checkpoint} +#' (\code{\link{ltx23_open_fp8_checkpoint}} reads any shard artifact). +#' @param device Character. +#' @param verbose Logical. +#' @param ... Passed to \code{\link{ltx23_transformer}} (tiny test configs). +#' +#' @return The loaded \code{ltx23_transformer}. +#' +#' @export +ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE, + ...) { + stopifnot(inherits(ckpt, "ltx23_checkpoint")) + model <- ltx23_transformer(...) + model$to(dtype = torch::torch_bfloat16()) + + groups <- ltx23_split_keys(ckpt$keys) + dit_keys <- groups$dit + absmax_keys <- dit_keys[endsWith(dit_keys, ".weight_absmax")] + main_keys <- setdiff(dit_keys, absmax_keys) + + dests <- c(model$named_parameters(), model$named_buffers()) + filled <- character(0) + unmapped <- character(0) + + torch::with_no_grad({ + for (i in seq_along(main_keys)) { + key <- main_keys[[i]] + mapped <- ltx23_map_dit_key(key) + + if (ltx23_is_fp8_cast_key(mapped) && + paste0(key, "_absmax") %in% absmax_keys) { + segments <- strsplit(mapped, ".", fixed = TRUE)[[1]] + parent <- .ltx23_walk_module(model, utils::head(segments, -2L)) + leaf <- segments[length(segments) - 1L] + old <- .ltx23_walk_module(parent, leaf) + if (is.null(old)) { + unmapped <- c(unmapped, key) + next + } + w_shape <- old$weight$shape + nf4_mod <- ltx23_nf4_linear(w_shape[1], w_shape[2], + bias = !is.null(old$bias)) + if (!is.null(old$bias)) { + # Adopt the original bias parameter; its checkpoint key + # loads through the pre-swap destination map + nf4_mod$bias <- old$bias + } + nf4_mod$set_nf4_weight( + ckpt$handle$get_tensor(key), + ckpt$handle$get_tensor(paste0(key, "_absmax")) + ) + do.call(`$<-`, list(parent, leaf, nf4_mod)) + filled <- c(filled, mapped) + } else { + dest <- dests[[mapped]] + if (is.null(dest)) { + unmapped <- c(unmapped, key) + next + } + dest$copy_(ckpt$handle$get_tensor(key)) + filled <- c(filled, mapped) + } + + if (i %% 100L == 0L) { + gc(verbose = FALSE) + if (verbose && i %% 500L == 0L) { + message(sprintf(" loaded %d/%d transformer tensors", i, + length(main_keys))) + } + } + } + }) + gc(verbose = FALSE) + + if (length(unmapped)) { + stop("NF4 transformer load: ", length(unmapped), " unmapped keys, e.g. ", + paste(utils::head(unmapped, 3), collapse = ", ")) + } + expected_missing <- grepl(.ltx23_fp8_cast_pattern, names(dests)) + unfilled <- setdiff(names(dests)[!expected_missing], filled) + if (length(unfilled)) { + stop("NF4 transformer load: ", length(unfilled), " unfilled params, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + + # Everything (packed weights included, as buffers) onto the GPU + model$to(device = device) + model$eval() + # Block intermediates (norms, projections, FF activations) are still + # ~1.5GB per block at high resolution; per-block gc keeps them bounded + options(diffuseR.block_gc = TRUE) + if (verbose) { + message("Transformer ready: NF4 weights resident on ", device) + } + model +} diff --git a/R/post_quant_conv.R b/R/post_quant_conv.R index 91bd7b2..4618e52 100644 --- a/R/post_quant_conv.R +++ b/R/post_quant_conv.R @@ -11,26 +11,21 @@ #' @param device Device on which the tensor is located (e.g., "cpu" or "cuda"). #' @return Processed tensor after applying the quantized convolution. #' @export -post_quant_conv <- function( - x, - dtype, - device -) { - params_path <- system.file("post_quant_conv/", package = "diffuseR") - qc_weights <- as.matrix(utils::read.csv(paste0(params_path, - "/post_quant_conv_weights.csv"), - header = FALSE)) - qc_bias <- as.numeric(utils::read.csv(paste0(params_path, - "/post_quant_conv_bias.csv"), - header = FALSE) [[1]]) +post_quant_conv <- function(x, dtype, device) { + params_path <- system.file("post_quant_conv/", package = "diffuseR") + qc_weights <- as.matrix(utils::read.csv(paste0(params_path, + "/post_quant_conv_weights.csv"), + header = FALSE)) + qc_bias <- as.numeric(utils::read.csv(paste0(params_path, + "/post_quant_conv_bias.csv"), + header = FALSE)[[1]]) - # Convert to torch tensors and reshape weights for conv2d - # Move weights to same device/dtype as input before convolution - qc_weights_tensor <- torch::torch_tensor(qc_weights)$view(c(4, 4, 1, 1)) - qc_weights_tensor <- qc_weights_tensor$to(dtype = dtype, device = torch::torch_device(device)) - qc_bias_tensor <- torch::torch_tensor(qc_bias) - qc_bias_tensor <- qc_bias_tensor$to(dtype = dtype, device = torch::torch_device(device)) + # Convert to torch tensors and reshape weights for conv2d + # Move weights to same device/dtype as input before convolution + qc_weights_tensor <- torch::torch_tensor(qc_weights)$view(c(4, 4, 1, 1)) + qc_weights_tensor <- qc_weights_tensor$to(dtype = dtype, device = torch::torch_device(device)) + qc_bias_tensor <- torch::torch_tensor(qc_bias) + qc_bias_tensor <- qc_bias_tensor$to(dtype = dtype, device = torch::torch_device(device)) - torch::nnf_conv2d(x, weight = qc_weights_tensor, bias = qc_bias_tensor) + torch::nnf_conv2d(x, weight = qc_weights_tensor, bias = qc_bias_tensor) } - diff --git a/R/preprocess_image.R b/R/preprocess_image.R index 8a74858..08a8f3b 100644 --- a/R/preprocess_image.R +++ b/R/preprocess_image.R @@ -1,5 +1,3 @@ - - #' Preprocess image for Stable Diffusion #' #' @param input File path to .jpg or .png, or a 3D array @@ -9,61 +7,53 @@ #' #' @return Torch tensor of shape c(1, 3, 512, 512), scaled to c(-1, 1) #' @export -preprocess_image <- function( - input, - device = "cpu", - width = 512, - height = 512 -) { - # Load JPEG/PNG if path - if (is.character(input)) { - if (grepl("\\.jpg$|\\.jpeg$", input, ignore.case = TRUE)) { - img <- jpeg::readJPEG(input) - } else if (grepl("\\.png$", input, ignore.case = TRUE)) { - img <- png::readPNG(input) +preprocess_image <- function(input, device = "cpu", width = 512, height = 512) { + # Load JPEG/PNG if path + if (is.character(input)) { + if (grepl("\\.jpg$|\\.jpeg$", input, ignore.case = TRUE)) { + img <- jpeg::readJPEG(input) + } else if (grepl("\\.png$", input, ignore.case = TRUE)) { + img <- png::readPNG(input) + } else { + stop("Unsupported image format: only .jpg/.jpeg/.png allowed") + } + } else if (is.array(input)) { + img <- input } else { - stop("Unsupported image format: only .jpg/.jpeg/.png allowed") + stop("Input must be a file path or 3D array") } - } else if (is.array(input)) { - img <- input - } else { - stop("Input must be a file path or 3D array") - } - # Handle alpha channel if present (RGBA -> RGB) - if (length(dim(img)) == 3 && dim(img)[3] == 4) { - img <- img[,, 1:3]# Drop alpha channel - } + # Handle alpha channel if present (RGBA -> RGB) + if (length(dim(img)) == 3 && dim(img)[3] == 4) { + img <- img[,, 1:3] # Drop alpha channel + } - # Handle grayscale images (add channels dimension) - if (length(dim(img)) == 2) { - img <- array(img, dim = c(dim(img), 1)) - # Optionally repeat to make it RGB: - img <- array(rep(img, 3), dim = c(dim(img)[1], dim(img)[2], 3)) - } + # Handle grayscale images (add channels dimension) + if (length(dim(img)) == 2) { + img <- array(img, dim = c(dim(img), 1)) + # Optionally repeat to make it RGB: + img <- array(rep(img, 3), dim = c(dim(img)[1], dim(img)[2], 3)) + } - # Convert to torch tensor - img_tensor <- torch::torch_tensor(img) + # Convert to torch tensor + img_tensor <- torch::torch_tensor(img) - # Add batch dimension and rearrange to [batch, channel, height, width] - if (length(dim(img_tensor)) == 3) { - # [H, W, C] -> [1, C, H, W] - img_tensor <- img_tensor$permute(c(3, 1, 2))$unsqueeze(1) - } else if (length(dim(img_tensor)) == 4) { - # Assume [B, H, W, C] -> [B, C, H, W] - img_tensor <- img_tensor$permute(c(1, 4, 2, 3)) - } + # Add batch dimension and rearrange to [batch, channel, height, width] + if (length(dim(img_tensor)) == 3) { + # [H, W, C] -> [1, C, H, W] + img_tensor <- img_tensor$permute(c(3, 1, 2))$unsqueeze(1) + } else if (length(dim(img_tensor)) == 4) { + # Assume [B, H, W, C] -> [B, C, H, W] + img_tensor <- img_tensor$permute(c(1, 4, 2, 3)) + } - # Use torch's interpolation for resizing - img_tensor <- torch::nnf_interpolate(img_tensor, - size = c(height, width), - mode = "bilinear", - align_corners = FALSE) # Usually FALSE for stable diffusion + # Use torch's interpolation for resizing + img_tensor <- torch::nnf_interpolate(img_tensor, size = c(height, width), + mode = "bilinear", align_corners = FALSE) # Usually FALSE for stable diffusion - # Convert to float and normalize to [-1, 1] - img_tensor <- img_tensor$to(dtype = torch::torch_float())$to(device = device) - img_tensor <- (img_tensor * 2) - 1 + # Convert to float and normalize to [-1, 1] + img_tensor <- img_tensor$to(dtype = torch::torch_float())$to(device = device) + img_tensor <- (img_tensor * 2) - 1 - return(img_tensor) + return(img_tensor) } - diff --git a/R/quant_conv.R b/R/quant_conv.R index 9cc3914..23903f0 100644 --- a/R/quant_conv.R +++ b/R/quant_conv.R @@ -11,24 +11,19 @@ #' @param device Device on which the tensor is located (e.g., "cpu" or "cuda"). #' @return Processed tensor after applying the quantized convolution. #' @export -quant_conv <- function( - x, - dtype, - device -) { - params_path <- system.file("quant_conv/", package = "diffuseR") - qc_weights <- as.matrix(utils::read.csv(paste0(params_path, - "/quant_conv_weights.csv"), - header = FALSE)) - qc_bias <- as.numeric(utils::read.csv(paste0(params_path, - "/quant_conv_bias.csv"), - header = FALSE) [[1]]) +quant_conv <- function(x, dtype, device) { + params_path <- system.file("quant_conv/", package = "diffuseR") + qc_weights <- as.matrix(utils::read.csv(paste0(params_path, + "/quant_conv_weights.csv"), + header = FALSE)) + qc_bias <- as.numeric(utils::read.csv(paste0(params_path, + "/quant_conv_bias.csv"), + header = FALSE)[[1]]) - # Convert to torch tensors and reshape weights for conv2d - qc_weights_tensor <- torch::torch_tensor(qc_weights)$view(c(8, 8, 1, 1)) - qc_bias_tensor <- torch::torch_tensor(qc_bias) + # Convert to torch tensors and reshape weights for conv2d + qc_weights_tensor <- torch::torch_tensor(qc_weights)$view(c(8, 8, 1, 1)) + qc_bias_tensor <- torch::torch_tensor(qc_bias) - conv <- torch::nnf_conv2d(x, weight = qc_weights_tensor, bias = qc_bias_tensor) - conv$to(dtype = dtype, device = torch::torch_device(device)) + conv <- torch::nnf_conv2d(x, weight = qc_weights_tensor, bias = qc_bias_tensor) + conv$to(dtype = dtype, device = torch::torch_device(device)) } - diff --git a/R/rope.R b/R/rope.R deleted file mode 100644 index 925ecf9..0000000 --- a/R/rope.R +++ /dev/null @@ -1,360 +0,0 @@ -#' Rotary Position Embeddings (RoPE) for LTX-2 Video Models -#' -#' Implementation of rotary positional embeddings for 3D video (spatiotemporal) -#' coordinates as used in LTX-2. RoPE encodes position information directly -#' into the attention queries and keys without adding position embeddings. -#' -#' @name rope -#' @references -#' Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). -#' "RoFormer: Enhanced Transformer with Rotary Position Embedding." -#' \url{https://arxiv.org/abs/2104.09864} -NULL - -#' Create RoPE position embedder for video -#' -#' Creates a RoPE embedder configured for video generation, handling -#' spatiotemporal coordinates (frames, height, width). -#' -#' @param dim Integer. Dimension for RoPE (typically attention head dim * num_heads). -#' @param patch_size Integer. Spatial patch size. Default: 1 -#' @param patch_size_t Integer. Temporal patch size. Default: 1 -#' @param base_num_frames Integer. Base number of frames for normalization. Default: 20 -#' @param base_height Integer. Base height for normalization. Default: 2048 -#' @param base_width Integer. Base width for normalization. Default: 2048 -#' @param scale_factors Numeric vector of length 3. VAE scale factors for -#' (temporal, height, width). Default: c(8, 32, 32) -#' @param theta Numeric. Base frequency for RoPE. Default: 10000.0 -#' @param causal_offset Integer. Offset for causal VAE modeling. Default: 1 -#' @param double_precision Logical. Whether to use float64 for frequency -#' computation. Default: TRUE -#' @param rope_type Character. Type of RoPE: "interleaved" or "split". Default: "interleaved" -#' @param num_attention_heads Integer. Number of attention heads (for split RoPE). Default: 32 -#' -#' @return A list containing RoPE configuration and methods. -#' @export -rope_embedder_create <- function( - dim, - patch_size = 1L, - patch_size_t = 1L, - base_num_frames = 20L, - base_height = 2048L, - base_width = 2048L, - scale_factors = c(8, 32, 32), - theta = 10000.0, - causal_offset = 1L, - double_precision = TRUE, - rope_type = c("interleaved", "split"), - num_attention_heads = 32L -) { - rope_type <- match.arg(rope_type) - - list( - dim = dim, - patch_size = patch_size, - patch_size_t = patch_size_t, - base_num_frames = base_num_frames, - base_height = base_height, - base_width = base_width, - scale_factors = scale_factors, - theta = theta, - causal_offset = causal_offset, - double_precision = double_precision, - rope_type = rope_type, - num_attention_heads = num_attention_heads - ) -} - -#' Prepare video coordinates for RoPE -#' -#' Creates per-dimension patch boundaries for video coordinates in pixel space. -#' Returns tensor of shape (batch_size, 3, num_patches, 2) where dimension 1 -#' represents (frame, height, width) and dimension 3 represents (start, end). -#' -#' @param embedder List. RoPE embedder configuration. -#' @param batch_size Integer. Batch size. -#' @param num_frames Integer. Number of latent frames. -#' @param height Integer. Latent height. -#' @param width Integer. Latent width. -#' @param device Character or torch device. Device for tensors. -#' @param fps Numeric. Video frames per second. Default: 24.0 -#' -#' @return torch tensor of shape (batch_size, 3, num_patches, 2). -#' @export -rope_prepare_video_coords <- function( - embedder, - batch_size, - num_frames, - height, - width, - device = "cpu", - fps = 24.0 -) { - patch_size <- embedder$patch_size - patch_size_t <- embedder$patch_size_t - scale_factors <- embedder$scale_factors - causal_offset <- embedder$causal_offset - - # 1. Generate grid coordinates for each spatiotemporal dimension - # Note: R's torch_arange is inclusive, so use end-step to match Python's exclusive behavior - grid_f <- torch::torch_arange( - start = 0, end = num_frames - patch_size_t, step = patch_size_t, - dtype = torch::torch_float32(), device = device - ) - grid_h <- torch::torch_arange( - start = 0, end = height - patch_size, step = patch_size, - dtype = torch::torch_float32(), device = device - ) - grid_w <- torch::torch_arange( - start = 0, end = width - patch_size, step = patch_size, - dtype = torch::torch_float32(), device = device - ) - - # Create meshgrid (ij indexing keeps order as frames, height, width) - grid <- torch::torch_meshgrid(list(grid_f, grid_h, grid_w), indexing = "ij") - grid <- torch::torch_stack(grid, dim = 1) # [3, N_F, N_H, N_W] - - # 2. Get patch boundaries with respect to latent grid - patch_size_vec <- c(patch_size_t, patch_size, patch_size) - patch_size_delta <- torch::torch_tensor( - patch_size_vec, dtype = torch::torch_float32(), device = device - )$view(c(3, 1, 1, 1)) - - patch_ends <- grid + patch_size_delta - - # Combine start and end coordinates [3, N_F, N_H, N_W, 2] - latent_coords <- torch::torch_stack(list(grid, patch_ends), dim = - 1) - - # Reshape to (1, 3, num_patches, 2) - latent_coords <- torch::torch_flatten(latent_coords, start_dim = 2, end_dim = 4) # [3, num_patches, 2] - latent_coords <- latent_coords$unsqueeze(1) # [1, 3, num_patches, 2] - latent_coords <- latent_coords$`repeat`(c(batch_size, 1, 1, 1)) - - # 3. Convert to pixel space using VAE scale factors - scale_tensor <- torch::torch_tensor( - scale_factors, dtype = torch::torch_float32(), device = device - )$view(c(1, 3, 1, 1)) - - pixel_coords <- latent_coords * scale_tensor - - # Handle causal offset for first frame - # Frame coordinates need special handling for causal VAE - frame_coords <- pixel_coords[, 1,,] - frame_coords <- (frame_coords + causal_offset - scale_factors[1])$clamp(min = 0) - pixel_coords[, 1,,] <- frame_coords - - # Scale temporal coordinates by FPS (convert to seconds) - pixel_coords[, 1,,] <- pixel_coords[, 1,,] / fps - - pixel_coords -} - -#' Compute RoPE frequencies from coordinates -#' -#' Converts spatiotemporal coordinates to (cos, sin) frequency tensors -#' for applying rotary embeddings. -#' -#' @param embedder List. RoPE embedder configuration. -#' @param coords torch tensor. Coordinate tensor from rope_prepare_video_coords(). -#' @param device Character or torch device. Device for output tensors. -#' -#' @return A list with: -#' \describe{ -#' \item{cos_freqs}{Cosine frequencies tensor} -#' \item{sin_freqs}{Sine frequencies tensor} -#' } -#' @export -rope_forward <- function( - embedder, - coords, - device = NULL -) { - if (is.null(device)) { - device <- coords$device - } - - dim <- embedder$dim - theta <- embedder$theta - base_num_frames <- embedder$base_num_frames - base_height <- embedder$base_height - base_width <- embedder$base_width - double_precision <- embedder$double_precision - rope_type <- embedder$rope_type - num_attention_heads <- embedder$num_attention_heads - - # Number of spatiotemporal dimensions (3 for video) - num_pos_dims <- as.numeric(coords$shape[2]) - - # 1. If coords are patch boundaries [start, end), use midpoint - if (length(coords$shape) == 4) { - # Split into start and end - coords_split <- coords$chunk(2, dim = - 1) - coords_start <- coords_split[[1]] - coords_end <- coords_split[[2]] - coords <- (coords_start + coords_end) / 2.0 - coords <- coords$squeeze(- 1) # [B, num_pos_dims, num_patches] - } - - # 2. Get coordinates as fraction of base shape - max_positions <- c(base_num_frames, base_height, base_width) - - # Create grid tensor [B, num_patches, num_pos_dims] - grid_list <- list() - for (i in seq_len(num_pos_dims)) { - grid_list[[i]] <- coords[, i,] / max_positions[i] - } - grid <- torch::torch_stack(grid_list, dim = - 1)$to(device = device) - - # Number of RoPE elements (3 dims * 2 for cos/sin) - num_rope_elems <- num_pos_dims * 2L - - # 3. Create 1D grid of frequencies - if (double_precision) { - freqs_dtype <- torch::torch_float64() - } else { - freqs_dtype <- torch::torch_float32() - } - - pow_indices <- torch::torch_pow( - theta, - torch::torch_linspace( - start = 0.0, end = 1.0, - steps = as.integer(dim %/% num_rope_elems), - dtype = freqs_dtype, device = device - ) - ) - freqs <- (pow_indices * pi / 2.0)$to(dtype = torch::torch_float32()) - - # 4. Compute position-specific frequencies - # [B, num_patches, num_pos_dims, dim // num_rope_elems] - freqs <- (grid$unsqueeze(- 1) * 2 - 1) * freqs - - # Transpose and flatten [B, num_patches, dim // 2] - freqs <- freqs$transpose(- 1, - 2)$flatten(start_dim = 3) - - # 5. Get interleaved (cos, sin) frequencies - if (rope_type == "interleaved") { - cos_freqs <- freqs$cos()$repeat_interleave(2L, dim = - 1L) - sin_freqs <- freqs$sin()$repeat_interleave(2L, dim = - 1L) - - # Handle padding if dim not divisible by num_rope_elems - if (dim %% num_rope_elems != 0) { - pad_size <- dim %% num_rope_elems - cos_padding <- torch::torch_ones_like(cos_freqs[,, 1:pad_size]) - sin_padding <- torch::torch_zeros_like(sin_freqs[,, 1:pad_size]) - cos_freqs <- torch::torch_cat(list(cos_padding, cos_freqs), dim = - 1) - sin_freqs <- torch::torch_cat(list(sin_padding, sin_freqs), dim = - 1) - } - } else if (rope_type == "split") { - expected_freqs <- dim %/% 2L - current_freqs <- as.numeric(freqs$shape[length(freqs$shape)]) - pad_size <- expected_freqs - current_freqs - - cos_freq <- freqs$cos() - sin_freq <- freqs$sin() - - if (pad_size > 0) { - cos_padding <- torch::torch_ones_like(cos_freq[,, 1:pad_size]) - sin_padding <- torch::torch_zeros_like(sin_freq[,, 1:pad_size]) - cos_freq <- torch::torch_cat(list(cos_padding, cos_freq), dim = - 1) - sin_freq <- torch::torch_cat(list(sin_padding, sin_freq), dim = - 1) - } - - # Reshape for multi-head attention - b <- as.numeric(cos_freq$shape[1]) - t <- as.numeric(cos_freq$shape[2]) - - cos_freq <- cos_freq$reshape(c(b, t, num_attention_heads, - 1)) - sin_freq <- sin_freq$reshape(c(b, t, num_attention_heads, - 1)) - - cos_freqs <- cos_freq$swapaxes(2, 3) # (B, H, T, D//2) - sin_freqs <- sin_freq$swapaxes(2, 3) # (B, H, T, D//2) - } - - list(cos_freqs = cos_freqs, sin_freqs = sin_freqs) -} - -#' Apply interleaved rotary embeddings -#' -#' Applies rotary position embeddings to query or key tensors using the -#' interleaved format where real and imaginary components alternate. -#' -#' @param x torch tensor. Query or key tensor of shape (B, S, C). -#' @param freqs List. Contains cos_freqs and sin_freqs from rope_forward(). -#' -#' @return torch tensor. Rotated tensor with same shape as input. -#' @export -apply_interleaved_rotary_emb <- function( - x, - freqs -) { - cos_freqs <- freqs$cos_freqs - sin_freqs <- freqs$sin_freqs - - # Split x into interleaved real/imaginary pairs - # x has shape [B, S, C], unflatten to [B, S, C//2, 2] - x_unflat <- x$unflatten(3, c(- 1, 2)) - x_split <- x_unflat$unbind(- 1) - x_real <- x_split[[1]]# [B, S, C // 2] - x_imag <- x_split[[2]]# [B, S, C // 2] - - # Rotate: stack [-x_imag, x_real] and flatten back - x_rotated <- torch::torch_stack(list(- x_imag, x_real), dim = - 1)$flatten(start_dim = 3) - - # Apply rotation formula: x * cos + rotate(x) * sin - out <- (x$to(dtype = torch::torch_float32()) * cos_freqs + - x_rotated$to(dtype = torch::torch_float32()) * sin_freqs) - out <- out$to(dtype = x$dtype) - - out -} - -#' Apply split rotary embeddings -#' -#' Applies rotary position embeddings to query or key tensors using the -#' split format where first half is real and second half is imaginary. -#' -#' @param x torch tensor. Query or key tensor. -#' @param freqs List. Contains cos_freqs and sin_freqs from rope_forward(). -#' -#' @return torch tensor. Rotated tensor with same shape as input. -#' @export -apply_split_rotary_emb <- function( - x, - freqs -) { - cos_freqs <- freqs$cos_freqs - sin_freqs <- freqs$sin_freqs - - x_dtype <- x$dtype - needs_reshape <- FALSE - - # Handle dimension mismatch - if (length(x$shape) != 4 && length(cos_freqs$shape) == 4) { - b <- as.numeric(x$shape[1]) - cos_shape <- as.numeric(cos_freqs$shape) - h <- cos_shape[2] - t <- cos_shape[3] - x <- x$reshape(c(b, t, h, - 1))$swapaxes(2, 3) - needs_reshape <- TRUE - } - - # Split into real and imaginary halves - half_dim <- as.numeric(x$shape[length(x$shape)]) %/% 2L - x_real <- x[,,, 1:half_dim] - x_imag <- x[,,, (half_dim + 1) :(half_dim * 2)] - - # Rotate: [-x_imag, x_real] - x_rotated <- torch::torch_cat(list(- x_imag, x_real), dim = - 1) - - # Apply rotation - out <- (x$to(dtype = torch::torch_float32()) * cos_freqs + - x_rotated$to(dtype = torch::torch_float32()) * sin_freqs) - - if (needs_reshape) { - out <- out$swapaxes(2, 3)$flatten(start_dim = 2, end_dim = 3) - } - - out$to(dtype = x_dtype) -} - diff --git a/R/rope_ltx23.R b/R/rope_ltx23.R new file mode 100644 index 0000000..64b5e8a --- /dev/null +++ b/R/rope_ltx23.R @@ -0,0 +1,292 @@ +#' LTX-2.3 Rotary Positional Embeddings +#' +#' Fresh R port of the LTX rotary positional embedding scheme from the +#' diffusers reference implementation (Apache-2.0, +#' src/diffusers/models/transformers/transformer_ltx2.py). LTX 2.3 uses +#' the "split" RoPE layout everywhere; "interleaved" is kept for +#' completeness. Frequencies are computed in float64 per the checkpoint +#' config (frequencies_precision) and applied in float32. +#' +#' @name rope_ltx23 +NULL + +#' Apply interleaved rotary embeddings +#' +#' Rotates adjacent element pairs of the last dimension: +#' \code{out = x * cos + rotate_half(x) * sin} with pairs interleaved +#' (elements 1,2 form the first complex pair). +#' +#' @param x Tensor of shape [B, S, C]. +#' @param freqs List of two tensors (cos, sin), each [B, S, C]. +#' +#' @return Tensor with the same shape and dtype as \code{x}. +#' +#' @export +ltx23_apply_interleaved_rotary_emb <- function(x, freqs) { + cos <- freqs[[1]] + sin <- freqs[[2]] + + # Split last dim into (C/2, 2) pairs: real = even positions, imag = odd + pairs <- x$unflatten(3L, c(-1L, 2L)) + x_real <- pairs[,,, 1] # [B, S, C/2] + x_imag <- pairs[,,, 2] + x_rotated <- torch::torch_stack(list(-x_imag, x_real), dim = -1L)$flatten(start_dim = 3L) + + out <- x$to(dtype = torch::torch_float32()) * cos + + x_rotated$to(dtype = torch::torch_float32()) * sin + out$to(dtype = x$dtype) +} + +#' Apply split rotary embeddings +#' +#' Rotates element pairs formed by splitting the last dimension in half: +#' element i pairs with element i + d/2. The cos/sin tensors carry half +#' the head dimension. +#' +#' @param x Tensor of shape [B, H, T, D] (per-head layout), or [B, T, H*D] +#' which is reshaped per-head when \code{freqs} is 4D. +#' @param freqs List of two tensors (cos, sin), each [B, H, T, D/2]. +#' +#' @return Tensor with the same shape and dtype as \code{x}. +#' +#' @export +ltx23_apply_split_rotary_emb <- function(x, freqs) { + cos <- freqs[[1]] + sin <- freqs[[2]] + + x_dtype <- x$dtype + needs_reshape <- FALSE + if (x$ndim != 4L && cos$ndim == 4L) { + # cos is [B, H, T, r] -> reshape x [B, T, H*D] to [B, H, T, D] + b <- cos$shape[1] + h <- cos$shape[2] + t <- cos$shape[3] + x <- x$reshape(c(b, t, h, -1L))$transpose(2L, 3L) + needs_reshape <- TRUE + } + + d <- x$shape[length(x$shape)] + if (d %% 2L != 0L) { + stop("Expected the last dimension of x to be even for split rotary, got ", + d) + } + r <- d %/% 2L + + # First/second halves of the last dim form the rotation pairs + x_first <- x$narrow(-1L, 1L, r)$to(dtype = torch::torch_float32()) + x_second <- x$narrow(-1L, r + 1L, r)$to(dtype = torch::torch_float32()) + + out_first <- x_first * cos - x_second * sin + out_second <- x_second * cos + x_first * sin + out <- torch::torch_cat(list(out_first, out_second), dim = -1L) + + if (needs_reshape) { + out <- out$transpose(2L, 3L)$reshape(c(b, t, -1L)) + } + + out$to(dtype = x_dtype) +} + +#' LTX-2.3 audio/video rotary position embedder +#' +#' Computes RoPE cos/sin frequency tensors from spatiotemporal patch +#' coordinates. Video coordinates are 3D (frames scaled to seconds via +#' fps, height, width in pixel space); audio coordinates are 1D +#' (seconds). Coordinates are patch boundaries [start, end); the midpoint +#' is used as the position. +#' +#' @param dim Integer. Rotary dimension (attention head dim x heads for +#' split type at model level; see reference). +#' @param patch_size,patch_size_t Integers. Spatial/temporal patch sizes. +#' @param base_num_frames,base_height,base_width Integers. Base grid the +#' coordinates are normalized against. +#' @param sampling_rate,hop_length Integers. Audio spectrogram params. +#' @param scale_factors Integer vector. VAE (time, height, width) scale +#' factors. +#' @param theta Numeric. RoPE theta. +#' @param causal_offset Integer. Temporal offset for the causal VAE +#' (first frame has stride 1). +#' @param modality "video" or "audio". +#' @param double_precision Logical. Compute base frequencies in float64. +#' @param rope_type "split" (LTX 2.3) or "interleaved". +#' @param num_attention_heads Integer. Needed for the split layout. +#' +#' @export +ltx23_rotary_pos_embed <- torch::nn_module( + "ltx23_rotary_pos_embed", + initialize = function( + dim, + patch_size = 1L, + patch_size_t = 1L, + base_num_frames = 20L, + base_height = 2048L, + base_width = 2048L, + sampling_rate = 16000L, + hop_length = 160L, + scale_factors = c(8L, 32L, 32L), + theta = 10000.0, + causal_offset = 1L, + modality = "video", + double_precision = TRUE, + rope_type = "split", + num_attention_heads = 32L + ) { + if (!rope_type %in% c("interleaved", "split")) { + stop("rope_type must be 'interleaved' or 'split', got: ", rope_type) + } + if (!modality %in% c("video", "audio")) { + stop("modality must be 'video' or 'audio', got: ", modality) + } + self$dim <- as.integer(dim) + self$patch_size <- as.integer(patch_size) + self$patch_size_t <- as.integer(patch_size_t) + self$base_num_frames <- base_num_frames + self$base_height <- base_height + self$base_width <- base_width + self$sampling_rate <- sampling_rate + self$hop_length <- hop_length + self$scale_factors <- as.integer(scale_factors) + self$theta <- theta + self$causal_offset <- causal_offset + self$modality <- modality + self$double_precision <- double_precision + self$rope_type <- rope_type + self$num_attention_heads <- as.integer(num_attention_heads) +}, + + # Patch boundaries [B, 3, num_patches, 2] in pixel/second space + prepare_video_coords = function(batch_size, num_frames, height, width, + device, fps = 24.0) { + f32 <- torch::torch_float32() + # torch_arange has an inclusive end; end = n - 1 matches Python's + # exclusive arange for integer grids + grid_f <- torch::torch_arange(start = 0, end = num_frames - 1, + step = self$patch_size_t, dtype = f32, + device = device) + grid_h <- torch::torch_arange( + start = 0, end = height - 1, step = self$patch_size, + dtype = f32, device = device + ) + grid_w <- torch::torch_arange( + start = 0, end = width - 1, step = self$patch_size, + dtype = f32, device = device + ) + grid <- torch::torch_meshgrid(list(grid_f, grid_h, grid_w), indexing = "ij") + grid <- torch::torch_stack(grid, dim = 1L) # [3, NF, NH, NW] + + patch_size <- c(self$patch_size_t, self$patch_size, self$patch_size) + patch_delta <- torch::torch_tensor(patch_size, dtype = grid$dtype, device = grid$device) + patch_ends <- grid + patch_delta$view(c(3L, 1L, 1L, 1L)) + + latent_coords <- torch::torch_stack(list(grid, patch_ends), dim = -1L) # [3,NF,NH,NW,2] + latent_coords <- latent_coords$flatten(start_dim = 2L, end_dim = 4L) # [3,N,2] + latent_coords <- latent_coords$unsqueeze(1L)$`repeat`(c(batch_size, 1L, 1L, 1L)) + + scale_tensor <- torch::torch_tensor(self$scale_factors, device = latent_coords$device) + pixel_coords <- latent_coords * scale_tensor$view(c(1L, -1L, 1L, 1L)) + + # First latent frame has temporal stride 1 in the causal VAE: shift and + # clamp so timestamps stay causal and non-negative + pixel_coords[, 1,,] <- (pixel_coords[, 1,,] + self$causal_offset - + self$scale_factors[1])$clamp(min = 0) + # Temporal coordinates in seconds + pixel_coords[, 1,,] <- pixel_coords[, 1,,] / fps + + pixel_coords +}, + + # Patch boundaries [B, 1, num_patches, 2] in seconds + prepare_audio_coords = function(batch_size, num_frames, device, shift = 0L) { + f32 <- torch::torch_float32() + grid_f <- torch::torch_arange( + start = shift, end = num_frames + shift - 1, step = self$patch_size_t, + dtype = f32, device = device + ) + + audio_scale_factor <- self$scale_factors[1] + grid_start_mel <- grid_f * audio_scale_factor + grid_start_mel <- (grid_start_mel + self$causal_offset - audio_scale_factor)$clamp(min = 0) + grid_start_s <- grid_start_mel * self$hop_length / self$sampling_rate + + grid_end_mel <- (grid_f + self$patch_size_t) * audio_scale_factor + grid_end_mel <- (grid_end_mel + self$causal_offset - audio_scale_factor)$clamp(min = 0) + grid_end_s <- grid_end_mel * self$hop_length / self$sampling_rate + + audio_coords <- torch::torch_stack(list(grid_start_s, grid_end_s), dim = -1L) # [N, 2] + audio_coords <- audio_coords$unsqueeze(1L)$expand(c(batch_size, -1L, -1L)) + audio_coords$unsqueeze(2L) # [B, 1, N, 2] +}, + + forward = function(coords, device = NULL) { + device <- device %||% coords$device + num_pos_dims <- coords$shape[2] + + # Patch boundaries [start, end) -> midpoint position + if (coords$ndim == 4L) { + chunks <- torch::torch_chunk(coords, 2L, dim = -1L) + coords <- ((chunks[[1]] + chunks[[2]]) / 2.0)$squeeze(-1L) # [B, dims, N] + } + + max_positions <- if (self$modality == "video") { + c(self$base_num_frames, self$base_height, self$base_width) + } else { + self$base_num_frames + } + # [B, dims, N] -> [B, N, dims], each dim normalized to its base size + grid <- torch::torch_stack( + lapply(seq_len(num_pos_dims), function(i) coords[, i,] / max_positions[i]), + dim = -1L + )$to(device = device) + + num_rope_elems <- num_pos_dims * 2L + + freqs_dtype <- if (self$double_precision) torch::torch_float64() else torch::torch_float32() + pow_indices <- torch::torch_pow( + self$theta, + torch::torch_linspace( + start = 0.0, end = 1.0, steps = self$dim %/% num_rope_elems, + dtype = freqs_dtype, device = device + ) + ) + freqs <- (pow_indices * pi / 2.0)$to(dtype = torch::torch_float32()) + + # Outer product of normalized positions (mapped to [-1, 1]) and freqs: + # [B, N, dims, dim / num_rope_elems] + freqs <- (grid$unsqueeze(-1L) * 2 - 1) * freqs + freqs <- freqs$transpose(-1L, -2L)$flatten(start_dim = 3L) # [B, N, dim/2] + + if (self$rope_type == "interleaved") { + cos_freqs <- freqs$cos()$repeat_interleave(2L, dim = -1L) + sin_freqs <- freqs$sin()$repeat_interleave(2L, dim = -1L) + + pad <- self$dim %% num_rope_elems + if (pad != 0L) { + cos_padding <- torch::torch_ones_like(cos_freqs[,, 1:pad]) + sin_padding <- torch::torch_zeros_like(cos_freqs[,, 1:pad]) + cos_freqs <- torch::torch_cat(list(cos_padding, cos_freqs), dim = -1L) + sin_freqs <- torch::torch_cat(list(sin_padding, sin_freqs), dim = -1L) + } + } else { + expected_freqs <- self$dim %/% 2L + current_freqs <- freqs$shape[length(freqs$shape)] + pad_size <- expected_freqs - current_freqs + cos_freqs <- freqs$cos() + sin_freqs <- freqs$sin() + + if (pad_size != 0L) { + cos_padding <- torch::torch_ones_like(cos_freqs[,, 1:pad_size]) + sin_padding <- torch::torch_zeros_like(sin_freqs[,, 1:pad_size]) + cos_freqs <- torch::torch_cat(list(cos_padding, cos_freqs), dim = -1L) + sin_freqs <- torch::torch_cat(list(sin_padding, sin_freqs), dim = -1L) + } + + # Per-head layout for split application: [B, H, N, dim/(2H)] + b <- cos_freqs$shape[1] + t <- cos_freqs$shape[2] + cos_freqs <- cos_freqs$reshape(c(b, t, self$num_attention_heads, -1L))$transpose(2L, 3L) + sin_freqs <- sin_freqs$reshape(c(b, t, self$num_attention_heads, -1L))$transpose(2L, 3L) + } + + list(cos_freqs, sin_freqs) +} +) diff --git a/R/save_image.R b/R/save_image.R index e646f35..457d857 100644 --- a/R/save_image.R +++ b/R/save_image.R @@ -14,23 +14,18 @@ #' \dontrun{ #' save_image(output_tensor, "sample.png") #' } -save_image <- function( - img, - save_to = "output.png", - normalize = TRUE -) { - # img_array <- tensor2image(img, normalize = normalize) - dims <- dim(img) +save_image <- function(img, save_to = "output.png", normalize = TRUE) { + # img_array <- tensor2image(img, normalize = normalize) + dims <- dim(img) - grDevices::png(filename = save_to, width = dims[2], height = dims[1]) - grid::grid.raster(img) - grDevices::dev.off() - - if (interactive()) { + grDevices::png(filename = save_to, width = dims[2], height = dims[1]) grid::grid.raster(img) - } + grDevices::dev.off() - cat("Image saved to", save_to, "\n") - invisible(save_to) -} + if (interactive()) { + grid::grid.raster(img) + } + cat("Image saved to", save_to, "\n") + invisible(save_to) +} diff --git a/R/save_video.R b/R/save_video.R index 9f4b911..5456364 100644 --- a/R/save_video.R +++ b/R/save_video.R @@ -27,74 +27,67 @@ #' # Save as individual frames #' save_video(video_array, "frames/", format = "frames") #' } -save_video <- function( - video, - file, - fps = 24, - format = NULL, - backend = "auto", - quality = 85, - verbose = TRUE -) { - # Validate input - if (!is.array(video)) { - stop("video must be an array") - } - - dims <- dim(video) - if (length(dims) != 4 || dims[4] != 3) { - stop("video must have shape [T, H, W, 3] (RGB)") - } - - num_frames <- dims[1] - height <- dims[2] - width <- dims[3] - - if (verbose) { - message(sprintf("Video: %d frames, %dx%d", num_frames, width, height)) - } - - # Infer format from extension if not specified - if (is.null(format)) { - ext <- tolower(tools::file_ext(file)) - format <- switch(ext, - "mp4" = "mp4", - "gif" = "gif", - "webm" = "webm", - "png" = "frames", - "frames"# default to frames for directories - ) - if (ext == "" && dir.exists(dirname(file))) { - format <- "frames" +save_video <- function(video, file, fps = 24, format = NULL, + backend = "auto", quality = 85, verbose = TRUE) { + # Validate input + if (!is.array(video)) { + stop("video must be an array") + } + + dims <- dim(video) + if (length(dims) != 4 || dims[4] != 3) { + stop("video must have shape [T, H, W, 3] (RGB)") + } + + num_frames <- dims[1] + height <- dims[2] + width <- dims[3] + + if (verbose) { + message(sprintf("Video: %d frames, %dx%d", num_frames, width, height)) } - } - - # Clamp values to [0, 1] - video <- pmax(pmin(video, 1), 0) - - # Dispatch to backend - if (format == "frames") { - save_frames(video, file, verbose = verbose) - } else if (backend == "auto") { - # Try ffmpeg first, fall back to av - if (.ffmpeg_available()) { - save_video_ffmpeg(video, file, fps = fps, format = format, - quality = quality, verbose = verbose) - } else if (requireNamespace("av", quietly = TRUE)) { - save_video_av(video, file, fps = fps, verbose = verbose) + + # Infer format from extension if not specified + if (is.null(format)) { + ext <- tolower(tools::file_ext(file)) + format <- switch(ext, + "mp4" = "mp4", + "gif" = "gif", + "webm" = "webm", + "png" = "frames", + "frames" # default to frames for directories + ) + if (ext == "" && dir.exists(dirname(file))) { + format <- "frames" + } + } + + # Clamp values to [0, 1] + video <- pmax(pmin(video, 1), 0) + + # Dispatch to backend + if (format == "frames") { + save_frames(video, file, verbose = verbose) + } else if (backend == "auto") { + # Try ffmpeg first, fall back to av + if (.ffmpeg_available()) { + save_video_ffmpeg(video, file, fps = fps, format = format, + quality = quality, verbose = verbose) + } else if (requireNamespace("av", quietly = TRUE)) { + save_video_av(video, file, fps = fps, verbose = verbose) + } else { + stop("No video backend available. Install ffmpeg or the 'av' R package.") + } + } else if (backend == "ffmpeg") { + save_video_ffmpeg(video, file, fps = fps, format = format, + quality = quality, verbose = verbose) + } else if (backend == "av") { + save_video_av(video, file, fps = fps, verbose = verbose) } else { - stop("No video backend available. Install ffmpeg or the 'av' R package.") + stop("Unknown backend: ", backend) } - } else if (backend == "ffmpeg") { - save_video_ffmpeg(video, file, fps = fps, format = format, - quality = quality, verbose = verbose) - } else if (backend == "av") { - save_video_av(video, file, fps = fps, verbose = verbose) - } else { - stop("Unknown backend: ", backend) - } - - invisible(file) + + invisible(file) } #' Save Video Frames as Individual Images @@ -108,44 +101,39 @@ save_video <- function( #' @return Invisibly returns vector of saved file paths. #' #' @export -save_frames <- function( - video, - dir, - prefix = "frame_", - format = "png", - verbose = TRUE -) { - # Create directory if needed - if (!dir.exists(dir)) { - dir.create(dir, recursive = TRUE) - } - - num_frames <- dim(video)[1] - files <- character(num_frames) - - for (i in seq_len(num_frames)) { - frame <- video[i,,,] - filename <- file.path(dir, sprintf("%s%04d.%s", prefix, i, format)) - - # Convert to 8-bit integer array - frame_int <- as.integer(frame * 255) - dim(frame_int) <- dim(frame) - - # Save using png or jpeg - if (format == "png") { - png::writePNG(frame_int / 255, filename) - } else { - jpeg::writeJPEG(frame_int / 255, filename) +save_frames <- function(video, dir, prefix = "frame_", format = "png", + verbose = TRUE) { + # Create directory if needed + if (!dir.exists(dir)) { + dir.create(dir, recursive = TRUE) } - files[i] <- filename - } + num_frames <- dim(video)[1] + files <- character(num_frames) + + for (i in seq_len(num_frames)) { + frame <- video[i,,,] + filename <- file.path(dir, sprintf("%s%04d.%s", prefix, i, format)) - if (verbose) { - message(sprintf("Saved %d frames to %s", num_frames, dir)) - } + # Convert to 8-bit integer array + frame_int <- as.integer(frame * 255) + dim(frame_int) <- dim(frame) - invisible(files) + # Save using png or jpeg + if (format == "png") { + png::writePNG(frame_int / 255, filename) + } else { + jpeg::writeJPEG(frame_int / 255, filename) + } + + files[i] <- filename + } + + if (verbose) { + message(sprintf("Saved %d frames to %s", num_frames, dir)) + } + + invisible(files) } #' Save Video using FFmpeg @@ -158,75 +146,77 @@ save_frames <- function( #' @param verbose Logical. #' #' @keywords internal -save_video_ffmpeg <- function( - video, - file, - fps = 24, - format = "mp4", - quality = 85, - verbose = TRUE -) { - if (!.ffmpeg_available()) { - stop("ffmpeg not found in PATH") - } - - # Create temp directory for frames - temp_dir <- tempfile("frames_") - dir.create(temp_dir) - on.exit(unlink(temp_dir, recursive = TRUE), add = TRUE) - - # Save frames as PNG - if (verbose) message("Writing frames...") - num_frames <- dim(video)[1] - - for (i in seq_len(num_frames)) { - frame <- video[i,,,] - filename <- file.path(temp_dir, sprintf("frame_%04d.png", i)) - png::writePNG(frame, filename) - } - - # Build ffmpeg command - input_pattern <- file.path(temp_dir, "frame_%04d.png") - - # Quality mapping (CRF for H.264: 0-51, lower is better) - crf <- as.integer(51 - (quality / 100) * 51) - - if (format == "mp4") { - cmd <- sprintf( - 'ffmpeg -y -framerate %d -i "%s" -c:v libx264 -crf %d -pix_fmt yuv420p "%s"', - fps, input_pattern, crf, file - ) - } else if (format == "webm") { - cmd <- sprintf( - 'ffmpeg -y -framerate %d -i "%s" -c:v libvpx-vp9 -crf %d -b:v 0 "%s"', - fps, input_pattern, crf, file - ) - } else if (format == "gif") { - # GIF requires palette generation for quality - palette <- file.path(temp_dir, "palette.png") - cmd1 <- sprintf( - 'ffmpeg -y -framerate %d -i "%s" -vf "palettegen=stats_mode=diff" "%s"', - fps, input_pattern, palette - ) - cmd2 <- sprintf( - 'ffmpeg -y -framerate %d -i "%s" -i "%s" -lavfi "paletteuse=dither=bayer" "%s"', - fps, input_pattern, palette, file - ) - if (verbose) message("Generating palette...") - system(cmd1, ignore.stdout = !verbose, ignore.stderr = !verbose) - cmd <- cmd2 - } else { - stop("Unsupported format for ffmpeg: ", format) - } +save_video_ffmpeg <- function(video, file, fps = 24, format = "mp4", + quality = 85, verbose = TRUE) { + if (!.ffmpeg_available()) { + stop("ffmpeg not found in PATH") + } + + # Create temp directory for frames + temp_dir <- tempfile("frames_") + dir.create(temp_dir) + on.exit(unlink(temp_dir, recursive = TRUE), add = TRUE) + + # Save frames as PNG + if (verbose) { + message("Writing frames...") + } + num_frames <- dim(video)[1] + + for (i in seq_len(num_frames)) { + frame <- video[i,,,] + filename <- file.path(temp_dir, sprintf("frame_%04d.png", i)) + png::writePNG(frame, filename) + } + + # Build ffmpeg command + input_pattern <- file.path(temp_dir, "frame_%04d.png") + + # Quality mapping (CRF for H.264: 0-51, lower is better) + crf <- as.integer(51 - (quality / 100) * 51) + + if (format == "mp4") { + cmd <- sprintf( + 'ffmpeg -y -framerate %d -i "%s" -c:v libx264 -crf %d -pix_fmt yuv420p "%s"', + fps, input_pattern, crf, file + ) + } else if (format == "webm") { + cmd <- sprintf( + 'ffmpeg -y -framerate %d -i "%s" -c:v libvpx-vp9 -crf %d -b:v 0 "%s"', + fps, input_pattern, crf, file + ) + } else if (format == "gif") { + # GIF requires palette generation for quality + palette <- file.path(temp_dir, "palette.png") + cmd1 <- sprintf( + 'ffmpeg -y -framerate %d -i "%s" -vf "palettegen=stats_mode=diff" "%s"', + fps, input_pattern, palette + ) + cmd2 <- sprintf( + 'ffmpeg -y -framerate %d -i "%s" -i "%s" -lavfi "paletteuse=dither=bayer" "%s"', + fps, input_pattern, palette, file + ) + if (verbose) { + message("Generating palette...") + } + system(cmd1, ignore.stdout = !verbose, ignore.stderr = !verbose) + cmd <- cmd2 + } else { + stop("Unsupported format for ffmpeg: ", format) + } - if (verbose) message(sprintf("Encoding %s...", format)) - result <- system(cmd, ignore.stdout = !verbose, ignore.stderr = !verbose) + if (verbose) { + message(sprintf("Encoding %s...", format)) + } + result <- system(cmd, ignore.stdout = !verbose, ignore.stderr = !verbose) - if (result != 0) { - stop("ffmpeg encoding failed") - } + if (result != 0) { + stop("ffmpeg encoding failed") + } - if (verbose) message(sprintf("Saved: %s", file)) + if (verbose) { + message(sprintf("Saved: %s", file)) + } } #' Save Video using av Package @@ -237,47 +227,42 @@ save_video_ffmpeg <- function( #' @param verbose Logical. #' #' @keywords internal -save_video_av <- function( - video, - file, - fps = 24, - verbose = TRUE -) { - if (!requireNamespace("av", quietly = TRUE)) { - stop("Package 'av' is required for this backend") - } - - num_frames <- dim(video)[1] - height <- dim(video)[2] - width <- dim(video)[3] - - if (verbose) message("Encoding video with av...") - - # av::av_encode_video expects a function that returns frames - # or a matrix/array. We need to convert our [T,H,W,C] to what av expects. - - # Create temp directory for frames - temp_dir <- tempfile("frames_") - dir.create(temp_dir) - on.exit(unlink(temp_dir, recursive = TRUE), add = TRUE) - - frame_files <- character(num_frames) - for (i in seq_len(num_frames)) { - frame <- video[i,,,] - filename <- file.path(temp_dir, sprintf("frame_%04d.png", i)) - png::writePNG(frame, filename) - frame_files[i] <- filename - } - - # Use av to encode from image files - av::av_encode_video( - input = frame_files, - output = file, - framerate = fps, - verbose = verbose - ) - - if (verbose) message(sprintf("Saved: %s", file)) +save_video_av <- function(video, file, fps = 24, verbose = TRUE) { + if (!requireNamespace("av", quietly = TRUE)) { + stop("Package 'av' is required for this backend") + } + + num_frames <- dim(video)[1] + height <- dim(video)[2] + width <- dim(video)[3] + + if (verbose) { + message("Encoding video with av...") + } + + # av::av_encode_video expects a function that returns frames + # or a matrix/array. We need to convert our [T,H,W,C] to what av expects. + + # Create temp directory for frames + temp_dir <- tempfile("frames_") + dir.create(temp_dir) + on.exit(unlink(temp_dir, recursive = TRUE), add = TRUE) + + frame_files <- character(num_frames) + for (i in seq_len(num_frames)) { + frame <- video[i,,,] + filename <- file.path(temp_dir, sprintf("frame_%04d.png", i)) + png::writePNG(frame, filename) + frame_files[i] <- filename + } + + # Use av to encode from image files + av::av_encode_video(input = frame_files, output = file, framerate = fps, + verbose = verbose) + + if (verbose) { + message(sprintf("Saved: %s", file)) + } } #' Check if FFmpeg is Available @@ -285,11 +270,11 @@ save_video_av <- function( #' @return Logical. TRUE if ffmpeg is in PATH. #' @keywords internal .ffmpeg_available <- function() { - result <- tryCatch( - system2("ffmpeg", "-version", stdout = FALSE, stderr = FALSE), - error = function(e) 1 - ) - result == 0 + result <- tryCatch( + system2("ffmpeg", "-version", stdout = FALSE, stderr = FALSE), + error = function(e) 1 + ) + result == 0 } #' Create Video from Latents (Helper) @@ -305,25 +290,18 @@ save_video_av <- function( #' @return Invisibly returns the output file path. #' #' @export -latents_to_video <- function( - latents, - vae, - file, - fps = 24, - ... -) { - torch::with_no_grad({ - # Decode latents - video_tensor <- vae$decode(latents) - - # Convert to array [T, H, W, C] - video_array <- video_tensor$squeeze(1L)$permute(c(2, 3, 4, 1))$cpu()$numpy() - - # Clamp to [0, 1] - video_array <- pmax(pmin(video_array, 1), 0) +latents_to_video <- function(latents, vae, file, fps = 24, ...) { + torch::with_no_grad({ + # Decode latents + video_tensor <- vae$decode(latents) + + # Convert to array [T, H, W, C] + video_array <- video_tensor$squeeze(1L)$permute(c(2, 3, 4, 1))$cpu()$numpy() + + # Clamp to [0, 1] + video_array <- pmax(pmin(video_array, 1), 0) }) - # Save video - save_video(video_array, file, fps = fps, ...) + # Save video + save_video(video_array, file, fps = fps, ...) } - diff --git a/R/staging_ltx23.R b/R/staging_ltx23.R new file mode 100644 index 0000000..c1999d5 --- /dev/null +++ b/R/staging_ltx23.R @@ -0,0 +1,76 @@ +#' Pinned Staging for Phase-Sequential Components +#' +#' Phase offloading moves each large component (transformer, +#' connectors, VAEs, vocoder) between CPU and GPU every render. From +#' pageable memory those copies run through the driver's bounce +#' buffer at a fraction of PCIe speed; page-locked (pinned) host +#' memory transfers by DMA at full rate. Each component's parameters +#' and buffers are pinned once at load; onload swaps every tensor to +#' a non-blocking GPU copy of its pinned source, and offload simply +#' re-points the tensors at the still-valid pinned copies — weights +#' are immutable during inference, so offload moves no bytes at all. +#' +#' Costs: the pinned copies are non-swappable host RAM for the life +#' of the pipeline (about the model's CPU footprint), and page-locking +#' ~20GB adds ~30s to pipeline load. Measured per render the win is +#' small (~2.7s: warm pinned onload 0.54s vs pageable 0.99s, offload +#' 0.06s vs 2.36s) because the real phase-transition cost was +#' allocator pool regrowth, fixed separately by the loader's pool +#' pre-warm and by not emptying the CUDA cache between phases. Off by +#' default; enable for long multi-render sessions with +#' \code{options(diffuseR.pin_staging = TRUE)} before +#' \code{ltx23_load_pipeline} (breaks even after ~11 renders). +#' +#' @name staging_ltx23 +NULL + +#' Pin a component's tensors for fast phase transfer +#' +#' @param module An nn_module on the CPU. +#' +#' @return A list of \code{list(live, pinned)} tensor pairs, or NULL +#' if pinning is unavailable (no CUDA, or page-locking failed). +#' +#' @keywords internal +.ltx23_pin_component <- function(module) { + if (!torch::cuda_is_available()) { + return(NULL) + } + cuda_dev <- torch::torch_device("cuda") + tryCatch({ + tensors <- c(module$parameters, module$buffers) + suppressWarnings(lapply(tensors, function(p) { + # The device argument is deprecated upstream but this + # torch build requires it + pinned <- p$pin_memory(device = cuda_dev) + p$set_data(pinned) + list(live = p, pinned = pinned) + })) + }, error = function(e) NULL) +} + +#' Move a pinned component onto the compute device +#' +#' Non-blocking copies from pinned memory share the default stream, +#' so later kernels are ordered after them; no explicit sync needed. +#' +#' @keywords internal +.ltx23_staged_onload <- function(staging, device) { + for (pair in staging) { + pair$live$set_data(pair$pinned$to(device = device, non_blocking = TRUE)) + } + invisible(NULL) +} + +#' Return a pinned component to the CPU +#' +#' Weights are immutable during inference, so the pinned host copies +#' are still current: offload is a pointer swap, no transfer. +#' +#' @keywords internal +.ltx23_staged_offload <- function(staging) { + for (pair in staging) { + pair$live$set_data(pair$pinned) + } + invisible(NULL) +} diff --git a/R/text_encoder.R b/R/text_encoder.R index ad40c91..6cd04e3 100644 --- a/R/text_encoder.R +++ b/R/text_encoder.R @@ -4,7 +4,7 @@ #' @param x Input tensor #' @keywords internal quick_gelu <- function(x) { - x * torch::torch_sigmoid(1.702 * x) + x * torch::torch_sigmoid(1.702 * x) } #' CLIP MLP Block @@ -15,30 +15,30 @@ quick_gelu <- function(x) { #' @param gelu_type GELU variant: "tanh" (tanh approximation), "quick" (QuickGELU), "exact" (standard GELU) #' @keywords internal CLIPMLP <- torch::nn_module( - "CLIPMLP", + "CLIPMLP", - initialize = function( - in_dim, - hidden_dim, - gelu_type = "tanh" - ) { + initialize = function( + in_dim, + hidden_dim, + gelu_type = "tanh" + ) { self$fc1 <- torch::nn_linear(in_dim, hidden_dim) self$fc2 <- torch::nn_linear(hidden_dim, in_dim) self$gelu_type <- gelu_type - }, +}, - forward = function(x) { + forward = function(x) { x <- self$fc1(x) if (self$gelu_type == "quick") { - x <- quick_gelu(x) + x <- quick_gelu(x) } else if (self$gelu_type == "tanh") { - x <- torch::nnf_gelu(x, approximate = "tanh") + x <- torch::nnf_gelu(x, approximate = "tanh") } else { - x <- torch::nnf_gelu(x) + x <- torch::nnf_gelu(x) } x <- self$fc2(x) x - } +} ) #' CLIP Attention Block @@ -48,28 +48,28 @@ CLIPMLP <- torch::nn_module( #' @param num_heads Number of attention heads #' @keywords internal CLIPAttention <- torch::nn_module( - "CLIPAttention", + "CLIPAttention", - initialize = function( - embed_dim, - num_heads - ) { + initialize = function( + embed_dim, + num_heads + ) { self$embed_dim <- embed_dim self$num_heads <- num_heads self$head_dim <- embed_dim %/% num_heads - self$scale <- self$head_dim ^ (- 0.5) + self$scale <- self$head_dim ^ (-0.5) # Separate Q/K/V projections (HuggingFace CLIPTextModel style) self$q_proj <- torch::nn_linear(embed_dim, embed_dim) self$k_proj <- torch::nn_linear(embed_dim, embed_dim) self$v_proj <- torch::nn_linear(embed_dim, embed_dim) self$out_proj <- torch::nn_linear(embed_dim, embed_dim) - }, +}, - forward = function( - x, - causal_mask = TRUE - ) { + forward = function( + x, + causal_mask = TRUE + ) { batch_size <- x$shape[1] seq_len <- x$shape[2] @@ -79,20 +79,23 @@ CLIPAttention <- torch::nn_module( v <- self$v_proj(x) # Reshape for multi-head attention - q <- q$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, 3) - k <- k$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, 3) - v <- v$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, 3) + q <- q$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, + 3) + k <- k$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, + 3) + v <- v$view(c(batch_size, seq_len, self$num_heads, self$head_dim))$transpose(2, + 3) # Scaled dot-product attention attn_weights <- torch::torch_matmul(q, k$transpose(3, 4)) * self$scale # Apply causal mask if (causal_mask) { - mask <- torch::torch_ones(seq_len, seq_len, device = x$device)$triu(diagonal = 1)$bool() - attn_weights <- attn_weights$masked_fill(mask, - Inf) + mask <- torch::torch_ones(seq_len, seq_len, device = x$device)$triu(diagonal = 1)$bool() + attn_weights <- attn_weights$masked_fill(mask, -Inf) } - attn_weights <- torch::nnf_softmax(attn_weights, dim = - 1) + attn_weights <- torch::nnf_softmax(attn_weights, dim = -1) attn_output <- torch::torch_matmul(attn_weights, v) # Reshape back @@ -100,7 +103,7 @@ CLIPAttention <- torch::nn_module( attn_output <- self$out_proj(attn_output) attn_output - } +} ) #' CLIP Transformer Block @@ -112,21 +115,21 @@ CLIPAttention <- torch::nn_module( #' @param gelu_type GELU variant: "tanh", "quick", or "exact" #' @keywords internal CLIPTransformerBlock <- torch::nn_module( - "CLIPTransformerBlock", - - initialize = function( - embed_dim, - num_heads, - mlp_dim, - gelu_type = "tanh" - ) { + "CLIPTransformerBlock", + + initialize = function( + embed_dim, + num_heads, + mlp_dim, + gelu_type = "tanh" + ) { self$attention <- CLIPAttention(embed_dim, num_heads) self$layernorm_1 <- torch::nn_layer_norm(embed_dim) self$mlp <- CLIPMLP(embed_dim, mlp_dim, gelu_type) self$layernorm_2 <- torch::nn_layer_norm(embed_dim) - }, +}, - forward = function(x) { + forward = function(x) { # Pre-norm attention: norm -> attn -> residual residual <- x x <- self$layernorm_1(x) @@ -138,7 +141,7 @@ CLIPTransformerBlock <- torch::nn_module( x <- residual + self$mlp(x) x - } +} ) #' Native CLIP Text Encoder @@ -158,17 +161,17 @@ CLIPTransformerBlock <- torch::nn_module( #' @return An nn_module representing the text encoder #' @export text_encoder_native <- torch::nn_module( - "TextEncoderNative", - - initialize = function( - vocab_size = 49408, - context_length = 77, - embed_dim = 768, - num_layers = 12, - num_heads = 12, - mlp_dim = 3072, - apply_final_ln = TRUE - ) { + "TextEncoderNative", + + initialize = function( + vocab_size = 49408, + context_length = 77, + embed_dim = 768, + num_layers = 12, + num_heads = 12, + mlp_dim = 3072, + apply_final_ln = TRUE + ) { self$context_length <- context_length self$embed_dim <- embed_dim self$num_layers <- num_layers @@ -177,22 +180,23 @@ text_encoder_native <- torch::nn_module( # Embeddings self$token_embedding <- torch::nn_embedding(vocab_size, embed_dim) self$position_embedding <- torch::nn_parameter( - torch::torch_zeros(context_length, embed_dim) + torch::torch_zeros(context_length, embed_dim) ) # Transformer blocks - tanh GELU approximation matches TorchScript export best self$transformer_blocks <- torch::nn_module_list() for (i in seq_len(num_layers)) { - self$transformer_blocks$append( - CLIPTransformerBlock(embed_dim, num_heads, mlp_dim, gelu_type = "tanh") - ) + self$transformer_blocks$append( + CLIPTransformerBlock(embed_dim, num_heads, mlp_dim, + gelu_type = "tanh") + ) } # Final layer norm self$final_layer_norm <- torch::nn_layer_norm(embed_dim) - }, +}, - forward = function(input_ids) { + forward = function(input_ids) { # Move input to model's device input_ids <- input_ids$to(device = self$token_embedding$weight$device) @@ -202,21 +206,21 @@ text_encoder_native <- torch::nn_module( # Token + position embeddings # Add 1 because R torch is 1-indexed but tokens are 0-indexed (Python convention) token_embeds <- self$token_embedding(input_ids + 1L) - pos_embeds <- self$position_embedding[1:seq_length,]$unsqueeze(1)$expand(c(batch_size, - 1, - 1)) + pos_embeds <- self$position_embedding[1:seq_length,]$unsqueeze(1)$expand(c(batch_size, -1, -1)) hidden_states <- token_embeds + pos_embeds # Transformer layers for (i in seq_len(self$num_layers)) { - hidden_states <- self$transformer_blocks[[i]](hidden_states) + hidden_states <- self$transformer_blocks[[i]](hidden_states) } # Conditionally apply final layer norm to match TorchScript behavior if (self$apply_final_ln) { - hidden_states <- self$final_layer_norm(hidden_states) + hidden_states <- self$final_layer_norm(hidden_states) } hidden_states - } +} ) #' Detect text encoder architecture from TorchScript file @@ -225,64 +229,65 @@ text_encoder_native <- torch::nn_module( #' @return List with vocab_size, context_length, embed_dim, num_layers, num_heads, mlp_dim #' @keywords internal detect_text_encoder_architecture <- function(torchscript_path) { - ts_encoder <- torch::jit_load(torchscript_path) - ts_params <- ts_encoder$parameters - param_names <- names(ts_params) - - # Detect prefix style (SD21 vs SDXL) - if (any(grepl("^text_encoder\\.", param_names))) { - prefix <- "text_encoder.text_model." - } else if (any(grepl("^enc\\.", param_names))) { - prefix <- "enc.text_model." - } else { - stop("Unknown TorchScript parameter prefix") - } - - # Get dimensions from embeddings - tok_emb <- ts_params[[paste0(prefix, "embeddings.token_embedding.weight")]] - pos_emb <- ts_params[[paste0(prefix, "embeddings.position_embedding.weight")]] - fc1 <- ts_params[[paste0(prefix, "encoder.layers.0.mlp.fc1.weight")]] - - vocab_size <- as.integer(tok_emb$shape[1]) - embed_dim <- as.integer(tok_emb$shape[2]) - context_length <- as.integer(pos_emb$shape[1]) - mlp_dim <- as.integer(fc1$shape[1]) - - # Count layers - layer_params <- grep("encoder\\.layers\\.[0-9]+\\.", param_names, value = TRUE) - layer_nums <- unique(as.integer(gsub(".*encoder\\.layers\\.([0-9]+)\\..*", "\\1", layer_params))) - num_layers <- length(layer_nums) - - # Infer heads from embed_dim (typical head_dim is 64) - num_heads <- embed_dim %/% 64L - - # Detect if final layer norm is applied in TorchScript output - # by checking output range with test input - tokens <- torch::torch_tensor(matrix(c(49406, 320, 4380, 49407, rep(49407, 73)), nrow = 1), - dtype = torch::torch_long()) - torch::with_no_grad({ - test_out <- ts_encoder(tokens) + ts_encoder <- torch::jit_load(torchscript_path) + ts_params <- ts_encoder$parameters + param_names <- names(ts_params) + + # Detect prefix style (SD21 vs SDXL) + if (any(grepl("^text_encoder\\.", param_names))) { + prefix <- "text_encoder.text_model." + } else if (any(grepl("^enc\\.", param_names))) { + prefix <- "enc.text_model." + } else { + stop("Unknown TorchScript parameter prefix") + } + + # Get dimensions from embeddings + tok_emb <- ts_params[[paste0(prefix, "embeddings.token_embedding.weight")]] + pos_emb <- ts_params[[paste0(prefix, + "embeddings.position_embedding.weight")]] + fc1 <- ts_params[[paste0(prefix, "encoder.layers.0.mlp.fc1.weight")]] + + vocab_size <- as.integer(tok_emb$shape[1]) + embed_dim <- as.integer(tok_emb$shape[2]) + context_length <- as.integer(pos_emb$shape[1]) + mlp_dim <- as.integer(fc1$shape[1]) + + # Count layers + layer_params <- grep("encoder\\.layers\\.[0-9]+\\.", param_names, value = TRUE) + layer_nums <- unique(as.integer(gsub(".*encoder\\.layers\\.([0-9]+)\\..*", "\\1", layer_params))) + num_layers <- length(layer_nums) + + # Infer heads from embed_dim (typical head_dim is 64) + num_heads <- embed_dim %/% 64L + + # Detect if final layer norm is applied in TorchScript output + # by checking output range with test input + tokens <- torch::torch_tensor(matrix(c(49406, 320, 4380, 49407, rep(49407, 73)), nrow = 1), + dtype = torch::torch_long()) + torch::with_no_grad({ + test_out <- ts_encoder(tokens) }) - # Handle both single tensor (text_encoder) and list output (text_encoder2) - if (is.list(test_out)) { - hidden_states <- test_out[[1]] - } else { - hidden_states <- test_out - } - # If output range is large (>100), final LN is NOT applied - max_abs <- max(abs(as.numeric(hidden_states$min())), abs(as.numeric(hidden_states$max()))) - apply_final_ln <- max_abs < 100 - - list( - vocab_size = vocab_size, - context_length = context_length, - embed_dim = embed_dim, - num_layers = num_layers, - num_heads = num_heads, - mlp_dim = mlp_dim, - prefix = prefix, - apply_final_ln = apply_final_ln - ) + # Handle both single tensor (text_encoder) and list output (text_encoder2) + if (is.list(test_out)) { + hidden_states <- test_out[[1]] + } else { + hidden_states <- test_out + } + # If output range is large (>100), final LN is NOT applied + max_abs <- max(abs(as.numeric(hidden_states$min())), abs(as.numeric(hidden_states$max()))) + apply_final_ln <- max_abs < 100 + + list( + vocab_size = vocab_size, + context_length = context_length, + embed_dim = embed_dim, + num_layers = num_layers, + num_heads = num_heads, + mlp_dim = mlp_dim, + prefix = prefix, + apply_final_ln = apply_final_ln + ) } #' Load weights from TorchScript text encoder into native encoder @@ -293,93 +298,95 @@ detect_text_encoder_architecture <- function(torchscript_path) { #' #' @return The native encoder with loaded weights (invisibly) #' @export -load_text_encoder_weights <- function( - native_encoder, - torchscript_path, - verbose = TRUE -) { - ts_encoder <- torch::jit_load(torchscript_path) - ts_params <- ts_encoder$parameters - param_names <- names(ts_params) - - # Build mapping from TorchScript names to native names - # TorchScript format: enc.text_model.encoder.layers.0.self_attn.q_proj.weight - # or: text_encoder.text_model.encoder.layers.0.self_attn.q_proj.weight - # Native format: transformer_blocks.1.attention.q_proj.weight - remap_key <- function(key) { - # Strip text_encoder. prefix (SD21 style) - key <- sub("^text_encoder\\.", "", key) - - # Strip enc. prefix (SDXL style) - key <- sub("^enc\\.", "", key) - - # Strip text_model. prefix - key <- sub("^text_model\\.", "", key) +load_text_encoder_weights <- function(native_encoder, torchscript_path, + verbose = TRUE) { + ts_encoder <- torch::jit_load(torchscript_path) + ts_params <- ts_encoder$parameters + param_names <- names(ts_params) - # Embeddings - key <- sub("^embeddings\\.token_embedding\\.", "token_embedding.", key) - # position_embedding.weight -> position_embedding (it's an nn_parameter, not nn_embedding) - key <- sub("^embeddings\\.position_embedding\\.weight$", "position_embedding", key) + # Build mapping from TorchScript names to native names + # TorchScript format: enc.text_model.encoder.layers.0.self_attn.q_proj.weight + # or: text_encoder.text_model.encoder.layers.0.self_attn.q_proj.weight + # Native format: transformer_blocks.1.attention.q_proj.weight + remap_key <- function(key) { + # Strip text_encoder. prefix (SD21 style) + key <- sub("^text_encoder\\.", "", key) - # Encoder layers -> transformer_blocks (keep 0-indexed) - key <- gsub("^encoder\\.layers\\.", "transformer_blocks.", key) + # Strip enc. prefix (SDXL style) + key <- sub("^enc\\.", "", key) - # self_attn -> attention - key <- gsub("\\.self_attn\\.", ".attention.", key) + # Strip text_model. prefix + key <- sub("^text_model\\.", "", key) - # Layer norms - key <- gsub("\\.layer_norm1\\.", ".layernorm_1.", key) - key <- gsub("\\.layer_norm2\\.", ".layernorm_2.", key) + # Embeddings + key <- sub("^embeddings\\.token_embedding\\.", "token_embedding.", key) + # position_embedding.weight -> position_embedding (it's an nn_parameter, not nn_embedding) + key <- sub("^embeddings\\.position_embedding\\.weight$", + "position_embedding", key) - # Final layer norm - key <- sub("^final_layer_norm\\.", "final_layer_norm.", key) - - # MLP (HuggingFace uses fc1/fc2 already) - # No change needed - - key - } - - loaded <- 0 - skipped <- 0 - unmapped <- character(0) - - torch::with_no_grad({ - for (ts_name in names(ts_params)) { - native_name <- remap_key(ts_name) - - if (native_name %in% names(native_encoder$parameters)) { - ts_tensor <- ts_params[[ts_name]] - native_tensor <- native_encoder$parameters[[native_name]] - - if (all(ts_tensor$shape == native_tensor$shape)) { - native_tensor$copy_(ts_tensor) - loaded <- loaded + 1 - } else if (verbose) { - message("Shape mismatch: ", native_name, - " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", - paste(as.integer(native_tensor$shape), collapse = "x"), ")") - skipped <- skipped + 1 - } - } else { - skipped <- skipped + 1 - unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + # Encoder layers -> transformer_blocks (keep 0-indexed) + key <- gsub("^encoder\\.layers\\.", "transformer_blocks.", key) + + # self_attn -> attention + key <- gsub("\\.self_attn\\.", ".attention.", key) + + # Layer norms + key <- gsub("\\.layer_norm1\\.", ".layernorm_1.", key) + key <- gsub("\\.layer_norm2\\.", ".layernorm_2.", key) + + # Final layer norm + key <- sub("^final_layer_norm\\.", "final_layer_norm.", key) + + # MLP (HuggingFace uses fc1/fc2 already) + # No change needed + + key + } + + loaded <- 0 + skipped <- 0 + unmapped <- character(0) + + torch::with_no_grad({ + for (ts_name in names(ts_params)) { + native_name <- remap_key(ts_name) + + if (native_name %in% names(native_encoder$parameters)) { + ts_tensor <- ts_params[[ts_name]] + native_tensor <- native_encoder$parameters[[native_name]] + + if (all(ts_tensor$shape == native_tensor$shape)) { + native_tensor$copy_(ts_tensor) + loaded <- loaded + 1 + } else if (verbose) { + message("Shape mismatch: ", native_name, + " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", + paste(as.integer(native_tensor$shape), collapse = "x"), ")") + skipped <- skipped + 1 + } + } else { + skipped <- skipped + 1 + unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + } } - } }) - if (verbose) { - if (length(unmapped) > 0 && length(unmapped) <= 10) { - message("Unmapped parameters:") - for (u in unmapped) message(" ", u) - } else if (length(unmapped) > 10) { - message("Unmapped: ", length(unmapped), " parameters (showing first 5)") - for (u in head(unmapped, 5)) message(" ", u) + if (verbose) { + if (length(unmapped) > 0 && length(unmapped) <= 10) { + message("Unmapped parameters:") + for (u in unmapped) { + message(" ", u) + } + } else if (length(unmapped) > 10) { + message("Unmapped: ", length(unmapped), " parameters (showing first 5)") + for (u in head(unmapped, 5)) { + message(" ", u) + } + } + message("Loaded ", loaded, "/", loaded + skipped, " parameters") } - message("Loaded ", loaded, "/", loaded + skipped, " parameters") - } - invisible(native_encoder) + invisible(native_encoder) } #' Native CLIP Text Encoder 2 (OpenCLIP ViT-bigG for SDXL) @@ -397,16 +404,16 @@ load_text_encoder_weights <- function( #' @return An nn_module representing the text encoder #' @export text_encoder2_native <- torch::nn_module( - "TextEncoder2Native", - - initialize = function( - vocab_size = 49408, - context_length = 77, - embed_dim = 1280, - num_layers = 32, - num_heads = 20, - mlp_dim = 5120 - ) { + "TextEncoder2Native", + + initialize = function( + vocab_size = 49408, + context_length = 77, + embed_dim = 1280, + num_layers = 32, + num_heads = 20, + mlp_dim = 5120 + ) { self$context_length <- context_length self$embed_dim <- embed_dim self$num_layers <- num_layers @@ -414,15 +421,16 @@ text_encoder2_native <- torch::nn_module( # Embeddings self$token_embedding <- torch::nn_embedding(vocab_size, embed_dim) self$position_embedding <- torch::nn_parameter( - torch::torch_zeros(context_length, embed_dim) + torch::torch_zeros(context_length, embed_dim) ) # Transformer blocks with standard GELU (OpenCLIP style) self$transformer_blocks <- torch::nn_module_list() for (i in seq_len(num_layers)) { - self$transformer_blocks$append( - CLIPTransformerBlock(embed_dim, num_heads, mlp_dim, gelu_type = "exact") - ) + self$transformer_blocks$append( + CLIPTransformerBlock(embed_dim, num_heads, mlp_dim, + gelu_type = "exact") + ) } # Final layer norm @@ -430,9 +438,9 @@ text_encoder2_native <- torch::nn_module( # Text projection for pooled output self$text_projection <- torch::nn_linear(embed_dim, embed_dim, bias = FALSE) - }, +}, - forward = function(input_ids) { + forward = function(input_ids) { # Move input to model's device input_ids <- input_ids$to(device = self$token_embedding$weight$device) @@ -442,12 +450,12 @@ text_encoder2_native <- torch::nn_module( # Token + position embeddings # Add 1 because R torch is 1-indexed but tokens are 0-indexed (Python convention) token_embeds <- self$token_embedding(input_ids + 1L) - pos_embeds <- self$position_embedding[1:seq_length,]$unsqueeze(1)$expand(c(batch_size, - 1, - 1)) + pos_embeds <- self$position_embedding[1:seq_length,]$unsqueeze(1)$expand(c(batch_size, -1, -1)) hidden_states <- token_embeds + pos_embeds # Transformer layers for (i in seq_len(self$num_layers)) { - hidden_states <- self$transformer_blocks[[i]](hidden_states) + hidden_states <- self$transformer_blocks[[i]](hidden_states) } # Note: TorchScript does NOT apply final_layer_norm to hidden_states output @@ -459,8 +467,8 @@ text_encoder2_native <- torch::nn_module( # In practice, SDXL tokenizer puts EOS at the position after the last real token eos_indices <- torch::torch_argmax(input_ids, dim = 2L, keepdim = TRUE) pooled_output <- hidden_states_normalized$gather( - dim = 2L, - index = eos_indices$unsqueeze(- 1L)$expand(c(- 1L, - 1L, self$embed_dim)) + dim = 2L, + index = eos_indices$unsqueeze(-1L)$expand(c(-1L, -1L, self$embed_dim)) )$squeeze(2L) # Apply text projection @@ -468,7 +476,7 @@ text_encoder2_native <- torch::nn_module( # Return hidden_states WITHOUT final LN (matches TorchScript), pooled WITH LN + projection list(hidden_states, pooled_output) - } +} ) #' Load weights from TorchScript text encoder 2 into native encoder @@ -479,92 +487,93 @@ text_encoder2_native <- torch::nn_module( #' #' @return The native encoder with loaded weights (invisibly) #' @export -load_text_encoder2_weights <- function( - native_encoder, - torchscript_path, - verbose = TRUE -) { - ts_encoder <- torch::jit_load(torchscript_path) - ts_params <- ts_encoder$parameters - - # Build mapping from TorchScript names to native names - # TorchScript format: enc.text_model.encoder.layers.0.self_attn.q_proj.weight - # Native format: transformer_blocks.1.attention.q_proj.weight - remap_key <- function(key) { - # Strip enc. prefix - key <- sub("^enc\\.", "", key) - - # text_projection is at enc level, not text_model level - if (grepl("^text_projection", key)) { - return(key) - } +load_text_encoder2_weights <- function(native_encoder, torchscript_path, + verbose = TRUE) { + ts_encoder <- torch::jit_load(torchscript_path) + ts_params <- ts_encoder$parameters + + # Build mapping from TorchScript names to native names + # TorchScript format: enc.text_model.encoder.layers.0.self_attn.q_proj.weight + # Native format: transformer_blocks.1.attention.q_proj.weight + remap_key <- function(key) { + # Strip enc. prefix + key <- sub("^enc\\.", "", key) + + # text_projection is at enc level, not text_model level + if (grepl("^text_projection", key)) { + return(key) + } - # Strip text_model. prefix - key <- sub("^text_model\\.", "", key) + # Strip text_model. prefix + key <- sub("^text_model\\.", "", key) - # Embeddings - key <- sub("^embeddings\\.token_embedding\\.", "token_embedding.", key) - # position_embedding.weight -> position_embedding (it's an nn_parameter, not nn_embedding) - key <- sub("^embeddings\\.position_embedding\\.weight$", "position_embedding", key) + # Embeddings + key <- sub("^embeddings\\.token_embedding\\.", "token_embedding.", key) + # position_embedding.weight -> position_embedding (it's an nn_parameter, not nn_embedding) + key <- sub("^embeddings\\.position_embedding\\.weight$", + "position_embedding", key) - # Encoder layers -> transformer_blocks (keep 0-indexed) - key <- gsub("^encoder\\.layers\\.", "transformer_blocks.", key) + # Encoder layers -> transformer_blocks (keep 0-indexed) + key <- gsub("^encoder\\.layers\\.", "transformer_blocks.", key) - # self_attn -> attention - key <- gsub("\\.self_attn\\.", ".attention.", key) + # self_attn -> attention + key <- gsub("\\.self_attn\\.", ".attention.", key) - # Layer norms - key <- gsub("\\.layer_norm1\\.", ".layernorm_1.", key) - key <- gsub("\\.layer_norm2\\.", ".layernorm_2.", key) + # Layer norms + key <- gsub("\\.layer_norm1\\.", ".layernorm_1.", key) + key <- gsub("\\.layer_norm2\\.", ".layernorm_2.", key) - # Final layer norm - key <- sub("^final_layer_norm\\.", "final_layer_norm.", key) - - # MLP (HuggingFace uses fc1/fc2 already) - # No change needed - - key - } - - loaded <- 0 - skipped <- 0 - unmapped <- character(0) - - torch::with_no_grad({ - for (ts_name in names(ts_params)) { - native_name <- remap_key(ts_name) - - if (native_name %in% names(native_encoder$parameters)) { - ts_tensor <- ts_params[[ts_name]] - native_tensor <- native_encoder$parameters[[native_name]] - - if (all(ts_tensor$shape == native_tensor$shape)) { - native_tensor$copy_(ts_tensor) - loaded <- loaded + 1 - } else if (verbose) { - message("Shape mismatch: ", native_name, - " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", - paste(as.integer(native_tensor$shape), collapse = "x"), ")") - skipped <- skipped + 1 - } - } else { - skipped <- skipped + 1 - unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + # Final layer norm + key <- sub("^final_layer_norm\\.", "final_layer_norm.", key) + + # MLP (HuggingFace uses fc1/fc2 already) + # No change needed + + key + } + + loaded <- 0 + skipped <- 0 + unmapped <- character(0) + + torch::with_no_grad({ + for (ts_name in names(ts_params)) { + native_name <- remap_key(ts_name) + + if (native_name %in% names(native_encoder$parameters)) { + ts_tensor <- ts_params[[ts_name]] + native_tensor <- native_encoder$parameters[[native_name]] + + if (all(ts_tensor$shape == native_tensor$shape)) { + native_tensor$copy_(ts_tensor) + loaded <- loaded + 1 + } else if (verbose) { + message("Shape mismatch: ", native_name, + " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", + paste(as.integer(native_tensor$shape), collapse = "x"), ")") + skipped <- skipped + 1 + } + } else { + skipped <- skipped + 1 + unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + } } - } }) - if (verbose) { - if (length(unmapped) > 0 && length(unmapped) <= 10) { - message("Unmapped parameters:") - for (u in unmapped) message(" ", u) - } else if (length(unmapped) > 10) { - message("Unmapped: ", length(unmapped), " parameters (showing first 5)") - for (u in head(unmapped, 5)) message(" ", u) + if (verbose) { + if (length(unmapped) > 0 && length(unmapped) <= 10) { + message("Unmapped parameters:") + for (u in unmapped) { + message(" ", u) + } + } else if (length(unmapped) > 10) { + message("Unmapped: ", length(unmapped), " parameters (showing first 5)") + for (u in head(unmapped, 5)) { + message(" ", u) + } + } + message("Loaded ", loaded, "/", loaded + skipped, " parameters") } - message("Loaded ", loaded, "/", loaded + skipped, " parameters") - } - invisible(native_encoder) + invisible(native_encoder) } - diff --git a/R/text_encoder_ltx2.R b/R/text_encoder_ltx2.R deleted file mode 100644 index 958e26a..0000000 --- a/R/text_encoder_ltx2.R +++ /dev/null @@ -1,714 +0,0 @@ -# LTX2 Text Encoder and Connectors -# -# This file implements: -# 1. LTX2TextConnectors - Transforms text embeddings for video/audio streams -# 2. Text encoder wrapper - Supports pre-computed embeddings and API-based encoding - -# ----------------------------------------------------------------------------- -# 1D RoPE for text connectors -# ----------------------------------------------------------------------------- - -#' 1D Rotary Position Embeddings for LTX2 Text Connectors -#' @keywords internal -ltx2_rotary_pos_embed_1d <- torch::nn_module( - "LTX2RotaryPosEmbed1d", - initialize = function( - dim, - base_seq_len = 4096L, - theta = 10000.0, - double_precision = TRUE, - rope_type = "interleaved", - num_attention_heads = 32L - ) { - self$dim <- dim - self$base_seq_len <- base_seq_len - self$theta <- theta - self$double_precision <- double_precision - self$rope_type <- rope_type - self$num_attention_heads <- num_attention_heads - }, - - forward = function( - batch_size, - seq_len, - device - ) { - # 1. Get 1D position ids as fractions of base_seq_len - grid_1d <- torch::torch_arange(start = 0, end = seq_len - 1L, - dtype = torch::torch_float32(), device = device) - grid_1d <- grid_1d / self$base_seq_len - grid <- grid_1d$unsqueeze(1L)$`repeat`(c(batch_size, 1L)) # [batch_size, seq_len] - - # 2. Calculate 1D RoPE frequencies - num_rope_elems <- 2L# 1D * 2 (for cos, sin) - if (self$double_precision) { - freqs_dtype <- torch::torch_float64() - } else { - freqs_dtype <- torch::torch_float32() - } - - pow_indices <- torch::torch_pow( - self$theta, - torch::torch_linspace(start = 0.0, end = 1.0, steps = self$dim %/% num_rope_elems, - dtype = freqs_dtype, device = device) - ) - freqs <- (pow_indices * pi / 2.0)$to(dtype = torch::torch_float32()) - - # 3. Outer product: [batch_size, seq_len] x [dim/2] -> [batch_size, seq_len, dim/2] - freqs_outer <- torch::torch_einsum("bs,d->bsd", list(grid, freqs)) - - # 4. Compute cos and sin - cos_freqs <- torch::torch_cos(freqs_outer) - sin_freqs <- torch::torch_sin(freqs_outer) - - # 5. Interleave or split based on rope_type - if (self$rope_type == "interleaved") { - # Repeat each element: [B, S, D/2] -> [B, S, D] - cos_freqs <- cos_freqs$unsqueeze(- 1L)$`repeat`(c(1L, 1L, 1L, 2L))$flatten(start_dim = 3L) - sin_freqs <- sin_freqs$unsqueeze(- 1L)$`repeat`(c(1L, 1L, 1L, 2L))$flatten(start_dim = 3L) - } else { - # Concatenate: [B, S, D/2] -> [B, S, D] - cos_freqs <- torch::torch_cat(list(cos_freqs, cos_freqs), dim = - 1L) - sin_freqs <- torch::torch_cat(list(sin_freqs, sin_freqs), dim = - 1L) - } - - list(cos_freqs, sin_freqs) - } -) - -# ----------------------------------------------------------------------------- -# 1D Transformer Block for Connectors -# ----------------------------------------------------------------------------- - -#' 1D Transformer Block for LTX2 Text Connectors -#' @keywords internal -ltx2_transformer_block_1d <- torch::nn_module( - "LTX2TransformerBlock1d", - initialize = function( - dim, - num_attention_heads, - attention_head_dim, - activation_fn = "gelu-approximate", - eps = 1e-6, - rope_type = "interleaved" - ) { - self$norm1 <- rms_norm(dim, eps = eps) - self$attn1 <- ltx2_attention( - query_dim = dim, - heads = num_attention_heads, - kv_heads = num_attention_heads, - dim_head = attention_head_dim, - rope_type = rope_type - ) - - self$norm2 <- rms_norm(dim, eps = eps) - self$ff <- feed_forward(dim, mult = 4L, activation_fn = activation_fn) - }, - - forward = function( - hidden_states, - attention_mask = NULL, - rotary_emb = NULL - ) { - norm_hidden_states <- self$norm1(hidden_states) - attn_hidden_states <- self$attn1(norm_hidden_states, - attention_mask = attention_mask, - query_rotary_emb = rotary_emb) - hidden_states <- hidden_states + attn_hidden_states - - norm_hidden_states <- self$norm2(hidden_states) - ff_hidden_states <- self$ff(norm_hidden_states) - hidden_states <- hidden_states + ff_hidden_states - - hidden_states - } -) - -# ----------------------------------------------------------------------------- -# Connector Transformer 1D -# ----------------------------------------------------------------------------- - -#' 1D Connector Transformer for LTX2 -#' @keywords internal -ltx2_connector_transformer_1d <- torch::nn_module( - "LTX2ConnectorTransformer1d", - initialize = function( - num_attention_heads = 30L, - attention_head_dim = 128L, - num_layers = 2L, - num_learnable_registers = 128L, - rope_base_seq_len = 4096L, - rope_theta = 10000.0, - rope_double_precision = TRUE, - eps = 1e-6, - causal_temporal_positioning = FALSE, - rope_type = "interleaved" - ) { - self$num_attention_heads <- num_attention_heads - self$inner_dim <- num_attention_heads * attention_head_dim - self$causal_temporal_positioning <- causal_temporal_positioning - self$num_learnable_registers <- num_learnable_registers - - # Learnable registers (replaces padding tokens) - if (!is.null(num_learnable_registers) && num_learnable_registers > 0L) { - init_registers <- torch::torch_rand(c(num_learnable_registers, self$inner_dim)) * 2.0 - 1.0 - self$learnable_registers <- torch::nn_parameter(init_registers) - } else { - self$learnable_registers <- NULL - } - - # 1D RoPE - self$rope <- ltx2_rotary_pos_embed_1d( - dim = self$inner_dim, - base_seq_len = rope_base_seq_len, - theta = rope_theta, - double_precision = rope_double_precision, - rope_type = rope_type, - num_attention_heads = num_attention_heads - ) - - # Transformer blocks - self$transformer_blocks <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { - ltx2_transformer_block_1d( - dim = self$inner_dim, - num_attention_heads = num_attention_heads, - attention_head_dim = attention_head_dim, - rope_type = rope_type - ) - })) - - self$norm_out <- rms_norm(self$inner_dim, eps = eps) - }, - - forward = function( - hidden_states, - attention_mask = NULL, - attn_mask_binarize_threshold = - 9000.0 - ) { - batch_size <- hidden_states$shape[1] - sequence_length <- hidden_states$shape[2] - - # 1. Replace padding with learned registers, if using - if (!is.null(self$learnable_registers)) { - if (sequence_length %% self$num_learnable_registers != 0L) { - stop(sprintf("Sequence length %d must be divisible by num_learnable_registers %d", - sequence_length, self$num_learnable_registers)) - } - - num_register_repeats <- sequence_length %/% self$num_learnable_registers - registers <- self$learnable_registers$`repeat`(c(num_register_repeats, 1L)) # [seq_len, inner_dim] - - # Binarize attention mask - binary_attn_mask <- (attention_mask >= attn_mask_binarize_threshold)$to(dtype = torch::torch_int32()) - if (binary_attn_mask$ndim == 4L) { - binary_attn_mask <- binary_attn_mask$squeeze(2L)$squeeze(2L) # [B, 1, 1, L] -> [B, L] - } - - # Extract non-padded tokens and re-pad with registers - padded_list <- list() - valid_seq_lens <- numeric(batch_size) - - for (i in seq_len(batch_size)) { - mask_i <- binary_attn_mask[i,]$to(dtype = torch::torch_bool()) - hs_i <- hidden_states[i, mask_i,] - valid_len <- as.integer(hs_i$shape[1]) - valid_seq_lens[i] <- valid_len - pad_len <- sequence_length - valid_len - - if (pad_len > 0L) { - # Pad with zeros on the right - hs_i <- torch::nnf_pad(hs_i, c(0L, 0L, 0L, pad_len)) - } - padded_list[[i]] <- hs_i$unsqueeze(1L) - } - - padded_hidden_states <- torch::torch_cat(padded_list, dim = 1L) # [B, L, D] - - # Flip mask along sequence dimension and blend with registers - # In R torch, flip requires a vector for dims - flipped_mask <- torch::torch_flip(binary_attn_mask, c(2L))$unsqueeze(- 1L)$to(dtype = hidden_states$dtype) # [B, L, 1] - # Expand registers to batch dimension for broadcasting - registers_expanded <- registers$unsqueeze(1L) # [L, D] -> [1, L, D] - broadcasts to [B, L, D] - hidden_states <- flipped_mask * padded_hidden_states + (1 - flipped_mask) * registers_expanded - - # Zero out attention mask when using registers - attention_mask <- torch::torch_zeros_like(attention_mask) - } - - # 2. Calculate 1D RoPE - rotary_emb <- self$rope(batch_size, sequence_length, device = hidden_states$device) - - # 3. Run transformer blocks - for (i in seq_along(self$transformer_blocks)) { - block <- self$transformer_blocks[[i]] - hidden_states <- block(hidden_states, attention_mask = attention_mask, rotary_emb = rotary_emb) - } - - hidden_states <- self$norm_out(hidden_states) - - list(hidden_states, attention_mask) - } -) - -# ----------------------------------------------------------------------------- -# Full Text Connectors -# ----------------------------------------------------------------------------- - -#' LTX2 Text Connectors -#' -#' Transforms packed text encoder hidden states for video and audio streams. -#' -#' @param caption_channels Integer. Dimension of caption embeddings (default 3840). -#' @param text_proj_in_factor Integer. Factor for input projection (default 1). -#' @param video_connector_num_attention_heads Integer. Number of attention heads for video connector. -#' @param video_connector_attention_head_dim Integer. Attention head dimension for video. -#' @param video_connector_num_layers Integer. Number of transformer layers for video. -#' @param video_connector_num_learnable_registers Integer. Number of learnable registers for video. -#' @param audio_connector_num_attention_heads Integer. Number of attention heads for audio connector. -#' @param audio_connector_attention_head_dim Integer. Attention head dimension for audio. -#' @param audio_connector_num_layers Integer. Number of transformer layers for audio. -#' @param audio_connector_num_learnable_registers Integer. Number of learnable registers for audio. -#' @param connector_rope_base_seq_len Integer. Base sequence length for RoPE. -#' @param rope_theta Numeric. RoPE theta parameter. -#' @param rope_double_precision Logical. Use double precision for RoPE. -#' @param causal_temporal_positioning Logical. Use causal temporal positioning. -#' @param rope_type Character. RoPE type ("interleaved" or "split"). -#' -#' @return nn_module for text connectors. -#' @export -ltx2_text_connectors <- torch::nn_module( - "LTX2TextConnectors", - initialize = function( - caption_channels = 3840L, - text_proj_in_factor = 49L, - video_connector_num_attention_heads = 30L, - video_connector_attention_head_dim = 128L, - video_connector_num_layers = 2L, - video_connector_num_learnable_registers = NULL, - audio_connector_num_attention_heads = 30L, - audio_connector_attention_head_dim = 128L, - audio_connector_num_layers = 2L, - audio_connector_num_learnable_registers = NULL, - connector_rope_base_seq_len = 4096L, - rope_theta = 10000.0, - rope_double_precision = TRUE, - causal_temporal_positioning = FALSE, - rope_type = "split" - ) { - - self$caption_channels <- caption_channels - - # Input projection (projects packed embeddings to caption_channels) - self$text_proj_in <- torch::nn_linear( - in_features = caption_channels * text_proj_in_factor, - out_features = caption_channels, - bias = FALSE - ) - - # Video connector - self$video_connector <- ltx2_connector_transformer_1d( - num_attention_heads = video_connector_num_attention_heads, - attention_head_dim = video_connector_attention_head_dim, - num_layers = video_connector_num_layers, - num_learnable_registers = video_connector_num_learnable_registers, - rope_base_seq_len = connector_rope_base_seq_len, - rope_theta = rope_theta, - rope_double_precision = rope_double_precision, - causal_temporal_positioning = causal_temporal_positioning, - rope_type = rope_type - ) - - # Audio connector - self$audio_connector <- ltx2_connector_transformer_1d( - num_attention_heads = audio_connector_num_attention_heads, - attention_head_dim = audio_connector_attention_head_dim, - num_layers = audio_connector_num_layers, - num_learnable_registers = audio_connector_num_learnable_registers, - rope_base_seq_len = connector_rope_base_seq_len, - rope_theta = rope_theta, - rope_double_precision = rope_double_precision, - causal_temporal_positioning = causal_temporal_positioning, - rope_type = rope_type - ) - }, - - forward = function( - text_encoder_hidden_states, - attention_mask, - additive_mask = FALSE - ) { - # Convert to additive attention mask if necessary - if (!additive_mask) { - text_dtype <- text_encoder_hidden_states$dtype - attention_mask <- (attention_mask - 1)$reshape(c(attention_mask$shape[1], 1L, - 1L, attention_mask$shape[length(attention_mask$shape)])) - attention_mask <- attention_mask$to(dtype = text_dtype) * torch::torch_finfo(text_dtype)$max - } - - # Project input - text_encoder_hidden_states <- self$text_proj_in(text_encoder_hidden_states) - - # Video connector - video_result <- self$video_connector(text_encoder_hidden_states, attention_mask) - video_text_embedding <- video_result[[1]] - new_attn_mask <- video_result[[2]] - - # Apply attention mask - attn_mask <- (new_attn_mask < 1e-6)$to(dtype = torch::torch_int64()) - attn_mask <- attn_mask$reshape(c(video_text_embedding$shape[1], video_text_embedding$shape[2], 1L)) - video_text_embedding <- video_text_embedding * attn_mask - new_attn_mask <- attn_mask$squeeze(- 1L) - - # Audio connector - audio_result <- self$audio_connector(text_encoder_hidden_states, attention_mask) - audio_text_embedding <- audio_result[[1]] - - list(video_text_embedding, audio_text_embedding, new_attn_mask) - } -) - -# ----------------------------------------------------------------------------- -# Text Encoder Wrapper -# ----------------------------------------------------------------------------- - -#' Encode Text for LTX2 -#' -#' Encodes text prompts for LTX2 video generation. Supports multiple backends: -#' - "gemma3": Native R torch Gemma3 text encoder -#' - "precomputed": Load pre-computed embeddings from file -#' - "api": Call an HTTP API for text encoding -#' - "random": Generate random embeddings (for testing only) -#' -#' @param prompt Character vector of prompts. -#' @param backend Character. Backend to use ("gemma3", "precomputed", "api", "random"). -#' @param model_path Character. Path to Gemma3 model directory (for "gemma3" backend). -#' @param tokenizer_path Character. Path to tokenizer (for "gemma3" backend, defaults to model_path). -#' @param text_encoder Pre-loaded Gemma3 text encoder module (for "gemma3" backend). -#' @param embeddings_file Character. Path to pre-computed embeddings (for "precomputed" backend). -#' @param api_url Character. URL of text encoding API (for "api" backend). -#' @param max_sequence_length Integer. Maximum sequence length (default 1024). -#' @param caption_channels Integer. Caption embedding dimension (default 3840). -#' @param device Character. Device for tensors. -#' @param dtype torch_dtype. Data type for tensors. -#' -#' @return List with prompt_embeds and prompt_attention_mask tensors. -#' @export -encode_text_ltx2 <- function( - prompt, - backend = "random", - model_path = NULL, - tokenizer_path = NULL, - text_encoder = NULL, - embeddings_file = NULL, - api_url = NULL, - max_sequence_length = 1024L, - caption_channels = 3840L, - device = "cpu", - dtype = torch::torch_float32() -) { - - if (is.character(prompt) && length(prompt) == 1) { - prompt <- list(prompt) - } else { - prompt <- as.list(prompt) - } - batch_size <- length(prompt) - - if (backend == "gemma3") { - # Native Gemma3 text encoding - if (identical(dtype, torch::torch_float16())) { - dtype_str <- "float16" - } else { - dtype_str <- "float32" - } - - result <- encode_with_gemma3( - prompts = unlist(prompt), - model = text_encoder %||% model_path, - tokenizer = tokenizer_path %||% model_path, - max_sequence_length = max_sequence_length, - device = device, - dtype = dtype_str, - verbose = FALSE - ) - - prompt_embeds <- result$prompt_embeds$to(dtype = dtype) - prompt_attention_mask <- result$prompt_attention_mask - - } else if (backend == "precomputed") { - if (is.null(embeddings_file)) { - stop("embeddings_file required for precomputed backend") - } - # Load pre-computed embeddings - data <- readRDS(embeddings_file) - prompt_embeds <- torch::torch_tensor(data$embeddings, device = device, dtype = dtype) - prompt_attention_mask <- torch::torch_tensor(data$attention_mask, device = device, dtype = torch::torch_int64()) - - } else if (backend == "api") { - if (is.null(api_url)) { - stop("api_url required for api backend") - } - # Call HTTP API - response <- httr::POST( - api_url, - body = jsonlite::toJSON(list( - prompts = prompt, - max_sequence_length = max_sequence_length - ), auto_unbox = TRUE), - httr::content_type_json() - ) - if (httr::status_code(response) != 200) { - stop("Text encoding API failed: ", httr::content(response, "text")) - } - data <- jsonlite::fromJSON(httr::content(response, "text")) - prompt_embeds <- torch::torch_tensor(data$embeddings, device = device, dtype = dtype) - prompt_attention_mask <- torch::torch_tensor(data$attention_mask, device = device, dtype = torch::torch_int64()) - - } else if (backend == "random") { - # Generate random embeddings (for testing) - # Shape: [batch, seq_len, caption_channels * num_layers] = [B, L, 3840*49] - # This mimics packed Gemma3 output for testing connectors - message("Using random embeddings - for testing only") - packed_dim <- caption_channels * 49L# 49 layers from Gemma3 - prompt_embeds <- torch::torch_randn(c(batch_size, max_sequence_length, packed_dim), - device = device, dtype = dtype) - prompt_attention_mask <- torch::torch_ones(c(batch_size, max_sequence_length), - device = device, dtype = torch::torch_int64()) - - } else { - stop("Unknown backend: ", backend, ". Use 'gemma3', 'precomputed', 'api', or 'random'") - } - - list( - prompt_embeds = prompt_embeds, - prompt_attention_mask = prompt_attention_mask - ) -} - -#' Pack Text Embeddings (Gemma-style) -#' -#' Normalizes and packs text encoder hidden states from multiple layers. -#' This is used when working with raw Gemma outputs. -#' -#' @param text_hidden_states Tensor of shape [batch, seq_len, hidden_dim, num_layers]. -#' @param sequence_lengths Integer vector of valid sequence lengths per batch item. -#' @param padding_side Character. "left" or "right". -#' @param scale_factor Numeric. Scale factor for normalization (default 8). -#' @param eps Numeric. Epsilon for numerical stability. -#' @param device Character. Device for tensors. -#' -#' @return Tensor of shape [batch, seq_len, hidden_dim * num_layers]. -#' @export -pack_text_embeds <- function( - text_hidden_states, - sequence_lengths, - padding_side = "left", - scale_factor = 8, - eps = 1e-6, - device = "cpu" -) { - - dims <- text_hidden_states$shape - batch_size <- dims[1] - seq_len <- dims[2] - hidden_dim <- dims[3] - num_layers <- dims[4] - - original_dtype <- text_hidden_states$dtype - - # Create padding mask - token_indices <- torch::torch_arange(start = 0, end = seq_len - 1L, device = device)$unsqueeze(1L) - sequence_lengths_t <- torch::torch_tensor(sequence_lengths, device = device) - - if (padding_side == "right") { - mask <- token_indices < sequence_lengths_t$unsqueeze(2L) - } else if (padding_side == "left") { - start_indices <- seq_len - sequence_lengths_t$unsqueeze(2L) - mask <- token_indices >= start_indices - } else { - stop("padding_side must be 'left' or 'right'") - } - mask <- mask$unsqueeze(- 1L)$unsqueeze(- 1L) # [B, seq_len, 1, 1] - - # Compute masked mean - masked_states <- text_hidden_states$masked_fill(!mask, 0.0) - num_valid <- (sequence_lengths_t * hidden_dim)$view(c(batch_size, 1L, 1L, 1L)) - masked_mean <- masked_states$sum(dim = c(2L, 3L), keepdim = TRUE) / (num_valid + eps) - - # Compute min/max - x_min <- text_hidden_states$masked_fill(!mask, Inf)$amin(dim = c(2L, 3L), keepdim = TRUE) - x_max <- text_hidden_states$masked_fill(!mask, - Inf)$amax(dim = c(2L, 3L), keepdim = TRUE) - - # Normalize - normalized <- (text_hidden_states - masked_mean) / (x_max - x_min + eps) - normalized <- normalized * scale_factor - - # Flatten layers dimension - normalized <- normalized$flatten(start_dim = 3L) - mask_flat <- mask$squeeze(- 1L)$expand(c(- 1L, - 1L, hidden_dim * num_layers)) - normalized <- normalized$masked_fill(!mask_flat, 0.0) - normalized <- normalized$to(dtype = original_dtype) - - normalized -} - -# ----------------------------------------------------------------------------- -# Weight Loading -# ----------------------------------------------------------------------------- - -#' Load LTX2 Text Connectors from safetensors -#' -#' Load pre-trained LTX2 connector weights from HuggingFace safetensors file. -#' -#' @param weights_path Character. Path to safetensors file. -#' @param config_path Character. Optional path to config.json. -#' @param device Character. Device to load weights to. Default: "cpu" -#' @param dtype Character. Data type ("float32", "float16"). Default: "float32" -#' @param verbose Logical. Print loading progress. Default: TRUE -#' @return Initialized ltx2_text_connectors module -#' @export -load_ltx2_connectors <- function( - weights_path, - config_path = NULL, - device = "cpu", - dtype = "float32", - verbose = TRUE -) { - if (!file.exists(weights_path)) { - stop("Weights file not found: ", weights_path) - } - - # Load config - config <- NULL - # Auto-detect config.json in same directory if not specified - if (is.null(config_path)) { - auto_config <- file.path(dirname(weights_path), "config.json") - if (file.exists(auto_config)) { - config_path <- auto_config - } - } - if (!is.null(config_path) && file.exists(config_path)) { - config <- jsonlite::fromJSON(config_path) - if (verbose) message("Loaded config from: ", config_path) - } - - # Create connectors with config or defaults - if (!is.null(config)) { - connectors <- ltx2_text_connectors( - caption_channels = config$caption_channels %||% 3840L, - text_proj_in_factor = config$text_proj_in_factor %||% 49L, - video_connector_num_attention_heads = config$video_connector_num_attention_heads %||% 30L, - video_connector_attention_head_dim = config$video_connector_attention_head_dim %||% 128L, - video_connector_num_layers = config$video_connector_num_layers %||% 2L, - video_connector_num_learnable_registers = as.integer(config$video_connector_num_learnable_registers), - audio_connector_num_attention_heads = config$audio_connector_num_attention_heads %||% 30L, - audio_connector_attention_head_dim = config$audio_connector_attention_head_dim %||% 128L, - audio_connector_num_layers = config$audio_connector_num_layers %||% 2L, - audio_connector_num_learnable_registers = as.integer(config$audio_connector_num_learnable_registers), - connector_rope_base_seq_len = config$connector_rope_base_seq_len %||% 4096L, - rope_theta = config$rope_theta %||% 10000.0, - rope_double_precision = config$rope_double_precision %||% TRUE, - causal_temporal_positioning = config$causal_temporal_positioning %||% FALSE, - rope_type = config$rope_type %||% "split" - ) - } else { - connectors <- ltx2_text_connectors() - } - - # Load weights - if (verbose) message("Loading weights from: ", weights_path) - weights <- safetensors::safe_load_file(weights_path, framework = "torch") - - load_ltx2_connector_weights(connectors, weights, verbose = verbose) - - # Move to device - torch_dtype <- switch(dtype, - "float32" = torch::torch_float32(), - "float16" = torch::torch_float16(), - "bfloat16" = torch::torch_bfloat16(), - torch::torch_float32() - ) - - connectors$to(device = device, dtype = torch_dtype) - - if (verbose) message("Connectors loaded successfully on device: ", device) - connectors -} - -#' Load weights into LTX2 connectors module -#' -#' @param connectors LTX2 connectors module -#' @param weights Named list of weight tensors -#' @param verbose Print progress -#' @keywords internal -load_ltx2_connector_weights <- function( - connectors, - weights, - verbose = TRUE -) { - native_params <- names(connectors$parameters) - - remap_connector_key <- function(key) { - # HuggingFace uses nn.ModuleList for FeedForward: - # ff.net.0.proj.weight -> ff.act_fn.proj.weight - # ff.net.2.weight -> ff.proj_out.weight - key <- gsub("\\.ff\\.net\\.0\\.", ".ff.act_fn.", key) - key <- gsub("\\.ff\\.net\\.2\\.", ".ff.proj_out.", key) - - # to_out.0 is correct - both HF and our module use ModuleList - key - } - - loaded <- 0L - skipped <- 0L - unmapped <- character(0) - - torch::with_no_grad({ - for (hf_name in names(weights)) { - native_name <- remap_connector_key(hf_name) - - if (native_name %in% native_params) { - hf_tensor <- weights[[hf_name]] - native_tensor <- connectors$parameters[[native_name]] - - if (all(as.integer(hf_tensor$shape) == as.integer(native_tensor$shape))) { - native_tensor$copy_(hf_tensor) - loaded <- loaded + 1L - } else { - if (verbose) { - message("Shape mismatch: ", native_name, - " (HF: ", paste(as.integer(hf_tensor$shape), collapse = "x"), - " vs R: ", paste(as.integer(native_tensor$shape), collapse = "x"), ")") - } - skipped <- skipped + 1L - } - } else { - skipped <- skipped + 1L - unmapped <- c(unmapped, paste0(hf_name, " -> ", native_name)) - } - } - }) - - if (verbose) { - message(sprintf("Connector weights: %d loaded, %d skipped", loaded, skipped)) - if (length(unmapped) > 0 && length(unmapped) <= 20) { - message("Unmapped parameters:") - for (u in unmapped[1:min(20, length(unmapped))]) { - message(" ", u) - } - } - if (length(unmapped) > 20) { - message(" ... and ", length(unmapped) - 20, " more") - } - } - - invisible(list(loaded = loaded, skipped = skipped, unmapped = unmapped)) -} - -# Null-coalescing operator (if not already defined) -if (!exists("%||%", mode = "function")) { - `%||%` <- function( - x, - y - ) if (is.null(x)) y else x -} - diff --git a/R/tokenizer_bpe.R b/R/tokenizer_bpe.R index 9f3fcc8..edaa862 100644 --- a/R/tokenizer_bpe.R +++ b/R/tokenizer_bpe.R @@ -15,110 +15,105 @@ #' @return A bpe_tokenizer object. #' @export bpe_tokenizer <- function(tokenizer_path) { - # Find tokenizer.json - - if (dir.exists(tokenizer_path)) { - json_path <- file.path(tokenizer_path, "tokenizer.json") - } else { - json_path <- tokenizer_path - } - - if (!file.exists(json_path)) { - stop("tokenizer.json not found at: ", json_path) - } - - # Load JSON - config <- jsonlite::fromJSON(json_path, simplifyVector = FALSE) - - # Extract model configuration - model <- config$model - if (is.null(model) || model$type != "BPE") { - stop("Only BPE tokenizers are supported") - } - - # Build vocabulary lookup (token -> id) - vocab <- model$vocab - if (is.null(vocab)) { - stop("No vocabulary found in tokenizer.json") - } - - # Convert vocab list to named vector for fast lookup - vocab_ids <- unlist(vocab) - names(vocab_ids) <- names(vocab) - - # Build reverse vocabulary (id -> token) - id_to_token <- names(vocab_ids) - names(id_to_token) <- as.character(vocab_ids) - - # Parse merges into a priority map - merges <- model$merges - if (is.null(merges)) { - merges <- character(0) - } - - # Create merge priority lookup (pair -> priority) - # Lower priority number = merge first - merge_priority <- seq_along(merges) - names(merge_priority) <- merges - - # Extract special tokens - added_tokens <- config$added_tokens - special_tokens <- list() - if (!is.null(added_tokens)) { - for (tok in added_tokens) { - special_tokens[[tok$content]] <- list( - id = tok$id, - special = tok$special %||% FALSE - ) + # Find tokenizer.json + + if (dir.exists(tokenizer_path)) { + json_path <- file.path(tokenizer_path, "tokenizer.json") + } else { + json_path <- tokenizer_path + } + + if (!file.exists(json_path)) { + stop("tokenizer.json not found at: ", json_path) + } + + # Load JSON + config <- jsonlite::fromJSON(json_path, simplifyVector = FALSE) + + # Extract model configuration + model <- config$model + if (is.null(model) || model$type != "BPE") { + stop("Only BPE tokenizers are supported") + } + + # Build vocabulary lookup (token -> id) + vocab <- model$vocab + if (is.null(vocab)) { + stop("No vocabulary found in tokenizer.json") } - } - - # Extract configuration - byte_fallback <- model$byte_fallback %||% FALSE - fuse_unk <- model$fuse_unk %||% FALSE - unk_token <- model$unk_token - - # Pre-tokenizer configuration - pre_tokenizer <- config$pre_tokenizer - add_prefix_space <- FALSE - if (!is.null(pre_tokenizer)) { - # Check for Metaspace pre-tokenizer (adds space prefix) - if (!is.null(pre_tokenizer$type) && pre_tokenizer$type == "Metaspace") { - add_prefix_space <- pre_tokenizer$add_prefix_space %||% TRUE + + # Convert vocab list to named vector for fast lookup + vocab_ids <- unlist(vocab) + names(vocab_ids) <- names(vocab) + + # Build reverse vocabulary (id -> token) + id_to_token <- names(vocab_ids) + names(id_to_token) <- as.character(vocab_ids) + + # Parse merges into a priority map + merges <- model$merges + if (is.null(merges)) { + merges <- character(0) + } + + # Create merge priority lookup (pair -> priority) + # Lower priority number = merge first + merge_priority <- seq_along(merges) + names(merge_priority) <- merges + + # Extract special tokens + added_tokens <- config$added_tokens + special_tokens <- list() + if (!is.null(added_tokens)) { + for (tok in added_tokens) { + special_tokens[[tok$content]] <- list(id = tok$id, + special = tok$special %||% FALSE) + } + } + + # Extract configuration + byte_fallback <- model$byte_fallback %||% FALSE + fuse_unk <- model$fuse_unk %||% FALSE + unk_token <- model$unk_token + + # Pre-tokenizer configuration + pre_tokenizer <- config$pre_tokenizer + add_prefix_space <- FALSE + if (!is.null(pre_tokenizer)) { + # Check for Metaspace pre-tokenizer (adds space prefix) + if (!is.null(pre_tokenizer$type) && pre_tokenizer$type == "Metaspace") { + add_prefix_space <- pre_tokenizer$add_prefix_space %||% TRUE + } } - } - - structure( - list( - vocab = vocab_ids, - id_to_token = id_to_token, - merges = merges, - merge_priority = merge_priority, - special_tokens = special_tokens, - byte_fallback = byte_fallback, - fuse_unk = fuse_unk, - unk_token = unk_token, - add_prefix_space = add_prefix_space, - vocab_size = length(vocab_ids) - ), - class = "bpe_tokenizer" - ) + + structure( + list( + vocab = vocab_ids, + id_to_token = id_to_token, + merges = merges, + merge_priority = merge_priority, + special_tokens = special_tokens, + byte_fallback = byte_fallback, + fuse_unk = fuse_unk, + unk_token = unk_token, + add_prefix_space = add_prefix_space, + vocab_size = length(vocab_ids) + ), + class = "bpe_tokenizer" + ) } #' Print BPE Tokenizer #' @param x A bpe_tokenizer object. #' @param ... Additional arguments (ignored). #' @export -print.bpe_tokenizer <- function( - x, - ... -) { - cat("BPE Tokenizer\n") - cat(" Vocabulary size:", x$vocab_size, "\n") - cat(" Merge rules:", length(x$merges), "\n") - cat(" Special tokens:", length(x$special_tokens), "\n") - cat(" Byte fallback:", x$byte_fallback, "\n") - invisible(x) +print.bpe_tokenizer <- function(x, ...) { + cat("BPE Tokenizer\n") + cat(" Vocabulary size:", x$vocab_size, "\n") + cat(" Merge rules:", length(x$merges), "\n") + cat(" Special tokens:", length(x$special_tokens), "\n") + cat(" Byte fallback:", x$byte_fallback, "\n") + invisible(x) } # ----------------------------------------------------------------------------- @@ -136,226 +131,216 @@ print.bpe_tokenizer <- function( #' @param return_tensors Character. Return type: "list" or "pt" (torch tensors). #' @return List with input_ids and attention_mask. #' @export -encode_bpe <- function( - tokenizer, - text, - add_special_tokens = TRUE, - max_length = NULL, - padding = "none", - truncation = FALSE, - return_tensors = "list" -) { - if (!inherits(tokenizer, "bpe_tokenizer")) { - stop("tokenizer must be a bpe_tokenizer object") - } - - # Handle single string - if (length(text) == 1) { - text <- list(text) - } - - # Encode each text - encoded <- lapply(text, function(t) { - encode_single(tokenizer, t, add_special_tokens = add_special_tokens) +encode_bpe <- function(tokenizer, text, add_special_tokens = TRUE, + max_length = NULL, padding = "none", + truncation = FALSE, return_tensors = "list") { + if (!inherits(tokenizer, "bpe_tokenizer")) { + stop("tokenizer must be a bpe_tokenizer object") + } + + # Handle single string + if (length(text) == 1) { + text <- list(text) + } + + # Encode each text + encoded <- lapply(text, function(t) { + encode_single(tokenizer, t, add_special_tokens = add_special_tokens) }) - # Get lengths - lengths <- vapply(encoded, length, integer(1)) - max_len <- max(lengths) + # Get lengths + lengths <- vapply(encoded, length, integer(1)) + max_len <- max(lengths) - # Apply max_length constraint - if (!is.null(max_length)) { - if (truncation) { - encoded <- lapply(encoded, function(ids) { - if (length(ids) > max_length) ids[1:max_length] else ids - }) + # Apply max_length constraint + if (!is.null(max_length)) { + if (truncation) { + encoded <- lapply(encoded, function(ids) { + if (length(ids) > max_length) ids[1:max_length] else ids + }) + } + if (padding == "max_length") { + max_len <- max_length + } } - if (padding == "max_length") { - max_len <- max_length + + # Apply padding + if (padding != "none") { + pad_id <- get_pad_id(tokenizer) + encoded <- lapply(encoded, function(ids) { + if (length(ids) < max_len) { + # Left padding (Gemma style) + c(rep(pad_id, max_len - length(ids)), ids) + } else { + ids + } + }) } - } - - # Apply padding - if (padding != "none") { - pad_id <- get_pad_id(tokenizer) - encoded <- lapply(encoded, function(ids) { - if (length(ids) < max_len) { - # Left padding (Gemma style) - c(rep(pad_id, max_len - length(ids)), ids) - } else { - ids - } - }) - } - # Create attention mask - attention_mask <- lapply(encoded, function(ids) { - pad_id <- get_pad_id(tokenizer) - as.integer(ids != pad_id) + # Create attention mask + attention_mask <- lapply(encoded, function(ids) { + pad_id <- get_pad_id(tokenizer) + as.integer(ids != pad_id) }) - # Convert to matrix - input_ids <- do.call(rbind, encoded) - attention_mask <- do.call(rbind, attention_mask) + # Convert to matrix + input_ids <- do.call(rbind, encoded) + attention_mask <- do.call(rbind, attention_mask) - # Convert to tensors if requested - if (return_tensors == "pt") { - input_ids <- torch::torch_tensor(input_ids, dtype = torch::torch_long()) - attention_mask <- torch::torch_tensor(attention_mask, dtype = torch::torch_long()) - } + # Convert to tensors if requested + if (return_tensors == "pt") { + input_ids <- torch::torch_tensor(input_ids, dtype = torch::torch_long()) + attention_mask <- torch::torch_tensor(attention_mask, + dtype = torch::torch_long()) + } - list( - input_ids = input_ids, - attention_mask = attention_mask - ) + list(input_ids = input_ids, attention_mask = attention_mask) } #' Encode a single text string #' @keywords internal -encode_single <- function( - tokenizer, - text, - add_special_tokens = TRUE -) { - # Add prefix space if configured (Gemma style) - if (tokenizer$add_prefix_space && !startsWith(text, " ")) { - text <- paste0(" ", text) - } - - # Pre-tokenize: split on whitespace while keeping track of spaces - # Replace spaces with special character (SentencePiece style) - text <- gsub(" ", "\u2581", text) # LOWER ONE EIGHTH BLOCK - - # Check for special tokens first - for (tok_name in names(tokenizer$special_tokens)) { - if (text == tok_name) { - return(tokenizer$special_tokens[[tok_name]]$id) +encode_single <- function(tokenizer, text, add_special_tokens = TRUE) { + # Add prefix space if configured (Gemma style) + if (tokenizer$add_prefix_space && !startsWith(text, " ")) { + text <- paste0(" ", text) + } + + # Pre-tokenize: split on whitespace while keeping track of spaces + # Replace spaces with special character (SentencePiece style) + text <- gsub(" ", "\u2581", text) # LOWER ONE EIGHTH BLOCK + + # Check for special tokens first + for (tok_name in names(tokenizer$special_tokens)) { + if (text == tok_name) { + return(tokenizer$special_tokens[[tok_name]]$id) + } } - } - # Use greedy longest match tokenization (for large vocabs like Gemma) - # This is more efficient than BPE merging for pre-trained vocabs - ids <- greedy_tokenize(text, tokenizer$vocab, tokenizer$byte_fallback, tokenizer$unk_token) + # Use greedy longest match tokenization (for large vocabs like Gemma) + # This is more efficient than BPE merging for pre-trained vocabs + ids <- greedy_tokenize(text, tokenizer$vocab, tokenizer$byte_fallback, + tokenizer$unk_token) - # Add special tokens - if (add_special_tokens) { - bos_id <- get_bos_id(tokenizer) - if (!is.null(bos_id)) { - ids <- c(bos_id, ids) + # Add special tokens + if (add_special_tokens) { + bos_id <- get_bos_id(tokenizer) + if (!is.null(bos_id)) { + ids <- c(bos_id, ids) + } } - } - ids + ids } #' Greedy longest match tokenization #' @keywords internal -greedy_tokenize <- function( - text, - vocab, - byte_fallback = FALSE, - unk_token = NULL -) { - ids <- integer(0) - i <- 1 - n <- nchar(text) - vocab_names <- names(vocab) - - while (i <= n) { - # Try to match the longest token starting at position i - matched <- FALSE - - # Try decreasing lengths - for (len in min(n - i + 1, 50) :1) { # Cap at 50 chars max - candidate <- substr(text, i, i + len - 1) - - if (candidate %in% vocab_names) { - ids <- c(ids, vocab[[candidate]]) - i <- i + len - matched <- TRUE - break - } - } +greedy_tokenize <- function(text, vocab, byte_fallback = FALSE, + unk_token = NULL) { + ids <- integer(0) + i <- 1 + n <- nchar(text) + vocab_names <- names(vocab) + + while (i <= n) { + # Try to match the longest token starting at position i + matched <- FALSE + + # Try decreasing lengths + for (len in min(n - i + 1, 50):1) { # Cap at 50 chars max + candidate <- substr(text, i, i + len - 1) + + if (candidate %in% vocab_names) { + ids <- c(ids, vocab[[candidate]]) + i <- i + len + matched <- TRUE + break + } + } - if (!matched) { - # No match found - use byte fallback or UNK - char <- substr(text, i, i) - if (byte_fallback) { - # Try single character in vocab first - if (char %in% vocab_names) { - ids <- c(ids, vocab[[char]]) - } else { - # Byte fallback: convert to <0xHH> format - bytes <- charToRaw(char) - for (b in bytes) { - byte_token <- sprintf("<0x%02X>", as.integer(b)) - if (byte_token %in% vocab_names) { - ids <- c(ids, vocab[[byte_token]]) + if (!matched) { + # No match found - use byte fallback or UNK + char <- substr(text, i, i) + if (byte_fallback) { + # Try single character in vocab first + if (char %in% vocab_names) { + ids <- c(ids, vocab[[char]]) + } else { + # Byte fallback: convert to <0xHH> format + bytes <- charToRaw(char) + for (b in bytes) { + byte_token <- sprintf("<0x%02X>", as.integer(b)) + if (byte_token %in% vocab_names) { + ids <- c(ids, vocab[[byte_token]]) + } else if (!is.null(unk_token) && + unk_token %in% vocab_names) { + ids <- c(ids, vocab[[unk_token]]) + } else { + ids <- c(ids, 3L) # Default UNK + } + } + } } else if (!is.null(unk_token) && unk_token %in% vocab_names) { - ids <- c(ids, vocab[[unk_token]]) + ids <- c(ids, vocab[[unk_token]]) } else { - ids <- c(ids, 3L) # Default UNK + ids <- c(ids, 3L) # Default UNK } - } + i <- i + 1 } - } else if (!is.null(unk_token) && unk_token %in% vocab_names) { - ids <- c(ids, vocab[[unk_token]]) - } else { - ids <- c(ids, 3L) # Default UNK - } - i <- i + 1 } - } - ids + ids } #' Apply BPE merge rules #' @keywords internal -apply_bpe_merges <- function( - tokens, - merge_priority, - vocab -) { - if (length(tokens) <= 1) { - return(tokens) - } - - # Iteratively merge token pairs - changed <- TRUE - while (changed && length(tokens) > 1) { - changed <- FALSE - best_idx <- NULL - best_priority <- Inf - - # Find the highest priority merge - for (i in seq_len(length(tokens) - 1)) { - pair <- paste(tokens[i], tokens[i + 1]) - if (pair %in% names(merge_priority)) { - priority <- merge_priority[[pair]] - if (priority < best_priority) { - best_priority <- priority - best_idx <- i - } - } +apply_bpe_merges <- function(tokens, merge_priority, vocab) { + if (length(tokens) <= 1) { + return(tokens) } - # Apply the best merge - if (!is.null(best_idx)) { - merged <- paste0(tokens[best_idx], tokens[best_idx + 1]) - # Only merge if result is in vocabulary - if (merged %in% names(vocab)) { - tokens <- c( - if (best_idx > 1) tokens[1:(best_idx - 1)] else character(0), - merged, - if (best_idx + 2 <= length(tokens)) tokens[(best_idx + 2) :length(tokens)] else character(0) - ) - changed <- TRUE - } + # Iteratively merge token pairs + changed <- TRUE + while (changed && length(tokens) > 1) { + changed <- FALSE + best_idx <- NULL + best_priority <- Inf + + # Find the highest priority merge + for (i in seq_len(length(tokens) - 1)) { + pair <- paste(tokens[i], tokens[i + 1]) + if (pair %in% names(merge_priority)) { + priority <- merge_priority[[pair]] + if (priority < best_priority) { + best_priority <- priority + best_idx <- i + } + } + } + + # Apply the best merge + if (!is.null(best_idx)) { + merged <- paste0(tokens[best_idx], tokens[best_idx + 1]) + # Only merge if result is in vocabulary + if (merged %in% names(vocab)) { + tokens <- c( + if (best_idx > 1) { + tokens[1:(best_idx - 1)] + } else { + character(0) + }, + merged, + if (best_idx + 2 <= length(tokens)) { + tokens[(best_idx + 2):length(tokens)] + } else { + character(0) + } + ) + changed <- TRUE + } + } } - } - tokens + tokens } # ----------------------------------------------------------------------------- @@ -369,72 +354,69 @@ apply_bpe_merges <- function( #' @param skip_special_tokens Logical. Skip special tokens in output. #' @return Character string or vector. #' @export -decode_bpe <- function( - tokenizer, - ids, - skip_special_tokens = TRUE -) { - if (!inherits(tokenizer, "bpe_tokenizer")) { - stop("tokenizer must be a bpe_tokenizer object") - } - - # Handle matrix input - if (is.matrix(ids)) { - return(apply(ids, 1, function(row) { - decode_bpe(tokenizer, row, skip_special_tokens = skip_special_tokens) +decode_bpe <- function(tokenizer, ids, skip_special_tokens = TRUE) { + if (!inherits(tokenizer, "bpe_tokenizer")) { + stop("tokenizer must be a bpe_tokenizer object") + } + + # Handle matrix input + if (is.matrix(ids)) { + return(apply(ids, 1, function(row) { + decode_bpe(tokenizer, row, + skip_special_tokens = skip_special_tokens) })) - } - - # Handle torch tensor - if (inherits(ids, "torch_tensor")) { - ids <- as.integer(ids$cpu()$numpy()) - } - - # Get special token IDs to skip - special_ids <- integer(0) - if (skip_special_tokens) { - special_ids <- vapply(tokenizer$special_tokens, function(tok) tok$id, integer(1)) - } - - # Convert IDs to tokens - tokens <- vapply(ids, function(id) { - if (id %in% special_ids) { - "" - } else { - id_str <- as.character(id) - if (id_str %in% names(tokenizer$id_to_token)) { - tokenizer$id_to_token[[id_str]] + } + + # Handle torch tensor + if (inherits(ids, "torch_tensor")) { + ids <- as.integer(ids$cpu()$numpy()) + } + + # Get special token IDs to skip + special_ids <- integer(0) + if (skip_special_tokens) { + special_ids <- vapply(tokenizer$special_tokens, function(tok) tok$id, integer(1)) + } + + # Convert IDs to tokens + tokens <- vapply(ids, function(id) { + if (id %in% special_ids) { + "" } else { - "" + id_str <- as.character(id) + if (id_str %in% names(tokenizer$id_to_token)) { + tokenizer$id_to_token[[id_str]] + } else { + "" + } } - } }, character(1)) - # Join tokens and decode - text <- paste(tokens, collapse = "") - - # Replace SentencePiece space marker with actual space - text <- gsub("\u2581", " ", text) - - # Decode byte tokens like <0x41> - # Find all byte tokens and replace them - byte_pattern <- "<0x([0-9A-Fa-f]{2})>" - while (grepl(byte_pattern, text)) { - match <- regmatches(text, regexpr(byte_pattern, text)) - if (length(match) > 0) { - hex_str <- sub("<0x([0-9A-Fa-f]{2})>", "\\1", match) - byte_val <- strtoi(hex_str, base = 16) - char <- rawToChar(as.raw(byte_val)) - text <- sub(byte_pattern, char, text, fixed = FALSE) - } else { - break + # Join tokens and decode + text <- paste(tokens, collapse = "") + + # Replace SentencePiece space marker with actual space + text <- gsub("\u2581", " ", text) + + # Decode byte tokens like <0x41> + # Find all byte tokens and replace them + byte_pattern <- "<0x([0-9A-Fa-f]{2})>" + while (grepl(byte_pattern, text)) { + match <- regmatches(text, regexpr(byte_pattern, text)) + if (length(match) > 0) { + hex_str <- sub("<0x([0-9A-Fa-f]{2})>", "\\1", match) + byte_val <- strtoi(hex_str, base = 16) + char <- rawToChar(as.raw(byte_val)) + text <- sub(byte_pattern, char, text, fixed = FALSE) + } else { + break + } } - } - # Trim leading space if prefix was added - text <- sub("^ ", "", text) + # Trim leading space if prefix was added + text <- sub("^ ", "", text) - text + text } # ----------------------------------------------------------------------------- @@ -444,37 +426,37 @@ decode_bpe <- function( #' Get padding token ID #' @keywords internal get_pad_id <- function(tokenizer) { - if ("" %in% names(tokenizer$special_tokens)) { - tokenizer$special_tokens[[""]]$id - } else if ("" %in% names(tokenizer$vocab)) { - tokenizer$vocab[[""]] - } else { - 0L - } + if ("" %in% names(tokenizer$special_tokens)) { + tokenizer$special_tokens[[""]]$id + } else if ("" %in% names(tokenizer$vocab)) { + tokenizer$vocab[[""]] + } else { + 0L + } } #' Get BOS token ID #' @keywords internal get_bos_id <- function(tokenizer) { - if ("" %in% names(tokenizer$special_tokens)) { - tokenizer$special_tokens[[""]]$id - } else if ("" %in% names(tokenizer$vocab)) { - tokenizer$vocab[[""]] - } else { - NULL - } + if ("" %in% names(tokenizer$special_tokens)) { + tokenizer$special_tokens[[""]]$id + } else if ("" %in% names(tokenizer$vocab)) { + tokenizer$vocab[[""]] + } else { + NULL + } } #' Get EOS token ID #' @keywords internal get_eos_id <- function(tokenizer) { - if ("" %in% names(tokenizer$special_tokens)) { - tokenizer$special_tokens[[""]]$id - } else if ("" %in% names(tokenizer$vocab)) { - tokenizer$vocab[[""]] - } else { - NULL - } + if ("" %in% names(tokenizer$special_tokens)) { + tokenizer$special_tokens[[""]]$id + } else if ("" %in% names(tokenizer$vocab)) { + tokenizer$vocab[[""]] + } else { + NULL + } } #' Get vocabulary size @@ -482,14 +464,14 @@ get_eos_id <- function(tokenizer) { #' @return Integer vocabulary size. #' @export vocab_size <- function(tokenizer) { - tokenizer$vocab_size + tokenizer$vocab_size } # Null-coalescing operator if (!exists("%||%", mode = "function")) { - `%||%` <- function( - x, - y - ) if (is.null(x)) y else x + `%||%` <- function(x, y) if (is.null(x)) { + y + } else { + x + } } - diff --git a/R/txt2img.R b/R/txt2img.R index 0684836..445f3dc 100644 --- a/R/txt2img.R +++ b/R/txt2img.R @@ -11,17 +11,12 @@ #' \dontrun{ #' img <- txt2img("a cat wearing sunglasses in space", device = "cuda") #' } -txt2img <- function( - prompt, - model_name = c("sd21", "sdxl"), - ... -) { - switch(model_name, - # "sd15" = txt2img_sd15(prompt, ...), - "sd21" = txt2img_sd21(prompt, ...), - "sdxl" = txt2img_sdxl(prompt, ...), - # "sd3" = txt2img_sd3(prompt, ...), - stop("Unsupported model: ", model_name) - ) +txt2img <- function(prompt, model_name = c("sd21", "sdxl"), ...) { + switch(model_name, + # "sd15" = txt2img_sd15(prompt, ...), + "sd21" = txt2img_sd21(prompt, ...), + "sdxl" = txt2img_sdxl(prompt, ...), + # "sd3" = txt2img_sd3(prompt, ...), + stop("Unsupported model: ", model_name) + ) } - diff --git a/R/txt2img_sd21.R b/R/txt2img_sd21.R index 6edce0e..0241cd6 100644 --- a/R/txt2img_sd21.R +++ b/R/txt2img_sd21.R @@ -33,17 +33,17 @@ #' \dontrun{ #' img <- txt2img("a cat wearing sunglasses in space", device = "cuda") #' } -txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, - pipeline = NULL, devices = "auto", - unet_dtype_str = NULL, download_models = FALSE, - scheduler = "ddim", timesteps = NULL, - initial_latents = NULL, num_inference_steps = 50, - guidance_scale = 7.5, seed = NULL, save_file = TRUE, - filename = NULL, metadata_path = NULL, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, ...) { - model_name = "sd21" +txt2img_sd21 <- function(prompt, negative_prompt = NULL, img_dim = 768, + pipeline = NULL, devices = "auto", + unet_dtype_str = NULL, download_models = FALSE, + scheduler = "ddim", timesteps = NULL, + initial_latents = NULL, num_inference_steps = 50, + guidance_scale = 7.5, seed = NULL, save_file = TRUE, + filename = NULL, metadata_path = NULL, + use_native_decoder = FALSE, + use_native_text_encoder = FALSE, + use_native_unet = FALSE, ...) { + model_name <- "sd21" # Handle "auto" devices if (identical(devices, "auto")) { @@ -51,8 +51,8 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, } m2d <- models2devices(model_name = model_name, devices = devices, - unet_dtype_str = unet_dtype_str, - download_models = download_models) + unet_dtype_str = unet_dtype_str, + download_models = download_models) devices <- m2d$devices unet_dtype <- m2d$unet_dtype device_cpu <- m2d$device_cpu @@ -60,10 +60,10 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, if (is.null(pipeline)) { pipeline <- load_pipeline(model_name = model_name, m2d = m2d, - unet_dtype_str = unet_dtype_str, - use_native_decoder = use_native_decoder, - use_native_text_encoder = use_native_text_encoder, - use_native_unet = use_native_unet) + unet_dtype_str = unet_dtype_str, + use_native_decoder = use_native_decoder, + use_native_text_encoder = use_native_text_encoder, + use_native_unet = use_native_unet) } # Start timing @@ -86,13 +86,13 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, empty_prompt_embed <- empty_prompt_embed$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) prompt_embed <- prompt_embed$to(dtype = unet_dtype, - device = torch::torch_device(devices$unet)) + device = torch::torch_device(devices$unet)) message("Creating schedule...") # Load scheduler schedule <- ddim_scheduler_create(num_inference_steps = num_inference_steps, - beta_schedule = "scaled_linear", - device = torch::torch_device(devices$unet)) + beta_schedule = "scaled_linear", + device = torch::torch_device(devices$unet)) if (is.null(timesteps)) { timesteps <- schedule$timesteps } @@ -111,42 +111,42 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, } else { # Create random latents latents <- torch::torch_randn(c(1, 4, latent_dim, latent_dim), - dtype = unet_dtype, - device = torch::torch_device(devices$unet)) + dtype = unet_dtype, + device = torch::torch_device(devices$unet)) } # Denoising loop (no gradients needed for inference) pb <- utils::txtProgressBar(min = 0, max = length(timesteps), style = 3) torch::with_no_grad({ - for (i in seq_along(timesteps)) { - timestep <- torch::torch_tensor(timesteps[i], - dtype = torch::torch_long(), - device = torch::torch_device(devices$unet)) - - # Get both conditional and unconditional predictions - noise_pred_uncond <- pipeline$unet(latents, timestep, empty_prompt_embed) - noise_pred_cond <- pipeline$unet(latents, timestep, prompt_embed) - - # CFG step - noise_pred <- noise_pred_uncond + guidance_scale * - (noise_pred_cond - noise_pred_uncond) - - # Calculating latent - latents <- ddim_scheduler_step(model_output = noise_pred, - timestep = timestep, - sample = latents, - schedule = schedule, - prediction_type = "v_prediction", - device = devices$unet) - latents <- latents$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) - utils::setTxtProgressBar(pb, i) - } - }) + for (i in seq_along(timesteps)) { + timestep <- torch::torch_tensor(timesteps[i], + dtype = torch::torch_long(), + device = torch::torch_device(devices$unet)) + + # Get both conditional and unconditional predictions + noise_pred_uncond <- pipeline$unet(latents, timestep, empty_prompt_embed) + noise_pred_cond <- pipeline$unet(latents, timestep, prompt_embed) + + # CFG step + noise_pred <- noise_pred_uncond + guidance_scale * + (noise_pred_cond - noise_pred_uncond) + + # Calculating latent + latents <- ddim_scheduler_step(model_output = noise_pred, + timestep = timestep, + sample = latents, + schedule = schedule, + prediction_type = "v_prediction", + device = devices$unet) + latents <- latents$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) + utils::setTxtProgressBar(pb, i) + } + }) close(pb) # Decode latents to image scaled_latent <- latents / 0.18215 scaled_latent <- scaled_latent$to(dtype = torch::torch_float32(), - device = torch::torch_device(devices$decoder)) + device = torch::torch_device(devices$decoder)) message("Decoding image...") decoded_output <- pipeline$decoder(scaled_latent) # Ensure tensor is on CPU @@ -161,7 +161,7 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, img <- img$permute(c(2, 3, 1)) # Normalize if needed - img <- (img + 1) / 2# scale from [-1, 1] → [0, 1] + img <- (img + 1) / 2 # scale from [-1, 1] → [0, 1] img <- torch::torch_clamp(img, min = 0, max = 1) # Convert to R array @@ -182,15 +182,15 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, } # Save metadata metadata <- list( - prompt = prompt, - negative_prompt = negative_prompt, - width = img_dim, - height = img_dim, - num_inference_steps = num_inference_steps, - guidance_scale = guidance_scale, - seed = seed, - scheduler = scheduler, - model = model_name + prompt = prompt, + negative_prompt = negative_prompt, + width = img_dim, + height = img_dim, + num_inference_steps = num_inference_steps, + guidance_scale = guidance_scale, + seed = seed, + scheduler = scheduler, + model = model_name ) if (!is.null(metadata_path)) { utils::write.csv(metadata, file = metadata_path, row.names = FALSE) @@ -201,9 +201,5 @@ txt2img_sd21 <- function (prompt, negative_prompt = NULL, img_dim = 768, message(sprintf("Image generated in %.2f seconds", elapsed[3])) # Return the generated image and metadata - return(list( - image = img_array, - metadata = metadata - )) + return(list(image = img_array, metadata = metadata)) } - diff --git a/R/txt2img_sdxl.R b/R/txt2img_sdxl.R index 542a7ab..4555fb7 100644 --- a/R/txt2img_sdxl.R +++ b/R/txt2img_sdxl.R @@ -44,16 +44,16 @@ #' profile <- sdxl_memory_profile(vram_gb = 8) #' img <- txt2img_sdxl("a forest path", memory_profile = profile) #' } -txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, - pipeline = NULL, devices = "auto", - memory_profile = NULL, unet_dtype_str = NULL, - download_models = FALSE, scheduler = "ddim", - timesteps = NULL, initial_latents = NULL, - num_inference_steps = 30, guidance_scale = 7.5, - seed = NULL, save_file = TRUE, filename = NULL, - metadata_path = NULL, use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, verbose = TRUE, ...) { +txt2img_sdxl <- function(prompt, negative_prompt = NULL, img_dim = 1024, + pipeline = NULL, devices = "auto", + memory_profile = NULL, unet_dtype_str = NULL, + download_models = FALSE, scheduler = "ddim", + timesteps = NULL, initial_latents = NULL, + num_inference_steps = 30, guidance_scale = 7.5, + seed = NULL, save_file = TRUE, filename = NULL, + metadata_path = NULL, use_native_decoder = FALSE, + use_native_text_encoder = FALSE, + use_native_unet = FALSE, verbose = TRUE, ...) { model_name <- "sdxl" # Resolve memory profile if provided @@ -63,13 +63,8 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, profile <- sdxl_memory_profile() } else if (is.character(memory_profile)) { # Named profile - vram <- switch(memory_profile, - "full_gpu" = 20, - "balanced" = 12, - "unet_gpu" = 8, - "cpu_only" = 0, - NULL - ) + vram <- switch(memory_profile, "full_gpu" = 20, "balanced" = 12, + "unet_gpu" = 8, "cpu_only" = 0, NULL) if (!is.null(vram)) { profile <- sdxl_memory_profile(vram_gb = vram) } @@ -87,7 +82,7 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, # Validate resolution if (img_dim > profile$max_resolution) { warning(sprintf("Resolution %d exceeds profile max %d, reducing", - img_dim, profile$max_resolution)) + img_dim, profile$max_resolution)) img_dim <- profile$max_resolution } # Set dtype from profile if not explicitly provided @@ -99,7 +94,7 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, } m2d <- models2devices(model_name = model_name, devices = devices, - unet_dtype_str = unet_dtype_str) + unet_dtype_str = unet_dtype_str) devices <- m2d$devices unet_dtype <- m2d$unet_dtype device_cpu <- m2d$device_cpu @@ -107,16 +102,18 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, if (is.null(pipeline)) { pipeline <- load_pipeline(model_name = model_name, m2d = m2d, - unet_dtype_str = unet_dtype_str, - use_native_decoder = use_native_decoder, - use_native_text_encoder = use_native_text_encoder, - use_native_unet = use_native_unet) + unet_dtype_str = unet_dtype_str, + use_native_decoder = use_native_decoder, + use_native_text_encoder = use_native_text_encoder, + use_native_unet = use_native_unet) } # Start timing start_time <- proc.time() # Process text prompt - if (verbose) message("Processing prompt...") + if (verbose) { + message("Processing prompt...") + } ## Tokenizer tokens <- CLIPTokenizer(prompt) prompt_embed1 <- pipeline$text_encoder(tokens) @@ -126,10 +123,10 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, text_embeds <- te2_output[[2]] text_embeds <- text_embeds$to(dtype = unet_dtype, - device = torch::torch_device(devices$unet)) - time_ids = torch::torch_tensor(c(img_dim, img_dim, 0, 0, img_dim, img_dim), # zero indexed as python - dtype = unet_dtype, - device = torch::torch_device(devices$unet))$unsqueeze(1) + device = torch::torch_device(devices$unet)) + time_ids <- torch::torch_tensor(c(img_dim, img_dim, 0, 0, img_dim, img_dim), # zero indexed as python + dtype = unet_dtype, + device = torch::torch_device(devices$unet))$unsqueeze(1) # clip-vit-large-patch14 if (is.null(negative_prompt)) { empty_tokens <- CLIPTokenizer("") @@ -147,30 +144,36 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, empty_prompt_embed <- empty_prompt_embed$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) prompt_embed <- prompt_embed$to(dtype = unet_dtype, - device = torch::torch_device(devices$unet)) + device = torch::torch_device(devices$unet)) empty_text_embeds <- empty_text_embeds$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) # Phase cleanup after text encoding (for GPU-poor mode) if (!is.null(profile) && profile$cleanup == "phase") { - if (verbose) message("Clearing VRAM after text encoding...") + if (verbose) { + message("Clearing VRAM after text encoding...") + } clear_vram(verbose = FALSE) } - if (verbose) message("Creating schedule...") + if (verbose) { + message("Creating schedule...") + } # Load scheduler schedule <- ddim_scheduler_create(num_inference_steps = num_inference_steps, - beta_schedule = "scaled_linear", - beta_start = 0.00085, - beta_end = 0.012, - rescale_betas_zero_snr = FALSE, - device = torch::torch_device(devices$unet)) + beta_schedule = "scaled_linear", + beta_start = 0.00085, + beta_end = 0.012, + rescale_betas_zero_snr = FALSE, + device = torch::torch_device(devices$unet)) if (is.null(timesteps)) { timesteps <- schedule$timesteps } # Run diffusion process - if (verbose) message("Generating image...") + if (verbose) { + message("Generating image...") + } if (!is.null(seed)) { set.seed(seed) torch::torch_manual_seed(seed = seed) @@ -183,70 +186,74 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, } else { # Create random latents latents <- torch::torch_randn(c(1, 4, latent_dim, latent_dim), - dtype = unet_dtype, - device = torch::torch_device(devices$unet)) + dtype = unet_dtype, + device = torch::torch_device(devices$unet)) } # Denoising loop (no gradients needed for inference) pb <- utils::txtProgressBar(min = 0, max = length(timesteps), style = 3) torch::with_no_grad({ - for (i in seq_along(timesteps)) { - timestep <- torch::torch_tensor(timesteps[i], - dtype = torch::torch_long(), - device = torch::torch_device(devices$unet)) + for (i in seq_along(timesteps)) { + timestep <- torch::torch_tensor(timesteps[i], + dtype = torch::torch_long(), + device = torch::torch_device(devices$unet)) - # Get both conditional and unconditional predictions - noise_pred_cond <- pipeline$unet(latents, timestep, prompt_embed, - text_embeds, time_ids) + # Get both conditional and unconditional predictions + noise_pred_cond <- pipeline$unet(latents, timestep, prompt_embed, + text_embeds, time_ids) - if (guidance_scale != 1) { - # If guidance scale is not 1, we need to calculate the unconditional prediction - # with an empty prompt - noise_pred_uncond <- pipeline$unet(latents, timestep, empty_prompt_embed, - empty_text_embeds, time_ids) - # CFG step - noise_pred <- noise_pred_uncond + guidance_scale * - (noise_pred_cond - noise_pred_uncond) - } else { - # If guidance scale is 1, we can use the conditional prediction directly - noise_pred <- noise_pred_cond - } + if (guidance_scale != 1) { + # If guidance scale is not 1, we need to calculate the unconditional prediction + # with an empty prompt + noise_pred_uncond <- pipeline$unet(latents, timestep, empty_prompt_embed, + empty_text_embeds, time_ids) + # CFG step + noise_pred <- noise_pred_uncond + guidance_scale * + (noise_pred_cond - noise_pred_uncond) + } else { + # If guidance scale is 1, we can use the conditional prediction directly + noise_pred <- noise_pred_cond + } - # Calculating latent - latents <- ddim_scheduler_step(model_output = noise_pred, - timestep = timestep, - sample = latents, - schedule = schedule, - prediction_type = "epsilon", - device = devices$unet) - latents <- latents$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) - utils::setTxtProgressBar(pb, i) + # Calculating latent + latents <- ddim_scheduler_step(model_output = noise_pred, + timestep = timestep, + sample = latents, + schedule = schedule, + prediction_type = "epsilon", + device = devices$unet) + latents <- latents$to(dtype = unet_dtype, device = torch::torch_device(devices$unet)) + utils::setTxtProgressBar(pb, i) - # Step cleanup for GPU-poor mode - if (!is.null(profile) && profile$step_cleanup_interval > 0) { - if (i %% profile$step_cleanup_interval == 0) { - clear_vram(verbose = FALSE) - } + # Step cleanup for GPU-poor mode + if (!is.null(profile) && profile$step_cleanup_interval > 0) { + if (i %% profile$step_cleanup_interval == 0) { + clear_vram(verbose = FALSE) } } - }) + } + }) close(pb) # Phase cleanup before decode (for GPU-poor mode) if (!is.null(profile) && profile$cleanup == "phase") { - if (verbose) message("Clearing VRAM before decode...") + if (verbose) { + message("Clearing VRAM before decode...") + } clear_vram(verbose = FALSE) } # Decode latents to image scaled_latent <- latents / 0.18215 scaled_latent <- scaled_latent$to(dtype = torch::torch_float32(), - device = torch::torch_device(devices$decoder)) + device = torch::torch_device(devices$decoder)) # message("Loading post_quant_conv...") post_conv_latent <- post_quant_conv(x = scaled_latent, - dtype = torch::torch_float32(), - device = devices$decoder) - if (verbose) message("Decoding image...") + dtype = torch::torch_float32(), + device = devices$decoder) + if (verbose) { + message("Decoding image...") + } decoded_output <- pipeline$decoder(post_conv_latent) # Ensure tensor is on CPU img <- decoded_output$cpu() @@ -260,7 +267,7 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, img <- img$permute(c(2, 3, 1)) # Normalize - img <- (img + 1) / 2# scale from [-1, 1] → [0, 1] + img <- (img + 1) / 2 # scale from [-1, 1] → [0, 1] img <- torch::torch_clamp(img, min = 0, max = 1) # Convert to R array @@ -281,16 +288,16 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, } # Save metadata metadata <- list( - prompt = prompt, - negative_prompt = negative_prompt, - width = img_dim, - height = img_dim, - num_inference_steps = num_inference_steps, - guidance_scale = guidance_scale, - seed = seed, - scheduler = scheduler, - model = model_name, - filename = filename + prompt = prompt, + negative_prompt = negative_prompt, + width = img_dim, + height = img_dim, + num_inference_steps = num_inference_steps, + guidance_scale = guidance_scale, + seed = seed, + scheduler = scheduler, + model = model_name, + filename = filename ) if (!is.null(metadata_path)) { utils::write.csv(metadata, file = metadata_path, row.names = FALSE) @@ -298,12 +305,10 @@ txt2img_sdxl <- function (prompt, negative_prompt = NULL, img_dim = 1024, } # Report timing elapsed <- proc.time() - start_time - if (verbose) message(sprintf("Image generated in %.2f seconds", elapsed[3])) + if (verbose) { + message(sprintf("Image generated in %.2f seconds", elapsed[3])) + } # Return the generated image and metadata - return(list( - image = img_array, - metadata = metadata - )) + return(list(image = img_array, metadata = metadata)) } - diff --git a/R/txt2vid_ltx2.R b/R/txt2vid_ltx2.R deleted file mode 100644 index 68517b6..0000000 --- a/R/txt2vid_ltx2.R +++ /dev/null @@ -1,615 +0,0 @@ -#' Generate Video from Text Prompt using LTX-2 -#' -#' Generates video using the LTX-2 diffusion transformer model. -#' -#' @param prompt Character. Text prompt describing the video to generate. -#' @param negative_prompt Character. Optional negative prompt. -#' @param width Integer. Video width in pixels (default 768). -#' @param height Integer. Video height in pixels (default 512). -#' @param num_frames Integer. Number of frames to generate (default 121). -#' @param fps Numeric. Frames per second (default 24). -#' @param num_inference_steps Integer. Number of denoising steps (default 8 for distilled). -#' @param guidance_scale Numeric. CFG scale (default 4.0). -#' @param memory_profile Character or list. Memory profile: "auto" for auto-detection, -#' or a profile from `ltx2_memory_profile()`. -#' @param text_backend Character. Text encoding backend: "gemma3" (native), "api", "precomputed", or "random". -#' @param text_model_path Character. Path to Gemma3 model (for "gemma3" backend). Supports glob patterns. -#' @param text_api_url Character. URL for text encoding API (if backend = "api"). -#' @param vae Optional. Pre-loaded VAE module. -#' @param dit Optional. Pre-loaded DiT transformer module. -#' @param connectors Optional. Pre-loaded text connectors module. -#' @param seed Integer. Random seed for reproducibility. -#' @param output_file Character. Path to save output video (NULL for no save). -#' @param output_format Character. Output format: "mp4", "gif", or "frames". -#' @param return_latents Logical. If TRUE, also return final latents. -#' @param verbose Logical. Print progress messages. -#' -#' @return A list with: -#' - `video`: Array of video frames [frames, height, width, channels] -#' - `latents`: (if return_latents=TRUE) Final latent tensor -#' - `metadata`: Generation metadata -#' -#' @export -#' -#' @examples -#' \dontrun{ -#' # Basic usage -#' result <- txt2vid_ltx2("A cat walking on a beach at sunset") -#' -#' # With specific settings -#' result <- txt2vid_ltx2( -#' prompt = "A timelapse of clouds moving over mountains", -#' width = 512, -#' height = 320, -#' num_frames = 61, -#' num_inference_steps = 8, -#' seed = 42, -#' output_file = "clouds.mp4" -#' ) -#' } -txt2vid_ltx2 <- function( - prompt, - negative_prompt = NULL, - width = 768L, - height = 512L, - num_frames = 121L, - fps = 24.0, - num_inference_steps = 8L, - guidance_scale = 4.0, - memory_profile = "auto", - text_backend = "gemma3", - text_model_path = NULL, - text_api_url = NULL, - vae = NULL, - dit = NULL, - connectors = NULL, - seed = NULL, - output_file = NULL, - output_format = "mp4", - return_latents = FALSE, - verbose = TRUE -) { - # Start timing - start_time <- Sys.time() - - # Ensure integers - width <- as.integer(width) - height <- as.integer(height) - num_frames <- as.integer(num_frames) - num_inference_steps <- as.integer(num_inference_steps) - - # Set seed if provided - if (!is.null(seed)) { - torch::torch_manual_seed(seed) - # torch_manual_seed sets both CPU and CUDA seeds - } - - # Resolve memory profile - if (identical(memory_profile, "auto")) { - memory_profile <- ltx2_memory_profile() - } else if (is.character(memory_profile)) { - # Named profile - vram <- switch(memory_profile, - "high" = 20, - "medium" = 12, - "low" = 8, - "very_low" = 6, - "cpu_only" = 0, - 8# default - ) - memory_profile <- ltx2_memory_profile(vram_gb = vram) - } - - if (verbose) { - message(sprintf("Using memory profile: %s", memory_profile$name)) - } - - # Validate and adjust resolution for profile - validated <- validate_resolution(height, width, num_frames, memory_profile) - if (validated$adjusted) { - height <- validated$height - width <- validated$width - num_frames <- validated$num_frames - } - - # LTX-2 VAE compression ratios - spatial_ratio <- 32L - temporal_ratio <- 8L - - # Calculate latent dimensions - latent_height <- height %/% spatial_ratio - latent_width <- width %/% spatial_ratio - latent_frames <- (num_frames - 1L) %/% temporal_ratio + 1L - - if (verbose) { - message(sprintf("Video: %dx%d, %d frames @ %.1f fps", width, height, num_frames, fps)) - message(sprintf("Latents: %dx%d, %d frames", latent_width, latent_height, latent_frames)) - } - - # Device setup - dit_device <- memory_profile$dit_device - vae_device <- memory_profile$vae_device - - torch::with_no_grad({ - - # ---- Step 1: Text Encoding ---- - if (verbose) message("Encoding text prompt...") - - # Determine dtype based on profile - latent_dtype <- if (memory_profile$dtype == "float16") { - torch::torch_float16() - } else { - torch::torch_float32() - } - - # Resolve model path (use hfhub or explicit path) - resolved_model_path <- NULL - if (text_backend == "gemma3") { - if (!is.null(text_model_path)) { - # Explicit path provided - expanded_path <- path.expand(text_model_path) - if (grepl("\\*", expanded_path)) { - # Glob pattern - find matching directories - matches <- Sys.glob(expanded_path) - if (length(matches) > 0) { - resolved_model_path <- matches[1] - } - } else if (dir.exists(expanded_path)) { - resolved_model_path <- expanded_path - } - if (is.null(resolved_model_path)) { - stop("Gemma3 model not found at: ", text_model_path) - } - } else { - # Use hfhub to find model via config.json - if (!requireNamespace("hfhub", quietly = TRUE)) { - stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") - } - gemma_repo <- "google/gemma-3-12b-it" - config_path <- tryCatch({ - hfhub::hub_download(gemma_repo, "config.json", local_files_only = TRUE) - }, error = function(e) NULL) - - if (is.null(config_path)) { - stop("Gemma3 model not found in HuggingFace cache.\n", - "Download with: huggingface-cli download ", gemma_repo) - } - resolved_model_path <- dirname(config_path) - } - } - - # Encode prompt - text_result <- encode_text_ltx2( - prompt = prompt, - backend = text_backend, - model_path = resolved_model_path, - tokenizer_path = resolved_model_path, - api_url = text_api_url, - max_sequence_length = 128L, - caption_channels = 3840L, - device = "cpu", # Text encoder on CPU (GPU-poor) - dtype = torch::torch_float32() # CPU always float32 - ) - prompt_embeds <- text_result$prompt_embeds - prompt_attention_mask <- text_result$prompt_attention_mask - - # Encode negative prompt - if (is.null(negative_prompt)) { - negative_prompt <- "" - } - neg_result <- encode_text_ltx2( - prompt = negative_prompt, - backend = text_backend, - model_path = resolved_model_path, - tokenizer_path = resolved_model_path, - api_url = text_api_url, - max_sequence_length = 128L, - caption_channels = 3840L, - device = "cpu", - dtype = torch::torch_float32() - ) - negative_prompt_embeds <- neg_result$prompt_embeds - negative_attention_mask <- neg_result$prompt_attention_mask - - # ---- Step 1b: Apply Connectors ---- - # Connectors transform packed text embeddings for video/audio cross-attention - if (is.null(connectors)) { - # Load connectors from HuggingFace using hfhub - connector_path <- tryCatch({ - if (!requireNamespace("hfhub", quietly = TRUE)) NULL - else hfhub::hub_download( - "Lightricks/LTX-2", - "connectors/diffusion_pytorch_model.safetensors", - local_files_only = TRUE - ) - }, error = function(e) NULL) - - if (!is.null(connector_path) && file.exists(connector_path)) { - if (verbose) message("Loading text connectors...") - connectors <- load_ltx2_connectors( - weights_path = connector_path, - device = "cpu", - dtype = "float32", - verbose = verbose - ) - } else { - if (verbose) message("Text connectors not found - using embeddings directly") - } - } - - if (!is.null(connectors)) { - # Run connectors to get video/audio conditioning - if (verbose) message("Applying text connectors...") - connector_result <- connectors(prompt_embeds, prompt_attention_mask) - video_embeds <- connector_result[[1]] - audio_embeds <- connector_result[[2]] - - neg_connector_result <- connectors(negative_prompt_embeds, negative_attention_mask) - neg_video_embeds <- neg_connector_result[[1]] - neg_audio_embeds <- neg_connector_result[[2]] - } else { - # Fallback: use packed embeddings directly (may not match DiT dimensions) - video_embeds <- prompt_embeds - audio_embeds <- prompt_embeds - neg_video_embeds <- negative_prompt_embeds - neg_audio_embeds <- negative_prompt_embeds - } - - # Move to GPU with correct dtype - video_embeds <- video_embeds$to(device = dit_device, dtype = latent_dtype) - audio_embeds <- audio_embeds$to(device = dit_device, dtype = latent_dtype) - neg_video_embeds <- neg_video_embeds$to(device = dit_device, dtype = latent_dtype) - neg_audio_embeds <- neg_audio_embeds$to(device = dit_device, dtype = latent_dtype) - - # ---- Step 2: Initialize Latents ---- - if (verbose) message("Initializing latents...") - - # LTX-2 has 128 latent channels - latent_channels <- 128L - batch_size <- 1L - - # Random noise - latents <- torch::torch_randn( - c(batch_size, latent_channels, latent_frames, latent_height, latent_width), - device = dit_device, - dtype = latent_dtype - ) - - # Flatten spatial dims for transformer: [B, C, T, H, W] -> [B, T*H*W, C] - num_patches <- latent_frames * latent_height * latent_width - latents <- latents$permute(c(1, 3, 4, 5, 2)) # [B, T, H, W, C] - latents <- latents$reshape(c(batch_size, num_patches, latent_channels)) - - # ---- Step 3: Create Scheduler ---- - if (verbose) message("Setting up FlowMatch scheduler...") - - schedule <- flowmatch_set_timesteps( - flowmatch_scheduler_create(shift = 9.0), - num_inference_steps = num_inference_steps, - device = dit_device - ) - - # ---- Step 4: Load/Create DiT if needed ---- - if (is.null(dit)) { - # Try to load INT4 quantized weights from R_user_dir cache - cache_dir <- tools::R_user_dir("diffuseR", "cache") - int4_path <- file.path(cache_dir, "ltx2_transformer_int4.safetensors") - # Check for single file or sharded files - int4_shards <- list.files(cache_dir, pattern = "ltx2_transformer_int4.*\\.safetensors$") - if (file.exists(int4_path) || length(int4_shards) > 0) { - if (verbose) message("Loading INT4 quantized DiT...") - - # Enable INT4-native model creation - old_use_int4 <- getOption("diffuseR.use_int4", FALSE) - old_int4_device <- getOption("diffuseR.int4_device", "cuda") - old_int4_dtype <- getOption("diffuseR.int4_dtype", torch::torch_float16()) - - options(diffuseR.use_int4 = TRUE) - options(diffuseR.int4_device = dit_device) - options(diffuseR.int4_dtype = latent_dtype) - - # Create model with int4_linear layers via make_linear() - dit <- ltx2_video_transformer_3d_model( - in_channels = latent_channels, - out_channels = latent_channels, - num_attention_heads = 32L, - attention_head_dim = 128L, - cross_attention_dim = 4096L, - audio_in_channels = 128L, - audio_out_channels = 128L, - audio_num_attention_heads = 32L, - audio_attention_head_dim = 64L, - audio_cross_attention_dim = 2048L, - num_layers = 48L, - caption_channels = 3840L, - vae_scale_factors = c(temporal_ratio, spatial_ratio, spatial_ratio) - ) - - # Load INT4 weights (keeps them compressed, dequantizes during forward) - int4_weights <- load_int4_weights(int4_path, verbose = verbose) - load_int4_weights_into_model(dit, int4_weights, verbose = verbose) - - # Move model to GPU (non-INT4 params like biases, norms) - dit <- dit$to(device = dit_device, dtype = latent_dtype) - - # Restore options - options(diffuseR.use_int4 = old_use_int4) - options(diffuseR.int4_device = old_int4_device) - options(diffuseR.int4_dtype = old_int4_dtype) - } else { - if (verbose) message("NOTE: INT4 weights not found - run quantize_ltx2_transformer() first") - stop("DiT model required. Run: quantize_ltx2_transformer()") - } - } - - # ---- Step 5: Denoising Loop ---- - if (verbose) message(sprintf("Denoising (%d steps)...", num_inference_steps)) - - # Audio placeholder (zeros for video-only generation) - audio_latents <- torch::torch_zeros( - c(batch_size, 50L, 128L), # Placeholder audio: [B, seq, audio_channels=128] - device = dit_device, - dtype = latent_dtype - ) - - timesteps_vec <- schedule$timesteps - sigmas <- schedule$sigmas - - for (i in seq_len(num_inference_steps)) { - t_idx <- i - t <- timesteps_vec[t_idx] - sigma <- sigmas[t_idx] - if (i < num_inference_steps) { - sigma_next <- sigmas[t_idx + 1L] - } else { - sigma_next <- 0 - } - - if (verbose && i %% max(1, num_inference_steps %/% 4) == 1) { - message(sprintf(" Step %d/%d (sigma=%.3f)", i, num_inference_steps, as.numeric(sigma))) - } - - # Prepare timestep tensor - timestep <- torch::torch_tensor(c(as.numeric(t)))$unsqueeze(2L) - timestep <- timestep$to(device = dit_device, dtype = latent_dtype) - - # CFG: conditional and unconditional pass - if (memory_profile$cfg_mode == "sequential") { - # Sequential CFG (memory efficient) - noise_pred <- sequential_cfg_forward( - model = dit, - latents = latents, - timestep = timestep, - prompt_embeds = video_embeds, - negative_prompt_embeds = neg_video_embeds, - guidance_scale = guidance_scale, - audio_hidden_states = audio_latents, - audio_encoder_hidden_states = audio_embeds, - num_frames = latent_frames, - height = latent_height, - width = latent_width, - fps = fps, - audio_num_frames = 50L - ) - } else { - # Batched CFG - latents_input <- torch::torch_cat(list(latents, latents), dim = 1L) - video_input <- torch::torch_cat(list(neg_video_embeds, video_embeds), dim = 1L) - audio_input <- torch::torch_cat(list(neg_audio_embeds, audio_embeds), dim = 1L) - timestep_input <- torch::torch_cat(list(timestep, timestep), dim = 1L) - - output <- dit( - hidden_states = latents_input, - audio_hidden_states = torch::torch_cat(list(audio_latents, audio_latents), dim = 1L), - encoder_hidden_states = video_input, - audio_encoder_hidden_states = audio_input, - timestep = timestep_input, - num_frames = latent_frames, - height = latent_height, - width = latent_width, - fps = fps, - audio_num_frames = 50L - ) - - noise_pred_all <- output$sample - noise_pred_uncond <- noise_pred_all[1,,]$unsqueeze(1L) - noise_pred_cond <- noise_pred_all[2,,]$unsqueeze(1L) - # CFG: use tensor method to preserve dtype - noise_pred <- noise_pred_uncond + (noise_pred_cond - noise_pred_uncond)$mul(guidance_scale) - } - - # FlowMatch step - dt <- torch::torch_tensor(sigma_next - sigma, dtype = latent_dtype, device = dit_device) - latents <- latents + dt * noise_pred - - # Cleanup for low memory - if (memory_profile$name %in% c("low", "very_low") && i %% 2 == 0) { - clear_vram() - } - } - - # ---- Step 6: Decode Latents ---- - if (verbose) message("Decoding video...") - - # Reshape latents back to spatial: [B, T*H*W, C] -> [B, C, T, H, W] - latents <- latents$reshape(c(batch_size, latent_frames, latent_height, latent_width, latent_channels)) - latents <- latents$permute(c(1, 5, 2, 3, 4)) # [B, C, T, H, W] - - # Load/create VAE if needed - if (is.null(vae)) { - if (verbose) message("Loading VAE...") - - # Try to find VAE in HuggingFace cache - vae_path <- tryCatch({ - if (requireNamespace("hfhub", quietly = TRUE)) { - config_path <- hfhub::hub_download("Lightricks/LTX-2", "vae/config.json", - local_files_only = TRUE) - dirname(config_path) - } else { - NULL - } - }, error = function(e) NULL) - - if (is.null(vae_path)) { - if (verbose) message("NOTE: VAE not found in cache - skipping decode") - video_tensor <- latents - } else { - # Determine VAE dtype based on latent dtype - vae_dtype <- if (identical(latent_dtype, torch::torch_bfloat16()) || - identical(latent_dtype, torch::torch_float16())) { - "float16" - } else { - "float32" - } - - vae <- load_ltx2_vae( - weights_path = vae_path, - device = vae_device, - dtype = vae_dtype, - verbose = verbose - ) - } - } - - if (!is.null(vae)) { - # Configure VAE for memory profile - configure_vae_for_profile(vae, memory_profile) - - # Move VAE to device - vae <- vae$to(device = vae_device) - - # Decode - video_tensor <- vae$decode(latents) - } - - # Prepare tensor for conversion to R array - # NOTE: Must use as.array() instead of $numpy() due to R torch bug where - - # tensors returned from with_no_grad() have corrupted method references - # (error: "could not find function 'fn'"). See cornyverse CLAUDE.md. - video_cpu <- video_tensor$squeeze(1L)$permute(c(2, 3, 4, 1))$cpu() - - }) # end with_no_grad - - # Convert to R array (as.array works, $numpy() fails on tensors from with_no_grad) - video_array <- as.array(video_cpu) - - # Clamp to [0, 1] - video_array <- pmax(pmin(video_array, 1), 0) - - # ---- Step 7: Save Output ---- - if (!is.null(output_file)) { - if (verbose) message(sprintf("Saving to %s...", output_file)) - save_video_frames(video_array, output_file, fps = fps, verbose = verbose) - } - - # Build result - elapsed <- as.numeric(difftime(Sys.time(), start_time, units = "secs")) - if (verbose) { - message(sprintf("Generation complete in %.1f seconds", elapsed)) - } - - result <- list( - video = video_array, - metadata = list( - prompt = prompt, - negative_prompt = negative_prompt, - width = width, - height = height, - num_frames = num_frames, - fps = fps, - num_inference_steps = num_inference_steps, - guidance_scale = guidance_scale, - seed = seed, - memory_profile = memory_profile$name, - elapsed_seconds = elapsed - ) - ) - - if (return_latents) { - result$latents <- latents$cpu() - } - - result -} - -#' Save video frames to file -#' -#' Saves an array of video frames to an MP4 file using ffmpeg via the av package. -#' -#' @param video_array Numeric array with dimensions [frames, height, width, channels]. -#' Values should be in range [0, 1]. -#' @param output_file Character. Path to output video file. -#' @param fps Numeric. Frames per second. Default: 24. -#' @param verbose Logical. Print progress messages. Default: TRUE. -#' @return Invisibly returns the output file path. -#' @keywords internal -save_video_frames <- function( - video_array, - output_file, - fps = 24, - verbose = TRUE -) { - if (!requireNamespace("av", quietly = TRUE)) { - stop("Package 'av' is required for video saving. Install with: install.packages('av')") - } - - # video_array should be [frames, height, width, channels] - dims <- dim(video_array) - if (length(dims) != 4) { - stop("video_array must have 4 dimensions: [frames, height, width, channels]") - } - - num_frames <- dims[1] - height <- dims[2] - width <- dims[3] - channels <- dims[4] - - if (channels != 3) { - stop("Expected 3 color channels, got ", channels) - } - - # Create temporary directory for frames - - temp_dir <- tempfile("video_frames_") - dir.create(temp_dir) - on.exit(unlink(temp_dir, recursive = TRUE), add = TRUE) - - # Save each frame as PNG - if (verbose) message(sprintf(" Writing %d frames...", num_frames)) - - for (i in seq_len(num_frames)) { - # Extract frame and convert to [0, 255] uint8 - frame <- video_array[i,,,] - frame <- pmax(pmin(frame, 1), 0) * 255 - - # Convert to integer matrix for png - # frame is [height, width, channels] - frame_int <- array(as.integer(round(frame)), dim = dim(frame)) - - # Save as PNG - frame_path <- file.path(temp_dir, sprintf("frame_%05d.png", i)) - png::writePNG(frame_int / 255, frame_path) - } - - # Encode video using av - if (verbose) message(" Encoding video...") - - # Get list of frame files - frame_files <- list.files(temp_dir, pattern = "frame_.*\\.png$", full.names = TRUE) - frame_files <- sort(frame_files) - - # Use av to encode - av::av_encode_video( - input = frame_files, - output = output_file, - framerate = fps, - codec = "libx264", - verbose = FALSE - ) - - if (verbose) message(sprintf(" Saved: %s", output_file)) - - invisible(output_file) -} - diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R new file mode 100644 index 0000000..9b66bff --- /dev/null +++ b/R/txt2vid_ltx23.R @@ -0,0 +1,922 @@ +#' LTX-2.3 Text-to-Video Pipeline +#' +#' Fresh R port of the LTX-2 text-to-video flow from the diffusers +#' reference (Apache-2.0, pipelines/ltx2/pipeline_ltx2.py), specialized +#' for the distilled LTX 2.3 checkpoints: 8-step official sigma schedule, +#' no classifier-free guidance, joint audio-video denoising with an Euler +#' velocity step, and audio decoding through the audio VAE and BWE +#' vocoder to 48 kHz stereo. +#' +#' @name txt2vid_ltx23 +NULL + +#' Official distilled sigma schedule +#' +#' The distilled LTX sigma values (with terminal zero appended), as +#' published in the Apache-2.0 diffusers reference +#' (pipelines/ltx2/utils.py). +#' +#' @return Numeric vector of length 9. +#' +#' @export +ltx23_distilled_sigmas <- function() { + c(1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) +} + +#' Stage-2 distilled sigma schedule (two-stage refinement) +#' +#' @return Numeric vector of length 4. +#' +#' @export +ltx23_stage2_distilled_sigmas <- function() { + c(0.909375, 0.725, 0.421875, 0.0) +} + +# Pack video latents [B, C, F, H, W] -> [B, F*H*W, C] (patch sizes are 1 +# for the LTX transformer; kept general per the reference) +ltx23_pack_video_latents <- function(latents, patch_size = 1L, + patch_size_t = 1L) { + d <- latents$shape + pf <- d[3] %/% patch_size_t + ph <- d[4] %/% patch_size + pw <- d[5] %/% patch_size + latents <- latents$reshape(c(d[1], -1L, pf, patch_size_t, ph, + patch_size, pw, patch_size)) + # [B, C, F', pt, H', p, W', p] -> [B, F', H', W', C, pt, p, p] + latents <- latents$permute(c(1L, 3L, 5L, 7L, 2L, 4L, 6L, 8L)) + latents$flatten(start_dim = 5L, end_dim = 8L)$flatten(start_dim = 2L, end_dim = 4L) +} + +# Unpack video latents [B, S, D] -> [B, C, F, H, W] +ltx23_unpack_video_latents <- function(latents, num_frames, height, width, + patch_size = 1L, patch_size_t = 1L) { + b <- latents$shape[1] + latents <- latents$reshape(c(b, num_frames, height, width, -1L, + patch_size_t, patch_size, patch_size)) + latents <- latents$permute(c(1L, 5L, 2L, 6L, 3L, 7L, 4L, 8L)) + latents$flatten(start_dim = 7L, end_dim = 8L)$ + flatten(start_dim = 5L, end_dim = 6L)$ + flatten(start_dim = 3L, end_dim = 4L) +} + +# Pack audio latents [B, C, L, M] -> [B, L, C * M] +ltx23_pack_audio_latents <- function(latents) { + latents$transpose(2L, 3L)$flatten(start_dim = 3L, end_dim = 4L) +} + +# Unpack audio latents [B, L, C * M] -> [B, C, L, M] +ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { + latents$unflatten(3L, c(-1L, num_mel_bins))$transpose(2L, 3L) +} + +# Audio statistics apply to the PACKED [B, L, C*M] representation +# (broadcast over the trailing feature dim) +.ltx23_denormalize_audio <- function(latents, latents_mean, latents_std) { + mean <- latents_mean$to(device = latents$device, dtype = latents$dtype) + std <- latents_std$to(device = latents$device, dtype = latents$dtype) + latents * std + mean +} + +# Joint audio/video Euler denoise loop over a sigma schedule (CFG-free). +# With a conditioning_mask [B, S] (prefix conditioning), conditioned +# video tokens see a per-token timestep of zero and are frozen through +# the Euler updates (reference i2v semantics at strength 1). +.ltx23_denoise <- function(transformer, latents, audio_latents, sigmas, + video_text_embeds, audio_text_embeds, text_mask, + latent_frames, latent_height, latent_width, + audio_num_frames, frame_rate, device, + compute_dtype, verbose = TRUE, stage = "", + conditioning_mask = NULL, + audio_conditioned = FALSE) { + f32 <- torch::torch_float32() + keep_mask <- if (!is.null(conditioning_mask)) { + conditioning_mask$unsqueeze(3L) # [B, S, 1] for latent blending + } else { + NULL + } + t_zero <- if (audio_conditioned) { + torch::torch_tensor(0, device = device, dtype = f32) + } else { + NULL + } + video_coords <- transformer$rope$prepare_video_coords(latents$shape[1], + latent_frames, latent_height, latent_width, device, + fps = frame_rate) + audio_coords <- transformer$audio_rope$prepare_audio_coords( + audio_latents$shape[1], audio_num_frames, device + ) + + n_steps <- length(sigmas) - 1L + scale_mult <- transformer$timestep_scale_multiplier + step_t0 <- Sys.time() + torch::with_no_grad({ + for (i in seq_len(n_steps)) { + sigma <- sigmas[i] + t <- torch::torch_tensor(sigma * scale_mult, device = device, dtype = f32) + t_video <- if (is.null(conditioning_mask)) { + t + } else { + # Per-token video timestep: conditioned tokens see zero + t * (1 - conditioning_mask) + } + + out <- transformer( + hidden_states = latents$to(dtype = compute_dtype), + audio_hidden_states = audio_latents$to(dtype = compute_dtype), + encoder_hidden_states = video_text_embeds, + audio_encoder_hidden_states = audio_text_embeds, + timestep = t_video, + audio_timestep = if (audio_conditioned) t_zero else t, + sigma = t, + audio_sigma = if (audio_conditioned) t_zero else t, + encoder_attention_mask = text_mask, + audio_encoder_attention_mask = text_mask, + num_frames = latent_frames, + height = latent_height, + width = latent_width, + fps = frame_rate, + audio_num_frames = audio_num_frames, + video_coords = video_coords, + audio_coords = audio_coords, + use_cross_timestep = TRUE + ) + + # Euler velocity step in float32; dt is negative (sigma decreasing) + dt <- torch::torch_tensor(sigmas[i + 1L] - sigma, device = device, dtype = f32) + stepped <- latents + dt * out$sample$to(dtype = f32) + latents <- if (is.null(keep_mask)) { + stepped + } else { + # Conditioned tokens stay clean + latents * keep_mask + stepped * (1 - keep_mask) + } + if (!audio_conditioned) { + audio_latents <- audio_latents + dt * out$audio_sample$to(dtype = f32) + } + rm(out) + gc(verbose = FALSE) + if (verbose) { + message(sprintf(" %sstep %d/%d (sigma %.4f, %.1fs)", + stage, i, n_steps, sigma, + as.numeric(difftime(Sys.time(), step_t0, units = "secs")))) + step_t0 <- Sys.time() + } + } + }) + list(latents = latents, audio_latents = audio_latents) +} + +#' Load the LTX-2.3 generation components from a single-file checkpoint +#' +#' Builds the transformer, connectors, video VAE, audio VAE, and vocoder +#' with the LTX 2.3 configuration and streams the checkpoint weights into +#' them. The Gemma3 text encoder ships separately (see +#' \code{\link{load_gemma3_text_encoder}}). +#' +#' @param checkpoint_path Path to the single-file checkpoint (e.g. +#' \code{ltx-2.3-22b-distilled-1.1.safetensors}) or to an fp8 artifact +#' directory produced by \code{\link{ltx23_quantize_fp8}}. With the fp8 +#' artifact, the transformer loads with CPU-resident fp8 weights that +#' stream to \code{device} during the forward pass. +#' @param device Character. Device for the small components (VAEs, +#' vocoder, connectors) and, with fp8, the transformer residents. +#' @param dtype Character. "bfloat16" (checkpoint native) or "float32". +#' @param transformer_device Character. Device for the transformer +#' weights when loading the plain (non-fp8) checkpoint. +#' @param components Character vector. Which components to load. +#' @param pin Logical. Pin fp8 host memory (fp8 artifact only). +#' @param attn_chunk Integer or NULL. Query-chunk size for attention +#' (see \code{\link{ltx23_set_attn_chunk}}). +#' @param phase_offload Logical. Load the small components (connectors, +#' VAEs, vocoder) to the CPU; the pipeline moves each onto the compute +#' device only for its phase. +#' @param verbose Logical. +#' +#' @return A list with the loaded modules and the checkpoint config, +#' class \code{ltx23_pipeline}. +#' +#' @export +ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", + dtype = "bfloat16", + transformer_device = "cpu", + components = c("dit", "connectors", "vae", "audio_vae", "vocoder"), + pin = TRUE, attn_chunk = NULL, + phase_offload = TRUE, verbose = TRUE) { + if (phase_offload) { + component_device <- "cpu" + } else { + component_device <- device + } + fp8 <- dir.exists(checkpoint_path) + ckpt <- if (fp8) { + ltx23_open_fp8_checkpoint(checkpoint_path) + } else { + ltx23_open_checkpoint(checkpoint_path) + } + if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { + # Must be set before the first CUDA allocation. NF4 runs the + # compiled block stack whose intermediates never become R + # handles: the native backend is faster there (expandable + # frees are page-unmaps, which is what made the decode gc + # storm expensive) and its 1280x704 fit is validated. The + # fp8/eager paths were only ever validated under + # expandable_segments (fragmentation from eager handle churn), + # so they keep it. + conf <- if (identical(ckpt$format, "nf4")) { + "backend:native" + } else { + "expandable_segments:True" + } + Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = conf) + } + if (device == "cuda") { + # Stop the allocator gc storm before the first CUDA op: the + # decode phases measured 86% of wall time in callback-driven + # R gc at the 0.20 default reserved rate. Only-if-unset, so an + # explicit user option wins. + if (identical(ckpt$format, "nf4")) { + footprint <- 12 + } else { + footprint <- 8 + } + ltx23_tune_gc(footprint_gb = footprint) + # Pre-warm the caching allocator with one large allocation + # freed straight into the pool: the phase onloads then carve + # from cached blocks instead of growing the pool one + # cudaMalloc per tensor (83.5s -> 4.6s for the NF4 + # transformer's first onload, measured). This is also the + # first CUDA op, so it runs after the options above are set. + tryCatch({ + warm <- torch::torch_empty(as.integer(footprint + 1) * 1e9, + dtype = torch::torch_uint8(), + device = "cuda") + rm(warm) + gc(verbose = FALSE) + }, error = function(e) invisible(NULL)) + } + groups <- ltx23_split_keys(ckpt$keys) + torch_dtype <- switch(dtype, bfloat16 = torch::torch_bfloat16(), + float16 = torch::torch_float16(), + float32 = torch::torch_float32(), + stop("Unsupported dtype: ", dtype)) + + pipe <- list(config = ckpt$config, checkpoint_path = ckpt$path) + + load_component <- function(name, module, map_key, dev) { + if (verbose) { + message("Loading ", name, " (", length(groups[[name]]), " tensors)...") + } + module$to(dtype = torch_dtype) + res <- ltx23_load_group(ckpt, groups[[name]], module, + map_key = map_key, verbose = verbose) + if (length(res$unmapped) || length(res$unfilled)) { + stop(name, ": incomplete load (", length(res$unmapped), " unmapped, ", + length(res$unfilled), " unfilled)") + } + module$to(device = dev) + module$eval() + module + } + + if ("dit" %in% components) { + if (fp8 && identical(ckpt$format, "nf4")) { + pipe$transformer <- ltx23_load_transformer_nf4( + ckpt, device = component_device, verbose = verbose + ) + } else if (fp8) { + pipe$transformer <- ltx23_load_transformer_fp8( + ckpt, device = component_device, pin = pin, verbose = verbose + ) + } else { + pipe$transformer <- load_component( + "dit", ltx23_transformer(), ltx23_map_dit_key, transformer_device + ) + } + if (!is.null(attn_chunk)) { + ltx23_set_attn_chunk(pipe$transformer, as.integer(attn_chunk)) + } + } + if ("connectors" %in% components) { + pipe$connectors <- load_component( + "connectors", ltx23_text_connectors(), ltx23_map_connector_key, component_device + ) + } + if ("vae" %in% components) { + pipe$vae <- load_component( + "vae", ltx23_video_vae(), ltx23_map_vae_key, component_device + ) + pipe$vae$enable_tiling() + } + if ("audio_vae" %in% components) { + pipe$audio_vae <- load_component( + "audio_vae", ltx23_audio_vae(), ltx23_map_audio_vae_key, component_device + ) + } + if ("vocoder" %in% components) { + # The vocoder is small and precision-sensitive: run in float32 + voc <- ltx23_vocoder_with_bwe() + voc$to(dtype = torch_dtype) + res <- ltx23_load_group(ckpt, groups$vocoder, voc, + map_key = ltx23_map_vocoder_key, verbose = verbose) + if (length(res$unmapped) || length(res$unfilled)) { + stop("vocoder: incomplete load") + } + voc$to(device = component_device, dtype = torch::torch_float32()) + voc$eval() + pipe$vocoder <- voc + } + + if (phase_offload && device == "cuda" && + isTRUE(getOption("diffuseR.pin_staging", FALSE))) { + # Page-lock every phase-offloaded component once so the + # per-render CPU<->GPU moves run at full PCIe rate (offload + # becomes a pointer swap; see staging_ltx23.R). Falls back + # silently per component if page-locking fails. + if (verbose) { + message("Pinning host staging buffers...") + } + staging <- list() + for (nm in intersect(c("transformer", "connectors", "vae", + "audio_vae", "vocoder"), names(pipe))) { + st <- .ltx23_pin_component(pipe[[nm]]) + if (!is.null(st)) { + staging[[nm]] <- st + } + } + if (length(staging)) { + pipe$staging <- staging + } + } + + structure(pipe, class = "ltx23_pipeline") +} + +#' Generate video (and audio) with LTX-2.3 +#' +#' Distilled text-to-video generation: encodes the prompt with Gemma3 + +#' connectors, denoises joint audio/video latents over the official +#' 8-step distilled schedule (no classifier-free guidance), decodes the +#' video with the causal VAE and the audio with the audio VAE + BWE +#' vocoder, and optionally muxes both into an MP4. +#' +#' @param prompt Character. The prompt. +#' @param pipeline An \code{ltx23_pipeline} from +#' \code{\link{ltx23_load_pipeline}}. +#' @param text_encoder,tokenizer Gemma3 model and tokenizer (or paths; +#' see \code{\link{load_gemma3_text_encoder}} and +#' \code{\link{gemma3_tokenizer}}). Ignored when \code{prompt_embeds} +#' is supplied. +#' @param prompt_embeds Optional precomputed list with +#' \code{prompt_embeds} (raw stacked Gemma3 states) and +#' \code{prompt_attention_mask}; bypasses the text encoder. +#' @param width,height Integers. Output resolution (multiples of 32). +#' @param num_frames Integer. 8k + 1 frames (e.g. 121). +#' @param frame_rate Numeric. Frames per second. +#' @param sigmas Numeric vector. Denoising schedule (default: official +#' distilled schedule; must end in 0). +#' @param guidance_scale Numeric. Only 1 (no CFG) is supported; the +#' distilled checkpoints are trained for CFG-free sampling. +#' @param seed Integer or NULL. +#' @param device Character. Compute device for the denoising loop. +#' @param dtype Character. Model compute dtype ("bfloat16" or "float32"). +#' @param filename Character or NULL. Output video path (.mp4). Audio is +#' muxed in when the av package is available. +#' @param max_sequence_length Integer. Text token length (multiple of 128). +#' @param decode_video,decode_audio Logicals. Decode the respective +#' latents (disable for latent-space work). +#' @param two_stage Logical. Generate at half resolution, upsample the +#' latents 2x spatially, and refine over the stage-2 schedule +#' (resolution must then be a multiple of 64; requires +#' \code{upsampler}). +#' @param upsampler An \code{\link{ltx23_latent_upsampler}} (see +#' \code{\link{ltx23_load_upsampler}}). +#' @param adain_factor Numeric. AdaIN blend of the upsampled latents +#' toward the stage-1 statistics (0 disables). +#' @param tone_map_compression Numeric in [0, 1]. Optional latent tone +#' mapping before stage 2. +#' @param phase_offload Logical. Move each small component to the compute +#' device only for its phase (text encoding, upsampling, decoding) and +#' back to the CPU afterwards, keeping the denoise phase as the sole +#' GPU tenant. +#' @param image Optional start image for image-to-video: a PNG/JPEG path +#' or an [H, W, 3] array in [0, 1]. The image conditions the first +#' frame; the rest of the video is generated (reference i2v). +#' @param condition_video Optional continuation source: a video path +#' (its trailing \code{conditioning_frames} frames are used) or an +#' [F, H, W, 3] array. The clip's tail becomes the frozen prefix of +#' the new video, so the output's first \code{conditioning_frames} +#' frames overlap the source (trim or crossfade when concatenating). +#' @param conditioning_frames Integer. Trailing pixel frames taken from +#' \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 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 +#' clean, frozen audio latents that the video attends to while +#' denoising, and the original samples are muxed into the output +#' (audio decoding is skipped). +#' @param verbose Logical. +#' +#' @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. +#' +#' @export +txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, + tokenizer = NULL, prompt_embeds = NULL, + width = 768L, height = 512L, num_frames = 121L, + frame_rate = 24, sigmas = ltx23_distilled_sigmas(), + guidance_scale = 1, seed = NULL, device = "cuda", + dtype = "bfloat16", filename = NULL, + max_sequence_length = 1024L, decode_video = TRUE, + decode_audio = TRUE, two_stage = FALSE, + upsampler = NULL, adain_factor = 1.0, + tone_map_compression = 0, phase_offload = TRUE, + image = NULL, condition_video = NULL, + conditioning_frames = 9L, cond_noise_scale = 0, + audio = NULL, verbose = TRUE) { + stopifnot(inherits(pipeline, "ltx23_pipeline")) + 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.") + } + conditioned <- !is.null(image) || !is.null(condition_video) + if (conditioned && two_stage) { + stop("Prefix conditioning with two_stage = TRUE is not supported yet.") + } + if (two_stage) { + if (is.null(upsampler)) { + stop("two_stage = TRUE requires an upsampler (see ltx23_load_upsampler).") + } + if (width %% 64L != 0L || height %% 64L != 0L) { + stop("width and height must be multiples of 64 for two-stage generation") + } + } + if (width %% 32L != 0L || height %% 32L != 0L) { + stop("width and height must be multiples of 32") + } + if ((num_frames - 1L) %% 8L != 0L) { + stop("num_frames must be 8k + 1 (e.g. 121)") + } + if (utils::tail(sigmas, 1) != 0) { + stop("sigmas must end with 0") + } + if (!is.null(seed)) { + torch::torch_manual_seed(seed) + } + + compute_dtype <- switch(dtype, bfloat16 = torch::torch_bfloat16(), + float16 = torch::torch_float16(), + float32 = torch::torch_float32(), + stop("Unsupported dtype: ", dtype)) + f32 <- torch::torch_float32() + + # --- Phase 1: text encoding ------------------------------------------------- + if (is.null(prompt_embeds)) { + if (is.null(text_encoder) || is.null(tokenizer)) { + stop("Supply text_encoder + tokenizer, or precomputed prompt_embeds.") + } + if (verbose) { + message("Encoding prompt...") + } + prompt_embeds <- encode_with_gemma3( + prompt, + model = text_encoder, tokenizer = tokenizer, + max_sequence_length = max_sequence_length, + device = if (is.character(text_encoder)) device else "cpu", + verbose = verbose + ) + } + + # Each phase is the sole GPU tenant: components move on for their + # phase and back off afterwards. Pipeline components are referred + # to by name so pinned staging (see staging_ltx23.R) can be used + # when the loader prepared it; plain modules (the upsampler) take + # the pageable path. + phase_offload <- phase_offload && device != "cpu" + staging <- pipeline$staging %||% list() + onload <- function(what) { + if (is.character(what)) { + module <- pipeline[[what]] + } else { + module <- what + } + if (phase_offload) { + if (is.character(what)) { + st <- staging[[what]] + } else { + st <- NULL + } + if (is.null(st)) { + module$to(device = device) + } else { + .ltx23_staged_onload(st, device) + } + } + module + } + offload <- function(what) { + if (is.character(what)) { + module <- pipeline[[what]] + } else { + module <- what + } + if (phase_offload) { + # Decode traces capture weight tensors; drop them so the + # module's GPU memory actually frees + .ltx23_release_vae_traces() + if (is.character(what)) { + st <- staging[[what]] + } else { + st <- NULL + } + if (is.null(st)) { + module$to(device = "cpu") + } else { + .ltx23_staged_offload(st) + } + # gc only -- NO cuda_empty_cache between phases: returning + # blocks to the driver forces the next phase to regrow the + # pool through cudaMalloc at ~15ms per allocation (~83s + # for the transformer's 5530 tensors, measured); the + # caching allocator reuses the freed blocks directly + gc(verbose = FALSE) + } + invisible(module) + } + + onload("connectors") + connectors_device <- pipeline$connectors$video_text_proj_in$weight$device + torch::with_no_grad({ + conn <- pipeline$connectors( + prompt_embeds$prompt_embeds$to(device = connectors_device, dtype = compute_dtype), + prompt_embeds$prompt_attention_mask$to(device = connectors_device) + ) + video_text_embeds <- conn$video_text_embedding$to(device = device) + audio_text_embeds <- conn$audio_text_embedding$to(device = device) + text_mask <- conn$attention_mask$to(device = device) + }) + rm(conn) + offload("connectors") + gc(verbose = FALSE) + + # --- Phase 2: latent preparation --------------------------------------------- + latent_frames <- (num_frames - 1L) %/% 8L + 1L + latent_height <- height %/% 32L + latent_width <- width %/% 32L + # Two-stage: stage 1 generates at half resolution + if (two_stage) { + s1_height <- latent_height %/% 2L + } else { + s1_height <- latent_height + } + if (two_stage) { + s1_width <- latent_width %/% 2L + } else { + s1_width <- latent_width + } + + # 25 audio latent frames per second: sampling_rate / hop / downsample + audio_num_frames <- as.integer(round(num_frames / frame_rate * 25)) + latent_mel_bins <- 16L # 64 mel bins / 4 + + # Prefix conditioning: encode the start image (i2v) or the tail of + # a previous clip (continuation) into normalized latents + cond_latents <- NULL + if (conditioned) { + vae <- onload("vae") + cond_frames <- if (!is.null(image)) { + ltx23_preprocess_frames(image, height, width) + } else { + arr <- if (is.character(condition_video)) { + ltx23_read_tail_frames(condition_video, conditioning_frames) + } else { + condition_video + } + ltx23_preprocess_frames(arr, height, width) + } + if (verbose) { + message(sprintf("Encoding %d conditioning frame(s)...", + cond_frames$shape[3])) + } + cond_latents <- ltx23_encode_video_frames(vae, cond_frames)$ + to(device = device) + offload("vae") + gc(verbose = FALSE) + } + + noise <- torch::torch_randn( + c(1L, 128L, latent_frames, s1_height, s1_width), + device = device, dtype = f32 + ) + conditioning_mask <- NULL + if (conditioned) { + prep <- ltx23_prepare_conditioned_latents( + cond_latents, latent_frames, s1_height, s1_width, + noise, cond_noise_scale = cond_noise_scale + ) + latents <- prep$latents + conditioning_mask <- prep$conditioning_mask + rm(prep, cond_latents) + } else { + latents <- ltx23_pack_video_latents(noise) + } + rm(noise) + + audio_conditioned <- !is.null(audio) + input_audio <- NULL + if (audio_conditioned) { + # Audio-driven generation: the user audio becomes clean, + # frozen audio latents; the video denoises while attending to + # them (lip sync). The output carries the original audio. + input_audio <- if (is.character(audio)) { + ltx23_read_audio(audio) + } else { + audio + } + audio_vae <- onload("audio_vae") + if (verbose) { + message("Encoding conditioning audio...") + } + audio_latents <- ltx23_encode_audio(audio_vae, input_audio, + audio_num_frames)$ + to(device = device) + offload("audio_vae") + gc(verbose = FALSE) + decode_audio <- FALSE + } else { + audio_latents <- torch::torch_randn( + c(1L, 8L, audio_num_frames, latent_mel_bins), + device = device, dtype = f32 + ) + audio_latents <- ltx23_pack_audio_latents(audio_latents) + } + + # --- Phase 3: denoising ------------------------------------------------------- + transformer <- onload("transformer") + if (verbose) message(sprintf("Denoising: %d steps at %dx%dx%d...", + length(sigmas) - 1L, width %/% (if (two_stage) 2L else 1L), + height %/% (if (two_stage) 2L else 1L), num_frames)) + + denoised <- .ltx23_denoise( + transformer, latents, audio_latents, sigmas, + video_text_embeds, audio_text_embeds, text_mask, + latent_frames, s1_height, s1_width, + audio_num_frames, frame_rate, + device, compute_dtype, verbose = verbose, + stage = if (two_stage) "stage 1 " else "", + conditioning_mask = conditioning_mask, + audio_conditioned = audio_conditioned + ) + latents <- denoised$latents + audio_latents <- denoised$audio_latents + + if (two_stage) { + vae <- pipeline$vae + if (verbose) { + message("Upsampling latents 2x...") + } + onload(upsampler) + torch::with_no_grad({ + up_device <- upsampler$final_conv$weight$device + up_dtype <- upsampler$final_conv$weight$dtype + s1_latents <- ltx23_unpack_video_latents( + latents, latent_frames, s1_height, s1_width + )$to(device = up_device) + # The upsampler operates on unnormalized latents + s1_latents <- ltx23_denormalize_latents( + s1_latents, vae$latents_mean, vae$latents_std + ) + up_latents <- upsampler(s1_latents$to(dtype = up_dtype))$to(dtype = f32) + if (adain_factor != 0) { + up_latents <- ltx23_adain_filter_latent( + up_latents, s1_latents$to(dtype = f32), adain_factor + ) + } + if (tone_map_compression > 0) { + up_latents <- ltx23_tone_map_latents(up_latents, tone_map_compression) + } + up_latents <- ltx23_normalize_latents( + up_latents, vae$latents_mean, vae$latents_std + ) + latents <- ltx23_pack_video_latents(up_latents$to(device = device)) + rm(s1_latents, up_latents) + }) + offload(upsampler) + gc(verbose = FALSE) + + # Re-noise BOTH modalities at the stage-2 entry sigma + s2_sigmas <- ltx23_stage2_distilled_sigmas() + noise_scale <- s2_sigmas[1] + torch::with_no_grad({ + latents <- torch::torch_randn_like(latents)$mul(noise_scale) + + latents$mul(1 - noise_scale) + audio_latents <- torch::torch_randn_like(audio_latents)$mul(noise_scale) + + audio_latents$mul(1 - noise_scale) + }) + + if (verbose) message(sprintf("Refining: %d steps at %dx%d...", + length(s2_sigmas) - 1L, width, height)) + denoised <- .ltx23_denoise( + transformer, latents, audio_latents, s2_sigmas, + video_text_embeds, audio_text_embeds, text_mask, + latent_frames, latent_height, latent_width, + audio_num_frames, frame_rate, + device, compute_dtype, verbose = verbose, stage = "stage 2 " + ) + latents <- denoised$latents + audio_latents <- denoised$audio_latents + } + + # The transformer and its dequant buffers are not needed past + # denoising; free the VRAM before the decoders claim it + offload("transformer") + ltx23_release_dequant_buffers() + + result <- list( + latents = latents, + audio_latents = audio_latents, + sample_rate = 48000L + ) + if (audio_conditioned) { + # The conditioning audio IS the soundtrack: carry the original + # samples through to the result and the mux + result$audio <- as.matrix(input_audio) + result$sample_rate <- 16000L + } + + # --- Phase 4: decoding ----------------------------------------------------------- + if (decode_video) { + if (verbose) { + message("Decoding video...") + } + phase_t0 <- Sys.time() + vae <- onload("vae") + vae_device <- vae$latents_mean$device + vae_dtype <- vae$decoder$conv_in$conv$weight$dtype + torch::with_no_grad({ + video_latents <- ltx23_unpack_video_latents( + latents, latent_frames, latent_height, latent_width + )$to(device = vae_device) + video_latents <- ltx23_denormalize_latents( + video_latents, vae$latents_mean, vae$latents_std + ) + video <- vae$decode(video_latents$to(dtype = vae_dtype)) + # [-1, 1] -> [0, 1], [B, 3, F, H, W] -> [F, H, W, 3] + video <- ((video$to(dtype = f32) / 2 + 0.5)$clamp(0, 1))[1,,,,] + video <- video$permute(c(2L, 3L, 4L, 1L))$cpu() + }) + if (verbose) { + message(sprintf(" video decode: %.1fs", + as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) + phase_t0 <- Sys.time() + } + result$video <- as.array(video) + rm(video, video_latents) + offload("vae") + gc(verbose = FALSE) + if (verbose) { + message(sprintf(" video to R array: %.1fs", + as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) + } + } + + if (decode_audio) { + if (verbose) { + message("Decoding audio...") + } + phase_t0 <- Sys.time() + audio_vae <- onload("audio_vae") + onload("vocoder") + av_device <- audio_vae$latents_mean$device + av_dtype <- audio_vae$decoder$conv_in$conv$weight$dtype + torch::with_no_grad({ + audio_packed <- .ltx23_denormalize_audio( + audio_latents$to(device = av_device), + audio_vae$latents_mean, audio_vae$latents_std + ) + audio_lat <- ltx23_unpack_audio_latents(audio_packed, latent_mel_bins) + mel <- audio_vae$decode(audio_lat$to(dtype = av_dtype)) + waveform <- .ltx23_traced_call( + pipeline$vocoder, + mel$to(dtype = torch::torch_float32()) + ) + waveform <- waveform[1,,]$cpu() + }) + result$audio <- as.matrix(as.array(waveform)) + rm(waveform, mel, audio_lat, audio_packed) + offload("audio_vae") + offload("vocoder") + gc(verbose = FALSE) + if (verbose) { + message(sprintf(" audio chain: %.1fs", + as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) + } + } + + if (!is.null(filename) && decode_video) { + phase_t0 <- Sys.time() + save_video_ltx23(result$video, filename, + fps = frame_rate, + audio = result$audio, + sample_rate = result$sample_rate, + verbose = verbose + ) + result$filename <- filename + if (verbose) { + message(sprintf(" encode + mux: %.1fs", + as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) + } + } + + invisible(result) +} + +#' Write a 16-bit PCM WAV file +#' +#' Minimal RIFF writer in base R. +#' +#' @param audio Numeric matrix [channels, samples] in [-1, 1]. +#' @param path Output path. +#' @param sample_rate Integer. +#' +#' @return Invisibly, the path. +#' +#' @export +write_wav <- function(audio, path, sample_rate = 48000L) { + if (is.null(dim(audio))) { + audio <- matrix(audio, nrow = 1L) + } + n_channels <- nrow(audio) + n_samples <- ncol(audio) + + # Interleave channels, scale to int16 + pcm <- as.integer(round(pmax(pmin(as.vector(audio), 1), -1) * 32767)) + byte_rate <- sample_rate * n_channels * 2L + data_size <- n_samples * n_channels * 2L + + con <- file(path, "wb") + on.exit(close(con), add = TRUE) + writeChar("RIFF", con, eos = NULL) + writeBin(as.integer(36L + data_size), con, size = 4, endian = "little") + writeChar("WAVEfmt ", con, eos = NULL) + writeBin(16L, con, size = 4, endian = "little") + writeBin(1L, con, size = 2, endian = "little") # PCM + writeBin(n_channels, con, size = 2, endian = "little") + writeBin(as.integer(sample_rate), con, size = 4, endian = "little") + writeBin(byte_rate, con, size = 4, endian = "little") + writeBin(n_channels * 2L, con, size = 2, endian = "little") # block align + writeBin(16L, con, size = 2, endian = "little") # bits per sample + writeChar("data", con, eos = NULL) + writeBin(data_size, con, size = 4, endian = "little") + writeBin(pcm, con, size = 2, endian = "little") + invisible(path) +} + +#' Save an LTX video (optionally with audio) to MP4 +#' +#' Uses the av package (Suggests) to encode frames and mux the audio +#' track. +#' +#' @param video Array [frames, height, width, 3] in [0, 1]. +#' @param filename Output path (.mp4). +#' @param fps Numeric. +#' @param audio Optional numeric matrix [channels, samples] in [-1, 1]. +#' @param sample_rate Integer. +#' @param verbose Logical. +#' +#' @return Invisibly, the filename. +#' +#' @export +save_video_ltx23 <- function(video, filename, fps = 24, audio = NULL, + sample_rate = 48000L, verbose = TRUE) { + if (!requireNamespace("av", quietly = TRUE)) { + stop("The av package is required to write video files.") + } + tmp_dir <- tempfile("ltx23_frames_") + dir.create(tmp_dir) + on.exit(unlink(tmp_dir, recursive = TRUE), add = TRUE) + + n_frames <- dim(video)[1] + frame_paths <- character(n_frames) + for (i in seq_len(n_frames)) { + frame_paths[i] <- file.path(tmp_dir, sprintf("frame_%05d.png", i)) + png::writePNG(video[i,,,], frame_paths[i]) + } + + audio_path <- NULL + if (!is.null(audio)) { + audio_path <- file.path(tmp_dir, "audio.wav") + write_wav(audio, audio_path, sample_rate = sample_rate) + } + + av::av_encode_video(frame_paths, output = filename, framerate = fps, + audio = audio_path, verbose = verbose) + if (verbose) { + message("Wrote ", filename) + } + invisible(filename) +} diff --git a/R/unet.R b/R/unet.R index f621f80..6ed7cf2 100644 --- a/R/unet.R +++ b/R/unet.R @@ -13,27 +13,27 @@ #' @return An nn_module representing the UNet #' @export unet_native <- torch::nn_module( - "UNetNative", - - initialize = function( - in_channels = 4L, - out_channels = 4L, - block_out_channels = c(320L, 640L, 1280L, 1280L), - layers_per_block = 2L, - cross_attention_dim = 1024L, - attention_head_dim = 64L - ) { - + "UNetNative", + + initialize = function( + in_channels = 4L, + out_channels = 4L, + block_out_channels = c(320L, 640L, 1280L, 1280L), + layers_per_block = 2L, + cross_attention_dim = 1024L, + attention_head_dim = 64L + ) { self$in_channels <- in_channels self$out_channels <- out_channels self$block_out_channels <- block_out_channels self$layers_per_block <- layers_per_block # Time embedding dimension - time_embed_dim <- block_out_channels[1] * 4L# 320 * 4 = 1280 + time_embed_dim <- block_out_channels[1] * 4L # 320 * 4 = 1280 # Input convolution - self$conv_in <- torch::nn_conv2d(in_channels, block_out_channels[1], 3L, padding = 1L) + self$conv_in <- torch::nn_conv2d(in_channels, block_out_channels[1], 3L, + padding = 1L) # Time embedding MLP self$time_embedding_linear_1 <- torch::nn_linear(block_out_channels[1], time_embed_dim) @@ -45,26 +45,26 @@ unet_native <- torch::nn_module( output_channel <- block_out_channels[1] for (i in seq_along(block_out_channels)) { - input_channel <- output_channel - output_channel <- block_out_channels[i] - is_final_block <- (i == length(block_out_channels)) - # Last block has no attention - has_attention <- !is_final_block - - down_block <- self$create_down_block( - input_channel, output_channel, time_embed_dim, - cross_attention_dim, attention_head_dim, - layers_per_block, - add_downsample = !is_final_block, - add_attention = has_attention - ) - self$down_blocks$append(down_block) + input_channel <- output_channel + output_channel <- block_out_channels[i] + is_final_block <- (i == length(block_out_channels)) + # Last block has no attention + has_attention <- !is_final_block + + down_block <- self$create_down_block( + input_channel, output_channel, time_embed_dim, + cross_attention_dim, attention_head_dim, + layers_per_block, + add_downsample = !is_final_block, + add_attention = has_attention + ) + self$down_blocks$append(down_block) } # Mid block mid_channels <- block_out_channels[length(block_out_channels)] self$mid_block <- self$create_mid_block( - mid_channels, time_embed_dim, cross_attention_dim, attention_head_dim + mid_channels, time_embed_dim, cross_attention_dim, attention_head_dim ) # Up blocks (reverse order) @@ -78,109 +78,109 @@ unet_native <- torch::nn_module( # down0: [in, r0, r1, ds], down1: [r0, r1, ds], down2: [r0, r1, ds], down3: [r0, r1] # This gives us: [320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280] skip_channels <- c( - block_out_channels[1], # conv_in - rep(block_out_channels[1], layers_per_block), # down0 resnets - block_out_channels[1], # down0 downsample - rep(block_out_channels[2], layers_per_block), # down1 resnets - block_out_channels[2], # down1 downsample - rep(block_out_channels[3], layers_per_block), # down2 resnets - block_out_channels[3], # down2 downsample - rep(block_out_channels[4], layers_per_block) # down3 resnets (no downsample) + block_out_channels[1], # conv_in + rep(block_out_channels[1], layers_per_block), # down0 resnets + block_out_channels[1], # down0 downsample + rep(block_out_channels[2], layers_per_block), # down1 resnets + block_out_channels[2], # down1 downsample + rep(block_out_channels[3], layers_per_block), # down2 resnets + block_out_channels[3], # down2 downsample + rep(block_out_channels[4], layers_per_block) # down3 resnets (no downsample) ) self$skip_channels <- rev(skip_channels) # Reverse for popping order for (i in seq_along(reversed_channels)) { - output_channel <- reversed_channels[i] - if (i == 1) { - prev_output <- mid_channels - } else { - prev_output <- reversed_channels[i - 1] - } - - is_final_block <- (i == length(reversed_channels)) - is_first_block <- (i == 1) - has_attention <- !is_first_block - - # Calculate per-resnet skip channels for this up_block - num_resnets <- layers_per_block + 1L - resnet_skip_channels <- integer(num_resnets) - for (j in seq_len(num_resnets)) { - skip_idx <- (i - 1) * num_resnets + j - resnet_skip_channels[j] <- self$skip_channels[skip_idx] - } - - up_block <- self$create_up_block_v2( - prev_output, output_channel, resnet_skip_channels, time_embed_dim, - cross_attention_dim, attention_head_dim, - add_upsample = !is_final_block, - add_attention = has_attention - ) - self$up_blocks$append(up_block) + output_channel <- reversed_channels[i] + if (i == 1) { + prev_output <- mid_channels + } else { + prev_output <- reversed_channels[i - 1] + } + + is_final_block <- (i == length(reversed_channels)) + is_first_block <- (i == 1) + has_attention <- !is_first_block + + # Calculate per-resnet skip channels for this up_block + num_resnets <- layers_per_block + 1L + resnet_skip_channels <- integer(num_resnets) + for (j in seq_len(num_resnets)) { + skip_idx <- (i - 1) * num_resnets + j + resnet_skip_channels[j] <- self$skip_channels[skip_idx] + } + + up_block <- self$create_up_block_v2( + prev_output, output_channel, resnet_skip_channels, time_embed_dim, + cross_attention_dim, attention_head_dim, + add_upsample = !is_final_block, + add_attention = has_attention + ) + self$up_blocks$append(up_block) } # Output self$conv_norm_out <- group_norm_32(block_out_channels[1]) self$conv_out <- torch::nn_conv2d(block_out_channels[1], out_channels, 3L, padding = 1L) - }, - - create_down_block = function( - in_channels, - out_channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim, - num_layers, - add_downsample = TRUE, - add_attention = TRUE - ) { +}, + + create_down_block = function( + in_channels, + out_channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim, + num_layers, + add_downsample = TRUE, + add_attention = TRUE + ) { # Create block as a simple nn_module DownBlock <- torch::nn_module( - "DownBlock", - initialize = function() { + "DownBlock", + initialize = function() { # ResNets self$resnets <- torch::nn_module_list() for (i in seq_len(num_layers)) { - if (i == 1) { - in_ch <- in_channels - } else { - in_ch <- out_channels - } - self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) + if (i == 1) { + in_ch <- in_channels + } else { + in_ch <- out_channels + } + self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) } # Attention blocks (optional) self$has_attentions <- add_attention if (add_attention) { - self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim - for (i in seq_len(num_layers)) { - self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, - depth = 1L, context_dim = cross_attention_dim) - ) - } + self$attentions <- torch::nn_module_list() + n_heads <- out_channels %/% attention_head_dim + for (i in seq_len(num_layers)) { + self$attentions$append( + SpatialTransformer(out_channels, n_heads, attention_head_dim, + depth = 1L, context_dim = cross_attention_dim) + ) + } } # Downsample self$has_downsamplers <- add_downsample if (add_downsample) { - self$downsamplers <- torch::nn_module_list() - self$downsamplers$append(Downsample2D(out_channels)) + self$downsamplers <- torch::nn_module_list() + self$downsamplers$append(Downsample2D(out_channels)) } - } + } ) DownBlock() - }, - - create_mid_block = function( - channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim - ) { +}, + + create_mid_block = function( + channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim + ) { MidBlock <- torch::nn_module( - "MidBlock", - initialize = function() { + "MidBlock", + initialize = function() { # Two resnets with attention in between self$resnets <- torch::nn_module_list() self$resnets$append(UNetResBlock(channels, channels, time_embed_dim)) @@ -190,78 +190,78 @@ unet_native <- torch::nn_module( n_heads <- channels %/% attention_head_dim self$attentions <- torch::nn_module_list() self$attentions$append( - SpatialTransformer(channels, n_heads, attention_head_dim, - depth = 1L, context_dim = cross_attention_dim) + SpatialTransformer(channels, n_heads, attention_head_dim, + depth = 1L, context_dim = cross_attention_dim) ) - } + } ) MidBlock() - }, - - create_up_block_v2 = function( - in_channels, - out_channels, - resnet_skip_channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim, - add_upsample = TRUE, - add_attention = TRUE - ) { +}, + + create_up_block_v2 = function( + in_channels, + out_channels, + resnet_skip_channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim, + add_upsample = TRUE, + add_attention = TRUE + ) { # resnet_skip_channels is a vector with one entry per resnet num_resnets <- length(resnet_skip_channels) UpBlock <- torch::nn_module( - "UpBlock", - initialize = function() { + "UpBlock", + initialize = function() { # ResNets (handle skip connections with variable channels) self$resnets <- torch::nn_module_list() for (i in seq_len(num_resnets)) { - # First resnet takes prev_output, rest take out_channels - if (i == 1) { - in_ch <- in_channels + resnet_skip_channels[i] - } else { - in_ch <- out_channels + resnet_skip_channels[i] - } - self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) + # First resnet takes prev_output, rest take out_channels + if (i == 1) { + in_ch <- in_channels + resnet_skip_channels[i] + } else { + in_ch <- out_channels + resnet_skip_channels[i] + } + self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) } # Attention blocks (optional) self$has_attentions <- add_attention if (add_attention) { - self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim - for (i in seq_len(num_resnets)) { - self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, - depth = 1L, context_dim = cross_attention_dim) - ) - } + self$attentions <- torch::nn_module_list() + n_heads <- out_channels %/% attention_head_dim + for (i in seq_len(num_resnets)) { + self$attentions$append( + SpatialTransformer(out_channels, n_heads, attention_head_dim, + depth = 1L, context_dim = cross_attention_dim) + ) + } } # Upsample self$has_upsamplers <- add_upsample if (add_upsample) { - self$upsamplers <- torch::nn_module_list() - self$upsamplers$append(Upsample2D(out_channels)) + self$upsamplers <- torch::nn_module_list() + self$upsamplers$append(Upsample2D(out_channels)) } - } + } ) UpBlock() - }, +}, - forward = function( - sample, - timestep, - encoder_hidden_states - ) { + forward = function( + sample, + timestep, + encoder_hidden_states + ) { # Get model dtype from weights model_dtype <- self$time_embedding_linear_1$weight$dtype # Time embedding (computed in float32, then cast to model dtype) # SD21 uses flip_sin_to_cos=FALSE, downscale_freq_shift=1 t_emb <- timestep_embedding(timestep, self$block_out_channels[1], - flip_sin_to_cos = FALSE, downscale_freq_shift = 1L) + flip_sin_to_cos = FALSE, downscale_freq_shift = 1L) t_emb <- t_emb$to(dtype = model_dtype) t_emb <- self$time_embedding_linear_1(t_emb) t_emb <- torch::nnf_silu(t_emb) @@ -275,20 +275,20 @@ unet_native <- torch::nn_module( # Down blocks for (i in seq_along(self$down_blocks)) { - block <- self$down_blocks[[i]] - - for (j in seq_along(block$resnets)) { - sample <- block$resnets[[j]](sample, t_emb) - if (block$has_attentions) { - sample <- block$attentions[[j]](sample, encoder_hidden_states) + block <- self$down_blocks[[i]] + + for (j in seq_along(block$resnets)) { + sample <- block$resnets[[j]](sample, t_emb) + if (block$has_attentions) { + sample <- block$attentions[[j]](sample, encoder_hidden_states) + } + down_block_res_samples <- c(down_block_res_samples, list(sample)) } - down_block_res_samples <- c(down_block_res_samples, list(sample)) - } - if (block$has_downsamplers) { - sample <- block$downsamplers[[1]](sample) - down_block_res_samples <- c(down_block_res_samples, list(sample)) - } + if (block$has_downsamplers) { + sample <- block$downsamplers[[1]](sample) + down_block_res_samples <- c(down_block_res_samples, list(sample)) + } } # Mid block @@ -298,25 +298,25 @@ unet_native <- torch::nn_module( # Up blocks for (i in seq_along(self$up_blocks)) { - block <- self$up_blocks[[i]] + block <- self$up_blocks[[i]] - for (j in seq_along(block$resnets)) { - # Pop skip connection - res_sample <- down_block_res_samples[[length(down_block_res_samples)]] - down_block_res_samples <- down_block_res_samples[- length(down_block_res_samples)] + for (j in seq_along(block$resnets)) { + # Pop skip connection + res_sample <- down_block_res_samples[[length(down_block_res_samples)]] + down_block_res_samples <- down_block_res_samples[-length(down_block_res_samples)] - # Concatenate - sample <- torch::torch_cat(list(sample, res_sample), dim = 2L) + # Concatenate + sample <- torch::torch_cat(list(sample, res_sample), dim = 2L) - sample <- block$resnets[[j]](sample, t_emb) - if (block$has_attentions) { - sample <- block$attentions[[j]](sample, encoder_hidden_states) + sample <- block$resnets[[j]](sample, t_emb) + if (block$has_attentions) { + sample <- block$attentions[[j]](sample, encoder_hidden_states) + } } - } - if (block$has_upsamplers) { - sample <- block$upsamplers[[1]](sample) - } + if (block$has_upsamplers) { + sample <- block$upsamplers[[1]](sample) + } } # Output @@ -325,7 +325,7 @@ unet_native <- torch::nn_module( sample <- self$conv_out(sample) sample - } +} ) #' Detect UNet architecture from TorchScript file @@ -334,36 +334,37 @@ unet_native <- torch::nn_module( #' @return List with architecture parameters #' @keywords internal detect_unet_architecture <- function(torchscript_path) { - ts_unet <- torch::jit_load(torchscript_path) - ts_params <- ts_unet$parameters - param_names <- names(ts_params) - - # Get block_out_channels from down_blocks resnets - # Use norm2 (not norm1) to get output channels - block_channels <- integer(0) - for (i in 0:3) { - key <- sprintf("unet.down_blocks.%d.resnets.0.norm2.weight", i) - if (key %in% param_names) { - block_channels <- c(block_channels, as.integer(ts_params[[key]]$shape[1])) + ts_unet <- torch::jit_load(torchscript_path) + ts_params <- ts_unet$parameters + param_names <- names(ts_params) + + # Get block_out_channels from down_blocks resnets + # Use norm2 (not norm1) to get output channels + block_channels <- integer(0) + for (i in 0:3) { + key <- sprintf("unet.down_blocks.%d.resnets.0.norm2.weight", i) + if (key %in% param_names) { + block_channels <- c(block_channels, + as.integer(ts_params[[key]]$shape[1])) + } } - } - - # Get cross_attention_dim from attn2.to_k - attn_key <- "unet.down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight" - cross_attention_dim <- as.integer(ts_params[[attn_key]]$shape[2]) - - # attention_head_dim is a fixed architectural choice for SD models - # SD21 and SDXL use 64 - attention_head_dim <- 64L - - list( - in_channels = 4L, - out_channels = 4L, - block_out_channels = block_channels, - layers_per_block = 2L, - cross_attention_dim = cross_attention_dim, - attention_head_dim = attention_head_dim - ) + + # Get cross_attention_dim from attn2.to_k + attn_key <- "unet.down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight" + cross_attention_dim <- as.integer(ts_params[[attn_key]]$shape[2]) + + # attention_head_dim is a fixed architectural choice for SD models + # SD21 and SDXL use 64 + attention_head_dim <- 64L + + list( + in_channels = 4L, + out_channels = 4L, + block_out_channels = block_channels, + layers_per_block = 2L, + cross_attention_dim = cross_attention_dim, + attention_head_dim = attention_head_dim + ) } #' Create native UNet from TorchScript @@ -375,34 +376,32 @@ detect_unet_architecture <- function(torchscript_path) { #' #' @return A native UNet module with loaded weights #' @export -unet_native_from_torchscript <- function( - torchscript_path, - verbose = TRUE -) { - # Detect architecture - arch <- detect_unet_architecture(torchscript_path) - - if (verbose) { - message("Detected UNet architecture:") - message(" block_out_channels: ", paste(arch$block_out_channels, collapse = ", ")) - message(" cross_attention_dim: ", arch$cross_attention_dim) - message(" attention_head_dim: ", arch$attention_head_dim) - } - - # Create native UNet - unet <- unet_native( - in_channels = arch$in_channels, - out_channels = arch$out_channels, - block_out_channels = arch$block_out_channels, - layers_per_block = arch$layers_per_block, - cross_attention_dim = arch$cross_attention_dim, - attention_head_dim = arch$attention_head_dim - ) - - # Load weights - load_unet_weights(unet, torchscript_path, verbose = verbose) - - unet +unet_native_from_torchscript <- function(torchscript_path, verbose = TRUE) { + # Detect architecture + arch <- detect_unet_architecture(torchscript_path) + + if (verbose) { + message("Detected UNet architecture:") + message(" block_out_channels: ", + paste(arch$block_out_channels, collapse = ", ")) + message(" cross_attention_dim: ", arch$cross_attention_dim) + message(" attention_head_dim: ", arch$attention_head_dim) + } + + # Create native UNet + unet <- unet_native( + in_channels = arch$in_channels, + out_channels = arch$out_channels, + block_out_channels = arch$block_out_channels, + layers_per_block = arch$layers_per_block, + cross_attention_dim = arch$cross_attention_dim, + attention_head_dim = arch$attention_head_dim + ) + + # Load weights + load_unet_weights(unet, torchscript_path, verbose = verbose) + + unet } #' Load weights from TorchScript UNet into native UNet @@ -413,69 +412,70 @@ unet_native_from_torchscript <- function( #' #' @return The native UNet with loaded weights (invisibly) #' @export -load_unet_weights <- function( - native_unet, - torchscript_path, - verbose = TRUE -) { - ts_unet <- torch::jit_load(torchscript_path) - ts_params <- ts_unet$parameters - - # Build weight mapping - remap_key <- function(key) { - # Strip unet. prefix - key <- sub("^unet\\.", "", key) - - # Time embedding: time_embedding.linear_1 -> time_embedding_linear_1 - key <- sub("^time_embedding\\.linear_1", "time_embedding_linear_1", key) - key <- sub("^time_embedding\\.linear_2", "time_embedding_linear_2", key) - - # Note: nn_sequential children use 0-indexing (matching Python/TorchScript) - # so no conversion needed for to_out.0, net.0, net.2 - - key - } - - loaded <- 0 - skipped <- 0 - unmapped <- character(0) - - torch::with_no_grad({ - for (ts_name in names(ts_params)) { - native_name <- remap_key(ts_name) - - if (native_name %in% names(native_unet$parameters)) { - ts_tensor <- ts_params[[ts_name]] - native_tensor <- native_unet$parameters[[native_name]] - - if (all(ts_tensor$shape == native_tensor$shape)) { - native_tensor$copy_(ts_tensor) - loaded <- loaded + 1 - } else if (verbose) { - message("Shape mismatch: ", native_name, - " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", - paste(as.integer(native_tensor$shape), collapse = "x"), ")") - skipped <- skipped + 1 - } - } else { - skipped <- skipped + 1 - unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) +load_unet_weights <- function(native_unet, torchscript_path, verbose = TRUE) { + ts_unet <- torch::jit_load(torchscript_path) + ts_params <- ts_unet$parameters + + # Build weight mapping + remap_key <- function(key) { + # Strip unet. prefix + key <- sub("^unet\\.", "", key) + + # Time embedding: time_embedding.linear_1 -> time_embedding_linear_1 + key <- sub("^time_embedding\\.linear_1", "time_embedding_linear_1", key) + key <- sub("^time_embedding\\.linear_2", "time_embedding_linear_2", key) + + # Note: nn_sequential children use 0-indexing (matching Python/TorchScript) + # so no conversion needed for to_out.0, net.0, net.2 + + key + } + + loaded <- 0 + skipped <- 0 + unmapped <- character(0) + + torch::with_no_grad({ + for (ts_name in names(ts_params)) { + native_name <- remap_key(ts_name) + + if (native_name %in% names(native_unet$parameters)) { + ts_tensor <- ts_params[[ts_name]] + native_tensor <- native_unet$parameters[[native_name]] + + if (all(ts_tensor$shape == native_tensor$shape)) { + native_tensor$copy_(ts_tensor) + loaded <- loaded + 1 + } else if (verbose) { + message("Shape mismatch: ", native_name, + " (", paste(as.integer(ts_tensor$shape), + collapse = "x"), " vs ", + paste(as.integer(native_tensor$shape), collapse = "x"), ")") + skipped <- skipped + 1 + } + } else { + skipped <- skipped + 1 + unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + } } - } }) - if (verbose) { - if (length(unmapped) > 0 && length(unmapped) <= 10) { - message("Unmapped parameters:") - for (u in unmapped) message(" ", u) - } else if (length(unmapped) > 10) { - message("Unmapped: ", length(unmapped), " parameters (showing first 10)") - for (u in head(unmapped, 10)) message(" ", u) + if (verbose) { + if (length(unmapped) > 0 && length(unmapped) <= 10) { + message("Unmapped parameters:") + for (u in unmapped) { + message(" ", u) + } + } else if (length(unmapped) > 10) { + message("Unmapped: ", length(unmapped), " parameters (showing first 10)") + for (u in head(unmapped, 10)) { + message(" ", u) + } + } + message("Loaded ", loaded, "/", loaded + skipped, " parameters") } - message("Loaded ", loaded, "/", loaded + skipped, " parameters") - } - invisible(native_unet) + invisible(native_unet) } #' Native SDXL UNet @@ -499,20 +499,19 @@ load_unet_weights <- function( #' @return An nn_module representing the SDXL UNet #' @export unet_sdxl_native <- torch::nn_module( - "UNetSDXLNative", - - initialize = function( - in_channels = 4L, - out_channels = 4L, - block_out_channels = c(320L, 640L, 1280L), - layers_per_block = 2L, - transformer_layers_per_block = c(0L, 2L, 10L), - cross_attention_dim = 2048L, - attention_head_dim = 64L, - addition_embed_dim = 1280L, - addition_time_embed_dim = 256L - ) { - + "UNetSDXLNative", + + initialize = function( + in_channels = 4L, + out_channels = 4L, + block_out_channels = c(320L, 640L, 1280L), + layers_per_block = 2L, + transformer_layers_per_block = c(0L, 2L, 10L), + cross_attention_dim = 2048L, + attention_head_dim = 64L, + addition_embed_dim = 1280L, + addition_time_embed_dim = 256L + ) { self$in_channels <- in_channels self$out_channels <- out_channels self$block_out_channels <- block_out_channels @@ -520,10 +519,11 @@ unet_sdxl_native <- torch::nn_module( self$transformer_layers_per_block <- transformer_layers_per_block # Time embedding dimension (same as SD21) - time_embed_dim <- block_out_channels[1] * 4L# 320 * 4 = 1280 + time_embed_dim <- block_out_channels[1] * 4L # 320 * 4 = 1280 # Input convolution - self$conv_in <- torch::nn_conv2d(in_channels, block_out_channels[1], 3L, padding = 1L) + self$conv_in <- torch::nn_conv2d(in_channels, block_out_channels[1], 3L, + padding = 1L) # Time embedding MLP self$time_embedding_linear_1 <- torch::nn_linear(block_out_channels[1], time_embed_dim) @@ -544,26 +544,26 @@ unet_sdxl_native <- torch::nn_module( output_channel <- block_out_channels[1] for (i in seq_along(block_out_channels)) { - input_channel <- output_channel - output_channel <- block_out_channels[i] - is_final_block <- (i == length(block_out_channels)) - transformer_depth <- transformer_layers_per_block[i] - - down_block <- self$create_down_block_sdxl( - input_channel, output_channel, time_embed_dim, - cross_attention_dim, attention_head_dim, - layers_per_block, transformer_depth, - add_downsample = !is_final_block - ) - self$down_blocks$append(down_block) + input_channel <- output_channel + output_channel <- block_out_channels[i] + is_final_block <- (i == length(block_out_channels)) + transformer_depth <- transformer_layers_per_block[i] + + down_block <- self$create_down_block_sdxl( + input_channel, output_channel, time_embed_dim, + cross_attention_dim, attention_head_dim, + layers_per_block, transformer_depth, + add_downsample = !is_final_block + ) + self$down_blocks$append(down_block) } # Mid block (SDXL: 10 transformer blocks) mid_channels <- block_out_channels[length(block_out_channels)] mid_transformer_depth <- transformer_layers_per_block[length(transformer_layers_per_block)] self$mid_block <- self$create_mid_block_sdxl( - mid_channels, time_embed_dim, cross_attention_dim, attention_head_dim, - mid_transformer_depth + mid_channels, time_embed_dim, cross_attention_dim, attention_head_dim, + mid_transformer_depth ) # Up blocks (SDXL: reverse order, 3 blocks) @@ -574,105 +574,105 @@ unet_sdxl_native <- torch::nn_module( # Calculate skip channels for SDXL (3 blocks) # down0: [in, r0, r1, ds], down1: [r0, r1, ds], down2: [r0, r1] skip_channels <- c( - block_out_channels[1], # conv_in - rep(block_out_channels[1], layers_per_block), # down0 resnets - block_out_channels[1], # down0 downsample - rep(block_out_channels[2], layers_per_block), # down1 resnets - block_out_channels[2], # down1 downsample - rep(block_out_channels[3], layers_per_block) # down2 resnets (no downsample) + block_out_channels[1], # conv_in + rep(block_out_channels[1], layers_per_block), # down0 resnets + block_out_channels[1], # down0 downsample + rep(block_out_channels[2], layers_per_block), # down1 resnets + block_out_channels[2], # down1 downsample + rep(block_out_channels[3], layers_per_block) # down2 resnets (no downsample) ) self$skip_channels <- rev(skip_channels) for (i in seq_along(reversed_channels)) { - output_channel <- reversed_channels[i] - if (i == 1) { - prev_output <- mid_channels - } else { - prev_output <- reversed_channels[i - 1] - } - - is_final_block <- (i == length(reversed_channels)) - transformer_depth <- reversed_transformer_depths[i] - - # Calculate per-resnet skip channels - num_resnets <- layers_per_block + 1L - resnet_skip_channels <- integer(num_resnets) - for (j in seq_len(num_resnets)) { - skip_idx <- (i - 1) * num_resnets + j - resnet_skip_channels[j] <- self$skip_channels[skip_idx] - } - - up_block <- self$create_up_block_sdxl( - prev_output, output_channel, resnet_skip_channels, time_embed_dim, - cross_attention_dim, attention_head_dim, transformer_depth, - add_upsample = !is_final_block - ) - self$up_blocks$append(up_block) + output_channel <- reversed_channels[i] + if (i == 1) { + prev_output <- mid_channels + } else { + prev_output <- reversed_channels[i - 1] + } + + is_final_block <- (i == length(reversed_channels)) + transformer_depth <- reversed_transformer_depths[i] + + # Calculate per-resnet skip channels + num_resnets <- layers_per_block + 1L + resnet_skip_channels <- integer(num_resnets) + for (j in seq_len(num_resnets)) { + skip_idx <- (i - 1) * num_resnets + j + resnet_skip_channels[j] <- self$skip_channels[skip_idx] + } + + up_block <- self$create_up_block_sdxl( + prev_output, output_channel, resnet_skip_channels, time_embed_dim, + cross_attention_dim, attention_head_dim, transformer_depth, + add_upsample = !is_final_block + ) + self$up_blocks$append(up_block) } # Output self$conv_norm_out <- group_norm_32(block_out_channels[1]) self$conv_out <- torch::nn_conv2d(block_out_channels[1], out_channels, 3L, padding = 1L) - }, - - create_down_block_sdxl = function( - in_channels, - out_channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim, - num_layers, - transformer_depth, - add_downsample = TRUE - ) { +}, + + create_down_block_sdxl = function( + in_channels, + out_channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim, + num_layers, + transformer_depth, + add_downsample = TRUE + ) { DownBlock <- torch::nn_module( - "DownBlockSDXL", - initialize = function() { + "DownBlockSDXL", + initialize = function() { # ResNets self$resnets <- torch::nn_module_list() for (i in seq_len(num_layers)) { - if (i == 1) { - in_ch <- in_channels - } else { - in_ch <- out_channels - } - self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) + if (i == 1) { + in_ch <- in_channels + } else { + in_ch <- out_channels + } + self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) } # Attention blocks (only if transformer_depth > 0) self$has_attentions <- transformer_depth > 0 if (self$has_attentions) { - self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim - for (i in seq_len(num_layers)) { - self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, - depth = transformer_depth, context_dim = cross_attention_dim) - ) - } + self$attentions <- torch::nn_module_list() + n_heads <- out_channels %/% attention_head_dim + for (i in seq_len(num_layers)) { + self$attentions$append( + SpatialTransformer(out_channels, n_heads, attention_head_dim, + depth = transformer_depth, context_dim = cross_attention_dim) + ) + } } # Downsample self$has_downsamplers <- add_downsample if (add_downsample) { - self$downsamplers <- torch::nn_module_list() - self$downsamplers$append(Downsample2D(out_channels)) + self$downsamplers <- torch::nn_module_list() + self$downsamplers$append(Downsample2D(out_channels)) } - } + } ) DownBlock() - }, - - create_mid_block_sdxl = function( - channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim, - transformer_depth - ) { +}, + + create_mid_block_sdxl = function( + channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim, + transformer_depth + ) { MidBlock <- torch::nn_module( - "MidBlockSDXL", - initialize = function() { + "MidBlockSDXL", + initialize = function() { self$resnets <- torch::nn_module_list() self$resnets$append(UNetResBlock(channels, channels, time_embed_dim)) self$resnets$append(UNetResBlock(channels, channels, time_embed_dim)) @@ -680,70 +680,70 @@ unet_sdxl_native <- torch::nn_module( n_heads <- channels %/% attention_head_dim self$attentions <- torch::nn_module_list() self$attentions$append( - SpatialTransformer(channels, n_heads, attention_head_dim, - depth = transformer_depth, context_dim = cross_attention_dim) + SpatialTransformer(channels, n_heads, attention_head_dim, + depth = transformer_depth, context_dim = cross_attention_dim) ) - } + } ) MidBlock() - }, - - create_up_block_sdxl = function( - in_channels, - out_channels, - resnet_skip_channels, - time_embed_dim, - cross_attention_dim, - attention_head_dim, - transformer_depth, - add_upsample = TRUE - ) { +}, + + create_up_block_sdxl = function( + in_channels, + out_channels, + resnet_skip_channels, + time_embed_dim, + cross_attention_dim, + attention_head_dim, + transformer_depth, + add_upsample = TRUE + ) { num_resnets <- length(resnet_skip_channels) UpBlock <- torch::nn_module( - "UpBlockSDXL", - initialize = function() { + "UpBlockSDXL", + initialize = function() { self$resnets <- torch::nn_module_list() for (i in seq_len(num_resnets)) { - if (i == 1) { - in_ch <- in_channels + resnet_skip_channels[i] - } else { - in_ch <- out_channels + resnet_skip_channels[i] - } - self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) + if (i == 1) { + in_ch <- in_channels + resnet_skip_channels[i] + } else { + in_ch <- out_channels + resnet_skip_channels[i] + } + self$resnets$append(UNetResBlock(in_ch, out_channels, time_embed_dim)) } # Attention blocks (only if transformer_depth > 0) self$has_attentions <- transformer_depth > 0 if (self$has_attentions) { - self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim - for (i in seq_len(num_resnets)) { - self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, - depth = transformer_depth, context_dim = cross_attention_dim) - ) - } + self$attentions <- torch::nn_module_list() + n_heads <- out_channels %/% attention_head_dim + for (i in seq_len(num_resnets)) { + self$attentions$append( + SpatialTransformer(out_channels, n_heads, attention_head_dim, + depth = transformer_depth, context_dim = cross_attention_dim) + ) + } } # Upsample self$has_upsamplers <- add_upsample if (add_upsample) { - self$upsamplers <- torch::nn_module_list() - self$upsamplers$append(Upsample2D(out_channels)) + self$upsamplers <- torch::nn_module_list() + self$upsamplers$append(Upsample2D(out_channels)) } - } + } ) UpBlock() - }, - - forward = function( - sample, - timestep, - encoder_hidden_states, - text_embeds, - time_ids - ) { +}, + + forward = function( + sample, + timestep, + encoder_hidden_states, + text_embeds, + time_ids + ) { # Accepts same signature as TorchScript: # unet(sample, timestep, encoder_hidden_states, text_embeds, time_ids) @@ -762,7 +762,7 @@ unet_sdxl_native <- torch::nn_module( # Uses same defaults as main time embedding time_ids_emb <- timestep_embedding(time_ids$flatten(), self$add_time_proj_dim) time_ids_emb <- time_ids_emb$to(dtype = model_dtype) - time_ids_emb <- time_ids_emb$reshape(c(text_embeds$shape[1], - 1L)) + time_ids_emb <- time_ids_emb$reshape(c(text_embeds$shape[1], -1L)) add_embeds <- torch::torch_cat(list(text_embeds, time_ids_emb), dim = 2L) add_emb <- self$add_embedding_linear_1(add_embeds) add_emb <- torch::nnf_silu(add_emb) @@ -779,20 +779,20 @@ unet_sdxl_native <- torch::nn_module( # Down blocks for (i in seq_along(self$down_blocks)) { - block <- self$down_blocks[[i]] - - for (j in seq_along(block$resnets)) { - sample <- block$resnets[[j]](sample, emb) - if (block$has_attentions) { - sample <- block$attentions[[j]](sample, encoder_hidden_states) + block <- self$down_blocks[[i]] + + for (j in seq_along(block$resnets)) { + sample <- block$resnets[[j]](sample, emb) + if (block$has_attentions) { + sample <- block$attentions[[j]](sample, encoder_hidden_states) + } + down_block_res_samples <- c(down_block_res_samples, list(sample)) } - down_block_res_samples <- c(down_block_res_samples, list(sample)) - } - if (block$has_downsamplers) { - sample <- block$downsamplers[[1]](sample) - down_block_res_samples <- c(down_block_res_samples, list(sample)) - } + if (block$has_downsamplers) { + sample <- block$downsamplers[[1]](sample) + down_block_res_samples <- c(down_block_res_samples, list(sample)) + } } # Mid block @@ -802,25 +802,25 @@ unet_sdxl_native <- torch::nn_module( # Up blocks for (i in seq_along(self$up_blocks)) { - block <- self$up_blocks[[i]] + block <- self$up_blocks[[i]] - for (j in seq_along(block$resnets)) { - # Pop skip connection - res_sample <- down_block_res_samples[[length(down_block_res_samples)]] - down_block_res_samples <- down_block_res_samples[- length(down_block_res_samples)] + for (j in seq_along(block$resnets)) { + # Pop skip connection + res_sample <- down_block_res_samples[[length(down_block_res_samples)]] + down_block_res_samples <- down_block_res_samples[-length(down_block_res_samples)] - # Concatenate - sample <- torch::torch_cat(list(sample, res_sample), dim = 2L) + # Concatenate + sample <- torch::torch_cat(list(sample, res_sample), dim = 2L) - sample <- block$resnets[[j]](sample, emb) - if (block$has_attentions) { - sample <- block$attentions[[j]](sample, encoder_hidden_states) + sample <- block$resnets[[j]](sample, emb) + if (block$has_attentions) { + sample <- block$attentions[[j]](sample, encoder_hidden_states) + } } - } - if (block$has_upsamplers) { - sample <- block$upsamplers[[1]](sample) - } + if (block$has_upsamplers) { + sample <- block$upsamplers[[1]](sample) + } } # Output @@ -829,7 +829,7 @@ unet_sdxl_native <- torch::nn_module( sample <- self$conv_out(sample) sample - } +} ) #' Detect SDXL UNet architecture from TorchScript file @@ -838,51 +838,54 @@ unet_sdxl_native <- torch::nn_module( #' @return List with architecture parameters #' @keywords internal detect_unet_sdxl_architecture <- function(torchscript_path) { - ts_unet <- torch::jit_load(torchscript_path, device = "cpu") - ts_params <- ts_unet$parameters - param_names <- names(ts_params) - - # Get block_out_channels from down_blocks resnets - block_channels <- integer(0) - for (i in 0:2) { # SDXL has 3 blocks - key <- sprintf("unet.down_blocks.%d.resnets.0.norm2.weight", i) - if (key %in% param_names) { - block_channels <- c(block_channels, as.integer(ts_params[[key]]$shape[1])) + ts_unet <- torch::jit_load(torchscript_path, device = "cpu") + ts_params <- ts_unet$parameters + param_names <- names(ts_params) + + # Get block_out_channels from down_blocks resnets + block_channels <- integer(0) + for (i in 0:2) { # SDXL has 3 blocks + key <- sprintf("unet.down_blocks.%d.resnets.0.norm2.weight", i) + if (key %in% param_names) { + block_channels <- c(block_channels, + as.integer(ts_params[[key]]$shape[1])) + } } - } - - # Get transformer depths per block - transformer_depths <- integer(length(block_channels)) - for (i in seq_along(block_channels)) { - depth <- 0L - for (j in 0:15) { - key <- sprintf("unet.down_blocks.%d.attentions.0.transformer_blocks.%d.attn1.to_q.weight", i - 1, j) - if (key %in% param_names) depth <- depth + 1L + + # Get transformer depths per block + transformer_depths <- integer(length(block_channels)) + for (i in seq_along(block_channels)) { + depth <- 0L + for (j in 0:15) { + key <- sprintf("unet.down_blocks.%d.attentions.0.transformer_blocks.%d.attn1.to_q.weight", i - 1, j) + if (key %in% param_names) { + depth <- depth + 1L + } + } + transformer_depths[i] <- depth } - transformer_depths[i] <- depth - } - - # Get cross_attention_dim from attn2.to_k (use block with attention) - cross_attention_dim <- NULL - for (i in 0:2) { - attn_key <- sprintf("unet.down_blocks.%d.attentions.0.transformer_blocks.0.attn2.to_k.weight", i) - if (attn_key %in% param_names) { - cross_attention_dim <- as.integer(ts_params[[attn_key]]$shape[2]) - break + + # Get cross_attention_dim from attn2.to_k (use block with attention) + cross_attention_dim <- NULL + for (i in 0:2) { + attn_key <- sprintf("unet.down_blocks.%d.attentions.0.transformer_blocks.0.attn2.to_k.weight", i) + if (attn_key %in% param_names) { + cross_attention_dim <- as.integer(ts_params[[attn_key]]$shape[2]) + break + } } - } - - list( - in_channels = 4L, - out_channels = 4L, - block_out_channels = block_channels, - layers_per_block = 2L, - transformer_layers_per_block = transformer_depths, - cross_attention_dim = cross_attention_dim, - attention_head_dim = 64L, - addition_embed_dim = 1280L, - addition_time_embed_dim = 256L - ) + + list( + in_channels = 4L, + out_channels = 4L, + block_out_channels = block_channels, + layers_per_block = 2L, + transformer_layers_per_block = transformer_depths, + cross_attention_dim = cross_attention_dim, + attention_head_dim = 64L, + addition_embed_dim = 1280L, + addition_time_embed_dim = 256L + ) } #' Create native SDXL UNet from TorchScript @@ -892,34 +895,33 @@ detect_unet_sdxl_architecture <- function(torchscript_path) { #' #' @return A native SDXL UNet module with loaded weights #' @export -unet_sdxl_native_from_torchscript <- function( - torchscript_path, - verbose = TRUE -) { - arch <- detect_unet_sdxl_architecture(torchscript_path) - - if (verbose) { - message("Detected SDXL UNet architecture:") - message(" block_out_channels: ", paste(arch$block_out_channels, collapse = ", ")) - message(" transformer_layers_per_block: ", paste(arch$transformer_layers_per_block, collapse = ", ")) - message(" cross_attention_dim: ", arch$cross_attention_dim) - } - - unet <- unet_sdxl_native( - in_channels = arch$in_channels, - out_channels = arch$out_channels, - block_out_channels = arch$block_out_channels, - layers_per_block = arch$layers_per_block, - transformer_layers_per_block = arch$transformer_layers_per_block, - cross_attention_dim = arch$cross_attention_dim, - attention_head_dim = arch$attention_head_dim, - addition_embed_dim = arch$addition_embed_dim, - addition_time_embed_dim = arch$addition_time_embed_dim - ) - - load_unet_sdxl_weights(unet, torchscript_path, verbose = verbose) - - unet +unet_sdxl_native_from_torchscript <- function(torchscript_path, + verbose = TRUE) { + arch <- detect_unet_sdxl_architecture(torchscript_path) + + if (verbose) { + message("Detected SDXL UNet architecture:") + message(" block_out_channels: ", + paste(arch$block_out_channels, collapse = ", ")) + message(" transformer_layers_per_block: ", paste(arch$transformer_layers_per_block, collapse = ", ")) + message(" cross_attention_dim: ", arch$cross_attention_dim) + } + + unet <- unet_sdxl_native( + in_channels = arch$in_channels, + out_channels = arch$out_channels, + block_out_channels = arch$block_out_channels, + layers_per_block = arch$layers_per_block, + transformer_layers_per_block = arch$transformer_layers_per_block, + cross_attention_dim = arch$cross_attention_dim, + attention_head_dim = arch$attention_head_dim, + addition_embed_dim = arch$addition_embed_dim, + addition_time_embed_dim = arch$addition_time_embed_dim + ) + + load_unet_sdxl_weights(unet, torchscript_path, verbose = verbose) + + unet } #' Load weights from TorchScript SDXL UNet into native SDXL UNet @@ -930,68 +932,69 @@ unet_sdxl_native_from_torchscript <- function( #' #' @return The native UNet with loaded weights (invisibly) #' @export -load_unet_sdxl_weights <- function( - native_unet, - torchscript_path, - verbose = TRUE -) { - ts_unet <- torch::jit_load(torchscript_path, device = "cpu") - ts_params <- ts_unet$parameters - - remap_key <- function(key) { - # Strip unet. prefix - key <- sub("^unet\\.", "", key) - - # Time embedding: time_embedding.linear_1 -> time_embedding_linear_1 - key <- sub("^time_embedding\\.linear_1", "time_embedding_linear_1", key) - key <- sub("^time_embedding\\.linear_2", "time_embedding_linear_2", key) - - # Add embedding: add_embedding.linear_1 -> add_embedding_linear_1 - key <- sub("^add_embedding\\.linear_1", "add_embedding_linear_1", key) - key <- sub("^add_embedding\\.linear_2", "add_embedding_linear_2", key) - - key - } - - loaded <- 0 - skipped <- 0 - unmapped <- character(0) - - torch::with_no_grad({ - for (ts_name in names(ts_params)) { - native_name <- remap_key(ts_name) - - if (native_name %in% names(native_unet$parameters)) { - ts_tensor <- ts_params[[ts_name]] - native_tensor <- native_unet$parameters[[native_name]] - - if (all(ts_tensor$shape == native_tensor$shape)) { - native_tensor$copy_(ts_tensor) - loaded <- loaded + 1 - } else if (verbose) { - message("Shape mismatch: ", native_name, - " (", paste(as.integer(ts_tensor$shape), collapse = "x"), " vs ", - paste(as.integer(native_tensor$shape), collapse = "x"), ")") - skipped <- skipped + 1 - } - } else { - skipped <- skipped + 1 - unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) +load_unet_sdxl_weights <- function(native_unet, torchscript_path, + verbose = TRUE) { + ts_unet <- torch::jit_load(torchscript_path, device = "cpu") + ts_params <- ts_unet$parameters + + remap_key <- function(key) { + # Strip unet. prefix + key <- sub("^unet\\.", "", key) + + # Time embedding: time_embedding.linear_1 -> time_embedding_linear_1 + key <- sub("^time_embedding\\.linear_1", "time_embedding_linear_1", key) + key <- sub("^time_embedding\\.linear_2", "time_embedding_linear_2", key) + + # Add embedding: add_embedding.linear_1 -> add_embedding_linear_1 + key <- sub("^add_embedding\\.linear_1", "add_embedding_linear_1", key) + key <- sub("^add_embedding\\.linear_2", "add_embedding_linear_2", key) + + key + } + + loaded <- 0 + skipped <- 0 + unmapped <- character(0) + + torch::with_no_grad({ + for (ts_name in names(ts_params)) { + native_name <- remap_key(ts_name) + + if (native_name %in% names(native_unet$parameters)) { + ts_tensor <- ts_params[[ts_name]] + native_tensor <- native_unet$parameters[[native_name]] + + if (all(ts_tensor$shape == native_tensor$shape)) { + native_tensor$copy_(ts_tensor) + loaded <- loaded + 1 + } else if (verbose) { + message("Shape mismatch: ", native_name, + " (", paste(as.integer(ts_tensor$shape), + collapse = "x"), " vs ", + paste(as.integer(native_tensor$shape), collapse = "x"), ")") + skipped <- skipped + 1 + } + } else { + skipped <- skipped + 1 + unmapped <- c(unmapped, paste0(ts_name, " -> ", native_name)) + } } - } }) - if (verbose) { - if (length(unmapped) > 0 && length(unmapped) <= 10) { - message("Unmapped parameters:") - for (u in unmapped) message(" ", u) - } else if (length(unmapped) > 10) { - message("Unmapped: ", length(unmapped), " parameters (showing first 10)") - for (u in head(unmapped, 10)) message(" ", u) + if (verbose) { + if (length(unmapped) > 0 && length(unmapped) <= 10) { + message("Unmapped parameters:") + for (u in unmapped) { + message(" ", u) + } + } else if (length(unmapped) > 10) { + message("Unmapped: ", length(unmapped), " parameters (showing first 10)") + for (u in head(unmapped, 10)) { + message(" ", u) + } + } + message("Loaded ", loaded, "/", loaded + skipped, " parameters") } - message("Loaded ", loaded, "/", loaded + skipped, " parameters") - } - invisible(native_unet) + invisible(native_unet) } - diff --git a/R/unet_modules.R b/R/unet_modules.R index 42bd703..7f4819a 100644 --- a/R/unet_modules.R +++ b/R/unet_modules.R @@ -8,7 +8,7 @@ NULL #' Group Normalization (32 groups) #' @keywords internal group_norm_32 <- function(channels) { - torch::nn_group_norm(32, channels, eps = 1e-5) + torch::nn_group_norm(32, channels, eps = 1e-5) } #' Sinusoidal Timestep Embedding @@ -21,54 +21,52 @@ group_norm_32 <- function(channels) { #' SD21 uses 1. With 0: exponent = log(10000) / half_dim. With 1: exponent = log(10000) / (half_dim - 1). #' @return Tensor (batch_size, dim) #' @keywords internal -timestep_embedding <- function( - timesteps, - dim, - flip_sin_to_cos = TRUE, - downscale_freq_shift = 0L -) { - half_dim <- dim %/% 2L - - # SDXL uses downscale_freq_shift=0, SD21 uses downscale_freq_shift=1 - emb_scale <- log(10000) / (half_dim - downscale_freq_shift) - - # Create frequency bands: [0, 1, ..., half_dim-1] - freqs <- torch::torch_exp(torch::torch_arange(0, half_dim - 1L) * - emb_scale) - freqs <- freqs$to(device = timesteps$device, dtype = torch::torch_float32()) - - # Expand and compute embeddings - # timesteps: [batch] -> [batch, 1] - # freqs: [half_dim] -> [1, half_dim] - timesteps_float <- timesteps$to(dtype = torch::torch_float32())$unsqueeze(2L) - args <- timesteps_float * freqs$unsqueeze(1L) - - # SDXL uses flip_sin_to_cos=True -> [cos, sin] - # SD21 uses flip_sin_to_cos=False -> [sin, cos] - if (flip_sin_to_cos) { - embedding <- torch::torch_cat(list(torch::torch_cos(args), torch::torch_sin(args)), dim = 2L) - } else { - embedding <- torch::torch_cat(list(torch::torch_sin(args), torch::torch_cos(args)), dim = 2L) - } - - # Pad if odd dimension - if (dim %% 2L == 1L) { - embedding <- torch::nnf_pad(embedding, c(0L, 1L, 0L, 0L)) - } - - embedding +timestep_embedding <- function(timesteps, dim, flip_sin_to_cos = TRUE, + downscale_freq_shift = 0L) { + half_dim <- dim %/% 2L + + # SDXL uses downscale_freq_shift=0, SD21 uses downscale_freq_shift=1 + emb_scale <- log(10000) / (half_dim - downscale_freq_shift) + + # Create frequency bands: [0, 1, ..., half_dim-1] + freqs <- torch::torch_exp(torch::torch_arange(0, half_dim - 1L) * -emb_scale) + freqs <- freqs$to(device = timesteps$device, dtype = torch::torch_float32()) + + # Expand and compute embeddings + # timesteps: [batch] -> [batch, 1] + # freqs: [half_dim] -> [1, half_dim] + timesteps_float <- timesteps$to(dtype = torch::torch_float32())$unsqueeze(2L) + args <- timesteps_float * freqs$unsqueeze(1L) + + # SDXL uses flip_sin_to_cos=True -> [cos, sin] + # SD21 uses flip_sin_to_cos=False -> [sin, cos] + if (flip_sin_to_cos) { + embedding <- torch::torch_cat(list(torch::torch_cos(args), + torch::torch_sin(args)), + dim = 2L) + } else { + embedding <- torch::torch_cat(list(torch::torch_sin(args), torch::torch_cos(args)), dim = 2L) + } + + # Pad if odd dimension + if (dim %% 2L == 1L) { + embedding <- torch::nnf_pad(embedding, c(0L, 1L, 0L, 0L)) + } + + embedding } #' ResNet Block for UNet #' #' @keywords internal UNetResBlock <- torch::nn_module( - "UNetResBlock", + "UNetResBlock", - initialize = function( - in_channels, - out_channels, - time_embed_dim - ) { + initialize = function( + in_channels, + out_channels, + time_embed_dim + ) { self$in_channels <- in_channels self$out_channels <- out_channels @@ -85,16 +83,16 @@ UNetResBlock <- torch::nn_module( # Skip connection (if channels differ) if (in_channels != out_channels) { - self$conv_shortcut <- torch::nn_conv2d(in_channels, out_channels, 1L) + self$conv_shortcut <- torch::nn_conv2d(in_channels, out_channels, 1L) } else { - self$conv_shortcut <- NULL + self$conv_shortcut <- NULL } - }, +}, - forward = function( - x, - temb - ) { + forward = function( + x, + temb + ) { h <- x # First block @@ -106,7 +104,7 @@ UNetResBlock <- torch::nn_module( temb_out <- torch::nnf_silu(temb) temb_out <- self$time_emb_proj(temb_out) # Expand to spatial dimensions: (B, C) -> (B, C, 1, 1) - temb_out <- temb_out$unsqueeze(- 1L)$unsqueeze(- 1L) + temb_out <- temb_out$unsqueeze(-1L)$unsqueeze(-1L) h <- h + temb_out # Second block @@ -116,60 +114,61 @@ UNetResBlock <- torch::nn_module( # Skip connection if (!is.null(self$conv_shortcut)) { - x <- self$conv_shortcut(x) + x <- self$conv_shortcut(x) } h + x - } +} ) #' Downsample Block #' @keywords internal Downsample2D <- torch::nn_module( - "Downsample2D", + "Downsample2D", - initialize = function(channels) { - self$conv <- torch::nn_conv2d(channels, channels, 3L, stride = 2L, padding = 1L) - }, + initialize = function(channels) { + self$conv <- torch::nn_conv2d(channels, channels, 3L, stride = 2L, + padding = 1L) +}, - forward = function(x) { + forward = function(x) { self$conv(x) - } +} ) #' Upsample Block #' @keywords internal Upsample2D <- torch::nn_module( - "Upsample2D", + "Upsample2D", - initialize = function(channels) { + initialize = function(channels) { self$conv <- torch::nn_conv2d(channels, channels, 3L, padding = 1L) - }, +}, - forward = function(x) { + forward = function(x) { # Upsample 2x then conv x <- torch::nnf_interpolate(x, scale_factor = 2, mode = "nearest") self$conv(x) - } +} ) #' Cross-Attention for UNet #' @keywords internal UNetCrossAttention <- torch::nn_module( - "UNetCrossAttention", - - initialize = function( - query_dim, - context_dim = NULL, - heads = 8L, - dim_head = 64L - ) { + "UNetCrossAttention", + + initialize = function( + query_dim, + context_dim = NULL, + heads = 8L, + dim_head = 64L + ) { self$heads <- heads self$dim_head <- dim_head inner_dim <- heads * dim_head context_dim <- context_dim %||% query_dim - self$scale <- dim_head ^ (- 0.5) + self$scale <- dim_head ^ (-0.5) # Query, Key, Value projections self$to_q <- torch::nn_linear(query_dim, inner_dim, bias = FALSE) @@ -177,16 +176,14 @@ UNetCrossAttention <- torch::nn_module( self$to_v <- torch::nn_linear(context_dim, inner_dim, bias = FALSE) # Output projection - self$to_out <- torch::nn_sequential( - torch::nn_linear(inner_dim, query_dim), - torch::nn_dropout(0) - ) - }, - - forward = function( - x, - context = NULL - ) { + self$to_out <- torch::nn_sequential(torch::nn_linear(inner_dim, query_dim), + torch::nn_dropout(0)) +}, + + forward = function( + x, + context = NULL + ) { if (is.null(context)) context <- x b <- x$shape[1] @@ -206,76 +203,74 @@ UNetCrossAttention <- torch::nn_module( # Scaled dot-product attention attn <- torch::torch_matmul(q, k$transpose(3L, 4L)) * self$scale - attn <- torch::nnf_softmax(attn, dim = - 1L) + attn <- torch::nnf_softmax(attn, dim = -1L) # Apply attention to values out <- torch::torch_matmul(attn, v) # Reshape back: (B, H, S, D) -> (B, S, H*D) - out <- out$transpose(2L, 3L)$contiguous()$view(c(b, seq_len, - 1L)) + out <- out$transpose(2L, 3L)$contiguous()$view(c(b, seq_len, -1L)) self$to_out(out) - } +} ) #' GEGLU Feedforward #' @keywords internal GEGLU <- torch::nn_module( - "GEGLU", + "GEGLU", - initialize = function( - dim_in, - dim_out - ) { + initialize = function( + dim_in, + dim_out + ) { self$proj <- torch::nn_linear(dim_in, dim_out * 2L) - }, +}, - forward = function(x) { + forward = function(x) { x_proj <- self$proj(x) - chunks <- torch::torch_chunk(x_proj, 2L, dim = - 1L) + chunks <- torch::torch_chunk(x_proj, 2L, dim = -1L) chunks[[1]] * torch::nnf_gelu(chunks[[2]]) - } +} ) #' FeedForward Network #' @keywords internal FeedForward <- torch::nn_module( - "FeedForward", + "FeedForward", - initialize = function( - dim, - mult = 4L - ) { + initialize = function( + dim, + mult = 4L + ) { inner_dim <- dim * mult - self$net <- torch::nn_sequential( - GEGLU(dim, inner_dim), - torch::nn_dropout(0), - torch::nn_linear(inner_dim, dim) - ) - }, - - forward = function(x) { + self$net <- torch::nn_sequential(GEGLU(dim, inner_dim), + torch::nn_dropout(0), + torch::nn_linear(inner_dim, dim)) +}, + + forward = function(x) { self$net(x) - } +} ) #' Basic Transformer Block #' @keywords internal BasicTransformerBlock <- torch::nn_module( - "BasicTransformerBlock", - - initialize = function( - dim, - n_heads, - d_head, - context_dim = NULL - ) { + "BasicTransformerBlock", + + initialize = function( + dim, + n_heads, + d_head, + context_dim = NULL + ) { # Self-attention self$attn1 <- UNetCrossAttention(dim, heads = n_heads, dim_head = d_head) # Cross-attention self$attn2 <- UNetCrossAttention(dim, context_dim = context_dim, - heads = n_heads, dim_head = d_head) + heads = n_heads, dim_head = d_head) # Feedforward self$ff <- FeedForward(dim) @@ -284,12 +279,12 @@ BasicTransformerBlock <- torch::nn_module( self$norm1 <- torch::nn_layer_norm(dim) self$norm2 <- torch::nn_layer_norm(dim) self$norm3 <- torch::nn_layer_norm(dim) - }, +}, - forward = function( - x, - context = NULL - ) { + forward = function( + x, + context = NULL + ) { # Self-attention x <- x + self$attn1(self$norm1(x)) @@ -300,21 +295,21 @@ BasicTransformerBlock <- torch::nn_module( x <- x + self$ff(self$norm3(x)) x - } +} ) #' Spatial Transformer (Attention Block) #' @keywords internal SpatialTransformer <- torch::nn_module( - "SpatialTransformer", - - initialize = function( - in_channels, - n_heads, - d_head, - depth = 1L, - context_dim = NULL - ) { + "SpatialTransformer", + + initialize = function( + in_channels, + n_heads, + d_head, + depth = 1L, + context_dim = NULL + ) { inner_dim <- n_heads * d_head self$in_channels <- in_channels @@ -328,16 +323,16 @@ SpatialTransformer <- torch::nn_module( # Transformer blocks self$transformer_blocks <- torch::nn_module_list() for (i in seq_len(depth)) { - self$transformer_blocks$append( - BasicTransformerBlock(inner_dim, n_heads, d_head, context_dim) - ) + self$transformer_blocks$append( + BasicTransformerBlock(inner_dim, n_heads, d_head, context_dim) + ) } - }, +}, - forward = function( - x, - context = NULL - ) { + forward = function( + x, + context = NULL + ) { b <- x$shape[1] c <- x$shape[2] h <- x$shape[3] @@ -356,7 +351,7 @@ SpatialTransformer <- torch::nn_module( # Transformer blocks for (i in seq_along(self$transformer_blocks)) { - x <- self$transformer_blocks[[i]](x, context) + x <- self$transformer_blocks[[i]](x, context) } # Project out @@ -367,6 +362,5 @@ SpatialTransformer <- torch::nn_module( # Residual x + x_in - } +} ) - diff --git a/R/upsampler_ltx23.R b/R/upsampler_ltx23.R new file mode 100644 index 0000000..9fff159 --- /dev/null +++ b/R/upsampler_ltx23.R @@ -0,0 +1,182 @@ +#' LTX-2.3 Spatial Latent Upsampler +#' +#' Fresh R port of the LTX latent upsampler from the diffusers reference +#' (Apache-2.0, pipelines/ltx2/latent_upsampler.py and +#' pipeline_ltx2_latent_upsample.py), with the LTX 2.3 configuration: +#' Conv3d ResBlock stages around a per-frame 2x pixel-shuffle spatial +#' upsampler (no rational resampler). Operates on unnormalized latents. +#' +#' @name upsampler_ltx23 +NULL + +# Conv -> GroupNorm(32) twice with the activation applied to the residual sum +ltx23_upsampler_res_block <- torch::nn_module( + "ltx23_upsampler_res_block", + initialize = function(channels, mid_channels = NULL) { + mid_channels <- mid_channels %||% channels + self$conv1 <- torch::nn_conv3d(channels, mid_channels, kernel_size = 3L, + padding = 1L) + self$norm1 <- torch::nn_group_norm(32L, mid_channels) + self$conv2 <- torch::nn_conv3d(mid_channels, channels, kernel_size = 3L, padding = 1L) + self$norm2 <- torch::nn_group_norm(32L, channels) +}, + forward = function(x) { + residual <- x + x <- torch::nnf_silu(self$norm1(self$conv1(x))) + x <- self$norm2(self$conv2(x)) + torch::nnf_silu(x + residual) +} +) + +#' LTX-2.3 latent upsampler model +#' +#' Latents [B, 128, F, H, W] -> [B, 128, F, 2H, 2W]. +#' +#' @param in_channels Integer. Latent channels. +#' @param mid_channels Integer. +#' @param num_blocks_per_stage Integer. +#' +#' @export +ltx23_latent_upsampler <- torch::nn_module( + "ltx23_latent_upsampler", + initialize = function( + in_channels = 128L, + mid_channels = 1024L, + num_blocks_per_stage = 4L + ) { + self$initial_conv <- torch::nn_conv3d(in_channels, mid_channels, + kernel_size = 3L, padding = 1L) + self$initial_norm <- torch::nn_group_norm(32L, mid_channels) + + self$res_blocks <- torch::nn_module_list(lapply(seq_len(num_blocks_per_stage), + function(i) ltx23_upsampler_res_block(mid_channels))) + + # Per-frame 2D conv + pixel shuffle (upsampler.0 in the checkpoint) + self$upsampler <- torch::nn_module_list(list( + torch::nn_conv2d(mid_channels, 4L * mid_channels, kernel_size = 3L, padding = 1L) + )) + + self$post_upsample_res_blocks <- torch::nn_module_list( + lapply(seq_len(num_blocks_per_stage), + function(i) ltx23_upsampler_res_block(mid_channels)) + ) + self$final_conv <- torch::nn_conv3d(mid_channels, in_channels, + kernel_size = 3L, padding = 1L) +}, + forward = function(hidden_states) { + batch_size <- hidden_states$shape[1] + + hidden_states <- torch::nnf_silu(self$initial_norm(self$initial_conv(hidden_states))) + for (i in seq_along(self$res_blocks)) { + hidden_states <- self$res_blocks[[i]](hidden_states) + } + + # [B, C, F, H, W] -> per-frame [B*F, C, H, W], 2x pixel shuffle, back + hidden_states <- hidden_states$permute(c(1L, 3L, 2L, 4L, 5L))$ + flatten(start_dim = 1L, end_dim = 2L) + hidden_states <- self$upsampler[[1]](hidden_states) + hidden_states <- torch::nnf_pixel_shuffle(hidden_states, 2L) + hidden_states <- hidden_states$unflatten(1L, c(batch_size, -1L))$ + permute(c(1L, 3L, 2L, 4L, 5L)) + + for (i in seq_along(self$post_upsample_res_blocks)) { + hidden_states <- self$post_upsample_res_blocks[[i]](hidden_states) + } + self$final_conv(hidden_states) +} +) + +#' Load the LTX-2.3 spatial upscaler weights +#' +#' The checkpoint keys match this module tree directly. +#' +#' @param path Path to e.g. \code{ltx-2.3-spatial-upscaler-x2-1.1.safetensors}. +#' @param device,dtype Placement for the loaded model. +#' @param verbose Logical. +#' +#' @return The loaded \code{ltx23_latent_upsampler}. +#' +#' @export +ltx23_load_upsampler <- function(path, device = "cuda", dtype = "bfloat16", + verbose = TRUE) { + torch_dtype <- switch(dtype, bfloat16 = torch::torch_bfloat16(), + float16 = torch::torch_float16(), + float32 = torch::torch_float32(), + stop("Unsupported dtype: ", dtype)) + handle <- safetensors::safetensors$new(path.expand(path), framework = "torch") + keys <- setdiff(handle$keys(), "__metadata__") + + model <- ltx23_latent_upsampler() + model$to(dtype = torch_dtype) + dests <- c(model$named_parameters(), model$named_buffers()) + + missing_dest <- setdiff(keys, names(dests)) + unfilled <- setdiff(names(dests), keys) + if (length(missing_dest) || length(unfilled)) { + stop("Upsampler load mismatch: ", length(missing_dest), " unmapped, ", + length(unfilled), " unfilled") + } + torch::with_no_grad({ + for (key in keys) { + dests[[key]]$copy_(handle$get_tensor(key)) + } + }) + model$to(device = device) + model$eval() + if (verbose) { + message("Loaded latent upsampler (", length(keys), " tensors)") + } + model +} + +#' Adaptive instance normalization between latent tensors +#' +#' Matches each (batch, channel) slice's mean/std to the reference +#' latents, blended by \code{factor} (cf. diffusers +#' \code{LTX2LatentUpsamplePipeline.adain_filter_latent}). +#' +#' @param latents Tensor [B, C, F, H, W]. +#' @param reference_latents Tensor with the target statistics. +#' @param factor Numeric blend in [-10, 10]; 0 is identity. +#' +#' @return Filtered latents. +#' +#' @export +ltx23_adain_filter_latent <- function(latents, reference_latents, + factor = 1.0) { + if (factor == 0) { + return(latents) + } + dims <- c(3L, 4L, 5L) + r_mean <- reference_latents$mean(dim = dims, keepdim = TRUE) + r_sd <- reference_latents$std(dim = dims, keepdim = TRUE) + i_mean <- latents$mean(dim = dims, keepdim = TRUE) + i_sd <- latents$std(dim = dims, keepdim = TRUE) + + result <- (latents - i_mean)$div(i_sd)$mul(r_sd) + r_mean + torch::torch_lerp(latents, result, factor) +} + +#' Sigmoid tone mapping for latents +#' +#' Compresses the latent dynamic range (cf. diffusers +#' \code{tone_map_latents}). \code{compression} 0 is identity, 1 is the +#' full effect. +#' +#' @param latents Tensor. +#' @param compression Numeric in [0, 1]. +#' +#' @return Tone-mapped latents. +#' +#' @export +ltx23_tone_map_latents <- function(latents, compression) { + if (compression == 0) { + return(latents) + } + stopifnot(compression >= 0, compression <= 1) + scale_factor <- compression * 0.75 + abs_latents <- torch::torch_abs(latents) + sigmoid_term <- torch::torch_sigmoid(abs_latents$add(-1)$mul(4 * scale_factor)) + scales <- sigmoid_term$mul(-0.8 * scale_factor)$add(1) + latents * scales +} diff --git a/R/vae_decoder.R b/R/vae_decoder.R index 58eae96..bf32dd5 100644 --- a/R/vae_decoder.R +++ b/R/vae_decoder.R @@ -4,24 +4,25 @@ #' @param out_channels Output channels #' @keywords internal VAEResnetBlock <- torch::nn_module( - "VAEResnetBlock", + "VAEResnetBlock", - initialize = function( - in_channels, - out_channels - ) { + initialize = function( + in_channels, + out_channels + ) { self$norm1 <- torch::nn_group_norm(32, in_channels, eps = 1e-6) - self$conv1 <- torch::nn_conv2d(in_channels, out_channels, kernel_size = 3, padding = 1) + self$conv1 <- torch::nn_conv2d(in_channels, out_channels, + kernel_size = 3, padding = 1) self$norm2 <- torch::nn_group_norm(32, out_channels, eps = 1e-6) self$conv2 <- torch::nn_conv2d(out_channels, out_channels, kernel_size = 3, padding = 1) # Shortcut if dimensions change if (in_channels != out_channels) { - self$conv_shortcut <- torch::nn_conv2d(in_channels, out_channels, kernel_size = 1) + self$conv_shortcut <- torch::nn_conv2d(in_channels, out_channels, kernel_size = 1) } - }, +}, - forward = function(x) { + forward = function(x) { h <- x h <- self$norm1(h) h <- torch::nnf_silu(h) @@ -33,11 +34,11 @@ VAEResnetBlock <- torch::nn_module( # Apply shortcut if it exists if (!is.null(self$conv_shortcut)) { - x <- self$conv_shortcut(x) + x <- self$conv_shortcut(x) } h + x - } +} ) #' VAE Attention Block @@ -46,20 +47,20 @@ VAEResnetBlock <- torch::nn_module( #' @param channels Number of channels #' @keywords internal VAEAttentionBlock <- torch::nn_module( - "VAEAttentionBlock", + "VAEAttentionBlock", - initialize = function(channels) { + initialize = function(channels) { self$group_norm <- torch::nn_group_norm(32, channels, eps = 1e-6) self$to_q <- torch::nn_linear(channels, channels) self$to_k <- torch::nn_linear(channels, channels) self$to_v <- torch::nn_linear(channels, channels) self$to_out <- torch::nn_module_list(list( - torch::nn_linear(channels, channels) - )) + torch::nn_linear(channels, channels) + )) self$channels <- channels - }, +}, - forward = function(x) { + forward = function(x) { residual <- x batch <- x$shape[1] channels <- x$shape[2] @@ -80,7 +81,7 @@ VAEAttentionBlock <- torch::nn_module( # Scaled dot-product attention scale <- 1.0 / sqrt(channels) attn <- torch::torch_bmm(q, k$transpose(2, 3)) * scale - attn <- torch::nnf_softmax(attn, dim = - 1) + attn <- torch::nnf_softmax(attn, dim = -1) out <- torch::torch_bmm(attn, v) # Project out @@ -90,7 +91,7 @@ VAEAttentionBlock <- torch::nn_module( out <- out$reshape(c(batch, height, width, channels))$permute(c(1, 4, 2, 3)) out + residual - } +} ) #' VAE Up Block @@ -101,50 +102,51 @@ VAEAttentionBlock <- torch::nn_module( #' @param add_upsample Whether to add upsampler #' @keywords internal VAEUpBlock <- torch::nn_module( - "VAEUpBlock", - - initialize = function( - in_channels, - out_channels, - num_resnets = 3, - add_upsample = TRUE - ) { + "VAEUpBlock", + + initialize = function( + in_channels, + out_channels, + num_resnets = 3, + add_upsample = TRUE + ) { self$resnets <- torch::nn_module_list() for (i in seq_len(num_resnets)) { - if (i == 1) { - res_in <- in_channels - } else { - res_in <- out_channels - } - self$resnets$append(VAEResnetBlock(res_in, out_channels)) + if (i == 1) { + res_in <- in_channels + } else { + res_in <- out_channels + } + self$resnets$append(VAEResnetBlock(res_in, out_channels)) } if (add_upsample) { - self$upsamplers <- torch::nn_module_list(list( - torch::nn_module( - "Upsampler", - initialize = function(channels) { - self$conv <- torch::nn_conv2d(channels, channels, kernel_size = 3, padding = 1) - }, - forward = function(x) { - x <- torch::nnf_interpolate(x, scale_factor = 2, mode = "nearest") - self$conv(x) - } - )(out_channels) - )) + self$upsamplers <- torch::nn_module_list(list( + torch::nn_module( + "Upsampler", + initialize = function(channels) { + self$conv <- torch::nn_conv2d(channels, channels, + kernel_size = 3, padding = 1) + }, + forward = function(x) { + x <- torch::nnf_interpolate(x, scale_factor = 2, mode = "nearest") + self$conv(x) + } + )(out_channels) + )) } - }, +}, - forward = function(x) { + forward = function(x) { for (i in seq_along(self$resnets)) { - x <- self$resnets[[i]](x) + x <- self$resnets[[i]](x) } if (!is.null(self$upsamplers)) { - x <- self$upsamplers[[1]](x) + x <- self$upsamplers[[1]](x) } x - } +} ) #' VAE Mid Block @@ -152,24 +154,22 @@ VAEUpBlock <- torch::nn_module( #' @param channels Number of channels #' @keywords internal VAEMidBlock <- torch::nn_module( - "VAEMidBlock", + "VAEMidBlock", - initialize = function(channels) { + initialize = function(channels) { self$resnets <- torch::nn_module_list(list( - VAEResnetBlock(channels, channels), - VAEResnetBlock(channels, channels) - )) - self$attentions <- torch::nn_module_list(list( - VAEAttentionBlock(channels) - )) - }, - - forward = function(x) { + VAEResnetBlock(channels, channels), + VAEResnetBlock(channels, channels) + )) + self$attentions <- torch::nn_module_list(list(VAEAttentionBlock(channels))) +}, + + forward = function(x) { x <- self$resnets[[1]](x) x <- self$attentions[[1]](x) x <- self$resnets[[2]](x) x - } +} ) #' Load weights from TorchScript decoder into native decoder @@ -180,41 +180,38 @@ VAEMidBlock <- torch::nn_module( #' #' @return The native decoder with loaded weights (invisibly) #' @export -load_decoder_weights <- function( - native_decoder, - torchscript_path, - verbose = TRUE -) { - ts_decoder <- torch::jit_load(torchscript_path) - ts_params <- ts_decoder$parameters - - loaded <- 0 - torch::with_no_grad({ - for (ts_name in names(ts_params)) { - # Strip dec. prefix - native_name <- sub("^dec\\.", "", ts_name) - - if (native_name %in% names(native_decoder$parameters)) { - ts_tensor <- ts_params[[ts_name]] - native_tensor <- native_decoder$parameters[[native_name]] - - if (all(ts_tensor$shape == native_tensor$shape)) { - native_tensor$copy_(ts_tensor) - loaded <- loaded + 1 - } else if (verbose) { - cat("Shape mismatch:", native_name, "\n") - } - } else if (verbose) { - cat("Missing param:", native_name, "\n") +load_decoder_weights <- function(native_decoder, torchscript_path, + verbose = TRUE) { + ts_decoder <- torch::jit_load(torchscript_path) + ts_params <- ts_decoder$parameters + + loaded <- 0 + torch::with_no_grad({ + for (ts_name in names(ts_params)) { + # Strip dec. prefix + native_name <- sub("^dec\\.", "", ts_name) + + if (native_name %in% names(native_decoder$parameters)) { + ts_tensor <- ts_params[[ts_name]] + native_tensor <- native_decoder$parameters[[native_name]] + + if (all(ts_tensor$shape == native_tensor$shape)) { + native_tensor$copy_(ts_tensor) + loaded <- loaded + 1 + } else if (verbose) { + cat("Shape mismatch:", native_name, "\n") + } + } else if (verbose) { + cat("Missing param:", native_name, "\n") + } } - } }) - if (verbose) { - cat("Loaded", loaded, "/", length(names(ts_params)), "parameters\n") - } + if (verbose) { + cat("Loaded", loaded, "/", length(names(ts_params)), "parameters\n") + } - invisible(native_decoder) + invisible(native_decoder) } #' Native VAE Decoder @@ -236,18 +233,19 @@ load_decoder_weights <- function( #' image <- decoder(latents) #' } vae_decoder_native <- torch::nn_module( - "VAEDecoderNative", + "VAEDecoderNative", - initialize = function( - latent_channels = 4, - out_channels = 3 - ) { + initialize = function( + latent_channels = 4, + out_channels = 3 + ) { # SDXL VAE decoder architecture: # Block channels: 512, 512, 256, 128 (reversed from encoder) block_channels <- c(512, 512, 256, 128) # Input conv: latent_channels -> 512 - self$conv_in <- torch::nn_conv2d(latent_channels, 512, kernel_size = 3, padding = 1) + self$conv_in <- torch::nn_conv2d(latent_channels, 512, kernel_size = 3, + padding = 1) # Mid block self$mid_block <- VAEMidBlock(512) @@ -270,9 +268,9 @@ vae_decoder_native <- torch::nn_module( # Output layers self$conv_norm_out <- torch::nn_group_norm(32, 128, eps = 1e-6) self$conv_out <- torch::nn_conv2d(128, out_channels, kernel_size = 3, padding = 1) - }, +}, - forward = function(x) { + forward = function(x) { # Input conv x <- self$conv_in(x) @@ -281,7 +279,7 @@ vae_decoder_native <- torch::nn_module( # Up blocks for (i in seq_along(self$up_blocks)) { - x <- self$up_blocks[[i]](x) + x <- self$up_blocks[[i]](x) } # Output @@ -290,6 +288,5 @@ vae_decoder_native <- torch::nn_module( x <- self$conv_out(x) x - } +} ) - diff --git a/R/vae_ltx2.R b/R/vae_ltx2.R deleted file mode 100644 index 9b9deeb..0000000 --- a/R/vae_ltx2.R +++ /dev/null @@ -1,913 +0,0 @@ -#' LTX2 Video VAE -#' -#' 3D causal VAE for video encoding/decoding as used in LTX-2. -#' Supports GPU-poor tiling for memory-efficient processing. -#' -#' @name vae_ltx2 -NULL - -#' LTX2 Video Encoder -#' -#' Encodes video frames into latent space with 3D causal convolutions. -#' -#' @param in_channels Integer. Input channels (typically 3 for RGB). -#' @param out_channels Integer. Latent channels. -#' @param block_out_channels Integer vector. Output channels per block. -#' @param spatio_temporal_scaling Logical vector. Whether each block downscales. -#' @param layers_per_block Integer vector. Number of layers per block. -#' @param downsample_type Character vector. Type of downsampling per block. -#' @param patch_size Integer. Spatial patch size. -#' @param patch_size_t Integer. Temporal patch size. -#' @param resnet_norm_eps Numeric. Epsilon for normalization. -#' @param is_causal Logical. Whether to use causal convolutions. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx2_video_encoder3d <- torch::nn_module( - "LTX2VideoEncoder3d", - - initialize = function( - in_channels = 3L, - out_channels = 128L, - block_out_channels = c(256L, 512L, 1024L, 2048L), - spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), - layers_per_block = c(4L, 6L, 6L, 2L, 2L), - downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - patch_size = 4L, - patch_size_t = 1L, - resnet_norm_eps = 1e-6, - is_causal = TRUE, - spatial_padding_mode = "zeros" - ) { - self$patch_size <- patch_size - self$patch_size_t <- patch_size_t - self$in_channels <- in_channels * patch_size ^ 2 - self$is_causal <- is_causal - - output_channel <- out_channels - - # Input convolution - self$conv_in <- ltx2_video_causal_conv3d( - in_channels = self$in_channels, - out_channels = output_channel, - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - - # Down blocks - num_blocks <- length(block_out_channels) - down_blocks <- list() - for (i in seq_len(num_blocks)) { - input_channel <- output_channel - output_channel <- block_out_channels[i] - - down_blocks[[i]] <- ltx2_video_down_block3d( - in_channels = input_channel, - out_channels = output_channel, - num_layers = layers_per_block[i], - resnet_eps = resnet_norm_eps, - spatio_temporal_scale = spatio_temporal_scaling[i], - downsample_type = downsample_type[i], - spatial_padding_mode = spatial_padding_mode - ) - } - self$down_blocks <- torch::nn_module_list(down_blocks) - - # Mid block - self$mid_block <- ltx2_video_mid_block3d( - in_channels = output_channel, - num_layers = layers_per_block[length(layers_per_block)], - resnet_eps = resnet_norm_eps, - spatial_padding_mode = spatial_padding_mode - ) - - # Output - self$norm_out <- per_channel_rms_norm() - self$conv_act <- torch::nn_silu() - self$conv_out <- ltx2_video_causal_conv3d( - in_channels = output_channel, - out_channels = out_channels + 1L, # +1 for log variance - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - }, - - forward = function( - hidden_states, - causal = NULL - ) { - p <- self$patch_size - p_t <- self$patch_size_t - - batch_size <- hidden_states$shape[1] - num_channels <- hidden_states$shape[2] - num_frames <- hidden_states$shape[3] - height <- hidden_states$shape[4] - width <- hidden_states$shape[5] - - post_patch_num_frames <- num_frames %/% p_t - post_patch_height <- height %/% p - post_patch_width <- width %/% p - - if (is.null(causal)) causal <- self$is_causal - - # Patchify: reshape to separate patches - hidden_states <- hidden_states$reshape(c( - batch_size, num_channels, - post_patch_num_frames, p_t, - post_patch_height, p, - post_patch_width, p - )) - # Permute to channel-first for patches - # [B, C, F', p_t, H', p, W', p] -> [B, C, p_t, p, p, F', H', W'] - hidden_states <- hidden_states$permute(c(1, 2, 4, 8, 6, 3, 5, 7)) - hidden_states <- hidden_states$flatten(start_dim = 2, end_dim = 5) - - hidden_states <- self$conv_in(hidden_states, causal = causal) - - for (i in seq_along(self$down_blocks)) { - hidden_states <- self$down_blocks[[i]](hidden_states, causal = causal) - } - - hidden_states <- self$mid_block(hidden_states, causal = causal) - - hidden_states <- self$norm_out(hidden_states) - hidden_states <- self$conv_act(hidden_states) - hidden_states <- self$conv_out(hidden_states, causal = causal) - - # Duplicate last channel for mean/logvar split - last_channel <- hidden_states[, - 1,,,]$unsqueeze(2) - last_channel <- last_channel$`repeat`(c(1L, hidden_states$shape[2] - 2L, 1L, 1L, 1L)) - hidden_states <- torch::torch_cat(list(hidden_states, last_channel), dim = 2) - - hidden_states - } -) - -#' LTX2 Video Decoder -#' -#' Decodes latent representations back to video frames. -#' -#' @param in_channels Integer. Latent channels. -#' @param out_channels Integer. Output channels (typically 3 for RGB). -#' @param block_out_channels Integer vector. Output channels per block. -#' @param spatio_temporal_scaling Logical vector. Whether each block upscales. -#' @param layers_per_block Integer vector. Number of layers per block. -#' @param patch_size Integer. Spatial patch size. -#' @param patch_size_t Integer. Temporal patch size. -#' @param resnet_norm_eps Numeric. Epsilon for normalization. -#' @param is_causal Logical. Whether to use causal convolutions. -#' @param inject_noise Logical vector. Whether to inject noise per block. -#' @param timestep_conditioning Logical. Whether to use timestep conditioning. -#' @param upsample_residual Logical vector. Whether upsamplers use residual. -#' @param upsample_factor Integer vector. Channel upscale factors. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx2_video_decoder3d <- torch::nn_module( - "LTX2VideoDecoder3d", - - initialize = function( - in_channels = 128L, - out_channels = 3L, - block_out_channels = c(256L, 512L, 1024L), - spatio_temporal_scaling = c(TRUE, TRUE, TRUE), - layers_per_block = c(5L, 5L, 5L, 5L), - patch_size = 4L, - patch_size_t = 1L, - resnet_norm_eps = 1e-6, - is_causal = FALSE, - inject_noise = c(FALSE, FALSE, FALSE, FALSE), - timestep_conditioning = FALSE, - upsample_residual = c(TRUE, TRUE, TRUE), - upsample_factor = c(2L, 2L, 2L), - spatial_padding_mode = "reflect" - ) { - self$patch_size <- patch_size - self$patch_size_t <- patch_size_t - self$out_channels <- out_channels * patch_size ^ 2 - self$is_causal <- is_causal - - # Reverse orders for decoder - block_out_channels <- rev(block_out_channels) - spatio_temporal_scaling <- rev(spatio_temporal_scaling) - layers_per_block <- rev(layers_per_block) - inject_noise <- rev(inject_noise) - upsample_residual <- rev(upsample_residual) - upsample_factor <- rev(upsample_factor) - - output_channel <- block_out_channels[1] - - # Input convolution - self$conv_in <- ltx2_video_causal_conv3d( - in_channels = in_channels, - out_channels = output_channel, - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - - # Mid block - self$mid_block <- ltx2_video_mid_block3d( - in_channels = output_channel, - num_layers = layers_per_block[1], - resnet_eps = resnet_norm_eps, - inject_noise = inject_noise[1], - timestep_conditioning = timestep_conditioning, - spatial_padding_mode = spatial_padding_mode - ) - - # Up blocks - num_blocks <- length(block_out_channels) - up_blocks <- list() - for (i in seq_len(num_blocks)) { - input_channel <- output_channel %/% upsample_factor[i] - output_channel <- block_out_channels[i] %/% upsample_factor[i] - - up_blocks[[i]] <- ltx2_video_up_block3d( - in_channels = input_channel, - out_channels = output_channel, - num_layers = layers_per_block[i + 1], - resnet_eps = resnet_norm_eps, - spatio_temporal_scale = spatio_temporal_scaling[i], - inject_noise = if (i + 1 <= length(inject_noise)) inject_noise[i + 1] else FALSE, - timestep_conditioning = timestep_conditioning, - upsample_residual = upsample_residual[i], - upscale_factor = upsample_factor[i], - spatial_padding_mode = spatial_padding_mode - ) - } - self$up_blocks <- torch::nn_module_list(up_blocks) - - # Output - self$norm_out <- per_channel_rms_norm() - self$conv_act <- torch::nn_silu() - self$conv_out <- ltx2_video_causal_conv3d( - in_channels = output_channel, - out_channels = self$out_channels, - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - - # Timestep embedding (optional) - self$time_embedder <- NULL - self$scale_shift_table <- NULL - self$timestep_scale_multiplier <- NULL - }, - - forward = function( - hidden_states, - temb = NULL, - causal = NULL - ) { - if (is.null(causal)) causal <- self$is_causal - - hidden_states <- self$conv_in(hidden_states, causal = causal) - - if (!is.null(self$timestep_scale_multiplier) && !is.null(temb)) { - temb <- temb * self$timestep_scale_multiplier - } - - hidden_states <- self$mid_block(hidden_states, temb, causal = causal) - - for (i in seq_along(self$up_blocks)) { - hidden_states <- self$up_blocks[[i]](hidden_states, temb, causal = causal) - } - - hidden_states <- self$norm_out(hidden_states) - hidden_states <- self$conv_act(hidden_states) - hidden_states <- self$conv_out(hidden_states, causal = causal) - - # Unpatchify: reshape to original spatial dims - p <- self$patch_size - p_t <- self$patch_size_t - - batch_size <- hidden_states$shape[1] - num_channels <- hidden_states$shape[2] - num_frames <- hidden_states$shape[3] - height <- hidden_states$shape[4] - width <- hidden_states$shape[5] - - # [B, C*p_t*p*p, F, H, W] -> [B, C, p_t, p, p, F, H, W] - hidden_states <- hidden_states$reshape(c( - batch_size, - 1, p_t, p, p, num_frames, height, width - )) - # Permute: [B, C, p_t, p, p, F, H, W] -> [B, C, F, p_t, H, p, W, p] - hidden_states <- hidden_states$permute(c(1, 2, 6, 3, 7, 5, 8, 4)) - # Flatten to full resolution - hidden_states <- hidden_states$flatten(start_dim = 7, end_dim = 8) # W*p - hidden_states <- hidden_states$flatten(start_dim = 5, end_dim = 6) # H*p - hidden_states <- hidden_states$flatten(start_dim = 3, end_dim = 4) # F*p_t - - hidden_states - } -) - -#' Diagonal Gaussian Distribution -#' -#' Represents a diagonal Gaussian distribution for VAE latents. -#' -#' @param parameters Tensor of concatenated mean and log variance. -#' @export -diagonal_gaussian_distribution <- function(parameters) { - # Split parameters into mean and logvar - chunk_dim <- 2# Channel dimension - mean_logvar <- parameters$chunk(2, dim = chunk_dim) - mean <- mean_logvar[[1]] - logvar <- mean_logvar[[2]] - - # Clamp logvar for numerical stability - logvar <- logvar$clamp(min = - 30.0, max = 20.0) - std <- torch::torch_exp(0.5 * logvar) - var <- torch::torch_exp(logvar) - - list( - mean = mean, - logvar = logvar, - std = std, - var = var, - sample = function(generator = NULL) { - sample <- torch::torch_randn_like(mean) - mean + std * sample - }, - mode = function() mean - ) -} - -#' LTX2 Video VAE -#' -#' Full VAE with encoder and decoder, supporting tiled encoding/decoding -#' for GPU-poor memory management. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer. Output channels. -#' @param latent_channels Integer. Latent space channels. -#' @param block_out_channels Integer vector. Encoder block channels. -#' @param decoder_block_out_channels Integer vector. Decoder block channels. -#' @param layers_per_block Integer vector. Encoder layers per block. -#' @param decoder_layers_per_block Integer vector. Decoder layers per block. -#' @param spatio_temporal_scaling Logical vector. Encoder scaling. -#' @param decoder_spatio_temporal_scaling Logical vector. Decoder scaling. -#' @param decoder_inject_noise Logical vector. Noise injection per decoder block. -#' @param downsample_type Character vector. Downsampling types. -#' @param upsample_residual Logical vector. Upsampler residual flags. -#' @param upsample_factor Integer vector. Upsampler factors. -#' @param timestep_conditioning Logical. Whether to use timestep conditioning. -#' @param patch_size Integer. Spatial patch size. -#' @param patch_size_t Integer. Temporal patch size. -#' @param resnet_norm_eps Numeric. Normalization epsilon. -#' @param scaling_factor Numeric. Latent scaling factor. -#' @param encoder_causal Logical. Encoder causality. -#' @param decoder_causal Logical. Decoder causality. -#' @param encoder_spatial_padding_mode Character. Encoder padding mode. -#' @param decoder_spatial_padding_mode Character. Decoder padding mode. -#' @export -ltx2_video_vae <- torch::nn_module( - "AutoencoderKLLTX2Video", - - initialize = function( - in_channels = 3L, - out_channels = 3L, - latent_channels = 128L, - block_out_channels = c(256L, 512L, 1024L, 2048L), - decoder_block_out_channels = c(256L, 512L, 1024L), - layers_per_block = c(4L, 6L, 6L, 2L, 2L), - decoder_layers_per_block = c(5L, 5L, 5L, 5L), - spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), - decoder_spatio_temporal_scaling = c(TRUE, TRUE, TRUE), - decoder_inject_noise = c(FALSE, FALSE, FALSE, FALSE), - downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - upsample_residual = c(TRUE, TRUE, TRUE), - upsample_factor = c(2L, 2L, 2L), - timestep_conditioning = FALSE, - patch_size = 4L, - patch_size_t = 1L, - resnet_norm_eps = 1e-6, - scaling_factor = 1.0, - encoder_causal = TRUE, - decoder_causal = TRUE, - encoder_spatial_padding_mode = "zeros", - decoder_spatial_padding_mode = "reflect" - ) { - self$encoder <- ltx2_video_encoder3d( - in_channels = in_channels, - out_channels = latent_channels, - block_out_channels = block_out_channels, - spatio_temporal_scaling = spatio_temporal_scaling, - layers_per_block = layers_per_block, - downsample_type = downsample_type, - patch_size = patch_size, - patch_size_t = patch_size_t, - resnet_norm_eps = resnet_norm_eps, - is_causal = encoder_causal, - spatial_padding_mode = encoder_spatial_padding_mode - ) - - self$decoder <- ltx2_video_decoder3d( - in_channels = latent_channels, - out_channels = out_channels, - block_out_channels = decoder_block_out_channels, - spatio_temporal_scaling = decoder_spatio_temporal_scaling, - layers_per_block = decoder_layers_per_block, - patch_size = patch_size, - patch_size_t = patch_size_t, - resnet_norm_eps = resnet_norm_eps, - is_causal = decoder_causal, - inject_noise = decoder_inject_noise, - timestep_conditioning = timestep_conditioning, - upsample_residual = upsample_residual, - upsample_factor = upsample_factor, - spatial_padding_mode = decoder_spatial_padding_mode - ) - - # Latent normalization buffers - self$latents_mean <- torch::nn_buffer(torch::torch_zeros(latent_channels)) - self$latents_std <- torch::nn_buffer(torch::torch_ones(latent_channels)) - - # Compression ratios - self$spatial_compression_ratio <- patch_size * 2 ^ sum(spatio_temporal_scaling) - self$temporal_compression_ratio <- patch_size_t * 2 ^ sum(spatio_temporal_scaling) - - self$scaling_factor <- scaling_factor - - # Tiling configuration (GPU-poor support) - self$use_slicing <- FALSE - self$use_tiling <- FALSE - self$use_framewise_encoding <- FALSE - self$use_framewise_decoding <- FALSE - - self$num_sample_frames_batch_size <- 16L - self$num_latent_frames_batch_size <- 2L - - self$tile_sample_min_height <- 512L - self$tile_sample_min_width <- 512L - self$tile_sample_min_num_frames <- 16L - - self$tile_sample_stride_height <- 448L - self$tile_sample_stride_width <- 448L - self$tile_sample_stride_num_frames <- 8L - }, - -#' Enable tiled encoding/decoding for memory efficiency - enable_tiling = function( - tile_sample_min_height = NULL, - tile_sample_min_width = NULL, - tile_sample_min_num_frames = NULL, - tile_sample_stride_height = NULL, - tile_sample_stride_width = NULL, - tile_sample_stride_num_frames = NULL - ) { - self$use_tiling <- TRUE - if (!is.null(tile_sample_min_height)) self$tile_sample_min_height <- tile_sample_min_height - if (!is.null(tile_sample_min_width)) self$tile_sample_min_width <- tile_sample_min_width - if (!is.null(tile_sample_min_num_frames)) self$tile_sample_min_num_frames <- tile_sample_min_num_frames - if (!is.null(tile_sample_stride_height)) self$tile_sample_stride_height <- tile_sample_stride_height - if (!is.null(tile_sample_stride_width)) self$tile_sample_stride_width <- tile_sample_stride_width - if (!is.null(tile_sample_stride_num_frames)) self$tile_sample_stride_num_frames <- tile_sample_stride_num_frames - invisible(self) - }, - -#' Disable tiling - disable_tiling = function() { - self$use_tiling <- FALSE - invisible(self) - }, - -#' Enable framewise decoding for long videos - enable_framewise_decoding = function() { - self$use_framewise_decoding <- TRUE - invisible(self) - }, - -#' Blend tiles vertically - blend_v = function( - a, - b, - blend_extent - ) { - blend_extent <- min(a$shape[4], b$shape[4], blend_extent) - for (y in seq_len(blend_extent)) { - weight <- (y - 1) / blend_extent - idx <- as.integer(a$shape[4] - blend_extent + y) - b[,,, y,] <- a[,,, idx,] * (1 - weight) + b[,,, y,] * weight - } - b - }, - -#' Blend tiles horizontally - blend_h = function( - a, - b, - blend_extent - ) { - blend_extent <- min(a$shape[5], b$shape[5], blend_extent) - for (x in seq_len(blend_extent)) { - weight <- (x - 1) / blend_extent - idx <- as.integer(a$shape[5] - blend_extent + x) - b[,,,, x] <- a[,,,, idx] * (1 - weight) + b[,,,, x] * weight - } - b - }, - -#' Blend tiles temporally - blend_t = function( - a, - b, - blend_extent - ) { - blend_extent <- min(a$shape[3], b$shape[3], blend_extent) - for (t in seq_len(blend_extent)) { - weight <- (t - 1) / blend_extent - idx <- as.integer(a$shape[3] - blend_extent + t) - b[,, t,,] <- a[,, idx,,] * (1 - weight) + b[,, t,,] * weight - } - b - }, - -#' Tiled encoding for large spatial dimensions - tiled_encode = function( - x, - causal = NULL - ) { - batch_size <- x$shape[1] - num_channels <- x$shape[2] - num_frames <- x$shape[3] - height <- x$shape[4] - width <- x$shape[5] - - latent_height <- height %/% self$spatial_compression_ratio - latent_width <- width %/% self$spatial_compression_ratio - - tile_latent_min_height <- self$tile_sample_min_height %/% self$spatial_compression_ratio - tile_latent_min_width <- self$tile_sample_min_width %/% self$spatial_compression_ratio - tile_latent_stride_height <- self$tile_sample_stride_height %/% self$spatial_compression_ratio - tile_latent_stride_width <- self$tile_sample_stride_width %/% self$spatial_compression_ratio - - blend_height <- tile_latent_min_height - tile_latent_stride_height - blend_width <- tile_latent_min_width - tile_latent_stride_width - - # Encode tiles - rows <- list() - for (i in seq(1, height, by = self$tile_sample_stride_height)) { - row <- list() - for (j in seq(1, width, by = self$tile_sample_stride_width)) { - i_end <- min(i + self$tile_sample_min_height - 1, height) - j_end <- min(j + self$tile_sample_min_width - 1, width) - tile <- self$encoder(x[,,, i:i_end, j:j_end], causal = causal) - row[[length(row) + 1]] <- tile - } - rows[[length(rows) + 1]] <- row - } - - # Blend tiles - result_rows <- list() - for (i in seq_along(rows)) { - result_row <- list() - for (j in seq_along(rows[[i]])) { - tile <- rows[[i]][[j]] - if (i > 1) { - tile <- self$blend_v(rows[[i - 1]][[j]], tile, blend_height) - } - if (j > 1) { - tile <- self$blend_h(rows[[i]][[j - 1]], tile, blend_width) - } - result_row[[length(result_row) + 1]] <- tile[,,, 1:tile_latent_stride_height, 1:tile_latent_stride_width] - } - result_rows[[length(result_rows) + 1]] <- torch::torch_cat(result_row, dim = 5) - } - - enc <- torch::torch_cat(result_rows, dim = 4)[,,, 1:latent_height, 1:latent_width] - enc - }, - -#' Tiled decoding for large spatial dimensions - tiled_decode = function( - z, - temb = NULL, - causal = NULL - ) { - batch_size <- z$shape[1] - num_channels <- z$shape[2] - num_frames <- z$shape[3] - height <- z$shape[4] - width <- z$shape[5] - - sample_height <- height * self$spatial_compression_ratio - sample_width <- width * self$spatial_compression_ratio - - tile_latent_min_height <- self$tile_sample_min_height %/% self$spatial_compression_ratio - tile_latent_min_width <- self$tile_sample_min_width %/% self$spatial_compression_ratio - tile_latent_stride_height <- self$tile_sample_stride_height %/% self$spatial_compression_ratio - tile_latent_stride_width <- self$tile_sample_stride_width %/% self$spatial_compression_ratio - - blend_height <- self$tile_sample_min_height - self$tile_sample_stride_height - blend_width <- self$tile_sample_min_width - self$tile_sample_stride_width - - # Decode tiles - rows <- list() - for (i in seq(1, height, by = tile_latent_stride_height)) { - row <- list() - for (j in seq(1, width, by = tile_latent_stride_width)) { - i_end <- min(i + tile_latent_min_height - 1, height) - j_end <- min(j + tile_latent_min_width - 1, width) - tile <- self$decoder(z[,,, i:i_end, j:j_end], temb, causal = causal) - row[[length(row) + 1]] <- tile - } - rows[[length(rows) + 1]] <- row - } - - # Blend tiles - result_rows <- list() - for (i in seq_along(rows)) { - result_row <- list() - for (j in seq_along(rows[[i]])) { - tile <- rows[[i]][[j]] - if (i > 1) { - tile <- self$blend_v(rows[[i - 1]][[j]], tile, blend_height) - } - if (j > 1) { - tile <- self$blend_h(rows[[i]][[j - 1]], tile, blend_width) - } - result_row[[length(result_row) + 1]] <- tile[,,, 1:self$tile_sample_stride_height, 1:self$tile_sample_stride_width] - } - result_rows[[length(result_rows) + 1]] <- torch::torch_cat(result_row, dim = 5) - } - - dec <- torch::torch_cat(result_rows, dim = 4)[,,, 1:sample_height, 1:sample_width] - dec - }, - -#' Internal encode with tiling support - .encode = function( - x, - causal = NULL - ) { - batch_size <- x$shape[1] - num_channels <- x$shape[2] - num_frames <- x$shape[3] - height <- x$shape[4] - width <- x$shape[5] - - if (self$use_tiling && (width > self$tile_sample_min_width || height > self$tile_sample_min_height)) { - return(self$tiled_encode(x, causal = causal)) - } - - self$encoder(x, causal = causal) - }, - -#' Encode video to latent space - encode = function( - x, - causal = NULL - ) { - if (self$use_slicing && x$shape[1] > 1) { - encoded_slices <- lapply(seq_len(x$shape[1]), function(i) { - self$.encode(x[i:i,,,,, drop = FALSE], causal = causal) - }) - h <- torch::torch_cat(encoded_slices, dim = 1) - } else { - h <- self$.encode(x, causal = causal) - } - diagonal_gaussian_distribution(h) - }, - -#' Internal decode with tiling support - .decode = function( - z, - temb = NULL, - causal = NULL - ) { - batch_size <- z$shape[1] - num_channels <- z$shape[2] - num_frames <- z$shape[3] - height <- z$shape[4] - width <- z$shape[5] - - tile_latent_min_height <- self$tile_sample_min_height %/% self$spatial_compression_ratio - tile_latent_min_width <- self$tile_sample_min_width %/% self$spatial_compression_ratio - - if (self$use_tiling && (width > tile_latent_min_width || height > tile_latent_min_height)) { - return(self$tiled_decode(z, temb, causal = causal)) - } - - self$decoder(z, temb, causal = causal) - }, - -#' Decode latent to video - decode = function( - z, - temb = NULL, - causal = NULL - ) { - if (self$use_slicing && z$shape[1] > 1) { - if (!is.null(temb)) { - decoded_slices <- lapply(seq_len(z$shape[1]), function(i) { - self$.decode(z[i:i,,,,, drop = FALSE], temb[i:i, drop = FALSE], causal = causal) - }) - } else { - decoded_slices <- lapply(seq_len(z$shape[1]), function(i) { - self$.decode(z[i:i,,,,, drop = FALSE], causal = causal) - }) - } - torch::torch_cat(decoded_slices, dim = 1) - } else { - self$.decode(z, temb, causal = causal) - } - }, - -#' Full forward pass (encode -> sample/mode -> decode) - forward = function( - sample, - temb = NULL, - sample_posterior = FALSE, - encoder_causal = NULL, - decoder_causal = NULL, - generator = NULL - ) { - posterior <- self$encode(sample, causal = encoder_causal) - if (sample_posterior) { - z <- posterior$sample(generator) - } else { - z <- posterior$mode() - } - self$decode(z, temb, causal = decoder_causal) - } -) - -#' Load LTX2 Video VAE from safetensors -#' -#' Load pre-trained LTX2 VAE weights from a HuggingFace safetensors file. -#' -#' @param weights_path Character. Path to safetensors file or directory containing weights. -#' @param config_path Character. Optional path to config.json. If NULL and weights_path -#' is a directory, looks for config.json in that directory. Otherwise uses default config. -#' @param device Character. Device to load weights to. Default: "cpu" -#' @param dtype Character or torch dtype. Data type. Default: "float32" -#' @param verbose Logical. Print loading progress. Default: TRUE -#' @return Initialized ltx2_video_vae module -#' @export -load_ltx2_vae <- function( - weights_path, - config_path = NULL, - device = "cpu", - dtype = "float32", - verbose = TRUE -) { - if (!file.exists(weights_path)) { - stop("Weights path not found: ", weights_path) - } - - # Handle directory vs file - if (dir.exists(weights_path)) { - vae_dir <- weights_path - # Look for config.json - if (is.null(config_path)) { - config_path <- file.path(vae_dir, "config.json") - if (!file.exists(config_path)) config_path <- NULL - } - # Look for weights file - weights_file <- file.path(vae_dir, "diffusion_pytorch_model.safetensors") - if (!file.exists(weights_file)) { - stop("Could not find diffusion_pytorch_model.safetensors in: ", vae_dir) - } - weights_path <- weights_file - } - - # Load config if provided - config <- NULL - if (!is.null(config_path) && file.exists(config_path)) { - config <- jsonlite::fromJSON(config_path) - if (verbose) message("Loaded config from: ", config_path) - } - - # Create VAE with config or defaults (matching HuggingFace LTX-2) - if (!is.null(config)) { - vae <- ltx2_video_vae( - in_channels = config$in_channels %||% 3L, - out_channels = config$out_channels %||% 3L, - latent_channels = config$latent_channels %||% 128L, - block_out_channels = as.integer(config$block_out_channels %||% c(256, 512, 1024, 2048)), - decoder_block_out_channels = as.integer(config$decoder_block_out_channels %||% c(256, 512, 1024)), - layers_per_block = as.integer(config$layers_per_block %||% c(4, 6, 6, 2, 2)), - decoder_layers_per_block = as.integer(config$decoder_layers_per_block %||% c(5, 5, 5, 5)), - spatio_temporal_scaling = as.logical(config$spatio_temporal_scaling %||% c(TRUE, TRUE, TRUE, TRUE)), - decoder_spatio_temporal_scaling = as.logical(config$decoder_spatio_temporal_scaling %||% c(TRUE, TRUE, TRUE)), - decoder_inject_noise = as.logical(config$decoder_inject_noise %||% c(FALSE, FALSE, FALSE, FALSE)), - downsample_type = config$downsample_type %||% c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - upsample_residual = as.logical(config$upsample_residual %||% c(TRUE, TRUE, TRUE)), - upsample_factor = as.integer(config$upsample_factor %||% c(2, 2, 2)), - timestep_conditioning = config$timestep_conditioning %||% FALSE, - patch_size = config$patch_size %||% 4L, - patch_size_t = config$patch_size_t %||% 1L, - resnet_norm_eps = config$resnet_norm_eps %||% 1e-6, - scaling_factor = config$scaling_factor %||% 1.0, - encoder_causal = config$encoder_causal %||% TRUE, - decoder_causal = config$decoder_causal %||% FALSE, - encoder_spatial_padding_mode = config$encoder_spatial_padding_mode %||% "zeros", - decoder_spatial_padding_mode = config$decoder_spatial_padding_mode %||% "reflect" - ) - } else { - vae <- ltx2_video_vae() - } - - # Load weights - if (verbose) message("Loading weights from: ", weights_path) - weights <- safetensors::safe_load_file(weights_path, framework = "torch") - - # Load weights into VAE - load_ltx2_vae_weights(vae, weights, verbose = verbose) - - # Move to device with dtype - if (dtype == "float16") { - torch_dtype <- torch::torch_float16() - } else { - torch_dtype <- torch::torch_float32() - } - vae$to(device = device, dtype = torch_dtype) - - if (verbose) message("VAE loaded successfully on device: ", device) - vae -} - -#' Load weights into LTX2 VAE module -#' -#' Maps HuggingFace safetensors parameter names to R module parameters. -#' -#' @param vae LTX2 VAE module -#' @param weights Named list of weight tensors from safetensors -#' @param verbose Logical. Print loading progress -#' @return The VAE with loaded weights (invisibly) -#' @keywords internal -load_ltx2_vae_weights <- function( - vae, - weights, - verbose = TRUE -) { - # Get native parameter names - native_params <- names(vae$parameters) - - # Build mapping from HF names to R names - remap_vae_key <- function(key) { - # HuggingFace VAE naming: - # encoder.conv_in.conv.weight -> encoder.conv_in.conv.weight - # encoder.down_blocks.0.resnets.0.norm.weight -> encoder.down_blocks.0.resnets.0.norm.weight - # The naming should be mostly 1:1 with our R module structure - - # No remapping needed for most keys - HF uses same structure - key - } - - loaded <- 0L - skipped <- 0L - unmapped <- character(0) - - torch::with_no_grad({ - for (hf_name in names(weights)) { - native_name <- remap_vae_key(hf_name) - - if (native_name %in% native_params) { - hf_tensor <- weights[[hf_name]] - native_tensor <- vae$parameters[[native_name]] - - if (all(as.integer(hf_tensor$shape) == as.integer(native_tensor$shape))) { - native_tensor$copy_(hf_tensor) - loaded <- loaded + 1L - } else { - if (verbose) { - message("Shape mismatch: ", native_name, - " (HF: ", paste(as.integer(hf_tensor$shape), collapse = "x"), - " vs R: ", paste(as.integer(native_tensor$shape), collapse = "x"), ")") - } - skipped <- skipped + 1L - } - } else { - skipped <- skipped + 1L - unmapped <- c(unmapped, paste0(hf_name, " -> ", native_name)) - } - } - }) - - if (verbose) { - message(sprintf("VAE weights: %d loaded, %d skipped", loaded, skipped)) - if (length(unmapped) > 0 && length(unmapped) <= 20) { - message("Unmapped parameters:") - for (u in unmapped[1:min(20, length(unmapped))]) { - message(" ", u) - } - } - if (length(unmapped) > 20) { - message(" ... and ", length(unmapped) - 20, " more") - } - } - - invisible(vae) -} - -#' Null-coalescing operator -#' @keywords internal -`%||%` <- function( - x, - y -) if (is.null(x)) y else x - diff --git a/R/vae_ltx23.R b/R/vae_ltx23.R new file mode 100644 index 0000000..b98a6dc --- /dev/null +++ b/R/vae_ltx23.R @@ -0,0 +1,603 @@ +#' LTX-2.3 Causal Video VAE +#' +#' Fresh R port of the LTX-2 video autoencoder from the diffusers +#' reference (Apache-2.0, autoencoder_kl_ltx2.py), with LTX 2.3 defaults: +#' encoder blocks (256, 512, 1024, 1024), a 4-up-block decoder with mixed +#' (spatiotemporal, spatiotemporal, temporal, spatial) upsampling, no +#' upsample residuals, and zeros spatial padding throughout. The encoder +#' is causal; the decoder is not. +#' +#' @name vae_ltx23 +NULL + +#' LTX-2.3 video encoder +#' +#' Pixel video [B, 3, F, H, W] -> latent statistics +#' [B, 2 * latent_channels, F/8, H/32, W/32] (mean and a uniform log-var +#' channel broadcast across the latent channels). +#' +#' @param in_channels,out_channels Integers. Pixel and latent channels. +#' @param block_out_channels Integer vector. Per-block output channels. +#' @param spatio_temporal_scaling Logical vector per block. +#' @param layers_per_block Integer vector (blocks then mid). +#' @param downsample_type Character vector per block. +#' @param patch_size,patch_size_t Integers. Pixel patchification. +#' @param resnet_norm_eps Numeric. +#' @param is_causal Logical. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_encoder3d <- torch::nn_module( + "ltx23_video_encoder3d", + initialize = function( + in_channels = 3L, + out_channels = 128L, + block_out_channels = c(256L, 512L, 1024L, 1024L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + downsample_type = c("spatial", "temporal", "spatiotemporal", + "spatiotemporal"), + patch_size = 4L, + patch_size_t = 1L, + resnet_norm_eps = 1e-6, + is_causal = TRUE, + spatial_padding_mode = "zeros" + ) { + self$patch_size <- as.integer(patch_size) + self$patch_size_t <- as.integer(patch_size_t) + self$in_channels <- in_channels * patch_size ^ 2 + self$is_causal <- is_causal + + output_channel <- out_channels + self$conv_in <- ltx23_causal_conv3d( + self$in_channels, output_channel, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) + + down_blocks <- list() + for (i in seq_along(block_out_channels)) { + input_channel <- output_channel + output_channel <- block_out_channels[i] + down_blocks[[i]] <- ltx23_video_down_block3d( + in_channels = input_channel, + out_channels = output_channel, + num_layers = layers_per_block[i], + resnet_eps = resnet_norm_eps, + spatio_temporal_scale = spatio_temporal_scaling[i], + downsample_type = downsample_type[i], + spatial_padding_mode = spatial_padding_mode + ) + } + self$down_blocks <- torch::nn_module_list(down_blocks) + + self$mid_block <- ltx23_video_mid_block3d( + in_channels = output_channel, + num_layers = layers_per_block[length(layers_per_block)], + resnet_eps = resnet_norm_eps, + spatial_padding_mode = spatial_padding_mode + ) + + self$norm_out <- ltx23_per_channel_rms_norm() + self$conv_out <- ltx23_causal_conv3d( + output_channel, out_channels + 1L, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) +}, + forward = function(hidden_states, causal = NULL) { + causal <- causal %||% self$is_causal + p <- self$patch_size + p_t <- self$patch_size_t + + dims <- hidden_states$shape + batch_size <- dims[1]; num_channels <- dims[2] + post_f <- dims[3] %/% p_t; post_h <- dims[4] %/% p; post_w <- dims[5] %/% p + + # Pixel patchification: space-to-channel with the LTX ordering + hidden_states <- hidden_states$reshape(c( + batch_size, num_channels, post_f, p_t, post_h, p, post_w, p + )) + # [B, C, F', pt, H', p, W', p] -> [B, C, pt, pw, ph, F', H', W'] + hidden_states <- hidden_states$permute(c(1L, 2L, 4L, 8L, 6L, 3L, 5L, 7L))$ + flatten(start_dim = 2L, end_dim = 5L) + hidden_states <- self$conv_in(hidden_states, causal = causal) + + for (i in seq_along(self$down_blocks)) { + hidden_states <- self$down_blocks[[i]](hidden_states, causal = causal) + } + hidden_states <- self$mid_block(hidden_states, causal = causal) + + hidden_states <- self$norm_out(hidden_states) + hidden_states <- torch::nnf_silu(hidden_states) + hidden_states <- self$conv_out(hidden_states, causal = causal) + + # Broadcast the single log-var channel across all latent channels + n_ch <- hidden_states$shape[2] + last_channel <- hidden_states$narrow(2L, n_ch, 1L)$`repeat`(c(1L, n_ch - 2L, 1L, 1L, 1L)) + torch::torch_cat(list(hidden_states, last_channel), dim = 2L) +} +) + +#' LTX-2.3 video decoder +#' +#' Latents [B, 128, F, H, W] -> pixel video [B, 3, 8F - 7, 32H, 32W]. +#' Block channel lists are given encoder-side (as in the config) and +#' reversed internally; \code{upsample_type} is indexed directly. +#' +#' @param in_channels,out_channels Integers. Latent and pixel channels. +#' @param block_out_channels Integer vector (config order). +#' @param spatio_temporal_scaling Logical vector per up block. +#' @param layers_per_block Integer vector (config order; first entry is +#' the mid block after reversal). +#' @param upsample_type Character vector per up block (not reversed). +#' @param patch_size,patch_size_t Integers. +#' @param resnet_norm_eps Numeric. +#' @param is_causal Logical. FALSE for LTX (symmetric temporal padding). +#' @param upsample_residual Logical vector per up block. +#' @param upsample_factor Integer vector per up block. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_decoder3d <- torch::nn_module( + "ltx23_video_decoder3d", + initialize = function( + in_channels = 128L, + out_channels = 3L, + block_out_channels = c(256L, 512L, 512L, 1024L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + upsample_type = c("spatiotemporal", "spatiotemporal", "temporal", + "spatial"), + patch_size = 4L, + patch_size_t = 1L, + resnet_norm_eps = 1e-6, + is_causal = FALSE, + upsample_residual = c(FALSE, FALSE, FALSE, FALSE), + upsample_factor = c(2L, 2L, 1L, 2L), + spatial_padding_mode = "zeros" + ) { + self$patch_size <- as.integer(patch_size) + self$patch_size_t <- as.integer(patch_size_t) + self$out_channels <- out_channels * patch_size ^ 2 + self$is_causal <- is_causal + + block_out_channels <- rev(block_out_channels) + spatio_temporal_scaling <- rev(spatio_temporal_scaling) + layers_per_block <- rev(layers_per_block) + upsample_residual <- rev(upsample_residual) + upsample_factor <- rev(upsample_factor) + # NOTE: upsample_type is deliberately NOT reversed (reference behavior) + + output_channel <- block_out_channels[1] + self$conv_in <- ltx23_causal_conv3d( + in_channels, output_channel, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) + + self$mid_block <- ltx23_video_mid_block3d( + in_channels = output_channel, + num_layers = layers_per_block[1], + resnet_eps = resnet_norm_eps, + spatial_padding_mode = spatial_padding_mode + ) + + up_blocks <- list() + for (i in seq_along(block_out_channels)) { + input_channel <- output_channel %/% upsample_factor[i] + output_channel <- block_out_channels[i] %/% upsample_factor[i] + up_blocks[[i]] <- ltx23_video_up_block3d( + in_channels = input_channel, + out_channels = output_channel, + num_layers = layers_per_block[i + 1L], + resnet_eps = resnet_norm_eps, + spatio_temporal_scale = spatio_temporal_scaling[i], + upsample_type = upsample_type[i], + upsample_residual = upsample_residual[i], + upscale_factor = upsample_factor[i], + spatial_padding_mode = spatial_padding_mode + ) + } + self$up_blocks <- torch::nn_module_list(up_blocks) + + self$norm_out <- ltx23_per_channel_rms_norm() + self$conv_out <- ltx23_causal_conv3d( + output_channel, self$out_channels, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) +}, + forward = function(hidden_states, causal = NULL) { + causal <- causal %||% self$is_causal + + hidden_states <- self$conv_in(hidden_states, causal = causal) + hidden_states <- self$mid_block(hidden_states, causal = causal) + for (i in seq_along(self$up_blocks)) { + hidden_states <- self$up_blocks[[i]](hidden_states, causal = causal) + } + + hidden_states <- self$norm_out(hidden_states) + hidden_states <- torch::nnf_silu(hidden_states) + hidden_states <- self$conv_out(hidden_states, causal = causal) + + # Un-patchify: channel-to-space with the LTX ordering + p <- self$patch_size + p_t <- self$patch_size_t + dims <- hidden_states$shape + hidden_states <- hidden_states$reshape(c( + dims[1], -1L, p_t, p, p, dims[3], dims[4], dims[5] + )) + # [B, C, pt, ph, pw, F, H, W] -> [B, C, F, pt, H, pw, W, ph] + hidden_states <- hidden_states$permute(c(1L, 2L, 6L, 3L, 7L, 5L, 8L, 4L)) + hidden_states$flatten(start_dim = 7L, end_dim = 8L)$ + flatten(start_dim = 5L, end_dim = 6L)$ + flatten(start_dim = 3L, end_dim = 4L) +} +) + +#' LTX-2.3 video VAE +#' +#' Encoder + decoder + per-channel latent statistics (loaded from the +#' checkpoint's \code{per_channel_statistics}). The checkpoint's +#' \code{scaling_factor} is 1.0, so latent (de)normalization is purely +#' the per-channel affine map. +#' +#' @param in_channels,out_channels Integers. Pixel channels. +#' @param latent_channels Integer. +#' @param block_out_channels,layers_per_block,spatio_temporal_scaling,downsample_type +#' Encoder configuration (see \code{\link{ltx23_video_encoder3d}}). +#' @param decoder_block_out_channels,decoder_layers_per_block,decoder_spatio_temporal_scaling,upsample_type,upsample_residual,upsample_factor +#' Decoder configuration (see \code{\link{ltx23_video_decoder3d}}). +#' @param patch_size,patch_size_t Integers. Pixel patchification. +#' @param resnet_norm_eps Numeric. +#' @param encoder_causal,decoder_causal Logicals. Temporal padding modes. +#' @param encoder_spatial_padding_mode,decoder_spatial_padding_mode Characters. +#' +#' @export +ltx23_video_vae <- torch::nn_module( + "ltx23_video_vae", + initialize = function( + in_channels = 3L, + out_channels = 3L, + latent_channels = 128L, + block_out_channels = c(256L, 512L, 1024L, 1024L), + decoder_block_out_channels = c(256L, 512L, 512L, 1024L), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + decoder_layers_per_block = c(4L, 6L, 4L, 2L, 2L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + decoder_spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + downsample_type = c("spatial", "temporal", "spatiotemporal", + "spatiotemporal"), + upsample_type = c("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + upsample_residual = c(FALSE, FALSE, FALSE, FALSE), + upsample_factor = c(2L, 2L, 1L, 2L), + patch_size = 4L, + patch_size_t = 1L, + resnet_norm_eps = 1e-6, + encoder_causal = TRUE, + decoder_causal = FALSE, + encoder_spatial_padding_mode = "zeros", + decoder_spatial_padding_mode = "zeros" + ) { + self$latent_channels <- as.integer(latent_channels) + + self$encoder <- ltx23_video_encoder3d( + in_channels = in_channels, + out_channels = latent_channels, + block_out_channels = block_out_channels, + spatio_temporal_scaling = spatio_temporal_scaling, + layers_per_block = layers_per_block, + downsample_type = downsample_type, + patch_size = patch_size, + patch_size_t = patch_size_t, + resnet_norm_eps = resnet_norm_eps, + is_causal = encoder_causal, + spatial_padding_mode = encoder_spatial_padding_mode + ) + self$decoder <- ltx23_video_decoder3d( + in_channels = latent_channels, + out_channels = out_channels, + block_out_channels = decoder_block_out_channels, + spatio_temporal_scaling = decoder_spatio_temporal_scaling, + layers_per_block = decoder_layers_per_block, + upsample_type = upsample_type, + patch_size = patch_size, + patch_size_t = patch_size_t, + resnet_norm_eps = resnet_norm_eps, + is_causal = decoder_causal, + upsample_residual = upsample_residual, + upsample_factor = upsample_factor, + spatial_padding_mode = decoder_spatial_padding_mode + ) + + self$latents_mean <- torch::nn_buffer(torch::torch_zeros(latent_channels)) + self$latents_std <- torch::nn_buffer(torch::torch_ones(latent_channels)) + + # Tiled decoding (reference defaults, in sample/pixel space) + self$use_tiling <- FALSE + self$use_framewise_decoding <- FALSE + self$tile_sample_min_height <- 512L + self$tile_sample_min_width <- 512L + self$tile_sample_min_num_frames <- 16L + self$tile_sample_stride_height <- 448L + self$tile_sample_stride_width <- 448L + self$tile_sample_stride_num_frames <- 8L +}, + encode = function(x, causal = NULL) { + moments <- self$encoder(x, causal = causal) + n <- self$latent_channels + list( + mean = moments$narrow(2L, 1L, n), + logvar = moments$narrow(2L, n + 1L, n) + ) +}, + enable_tiling = function(spatial = TRUE, temporal = TRUE) { + self$use_tiling <- spatial + self$use_framewise_decoding <- temporal + invisible(self) +}, + decode = function(z, causal = NULL) { + num_frames <- z$shape[3] + height <- z$shape[4] + width <- z$shape[5] + tile_lat_h <- self$tile_sample_min_height %/% 32L + 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) { + return(.ltx23_temporal_tiled_decode(self, z, causal = causal)) + } + if (self$use_tiling && (width > tile_lat_w || height > tile_lat_h)) { + return(.ltx23_tiled_decode(self, z, causal = causal)) + } + .ltx23_decode_tile(self, z, causal) +}, + forward = function(z) { + self$decode(z) +} +) + +# Linear crossfade of tile b's leading rows/columns/frames with tile a's +# trailing ones (in-place on b), per the diffusers reference +.ltx23_blend_v <- function(a, b, blend_extent) { + blend_extent <- min(a$shape[4], b$shape[4], blend_extent) + if (blend_extent <= 0L) { + return(b) + } + a_h <- a$shape[4] + for (y in seq_len(blend_extent)) { + w <- (y - 1) / blend_extent + b[,,, y,] <- a[,,, a_h - blend_extent + y,]$mul(1 - w) + + b[,,, y,]$mul(w) + } + b +} + +.ltx23_blend_h <- function(a, b, blend_extent) { + blend_extent <- min(a$shape[5], b$shape[5], blend_extent) + if (blend_extent <= 0L) { + return(b) + } + a_w <- a$shape[5] + for (x in seq_len(blend_extent)) { + w <- (x - 1) / blend_extent + b[,,,, x] <- a[,,,, a_w - blend_extent + x]$mul(1 - w) + + b[,,,, x]$mul(w) + } + b +} + +.ltx23_blend_t <- function(a, b, blend_extent) { + blend_extent <- min(a$shape[3], b$shape[3], blend_extent) + if (blend_extent <= 0L) { + return(b) + } + a_f <- a$shape[3] + for (x in seq_len(blend_extent)) { + w <- (x - 1) / blend_extent + b[,, x,,] <- a[,, a_f - blend_extent + x,,]$mul(1 - w) + + b[,, x,,]$mul(w) + } + b +} + +# Spatially tiled decode: overlapping latent tiles, decoded separately, +# crossfaded at the seams (diffusers AutoencoderKLLTX2Video.tiled_decode) +.ltx23_tiled_decode <- function(vae, z, causal = NULL) { + height <- z$shape[4] + width <- z$shape[5] + sample_height <- height * 32L + sample_width <- width * 32L + + tile_lat_h <- vae$tile_sample_min_height %/% 32L + tile_lat_w <- vae$tile_sample_min_width %/% 32L + stride_lat_h <- vae$tile_sample_stride_height %/% 32L + stride_lat_w <- vae$tile_sample_stride_width %/% 32L + + blend_height <- vae$tile_sample_min_height - vae$tile_sample_stride_height + blend_width <- vae$tile_sample_min_width - vae$tile_sample_stride_width + + rows <- list() + for (i in seq(0L, height - 1L, by = stride_lat_h)) { + row <- list() + for (j in seq(0L, width - 1L, by = stride_lat_w)) { + h_len <- min(tile_lat_h, height - i) + w_len <- min(tile_lat_w, width - j) + tile <- .ltx23_decode_tile( + vae, + z$narrow(4L, i + 1L, h_len)$narrow(5L, j + 1L, w_len), + causal + ) + 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 + } + + result_rows <- list() + for (i in seq_along(rows)) { + result_row <- list() + for (j in seq_along(rows[[i]])) { + tile <- rows[[i]][[j]] + if (i > 1L) { + tile <- .ltx23_blend_v(rows[[i - 1L]][[j]], tile, blend_height) + } + if (j > 1L) { + tile <- .ltx23_blend_h(rows[[i]][[j - 1L]], tile, blend_width) + } + keep_h <- min(vae$tile_sample_stride_height, tile$shape[4]) + keep_w <- min(vae$tile_sample_stride_width, tile$shape[5]) + result_row[[length(result_row) + 1L]] <- + tile$narrow(4L, 1L, keep_h)$narrow(5L, 1L, keep_w) + } + result_rows[[length(result_rows) + 1L]] <- torch::torch_cat(result_row, + dim = 5L) + } + + dec <- torch::torch_cat(result_rows, dim = 4L) + dec$narrow(4L, 1L, min(sample_height, dec$shape[4]))$ + narrow(5L, 1L, min(sample_width, dec$shape[5])) +} + +# Temporally tiled decode: overlapping latent frame windows (spatially +# tiled inside when the tile is large), crossfaded over time +# (diffusers AutoencoderKLLTX2Video._temporal_tiled_decode) +.ltx23_temporal_tiled_decode <- function(vae, z, causal = NULL) { + num_frames <- z$shape[3] + num_sample_frames <- (num_frames - 1L) * 8L + 1L + + tile_lat_h <- vae$tile_sample_min_height %/% 32L + tile_lat_w <- vae$tile_sample_min_width %/% 32L + tile_lat_f <- vae$tile_sample_min_num_frames %/% 8L + stride_lat_f <- vae$tile_sample_stride_num_frames %/% 8L + blend_num_frames <- vae$tile_sample_min_num_frames - vae$tile_sample_stride_num_frames + + row <- list() + for (i in seq(0L, num_frames - 1L, by = stride_lat_f)) { + f_len <- min(tile_lat_f + 1L, num_frames - i) + tile <- z$narrow(3L, i + 1L, f_len) + if (vae$use_tiling && + (tile$shape[5] > tile_lat_w || tile$shape[4] > tile_lat_h)) { + decoded <- .ltx23_tiled_decode(vae, tile, causal = causal) + } else { + decoded <- .ltx23_decode_tile(vae, tile, causal) + } + if (i > 0L) { + 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() + for (i in seq_along(row)) { + tile <- row[[i]] + if (i > 1L) { + tile <- .ltx23_blend_t(row[[i - 1L]], tile, blend_num_frames) + keep <- min(vae$tile_sample_stride_num_frames, tile$shape[3]) + result_row[[length(result_row) + 1L]] <- tile$narrow(3L, 1L, keep) + } else { + keep <- min(vae$tile_sample_stride_num_frames + 1L, tile$shape[3]) + result_row[[length(result_row) + 1L]] <- tile$narrow(3L, 1L, keep) + } + } + + dec <- torch::torch_cat(result_row, dim = 3L) + dec$narrow(3L, 1L, min(num_sample_frames, dec$shape[3])) +} + +#' Normalize latents with the VAE's per-channel statistics +#' +#' @param latents Tensor [B, C, F, H, W]. +#' @param latents_mean,latents_std Tensors [C]. +#' +#' @return Normalized latents. +#' +#' @export +ltx23_normalize_latents <- function(latents, latents_mean, latents_std) { + mean <- latents_mean$view(c(1L, -1L, 1L, 1L, 1L))$to(device = latents$device, + dtype = latents$dtype) + std <- latents_std$view(c(1L, -1L, 1L, 1L, 1L))$to(device = latents$device, + dtype = latents$dtype) + (latents - mean) / std +} + +#' Denormalize latents with the VAE's per-channel statistics +#' +#' @param latents Tensor [B, C, F, H, W]. +#' @param latents_mean,latents_std Tensors [C]. +#' +#' @return Denormalized latents ready for the decoder. +#' +#' @export +ltx23_denormalize_latents <- function(latents, latents_mean, latents_std) { + mean <- latents_mean$view(c(1L, -1L, 1L, 1L, 1L))$to(device = latents$device, + dtype = latents$dtype) + std <- latents_std$view(c(1L, -1L, 1L, 1L, 1L))$to(device = latents$device, + dtype = latents$dtype) + latents * std + mean +} + +#' Map an official VAE checkpoint key to the R module name +#' +#' The official checkpoint stores the encoder/decoder as flat block lists +#' (down_blocks.0-8 / up_blocks.0-8) where downsamplers/upsamplers and +#' the mid block are separate entries; diffusers (and this port) nest +#' them. Index mapping per diffusers convert_ltx2_to_diffusers.py. +#' +#' @param key Character. Checkpoint key (with or without "vae." prefix). +#' +#' @return Character. Module parameter/buffer name. +#' +#' @export +ltx23_map_vae_key <- function(key) { + key <- sub("^vae\\.", "", key) + + key <- sub("^per_channel_statistics\\.mean-of-means$", "latents_mean", key) + key <- sub("^per_channel_statistics\\.std-of-means$", "latents_std", key) + + down_map <- c("0" = "down_blocks.0", "1" = "down_blocks.0.downsamplers.0", + "2" = "down_blocks.1", "3" = "down_blocks.1.downsamplers.0", + "4" = "down_blocks.2", "5" = "down_blocks.2.downsamplers.0", + "6" = "down_blocks.3", "7" = "down_blocks.3.downsamplers.0", + "8" = "mid_block") + up_map <- c( + "0" = "mid_block", + "1" = "up_blocks.0.upsamplers.0", + "2" = "up_blocks.0", + "3" = "up_blocks.1.upsamplers.0", + "4" = "up_blocks.1", + "5" = "up_blocks.2.upsamplers.0", + "6" = "up_blocks.2", + "7" = "up_blocks.3.upsamplers.0", + "8" = "up_blocks.3" + ) + + m <- regmatches(key, regexec("^(encoder\\.down_blocks|decoder\\.up_blocks)\\.([0-9]+)\\.(.*)$", key))[[1]] + if (length(m) == 4L) { + if (startsWith(m[2], "encoder")) { + section <- "encoder" + } else { + section <- "decoder" + } + if (section == "encoder") { + map <- down_map + } else { + map <- up_map + } + idx <- m[3] + if (is.na(map[idx])) { + return(NA_character_) + } + rest <- m[4] + # Sampler entries map to the nested module directly (their inner + # structure is just .conv.conv / norms) + key <- paste0(section, ".", map[[idx]], ".", rest) + } + + key <- gsub("res_blocks", "resnets", key, fixed = TRUE) + key +} diff --git a/R/vae_ltx23_modules.R b/R/vae_ltx23_modules.R new file mode 100644 index 0000000..c23acce --- /dev/null +++ b/R/vae_ltx23_modules.R @@ -0,0 +1,429 @@ +#' LTX-2.3 Video VAE Building Blocks +#' +#' Fresh R port of the LTX-2 causal video autoencoder blocks from the +#' diffusers reference (Apache-2.0, +#' src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py). Training +#' and unused inference branches (noise injection, timestep conditioning, +#' plain-conv downsampling) are intentionally not ported; the 2.3 +#' checkpoints carry no such weights. +#' +#' @name vae_ltx23_modules +NULL + +#' Per-channel RMS normalization +#' +#' Normalizes by the root-mean-square across the channel dimension +#' (dim 2 of [B, C, F, H, W]); no learned parameters. +#' +#' @param eps Numeric. Stability epsilon. +#' +#' @export +ltx23_per_channel_rms_norm <- torch::nn_module( + "ltx23_per_channel_rms_norm", + initialize = function(eps = 1e-8) { + self$eps <- eps +}, + forward = function(x) { + mean_sq <- torch::torch_mean(x ^ 2, dim = 2L, keepdim = TRUE) + x$div(mean_sq$add(self$eps)$sqrt()) +} +) + +#' Causal 3D convolution +#' +#' Spatial padding is handled by the convolution; temporal padding +#' replicates the first frame (causal) or both edge frames (non-causal), +#' chosen at call time. +#' +#' @param in_channels,out_channels Integers. +#' @param kernel_size Integer or length-3 vector (t, h, w). +#' @param stride Integer or length-3 vector. +#' @param spatial_padding_mode Character. Conv padding mode. +#' +#' @export +ltx23_causal_conv3d <- torch::nn_module( + "ltx23_causal_conv3d", + initialize = function( + in_channels, + out_channels, + kernel_size = 3L, + stride = 1L, + spatial_padding_mode = "zeros" + ) { + if (length(kernel_size) == 1L) kernel_size <- rep(kernel_size, 3L) + if (length(stride) == 1L) stride <- rep(stride, 3L) + self$kernel_size <- as.integer(kernel_size) + + height_pad <- kernel_size[2] %/% 2L + width_pad <- kernel_size[3] %/% 2L + + self$conv <- torch::nn_conv3d(in_channels, out_channels, + self$kernel_size, + stride = as.integer(stride), + padding = c(0L, height_pad, width_pad), + padding_mode = spatial_padding_mode) +}, + forward = function(hidden_states, causal = TRUE) { + time_k <- self$kernel_size[1] + if (time_k > 1L) { + if (causal) { + pad_left <- hidden_states$narrow(3L, 1L, 1L)$`repeat`(c(1L, 1L, time_k - 1L, 1L, 1L)) + hidden_states <- torch::torch_cat(list(pad_left, hidden_states), dim = 3L) + } else { + half <- (time_k - 1L) %/% 2L + pad_left <- hidden_states$narrow(3L, 1L, 1L)$`repeat`(c(1L, 1L, half, 1L, 1L)) + pad_right <- hidden_states$narrow( + 3L, hidden_states$shape[3], 1L + )$`repeat`(c(1L, 1L, half, 1L, 1L)) + hidden_states <- torch::torch_cat(list(pad_left, hidden_states, pad_right), dim = 3L) + } + } + self$conv(hidden_states) +} +) + +#' LTX 3D ResNet block +#' +#' PerChannelRMSNorm -> SiLU -> causal conv, twice, with a LayerNorm + +#' 1x1 Conv3d shortcut when the channel count changes. +#' +#' @param in_channels,out_channels Integers. +#' @param eps Numeric. Shortcut LayerNorm epsilon. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_resnet_block3d <- torch::nn_module( + "ltx23_video_resnet_block3d", + initialize = function( + in_channels, + out_channels = NULL, + eps = 1e-6, + spatial_padding_mode = "zeros" + ) { + out_channels <- out_channels %||% in_channels + + self$norm1 <- ltx23_per_channel_rms_norm() + self$conv1 <- ltx23_causal_conv3d( + in_channels, out_channels, kernel_size = 3L, + spatial_padding_mode = spatial_padding_mode + ) + self$norm2 <- ltx23_per_channel_rms_norm() + self$conv2 <- ltx23_causal_conv3d( + out_channels, out_channels, kernel_size = 3L, + spatial_padding_mode = spatial_padding_mode + ) + + if (in_channels != out_channels) { + self$norm3 <- torch::nn_layer_norm(in_channels, eps = eps, + elementwise_affine = TRUE) + # A plain (non-causal) 1x1 Conv3d, per the reference + self$conv_shortcut <- torch::nn_conv3d(in_channels, out_channels, + kernel_size = 1L, stride = 1L) + } +}, + forward = function(inputs, causal = TRUE) { + hidden_states <- self$norm1(inputs) + hidden_states <- torch::nnf_silu(hidden_states) + hidden_states <- self$conv1(hidden_states, causal = causal) + + hidden_states <- self$norm2(hidden_states) + hidden_states <- torch::nnf_silu(hidden_states) + hidden_states <- self$conv2(hidden_states, causal = causal) + + if (!is.null(self$norm3)) { + # LayerNorm over channels: move C last, norm, move back + inputs <- self$norm3(inputs$permute(c(1L, 3L, 4L, 5L, 2L)))$permute(c(1L, 5L, 2L, 3L, 4L)) + } + if (!is.null(self$conv_shortcut)) { + inputs <- self$conv_shortcut(inputs) + } + hidden_states + inputs +} +) + +#' Pixel-unshuffle 3D downsampler +#' +#' Conv followed by space/time-to-channel rearrangement, plus a grouped +#' channel-mean residual of the same rearrangement. +#' +#' @param in_channels,out_channels Integers. +#' @param stride Length-3 integer vector (t, h, w). +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_downsampler3d <- torch::nn_module( + "ltx23_video_downsampler3d", + initialize = function( + in_channels, + out_channels, + stride = c(1L, 1L, 1L), + spatial_padding_mode = "zeros" + ) { + if (length(stride) == 1L) stride <- rep(stride, 3L) + self$stride <- as.integer(stride) + self$group_size <- (in_channels * prod(stride)) %/% out_channels + + conv_out <- out_channels %/% prod(stride) + self$conv <- ltx23_causal_conv3d( + in_channels, conv_out, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) +}, + forward = function(hidden_states, causal = TRUE) { + s <- self$stride + if (s[1] > 1L) { + pad <- hidden_states$narrow(3L, 1L, 1L)$`repeat`(c(1L, 1L, s[1] - 1L, 1L, 1L)) + hidden_states <- torch::torch_cat(list(pad, hidden_states), dim = 3L) + } + + # Space/time-to-channel rearrangement shared by both paths + to_channels <- function(x) { + x <- x$unflatten(5L, c(-1L, s[3]))$unflatten(4L, c(-1L, s[2]))$unflatten(3L, + c(-1L, s[1])) + # [B, C, F', st, H', sh, W', sw] -> [B, C, st, sh, sw, F', H', W'] + x$permute(c(1L, 2L, 4L, 6L, 8L, 3L, 5L, 7L))$flatten(start_dim = 2L, + end_dim = 5L) + } + + residual <- to_channels(hidden_states) + residual <- residual$unflatten(2L, c(-1L, self$group_size))$mean(dim = 3L) + + hidden_states <- self$conv(hidden_states, causal = causal) + hidden_states <- to_channels(hidden_states) + hidden_states + residual +} +) + +#' Pixel-shuffle 3D upsampler +#' +#' Conv followed by channel-to-space/time rearrangement, with an optional +#' channel-repeat residual and an upscale factor that divides the conv +#' output channels. +#' +#' @param in_channels Integer. +#' @param stride Length-3 integer vector (t, h, w). +#' @param residual Logical. Add the rearranged input as a residual. +#' @param upscale_factor Integer. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_upsampler3d <- torch::nn_module( + "ltx23_video_upsampler3d", + initialize = function( + in_channels, + stride = c(1L, 1L, 1L), + residual = FALSE, + upscale_factor = 1L, + spatial_padding_mode = "zeros" + ) { + if (length(stride) == 1L) stride <- rep(stride, 3L) + self$stride <- as.integer(stride) + self$residual <- residual + self$upscale_factor <- as.integer(upscale_factor) + + out_channels <- (in_channels * prod(stride)) %/% upscale_factor + self$conv <- ltx23_causal_conv3d( + in_channels, out_channels, kernel_size = 3L, stride = 1L, + spatial_padding_mode = spatial_padding_mode + ) +}, + forward = function(hidden_states, causal = TRUE) { + s <- self$stride + dims <- hidden_states$shape + num_frames <- dims[3]; height <- dims[4]; width <- dims[5] + + # Channel-to-space/time rearrangement shared by both paths + to_space <- function(x) { + x <- x$reshape(c(dims[1], -1L, s[1], s[2], s[3], num_frames, + height, width)) + # [B, C', st, sh, sw, F, H, W] -> [B, C', F, st, H, sh, W, sw] + x <- x$permute(c(1L, 2L, 6L, 3L, 7L, 4L, 8L, 5L)) + x <- x$flatten(start_dim = 7L, end_dim = 8L)$ + flatten(start_dim = 5L, end_dim = 6L)$ + flatten(start_dim = 3L, end_dim = 4L) + # Drop the causally duplicated leading frames + if (s[1] > 1L) x <- x$narrow(3L, s[1], x$shape[3] - s[1] + 1L) else x + } + + residual <- NULL + if (self$residual) { + residual <- to_space(hidden_states) + repeats <- prod(s) %/% self$upscale_factor + residual <- residual$`repeat`(c(1L, repeats, 1L, 1L, 1L)) + } + + hidden_states <- self$conv(hidden_states, causal = causal) + hidden_states <- to_space(hidden_states) + + if (!is.null(residual)) hidden_states <- hidden_states + residual + hidden_states +} +) + +#' LTX video down block +#' +#' ResNet stack (at the input channel count) followed by a +#' pixel-unshuffle downsampler that also changes the channel count. +#' +#' @param in_channels,out_channels Integers. +#' @param num_layers Integer. ResNet count. +#' @param resnet_eps Numeric. +#' @param spatio_temporal_scale Logical. Whether to downsample at all. +#' @param downsample_type "spatial", "temporal", or "spatiotemporal". +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_down_block3d <- torch::nn_module( + "ltx23_video_down_block3d", + initialize = function( + in_channels, + out_channels = NULL, + num_layers = 1L, + resnet_eps = 1e-6, + spatio_temporal_scale = TRUE, + downsample_type = "spatiotemporal", + spatial_padding_mode = "zeros" + ) { + out_channels <- out_channels %||% in_channels + + self$resnets <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { + ltx23_video_resnet_block3d(in_channels, in_channels, + eps = resnet_eps, + spatial_padding_mode = spatial_padding_mode) + })) + + if (spatio_temporal_scale) { + stride <- switch(downsample_type, + spatial = c(1L, 2L, 2L), + temporal = c(2L, 1L, 1L), + spatiotemporal = c(2L, 2L, 2L), + stop("Unknown downsample_type: ", downsample_type) + ) + self$downsamplers <- torch::nn_module_list(list( + ltx23_video_downsampler3d( + in_channels, out_channels, stride = stride, + spatial_padding_mode = spatial_padding_mode + ) + )) + } +}, + forward = function(hidden_states, causal = TRUE) { + for (i in seq_along(self$resnets)) { + hidden_states <- self$resnets[[i]](hidden_states, causal = causal) + } + if (!is.null(self$downsamplers)) { + for (i in seq_along(self$downsamplers)) { + hidden_states <- self$downsamplers[[i]](hidden_states, causal = causal) + } + } + hidden_states +} +) + +#' LTX video mid block +#' +#' A plain ResNet stack at a fixed channel count. +#' +#' @param in_channels Integer. +#' @param num_layers Integer. +#' @param resnet_eps Numeric. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_mid_block3d <- torch::nn_module( + "ltx23_video_mid_block3d", + initialize = function( + in_channels, + num_layers = 1L, + resnet_eps = 1e-6, + spatial_padding_mode = "zeros" + ) { + self$resnets <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { + ltx23_video_resnet_block3d(in_channels, in_channels, + eps = resnet_eps, + spatial_padding_mode = spatial_padding_mode) + })) +}, + forward = function(hidden_states, causal = TRUE) { + for (i in seq_along(self$resnets)) { + hidden_states <- self$resnets[[i]](hidden_states, causal = causal) + } + hidden_states +} +) + +#' LTX video up block +#' +#' Optional channel-changing conv-in ResNet, pixel-shuffle upsampler, +#' then a ResNet stack at the output channel count. +#' +#' @param in_channels,out_channels Integers. +#' @param num_layers Integer. +#' @param resnet_eps Numeric. +#' @param spatio_temporal_scale Logical. +#' @param upsample_type "spatial", "temporal", or "spatiotemporal". +#' @param upsample_residual Logical. +#' @param upscale_factor Integer. +#' @param spatial_padding_mode Character. +#' +#' @export +ltx23_video_up_block3d <- torch::nn_module( + "ltx23_video_up_block3d", + initialize = function( + in_channels, + out_channels = NULL, + num_layers = 1L, + resnet_eps = 1e-6, + spatio_temporal_scale = TRUE, + upsample_type = "spatiotemporal", + upsample_residual = FALSE, + upscale_factor = 1L, + spatial_padding_mode = "zeros" + ) { + out_channels <- out_channels %||% in_channels + + if (in_channels != out_channels) { + self$conv_in <- ltx23_video_resnet_block3d(in_channels, out_channels, + eps = resnet_eps, spatial_padding_mode = spatial_padding_mode) + } + + if (spatio_temporal_scale) { + stride <- switch(upsample_type, + spatial = c(1L, 2L, 2L), + temporal = c(2L, 1L, 1L), + spatiotemporal = c(2L, 2L, 2L), + stop("Unknown upsample_type: ", upsample_type) + ) + self$upsamplers <- torch::nn_module_list(list( + ltx23_video_upsampler3d( + in_channels = out_channels * upscale_factor, + stride = stride, + residual = upsample_residual, + upscale_factor = upscale_factor, + spatial_padding_mode = spatial_padding_mode + ) + )) + } + + self$resnets <- torch::nn_module_list(lapply(seq_len(num_layers), function(i) { + ltx23_video_resnet_block3d( + out_channels, out_channels, eps = resnet_eps, + spatial_padding_mode = spatial_padding_mode + ) + })) +}, + forward = function(hidden_states, causal = TRUE) { + if (!is.null(self$conv_in)) { + hidden_states <- self$conv_in(hidden_states, causal = causal) + } + if (!is.null(self$upsamplers)) { + for (i in seq_along(self$upsamplers)) { + hidden_states <- self$upsamplers[[i]](hidden_states, causal = causal) + } + } + for (i in seq_along(self$resnets)) { + hidden_states <- self$resnets[[i]](hidden_states, causal = causal) + } + hidden_states +} +) diff --git a/R/vae_ltx2_modules.R b/R/vae_ltx2_modules.R deleted file mode 100644 index 6706595..0000000 --- a/R/vae_ltx2_modules.R +++ /dev/null @@ -1,746 +0,0 @@ -#' LTX2 Video VAE Module Building Blocks -#' -#' Low-level nn_module components for the LTX2 3D causal VAE. -#' Used by vae_ltx2.R encoder/decoder. -#' -#' @name vae_ltx2_modules -NULL - -#' Per-channel RMS normalization -#' -#' Normalizes tensor by root-mean-square along the channel dimension: -#' y = x / sqrt(mean(x^2, dim=channel_dim, keepdim=TRUE) + eps) -#' -#' @param channel_dim Integer. Dimension for RMS computation (1-indexed). Default: 2 (channels) -#' @param eps Numeric. Small constant for numerical stability. Default: 1e-8 -#' @export -per_channel_rms_norm <- torch::nn_module( - - "PerChannelRMSNorm", - - initialize = function( - channel_dim = 2L, - eps = 1e-8 - ) { - self$channel_dim <- channel_dim - self$eps <- eps - }, - - forward = function( - x, - channel_dim = NULL - ) { - if (is.null(channel_dim)) { - channel_dim <- self$channel_dim - } - mean_sq <- torch::torch_mean(x ^ 2, dim = channel_dim, keepdim = TRUE) - rms <- torch::torch_sqrt(mean_sq + self$eps) - x / rms - } -) - -#' LTX2 Video Causal 3D Convolution -#' -#' 3D convolution with runtime-selectable causal or non-causal padding. -#' Causal mode pads temporally by repeating first frame. -#' Non-causal mode pads temporally by repeating first and last frames. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer. Output channels. -#' @param kernel_size Integer or vector of 3. Convolution kernel size. -#' @param stride Integer or vector of 3. Stride. -#' @param dilation Integer or vector of 3. Dilation. -#' @param groups Integer. Convolution groups. Default: 1 -#' @param spatial_padding_mode Character. Padding mode for spatial dims. Default: "zeros" -#' @export -ltx2_video_causal_conv3d <- torch::nn_module( - "LTX2VideoCausalConv3d", - - initialize = function( - in_channels, - out_channels, - kernel_size = 3L, - stride = 1L, - dilation = 1L, - groups = 1L, - spatial_padding_mode = "zeros" - ) { - self$in_channels <- in_channels - self$out_channels <- out_channels - - # Normalize kernel_size to tuple of 3 - if (length(kernel_size) == 1) { - self$kernel_size <- rep(as.integer(kernel_size), 3) - } else { - self$kernel_size <- as.integer(kernel_size) - } - - # Normalize dilation (default temporal dilation only) - if (length(dilation) == 1) { - dilation <- c(as.integer(dilation), 1L, 1L) - } - - # Normalize stride - if (length(stride) == 1) { - stride <- rep(as.integer(stride), 3) - } - - # Spatial padding (no temporal padding in the conv itself) - height_pad <- self$kernel_size[2] %/% 2L - width_pad <- self$kernel_size[3] %/% 2L - padding <- c(0L, height_pad, width_pad) - - self$conv <- torch::nn_conv3d( - in_channels = in_channels, - out_channels = out_channels, - kernel_size = self$kernel_size, - stride = stride, - dilation = dilation, - groups = groups, - padding = padding, - padding_mode = spatial_padding_mode - ) - }, - - forward = function( - hidden_states, - causal = TRUE - ) { - time_kernel_size <- self$kernel_size[1] - - if (causal) { - # Causal: pad by repeating first frame on left - pad_left <- hidden_states[,, 1:1,,]$`repeat`(c(1L, 1L, time_kernel_size - 1L, 1L, 1L)) - hidden_states <- torch::torch_cat(list(pad_left, hidden_states), dim = 3) - } else { - # Non-causal: pad by repeating first/last frames symmetrically - pad_amount <- (time_kernel_size - 1L) %/% 2L - if (pad_amount > 0) { - pad_left <- hidden_states[,, 1:1,,]$`repeat`(c(1L, 1L, pad_amount, 1L, 1L)) - n_frames <- hidden_states$shape[3] - pad_right <- hidden_states[,, n_frames:n_frames,,]$`repeat`(c(1L, 1L, pad_amount, 1L, 1L)) - hidden_states <- torch::torch_cat(list(pad_left, hidden_states, pad_right), dim = 3) - } - } - - self$conv(hidden_states) - } -) - -#' LTX2 Video ResNet Block 3D -#' -#' 3D ResNet block with per-channel RMS normalization and optional -#' noise injection and timestep conditioning. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer or NULL. Output channels (defaults to in_channels). -#' @param dropout Numeric. Dropout rate. Default: 0.0 -#' @param eps Numeric. Epsilon for normalization. Default: 1e-6 -#' @param non_linearity Character. Activation function. Default: "silu" -#' @param inject_noise Logical. Whether to inject noise. Default: FALSE -#' @param timestep_conditioning Logical. Whether to use timestep conditioning. Default: FALSE -#' @param spatial_padding_mode Character. Padding mode. Default: "zeros" -#' @export -ltx2_video_resnet_block3d <- torch::nn_module( - "LTX2VideoResnetBlock3d", - - initialize = function( - in_channels, - out_channels = NULL, - dropout = 0.0, - eps = 1e-6, - non_linearity = "silu", - inject_noise = FALSE, - timestep_conditioning = FALSE, - spatial_padding_mode = "zeros" - ) { - if (is.null(out_channels)) out_channels <- in_channels - - # Activation - self$nonlinearity <- if (non_linearity == "silu" || non_linearity == "swish") { - torch::nn_silu() - } else if (non_linearity == "gelu") { - torch::nn_gelu() - } else { - torch::nn_relu() - } - - self$norm1 <- per_channel_rms_norm() - self$conv1 <- ltx2_video_causal_conv3d( - in_channels = in_channels, - out_channels = out_channels, - kernel_size = 3L, - spatial_padding_mode = spatial_padding_mode - ) - - self$norm2 <- per_channel_rms_norm() - self$dropout <- torch::nn_dropout(p = dropout) - self$conv2 <- ltx2_video_causal_conv3d( - in_channels = out_channels, - out_channels = out_channels, - kernel_size = 3L, - spatial_padding_mode = spatial_padding_mode - ) - - # Shortcut connection if channels differ - self$norm3 <- NULL - self$conv_shortcut <- NULL - if (in_channels != out_channels) { - self$norm3 <- torch::nn_layer_norm(in_channels, eps = eps, elementwise_affine = TRUE) - # Regular Conv3d for shortcut (not causal) - self$conv_shortcut <- torch::nn_conv3d( - in_channels = in_channels, - out_channels = out_channels, - kernel_size = 1L, - stride = 1L - ) - } - - # Noise injection (optional) - self$per_channel_scale1 <- NULL - self$per_channel_scale2 <- NULL - if (inject_noise) { - self$per_channel_scale1 <- torch::nn_parameter(torch::torch_zeros(c(in_channels, 1, 1))) - self$per_channel_scale2 <- torch::nn_parameter(torch::torch_zeros(c(in_channels, 1, 1))) - } - - # Timestep conditioning (optional) - self$scale_shift_table <- NULL - if (timestep_conditioning) { - init_table <- torch::torch_randn(c(4, in_channels)) / sqrt(in_channels) - self$scale_shift_table <- torch::nn_parameter(init_table) - } - - self$in_channels <- in_channels - self$out_channels <- out_channels - }, - - forward = function( - inputs, - temb = NULL, - generator = NULL, - causal = TRUE - ) { - hidden_states <- inputs - - hidden_states <- self$norm1(hidden_states) - - # Timestep conditioning (shift/scale) - shift_1 <- NULL - scale_1 <- NULL - shift_2 <- NULL - scale_2 <- NULL - if (!is.null(self$scale_shift_table) && !is.null(temb)) { - # temb shape: [B, C*4, 1, 1, 1] -> unflatten to [B, 4, C, 1, 1, 1] - temb <- temb$unflatten(2, c(4, - 1)) + self$scale_shift_table[NULL, .., NULL, NULL, NULL] - splits <- temb$unbind(dim = 2) - shift_1 <- splits[[1]] - scale_1 <- splits[[2]] - shift_2 <- splits[[3]] - scale_2 <- splits[[4]] - hidden_states <- hidden_states * (1 + scale_1) + shift_1 - } - - hidden_states <- self$nonlinearity(hidden_states) - hidden_states <- self$conv1(hidden_states, causal = causal) - - # Noise injection after conv1 - if (!is.null(self$per_channel_scale1)) { - h <- hidden_states$shape[4] - w <- hidden_states$shape[5] - spatial_noise <- torch::torch_randn(c(h, w), device = hidden_states$device, - dtype = hidden_states$dtype)$unsqueeze(1) - hidden_states <- hidden_states + (spatial_noise * self$per_channel_scale1)[NULL,, NULL,,] - } - - hidden_states <- self$norm2(hidden_states) - - # Second timestep conditioning - if (!is.null(self$scale_shift_table) && !is.null(temb)) { - hidden_states <- hidden_states * (1 + scale_2) + shift_2 - } - - hidden_states <- self$nonlinearity(hidden_states) - hidden_states <- self$dropout(hidden_states) - hidden_states <- self$conv2(hidden_states, causal = causal) - - # Noise injection after conv2 - if (!is.null(self$per_channel_scale2)) { - h <- hidden_states$shape[4] - w <- hidden_states$shape[5] - spatial_noise <- torch::torch_randn(c(h, w), device = hidden_states$device, - dtype = hidden_states$dtype)$unsqueeze(1) - hidden_states <- hidden_states + (spatial_noise * self$per_channel_scale2)[NULL,, NULL,,] - } - - # Shortcut connection - if (!is.null(self$norm3)) { - # LayerNorm expects last dim to be features, so move channels to last - inputs <- inputs$movedim(2, - 1) - inputs <- self$norm3(inputs) - inputs <- inputs$movedim(- 1, 2) - } - - if (!is.null(self$conv_shortcut)) { - inputs <- self$conv_shortcut(inputs) - } - - hidden_states <- hidden_states + inputs - hidden_states - } -) - -#' LTX Video Downsampler 3D -#' -#' Spatiotemporal downsampling with strided pixel unshuffle + convolution. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer. Output channels. -#' @param stride Integer or vector of 3. Downsampling stride. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx_video_downsampler3d <- torch::nn_module( - "LTXVideoDownsampler3d", - - initialize = function( - in_channels, - out_channels, - stride = 1L, - spatial_padding_mode = "zeros" - ) { - if (length(stride) == 1) { - self$stride <- rep(as.integer(stride), 3) - } else { - self$stride <- as.integer(stride) - } - - # Calculate group size for averaging residual - stride_prod <- self$stride[1] * self$stride[2] * self$stride[3] - self$group_size <- (in_channels * stride_prod) %/% out_channels - - # Output channels after pixel unshuffle - conv_out_channels <- out_channels %/% stride_prod - - self$conv <- ltx2_video_causal_conv3d( - in_channels = in_channels, - out_channels = conv_out_channels, - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - }, - - forward = function( - hidden_states, - causal = TRUE - ) { - # Pad temporal dimension to handle stride - # cat with first (stride[1]-1) frames repeated - if (self$stride[1] > 1) { - pad_frames <- hidden_states[,, 1:(self$stride[1] - 1),,] - hidden_states <- torch::torch_cat(list(pad_frames, hidden_states), dim = 3) - } - - # Compute residual via pixel unshuffle pattern - # unflatten spatial dims, permute, flatten to channel-like - residual <- hidden_states$unflatten(5, c(- 1, self$stride[3])) # width - residual <- residual$unflatten(4, c(- 1, self$stride[2])) # height - residual <- residual$unflatten(3, c(- 1, self$stride[1])) # frames - - # Permute: [B, C, F//s, s, H//s, s, W//s, s] -> [B, C, s, s, s, F//s, H//s, W//s] - residual <- residual$permute(c(1, 2, 4, 6, 8, 3, 5, 7)) - residual <- residual$flatten(start_dim = 2, end_dim = 5) # [B, C*s^3, F', H', W'] - residual <- residual$unflatten(2, c(- 1, self$group_size)) # [B, out_c, group_size, F', H', W'] - residual <- residual$mean(dim = 3) # Average over groups - - # Convolution path - hidden_states <- self$conv(hidden_states, causal = causal) - - # Same pixel unshuffle for conv output - hidden_states <- hidden_states$unflatten(5, c(- 1, self$stride[3])) - hidden_states <- hidden_states$unflatten(4, c(- 1, self$stride[2])) - hidden_states <- hidden_states$unflatten(3, c(- 1, self$stride[1])) - hidden_states <- hidden_states$permute(c(1, 2, 4, 6, 8, 3, 5, 7)) - hidden_states <- hidden_states$flatten(start_dim = 2, end_dim = 5) - - hidden_states <- hidden_states + residual - hidden_states - } -) - -#' LTX Video Upsampler 3D -#' -#' Spatiotemporal upsampling with pixel shuffle + optional residual. -#' -#' @param in_channels Integer. Input channels. -#' @param stride Integer or vector of 3. Upsampling stride. -#' @param residual Logical. Whether to use residual connection. -#' @param upscale_factor Integer. Channel upscale factor. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx_video_upsampler3d <- torch::nn_module( - "LTXVideoUpsampler3d", - - initialize = function( - in_channels, - stride = 1L, - residual = FALSE, - upscale_factor = 1L, - spatial_padding_mode = "zeros" - ) { - if (length(stride) == 1) { - self$stride <- rep(as.integer(stride), 3) - } else { - self$stride <- as.integer(stride) - } - self$residual <- residual - self$upscale_factor <- upscale_factor - - stride_prod <- self$stride[1] * self$stride[2] * self$stride[3] - out_channels <- (in_channels * stride_prod) %/% upscale_factor - - self$conv <- ltx2_video_causal_conv3d( - in_channels = in_channels, - out_channels = out_channels, - kernel_size = 3L, - stride = 1L, - spatial_padding_mode = spatial_padding_mode - ) - }, - - forward = function( - hidden_states, - causal = TRUE - ) { - batch_size <- hidden_states$shape[1] - num_channels <- hidden_states$shape[2] - num_frames <- hidden_states$shape[3] - height <- hidden_states$shape[4] - width <- hidden_states$shape[5] - - stride_prod <- self$stride[1] * self$stride[2] * self$stride[3] - - # Residual path - residual_out <- NULL - if (self$residual) { - # Reshape for pixel shuffle: [B, C, F, H, W] -> [B, C//s^3, s, s, s, F, H, W] - residual_out <- hidden_states$reshape(c( - batch_size, - - 1, - self$stride[1], self$stride[2], self$stride[3], - num_frames, height, width - )) - # Permute: [B, C', s_t, s_h, s_w, F, H, W] -> [B, C', F, s_t, H, s_h, W, s_w] - residual_out <- residual_out$permute(c(1, 2, 6, 3, 7, 4, 8, 5)) - # Flatten spatial dims - residual_out <- residual_out$flatten(start_dim = 7, end_dim = 8) # W * s_w - residual_out <- residual_out$flatten(start_dim = 5, end_dim = 6) # H * s_h - residual_out <- residual_out$flatten(start_dim = 3, end_dim = 4) # F * s_t - - # Repeat channels for upscale factor - repeats <- stride_prod %/% self$upscale_factor - residual_out <- residual_out$`repeat`(c(1L, repeats, 1L, 1L, 1L)) - # Remove first (stride[1]-1) frames - residual_out <- residual_out[,, self$stride[1]:N,,] - } - - # Convolution path - hidden_states <- self$conv(hidden_states, causal = causal) - - # Pixel shuffle - hidden_states <- hidden_states$reshape(c( - batch_size, - - 1, - self$stride[1], self$stride[2], self$stride[3], - num_frames, height, width - )) - hidden_states <- hidden_states$permute(c(1, 2, 6, 3, 7, 4, 8, 5)) - hidden_states <- hidden_states$flatten(start_dim = 7, end_dim = 8) - hidden_states <- hidden_states$flatten(start_dim = 5, end_dim = 6) - hidden_states <- hidden_states$flatten(start_dim = 3, end_dim = 4) - # Remove first (stride[1]-1) frames - hidden_states <- hidden_states[,, self$stride[1]:N,,] - - if (self$residual && !is.null(residual_out)) { - hidden_states <- hidden_states + residual_out - } - - hidden_states - } -) - -#' LTX2 Video Down Block 3D -#' -#' Encoder down block with multiple ResNet layers and optional downsampling. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer or NULL. Output channels. -#' @param num_layers Integer. Number of ResNet layers. -#' @param dropout Numeric. Dropout rate. -#' @param resnet_eps Numeric. Epsilon for normalization. -#' @param resnet_act_fn Character. Activation function. -#' @param spatio_temporal_scale Logical. Whether to use downsampling. -#' @param downsample_type Character. Type of downsampling. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx2_video_down_block3d <- torch::nn_module( - "LTX2VideoDownBlock3D", - - initialize = function( - in_channels, - out_channels = NULL, - num_layers = 1L, - dropout = 0.0, - resnet_eps = 1e-6, - resnet_act_fn = "swish", - spatio_temporal_scale = TRUE, - downsample_type = "conv", - spatial_padding_mode = "zeros" - ) { - if (is.null(out_channels)) out_channels <- in_channels - - # ResNet layers - resnets <- list() - for (i in seq_len(num_layers)) { - resnets[[i]] <- ltx2_video_resnet_block3d( - in_channels = in_channels, - out_channels = in_channels, # Note: same channels within block - dropout = dropout, - eps = resnet_eps, - non_linearity = resnet_act_fn, - spatial_padding_mode = spatial_padding_mode - ) - } - self$resnets <- torch::nn_module_list(resnets) - - # Downsampler - self$downsamplers <- NULL - if (spatio_temporal_scale) { - if (downsample_type == "conv") { - self$downsamplers <- torch::nn_module_list(list( - ltx2_video_causal_conv3d( - in_channels = in_channels, - out_channels = in_channels, - kernel_size = 3L, - stride = c(2L, 2L, 2L), - spatial_padding_mode = spatial_padding_mode - ) - )) - } else if (downsample_type == "spatial") { - self$downsamplers <- torch::nn_module_list(list( - ltx_video_downsampler3d( - in_channels = in_channels, - out_channels = out_channels, - stride = c(1L, 2L, 2L), - spatial_padding_mode = spatial_padding_mode - ) - )) - } else if (downsample_type == "temporal") { - self$downsamplers <- torch::nn_module_list(list( - ltx_video_downsampler3d( - in_channels = in_channels, - out_channels = out_channels, - stride = c(2L, 1L, 1L), - spatial_padding_mode = spatial_padding_mode - ) - )) - } else if (downsample_type == "spatiotemporal") { - self$downsamplers <- torch::nn_module_list(list( - ltx_video_downsampler3d( - in_channels = in_channels, - out_channels = out_channels, - stride = c(2L, 2L, 2L), - spatial_padding_mode = spatial_padding_mode - ) - )) - } - } - }, - - forward = function( - hidden_states, - temb = NULL, - generator = NULL, - causal = TRUE - ) { - for (i in seq_along(self$resnets)) { - hidden_states <- self$resnets[[i]](hidden_states, temb, generator, causal = causal) - } - - if (!is.null(self$downsamplers)) { - for (i in seq_along(self$downsamplers)) { - hidden_states <- self$downsamplers[[i]](hidden_states, causal = causal) - } - } - - hidden_states - } -) - -#' LTX2 Video Mid Block 3D -#' -#' Middle block with ResNet layers and optional timestep conditioning. -#' -#' @param in_channels Integer. Input channels. -#' @param num_layers Integer. Number of ResNet layers. -#' @param dropout Numeric. Dropout rate. -#' @param resnet_eps Numeric. Epsilon for normalization. -#' @param resnet_act_fn Character. Activation function. -#' @param inject_noise Logical. Whether to inject noise. -#' @param timestep_conditioning Logical. Whether to use timestep conditioning. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx2_video_mid_block3d <- torch::nn_module( - "LTX2VideoMidBlock3d", - - initialize = function( - in_channels, - num_layers = 1L, - dropout = 0.0, - resnet_eps = 1e-6, - resnet_act_fn = "swish", - inject_noise = FALSE, - timestep_conditioning = FALSE, - spatial_padding_mode = "zeros" - ) { - # Time embedder for timestep conditioning (TODO: implement if needed) - self$time_embedder <- NULL - - resnets <- list() - for (i in seq_len(num_layers)) { - resnets[[i]] <- ltx2_video_resnet_block3d( - in_channels = in_channels, - out_channels = in_channels, - dropout = dropout, - eps = resnet_eps, - non_linearity = resnet_act_fn, - inject_noise = inject_noise, - timestep_conditioning = timestep_conditioning, - spatial_padding_mode = spatial_padding_mode - ) - } - self$resnets <- torch::nn_module_list(resnets) - }, - - forward = function( - hidden_states, - temb = NULL, - generator = NULL, - causal = TRUE - ) { - for (i in seq_along(self$resnets)) { - hidden_states <- self$resnets[[i]](hidden_states, temb, generator, causal = causal) - } - hidden_states - } -) - -#' LTX2 Video Up Block 3D -#' -#' Decoder up block with upsampling and ResNet layers. -#' -#' @param in_channels Integer. Input channels. -#' @param out_channels Integer or NULL. Output channels. -#' @param num_layers Integer. Number of ResNet layers. -#' @param dropout Numeric. Dropout rate. -#' @param resnet_eps Numeric. Epsilon for normalization. -#' @param resnet_act_fn Character. Activation function. -#' @param spatio_temporal_scale Logical. Whether to use upsampling. -#' @param inject_noise Logical. Whether to inject noise. -#' @param timestep_conditioning Logical. Whether to use timestep conditioning. -#' @param upsample_residual Logical. Whether upsampler uses residual. -#' @param upscale_factor Integer. Channel upscale factor. -#' @param spatial_padding_mode Character. Padding mode. -#' @export -ltx2_video_up_block3d <- torch::nn_module( - "LTX2VideoUpBlock3d", - - initialize = function( - in_channels, - out_channels = NULL, - num_layers = 1L, - dropout = 0.0, - resnet_eps = 1e-6, - resnet_act_fn = "swish", - spatio_temporal_scale = TRUE, - inject_noise = FALSE, - timestep_conditioning = FALSE, - upsample_residual = FALSE, - upscale_factor = 1L, - spatial_padding_mode = "zeros" - ) { - if (is.null(out_channels)) out_channels <- in_channels - - # Time embedder (TODO: implement if needed) - self$time_embedder <- NULL - - # Input conv if channels differ - self$conv_in <- NULL - if (in_channels != out_channels) { - self$conv_in <- ltx2_video_resnet_block3d( - in_channels = in_channels, - out_channels = out_channels, - dropout = dropout, - eps = resnet_eps, - non_linearity = resnet_act_fn, - inject_noise = inject_noise, - timestep_conditioning = timestep_conditioning, - spatial_padding_mode = spatial_padding_mode - ) - } - - # Upsampler - self$upsamplers <- NULL - if (spatio_temporal_scale) { - self$upsamplers <- torch::nn_module_list(list( - ltx_video_upsampler3d( - in_channels = out_channels * upscale_factor, - stride = c(2L, 2L, 2L), - residual = upsample_residual, - upscale_factor = upscale_factor, - spatial_padding_mode = spatial_padding_mode - ) - )) - } - - # ResNet layers - resnets <- list() - for (i in seq_len(num_layers)) { - resnets[[i]] <- ltx2_video_resnet_block3d( - in_channels = out_channels, - out_channels = out_channels, - dropout = dropout, - eps = resnet_eps, - non_linearity = resnet_act_fn, - inject_noise = inject_noise, - timestep_conditioning = timestep_conditioning, - spatial_padding_mode = spatial_padding_mode - ) - } - self$resnets <- torch::nn_module_list(resnets) - }, - - forward = function( - hidden_states, - temb = NULL, - generator = NULL, - causal = TRUE - ) { - if (!is.null(self$conv_in)) { - hidden_states <- self$conv_in(hidden_states, temb, generator, causal = causal) - } - - if (!is.null(self$upsamplers)) { - for (i in seq_along(self$upsamplers)) { - hidden_states <- self$upsamplers[[i]](hidden_states, causal = causal) - } - } - - for (i in seq_along(self$resnets)) { - hidden_states <- self$resnets[[i]](hidden_states, temb, generator, causal = causal) - } - - hidden_states - } -) - diff --git a/R/vocoder_ltx23.R b/R/vocoder_ltx23.R new file mode 100644 index 0000000..35c0a29 --- /dev/null +++ b/R/vocoder_ltx23.R @@ -0,0 +1,505 @@ +#' LTX-2.3 Vocoder with Bandwidth Extension +#' +#' Fresh R port of the LTX-2 BigVGAN-style vocoder from the diffusers +#' reference (Apache-2.0, pipelines/ltx2/vocoder.py). The 2.3 vocoder +#' runs a 16 kHz stage (hidden 1536, snakebeta activations with +#' anti-aliased up/downsampling), re-analyzes its output into a causal +#' log-mel spectrogram, and feeds a bandwidth-extension vocoder whose +#' residual is added to a Hann-resampled skip path for 48 kHz output. +#' The Kaiser sinc / Hann filters and STFT bases are checkpoint buffers. +#' Runs in float32 (small model; snakebeta is precision-sensitive). +#' +#' @name vocoder_ltx23 +NULL + +# Kaiser window (periodic = FALSE), base-R besselI implementation used +# when building filters at init; checkpoint buffers override them. +.ltx23_kaiser_window <- function(n, beta) { + if (n == 1L) { + return(torch::torch_ones(1L)) + } + k <- seq(0L, n - 1L) + ratio <- (2 * k / (n - 1)) - 1 + vals <- besselI(beta * sqrt(pmax(1 - ratio ^ 2, 0)), 0) / besselI(beta, 0) + torch::torch_tensor(vals, dtype = torch::torch_float32()) +} + +#' Kaiser sinc low-pass filter kernel +#' +#' @param cutoff Numeric. Normalized cutoff in (0, 0.5]. +#' @param half_width Numeric. Transition band half width. +#' @param kernel_size Integer. +#' +#' @return Tensor [kernel_size]. +#' +#' @export +ltx23_kaiser_sinc_filter1d <- function(cutoff, half_width, kernel_size) { + delta_f <- 4 * half_width + half_size <- kernel_size %/% 2L + amplitude <- 2.285 * (half_size - 1) * pi * delta_f + 7.95 + beta <- if (amplitude > 50) { + 0.1102 * (amplitude - 8.7) + } else if (amplitude >= 21) { + 0.5842 * (amplitude - 21) ^ 0.4 + 0.07886 * (amplitude - 21) + } else { + 0 + } + + window <- .ltx23_kaiser_window(kernel_size, beta) + + even <- kernel_size %% 2L == 0L + time <- if (even) { + torch::torch_arange(start = -half_size, end = half_size - 1, + dtype = torch::torch_float32()) + 0.5 + } else { + torch::torch_arange(start = 0, end = kernel_size - 1, + dtype = torch::torch_float32()) - half_size + } + + if (cutoff == 0) { + return(torch::torch_zeros_like(time)) + } + time <- 2 * cutoff * time + sinc <- torch::torch_where( + time == 0, + torch::torch_ones_like(time), + torch::torch_sin(pi * time) / pi / time + ) + filter <- 2 * cutoff * window * sinc + filter / filter$sum() +} + +#' Anti-aliasing 1D downsampler (low-pass then stride) +#' +#' @param ratio Integer. Downsampling ratio. +#' @param kernel_size Integer or NULL (default 6*ratio rounded even). +#' +#' @export +ltx23_downsample1d <- torch::nn_module( + "ltx23_downsample1d", + initialize = function(ratio = 2L, kernel_size = NULL) { + self$ratio <- as.integer(ratio) + self$kernel_size <- as.integer(kernel_size %||% (as.integer(6 * ratio %/% 2) * 2L)) + self$pad_left <- self$kernel_size %/% 2L + (self$kernel_size %% 2L) - 1L + self$pad_right <- self$kernel_size %/% 2L + + lp <- ltx23_kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, self$kernel_size) + self$filter <- torch::nn_buffer(lp$view(c(1L, 1L, self$kernel_size))) +}, + forward = function(x) { + num_channels <- x$shape[2] + x <- torch::nnf_pad(x, c(self$pad_left, self$pad_right), mode = "replicate") + torch::nnf_conv1d(x, self$filter$expand(c(num_channels, -1L, -1L)), + stride = self$ratio, groups = num_channels) +} +) + +#' Anti-aliasing 1D upsampler (transposed low-pass) +#' +#' @param ratio Integer. Upsampling ratio. +#' @param kernel_size Integer or NULL. +#' @param window_type "kaiser" (BigVGAN default) or "hann" (final resampler). +#' @param persistent Logical. Register the filter as a buffer (present in +#' checkpoints); FALSE stores the computed filter as a plain field. +#' +#' @export +ltx23_upsample1d <- torch::nn_module( + "ltx23_upsample1d", + initialize = function(ratio = 2L, kernel_size = NULL, window_type = "kaiser", + persistent = TRUE) { + self$ratio <- as.integer(ratio) + + if (window_type == "hann") { + rolloff <- 0.99 + lowpass_filter_width <- 6L + width <- as.integer(ceiling(lowpass_filter_width / rolloff)) + self$kernel_size <- 2L * width * self$ratio + 1L + self$pad <- width + self$pad_left <- 2L * width * self$ratio + self$pad_right <- self$kernel_size - self$ratio + + time_axis <- (torch::torch_arange(start = 0, + end = self$kernel_size - 1, dtype = torch::torch_float32()) / ratio - width) * rolloff + time_clamped <- time_axis$clamp(-lowpass_filter_width, lowpass_filter_width) + window <- torch::torch_cos(time_clamped * pi / lowpass_filter_width / 2) ^ 2 + # sinc(x) = sin(pi x) / (pi x), 1 at x = 0 + sinc <- torch::torch_where( + time_axis == 0, + torch::torch_ones_like(time_axis), + torch::torch_sin(pi * time_axis) / (pi * time_axis) + ) + filt <- (sinc * window * rolloff / ratio)$view(c(1L, 1L, -1L)) + } else { + self$kernel_size <- as.integer(kernel_size %||% (as.integer(6 * ratio %/% 2) * 2L)) + self$pad <- self$kernel_size %/% self$ratio - 1L + self$pad_left <- self$pad * self$ratio + (self$kernel_size - self$ratio) %/% 2L + self$pad_right <- self$pad * self$ratio + (self$kernel_size - self$ratio + 1L) %/% 2L + filt <- ltx23_kaiser_sinc_filter1d( + 0.5 / ratio, 0.6 / ratio, self$kernel_size + )$view(c(1L, 1L, -1L)) + } + if (persistent) { + self$filter <- torch::nn_buffer(filt) + } else { + # Computed filter, absent from checkpoints (moved at use time) + self$filter <- filt + } +}, + forward = function(x) { + num_channels <- x$shape[2] + x <- torch::nnf_pad(x, c(self$pad, self$pad), mode = "replicate") + lp <- self$filter$to(dtype = x$dtype, device = x$device)$expand(c(num_channels, -1L, -1L)) + x <- torch::nnf_conv_transpose1d(x, lp, stride = self$ratio, + groups = num_channels)$mul(self$ratio) + n <- x$shape[3] + x$narrow(3L, self$pad_left + 1L, n - self$pad_left - self$pad_right) +} +) + +#' SnakeBeta activation +#' +#' \code{x + (1 / (beta + eps)) * sin(x * alpha)^2} with per-channel +#' log-scale alpha/beta parameters. +#' +#' @param channels Integer. +#' @param eps Numeric. +#' +#' @export +ltx23_snake_beta <- torch::nn_module( + "ltx23_snake_beta", + initialize = function(channels, eps = 1e-9) { + self$eps <- eps + self$alpha <- torch::nn_parameter(torch::torch_zeros(channels)) + self$beta <- torch::nn_parameter(torch::torch_zeros(channels)) +}, + forward = function(hidden_states) { + shape <- rep(1L, hidden_states$ndim) + shape[2] <- -1L + alpha <- torch::torch_exp(self$alpha$view(shape)) + beta <- torch::torch_exp(self$beta$view(shape)) + hidden_states + torch::torch_sin(hidden_states * alpha)$pow(2) * + beta$add(self$eps)$reciprocal() +} +) + +#' Anti-aliased activation +#' +#' Upsample 2x, apply the activation, downsample 2x. +#' +#' @param channels Integer. Channels for the SnakeBeta activation. +#' @param ratio,kernel_size Integers. Resampling config. +#' +#' @export +ltx23_antialias_act1d <- torch::nn_module( + "ltx23_antialias_act1d", + initialize = function(channels, ratio = 2L, kernel_size = 12L) { + self$upsample <- ltx23_upsample1d(ratio = ratio, kernel_size = kernel_size) + self$act <- ltx23_snake_beta(channels) + self$downsample <- ltx23_downsample1d(ratio = ratio, + kernel_size = kernel_size) +}, + forward = function(x) { + self$downsample(self$act(self$upsample(x))) +} +) + +#' Vocoder ResNet block (AMP) +#' +#' Dilated conv pairs, each preceded by an anti-aliased SnakeBeta +#' activation, with residual connections. +#' +#' @param channels Integer. +#' @param kernel_size Integer. +#' @param dilations Integer vector. +#' @param antialias_ratio,antialias_kernel_size Integers. +#' +#' @export +ltx23_vocoder_resblock <- torch::nn_module( + "ltx23_vocoder_resblock", + initialize = function( + channels, + kernel_size = 3L, + dilations = c(1L, 3L, 5L), + antialias_ratio = 2L, + antialias_kernel_size = 12L + ) { + self$convs1 <- torch::nn_module_list(lapply(dilations, function(d) { + torch::nn_conv1d(channels, channels, kernel_size, stride = 1L, + dilation = d, padding = d * (kernel_size - 1L) %/% 2L) + })) + self$acts1 <- torch::nn_module_list(lapply(dilations, function(d) { + ltx23_antialias_act1d(channels, antialias_ratio, antialias_kernel_size) + })) + self$convs2 <- torch::nn_module_list(lapply(dilations, function(d) { + torch::nn_conv1d(channels, channels, kernel_size, stride = 1L, + dilation = 1L, padding = (kernel_size - 1L) %/% 2L) + })) + self$acts2 <- torch::nn_module_list(lapply(dilations, function(d) { + ltx23_antialias_act1d(channels, antialias_ratio, antialias_kernel_size) + })) +}, + forward = function(x) { + for (i in seq_along(self$convs1)) { + xt <- self$acts1[[i]](x) + xt <- self$convs1[[i]](xt) + xt <- self$acts2[[i]](xt) + xt <- self$convs2[[i]](xt) + x <- x + xt + } + x +} +) + +#' LTX-2.3 vocoder stage +#' +#' Mel spectrogram [B, C, T, M] -> waveform [B, out_channels, samples]. +#' Channel and mel dims are flattened into conv channels; each upsample +#' stage halves the channel count and averages three parallel ResNet +#' branches. +#' +#' @param in_channels Integer. Flattened input channels (C * mel bins / 1). +#' @param hidden_channels Integer. +#' @param out_channels Integer. +#' @param upsample_kernel_sizes,upsample_factors Integer vectors. +#' @param resnet_kernel_sizes Integer vector. +#' @param resnet_dilations List of integer vectors. +#' @param antialias_ratio,antialias_kernel_size Integers. +#' @param final_bias Logical. +#' +#' @export +ltx23_vocoder <- torch::nn_module( + "ltx23_vocoder", + initialize = function( + in_channels = 128L, + hidden_channels = 1536L, + out_channels = 2L, + upsample_kernel_sizes = c(11L, 4L, 4L, 4L, 4L, 4L), + upsample_factors = c(5L, 2L, 2L, 2L, 2L, 2L), + resnet_kernel_sizes = c(3L, 7L, 11L), + resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), c(1L, 3L, 5L)), + antialias_ratio = 2L, + antialias_kernel_size = 12L, + final_bias = FALSE + ) { + self$num_upsample_layers <- length(upsample_kernel_sizes) + self$resnets_per_upsample <- length(resnet_kernel_sizes) + + self$conv_in <- torch::nn_conv1d(in_channels, hidden_channels, + kernel_size = 7L, stride = 1L, + padding = 3L) + + upsamplers <- list() + resnets <- list() + input_channels <- hidden_channels + for (i in seq_along(upsample_factors)) { + output_channels <- input_channels %/% 2L + upsamplers[[i]] <- torch::nn_conv_transpose1d( + input_channels, output_channels, upsample_kernel_sizes[i], + stride = upsample_factors[i], + padding = (upsample_kernel_sizes[i] - upsample_factors[i]) %/% 2L + ) + for (j in seq_along(resnet_kernel_sizes)) { + resnets[[length(resnets) + 1L]] <- ltx23_vocoder_resblock( + channels = output_channels, + kernel_size = resnet_kernel_sizes[j], + dilations = resnet_dilations[[j]], + antialias_ratio = antialias_ratio, + antialias_kernel_size = antialias_kernel_size + ) + } + input_channels <- output_channels + } + self$upsamplers <- torch::nn_module_list(upsamplers) + self$resnets <- torch::nn_module_list(resnets) + + self$act_out <- ltx23_antialias_act1d(output_channels, antialias_ratio, + antialias_kernel_size) + self$conv_out <- torch::nn_conv1d(output_channels, out_channels, 7L, + stride = 1L, padding = 3L, bias = final_bias) +}, + forward = function(hidden_states, time_last = FALSE) { + if (!time_last) { + hidden_states <- hidden_states$transpose(3L, 4L) + } + hidden_states <- hidden_states$flatten(start_dim = 2L, end_dim = 3L) + + hidden_states <- self$conv_in(hidden_states) + for (i in seq_len(self$num_upsample_layers)) { + hidden_states <- self$upsamplers[[i]](hidden_states) + start <- (i - 1L) * self$resnets_per_upsample + branch_outputs <- lapply(seq_len(self$resnets_per_upsample), function(j) { + self$resnets[[start + j]](hidden_states) + }) + hidden_states <- torch::torch_mean( + torch::torch_stack(branch_outputs, dim = 1L), dim = 1L + ) + } + hidden_states <- self$act_out(hidden_states) + self$conv_out(hidden_states) +} +) + +# Causal STFT with checkpoint-loaded DFT bases +ltx23_causal_stft <- torch::nn_module( + "ltx23_causal_stft", + initialize = function(filter_length = 512L, hop_length = 80L, window_length = 512L) { + self$hop_length <- as.integer(hop_length) + self$window_length <- as.integer(window_length) + n_freqs <- filter_length %/% 2L + 1L + self$forward_basis <- torch::nn_buffer(torch::torch_zeros(n_freqs * 2L, + 1L, filter_length)) + self$inverse_basis <- torch::nn_buffer(torch::torch_zeros(n_freqs * 2L, 1L, filter_length)) +}, + forward = function(waveform) { + if (waveform$ndim == 2L) waveform <- waveform$unsqueeze(2L) + left_pad <- max(0L, self$window_length - self$hop_length) + waveform <- torch::nnf_pad(waveform, c(left_pad, 0L)) + + spec <- torch::nnf_conv1d(waveform, self$forward_basis, stride = self$hop_length) + n_freqs <- spec$shape[2] %/% 2L + real <- spec$narrow(2L, 1L, n_freqs) + imag <- spec$narrow(2L, n_freqs + 1L, n_freqs) + magnitude <- torch::torch_sqrt(real ^ 2 + imag ^ 2) + list(magnitude = magnitude) +} +) + +#' Causal log-mel spectrogram with checkpoint-loaded bases +#' +#' @param filter_length,hop_length,window_length,num_mel_channels Integers. +#' +#' @export +ltx23_mel_stft <- torch::nn_module( + "ltx23_mel_stft", + initialize = function( + filter_length = 512L, + hop_length = 80L, + window_length = 512L, + num_mel_channels = 64L + ) { + self$stft_fn <- ltx23_causal_stft(filter_length, hop_length, window_length) + num_freqs <- filter_length %/% 2L + 1L + self$mel_basis <- torch::nn_buffer(torch::torch_zeros(num_mel_channels, + num_freqs)) +}, + forward = function(waveform) { + magnitude <- self$stft_fn(waveform)$magnitude + mel <- torch::torch_matmul(self$mel_basis$to(dtype = magnitude$dtype), magnitude) + torch::torch_log(torch::torch_clamp(mel, min = 1e-5)) +} +) + +#' LTX-2.3 vocoder with bandwidth extension +#' +#' Full mel [B, 2, T, 64] -> 48 kHz stereo waveform pipeline: 16 kHz +#' vocoder, causal mel re-analysis, BWE vocoder residual added to a +#' Hann-resampled skip path, clamped to [-1, 1]. +#' +#' @param in_channels,bwe_in_channels Integers. Flattened mel input channels. +#' @param hidden_channels,bwe_hidden_channels Integers. +#' @param out_channels Integer. Audio channels. +#' @param upsample_kernel_sizes,upsample_factors,bwe_upsample_kernel_sizes,bwe_upsample_factors +#' Integer vectors. Per-stage transposed-conv configs. +#' @param resnet_kernel_sizes,bwe_resnet_kernel_sizes Integer vectors. +#' @param resnet_dilations,bwe_resnet_dilations Lists of integer vectors. +#' @param filter_length,window_length,num_mel_channels Integers. Mel +#' re-analysis configuration. +#' @param input_sampling_rate,output_sampling_rate Integers. +#' @param hop_length Integer. Mel analysis hop. +#' +#' @export +ltx23_vocoder_with_bwe <- torch::nn_module( + "ltx23_vocoder_with_bwe", + initialize = function( + in_channels = 128L, + hidden_channels = 1536L, + out_channels = 2L, + upsample_kernel_sizes = c(11L, 4L, 4L, 4L, 4L, 4L), + upsample_factors = c(5L, 2L, 2L, 2L, 2L, 2L), + resnet_kernel_sizes = c(3L, 7L, 11L), + resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), c(1L, 3L, 5L)), + bwe_in_channels = 128L, + bwe_hidden_channels = 512L, + bwe_upsample_kernel_sizes = c(12L, 11L, 4L, 4L, 4L), + bwe_upsample_factors = c(6L, 5L, 2L, 2L, 2L), + bwe_resnet_kernel_sizes = c(3L, 7L, 11L), + bwe_resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), + c(1L, 3L, 5L)), + filter_length = 512L, + hop_length = 80L, + window_length = 512L, + num_mel_channels = 64L, + input_sampling_rate = 16000L, + output_sampling_rate = 48000L + ) { + self$hop_length <- as.integer(hop_length) + self$rate_ratio <- as.integer(output_sampling_rate %/% input_sampling_rate) + + self$vocoder <- ltx23_vocoder( + in_channels = in_channels, + hidden_channels = hidden_channels, + out_channels = out_channels, + upsample_kernel_sizes = upsample_kernel_sizes, + upsample_factors = upsample_factors, + resnet_kernel_sizes = resnet_kernel_sizes, + resnet_dilations = resnet_dilations + ) + self$bwe_generator <- ltx23_vocoder( + in_channels = bwe_in_channels, + hidden_channels = bwe_hidden_channels, + out_channels = out_channels, + upsample_kernel_sizes = bwe_upsample_kernel_sizes, + upsample_factors = bwe_upsample_factors, + resnet_kernel_sizes = bwe_resnet_kernel_sizes, + resnet_dilations = bwe_resnet_dilations + ) + self$mel_stft <- ltx23_mel_stft( + filter_length = filter_length, + hop_length = hop_length, + window_length = window_length, + num_mel_channels = num_mel_channels + ) + self$resampler <- ltx23_upsample1d(ratio = self$rate_ratio, window_type = "hann", + persistent = FALSE) +}, + forward = function(mel_spec) { + x <- self$vocoder(mel_spec) + num_channels <- x$shape[2] + num_samples <- x$shape[3] + + remainder <- num_samples %% self$hop_length + if (remainder != 0L) { + x <- torch::nnf_pad(x, c(0L, self$hop_length - remainder)) + } + + mel <- self$mel_stft(x$flatten(start_dim = 1L, end_dim = 2L)) + mel <- mel$unflatten(1L, c(-1L, num_channels)) + + # [B, C, mel_bins, frames] -> [B, C, frames, mel_bins] + residual <- self$bwe_generator(mel$transpose(3L, 4L)) + + skip <- self$resampler(x) + waveform <- torch::torch_clamp(residual + skip, -1, 1) + output_samples <- num_samples * self$rate_ratio + waveform$narrow(3L, 1L, min(output_samples, waveform$shape[3])) +} +) + +#' Map an official vocoder checkpoint key to the R module name +#' +#' @param key Character. Checkpoint key. +#' +#' @return Character. Module parameter/buffer name. +#' +#' @export +ltx23_map_vocoder_key <- function(key) { + key <- sub("^vocoder\\.", "", key) + key <- gsub("resblocks", "resnets", key, fixed = TRUE) + key <- gsub("conv_pre", "conv_in", key, fixed = TRUE) + key <- gsub("conv_post", "conv_out", key, fixed = TRUE) + key <- gsub("act_post", "act_out", key, fixed = TRUE) + key <- gsub("downsample.lowpass", "downsample", key, fixed = TRUE) + key <- gsub(".ups.", ".upsamplers.", key, fixed = TRUE) + key <- sub("^ups\\.", "upsamplers.", key) + key +} diff --git a/R/vram.R b/R/vram.R new file mode 100644 index 0000000..a835ffa --- /dev/null +++ b/R/vram.R @@ -0,0 +1,204 @@ +#' VRAM Detection and Management Utilities +#' +#' Device detection, VRAM reporting, and module offloading helpers shared +#' by the image and video pipelines. +#' +#' @name vram +NULL + +#' Check if GPU is Blackwell Architecture +#' +#' Blackwell GPUs (RTX 50xx) may need special handling. +#' +#' @return Logical. TRUE if Blackwell GPU detected. +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' if (is_blackwell_gpu()) { +#' message("Using Blackwell-compatible settings") +#' } +#' } +is_blackwell_gpu <- function() { + # Use gpuctl if available + if (requireNamespace("gpu.ctl", quietly = TRUE)) { + return(gpu.ctl::gpu_is_blackwell()) + } + + # Fallback: check compute capability via torch + if (torch::cuda_is_available()) { + cap <- tryCatch(torch::cuda_get_device_capability(0L), + error = function(e) NULL) + if (!is.null(cap)) { + # Blackwell is compute 12.x + return(as.integer(cap[[1]]) >= 12L) + } + } + + FALSE +} + +#' Detect Available VRAM +#' +#' Uses gpuctl if available. +#' +#' @param use_free Logical. If TRUE, return free VRAM. If FALSE, return total. +#' +#' @return Numeric. VRAM in GB, or 0 if no GPU detected. +#' @keywords internal +.detect_vram <- function(use_free = FALSE) { + # Try gpuctl (preferred - uses nvidia-smi) + if (requireNamespace("gpu.ctl", quietly = TRUE)) { + info <- gpu.ctl::gpu_detect() + if (!is.null(info)) { + if (use_free && !is.null(info$vram_free_gb)) { + return(info$vram_free_gb) + } + if (!is.null(info$vram_total_gb)) { + return(info$vram_total_gb) + } + } + } + + # Fallback: ask nvidia-smi directly + smi <- suppressWarnings(tryCatch( + system2("nvidia-smi", + c(paste0("--query-gpu=memory.", + if (use_free) "free" else "total"), + "--format=csv,noheader,nounits"), + stdout = TRUE, stderr = FALSE), + error = function(e) character(0) + )) + mb <- suppressWarnings(as.numeric(smi[1])) + if (isTRUE(mb > 0)) { + return(mb / 1024) + } + + # Fallback: check if CUDA available but can't determine VRAM + if (torch::cuda_is_available()) { + # Conservative estimate - assume 8GB if we can't detect + message("Could not detect VRAM. Install gpuctl for accurate detection.") + return(8) + } + + # No GPU detected + 0 +} + +#' Offload Module to CPU +#' +#' Moves a torch module and all its parameters to CPU. +#' +#' @param module A torch nn_module. +#' @param gc Logical. Run garbage collection after offload. +#' +#' @return The module (modified in place). +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' model$to(device = "cuda") +#' output <- model(x) +#' offload_to_cpu(model) +#' } +offload_to_cpu <- function(module, gc = TRUE) { + module$to(device = "cpu") + if (gc && torch::cuda_is_available()) { + gc() + torch::cuda_empty_cache() + } + invisible(module) +} + +#' Load Module to GPU +#' +#' Moves a torch module and all its parameters to CUDA. +#' +#' @param module A torch nn_module. +#' @param device Character. Target device (default "cuda"). +#' +#' @return The module (modified in place). +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' load_to_gpu(model) +#' output <- model(x) +#' offload_to_cpu(model) +#' } +load_to_gpu <- function(module, device = "cuda") { + module$to(device = device) + invisible(module) +} + +#' Report VRAM Usage +#' +#' Prints current VRAM usage using gpuctl. +#' +#' @param label Character. Label for the report. +#' +#' @return Invisibly returns a list with used and free VRAM in GB. +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' vram_report("After model load") +#' } +vram_report <- function(label = "") { + if (!torch::cuda_is_available()) { + message("[", label, "] No CUDA available") + return(invisible(list(used = 0, free = 0))) + } + + # Use gpuctl for accurate reporting + if (requireNamespace("gpu.ctl", quietly = TRUE)) { + info <- gpu.ctl::gpu_detect() + if (!is.null(info)) { + used <- info$vram_used_gb + free <- info$vram_free_gb + message(sprintf("[%s] VRAM: %.2f GB used, %.2f GB free", label, + used, free)) + return(invisible(list(used = used, free = free))) + } + } + + message("[", label, "] VRAM: (install gpuctl for detailed stats)") + invisible(list(used = NA, free = NA)) +} + +#' Clear VRAM Cache +#' +#' Forces garbage collection and clears CUDA memory cache. +#' +#' @param verbose Logical. Print memory status before/after. +#' +#' @return Invisibly returns NULL. +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' clear_vram() +#' } +clear_vram <- function(verbose = FALSE) { + if (!torch::cuda_is_available()) { + return(invisible(NULL)) + } + + if (verbose) { + vram_report("Before clear") + } + + gc() + tryCatch(torch::cuda_empty_cache(), error = function(e) NULL) + + if (verbose) { + vram_report("After clear") + } + + invisible(NULL) +} diff --git a/R/zzz.R b/R/zzz.R index dc8f7e9..fde3513 100644 --- a/R/zzz.R +++ b/R/zzz.R @@ -3,4 +3,3 @@ NULL # torch indexing uses `..` as a Python-style ellipsis utils::globalVariables("..") - diff --git a/inst/COPYRIGHTS b/inst/COPYRIGHTS new file mode 100644 index 0000000..b561f16 --- /dev/null +++ b/inst/COPYRIGHTS @@ -0,0 +1,39 @@ +COPYRIGHT AND ATTRIBUTION NOTES FOR diffuseR + +diffuseR is licensed under the Apache License (>= 2). Portions of the +package are R ports of, or derived from, third-party work: + +1. HuggingFace diffusers (Apache License 2.0) + Copyright The HuggingFace Team. + https://github.com/huggingface/diffusers + The LTX video pipeline components (transformer, video/audio VAE, + text-embedding connectors, vocoder, latent upsampler, schedulers, and + pipeline logic including the AdaIN latent filter and tone mapping) are + R ports written from the diffusers reference implementation. + The FlowMatch Euler scheduler is ported from + diffusers/schedulers/scheduling_flow_match_euler_discrete.py. + +2. HuggingFace transformers (Apache License 2.0) + Copyright The HuggingFace Team. + https://github.com/huggingface/transformers + The Gemma3 text encoder (R/gemma3_text_encoder.R) is an R port of the + Gemma3 implementation in transformers. + +3. Lightricks Ltd. + https://github.com/Lightricks/LTX-2 + Numeric constants (e.g. distilled sigma schedules), the single-file + checkpoint key layout, and pipeline parameter defaults are facts taken + from the official LTX release. No code from the Lightricks repository + (LTX-2 Community License) is included in this package. + + LTX model weights are downloaded by the user directly from Lightricks' + HuggingFace repositories under the LTX-2 Community License and are not + part of, nor redistributed with, this package. + +4. Stability AI / Stable Diffusion (SD 2.1, SDXL) + Model weights are downloaded by the user under their respective + licenses (CreativeML Open RAIL++-M and similar) and are not part of, + nor redistributed with, this package. + +See inst/REFERENCES.md for a technique-by-technique account of where +each implementation idea comes from. diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md new file mode 100644 index 0000000..f985aba --- /dev/null +++ b/inst/REFERENCES.md @@ -0,0 +1,54 @@ +# LTX-2.3 implementation references + +Every technique in diffuseR's LTX-2.3 support traces to a public, +permissively licensed source or to our own measured engineering. None of +it derives from Wan2GP (WanGP Community License) or its `mmgp` module; +this file documents the actual lineage, idea by idea. + +## Model architecture and pipeline + +| What | Source | +|---|---| +| DiT, VAEs, connectors, vocoder, upsampler module code | Ported from HuggingFace **diffusers** (Apache-2.0): `models/transformers/transformer_ltx2.py`, `models/autoencoders/autoencoder_kl_ltx2.py`, `autoencoder_kl_ltx2_audio.py`, `pipelines/ltx2/{connectors,vocoder,latent_upsampler}.py` | +| Pipeline flow, Euler velocity steps, latent packing, AdaIN filter, tone mapping | diffusers `pipelines/ltx2/pipeline_ltx2.py`, `pipeline_ltx2_latent_upsample.py` (Apache-2.0) | +| Distilled sigma schedules, default negative prompt | diffusers `pipelines/ltx2/utils.py` (Apache-2.0); values originate from Lightricks' LTX-2 release (numeric facts) | +| Single-file checkpoint key layout, embedded config, `model_version` metadata | Format facts of the official Lightricks checkpoints, cross-checked against diffusers `scripts/convert_ltx2_to_diffusers.py` (Apache-2.0) | +| Gemma3 text encoder | Ported from HuggingFace **transformers** (Apache-2.0) | + +## Quantization + +| What | Source | +|---|---| +| FP8 e4m3fn weights with per-tensor scales, upcast-in-forward; the exact linear cast set (attention + FFN projections in transformer blocks) | Official Lightricks LTX-2 quantization policy (design facts from `ltx-core/quantization/fp8_cast.py`, LTX-2 Community License — policy observed, no code taken); fp8 storage-with-upcast is also diffusers' documented "layerwise casting" (Apache-2.0) | +| NF4 4-bit format: 16-level NormalFloat quantile code, per-64-block absmax, packed nibbles | **QLoRA**: Dettmers, Pagnoni, Holtzman, Zettlemoyer (2023), "QLoRA: Efficient Finetuning of Quantized LLMs", arXiv:2305.14314; format details as implemented in bitsandbytes (MIT) — reimplemented here in pure torch ops | +| 4-bit weights for consumer-GPU video DiTs as a practice | Community standard: GGUF Q4 checkpoints for LTX/ComfyUI (e.g. city96/ComfyUI-GGUF ecosystem), diffusers bitsandbytes integration docs | + +## Conditioning (image-to-video, continuation, audio-driven) + +| What | Source | +|---|---| +| Prefix frame conditioning: VAE-encoded pixels as frozen latent tokens, per-token timestep `t * (1 - mask)`, frozen Euler updates, partial-noise option | Ported from diffusers `pipelines/ltx2/pipeline_ltx2_image2video.py` and `pipeline_ltx2_condition.py` (Apache-2.0), restricted to index-0 / strength-1 | +| Audio-driven generation (lip sync): user audio encoded to clean, frozen audio latents; audio stream sees timestep/sigma zero while video denoises attending to it | Our own design — the audio-side mirror of the Apache i2v conditioning pattern above (diffusers has no audio-input conditioning) | +| Audio VAE encoder module | Ported from diffusers `autoencoder_kl_ltx2_audio.py` `LTX2AudioEncoder` (Apache-2.0) | +| 16 kHz log-mel frontend (causal STFT 1024/160, 64 slaney-normed mel bins to 8 kHz) | Parameters are the checkpoint's embedded preprocessing config (numeric facts); STFT/mel constructions are textbook DSP, verified against the checkpoint's stored vocoder bases | + +## Memory management + +| What | Source | +|---|---| +| Phase-sequential component offloading (each component on the GPU only for its phase) | diffusers memory docs: `enable_model_cpu_offload`, group offloading (`docs/source/en/optimization/memory.md`, Apache-2.0) | +| CPU-resident weights streamed per layer | diffusers leaf-level/group offloading with pinned memory (same doc); long-standing technique (e.g. DeepSpeed ZeRO-Offload, arXiv:2101.06840) | +| Tiled VAE decode (overlapping tiles, crossfaded seams; spatial + temporal) | Direct port of diffusers `AutoencoderKLLTX2Video.tiled_decode` / `_temporal_tiled_decode` (Apache-2.0); tiled VAE decoding dates to the original Stable Diffusion AutoencoderKL (`enable_tiling`) | +| Attention query chunking under a memory budget | diffusers attention slicing (`enable_attention_slicing`, Apache-2.0); memory-efficient attention lineage: Rabe & Staats (2021), arXiv:2112.05682 | +| Pinned host memory + non-blocking transfers | PyTorch CUDA semantics documentation | +| `PYTORCH_CUDA_ALLOC_CONF=expandable_segments` | PyTorch CUDA caching-allocator documentation | +| R-level allocator GC tuning (`torch.cuda_allocator_reserved_rate`, `torch.threshold_call_gc`) | mlverse torch memory-management article (`vignettes/articles/memory-management.Rmd`); A/B methodology and measurements from our own whisper/chatterbox work (both public cornball-ai packages) | +| Persistent scratch buffers for dequantization and attention temporaries | Our own engineering, driven by R's lazy garbage collection of tensor handles (documented in the mlverse article above) | +| Per-block `gc()` for streamed-weight temporaries | Our own engineering (see whisper/chatterbox inference loops) | +| TorchScript-compiled block stack (`torch::jit_compile`, weights as `List[Tensor]`, fused `scaled_dot_product_attention`) | mlverse torch JIT API; the JIT-decode pattern from our own whisper (`decode_jit.R`) and chatterbox (`t3_jit.R`) work, both public cornball-ai packages | + +## Measured on our hardware (RTX 5060 Ti 16GB) + +Resolution caps, VRAM peaks, profile thresholds, and the GC A/B results +in `tasks/todo.md` and commit messages are our own measurements of the +above techniques, not third-party numbers. diff --git a/inst/tinytest/fixtures/audio_ltx23.safetensors b/inst/tinytest/fixtures/audio_ltx23.safetensors new file mode 100644 index 0000000..e00f861 Binary files /dev/null and b/inst/tinytest/fixtures/audio_ltx23.safetensors differ diff --git a/inst/tinytest/fixtures/connectors_ltx23.safetensors b/inst/tinytest/fixtures/connectors_ltx23.safetensors new file mode 100644 index 0000000..d23555b Binary files /dev/null and b/inst/tinytest/fixtures/connectors_ltx23.safetensors differ diff --git a/inst/tinytest/fixtures/dit_ltx23.safetensors b/inst/tinytest/fixtures/dit_ltx23.safetensors new file mode 100644 index 0000000..4cc5b8b Binary files /dev/null and b/inst/tinytest/fixtures/dit_ltx23.safetensors differ diff --git a/inst/tinytest/fixtures/rope_ltx23.safetensors b/inst/tinytest/fixtures/rope_ltx23.safetensors new file mode 100644 index 0000000..ca61e4b Binary files /dev/null and b/inst/tinytest/fixtures/rope_ltx23.safetensors differ diff --git a/inst/tinytest/fixtures/vae_ltx23.safetensors b/inst/tinytest/fixtures/vae_ltx23.safetensors new file mode 100644 index 0000000..5b34255 Binary files /dev/null and b/inst/tinytest/fixtures/vae_ltx23.safetensors differ diff --git a/inst/tinytest/test_audio_encode_ltx23.R b/inst/tinytest/test_audio_encode_ltx23.R new file mode 100644 index 0000000..0c914ad --- /dev/null +++ b/inst/tinytest/test_audio_encode_ltx23.R @@ -0,0 +1,57 @@ +# Audio conditioning frontend (R/audio_encode_ltx23.R): STFT/mel basis +# constructors, mel geometry, and encode length math. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) +torch::torch_manual_seed(31) + +# --- Basis constructors --------------------------------------------------------------- + +fb <- diffuseR:::.ltx23_stft_basis(64L) +expect_equal(as.integer(fb$shape), c(66L, 1L, 64L)) +# Row 0 is the DC row: pure window (cos(0) = 1) +w <- diffuseR:::.ltx23_hann(64L) +expect_true(as.numeric((fb[1, 1, ] - torch::torch_tensor(w))$abs()$max()) < 1e-6) +# The imaginary DC row is all zeros (-sin(0)) +expect_true(as.numeric(fb[34, 1, ]$abs()$max()) < 1e-6) + +mel <- diffuseR:::.ltx23_mel_filterbank(16000, 512L, 64L, 0, 8000) +expect_equal(as.integer(mel$shape), c(64L, 257L)) +# Every filter is nonnegative with positive mass +expect_true(as.numeric(mel$min()) >= 0) +expect_true(all(as.numeric(mel$sum(dim = 2L)) > 0)) + +# --- Frontend geometry --------------------------------------------------------------- + +frontend <- ltx23_audio_mel_frontend(filter_length = 64L, hop_length = 16L, + n_mels = 8L, sample_rate = 1600L, fmin = 0, fmax = 800) +x <- torch::torch_randn(2L, 1L, 320L) +torch::with_no_grad(m <- frontend(x)) +expect_equal(as.integer(m$shape)[1:2], c(2L, 8L)) +expect_true(all(is.finite(as.numeric(m$min())))) +# Log-clamped: never below log(1e-5) +expect_true(as.numeric(m$min()) >= log(1e-5) - 1e-4) + +# --- Encode length math --------------------------------------------------------------- + +# base_channels 128 so the stats buffer matches the packed feature dim +# (latent_channels * latent_mel_bins = 128), as in the real checkpoint +avae <- ltx23_audio_vae( + base_channels = 128L, output_channels = 2L, num_res_blocks = 1L, + latent_channels = 8L, ch_mult = c(1L, 1L, 1L), mel_bins = 64L +) +avae$eval() +small_frontend <- ltx23_audio_mel_frontend() +wav <- matrix(runif(2L * 8000L, -0.5, 0.5), nrow = 2L) +lat <- ltx23_encode_audio(avae, wav, audio_num_frames = 9L, + frontend = small_frontend) +# Packed [1, L, latent_channels * latent_mel_bins]; L == audio_num_frames +expect_equal(as.integer(lat$shape), c(1L, 9L, 8L * 16L)) + +# Longer target than the audio covers: right-padded, still exact length +lat2 <- ltx23_encode_audio(avae, wav, audio_num_frames = 25L, + frontend = small_frontend) +expect_equal(as.integer(lat2$shape)[2], 25L) diff --git a/inst/tinytest/test_audio_ltx23.R b/inst/tinytest/test_audio_ltx23.R new file mode 100644 index 0000000..a12b6c9 --- /dev/null +++ b/inst/tinytest/test_audio_ltx23.R @@ -0,0 +1,204 @@ +# Parity tests for the LTX-2.3 audio VAE decoder and vocoder port +# against diffusers reference fixtures (tools/gen_fixtures_audio.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "audio_ltx23.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/audio_ltx23.safetensors" +if (!file.exists(fixture_path)) exit_file("audio fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +load_group <- function(module, prefix) { + keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE) + w <- fx[keys] + names(w) <- sub(paste0("^", prefix, "\\."), "", keys) + dests <- c(module$named_parameters(), module$named_buffers()) + if (!setequal(names(w), names(dests))) { + stop( + prefix, ": name mismatch. missing dest: ", + paste(utils::head(setdiff(names(w), names(dests)), 4), collapse = ", "), + " | unfilled: ", + paste(utils::head(setdiff(names(dests), names(w)), 4), collapse = ", ") + ) + } + torch::with_no_grad({ + for (name in names(w)) dests[[name]]$copy_(w[[name]]) + }) + module$eval() + module +} + +# --- Audio VAE decoder -------------------------------------------------------- + +adec <- load_group( + ltx23_audio_decoder( + base_channels = 8L, output_channels = 2L, num_res_blocks = 1L, + latent_channels = 4L, ch_mult = c(1L, 2L), mel_bins = 8L + ), + "adec" +) +torch::with_no_grad(out_adec <- adec(fx$adec_x)) +expect_equal(as.integer(out_adec$shape), c(1L, 2L, 17L, 8L))# 5 * 4 - 3 frames +expect_true(max_abs_diff(out_adec, fx$adec_out) < 1e-5) + +# --- Kaiser sinc filter + window ------------------------------------------------ + +expect_true(max_abs_diff( + ltx23_kaiser_sinc_filter1d(0.25, 0.3, 12L), fx$kaiser_filt_12 +) < 1e-5) +expect_true(max_abs_diff( + ltx23_kaiser_sinc_filter1d(0.1, 0.12, 13L), fx$kaiser_filt_13 +) < 1e-5) +expect_true(max_abs_diff( + diffuseR:::.ltx23_kaiser_window(12L, 4.7), fx$kaiser_window_12 +) < 1e-6) + +# --- Up/DownSample1d --------------------------------------------------------------- + +ds1 <- ltx23_downsample1d(ratio = 2L, kernel_size = 12L) +torch::with_no_grad(out_down <- ds1(fx$w_x)) +expect_true(max_abs_diff(out_down, fx$w_down) < 1e-5) + +us1 <- ltx23_upsample1d(ratio = 2L, kernel_size = 12L) +torch::with_no_grad(out_up <- us1(fx$w_x)) +expect_equal(as.integer(out_up$shape), as.integer(fx$w_up$shape)) +expect_true(max_abs_diff(out_up, fx$w_up) < 1e-5) + +ush <- ltx23_upsample1d(ratio = 4L, window_type = "hann") +torch::with_no_grad(out_uph <- ush(fx$w_x)) +expect_equal(as.integer(out_uph$shape), as.integer(fx$w_up_hann$shape)) +expect_true(max_abs_diff(out_uph, fx$w_up_hann) < 1e-5) + +# --- SnakeBeta ----------------------------------------------------------------------- + +sb <- ltx23_snake_beta(3L) +torch::with_no_grad({ + sb$alpha$copy_(fx$sb.alpha) + sb$beta$copy_(fx$sb.beta) + out_sb <- sb(fx$w_x) +}) +expect_true(max_abs_diff(out_sb, fx$sb_out) < 1e-5) + +# --- Tiny vocoder stage ----------------------------------------------------------------- + +voc <- load_group( + ltx23_vocoder( + in_channels = 8L, hidden_channels = 16L, out_channels = 2L, + upsample_kernel_sizes = c(4L, 4L), upsample_factors = c(2L, 2L), + resnet_kernel_sizes = c(3L), resnet_dilations = list(c(1L, 3L)), + final_bias = FALSE + ), + "voc" +) +torch::with_no_grad(out_voc <- voc(fx$voc_x)) +expect_equal(as.integer(out_voc$shape), c(1L, 2L, 24L)) +expect_true(max_abs_diff(out_voc, fx$voc_out) < 1e-4) + +# --- Tiny BWE wrapper -------------------------------------------------------------------- + +bwe <- load_group( + ltx23_vocoder_with_bwe( + in_channels = 8L, hidden_channels = 16L, out_channels = 2L, + upsample_kernel_sizes = c(4L, 4L), upsample_factors = c(2L, 2L), + resnet_kernel_sizes = c(3L), resnet_dilations = list(c(1L, 3L)), + bwe_in_channels = 16L, bwe_hidden_channels = 8L, + bwe_upsample_kernel_sizes = c(8L, 4L), bwe_upsample_factors = c(4L, 2L), + bwe_resnet_kernel_sizes = c(3L), bwe_resnet_dilations = list(c(1L, 3L)), + filter_length = 8L, hop_length = 2L, window_length = 8L, + num_mel_channels = 8L, + input_sampling_rate = 100L, output_sampling_rate = 400L + ), + "bwe" +) +torch::with_no_grad(out_bwe <- bwe(fx$voc_x)) +expect_equal(as.integer(out_bwe$shape), c(1L, 2L, 96L)) +expect_true(max_abs_diff(out_bwe, fx$bwe_out) < 1e-4) + +# --- Key mappers ---------------------------------------------------------------------------- + +expect_equal( + ltx23_map_vocoder_key("vocoder.vocoder.conv_pre.weight"), + "vocoder.conv_in.weight" +) +expect_equal( + ltx23_map_vocoder_key("vocoder.vocoder.resblocks.0.acts1.1.downsample.lowpass.filter"), + "vocoder.resnets.0.acts1.1.downsample.filter" +) +expect_equal( + ltx23_map_vocoder_key("vocoder.vocoder.ups.2.weight"), + "vocoder.upsamplers.2.weight" +) +expect_equal( + ltx23_map_vocoder_key("vocoder.bwe_generator.act_post.act.alpha"), + "bwe_generator.act_out.act.alpha" +) +expect_equal( + ltx23_map_vocoder_key("vocoder.mel_stft.stft_fn.forward_basis"), + "mel_stft.stft_fn.forward_basis" +) +expect_equal( + ltx23_map_audio_vae_key("audio_vae.decoder.up.1.upsample.conv.conv.weight"), + "decoder.up.1.upsample.conv.conv.weight" +) +expect_equal( + ltx23_map_audio_vae_key("audio_vae.per_channel_statistics.mean-of-means"), + "latents_mean" +) +expect_equal( + ltx23_map_audio_vae_key("audio_vae.encoder.conv_in.conv.weight"), + "encoder.conv_in.conv.weight" +) +expect_equal( + ltx23_map_audio_vae_key("audio_vae.encoder.down.1.downsample.conv.weight"), + "encoder.down.1.downsample.conv.weight" +) + +# --- Encoder: shapes and key-name census ------------------------------------------- + +enc <- ltx23_audio_encoder( + base_channels = 8L, in_channels = 2L, num_res_blocks = 2L, + latent_channels = 4L, ch_mult = c(1L, 2L) +) +enc$eval() +mel_in <- torch::torch_randn(1L, 2L, 12L, 8L) +torch::with_no_grad(out_enc <- enc(mel_in)) +# One downsample level: time and mel halved; moments = 2 * latent +expect_equal(as.integer(out_enc$shape), c(1L, 8L, 6L, 4L)) + +# Parameter names line up with the checkpoint layout +pn <- names(enc$named_parameters()) +expect_true("conv_in.conv.weight" %in% pn) +expect_true("down.0.block.0.conv1.conv.weight" %in% pn) +expect_true("down.0.block.1.conv2.conv.bias" %in% pn) +expect_true("down.0.downsample.conv.weight" %in% pn) +expect_true("down.1.block.0.nin_shortcut.conv.weight" %in% pn) +expect_true("mid.block_1.conv1.conv.weight" %in% pn) +expect_true("mid.block_2.conv2.conv.bias" %in% pn) +expect_true("conv_out.conv.weight" %in% pn) +# Pixel norms are parameterless: exactly the conv tensors +expect_true(all(grepl("conv", pn))) + +# encode() through the vae wrapper returns mean/logvar of latent size +avae <- ltx23_audio_vae( + base_channels = 8L, output_channels = 2L, num_res_blocks = 1L, + latent_channels = 4L, ch_mult = c(1L, 2L), mel_bins = 8L +) +avae$eval() +torch::with_no_grad(m <- avae$encode(torch::torch_randn(1L, 2L, 12L, 8L))) +expect_equal(as.integer(m$mean$shape)[2], 4L) +expect_equal(as.integer(m$logvar$shape)[2], 4L) diff --git a/inst/tinytest/test_checkpoint_ltx23.R b/inst/tinytest/test_checkpoint_ltx23.R new file mode 100644 index 0000000..4267b97 --- /dev/null +++ b/inst/tinytest/test_checkpoint_ltx23.R @@ -0,0 +1,180 @@ +# Tests for the LTX-2.3 single-file checkpoint reader. +# Uses a tiny fake checkpoint written with official-style key names and +# model_version metadata; no model downloads required. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +make_fake_checkpoint <- function(path, version = "2.3.0") { + tensors <- list( + # DiT group + "model.diffusion_model.patchify_proj.weight" = + torch::torch_randn(8, 4), + "model.diffusion_model.patchify_proj.bias" = + torch::torch_randn(8), + "model.diffusion_model.scale_shift_table" = + torch::torch_randn(6, 8), + "model.diffusion_model.transformer_blocks.0.attn1.to_q.weight" = + torch::torch_randn(8, 8), + # Connector group (inside and outside the diffusion_model prefix) + "model.diffusion_model.video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight" = + torch::torch_randn(8, 8), + "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight" = + torch::torch_randn(8, 8), + "text_embedding_projection.video_aggregate_embed.weight" = + torch::torch_randn(8, 12), + "text_embedding_projection.audio_aggregate_embed.weight" = + torch::torch_randn(4, 12), + # VAE / audio VAE / vocoder groups + "vae.decoder.conv_in.conv.weight" = torch::torch_randn(4, 4, 3, 3, 3), + "audio_vae.decoder.conv_in.weight" = torch::torch_randn(4, 2, 3, 3), + "vocoder.vocoder.conv_in.weight" = torch::torch_randn(4, 4, 7), + "vocoder.bwe_generator.conv_in.weight" = torch::torch_randn(4, 4, 7), + "vocoder.mel_stft.mel_basis" = torch::torch_randn(64, 257) + ) + metadata <- list(model_version = version) + safetensors::safe_save_file(tensors, path, metadata = metadata) + invisible(tensors) +} + +tmp <- tempfile(fileext = ".safetensors") +on.exit(unlink(tmp), add = TRUE) +tensors <- make_fake_checkpoint(tmp) + +# --- Open + version gate ----------------------------------------------------- + +ckpt <- ltx23_open_checkpoint(tmp) +expect_inherits(ckpt, "ltx23_checkpoint") +expect_equal(ckpt$version, "2.3.0") +expect_equal(sort(ckpt$keys), sort(names(tensors))) + +# Wrong version is rejected +tmp_old <- tempfile(fileext = ".safetensors") +on.exit(unlink(tmp_old), add = TRUE) +make_fake_checkpoint(tmp_old, version = "2.0.1") +expect_error(ltx23_open_checkpoint(tmp_old), pattern = "2\\.3") + +# Version check can be disabled +ckpt_old <- ltx23_open_checkpoint(tmp_old, require_version = NULL) +expect_equal(ckpt_old$version, "2.0.1") + +# Missing file +expect_error(ltx23_open_checkpoint(tempfile()), pattern = "not found") + +# --- Key group split --------------------------------------------------------- + +groups <- ltx23_split_keys(ckpt$keys) +expect_equal(length(groups$dit), 4L) +expect_equal(length(groups$connectors), 4L) +expect_equal(length(groups$vae), 1L) +expect_equal(length(groups$audio_vae), 1L) +expect_equal(length(groups$vocoder), 3L) +expect_equal(length(groups$other), 0L) + +# Every key lands in exactly one group +expect_equal(sort(unlist(groups, use.names = FALSE)), sort(ckpt$keys)) + +# Census matches +cen <- ltx23_census(ckpt) +expect_equal(cen$keys[cen$group == "vocoder"], 3L) + +# --- Streaming group load ---------------------------------------------------- + +toy <- torch::nn_module( + "toy", + initialize = function() { + self$patchify_proj <- torch::nn_linear(4, 8) + self$scale_shift_table <- torch::nn_parameter(torch::torch_zeros(6, 8)) + self$blocks <- torch::nn_module_list(list( + torch::nn_module( + "toy_block", + initialize = function() { + self$to_q <- torch::nn_linear(8, 8, bias = FALSE) + } + )() + )) + } +)() + +map_dit_toy <- function(key) { + key <- sub("^model\\.diffusion_model\\.", "", key) + key <- sub("^transformer_blocks\\.0\\.attn1\\.", "blocks.0.", key) + key +} + +res <- ltx23_load_group(ckpt, groups$dit, toy, map_key = map_dit_toy, verbose = FALSE) +expect_equal(length(res$unmapped), 0L) +expect_equal(length(res$unfilled), 0L) +expect_equal(length(res$skipped), 0L) + +# Values actually copied +expect_equal( + as.numeric(torch::torch_sum(toy$patchify_proj$weight)), + as.numeric(torch::torch_sum(tensors[["model.diffusion_model.patchify_proj.weight"]])), + tolerance = 1e-5 +) + +# copy_ converts dtype: load into a float64 module +toy64 <- torch::nn_module( + "toy64", + initialize = function() { + self$w <- torch::nn_parameter(torch::torch_zeros(8, 4, dtype = torch::torch_float64())) + } +)() +res64 <- ltx23_load_group( + ckpt, "model.diffusion_model.patchify_proj.weight", toy64, + map_key = function(key) "w", verbose = FALSE +) +expect_equal(length(res64$unmapped), 0L) +expect_equal(toy64$w$dtype$.type(), "Double") + +# Unmapped keys are reported, not fatal +res_bad <- ltx23_load_group( + ckpt, groups$vae, toy, + map_key = identity, verbose = FALSE +) +expect_equal(length(res_bad$unmapped), 1L) + +# Mapper can deliberately skip keys with NA +res_skip <- ltx23_load_group( + ckpt, groups$vae, toy, + map_key = function(key) NA_character_, verbose = FALSE +) +expect_equal(length(res_skip$skipped), 1L) +expect_equal(length(res_skip$unmapped), 0L) + +# Shape mismatch is a hard error +expect_error( + ltx23_load_group( + ckpt, "model.diffusion_model.scale_shift_table", toy, + map_key = function(key) "patchify_proj.weight", verbose = FALSE + ), + pattern = "Shape mismatch" +) + +# --- Real checkpoint census (local only, needs the 46GB file) ---------------- + +if (at_home()) { + real <- tryCatch( + hfhub::hub_download( + "Lightricks/LTX-2.3", "ltx-2.3-22b-distilled-1.1.safetensors", + local_files_only = TRUE + ), + error = function(e) NULL + ) + if (!is.null(real)) { + rc <- ltx23_open_checkpoint(real) + expect_equal(rc$version, "2.3.0") + expect_equal(length(rc$keys), 5947L) + rg <- ltx23_split_keys(rc$keys) + expect_equal(length(rg$other), 0L) + expect_equal(length(rg$vocoder), 1227L) + expect_true(!is.null(rc$config$transformer$num_layers)) + } +} diff --git a/inst/tinytest/test_condition_ltx23.R b/inst/tinytest/test_condition_ltx23.R new file mode 100644 index 0000000..2cd8e15 --- /dev/null +++ b/inst/tinytest/test_condition_ltx23.R @@ -0,0 +1,103 @@ +# Unit tests for the prefix-conditioning helpers (R/condition_ltx23.R): +# preprocessing geometry, i2v repeat semantics, prefix token placement, +# and the packed conditioning mask. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) +torch::torch_manual_seed(21) + +# --- Preprocessing: resize + center-crop + range ------------------------------------ + +img <- array(runif(48 * 80 * 3), dim = c(48L, 80L, 3L)) +fr <- ltx23_preprocess_frames(img, height = 32L, width = 32L) +expect_equal(as.integer(fr$shape), c(1L, 3L, 1L, 32L, 32L)) +expect_true(as.numeric(fr$max()) <= 1) +expect_true(as.numeric(fr$min()) >= -1) + +# Multi-frame input keeps frame count +vid <- array(runif(5 * 40 * 40 * 3), dim = c(5L, 40L, 40L, 3L)) +frv <- ltx23_preprocess_frames(vid, height = 32L, width = 32L) +expect_equal(as.integer(frv$shape), c(1L, 3L, 5L, 32L, 32L)) + +# RGBA alpha is dropped +rgba <- array(runif(32 * 32 * 4), dim = c(32L, 32L, 4L)) +fra <- ltx23_preprocess_frames(rgba, height = 32L, width = 32L) +expect_equal(as.integer(fra$shape)[2], 3L) + +# A [0,1] gray image maps to ~0 in [-1,1] +gray <- array(0.5, dim = c(32L, 32L, 3L)) +frg <- ltx23_preprocess_frames(gray, height = 32L, width = 32L) +expect_true(as.numeric(frg$abs()$max()) < 1e-6) + +# --- Conditioned latent init + mask ------------------------------------------------- + +lf <- 4L; lh <- 2L; lw <- 3L +noise <- torch::torch_randn(1L, 128L, lf, lh, lw) + +# i2v: one conditioning latent frame, repeated then masked to frame 0 +cond1 <- torch::torch_randn(1L, 128L, 1L, lh, lw) +prep1 <- ltx23_prepare_conditioned_latents(cond1, lf, lh, lw, noise) +expect_equal(as.integer(prep1$latents$shape), c(1L, lf * lh * lw, 128L)) +expect_equal(as.integer(prep1$conditioning_mask$shape), c(1L, lf * lh * lw)) +n_tok <- lh * lw +m <- as.numeric(prep1$conditioning_mask[1, ]) +expect_true(all(m[1:n_tok] == 1)) +expect_true(all(m[(n_tok + 1):(lf * n_tok)] == 0)) +# Conditioned tokens equal the packed condition; the rest equal noise +packed_cond <- diffuseR:::ltx23_pack_video_latents(cond1) +expect_true(as.numeric((prep1$latents$narrow(2L, 1L, n_tok) - + packed_cond)$abs()$max()) < 1e-6) +packed_noise <- diffuseR:::ltx23_pack_video_latents(noise) +expect_true(as.numeric((prep1$latents$narrow(2L, n_tok + 1L, (lf - 1L) * n_tok) - + packed_noise$narrow(2L, n_tok + 1L, (lf - 1L) * n_tok))$abs()$max()) < 1e-6) + +# Continuation: two conditioning latent frames +cond2 <- torch::torch_randn(1L, 128L, 2L, lh, lw) +prep2 <- ltx23_prepare_conditioned_latents(cond2, lf, lh, lw, noise) +m2 <- as.numeric(prep2$conditioning_mask[1, ]) +expect_true(all(m2[1:(2L * n_tok)] == 1)) +expect_true(all(m2[(2L * n_tok + 1L):(lf * n_tok)] == 0)) +expect_true(as.numeric((prep2$latents$narrow(2L, 1L, 2L * n_tok) - + diffuseR:::ltx23_pack_video_latents(cond2))$abs()$max()) < 1e-6) + +# cond_noise_scale partially noises ONLY the conditioned tokens +prep3 <- ltx23_prepare_conditioned_latents(cond1, lf, lh, lw, noise, + cond_noise_scale = 0.5) +d_cond <- as.numeric((prep3$latents$narrow(2L, 1L, n_tok) - + packed_cond)$abs()$max()) +expect_true(d_cond > 1e-3) +d_free <- as.numeric((prep3$latents$narrow(2L, n_tok + 1L, (lf - 1L) * n_tok) - + packed_noise$narrow(2L, n_tok + 1L, (lf - 1L) * n_tok))$abs()$max()) +expect_true(d_free < 1e-6) + +# Expected blend value: noise*0.5 + cond*0.5 on conditioned tokens +expected <- packed_noise$narrow(2L, 1L, n_tok)$mul(0.5) + +packed_cond$mul(0.5) +expect_true(as.numeric((prep3$latents$narrow(2L, 1L, n_tok) - + expected)$abs()$max()) < 1e-6) + +# --- Encoding: argmax + normalization round trip ------------------------------------ + +vae <- ltx23_video_vae( + latent_channels = 4L, + block_out_channels = c(8L, 8L, 8L, 8L), + decoder_block_out_channels = c(4L, 8L, 8L, 16L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + decoder_layers_per_block = c(1L, 1L, 1L, 1L, 1L) +) +vae$eval() +torch::with_no_grad({ + vae$latents_mean$copy_(torch::torch_randn(4L)) + vae$latents_std$copy_(torch::torch_rand(4L) + 0.5) +}) +px <- torch::torch_rand(1L, 3L, 1L, 32L, 32L)$mul(2)$sub(1) +lat <- ltx23_encode_video_frames(vae, px) +expect_equal(as.integer(lat$shape), c(1L, 4L, 1L, 1L, 1L)) +# Matches manual encode + normalize +torch::with_no_grad(ref_moments <- vae$encode(px)) +ref <- ltx23_normalize_latents(ref_moments$mean, vae$latents_mean, + vae$latents_std) +expect_true(as.numeric((lat - ref)$abs()$max()) < 1e-5) diff --git a/inst/tinytest/test_connectors_ltx23.R b/inst/tinytest/test_connectors_ltx23.R new file mode 100644 index 0000000..955dc55 --- /dev/null +++ b/inst/tinytest/test_connectors_ltx23.R @@ -0,0 +1,86 @@ +# Parity tests for the LTX-2.3 text connectors against diffusers +# reference fixtures (tools/gen_fixtures_connectors.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "connectors_ltx23.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/connectors_ltx23.safetensors" +if (!file.exists(fixture_path)) exit_file("connector fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +conn <- ltx23_text_connectors( + caption_channels = 8L, + text_proj_in_factor = 3L, + video_connector_num_attention_heads = 2L, + video_connector_attention_head_dim = 8L, + video_connector_num_layers = 2L, + video_connector_num_learnable_registers = 4L, + video_gated_attn = TRUE, + audio_connector_num_attention_heads = 2L, + audio_connector_attention_head_dim = 4L, + audio_connector_num_layers = 2L, + audio_connector_num_learnable_registers = 4L, + audio_gated_attn = TRUE, + rope_type = "split", + video_hidden_dim = 16L, + audio_hidden_dim = 8L, + proj_bias = TRUE +) +conn$eval() + +weights <- fx[grep("^conn\\.", names(fx))] +names(weights) <- sub("^conn\\.", "", names(weights)) +dests <- c(conn$named_parameters(), conn$named_buffers()) +expect_equal(sort(names(weights)), sort(names(dests))) +torch::with_no_grad({ + for (name in names(weights)) dests[[name]]$copy_(weights[[name]]) +}) + +torch::with_no_grad({ + res <- conn(fx$c_states, fx$c_mask) +}) +expect_equal(as.integer(res$video_text_embedding$shape), as.integer(fx$c_video_emb$shape)) +expect_true(max_abs_diff(res$video_text_embedding, fx$c_video_emb) < 1e-4) +expect_true(max_abs_diff(res$audio_text_embedding, fx$c_audio_emb) < 1e-4) +expect_true(max_abs_diff( + res$attention_mask$to(dtype = torch::torch_float32()), fx$c_out_mask +) < 1e-6) + +# Flattened 3D input gives identical output +torch::with_no_grad({ + res3 <- conn(fx$c_states$flatten(start_dim = 3L), fx$c_mask) +}) +expect_true(max_abs_diff(res3$video_text_embedding, fx$c_video_emb3) < 1e-5) + +# Key mapper +expect_equal( + ltx23_map_connector_key("model.diffusion_model.video_embeddings_connector.transformer_1d_blocks.0.attn1.q_norm.weight"), + "video_connector.transformer_blocks.0.attn1.norm_q.weight" +) +expect_equal( + ltx23_map_connector_key("text_embedding_projection.video_aggregate_embed.weight"), + "video_text_proj_in.weight" +) +expect_equal( + ltx23_map_connector_key("text_embedding_projection.audio_aggregate_embed.bias"), + "audio_text_proj_in.bias" +) +expect_equal( + ltx23_map_connector_key("model.diffusion_model.audio_embeddings_connector.learnable_registers"), + "audio_connector.learnable_registers" +) diff --git a/inst/tinytest/test_dit_ltx2.R b/inst/tinytest/test_dit_ltx2.R deleted file mode 100644 index 0f9fd68..0000000 --- a/inst/tinytest/test_dit_ltx2.R +++ /dev/null @@ -1,256 +0,0 @@ -# Tests for LTX2 DiT Transformer modules - -# Test 1: Timesteps module -cat("Test 1: Timesteps module\n") -ts <- diffuseR:::timesteps_module(num_channels = 256L, flip_sin_to_cos = TRUE, downscale_freq_shift = 0) -timesteps <- torch::torch_tensor(c(0.5, 0.8)) -emb <- ts(timesteps) -expect_equal(as.numeric(emb$shape[2]), 256L, info = "Timestep embedding should have 256 channels") - -# Test 2: TimestepEmbedding MLP -cat("Test 2: TimestepEmbedding MLP\n") -te <- diffuseR:::timestep_embedding_module(in_channels = 256L, time_embed_dim = 512L) -x <- torch::torch_randn(c(2, 256)) -y <- te(x) -expect_equal(as.numeric(y$shape[2]), 512L, info = "Output should have 512 channels") - -# Test 3: PixArtAlphaCombinedTimestepSizeEmbeddings -cat("Test 3: PixArtAlphaCombinedTimestepSizeEmbeddings\n") -emb_module <- diffuseR:::pixart_alpha_combined_timestep_size_embeddings( - embedding_dim = 512L, - size_emb_dim = 170L, - use_additional_conditions = FALSE -) -timesteps <- torch::torch_tensor(c(500.0, 750.0)) -cond <- emb_module(timesteps, batch_size = 2L, hidden_dtype = torch::torch_float32()) -expect_true(!is.null(cond), info = "Should return conditioning tensor") -cat(sprintf(" Conditioning shape: [%s]\n", paste(as.numeric(cond$shape), collapse = ", "))) - -# Test 4: PixArtAlphaTextProjection -cat("Test 4: PixArtAlphaTextProjection\n") -text_proj <- diffuseR:::pixart_alpha_text_projection( - in_features = 3840L, - hidden_size = 4096L -) -caption <- torch::torch_randn(c(2, 77, 3840)) -proj <- text_proj(caption) -expect_equal(as.numeric(proj$shape[3]), 4096L, info = "Projected caption should have 4096 channels") - -# Test 5: RMSNorm -cat("Test 5: RMSNorm\n") -norm <- diffuseR:::rms_norm(dim = 512L, eps = 1e-6) -x <- torch::torch_randn(c(2, 10, 512)) -y <- norm(x) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "RMS norm should preserve shape") - -# Test 6: FeedForward -cat("Test 6: FeedForward\n") -ff <- diffuseR:::feed_forward(dim = 512L, mult = 4L, activation_fn = "gelu-approximate") -x <- torch::torch_randn(c(2, 10, 512)) -y <- ff(x) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "FFN should preserve shape") - -# Test 7: LTX2AdaLayerNormSingle -cat("Test 7: LTX2AdaLayerNormSingle\n") -ada_ln <- diffuseR:::ltx2_ada_layer_norm_single(embedding_dim = 512L, num_mod_params = 6L) -timesteps <- torch::torch_tensor(c(500.0, 750.0)) -result <- ada_ln(timesteps, batch_size = 2L, hidden_dtype = torch::torch_float32()) -expect_equal(length(result), 2L, info = "Should return mod_params and embedded_timestep") -cat(sprintf(" Mod params shape: [%s]\n", paste(as.numeric(result[[1]]$shape), collapse = ", "))) - -# Test 8: LTX2Attention -cat("Test 8: LTX2Attention\n") -attn <- diffuseR:::ltx2_attention( - query_dim = 512L, - heads = 8L, - kv_heads = 8L, - dim_head = 64L, - rope_type = "interleaved" -) -x <- torch::torch_randn(c(2, 100, 512)) -y <- attn(x) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "Self-attention should preserve shape") - -# Test 9: LTX2Attention with RoPE -cat("Test 9: LTX2Attention with RoPE\n") -# Create mock RoPE frequencies -cos_freqs <- torch::torch_ones(c(2, 100, 512)) -sin_freqs <- torch::torch_zeros(c(2, 100, 512)) -rope_freqs <- list(cos_freqs, sin_freqs) - -y_rope <- attn(x, query_rotary_emb = rope_freqs) -expect_equal(as.numeric(y_rope$shape), as.numeric(x$shape), info = "Attention with RoPE should preserve shape") - -# Test 10: LTX2Attention cross-attention -cat("Test 10: LTX2Attention cross-attention\n") -cross_attn <- diffuseR:::ltx2_attention( - query_dim = 512L, - heads = 8L, - kv_heads = 8L, - dim_head = 64L, - cross_attention_dim = 2048L -) -q <- torch::torch_randn(c(2, 100, 512)) -kv <- torch::torch_randn(c(2, 77, 2048)) -y_cross <- cross_attn(q, encoder_hidden_states = kv) -expect_equal(as.numeric(y_cross$shape), as.numeric(q$shape), info = "Cross-attention should output query shape") - -# Test 11: LTX2AudioVideoRotaryPosEmbed (video) -cat("Test 11: LTX2AudioVideoRotaryPosEmbed (video)\n") -rope <- diffuseR:::ltx2_audio_video_rotary_pos_embed( - dim = 4096L, - patch_size = 1L, - patch_size_t = 1L, - base_num_frames = 20L, - base_height = 2048L, - base_width = 2048L, - scale_factors = c(8L, 32L, 32L), - modality = "video", - rope_type = "interleaved" -) - -video_coords <- rope$prepare_video_coords( - batch_size = 2L, - num_frames = 4L, - height = 16L, - width = 16L, - device = "cpu", - fps = 24.0 -) -expect_equal(video_coords$shape[2], 3L, info = "Video coords should have 3 dims (T, H, W)") - -freqs <- rope(video_coords) -expect_equal(length(freqs), 2L, info = "Should return cos and sin freqs") -cat(sprintf(" Video RoPE cos shape: [%s]\n", paste(as.numeric(freqs[[1]]$shape), collapse = ", "))) - -# Test 12: LTX2AudioVideoRotaryPosEmbed (audio) -cat("Test 12: LTX2AudioVideoRotaryPosEmbed (audio)\n") -audio_rope <- diffuseR:::ltx2_audio_video_rotary_pos_embed( - dim = 2048L, - patch_size_t = 1L, - base_num_frames = 20L, - scale_factors = c(4L), - modality = "audio", - rope_type = "interleaved" -) - -audio_coords <- audio_rope$prepare_audio_coords( - batch_size = 2L, - num_frames = 100L, - device = "cpu" -) -expect_equal(audio_coords$shape[2], 1L, info = "Audio coords should have 1 dim (T)") - -audio_freqs <- audio_rope(audio_coords) -expect_equal(length(audio_freqs), 2L, info = "Should return cos and sin freqs") -cat(sprintf(" Audio RoPE cos shape: [%s]\n", paste(as.numeric(audio_freqs[[1]]$shape), collapse = ", "))) - -# Test 13: LTX2VideoTransformerBlock -cat("Test 13: LTX2VideoTransformerBlock\n") -block <- diffuseR:::ltx2_video_transformer_block( - dim = 512L, - num_attention_heads = 8L, - attention_head_dim = 64L, - cross_attention_dim = 2048L, - audio_dim = 256L, - audio_num_attention_heads = 4L, - audio_attention_head_dim = 64L, - audio_cross_attention_dim = 1024L -) - -# Prepare inputs -hidden_states <- torch::torch_randn(c(2, 100, 512)) -audio_hidden_states <- torch::torch_randn(c(2, 50, 256)) -encoder_hidden_states <- torch::torch_randn(c(2, 77, 2048)) -audio_encoder_hidden_states <- torch::torch_randn(c(2, 77, 1024)) -temb <- torch::torch_randn(c(2, 1, 3072)) # 6 * 512 -temb_audio <- torch::torch_randn(c(2, 1, 1536)) # 6 * 256 -temb_ca_scale_shift <- torch::torch_randn(c(2, 1, 2048)) # 4 * 512 -temb_ca_audio_scale_shift <- torch::torch_randn(c(2, 1, 1024)) # 4 * 256 -temb_ca_gate <- torch::torch_randn(c(2, 1, 512)) # 1 * 512 -temb_ca_audio_gate <- torch::torch_randn(c(2, 1, 256)) # 1 * 256 - -# Mock RoPE (ones for cos, zeros for sin = identity) -video_rope_freqs <- list(torch::torch_ones(c(2, 100, 512)), torch::torch_zeros(c(2, 100, 512))) -audio_rope_freqs <- list(torch::torch_ones(c(2, 50, 256)), torch::torch_zeros(c(2, 50, 256))) -ca_video_rope <- list(torch::torch_ones(c(2, 100, 256)), torch::torch_zeros(c(2, 100, 256))) -ca_audio_rope <- list(torch::torch_ones(c(2, 50, 256)), torch::torch_zeros(c(2, 50, 256))) - -result <- block( - hidden_states = hidden_states, - audio_hidden_states = audio_hidden_states, - encoder_hidden_states = encoder_hidden_states, - audio_encoder_hidden_states = audio_encoder_hidden_states, - temb = temb, - temb_audio = temb_audio, - temb_ca_scale_shift = temb_ca_scale_shift, - temb_ca_audio_scale_shift = temb_ca_audio_scale_shift, - temb_ca_gate = temb_ca_gate, - temb_ca_audio_gate = temb_ca_audio_gate, - video_rotary_emb = video_rope_freqs, - audio_rotary_emb = audio_rope_freqs, - ca_video_rotary_emb = ca_video_rope, - ca_audio_rotary_emb = ca_audio_rope -) - -expect_equal(length(result), 2L, info = "Block should return video and audio hidden states") -expect_equal(as.numeric(result[[1]]$shape), as.numeric(hidden_states$shape), info = "Video shape preserved") -expect_equal(as.numeric(result[[2]]$shape), as.numeric(audio_hidden_states$shape), info = "Audio shape preserved") -cat(" Block forward pass successful\n") - -# Test 14: Full LTX2VideoTransformer3DModel (small config) -cat("Test 14: LTX2VideoTransformer3DModel (small config)\n") -model <- ltx2_video_transformer_3d_model( - in_channels = 32L, - out_channels = 32L, - num_attention_heads = 4L, - attention_head_dim = 32L, - cross_attention_dim = 128L, # Must equal inner_dim = 4*32 = 128 (after caption projection) - audio_in_channels = 16L, - audio_out_channels = 16L, - audio_num_attention_heads = 2L, - audio_attention_head_dim = 32L, - audio_cross_attention_dim = 64L, # Must equal audio_inner_dim = 2*32 = 64 - num_layers = 2L, - caption_channels = 512L, - vae_scale_factors = c(8L, 32L, 32L) -) -expect_true(!is.null(model), info = "Model should instantiate") -cat(" Model instantiated with 2 layers\n") - -# Test 15: Full model forward pass -cat("Test 15: Full model forward pass\n") -torch::with_no_grad({ - batch_size <- 1L - num_frames <- 4L - height <- 8L - width <- 8L - num_patches <- num_frames * height * width - - hidden_states <- torch::torch_randn(c(batch_size, num_patches, 32L)) - audio_hidden_states <- torch::torch_randn(c(batch_size, 50L, 16L)) - encoder_hidden_states <- torch::torch_randn(c(batch_size, 77L, 512L)) - audio_encoder_hidden_states <- torch::torch_randn(c(batch_size, 77L, 512L)) # Must match caption_channels - timestep <- torch::torch_tensor(c(500.0))$unsqueeze(2) - - output <- model( - hidden_states = hidden_states, - audio_hidden_states = audio_hidden_states, - encoder_hidden_states = encoder_hidden_states, - audio_encoder_hidden_states = audio_encoder_hidden_states, - timestep = timestep, - num_frames = num_frames, - height = height, - width = width, - fps = 24.0, - audio_num_frames = 50L - ) -}) - -expect_equal(length(output), 2L, info = "Model should return sample and audio_sample") -expect_equal(as.numeric(output$sample$shape), c(1L, num_patches, 32L), info = "Video output shape correct") -expect_equal(as.numeric(output$audio_sample$shape), c(1L, 50L, 16L), info = "Audio output shape correct") -cat(sprintf(" Video output shape: [%s]\n", paste(as.numeric(output$sample$shape), collapse = ", "))) -cat(sprintf(" Audio output shape: [%s]\n", paste(as.numeric(output$audio_sample$shape), collapse = ", "))) - -cat("\nAll LTX2 DiT transformer tests completed\n") diff --git a/inst/tinytest/test_dit_ltx23.R b/inst/tinytest/test_dit_ltx23.R new file mode 100644 index 0000000..8a32678 --- /dev/null +++ b/inst/tinytest/test_dit_ltx23.R @@ -0,0 +1,204 @@ +# Parity tests for the LTX-2.3 DiT port against diffusers reference +# fixtures (generated by tools/gen_fixtures_dit.py, checked in). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "dit_ltx23.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/dit_ltx23.safetensors" +if (!file.exists(fixture_path)) exit_file("dit fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# Copy fixture weights (diffusers state-dict names) into an R module. +# Errors if any name has no destination or any destination stays unfilled. +load_named_weights <- function(module, weights) { + dests <- c(module$named_parameters(), module$named_buffers()) + missing_dest <- setdiff(names(weights), names(dests)) + if (length(missing_dest)) { + stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", ")) + } + unfilled <- setdiff(names(dests), names(weights)) + if (length(unfilled)) { + stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", ")) + } + torch::with_no_grad({ + for (name in names(weights)) dests[[name]]$copy_(weights[[name]]) + }) + invisible(module) +} + +fixture_group <- function(prefix) { + keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE) + w <- fx[keys] + names(w) <- sub(paste0("^", prefix, "\\."), "", keys) + w +} + +# --- Gated attention with split RoPE ------------------------------------------ + +attn <- ltx23_attention( + query_dim = 16L, heads = 2L, dim_head = 8L, + rope_type = "split", apply_gated_attention = TRUE +) +attn$eval() +load_named_weights(attn, fixture_group("attn")) + +rope <- list(fx$attn_cos, fx$attn_sin) +torch::with_no_grad({ + out <- attn(fx$attn_x, attention_mask = fx$attn_mask, query_rotary_emb = rope) + out_nomask <- attn(fx$attn_x, query_rotary_emb = rope) +}) +expect_true(max_abs_diff(out, fx$attn_out) < 1e-4) +expect_true(max_abs_diff(out_nomask, fx$attn_out_nomask) < 1e-4) + +# Query chunking is numerically identical +attn$attn_chunk <- 2L +torch::with_no_grad({ + out_chunked <- attn(fx$attn_x, attention_mask = fx$attn_mask, query_rotary_emb = rope) +}) +expect_true(max_abs_diff(out_chunked, fx$attn_out) < 1e-5) +attn$attn_chunk <- NULL + +# --- AdaLN single (9 mod params) ----------------------------------------------- + +ada <- ltx23_ada_layer_norm_single(16L, num_mod_params = 9L) +ada$eval() +load_named_weights(ada, fixture_group("ada")) +torch::with_no_grad({ + ada_res <- ada(fx$ada_t, hidden_dtype = torch::torch_float32()) +}) +expect_true(max_abs_diff(ada_res[[1]], fx$ada_out) < 1e-4) +expect_true(max_abs_diff(ada_res[[2]], fx$ada_emb) < 1e-4) + +# --- Tiny full model ------------------------------------------------------------- + +model <- ltx23_transformer( + in_channels = 4L, + out_channels = 4L, + num_attention_heads = 2L, + attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 4L, + audio_out_channels = 4L, + audio_num_attention_heads = 2L, + audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 2L, + rope_type = "split" +) +model$eval() +load_named_weights(model, fixture_group("model")) + +run_model <- function(...) { + torch::with_no_grad({ + model( + hidden_states = fx$m_hidden, + audio_hidden_states = fx$m_audio_hidden, + encoder_hidden_states = fx$m_enc, + audio_encoder_hidden_states = fx$m_audio_enc, + timestep = fx$m_timestep, + sigma = fx$m_sigma, + num_frames = 2L, height = 3L, width = 4L, fps = 24.0, + audio_num_frames = 5L, + use_cross_timestep = TRUE, + ... + ) + }) +} + +res <- run_model( + encoder_attention_mask = fx$m_enc_mask, + audio_encoder_attention_mask = fx$m_enc_mask +) +expect_equal(as.integer(res$sample$shape), as.integer(fx$m_out_video$shape)) +expect_true(max_abs_diff(res$sample, fx$m_out_video) < 1e-4) +expect_true(max_abs_diff(res$audio_sample, fx$m_out_audio) < 1e-4) + +# Isolated modalities (a2v/v2a off) +res_iso <- run_model(isolate_modalities = TRUE) +expect_true(max_abs_diff(res_iso$sample, fx$m_out_video_iso) < 1e-4) +expect_true(max_abs_diff(res_iso$audio_sample, fx$m_out_audio_iso) < 1e-4) + +# STG perturbation on block index 1 (0-based, i.e. the second block) +res_stg <- run_model(spatio_temporal_guidance_blocks = c(1L)) +expect_true(max_abs_diff(res_stg$sample, fx$m_out_video_stg) < 1e-4) +expect_true(max_abs_diff(res_stg$audio_sample, fx$m_out_audio_stg) < 1e-4) + +# --- Key mapper ------------------------------------------------------------------ + +expect_equal( + ltx23_map_dit_key("model.diffusion_model.patchify_proj.weight"), + "proj_in.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.audio_patchify_proj.bias"), + "audio_proj_in.bias" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.weight"), + "time_embed.emb.timestep_embedder.linear_1.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.audio_adaln_single.linear.bias"), + "audio_time_embed.linear.bias" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.prompt_adaln_single.linear.weight"), + "prompt_adaln.linear.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.audio_prompt_adaln_single.linear.weight"), + "audio_prompt_adaln.linear.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.av_ca_a2v_gate_adaln_single.linear.weight"), + "av_cross_attn_video_a2v_gate.linear.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.transformer_blocks.0.attn1.q_norm.weight"), + "transformer_blocks.0.attn1.norm_q.weight" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.transformer_blocks.5.scale_shift_table_a2v_ca_video"), + "transformer_blocks.5.video_a2v_cross_attn_scale_shift_table" +) +expect_equal( + ltx23_map_dit_key("model.diffusion_model.transformer_blocks.0.ff.net.0.proj.weight"), + "transformer_blocks.0.ff.net.0.proj.weight" +) +expect_equal(ltx23_map_dit_key("model.diffusion_model.scale_shift_table"), "scale_shift_table") + +# --- In-place tanh GELU (large-activation feed-forward path) ------------------------ + +ff <- ltx23_feed_forward(8L) +x_ff <- torch::torch_randn(1L, 37L, 8L) +torch::with_no_grad({ + ref_ff <- ff(x_ff) + # Force the in-place chunked path (threshold below this tensor's numel, + # chunk smaller than the intermediate so several chunks run) + op <- options(diffuseR.gelu_inplace_min = 1) + inp_ff <- ff(x_ff) + options(op) +}) +expect_true(as.numeric((ref_ff - inp_ff)$abs()$max()) < 1e-6) + +h_g <- torch::torch_randn(3L, 41L) +g_ref <- torch::nnf_gelu(h_g, approximate = "tanh") +torch::with_no_grad( + g_inp <- diffuseR:::.ltx23_gelu_tanh_inplace(h_g$clone(), chunk_elements = 16) +) +expect_true(as.numeric((g_ref - g_inp)$abs()$max()) < 1e-6) diff --git a/inst/tinytest/test_fp8_ltx23.R b/inst/tinytest/test_fp8_ltx23.R new file mode 100644 index 0000000..fb86506 --- /dev/null +++ b/inst/tinytest/test_fp8_ltx23.R @@ -0,0 +1,166 @@ +# FP8 quantization round trip: quantize a tiny official-named checkpoint, +# reload with fp8 linears, and compare against the bf16 original. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} +# fp8 needs the F8-capable safetensors build +f8_ok <- tryCatch({ + x <- torch::torch_randn(2, 2)$to(dtype = torch::torch_float8_e4m3fn()) + tmp <- tempfile(fileext = ".safetensors") + safetensors::safe_save_file(list(w = x), tmp) + y <- safetensors::safe_load_file(tmp, framework = "torch") + unlink(tmp) + TRUE +}, error = function(e) FALSE) +if (!f8_ok) exit_file("safetensors build lacks F8 support") + +library(diffuseR) + +torch::torch_manual_seed(123) + +tiny_cfg <- list( + in_channels = 4L, out_channels = 4L, + num_attention_heads = 2L, attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 4L, audio_out_channels = 4L, + audio_num_attention_heads = 2L, audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 2L +) +ref_model <- do.call(ltx23_transformer, tiny_cfg) +ref_model$eval() + +# Reverse the diffusers renames to synthesize an official-named checkpoint +to_official <- function(name) { + name <- gsub("av_cross_attn_video_scale_shift", "av_ca_video_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_video_a2v_gate", "av_ca_a2v_gate_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_scale_shift", "av_ca_audio_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_v2a_gate", "av_ca_v2a_gate_adaln_single", name, fixed = TRUE) + name <- gsub("video_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_video", name, fixed = TRUE) + name <- gsub("audio_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_audio", name, fixed = TRUE) + name <- gsub("prompt_adaln", "prompt_adaln_single", name, fixed = TRUE) + name <- sub("^audio_time_embed\\.", "audio_adaln_single.", name) + name <- sub("^time_embed\\.", "adaln_single.", name) + name <- gsub("proj_in", "patchify_proj", name, fixed = TRUE) + name <- gsub("norm_q", "q_norm", name, fixed = TRUE) + name <- gsub("norm_k", "k_norm", name, fixed = TRUE) + paste0("model.diffusion_model.", name) +} + +params <- ref_model$named_parameters() +tensors <- list() +for (name in names(params)) { + tensors[[to_official(name)]] <- params[[name]]$detach() +} +# Round trip through the mapper must be lossless +expect_equal( + sort(vapply(names(tensors), ltx23_map_dit_key, character(1), USE.NAMES = FALSE)), + sort(names(params)) +) + +src <- tempfile(fileext = ".safetensors") +on.exit(unlink(src), add = TRUE) +safetensors::safe_save_file(tensors, src, metadata = list(model_version = "2.3.0")) + +# --- Quantize ----------------------------------------------------------------- + +fp8_dir <- tempfile("fp8_") +on.exit(unlink(fp8_dir, recursive = TRUE), add = TRUE) +manifest <- ltx23_quantize_fp8(src, fp8_dir, verbose = FALSE) +expect_true(file.exists(file.path(fp8_dir, "manifest.json"))) +expect_true(manifest$fp8_cast > 0) +# 2 layers x (3 video self + 3 audio self + 3 x-attn x 2 + 4 a2v/v2a) qkv/out +# + 4 ff projections: just assert the exact official policy count +n_expected <- sum(ltx23_is_fp8_cast_key(names(params))) +expect_equal(manifest$fp8_cast, n_expected) + +# Re-running is a no-op (skip-if-exists) +m2 <- ltx23_quantize_fp8(src, fp8_dir, verbose = FALSE) +expect_equal(m2$shards, manifest$shards) + +# --- Load fp8 + forward parity --------------------------------------------------- + +ckpt <- ltx23_open_fp8_checkpoint(fp8_dir) +expect_equal(ckpt$version, "2.3.0") + +old_opt <- getOption("diffuseR.use_fp8") +fp8_model <- do.call(ltx23_load_transformer_fp8, c( + list(ckpt = ckpt, device = "cpu", pin = FALSE, verbose = FALSE), + tiny_cfg +)) +options(diffuseR.use_fp8 = old_opt) + +# Cast linears were swapped +expect_inherits(fp8_model$transformer_blocks[[1]]$attn1$to_q, "ltx23_fp8_linear") +expect_equal( + fp8_model$transformer_blocks[[1]]$attn1$to_q$weight_fp8$dtype$.type(), + "Float8_e4m3fn" +) +# Non-cast params match exactly (bf16 storage of an fp32 original) +expect_true(as.numeric(torch::torch_max(torch::torch_abs( + fp8_model$proj_in$weight$to(dtype = torch::torch_float32()) - + ref_model$proj_in$weight +))) < 0.01) + +# Dequantized weights reconstruct the originals within fp8 e4m3 tolerance +# (a random tiny network amplifies chaotically, so end-to-end output error +# is not a meaningful metric; real-weight quality shows in the E2E render) +fp8_q <- fp8_model$transformer_blocks[[1]]$attn1$to_q +deq <- fp8_q$weight_fp8$to(dtype = torch::torch_float32()) * + fp8_q$weight_scale +orig <- ref_model$transformer_blocks[[1]]$attn1$to_q$weight +rel_w <- as.numeric((deq - orig)$abs()$mean() / orig$abs()$mean()) +expect_true(rel_w < 0.05) + +# The fp8_linear layer itself matches a plain linear closely +x_probe <- torch::torch_randn(3L, 16L) +torch::with_no_grad({ + y_fp8 <- fp8_q(x_probe$to(dtype = torch::torch_bfloat16())) + y_ref <- torch::nnf_linear(x_probe, orig, + ref_model$transformer_blocks[[1]]$attn1$to_q$bias) +}) +rel_y <- as.numeric( + (y_fp8$to(dtype = torch::torch_float32()) - y_ref)$abs()$mean() / + y_ref$abs()$mean() +) +expect_true(rel_y < 0.1) + +# Full fp8 model forward runs in bf16 and stays finite +run <- function(model, dtype) { + torch::torch_manual_seed(5) + hidden <- torch::torch_randn(1L, 24L, 4L, dtype = dtype) + audio <- torch::torch_randn(1L, 5L, 4L, dtype = dtype) + enc <- torch::torch_randn(1L, 7L, 16L, dtype = dtype) + aenc <- torch::torch_randn(1L, 7L, 8L, dtype = dtype) + t <- torch::torch_tensor(700, dtype = torch::torch_float32()) + torch::with_no_grad({ + model( + hidden_states = hidden, audio_hidden_states = audio, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + timestep = t, sigma = t, + num_frames = 2L, height = 3L, width = 4L, audio_num_frames = 5L, + use_cross_timestep = TRUE + ) + }) +} +fp8_out <- run(fp8_model, torch::torch_bfloat16()) +expect_equal(as.integer(fp8_out$sample$shape), c(1L, 24L, 4L)) +expect_true(as.logical(fp8_out$sample$isfinite()$all()$item())) +expect_true(as.logical(fp8_out$audio_sample$isfinite()$all()$item())) + +# --- attn chunking + memory profile ------------------------------------------------- + +ltx23_set_attn_chunk(fp8_model, 8L) +expect_equal(fp8_model$transformer_blocks[[1]]$attn2$attn_chunk, 8L) +ltx23_set_attn_chunk(fp8_model, NULL) +expect_null(fp8_model$transformer_blocks[[1]]$attn2$attn_chunk) + +prof <- ltx23_memory_profile(vram_gb = 16) +expect_equal(prof$name, "high") +expect_true(is.integer(prof$attn_chunk) || is.null(prof$attn_chunk)) +expect_equal(prof$precision, "nf4") +expect_equal(ltx23_memory_profile(vram_gb = 4)$name, "cpu_only") diff --git a/inst/tinytest/test_gemma3_states.R b/inst/tinytest/test_gemma3_states.R new file mode 100644 index 0000000..1fec51f --- /dev/null +++ b/inst/tinytest/test_gemma3_states.R @@ -0,0 +1,53 @@ +# Hidden-state collection semantics of the Gemma3 text encoder. +# +# HF transformers returns num_layers + 1 hidden states: the state BEFORE +# each decoder layer (the first being the embedding output) plus the +# post-final-norm output. The un-normed last-layer output is never +# included. The LTX connectors' input projection expects exactly +# hidden_size * (num_layers + 1) features, so an extra or missing state +# fails at the first matmul with real weights. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +config <- gemma3_config_ltx2() +config$vocab_size <- 64L +config$hidden_size <- 16L +config$intermediate_size <- 32L +config$num_hidden_layers <- 3L +config$num_attention_heads <- 2L +config$num_key_value_heads <- 1L +config$head_dim <- 8L +config$query_pre_attn_scalar <- 8L +config$sliding_window <- 4L + +model <- gemma3_text_model(config) +model$eval() + +ids <- torch::torch_randint(1L, 64L, size = c(1L, 7L), dtype = torch::torch_long()) +mask <- torch::torch_ones(1L, 7L) +torch::with_no_grad({ + out <- model(ids, attention_mask = mask, output_hidden_states = TRUE) +}) + +# num_layers + 1 states +expect_equal(length(out$hidden_states), config$num_hidden_layers + 1L) + +# Last collected state is the post-norm output (identical to last_hidden_state) +last_collected <- out$hidden_states[[length(out$hidden_states)]] +expect_true(as.numeric((last_collected - out$last_hidden_state)$abs()$max()) == 0) + +# The un-normed last-layer output must NOT be in the stack: the +# second-to-last entry is the input to the final layer, so running the +# final layer + norm over it must reproduce last_hidden_state. +pre_last <- out$hidden_states[[config$num_hidden_layers]] +expect_false(as.numeric((pre_last - out$last_hidden_state)$abs()$max()) == 0) + +# Without the flag, no states are collected +torch::with_no_grad({ + out2 <- model(ids, attention_mask = mask, output_hidden_states = FALSE) +}) +expect_null(out2$hidden_states) diff --git a/inst/tinytest/test_gpu_poor.R b/inst/tinytest/test_gpu_poor.R deleted file mode 100644 index a80030d..0000000 --- a/inst/tinytest/test_gpu_poor.R +++ /dev/null @@ -1,82 +0,0 @@ -# Tests for GPU-Poor Memory Management - -# Test 1: Memory profile detection -cat("Test 1: ltx2_memory_profile with explicit VRAM\n") -profile_high <- ltx2_memory_profile(vram_gb = 20) -expect_equal(profile_high$name, "high", info = "20GB should be high profile") - -profile_medium <- ltx2_memory_profile(vram_gb = 12) -expect_equal(profile_medium$name, "medium", info = "12GB should be medium profile") - -profile_low <- ltx2_memory_profile(vram_gb = 8) -expect_equal(profile_low$name, "low", info = "8GB should be low profile") - -profile_very_low <- ltx2_memory_profile(vram_gb = 6) -expect_equal(profile_very_low$name, "very_low", info = "6GB should be very_low profile") - -profile_cpu <- ltx2_memory_profile(vram_gb = 2) -expect_equal(profile_cpu$name, "cpu_only", info = "2GB should be cpu_only profile") - -# Test 2: Profile contains required fields -cat("Test 2: Profile structure\n") -profile <- ltx2_memory_profile(vram_gb = 8) -expect_true("dit_device" %in% names(profile), info = "Profile should have dit_device") -expect_true("vae_device" %in% names(profile), info = "Profile should have vae_device") -expect_true("vae_tiling" %in% names(profile), info = "Profile should have vae_tiling") -expect_true("vae_tile_size" %in% names(profile), info = "Profile should have vae_tile_size") -expect_true("max_resolution" %in% names(profile), info = "Profile should have max_resolution") -expect_true("max_frames" %in% names(profile), info = "Profile should have max_frames") -expect_true("cfg_mode" %in% names(profile), info = "Profile should have cfg_mode") - -# Test 3: Low profile enables tiling and offloading -cat("Test 3: Low profile settings\n") -low <- ltx2_memory_profile(vram_gb = 8) -expect_true(low$dit_offload, info = "Low profile should enable DiT offload") -expect_true(low$vae_tiling, info = "Low profile should enable VAE tiling") -expect_equal(low$cfg_mode, "sequential", info = "Low profile should use sequential CFG") - -# Test 4: High profile disables aggressive optimizations -cat("Test 4: High profile settings\n") -high <- ltx2_memory_profile(vram_gb = 20) -expect_false(high$dit_offload, info = "High profile should not offload DiT") -expect_false(high$vae_tiling, info = "High profile should not tile VAE") -expect_equal(high$cfg_mode, "batched", info = "High profile should use batched CFG") - -# Test 5: Resolution validation -cat("Test 5: validate_resolution\n") -profile <- ltx2_memory_profile(vram_gb = 8) -result <- validate_resolution(1080, 1920, 121, profile) -expect_true(result$adjusted, info = "Should adjust 1080p for low profile") -expect_true(result$height <= profile$max_resolution[1], info = "Height should be capped") -expect_true(result$width <= profile$max_resolution[2], info = "Width should be capped") -expect_true(result$num_frames <= profile$max_frames, info = "Frames should be capped") - -# Test 6: Resolution within limits -cat("Test 6: validate_resolution within limits\n") -result_ok <- validate_resolution(480, 640, 30, profile) -expect_false(result_ok$adjusted, info = "Should not adjust if within limits") - -# Test 7: VRAM report (just test it runs) -cat("Test 7: vram_report\n") -result <- vram_report("test") -expect_true(is.list(result), info = "vram_report should return a list") -expect_true("used" %in% names(result), info = "Result should have used") -expect_true("free" %in% names(result), info = "Result should have free") - -# Test 8: clear_vram (just test it runs) -cat("Test 8: clear_vram\n") -clear_vram() -expect_true(TRUE, info = "clear_vram should run without error") - -# Test 9: is_blackwell_gpu returns logical -cat("Test 9: is_blackwell_gpu\n") -is_bb <- is_blackwell_gpu() -expect_true(is.logical(is_bb), info = "is_blackwell_gpu should return logical") - -# Test 10: .detect_vram returns numeric -cat("Test 10: .detect_vram\n") -vram <- diffuseR:::.detect_vram() -expect_true(is.numeric(vram), info = ".detect_vram should return numeric") -expect_true(vram >= 0, info = "VRAM should be non-negative") - -cat("\nAll GPU-poor tests completed\n") diff --git a/inst/tinytest/test_jit_ltx23.R b/inst/tinytest/test_jit_ltx23.R new file mode 100644 index 0000000..e68f56c --- /dev/null +++ b/inst/tinytest/test_jit_ltx23.R @@ -0,0 +1,214 @@ +# Parity of the TorchScript block stack (R/jit_ltx23.R) against the +# eager NF4 transformer block. The compiled path must reproduce the +# eager path bit-for-bit up to SDPA-vs-materialized-attention rounding. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) +torch::torch_manual_seed(7) + +# The prompt KV modulation tables are dim-sized and apply directly to +# the text streams, so cross_attention_dim == dim (as in the real +# model: 4096/4096 video, 2048/2048 audio) +dim <- 32L; heads <- 2L; head_dim <- 16L; cross_dim <- 32L +adim <- 16L; aheads <- 2L; ahead_dim <- 8L; across_dim <- 16L +B <- 2L; Sv <- 6L; Sa <- 5L; St <- 4L + +make_block <- function() { + blk <- ltx23_transformer_block( + dim = dim, num_attention_heads = heads, attention_head_dim = head_dim, + cross_attention_dim = cross_dim, + audio_dim = adim, audio_num_attention_heads = aheads, + audio_attention_head_dim = ahead_dim, + audio_cross_attention_dim = across_dim + ) + blk$eval() + + # Swap every cast linear for an NF4 module exactly like the loader does + swap_lin <- function(parent, leaf) { + old <- diffuseR:::.ltx23_walk_module(parent, leaf) + q <- ltx23_nf4_quantize(old$weight) + m <- ltx23_nf4_linear(old$weight$shape[1], old$weight$shape[2], + bias = !is.null(old$bias)) + if (!is.null(old$bias)) m$bias <- old$bias + m$set_nf4_weight(q$packed, q$absmax) + do.call(`$<-`, list(parent, leaf, m)) + } + for (attn in list(blk$attn1, blk$audio_attn1, blk$attn2, blk$audio_attn2, + blk$audio_to_video_attn, blk$video_to_audio_attn)) { + swap_lin(attn, "to_q") + swap_lin(attn, "to_k") + swap_lin(attn, "to_v") + swap_lin(attn$to_out, "0") + } + for (ff in list(blk$ff, blk$audio_ff)) { + swap_lin(ff$net[[1]], "proj") + swap_lin(ff$net, "2") + } + blk +} + +blk1 <- make_block() +blk2 <- make_block() +expect_true(diffuseR:::.ltx23_jit_block_ok(blk1)) + +# The model's split-rope embedder emits per-head 4D freqs [B, H, T, r] +rope_pair <- function(s, n_heads, r) { + ang <- torch::torch_rand(B, n_heads, s, r) * 6.28 + list(torch::torch_cos(ang), torch::torch_sin(ang)) +} + +h <- torch::torch_randn(B, Sv, dim) +ah <- torch::torch_randn(B, Sa, adim) +enc <- torch::torch_randn(B, St, cross_dim) +aenc <- torch::torch_randn(B, St, across_dim) +temb <- torch::torch_randn(B, 1L, 9L * dim) +temb_a <- torch::torch_randn(B, 1L, 9L * adim) +tcss <- torch::torch_randn(B, 1L, 4L * dim) +tcass <- torch::torch_randn(B, 1L, 4L * adim) +tcg <- torch::torch_randn(B, 1L, dim) +tcag <- torch::torch_randn(B, 1L, adim) +tp <- torch::torch_randn(B, 1L, 2L * dim) +tpa <- torch::torch_randn(B, 1L, 2L * adim) +v_rope <- rope_pair(Sv, heads, head_dim %/% 2L) +a_rope <- rope_pair(Sa, aheads, ahead_dim %/% 2L) +cav_rope <- rope_pair(Sv, aheads, ahead_dim %/% 2L) +caa_rope <- rope_pair(Sa, aheads, ahead_dim %/% 2L) +# Additive [B, 1, S] mask with one text token masked out +enc_mask <- torch::torch_zeros(B, 1L, St) +enc_mask[, , St] <- -10000 +aenc_mask <- torch::torch_zeros(B, 1L, St) + +eager <- function(blk, h, ah) { + blk( + hidden_states = h, audio_hidden_states = ah, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + temb = temb, 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, + encoder_attention_mask = enc_mask, + audio_encoder_attention_mask = aenc_mask + ) +} + +jit <- function(blocks, h, ah) { + diffuseR:::.ltx23_jit_run_stack( + blocks, h, ah, enc, aenc, + temb, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v_rope, a_rope, cav_rope, caa_rope, + encoder_attention_mask = enc_mask, + audio_encoder_attention_mask = aenc_mask + ) +} + +max_abs_diff <- function(a, b) as.numeric((a - b)$abs()$max()) + +torch::with_no_grad({ + ref1 <- eager(blk1, h, ah) + out1 <- jit(list(blk1), h, ah) +}) +expect_equal(as.integer(out1[[1]]$shape), as.integer(ref1[[1]]$shape)) +expect_equal(as.integer(out1[[2]]$shape), as.integer(ref1[[2]]$shape)) +expect_true(max_abs_diff(out1[[1]], ref1[[1]]) < 1e-4) +expect_true(max_abs_diff(out1[[2]], ref1[[2]]) < 1e-4) + +# Two-block stack must equal two sequential eager blocks (checks the +# per-block base-offset arithmetic in the weight list) +torch::with_no_grad({ + mid <- eager(blk1, h, ah) + ref2 <- eager(blk2, mid[[1]], mid[[2]]) + out2 <- jit(list(blk1, blk2), h, ah) +}) +expect_true(max_abs_diff(out2[[1]], ref2[[1]]) < 1e-4) +expect_true(max_abs_diff(out2[[2]], ref2[[2]]) < 1e-4) + +# The mask must actually mask: moving the -10000 changes the output +enc_mask2 <- torch::torch_zeros(B, 1L, St) +enc_mask2[, , 1L] <- -10000 +torch::with_no_grad({ + out_m <- diffuseR:::.ltx23_jit_run_stack( + list(blk1), h, ah, enc, aenc, + temb, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v_rope, a_rope, cav_rope, caa_rope, + encoder_attention_mask = enc_mask2, + audio_encoder_attention_mask = aenc_mask + ) +}) +expect_true(max_abs_diff(out_m[[1]], out1[[1]]) > 1e-6) + +# NULL masks accepted +torch::with_no_grad({ + ref_nm <- blk1( + hidden_states = h, audio_hidden_states = ah, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + temb = temb, 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_nm <- diffuseR:::.ltx23_jit_run_stack( + list(blk1), h, ah, enc, aenc, + temb, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v_rope, a_rope, cav_rope, caa_rope + ) +}) +expect_true(max_abs_diff(out_nm[[1]], ref_nm[[1]]) < 1e-4) +expect_true(max_abs_diff(out_nm[[2]], ref_nm[[2]]) < 1e-4) + +# 3D whole-vector rope layout (the apply fn's other branch) also matches +rope3 <- function(s, r) { + ang <- torch::torch_rand(1L, s, r) * 6.28 + list(torch::torch_cos(ang), torch::torch_sin(ang)) +} +v3 <- rope3(Sv, (heads * head_dim) %/% 2L) +a3 <- rope3(Sa, (aheads * ahead_dim) %/% 2L) +cav3 <- rope3(Sv, (aheads * ahead_dim) %/% 2L) +caa3 <- rope3(Sa, (aheads * ahead_dim) %/% 2L) +torch::with_no_grad({ + ref3 <- blk1( + hidden_states = h, audio_hidden_states = ah, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + temb = temb, 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 = v3, audio_rotary_emb = a3, + ca_video_rotary_emb = cav3, ca_audio_rotary_emb = caa3 + ) + out3 <- diffuseR:::.ltx23_jit_run_stack( + list(blk1), h, ah, enc, aenc, + temb, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v3, a3, cav3, caa3 + ) +}) +expect_true(max_abs_diff(out3[[1]], ref3[[1]]) < 1e-4) +expect_true(max_abs_diff(out3[[2]], ref3[[2]]) < 1e-4) + +# Per-token temb (i2v conditioning: per-token video timestep) matches +temb_tok <- torch::torch_randn(B, Sv, 9L * dim) +torch::with_no_grad({ + ref_pt <- blk1( + hidden_states = h, audio_hidden_states = ah, + encoder_hidden_states = enc, audio_encoder_hidden_states = aenc, + temb = temb_tok, 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_pt <- diffuseR:::.ltx23_jit_run_stack( + list(blk1), h, ah, enc, aenc, + temb_tok, temb_a, tcss, tcass, tcg, tcag, tp, tpa, + v_rope, a_rope, cav_rope, caa_rope + ) +}) +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) diff --git a/inst/tinytest/test_jit_vae_ltx23.R b/inst/tinytest/test_jit_vae_ltx23.R new file mode 100644 index 0000000..e0b37ca --- /dev/null +++ b/inst/tinytest/test_jit_vae_ltx23.R @@ -0,0 +1,124 @@ +# Traced-decode parity (R/jit_vae_ltx23.R): the per-shape jit_trace +# cache must reproduce the eager video decoder, audio decoder, and +# vocoder exactly, re-trace on new shapes, and honor the off switch. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) +torch::torch_manual_seed(11) + +old_opt <- options(diffuseR.jit_vae = TRUE) +on.exit(options(old_opt), add = TRUE) + +max_abs_diff <- function(a, b) as.numeric((a - b)$abs()$max()) +trace_count <- function() length(ls(diffuseR:::.ltx23_vae_traces)) + +diffuseR:::.ltx23_release_vae_traces() + +# --- Video decoder through the vae wrapper (tiled + direct paths) ------------------- + +vae <- ltx23_video_vae( + latent_channels = 4L, + block_out_channels = c(8L, 8L, 8L, 8L), + decoder_block_out_channels = c(16L, 32L, 32L, 64L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + decoder_layers_per_block = c(1L, 1L, 1L, 1L, 1L) +) +vae$eval() + +z <- torch::torch_randn(1L, 4L, 3L, 4L, 6L) +torch::with_no_grad({ + op <- options(diffuseR.jit_vae = FALSE) + ref <- vae$decode(z) + options(op) + op <- options(diffuseR.jit_vae = TRUE) + out <- vae$decode(z) + options(op) +}) +expect_equal(as.integer(out$shape), as.integer(ref$shape)) +expect_true(max_abs_diff(out, ref) < 1e-6) +expect_equal(trace_count(), 1L) + +# New shape -> new trace, still exact +z2 <- torch::torch_randn(1L, 4L, 3L, 6L, 4L) +torch::with_no_grad({ + op <- options(diffuseR.jit_vae = FALSE) + ref2 <- vae$decode(z2) + options(op) + op <- options(diffuseR.jit_vae = TRUE) + out2 <- vae$decode(z2) + options(op) +}) +expect_true(max_abs_diff(out2, ref2) < 1e-6) +expect_equal(trace_count(), 2L) + +# Same shape again reuses the cache (no growth) +torch::with_no_grad({ + op <- options(diffuseR.jit_vae = TRUE) + invisible(vae$decode(z)) + options(op) +}) +expect_equal(trace_count(), 2L) + +# Tiled decode: traced tiles must equal eager tiles exactly +vae$enable_tiling(spatial = TRUE, temporal = TRUE) +vae$tile_sample_min_height <- 64L +vae$tile_sample_min_width <- 64L +vae$tile_sample_stride_height <- 32L +vae$tile_sample_stride_width <- 32L +vae$tile_sample_min_num_frames <- 16L +vae$tile_sample_stride_num_frames <- 8L +z_big <- torch::torch_randn(1L, 4L, 5L, 4L, 6L) +torch::with_no_grad({ + op <- options(diffuseR.jit_vae = FALSE) + ref_t <- vae$decode(z_big) + options(op) + op <- options(diffuseR.jit_vae = TRUE) + out_t <- vae$decode(z_big) + options(op) +}) +expect_true(max_abs_diff(out_t, ref_t) < 1e-6) + +# Release drops every cached trace +diffuseR:::.ltx23_release_vae_traces() +expect_equal(trace_count(), 0L) + +# --- Audio decoder --------------------------------------------------------------------- + +adec <- ltx23_audio_decoder( + base_channels = 8L, output_channels = 2L, num_res_blocks = 1L, + latent_channels = 4L, ch_mult = c(1L, 2L), mel_bins = 8L +) +adec$eval() +za <- torch::torch_randn(1L, 4L, 5L, 2L) +torch::with_no_grad({ + ref_a <- adec(za) + out_a <- diffuseR:::.ltx23_traced_call(adec, za) +}) +expect_true(max_abs_diff(out_a, ref_a) < 1e-6) + +# --- Vocoder with BWE ------------------------------------------------------------------ + +bwe <- ltx23_vocoder_with_bwe( + in_channels = 8L, hidden_channels = 16L, out_channels = 2L, + upsample_kernel_sizes = c(4L, 4L), upsample_factors = c(2L, 2L), + resnet_kernel_sizes = c(3L), resnet_dilations = list(c(1L, 3L)), + bwe_in_channels = 16L, bwe_hidden_channels = 8L, + bwe_upsample_kernel_sizes = c(8L, 4L), bwe_upsample_factors = c(4L, 2L), + bwe_resnet_kernel_sizes = c(3L), bwe_resnet_dilations = list(c(1L, 3L)), + filter_length = 8L, hop_length = 2L, window_length = 8L, + num_mel_channels = 8L, + input_sampling_rate = 100L, output_sampling_rate = 400L +) +bwe$eval() +# Vocoder input is [B, C, T, M] with C * M == in_channels +mel <- torch::torch_randn(1L, 2L, 6L, 4L) +torch::with_no_grad({ + ref_v <- bwe(mel) + out_v <- diffuseR:::.ltx23_traced_call(bwe, mel) +}) +expect_true(max_abs_diff(out_v, ref_v) < 1e-6) + +diffuseR:::.ltx23_release_vae_traces() diff --git a/inst/tinytest/test_ltx2_weights.R b/inst/tinytest/test_ltx2_weights.R deleted file mode 100644 index d63e579..0000000 --- a/inst/tinytest/test_ltx2_weights.R +++ /dev/null @@ -1,101 +0,0 @@ -# Tests for LTX2 weight loading functions - -# Test 1: load_ltx2_vae function exists -cat("Test 1: load_ltx2_vae function exists\n") -expect_true(is.function(load_ltx2_vae), info = "load_ltx2_vae should be a function") - -# Test 2: load_ltx2_transformer function exists -cat("Test 2: load_ltx2_transformer function exists\n") -expect_true(is.function(load_ltx2_transformer), info = "load_ltx2_transformer should be a function") - -# Test 3: load_ltx2_connectors function exists -cat("Test 3: load_ltx2_connectors function exists\n") -expect_true(is.function(load_ltx2_connectors), info = "load_ltx2_connectors should be a function") - -# Test 4: VAE module can be created with defaults -cat("Test 4: VAE module creation\n") -vae <- ltx2_video_vae() -expect_true(inherits(vae, "nn_module"), info = "VAE should be nn_module") - -# Test 5: VAE has expected structure -cat("Test 5: VAE structure\n") -expect_true("encoder" %in% names(vae$children), info = "VAE should have encoder") -expect_true("decoder" %in% names(vae$children), info = "VAE should have decoder") - -# Test 6: VAE parameter count is reasonable -cat("Test 6: VAE parameter count\n") -vae_params <- names(vae$parameters) -expect_true(length(vae_params) > 100, - info = sprintf("VAE should have many parameters (got %d)", length(vae_params))) - -# Test 7: Transformer module can be created -cat("Test 7: Transformer module creation\n") -# Create with smaller config for testing (less memory) -transformer <- ltx2_video_transformer_3d_model( - num_layers = 2L, # Minimal layers for testing - num_attention_heads = 4L, - attention_head_dim = 32L, - audio_num_attention_heads = 4L, - audio_attention_head_dim = 32L -) -expect_true(inherits(transformer, "nn_module"), info = "Transformer should be nn_module") - -# Test 8: Transformer has expected structure -cat("Test 8: Transformer structure\n") -expect_true("transformer_blocks" %in% names(transformer$children), - info = "Transformer should have transformer_blocks") -expect_true("proj_in" %in% names(transformer$children), - info = "Transformer should have proj_in") -expect_true("proj_out" %in% names(transformer$children), - info = "Transformer should have proj_out") - -# Test 9: Connectors module can be created -cat("Test 9: Connectors module creation\n") -connectors <- ltx2_text_connectors( - video_connector_num_attention_heads = 4L, - video_connector_attention_head_dim = 32L, - audio_connector_num_attention_heads = 4L, - audio_connector_attention_head_dim = 32L -) -expect_true(inherits(connectors, "nn_module"), info = "Connectors should be nn_module") - -# Test 10: Connectors has expected structure -cat("Test 10: Connectors structure\n") -expect_true("video_connector" %in% names(connectors$children), - info = "Connectors should have video_connector") -expect_true("audio_connector" %in% names(connectors$children), - info = "Connectors should have audio_connector") -expect_true("text_proj_in" %in% names(connectors$children), - info = "Connectors should have text_proj_in") - -# Test 11: load_ltx2_vae errors on missing file -cat("Test 11: load_ltx2_vae error handling\n") -expect_error(load_ltx2_vae("/nonexistent/path.safetensors"), - info = "Should error on missing file") - -# Test 12: load_ltx2_transformer errors on missing directory -cat("Test 12: load_ltx2_transformer error handling\n") -expect_error(load_ltx2_transformer("/nonexistent/dir"), - info = "Should error on missing directory") - -# Test 13: load_ltx2_connectors errors on missing file -cat("Test 13: load_ltx2_connectors error handling\n") -expect_error(load_ltx2_connectors("/nonexistent/path.safetensors"), - info = "Should error on missing file") - -# Test 14: Internal weight loading function exists for VAE -cat("Test 14: Internal VAE weight loading function\n") -expect_true(exists("load_ltx2_vae_weights", envir = asNamespace("diffuseR")), - info = "load_ltx2_vae_weights should exist") - -# Test 15: Internal weight loading function exists for transformer -cat("Test 15: Internal transformer weight loading function\n") -expect_true(exists("load_ltx2_transformer_weights", envir = asNamespace("diffuseR")), - info = "load_ltx2_transformer_weights should exist") - -# Test 16: Internal weight loading function exists for connectors -cat("Test 16: Internal connector weight loading function\n") -expect_true(exists("load_ltx2_connector_weights", envir = asNamespace("diffuseR")), - info = "load_ltx2_connector_weights should exist") - -cat("\nLTX2 weight loading tests completed\n") diff --git a/inst/tinytest/test_nf4_ltx23.R b/inst/tinytest/test_nf4_ltx23.R new file mode 100644 index 0000000..c018212 --- /dev/null +++ b/inst/tinytest/test_nf4_ltx23.R @@ -0,0 +1,134 @@ +# NF4 quantization: round trip quality, linear parity, and resident +# transformer loading from a tiny official-named checkpoint. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +torch::torch_manual_seed(321) + +# --- Quantize/dequantize round trip ------------------------------------------ + +w <- torch::torch_randn(64L, 128L) * 0.02 +q <- ltx23_nf4_quantize(w) +expect_equal(as.integer(q$packed$shape), 64L * 128L / 2L) +expect_equal(as.integer(q$absmax$shape), 64L * 128L / 64L) +expect_equal(q$packed$dtype$.type(), "Byte") + +w_rt <- ltx23_nf4_dequantize(q$packed, q$absmax, c(64L, 128L), + dtype = torch::torch_float32()) +expect_equal(as.integer(w_rt$shape), c(64L, 128L)) +# NF4 round-trip error on Gaussian data: a few percent relative +rel <- as.numeric((w_rt - w)$abs()$mean() / w$abs()$mean()) +expect_true(rel < 0.1) + +# Values land exactly on table levels x absmax +blocks <- w_rt$reshape(c(-1L, 64L)) / q$absmax$unsqueeze(2L) +dists <- torch::torch_min( + (blocks$flatten()$unsqueeze(2L) - + torch::torch_tensor(diffuseR:::.ltx23_nf4_table)$unsqueeze(1L))$abs(), + dim = 2L +)[[1]] +expect_true(as.numeric(dists$max()) < 1e-5) + +# Chunked dequant is identical to single-shot +w_chunked <- ltx23_nf4_dequantize(q$packed, q$absmax, c(64L, 128L), + dtype = torch::torch_float32(), chunk_elements = 512L) +expect_true(as.numeric((w_chunked - w_rt)$abs()$max()) == 0) + +# --- nf4_linear parity ----------------------------------------------------------- + +lin <- torch::nn_linear(128L, 64L) +qw <- ltx23_nf4_quantize(lin$weight) +nf4 <- ltx23_nf4_linear(64L, 128L, bias = TRUE) +nf4$set_nf4_weight(qw$packed, qw$absmax) +torch::with_no_grad(nf4$bias$copy_(lin$bias)) + +x <- torch::torch_randn(5L, 128L) +torch::with_no_grad({ + y_ref <- lin(x) + y_nf4 <- nf4(x) +}) +rel_y <- as.numeric((y_nf4 - y_ref)$abs()$mean() / y_ref$abs()$mean()) +expect_true(rel_y < 0.15) + +# Buffers move with the module and survive dtype casts +nf4$to(dtype = torch::torch_float64()) +expect_equal(nf4$weight_nf4$dtype$.type(), "Byte") + +# --- Tiny transformer through the NF4 artifact ------------------------------------ + +tiny_cfg <- list( + in_channels = 4L, out_channels = 4L, + num_attention_heads = 2L, attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 4L, audio_out_channels = 4L, + audio_num_attention_heads = 2L, audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 1L +) +ref_model <- do.call(ltx23_transformer, tiny_cfg) +ref_model$eval() + +to_official <- function(name) { + name <- gsub("av_cross_attn_video_scale_shift", "av_ca_video_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_video_a2v_gate", "av_ca_a2v_gate_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_scale_shift", "av_ca_audio_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_v2a_gate", "av_ca_v2a_gate_adaln_single", name, fixed = TRUE) + name <- gsub("video_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_video", name, fixed = TRUE) + name <- gsub("audio_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_audio", name, fixed = TRUE) + name <- gsub("prompt_adaln", "prompt_adaln_single", name, fixed = TRUE) + name <- sub("^audio_time_embed\\.", "audio_adaln_single.", name) + name <- sub("^time_embed\\.", "adaln_single.", name) + name <- gsub("proj_in", "patchify_proj", name, fixed = TRUE) + name <- gsub("norm_q", "q_norm", name, fixed = TRUE) + name <- gsub("norm_k", "k_norm", name, fixed = TRUE) + paste0("model.diffusion_model.", name) +} + +params <- ref_model$named_parameters() +tensors <- list() +for (name in names(params)) { + tensors[[to_official(name)]] <- params[[name]]$detach() +} +src <- tempfile(fileext = ".safetensors") +on.exit(unlink(src), add = TRUE) +safetensors::safe_save_file(tensors, src, metadata = list(model_version = "2.3.0")) + +nf4_dir <- tempfile("nf4_") +on.exit(unlink(nf4_dir, recursive = TRUE), add = TRUE) +manifest <- ltx23_quantize_nf4(src, nf4_dir, verbose = FALSE) +expect_equal(manifest$format, "nf4") +expect_equal(manifest$nf4_cast, sum(ltx23_is_fp8_cast_key(names(params)))) + +ckpt <- ltx23_open_fp8_checkpoint(nf4_dir) +expect_equal(ckpt$format, "nf4") + +nf4_model <- do.call(ltx23_load_transformer_nf4, c( + list(ckpt = ckpt, device = "cpu", verbose = FALSE), + tiny_cfg +)) +expect_inherits(nf4_model$transformer_blocks[[1]]$attn1$to_q, "ltx23_nf4_linear") + +# Forward runs in bf16 and stays finite +torch::torch_manual_seed(5) +bf <- torch::torch_bfloat16() +torch::with_no_grad({ + out <- nf4_model( + hidden_states = torch::torch_randn(1L, 24L, 4L, dtype = bf), + audio_hidden_states = torch::torch_randn(1L, 5L, 4L, dtype = bf), + encoder_hidden_states = torch::torch_randn(1L, 7L, 16L, dtype = bf), + audio_encoder_hidden_states = torch::torch_randn(1L, 7L, 8L, dtype = bf), + timestep = torch::torch_tensor(700, dtype = torch::torch_float32()), + sigma = torch::torch_tensor(700, dtype = torch::torch_float32()), + num_frames = 2L, height = 3L, width = 4L, audio_num_frames = 5L, + use_cross_timestep = TRUE + ) +}) +expect_true(as.logical(out$sample$isfinite()$all()$item())) +expect_true(as.logical(out$audio_sample$isfinite()$all()$item())) diff --git a/inst/tinytest/test_rope.R b/inst/tinytest/test_rope.R deleted file mode 100644 index 76b0b57..0000000 --- a/inst/tinytest/test_rope.R +++ /dev/null @@ -1,140 +0,0 @@ -# Tests for RoPE (Rotary Position Embeddings) - -# Test embedder creation -expect_silent({ - embedder <- rope_embedder_create( - dim = 2048, - patch_size = 1, - patch_size_t = 1 - ) -}) - -expect_equal(embedder$dim, 2048) -expect_equal(embedder$rope_type, "interleaved") -expect_equal(embedder$theta, 10000.0) - -# Test video coordinate preparation -embedder <- rope_embedder_create( - dim = 2048, - patch_size = 1, - patch_size_t = 1, - scale_factors = c(8, 32, 32) -) - -coords <- rope_prepare_video_coords( - embedder = embedder, - batch_size = 2, - num_frames = 4, - height = 16, - width = 16, - device = "cpu", - fps = 24.0 -) - -expect_true(inherits(coords, "torch_tensor")) - -# Check shape: (batch_size, 3, num_patches, 2) -# num_patches = num_frames * height * width = 4 * 16 * 16 = 1024 -expect_equal(as.numeric(coords$shape[1]), 2) # batch_size -expect_equal(as.numeric(coords$shape[2]), 3) # 3 dimensions (f, h, w) -expect_equal(as.numeric(coords$shape[3]), 4 * 16 * 16) # num_patches -expect_equal(as.numeric(coords$shape[4]), 2) # start, end - -# Test RoPE forward (frequency computation) -freqs <- rope_forward(embedder, coords, device = "cpu") - -expect_true(is.list(freqs)) -expect_true(inherits(freqs$cos_freqs, "torch_tensor")) -expect_true(inherits(freqs$sin_freqs, "torch_tensor")) - -# cos and sin should have same shape -expect_equal(as.numeric(freqs$cos_freqs$shape), as.numeric(freqs$sin_freqs$shape)) - -# cos values should be in [-1, 1] -expect_true(as.numeric(freqs$cos_freqs$min()) >= -1.0) -expect_true(as.numeric(freqs$cos_freqs$max()) <= 1.0) - -# sin values should be in [-1, 1] -expect_true(as.numeric(freqs$sin_freqs$min()) >= -1.0) -expect_true(as.numeric(freqs$sin_freqs$max()) <= 1.0) - -# Test apply_interleaved_rotary_emb -batch_size <- 2 -seq_len <- 1024 -channels <- 2048 - -x <- torch::torch_randn(c(batch_size, seq_len, channels)) - -# Need to ensure freqs match x's sequence length -embedder_small <- rope_embedder_create(dim = channels) -coords_small <- rope_prepare_video_coords( - embedder = embedder_small, - batch_size = batch_size, - num_frames = 4, - height = 16, - width = 16, - device = "cpu" -) -freqs_small <- rope_forward(embedder_small, coords_small, device = "cpu") - -# Apply rotation -rotated <- apply_interleaved_rotary_emb(x, freqs_small) - -expect_true(inherits(rotated, "torch_tensor")) -expect_equal(as.numeric(rotated$shape), c(batch_size, seq_len, channels)) - -# Rotated should differ from original -expect_false(torch::torch_allclose(x, rotated)) - -# Test that applying rotation twice with negated freqs returns original -# (RoPE is a rotation, so -rotation undoes it) -freqs_neg <- list( - cos_freqs = freqs_small$cos_freqs, - sin_freqs = -freqs_small$sin_freqs -) -double_rotated <- apply_interleaved_rotary_emb(rotated, freqs_neg) - -# Should be close to original (numerical precision) -diff <- (x - double_rotated)$abs()$max()$item() -expect_true(diff < 1e-5, info = sprintf("Double rotation diff: %f", diff)) - -# Test with different patch sizes -embedder_patched <- rope_embedder_create( - dim = 2048, - patch_size = 2, - patch_size_t = 2 -) - -coords_patched <- rope_prepare_video_coords( - embedder = embedder_patched, - batch_size = 1, - num_frames = 8, - height = 32, - width = 32, - device = "cpu" -) - -# num_patches = (8/2) * (32/2) * (32/2) = 4 * 16 * 16 = 1024 -expect_equal(as.numeric(coords_patched$shape[3]), 4 * 16 * 16) - -# Test GPU if available -if (torch::cuda_is_available()) { - coords_gpu <- rope_prepare_video_coords( - embedder = embedder, - batch_size = 1, - num_frames = 4, - height = 16, - width = 16, - device = "cuda" - ) - - freqs_gpu <- rope_forward(embedder, coords_gpu, device = "cuda") - - expect_equal(as.character(freqs_gpu$cos_freqs$device), "cuda:0") - - x_gpu <- torch::torch_randn(c(1, 1024, 2048), device = "cuda") - rotated_gpu <- apply_interleaved_rotary_emb(x_gpu, freqs_gpu) - expect_equal(as.character(rotated_gpu$device), "cuda:0") -} - -cat("All RoPE tests passed\n") diff --git a/inst/tinytest/test_rope_ltx23.R b/inst/tinytest/test_rope_ltx23.R new file mode 100644 index 0000000..15b76a5 --- /dev/null +++ b/inst/tinytest/test_rope_ltx23.R @@ -0,0 +1,89 @@ +# Parity tests for the LTX-2.3 RoPE port against diffusers reference +# fixtures (generated by tools/gen_fixtures_rope.py, checked in). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "rope_ltx23.safetensors", + package = "diffuseR") +if (fixture_path == "") { + # Running from the source tree (e.g. during development) + fixture_path <- "fixtures/rope_ltx23.safetensors" +} +if (!file.exists(fixture_path)) exit_file("rope fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# --- apply_split_rotary_emb --------------------------------------------------- + +out4 <- ltx23_apply_split_rotary_emb(fx$split4_x, list(fx$split4_cos, fx$split4_sin)) +expect_equal(as.integer(out4$shape), as.integer(fx$split4_out$shape)) +expect_true(max_abs_diff(out4, fx$split4_out) < 1e-5) + +# 3D input with 4D freqs reshapes per-head and back +out3 <- ltx23_apply_split_rotary_emb(fx$split3_x, list(fx$split4_cos, fx$split4_sin)) +expect_equal(as.integer(out3$shape), as.integer(fx$split3_out$shape)) +expect_true(max_abs_diff(out3, fx$split3_out) < 1e-5) + +# --- apply_interleaved_rotary_emb ---------------------------------------------- + +outi <- ltx23_apply_interleaved_rotary_emb(fx$inter_x, list(fx$inter_cos, fx$inter_sin)) +expect_true(max_abs_diff(outi, fx$inter_out) < 1e-5) + +# --- video coords + split freqs ------------------------------------------------- + +rope_v <- ltx23_rotary_pos_embed( + dim = 64L, scale_factors = c(8L, 32L, 32L), modality = "video", + rope_type = "split", num_attention_heads = 4L +) +vc <- rope_v$prepare_video_coords( + batch_size = 2L, num_frames = 3L, height = 4L, width = 6L, + device = "cpu", fps = 24.0 +) +expect_equal(as.integer(vc$shape), as.integer(fx$video_coords$shape)) +expect_true(max_abs_diff(vc, fx$video_coords) < 1e-6) + +vf <- rope_v(vc) +expect_true(max_abs_diff(vf[[1]], fx$video_cos) < 1e-6) +expect_true(max_abs_diff(vf[[2]], fx$video_sin) < 1e-6) + +# fp16 input stays fp16 and matches reference values +outv16 <- ltx23_apply_split_rotary_emb(fx$video_x_f16, vf) +expect_equal(outv16$dtype$.type(), "Half") +expect_true(max_abs_diff(outv16, fx$video_out_f16) < 1e-3) + +# --- audio coords + split freqs ------------------------------------------------- + +rope_a <- ltx23_rotary_pos_embed( + dim = 32L, scale_factors = c(8L, 32L, 32L), modality = "audio", + rope_type = "split", num_attention_heads = 4L +) +ac <- rope_a$prepare_audio_coords(batch_size = 2L, num_frames = 5L, device = "cpu") +expect_equal(as.integer(ac$shape), as.integer(fx$audio_coords$shape)) +expect_true(max_abs_diff(ac, fx$audio_coords) < 1e-6) + +af <- rope_a(ac) +expect_true(max_abs_diff(af[[1]], fx$audio_cos) < 1e-6) +expect_true(max_abs_diff(af[[2]], fx$audio_sin) < 1e-6) + +# --- interleaved freqs ----------------------------------------------------------- + +rope_i <- ltx23_rotary_pos_embed( + dim = 64L, modality = "video", rope_type = "interleaved", + num_attention_heads = 4L +) +ifr <- rope_i(vc) +expect_true(max_abs_diff(ifr[[1]], fx$video_inter_cos) < 1e-6) +expect_true(max_abs_diff(ifr[[2]], fx$video_inter_sin) < 1e-6) diff --git a/inst/tinytest/test_rope_python_validation.R b/inst/tinytest/test_rope_python_validation.R deleted file mode 100644 index b57b2c9..0000000 --- a/inst/tinytest/test_rope_python_validation.R +++ /dev/null @@ -1,132 +0,0 @@ -# Validation tests for RoPE against Python diffusers -# These tests ensure numerical equivalence with HuggingFace implementation - -# Load Python test cases -test_cases_path <- system.file("validation/rope_test_cases.json", package = "diffuseR") - -if (!file.exists(test_cases_path)) { - cat("Skipping Python validation tests - test cases file not found\n") -} else { - test_cases <- jsonlite::fromJSON(test_cases_path) - tol <- 1e-4 - - # Test 1: Video coordinate shape and bounds - cat("Test 1: Video coordinate shape and bounds\n") - embedder <- rope_embedder_create( - dim = 2048, - patch_size = 1, - patch_size_t = 1, - scale_factors = c(8, 32, 32) - ) - - coords <- rope_prepare_video_coords( - embedder = embedder, - batch_size = 2, - num_frames = 4, - height = 16, - width = 16, - device = "cpu", - fps = 24.0 - ) - - python_coords <- test_cases$video_coords - - # Check shape - r_shape <- as.numeric(coords$shape) - py_shape <- python_coords$coords_shape - expect_equal(r_shape, py_shape, info = "Coordinate shapes should match") - - # Check bounds - r_min <- as.numeric(coords$min()) - r_max <- as.numeric(coords$max()) - expect_true(abs(r_min - python_coords$coords_min) < tol, - info = sprintf("Min diff: %f", abs(r_min - python_coords$coords_min))) - expect_true(abs(r_max - python_coords$coords_max) < tol, - info = sprintf("Max diff: %f", abs(r_max - python_coords$coords_max))) - - cat(sprintf(" Shape: R=%s, Python=%s\n", - paste(r_shape, collapse=","), paste(py_shape, collapse=","))) - cat(sprintf(" Min: R=%.4f, Python=%.4f\n", r_min, python_coords$coords_min)) - cat(sprintf(" Max: R=%.4f, Python=%.4f\n", r_max, python_coords$coords_max)) - - # Test 2: RoPE frequency shapes and bounds - cat("Test 2: RoPE frequency computation\n") - freqs <- rope_forward(embedder, coords, device = "cpu") - - python_freqs <- test_cases$rope_freqs - - # Check shapes - cos_shape <- as.numeric(freqs$cos_freqs$shape) - sin_shape <- as.numeric(freqs$sin_freqs$shape) - expect_equal(cos_shape, python_freqs$cos_shape, info = "Cos shape mismatch") - expect_equal(sin_shape, python_freqs$sin_shape, info = "Sin shape mismatch") - - # Check bounds - cos_min <- as.numeric(freqs$cos_freqs$min()) - cos_max <- as.numeric(freqs$cos_freqs$max()) - sin_min <- as.numeric(freqs$sin_freqs$min()) - sin_max <- as.numeric(freqs$sin_freqs$max()) - - expect_true(abs(cos_min - python_freqs$cos_min) < tol) - expect_true(abs(cos_max - python_freqs$cos_max) < tol) - expect_true(abs(sin_min - python_freqs$sin_min) < tol) - expect_true(abs(sin_max - python_freqs$sin_max) < tol) - - cat(sprintf(" Cos shape: %s\n", paste(cos_shape, collapse=","))) - cat(sprintf(" Cos range: [%.4f, %.4f]\n", cos_min, cos_max)) - cat(sprintf(" Sin range: [%.4f, %.4f]\n", sin_min, sin_max)) - - # Check means (important for numerical equivalence) - cos_mean <- as.numeric(freqs$cos_freqs$mean()) - sin_mean <- as.numeric(freqs$sin_freqs$mean()) - - cat(sprintf(" Cos mean: R=%.6f, Python=%.6f, diff=%.2e\n", - cos_mean, python_freqs$cos_mean, abs(cos_mean - python_freqs$cos_mean))) - cat(sprintf(" Sin mean: R=%.6f, Python=%.6f, diff=%.2e\n", - sin_mean, python_freqs$sin_mean, abs(sin_mean - python_freqs$sin_mean))) - - # Test 3: Apply rotary embedding - cat("Test 3: Apply interleaved rotary embedding\n") - torch::torch_manual_seed(42) - x <- torch::torch_randn(c(2, 1024, 2048)) - - rotated <- apply_interleaved_rotary_emb(x, freqs) - - python_apply <- test_cases$apply_rope - r_mean <- as.numeric(rotated$mean()) - r_std <- as.numeric(rotated$std()) - - cat(sprintf(" Output mean: R=%.6f, Python=%.6f\n", r_mean, python_apply$output_mean)) - cat(sprintf(" Output std: R=%.6f, Python=%.6f\n", r_std, python_apply$output_std)) - - # Shape should match exactly - expect_equal(as.numeric(rotated$shape), python_apply$output_shape) - - # Test 4: Patched coordinates - cat("Test 4: Patched coordinates\n") - embedder_patched <- rope_embedder_create( - dim = 2048, - patch_size = 2, - patch_size_t = 2 - ) - - coords_patched <- rope_prepare_video_coords( - embedder = embedder_patched, - batch_size = 1, - num_frames = 8, - height = 32, - width = 32, - device = "cpu" - ) - - python_patched <- test_cases$patched_coords - r_num_patches <- as.numeric(coords_patched$shape[3]) - - expect_equal(r_num_patches, python_patched$expected_num_patches, - info = sprintf("Num patches: R=%d, expected=%d", - r_num_patches, python_patched$expected_num_patches)) - cat(sprintf(" Num patches: %d (expected: %d)\n", - r_num_patches, python_patched$expected_num_patches)) - - cat("\nAll RoPE Python validation tests completed\n") -} diff --git a/inst/tinytest/test_staging_ltx23.R b/inst/tinytest/test_staging_ltx23.R new file mode 100644 index 0000000..7345855 --- /dev/null +++ b/inst/tinytest/test_staging_ltx23.R @@ -0,0 +1,45 @@ +# Pinned staging round trip (R/staging_ltx23.R): pin -> onload -> +# offload -> onload must preserve outputs exactly. CUDA-only. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!torch::cuda_is_available()) exit_file("no CUDA") + +library(diffuseR) +torch::torch_manual_seed(13) + +m <- ltx23_feed_forward(16L) +m$eval() +x <- torch::torch_randn(2L, 5L, 16L) +torch::with_no_grad(ref <- m(x)) + +st <- diffuseR:::.ltx23_pin_component(m) +expect_false(is.null(st)) +expect_true(suppressWarnings( + st[[1]]$live$is_pinned(device = torch::torch_device("cuda")) +)) + +# Pinning must not change CPU outputs +torch::with_no_grad(out_pinned <- m(x)) +expect_true(as.numeric((out_pinned - ref)$abs()$max()) == 0) + +# Onload: GPU forward matches +diffuseR:::.ltx23_staged_onload(st, "cuda") +expect_equal(st[[1]]$live$device$type, "cuda") +torch::with_no_grad( + out_gpu <- m(x$to(device = "cuda"))$cpu() +) +expect_true(as.numeric((out_gpu - ref)$abs()$max()) < 1e-5) + +# Offload: pointer swap back to the pinned copies, exact outputs +diffuseR:::.ltx23_staged_offload(st) +expect_equal(st[[1]]$live$device$type, "cpu") +torch::with_no_grad(out_back <- m(x)) +expect_true(as.numeric((out_back - ref)$abs()$max()) == 0) + +# Second round trip still exact +diffuseR:::.ltx23_staged_onload(st, "cuda") +torch::with_no_grad(out_gpu2 <- m(x$to(device = "cuda"))$cpu()) +expect_true(as.numeric((out_gpu2 - out_gpu)$abs()$max()) == 0) +diffuseR:::.ltx23_staged_offload(st) diff --git a/inst/tinytest/test_text_encoder_ltx2.R b/inst/tinytest/test_text_encoder_ltx2.R deleted file mode 100644 index 0eeac12..0000000 --- a/inst/tinytest/test_text_encoder_ltx2.R +++ /dev/null @@ -1,155 +0,0 @@ -# Tests for LTX2 Text Encoder and Connectors - -# Test 1: 1D RoPE for connectors -cat("Test 1: LTX2RotaryPosEmbed1d\n") -rope <- diffuseR:::ltx2_rotary_pos_embed_1d( - dim = 512L, - base_seq_len = 4096L, - theta = 10000.0, - rope_type = "interleaved" -) -freqs <- rope(batch_size = 2L, seq_len = 128L, device = "cpu") -expect_equal(length(freqs), 2L, info = "Should return cos and sin") -expect_equal(as.numeric(freqs[[1]]$shape), c(2, 128, 512), info = "Cos shape correct") -expect_equal(as.numeric(freqs[[2]]$shape), c(2, 128, 512), info = "Sin shape correct") - -# Test 2: 1D Transformer Block -cat("Test 2: LTX2TransformerBlock1d\n") -block <- diffuseR:::ltx2_transformer_block_1d( - dim = 512L, - num_attention_heads = 8L, - attention_head_dim = 64L -) -x <- torch::torch_randn(c(2, 100, 512)) -y <- block(x) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "Block preserves shape") - -# Test 3: 1D Transformer Block with RoPE -cat("Test 3: LTX2TransformerBlock1d with RoPE\n") -rope_freqs <- rope(batch_size = 2L, seq_len = 100L, device = "cpu") -y_rope <- block(x, rotary_emb = rope_freqs) -expect_equal(as.numeric(y_rope$shape), as.numeric(x$shape), info = "Block with RoPE preserves shape") - -# Test 4: Connector Transformer 1D (without learnable registers) -cat("Test 4: LTX2ConnectorTransformer1d (no registers)\n") -connector <- diffuseR:::ltx2_connector_transformer_1d( - num_attention_heads = 8L, - attention_head_dim = 64L, - num_layers = 2L, - num_learnable_registers = NULL -) -x <- torch::torch_randn(c(2, 128, 512)) -attn_mask <- torch::torch_zeros(c(2, 128)) -result <- connector(x, attn_mask) -expect_equal(length(result), 2L, info = "Returns hidden_states and attention_mask") -expect_equal(as.numeric(result[[1]]$shape), c(2, 128, 512), info = "Output shape correct") - -# Test 5: Connector Transformer 1D with attention mask (no registers) -cat("Test 5: LTX2ConnectorTransformer1d with mask (no registers)\n") -connector_mask <- diffuseR:::ltx2_connector_transformer_1d( - num_attention_heads = 8L, - attention_head_dim = 64L, - num_layers = 2L, - num_learnable_registers = NULL # No registers - simpler for testing -) -x <- torch::torch_randn(c(2, 128, 512)) -# Additive attention mask (0 = valid, negative large = masked) -attn_mask <- torch::torch_zeros(c(2, 128)) -attn_mask[1, 65:128] <- -10000.0 # Mask second half of first batch -attn_mask[2, 100:128] <- -10000.0 # Mask last 28 of second batch -result <- connector_mask(x, attn_mask) -expect_equal(as.numeric(result[[1]]$shape), c(2, 128, 512), info = "Output shape with mask correct") - -# Test 6: Full Text Connectors (small config) -cat("Test 6: LTX2TextConnectors (small config)\n") -connectors <- ltx2_text_connectors( - caption_channels = 256L, - text_proj_in_factor = 1L, - video_connector_num_attention_heads = 4L, - video_connector_attention_head_dim = 64L, - video_connector_num_layers = 1L, - video_connector_num_learnable_registers = NULL, - audio_connector_num_attention_heads = 4L, - audio_connector_attention_head_dim = 64L, - audio_connector_num_layers = 1L, - audio_connector_num_learnable_registers = NULL -) -expect_true(!is.null(connectors), info = "Connectors instantiate") -cat(" Connectors instantiated\n") - -# Test 7: Text Connectors forward pass -cat("Test 7: Text Connectors forward\n") -text_embeds <- torch::torch_randn(c(2, 128, 256)) -attn_mask <- torch::torch_ones(c(2, 128)) # All valid -result <- connectors(text_embeds, attn_mask, additive_mask = FALSE) -expect_equal(length(result), 3L, info = "Returns video, audio, and attention mask") -cat(sprintf(" Video embedding shape: [%s]\n", paste(as.numeric(result[[1]]$shape), collapse = ", "))) -cat(sprintf(" Audio embedding shape: [%s]\n", paste(as.numeric(result[[2]]$shape), collapse = ", "))) - -# Test 8: encode_text_ltx2 with random backend -cat("Test 8: encode_text_ltx2 (random backend)\n") -result <- encode_text_ltx2( - prompt = c("A cat sitting on a mat", "A dog running in a field"), - backend = "random", - max_sequence_length = 128L, - caption_channels = 256L -) -expect_true(!is.null(result$prompt_embeds), info = "Returns prompt_embeds") -expect_true(!is.null(result$prompt_attention_mask), info = "Returns attention_mask") -expect_equal(as.numeric(result$prompt_embeds$shape), c(2, 128, 256), info = "Embeddings shape correct") -expect_equal(as.numeric(result$prompt_attention_mask$shape), c(2, 128), info = "Mask shape correct") - -# Test 9: pack_text_embeds -cat("Test 9: pack_text_embeds\n") -# Simulate Gemma output: [batch, seq_len, hidden_dim, num_layers] -hidden_states <- torch::torch_randn(c(2, 64, 128, 4)) # 4 layers -sequence_lengths <- c(50L, 60L) # Valid lengths -packed <- pack_text_embeds( - hidden_states, - sequence_lengths, - padding_side = "left" -) -expect_equal(as.numeric(packed$shape), c(2, 64, 512), info = "Packed shape correct (128 * 4 = 512)") - -# Test 10: Full integration - connectors with encoded text -cat("Test 10: Full integration test\n") -torch::with_no_grad({ - # 1. Get text embeddings (random for testing) - text_result <- encode_text_ltx2( - prompt = "A beautiful sunset over the ocean", - backend = "random", - max_sequence_length = 128L, - caption_channels = 256L - ) - - # 2. Create connectors - connectors <- ltx2_text_connectors( - caption_channels = 256L, - text_proj_in_factor = 1L, - video_connector_num_attention_heads = 4L, - video_connector_attention_head_dim = 64L, - video_connector_num_layers = 1L, - video_connector_num_learnable_registers = NULL, - audio_connector_num_attention_heads = 4L, - audio_connector_attention_head_dim = 64L, - audio_connector_num_layers = 1L, - audio_connector_num_learnable_registers = NULL - ) - - # 3. Process through connectors - connector_result <- connectors( - text_result$prompt_embeds, - text_result$prompt_attention_mask, - additive_mask = FALSE - ) - - video_embeds <- connector_result[[1]] - audio_embeds <- connector_result[[2]] -}) - -expect_true(!is.null(video_embeds), info = "Video embeddings produced") -expect_true(!is.null(audio_embeds), info = "Audio embeddings produced") -cat(sprintf(" Final video embeddings: [%s]\n", paste(as.numeric(video_embeds$shape), collapse = ", "))) -cat(sprintf(" Final audio embeddings: [%s]\n", paste(as.numeric(audio_embeds$shape), collapse = ", "))) - -cat("\nAll LTX2 Text Encoder tests completed\n") diff --git a/inst/tinytest/test_txt2vid_ltx2.R b/inst/tinytest/test_txt2vid_ltx2.R deleted file mode 100644 index bea23ed..0000000 --- a/inst/tinytest/test_txt2vid_ltx2.R +++ /dev/null @@ -1,46 +0,0 @@ -# Tests for LTX-2 Video Generation Pipeline -# Note: Full integration tests require model weights - -# Test 1: Pipeline function exists and has expected signature -cat("Test 1: txt2vid_ltx2 function signature\n") -expect_true(is.function(txt2vid_ltx2), info = "txt2vid_ltx2 should be a function") - -# Check key parameters exist -params <- names(formals(txt2vid_ltx2)) -expect_true("prompt" %in% params, info = "Should have prompt param") -expect_true("width" %in% params, info = "Should have width param") -expect_true("height" %in% params, info = "Should have height param") -expect_true("num_frames" %in% params, info = "Should have num_frames param") -expect_true("memory_profile" %in% params, info = "Should have memory_profile param") -expect_true("text_backend" %in% params, info = "Should have text_backend param") - -# Test 2: Default parameters are sensible -cat("Test 2: Default parameters\n") -defaults <- formals(txt2vid_ltx2) -expect_equal(defaults$width, 768L, info = "Default width should be 768") -expect_equal(defaults$height, 512L, info = "Default height should be 512") -expect_equal(defaults$num_frames, 121L, info = "Default frames should be 121") -expect_equal(defaults$num_inference_steps, 8L, info = "Default steps should be 8 (distilled)") -expect_equal(defaults$guidance_scale, 4.0, info = "Default CFG should be 4.0") - -# Test 3: Memory profile resolution -cat("Test 3: Memory profile parameter\n") -profile_str <- ltx2_memory_profile(vram_gb = 8) -expect_equal(profile_str$name, "low", info = "8GB should resolve to low profile") - -# Test 4: Latent dimension calculation -cat("Test 4: Latent dimensions\n") -# LTX-2 uses 32x spatial and 8x temporal compression -width <- 768L -height <- 512L -num_frames <- 121L - -latent_width <- width %/% 32L -latent_height <- height %/% 32L -latent_frames <- (num_frames - 1L) %/% 8L + 1L - -expect_equal(latent_width, 24L, info = "Latent width correct") -expect_equal(latent_height, 16L, info = "Latent height correct") -expect_equal(latent_frames, 16L, info = "Latent frames correct (121 -> 16)") - -cat("\nLTX-2 pipeline tests completed\n") diff --git a/inst/tinytest/test_txt2vid_ltx23.R b/inst/tinytest/test_txt2vid_ltx23.R new file mode 100644 index 0000000..c5aa7b6 --- /dev/null +++ b/inst/tinytest/test_txt2vid_ltx23.R @@ -0,0 +1,211 @@ +# End-to-end plumbing test for the LTX-2.3 pipeline with tiny +# random-weight components (no model downloads; CPU-safe). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +torch::torch_manual_seed(99) + +# Tiny components wired to the pipeline's structural constants: +# 128 video latent channels, 8 audio channels x 16 latent mel bins. +transformer <- ltx23_transformer( + in_channels = 128L, out_channels = 128L, + num_attention_heads = 2L, attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 128L, audio_out_channels = 128L, + audio_num_attention_heads = 2L, audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 1L +) +transformer$eval() + +connectors <- ltx23_text_connectors( + caption_channels = 8L, text_proj_in_factor = 3L, + video_connector_num_attention_heads = 2L, + video_connector_attention_head_dim = 8L, + video_connector_num_layers = 1L, + video_connector_num_learnable_registers = 4L, + audio_connector_num_attention_heads = 2L, + audio_connector_attention_head_dim = 4L, + audio_connector_num_layers = 1L, + audio_connector_num_learnable_registers = 4L, + video_hidden_dim = 16L, audio_hidden_dim = 8L +) +connectors$eval() + +vae <- ltx23_video_vae( + latent_channels = 128L, + block_out_channels = c(8L, 8L, 8L, 8L), + decoder_block_out_channels = c(4L, 8L, 8L, 16L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + decoder_layers_per_block = c(1L, 1L, 1L, 1L, 1L) +) +vae$eval() + +audio_vae <- ltx23_audio_vae( + base_channels = 128L, ch_mult = c(1L, 1L, 1L), num_res_blocks = 1L, + latent_channels = 8L, mel_bins = 64L +) +audio_vae$eval() + +vocoder <- ltx23_vocoder_with_bwe( + in_channels = 128L, hidden_channels = 16L, + upsample_kernel_sizes = c(4L, 4L), upsample_factors = c(2L, 2L), + resnet_kernel_sizes = c(3L), resnet_dilations = list(c(1L, 3L)), + bwe_in_channels = 128L, bwe_hidden_channels = 8L, + bwe_upsample_kernel_sizes = c(8L, 4L), bwe_upsample_factors = c(4L, 2L), + bwe_resnet_kernel_sizes = c(3L), bwe_resnet_dilations = list(c(1L, 3L)), + filter_length = 8L, hop_length = 2L, window_length = 8L, + num_mel_channels = 64L, + input_sampling_rate = 100L, output_sampling_rate = 400L +) +vocoder$eval() + +pipe <- structure( + list( + transformer = transformer, + connectors = connectors, + vae = vae, + audio_vae = audio_vae, + vocoder = vocoder + ), + class = "ltx23_pipeline" +) + +stub_embeds <- list( + prompt_embeds = torch::torch_randn(1L, 8L, 8L, 3L), + prompt_attention_mask = torch::torch_ones(1L, 8L) +) + +res <- txt2vid_ltx2( + prompt = "a tiny test", + pipeline = pipe, + prompt_embeds = stub_embeds, + width = 64L, height = 64L, num_frames = 9L, frame_rate = 24, + seed = 7L, + device = "cpu", + dtype = "float32", + verbose = FALSE +) + +# Video: [frames, height, width, 3] in [0, 1] +expect_equal(dim(res$video), c(9L, 64L, 64L, 3L)) +expect_true(all(res$video >= 0 & res$video <= 1)) +expect_true(all(is.finite(res$video))) + +# Audio: [2, samples] in [-1, 1]; latent L = round(9/24*25) = 9 -> +# mel frames 4*9-3 = 33 (padded to hop multiple 34) -> vocoder x4 -> BWE x4 +expect_equal(nrow(res$audio), 2L) +expect_true(ncol(res$audio) > 100L) +expect_true(all(abs(res$audio) <= 1)) +expect_true(all(is.finite(res$audio))) +expect_equal(res$sample_rate, 48000L) + +# Guardrails +expect_error( + txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, guidance_scale = 3, + device = "cpu", dtype = "float32"), + pattern = "guidance_scale" +) +expect_error( + txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, width = 50L, + device = "cpu", dtype = "float32"), + pattern = "multiples of 32" +) +expect_error( + txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, num_frames = 10L, + device = "cpu", dtype = "float32"), + pattern = "8k" +) + +# WAV writer round trip (header + size) +wav_path <- tempfile(fileext = ".wav") +write_wav(res$audio, wav_path, sample_rate = 48000L) +expect_true(file.exists(wav_path)) +expect_equal(readBin(wav_path, "raw", 4), charToRaw("RIFF")) +expect_equal(file.size(wav_path), 44 + 2 * 2 * ncol(res$audio)) +unlink(wav_path) + +# MP4 mux when av is available +if (requireNamespace("av", quietly = TRUE)) { + mp4 <- tempfile(fileext = ".mp4") + save_video_ltx23(res$video, mp4, fps = 24, audio = res$audio, + sample_rate = 48000L, verbose = FALSE) + expect_true(file.exists(mp4) && file.size(mp4) > 0) + unlink(mp4) +} + +# --- Prefix conditioning smoke tests ------------------------------------------------- + +# i2v: start image conditions frame 0; pipeline runs end to end +start_img <- array(runif(64 * 64 * 3), dim = c(64L, 64L, 3L)) +res_i2v <- txt2vid_ltx2( + prompt = "tiny i2v", + pipeline = pipe, + prompt_embeds = stub_embeds, + width = 64L, height = 64L, num_frames = 9L, frame_rate = 24, + seed = 7L, device = "cpu", dtype = "float32", + image = start_img, + decode_audio = FALSE, + verbose = FALSE +) +expect_equal(dim(res_i2v$video), c(9L, 64L, 64L, 3L)) +expect_true(all(is.finite(res_i2v$video))) + +# Same seed without conditioning gives a different video (mask engaged) +res_t2v <- txt2vid_ltx2( + prompt = "tiny i2v", + pipeline = pipe, + prompt_embeds = stub_embeds, + width = 64L, height = 64L, num_frames = 9L, frame_rate = 24, + seed = 7L, device = "cpu", dtype = "float32", + decode_audio = FALSE, + verbose = FALSE +) +expect_true(max(abs(res_i2v$video - res_t2v$video)) > 1e-4) + +# Continuation: 9-frame tail array as the frozen prefix +tail_arr <- array(runif(9 * 64 * 64 * 3), dim = c(9L, 64L, 64L, 3L)) +res_cont <- 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, conditioning_frames = 9L, + decode_audio = FALSE, + verbose = FALSE +) +expect_equal(dim(res_cont$video), c(17L, 64L, 64L, 3L)) +expect_true(all(is.finite(res_cont$video))) + +# Guardrails +expect_error( + txt2vid_ltx2("x", pipe, prompt_embeds = stub_embeds, + image = start_img, condition_video = tail_arr, + device = "cpu", dtype = "float32"), + pattern = "not both" +) + + +# --- Audio-conditioned generation (lip-sync plumbing) -------------------------------- + +wav_in <- matrix(runif(2L * 6000L, -0.5, 0.5), nrow = 2L) +res_audio <- txt2vid_ltx2( + prompt = "tiny audio-driven", + pipeline = pipe, + prompt_embeds = stub_embeds, + width = 64L, height = 64L, num_frames = 9L, frame_rate = 24, + seed = 7L, device = "cpu", dtype = "float32", + image = start_img, + audio = wav_in, + verbose = FALSE +) +expect_equal(dim(res_audio$video), c(9L, 64L, 64L, 3L)) +expect_true(all(is.finite(res_audio$video))) +# The original audio rides through untouched +expect_equal(res_audio$sample_rate, 16000L) +expect_equal(res_audio$audio, wav_in) diff --git a/inst/tinytest/test_vae_ltx2.R b/inst/tinytest/test_vae_ltx2.R deleted file mode 100644 index 7937478..0000000 --- a/inst/tinytest/test_vae_ltx2.R +++ /dev/null @@ -1,242 +0,0 @@ -# Tests for LTX2 Video VAE modules - -# Test 1: PerChannelRMSNorm -cat("Test 1: PerChannelRMSNorm initialization and forward\n") -norm <- per_channel_rms_norm() -expect_true(!is.null(norm), info = "RMS norm should initialize") - -x <- torch::torch_randn(c(2, 4, 3, 8, 8)) -y <- norm(x) -expect_equal(as.numeric(x$shape), as.numeric(y$shape), info = "Shape should be preserved") - -# RMS along channel dim should be approximately 1 -rms_out <- torch::torch_sqrt(torch::torch_mean(y^2, dim = 2, keepdim = TRUE) + 1e-8) -expect_true(abs(rms_out$mean()$item() - 1.0) < 0.1, info = "Output RMS should be ~1") - -# Test 2: LTX2VideoCausalConv3d -cat("Test 2: LTX2VideoCausalConv3d initialization\n") -conv <- ltx2_video_causal_conv3d( - in_channels = 4L, - out_channels = 8L, - kernel_size = 3L -) -expect_true(!is.null(conv), info = "Causal conv should initialize") -expect_equal(conv$kernel_size, c(3L, 3L, 3L), info = "Kernel size should be 3x3x3") - -# Test 3: Causal conv forward (causal mode) -cat("Test 3: LTX2VideoCausalConv3d forward (causal)\n") -x <- torch::torch_randn(c(2, 4, 5, 8, 8)) # B, C, T, H, W -y <- conv(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 2, info = "Batch dim preserved") -expect_equal(as.numeric(y$shape[2]), 8, info = "Output channels correct") -expect_equal(as.numeric(y$shape[3]), 5, info = "Temporal dim preserved (stride=1)") -expect_equal(as.numeric(y$shape[4]), 8, info = "Height preserved") -expect_equal(as.numeric(y$shape[5]), 8, info = "Width preserved") - -# Test 4: Causal conv forward (non-causal mode) -cat("Test 4: LTX2VideoCausalConv3d forward (non-causal)\n") -y_nc <- conv(x, causal = FALSE) -expect_equal(as.numeric(y_nc$shape), as.numeric(y$shape), info = "Non-causal shape matches causal") - -# Test 5: LTX2VideoResnetBlock3d -cat("Test 5: LTX2VideoResnetBlock3d initialization and forward\n") -resnet <- ltx2_video_resnet_block3d( - in_channels = 8L, - out_channels = 8L -) -expect_true(!is.null(resnet), info = "ResNet block should initialize") - -x <- torch::torch_randn(c(2, 8, 4, 8, 8)) -y <- resnet(x, causal = TRUE) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "ResNet should preserve shape") - -# Test 6: ResNet block with channel change -cat("Test 6: LTX2VideoResnetBlock3d with channel change\n") -resnet_change <- ltx2_video_resnet_block3d( - in_channels = 8L, - out_channels = 16L -) -y <- resnet_change(x, causal = TRUE) -expect_equal(as.numeric(y$shape[2]), 16, info = "Output channels should change") - -# Test 7: LTXVideoDownsampler3d -# Note: For stride (2,2,2), temporal dim T must satisfy (T + stride-1) % stride == 0 -# So T % 2 == 1 (T must be odd for temporal stride 2) -cat("Test 7: LTXVideoDownsampler3d\n") -downsampler <- ltx_video_downsampler3d( - in_channels = 8L, - out_channels = 16L, - stride = c(2L, 2L, 2L) -) -x <- torch::torch_randn(c(2, 8, 5, 16, 16)) # T=5 (odd for stride 2 compatibility) -y <- downsampler(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 2, info = "Batch preserved") -expect_equal(as.numeric(y$shape[2]), 16, info = "Channels increased") -# Output T = (T_in + stride - 1) / stride = (5 + 1) / 2 = 3 -expect_equal(as.numeric(y$shape[3]), 3, info = "Temporal: (5+1)/2=3") -expect_equal(as.numeric(y$shape[4]), 8, info = "Height halved") -expect_equal(as.numeric(y$shape[5]), 8, info = "Width halved") - -# Test 8: LTXVideoUpsampler3d -cat("Test 8: LTXVideoUpsampler3d\n") -upsampler <- ltx_video_upsampler3d( - in_channels = 16L, - stride = c(2L, 2L, 2L), - residual = TRUE, - upscale_factor = 1L -) -x <- torch::torch_randn(c(2, 16, 2, 8, 8)) -y <- upsampler(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 2, info = "Batch preserved") -# Output channels = in_channels * stride_prod / upscale_factor = 16 * 8 / 1 = 128 -expect_equal(as.numeric(y$shape[3]), 3, info = "Temporal doubled (minus 1 for causal)") -expect_equal(as.numeric(y$shape[4]), 16, info = "Height doubled") -expect_equal(as.numeric(y$shape[5]), 16, info = "Width doubled") - -# Test 9: LTX2VideoDownBlock3D -cat("Test 9: LTX2VideoDownBlock3D\n") -down_block <- ltx2_video_down_block3d( - in_channels = 8L, - out_channels = 16L, - num_layers = 2L, - spatio_temporal_scale = TRUE, - downsample_type = "conv" -) -x <- torch::torch_randn(c(2, 8, 4, 16, 16)) -y <- down_block(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 2, info = "Batch preserved") -expect_equal(as.numeric(y$shape[3]), 2, info = "Temporal halved") -expect_equal(as.numeric(y$shape[4]), 8, info = "Height halved") -expect_equal(as.numeric(y$shape[5]), 8, info = "Width halved") - -# Test 10: LTX2VideoMidBlock3d -cat("Test 10: LTX2VideoMidBlock3d\n") -mid_block <- ltx2_video_mid_block3d( - in_channels = 8L, - num_layers = 2L -) -x <- torch::torch_randn(c(2, 8, 4, 8, 8)) -y <- mid_block(x, causal = TRUE) -expect_equal(as.numeric(y$shape), as.numeric(x$shape), info = "Mid block preserves shape") - -# Test 11: LTX2VideoUpBlock3d -# Note: In LTX2 decoder, in_channels always equals out_channels. -# The upsampler expects out_channels * upscale_factor as input. -# So input tensor channels = out_channels * upscale_factor, output = out_channels -cat("Test 11: LTX2VideoUpBlock3d\n") -up_block <- ltx2_video_up_block3d( - in_channels = 8L, # Same as out_channels (normal LTX2 usage) - out_channels = 8L, - num_layers = 2L, - spatio_temporal_scale = TRUE, - upsample_residual = TRUE, - upscale_factor = 2L -) -# Input has out_channels * upscale_factor = 8 * 2 = 16 channels -x <- torch::torch_randn(c(2, 16, 2, 8, 8)) -y <- up_block(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 2, info = "Batch preserved") -# Output has out_channels = 8 channels -expect_equal(as.numeric(y$shape[2]), 8, info = "Channels: 16 -> 8") - -# Test 12: LTX2VideoEncoder3d (small config) -# Note: For spatiotemporal downsampling, T must be odd so that (T+1)/2 is integer. -# For 2 spatiotemporal downs: T=5 -> (5+1)/2=3 -> (3+1)/2=2 -cat("Test 12: LTX2VideoEncoder3d\n") -encoder <- ltx2_video_encoder3d( - in_channels = 3L, - out_channels = 32L, - block_out_channels = c(32L, 64L), - spatio_temporal_scaling = c(TRUE, TRUE), - layers_per_block = c(1L, 1L, 1L), - downsample_type = c("spatiotemporal", "spatiotemporal"), - patch_size = 2L, - patch_size_t = 1L -) -# Input: T=5 (odd for spatiotemporal compat), H=32, W=32 -x <- torch::torch_randn(c(1, 3, 5, 32, 32)) -y <- encoder(x, causal = TRUE) -expect_equal(as.numeric(y$shape[1]), 1, info = "Encoder batch preserved") -cat(sprintf(" Encoder output shape: [%s]\n", paste(as.numeric(y$shape), collapse=", "))) - -# Test 13: LTX2VideoDecoder3d (small config) -cat("Test 13: LTX2VideoDecoder3d\n") -decoder <- ltx2_video_decoder3d( - in_channels = 32L, - out_channels = 3L, - block_out_channels = c(32L, 64L), - spatio_temporal_scaling = c(TRUE, TRUE), - layers_per_block = c(1L, 1L, 1L), - patch_size = 2L, - patch_size_t = 1L, - upsample_residual = c(TRUE, TRUE), - upsample_factor = c(2L, 2L) -) -# Use a small latent input -z <- torch::torch_randn(c(1, 32, 1, 4, 4)) -out <- decoder(z, causal = TRUE) -expect_equal(as.numeric(out$shape[1]), 1, info = "Decoder batch preserved") -cat(sprintf(" Decoder output shape: [%s]\n", paste(as.numeric(out$shape), collapse=", "))) - -# Test 14: DiagonalGaussianDistribution -cat("Test 14: DiagonalGaussianDistribution\n") -# Create parameters tensor (mean and logvar concatenated along channel dim) -params <- torch::torch_randn(c(2, 64, 2, 4, 4)) # 32 mean + 32 logvar -dist <- diagonal_gaussian_distribution(params) -expect_equal(as.numeric(dist$mean$shape[2]), 32, info = "Mean has half channels") -expect_equal(as.numeric(dist$logvar$shape[2]), 32, info = "Logvar has half channels") -sample <- dist$sample() -expect_equal(as.numeric(sample$shape), as.numeric(dist$mean$shape), info = "Sample shape matches mean") -mode <- dist$mode() -expect_equal(as.numeric(mode$shape), as.numeric(dist$mean$shape), info = "Mode shape matches mean") - -# Test 15: Full VAE instantiation (small config) -cat("Test 15: Full ltx2_video_vae instantiation\n") -vae <- ltx2_video_vae( - in_channels = 3L, - out_channels = 3L, - latent_channels = 32L, - block_out_channels = c(32L, 64L), - decoder_block_out_channels = c(32L, 64L), - layers_per_block = c(1L, 1L, 1L), - decoder_layers_per_block = c(1L, 1L, 1L), - spatio_temporal_scaling = c(TRUE, TRUE), - decoder_spatio_temporal_scaling = c(TRUE, TRUE), - downsample_type = c("spatiotemporal", "spatiotemporal"), - upsample_residual = c(TRUE, TRUE), - upsample_factor = c(2L, 2L), - patch_size = 2L, - patch_size_t = 1L -) -expect_true(!is.null(vae), info = "VAE should instantiate") -expect_true(!is.null(vae$encoder), info = "VAE should have encoder") -expect_true(!is.null(vae$decoder), info = "VAE should have decoder") - -# Test 16: VAE encode -cat("Test 16: VAE encode\n") -x <- torch::torch_randn(c(1, 3, 5, 32, 32)) # T=5 (odd for spatiotemporal) -torch::with_no_grad({ - posterior <- vae$encode(x, causal = TRUE) -}) -expect_true(!is.null(posterior$mean), info = "Posterior should have mean") -expect_true(!is.null(posterior$std), info = "Posterior should have std") -cat(sprintf(" Latent mean shape: [%s]\n", paste(as.numeric(posterior$mean$shape), collapse=", "))) - -# Test 17: VAE decode -cat("Test 17: VAE decode\n") -torch::with_no_grad({ - z <- posterior$sample() - decoded <- vae$decode(z, causal = TRUE) -}) -cat(sprintf(" Decoded shape: [%s]\n", paste(as.numeric(decoded$shape), collapse=", "))) - -# Test 18: VAE enable_tiling -cat("Test 18: VAE enable_tiling\n") -vae$enable_tiling(tile_sample_min_height = 64L, tile_sample_min_width = 64L) -expect_true(vae$use_tiling, info = "Tiling should be enabled") -expect_equal(vae$tile_sample_min_height, 64L, info = "Tile height should be updated") - -vae$disable_tiling() -expect_false(vae$use_tiling, info = "Tiling should be disabled") - -cat("\nAll LTX2 VAE module tests completed\n") diff --git a/inst/tinytest/test_vae_ltx23.R b/inst/tinytest/test_vae_ltx23.R new file mode 100644 index 0000000..620076c --- /dev/null +++ b/inst/tinytest/test_vae_ltx23.R @@ -0,0 +1,213 @@ +# Parity tests for the LTX-2.3 video VAE port against diffusers +# reference fixtures (tools/gen_fixtures_vae.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "vae_ltx23.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/vae_ltx23.safetensors" +if (!file.exists(fixture_path)) exit_file("vae fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +load_group <- function(module, prefix) { + keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE) + w <- fx[keys] + names(w) <- sub(paste0("^", prefix, "\\."), "", keys) + dests <- c(module$named_parameters(), module$named_buffers()) + if (!setequal(names(w), names(dests))) { + stop( + prefix, ": name mismatch. missing dest: ", + paste(utils::head(setdiff(names(w), names(dests)), 3), collapse = ", "), + " | unfilled: ", + paste(utils::head(setdiff(names(dests), names(w)), 3), collapse = ", ") + ) + } + torch::with_no_grad({ + for (name in names(w)) dests[[name]]$copy_(w[[name]]) + }) + module$eval() + module +} + +# --- Causal conv ----------------------------------------------------------------- + +cc <- load_group(ltx23_causal_conv3d(4L, 6L, kernel_size = 3L), "cc") +torch::with_no_grad({ + out_c <- cc(fx$cc_x, causal = TRUE) + out_nc <- cc(fx$cc_x, causal = FALSE) +}) +expect_true(max_abs_diff(out_c, fx$cc_out_causal) < 1e-5) +expect_true(max_abs_diff(out_nc, fx$cc_out_noncausal) < 1e-5) + +# --- Resnet block with channel-change shortcut ------------------------------------- + +rb <- load_group(ltx23_video_resnet_block3d(4L, 8L), "rb") +torch::with_no_grad(out_rb <- rb(fx$cc_x, causal = TRUE)) +expect_true(max_abs_diff(out_rb, fx$rb_out) < 1e-5) + +# --- Down/upsamplers --------------------------------------------------------------- + +ds <- load_group( + ltx23_video_downsampler3d(8L, 16L, stride = c(2L, 1L, 1L)), "ds" +) +torch::with_no_grad(out_ds <- ds(fx$ds_x, causal = TRUE)) +expect_equal(as.integer(out_ds$shape), as.integer(fx$ds_out$shape)) +expect_true(max_abs_diff(out_ds, fx$ds_out) < 1e-5) + +us <- load_group( + ltx23_video_upsampler3d(16L, stride = c(2L, 2L, 2L), residual = TRUE, upscale_factor = 2L), + "us" +) +torch::with_no_grad(out_us <- us(fx$us_x, causal = FALSE)) +expect_equal(as.integer(out_us$shape), as.integer(fx$us_out$shape)) +expect_true(max_abs_diff(out_us, fx$us_out) < 1e-5) + +# --- Tiny encoder (2.3 structure) -------------------------------------------------- + +enc <- load_group( + ltx23_video_encoder3d( + in_channels = 3L, + out_channels = 4L, + block_out_channels = c(8L, 16L, 32L, 32L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + patch_size = 4L + ), + "enc" +) +torch::with_no_grad(out_enc <- enc(fx$enc_x)) +expect_equal(as.integer(out_enc$shape), as.integer(fx$enc_out$shape)) +expect_true(max_abs_diff(out_enc, fx$enc_out) < 1e-5) + +# --- Tiny decoder (2.3 structure) -------------------------------------------------- + +dec <- load_group( + ltx23_video_decoder3d( + in_channels = 4L, + out_channels = 3L, + block_out_channels = c(16L, 32L, 32L, 64L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + upsample_type = c("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + upsample_residual = c(FALSE, FALSE, FALSE, FALSE), + upsample_factor = c(2L, 2L, 1L, 2L), + patch_size = 4L + ), + "dec" +) +torch::with_no_grad(out_dec <- dec(fx$dec_x)) +expect_equal(as.integer(out_dec$shape), c(1L, 3L, 17L, 64L, 64L)) +expect_true(max_abs_diff(out_dec, fx$dec_out) < 1e-5) + +# --- Latent (de)normalization round trip ------------------------------------------- + +lat <- torch::torch_randn(1L, 4L, 2L, 2L, 2L) +mean <- torch::torch_randn(4L) +std <- torch::torch_rand(4L) + 0.5 +lat_n <- ltx23_normalize_latents(lat, mean, std) +lat_rt <- ltx23_denormalize_latents(lat_n, mean, std) +expect_true(max_abs_diff(lat, lat_rt) < 1e-5) + +# --- Key mapper -------------------------------------------------------------------- + +expect_equal( + ltx23_map_vae_key("vae.decoder.up_blocks.1.conv.conv.weight"), + "decoder.up_blocks.0.upsamplers.0.conv.conv.weight" +) +expect_equal( + ltx23_map_vae_key("vae.decoder.up_blocks.0.res_blocks.0.conv1.conv.weight"), + "decoder.mid_block.resnets.0.conv1.conv.weight" +) +expect_equal( + ltx23_map_vae_key("vae.decoder.up_blocks.8.res_blocks.1.norm3.weight"), + "decoder.up_blocks.3.resnets.1.norm3.weight" +) +expect_equal( + ltx23_map_vae_key("vae.encoder.down_blocks.8.res_blocks.0.conv2.conv.bias"), + "encoder.mid_block.resnets.0.conv2.conv.bias" +) +expect_equal( + ltx23_map_vae_key("vae.encoder.down_blocks.3.conv.conv.weight"), + "encoder.down_blocks.1.downsamplers.0.conv.conv.weight" +) +expect_equal( + ltx23_map_vae_key("vae.per_channel_statistics.mean-of-means"), + "latents_mean" +) +expect_equal(ltx23_map_vae_key("vae.encoder.conv_in.conv.weight"), "encoder.conv_in.conv.weight") + +# --- Tiled decode ------------------------------------------------------------------ + +# Reuse the tiny decoder inside a vae wrapper with small tile settings +vae_t <- ltx23_video_vae( + latent_channels = 4L, + block_out_channels = c(8L, 8L, 8L, 8L), + decoder_block_out_channels = c(16L, 32L, 32L, 64L), + layers_per_block = c(1L, 1L, 1L, 1L, 1L), + decoder_layers_per_block = c(1L, 1L, 1L, 1L, 1L) +) +# Adopt the fixture decoder weights (same architecture) +dests_t <- vae_t$decoder$named_parameters() +w_dec <- fx[grep("^dec\\.", names(fx))] +names(w_dec) <- sub("^dec\\.", "", names(w_dec)) +torch::with_no_grad({ + for (name in names(w_dec)) dests_t[[name]]$copy_(w_dec[[name]]) +}) +vae_t$eval() + +z_big <- torch::torch_randn(1L, 4L, 5L, 4L, 6L) +torch::with_no_grad(ref_full <- vae_t$decode(z_big)) +expect_equal(as.integer(ref_full$shape), c(1L, 3L, 33L, 128L, 192L)) + +# Spatial tiling: tiles of 2x2 latent with stride 1 +vae_t$enable_tiling(spatial = TRUE, temporal = FALSE) +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 +torch::with_no_grad(tiled <- vae_t$decode(z_big)) +expect_equal(as.integer(tiled$shape), as.integer(ref_full$shape)) +expect_true(as.logical(tiled$isfinite()$all()$item())) +# With random weights, tile borders dominate; value similarity to the +# full decode is only meaningful with trained weights. Instead check the +# degenerate case: a tile that covers the whole input must be identical. +vae_t$tile_sample_min_height <- 128L +vae_t$tile_sample_min_width <- 192L +vae_t$tile_sample_stride_height <- 96L +vae_t$tile_sample_stride_width <- 160L +torch::with_no_grad(tiled_one <- vae_t$decode(z_big)) +expect_true(max_abs_diff(tiled_one, ref_full) < 1e-5) +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 + +# Temporal + spatial tiling +vae_t$enable_tiling(spatial = TRUE, temporal = TRUE) +vae_t$tile_sample_min_num_frames <- 16L +vae_t$tile_sample_stride_num_frames <- 8L +torch::with_no_grad(tiled_t <- vae_t$decode(z_big)) +expect_equal(as.integer(tiled_t$shape), as.integer(ref_full$shape)) +expect_true(as.logical(tiled_t$isfinite()$all()$item())) + +# Small inputs bypass tiling entirely (identical output) +z_small <- torch::torch_randn(1L, 4L, 2L, 2L, 2L) +torch::with_no_grad({ + a <- vae_t$decode(z_small) + vae_t$enable_tiling(spatial = FALSE, temporal = FALSE) + b <- vae_t$decode(z_small) +}) +expect_true(as.numeric((a - b)$abs()$max()) == 0) diff --git a/inst/validation/.venv/bin/python b/inst/validation/.venv/bin/python deleted file mode 120000 index b8a0adb..0000000 --- a/inst/validation/.venv/bin/python +++ /dev/null @@ -1 +0,0 @@ -python3 \ No newline at end of file diff --git a/inst/validation/.venv/bin/python3 b/inst/validation/.venv/bin/python3 deleted file mode 120000 index ae65fda..0000000 --- a/inst/validation/.venv/bin/python3 +++ /dev/null @@ -1 +0,0 @@ -/usr/bin/python3 \ No newline at end of file diff --git a/inst/validation/.venv/bin/python3.12 b/inst/validation/.venv/bin/python3.12 deleted file mode 120000 index b8a0adb..0000000 --- a/inst/validation/.venv/bin/python3.12 +++ /dev/null @@ -1 +0,0 @@ -python3 \ No newline at end of file diff --git a/inst/validation/.venv/lib64 b/inst/validation/.venv/lib64 deleted file mode 120000 index 7951405..0000000 --- a/inst/validation/.venv/lib64 +++ /dev/null @@ -1 +0,0 @@ -lib \ No newline at end of file diff --git a/inst/validation/.venv/pyvenv.cfg b/inst/validation/.venv/pyvenv.cfg deleted file mode 100644 index d08feca..0000000 --- a/inst/validation/.venv/pyvenv.cfg +++ /dev/null @@ -1,5 +0,0 @@ -home = /usr/bin -include-system-site-packages = false -version = 3.12.3 -executable = /usr/bin/python3.12 -command = /usr/bin/python3 -m venv /home/troy/diffuseR/inst/validation/.venv diff --git a/man/CLIPTokenizer.Rd b/man/CLIPTokenizer.Rd index cddad68..37fc454 100644 --- a/man/CLIPTokenizer.Rd +++ b/man/CLIPTokenizer.Rd @@ -3,12 +3,10 @@ \alias{CLIPTokenizer} \title{Tokenize a prompt} \usage{ -CLIPTokenizer( - prompt, - merges = system.file("tokenizer/merges.txt", package = "diffuseR"), - vocab_file = system.file("tokenizer/vocab.json", package = "diffuseR"), - pad_token = 0L -) +CLIPTokenizer(prompt, + merges = system.file("tokenizer/merges.txt", package = "diffuseR"), + vocab_file = system.file("tokenizer/vocab.json", package = "diffuseR"), + pad_token = 0L) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} diff --git a/man/apply_interleaved_rotary_emb.Rd b/man/apply_interleaved_rotary_emb.Rd deleted file mode 100644 index bf98173..0000000 --- a/man/apply_interleaved_rotary_emb.Rd +++ /dev/null @@ -1,19 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{apply_interleaved_rotary_emb} -\alias{apply_interleaved_rotary_emb} -\title{Apply interleaved rotary embeddings} -\usage{ -apply_interleaved_rotary_emb(x, freqs) -} -\arguments{ -\item{x}{torch tensor. Query or key tensor of shape (B, S, C).} - -\item{freqs}{List. Contains cos_freqs and sin_freqs from rope_forward().} -} -\value{ -torch tensor. Rotated tensor with same shape as input. -} -\description{ -Applies rotary position embeddings to query or key tensors using the -interleaved format where real and imaginary components alternate. -} diff --git a/man/apply_interleaved_rotary_emb_list.Rd b/man/apply_interleaved_rotary_emb_list.Rd deleted file mode 100644 index 3c41a47..0000000 --- a/man/apply_interleaved_rotary_emb_list.Rd +++ /dev/null @@ -1,11 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{apply_interleaved_rotary_emb_list} -\alias{apply_interleaved_rotary_emb_list} -\title{Apply interleaved rotary embedding} -\usage{ -apply_interleaved_rotary_emb_list(x, freqs) -} -\description{ -Apply interleaved rotary embedding -} -\keyword{internal} diff --git a/man/apply_split_rotary_emb.Rd b/man/apply_split_rotary_emb.Rd deleted file mode 100644 index 2a1bc4d..0000000 --- a/man/apply_split_rotary_emb.Rd +++ /dev/null @@ -1,19 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{apply_split_rotary_emb} -\alias{apply_split_rotary_emb} -\title{Apply split rotary embeddings} -\usage{ -apply_split_rotary_emb(x, freqs) -} -\arguments{ -\item{x}{torch tensor. Query or key tensor.} - -\item{freqs}{List. Contains cos_freqs and sin_freqs from rope_forward().} -} -\value{ -torch tensor. Rotated tensor with same shape as input. -} -\description{ -Applies rotary position embeddings to query or key tensors using the -split format where first half is real and second half is imaginary. -} diff --git a/man/audio_encode_ltx23.Rd b/man/audio_encode_ltx23.Rd new file mode 100644 index 0000000..7e14436 --- /dev/null +++ b/man/audio_encode_ltx23.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{audio_encode_ltx23} +\alias{audio_encode_ltx23} +\title{Audio Conditioning Frontend for LTX-2.3} +\description{ +Turns user audio into the normalized, packed audio latents the joint +denoiser conditions on (lip sync): decode to 16 kHz stereo PCM, +log-mel via a causal STFT (filter 1024, hop 160, 64 slaney-normed +mel bins to 8 kHz — the checkpoint's preprocessing spec), then the +audio VAE encoder in argmax mode. The STFT and mel-filterbank +constructors were verified against the checkpoint's stored vocoder +bases (identical up to bf16 rounding), so the convention matches +training. +} diff --git a/man/audio_vae_ltx23.Rd b/man/audio_vae_ltx23.Rd new file mode 100644 index 0000000..ff8a478 --- /dev/null +++ b/man/audio_vae_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{audio_vae_ltx23} +\alias{audio_vae_ltx23} +\title{LTX-2.3 Audio VAE} +\description{ +Fresh R port of the LTX-2 audio autoencoder from the diffusers +reference (Apache-2.0, autoencoder_kl_ltx2_audio.py), configured per +the checkpoint: pixel norm, height-axis causality, base 128 channels +with multipliers (1, 2, 4), 8 latent channels, 64 mel bins, no +attention. The decoder produces mel for the vocoder; the encoder +turns user audio into conditioning latents (lip sync). +} diff --git a/man/checkpoint_ltx23.Rd b/man/checkpoint_ltx23.Rd new file mode 100644 index 0000000..ad9f141 --- /dev/null +++ b/man/checkpoint_ltx23.Rd @@ -0,0 +1,13 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{checkpoint_ltx23} +\alias{checkpoint_ltx23} +\title{LTX-2.3 Single-File Checkpoint Reader} +\description{ +LTX 2.3 checkpoints ship as one safetensors file containing every +component (transformer, connectors, video VAE, audio VAE, vocoder), +with the model version and full component configuration embedded in +the safetensors metadata. These helpers open the file, validate the +version, split the key space by component, and stream tensors into +R torch modules one at a time so the 46 GB file is never fully +materialized in memory. +} diff --git a/man/condition_ltx23.Rd b/man/condition_ltx23.Rd new file mode 100644 index 0000000..734b7b4 --- /dev/null +++ b/man/condition_ltx23.Rd @@ -0,0 +1,13 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{condition_ltx23} +\alias{condition_ltx23} +\title{LTX-2.3 Prefix Conditioning (image-to-video, video continuation)} +\description{ +Fresh R port of the frame-conditioning mechanics from the diffusers +reference (Apache-2.0, pipelines/ltx2/pipeline_ltx2_image2video.py +and pipeline_ltx2_condition.py), restricted to prefix conditioning +at latent index 0 with strength 1: a single start image (i2v) or the +leading pixel frames of a previous clip (continuation). Conditioned +latent tokens are initialized from the VAE-encoded pixels, see a +per-token timestep of zero, and are frozen through the Euler loop. +} diff --git a/man/configure_vae_for_profile.Rd b/man/configure_vae_for_profile.Rd deleted file mode 100644 index 92238ce..0000000 --- a/man/configure_vae_for_profile.Rd +++ /dev/null @@ -1,25 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{configure_vae_for_profile} -\alias{configure_vae_for_profile} -\title{Configure VAE for Memory Profile} -\usage{ -configure_vae_for_profile(vae, profile) -} -\arguments{ -\item{vae}{The LTX2 VAE module.} - -\item{profile}{Memory profile from `ltx2_memory_profile()`.} -} -\value{ -The VAE (modified in place). -} -\description{ -Sets VAE tiling parameters based on memory profile. -} -\examples{ -\dontrun{ -profile <- ltx2_memory_profile(vram_gb = 8) -vae <- load_ltx2_vae(...) -configure_vae_for_profile(vae, profile) -} -} diff --git a/man/connectors_ltx23.Rd b/man/connectors_ltx23.Rd new file mode 100644 index 0000000..0815294 --- /dev/null +++ b/man/connectors_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{connectors_ltx23} +\alias{connectors_ltx23} +\title{LTX-2.3 Text Embedding Connectors} +\description{ +Fresh R port of the LTX text connectors from the diffusers reference +(Apache-2.0, src/diffusers/pipelines/ltx2/connectors.py). The +connectors take raw stacked per-layer Gemma3 hidden states +[batch, seq, caption_channels, num_layers + 1], normalize and project +them per modality, replace padding with learnable registers, and run a +small 1D transformer per modality to produce the DiT text embeddings. +} diff --git a/man/ddim_scheduler_create.Rd b/man/ddim_scheduler_create.Rd index 7fff721..6cd646b 100644 --- a/man/ddim_scheduler_create.Rd +++ b/man/ddim_scheduler_create.Rd @@ -3,17 +3,13 @@ \alias{ddim_scheduler_create} \title{Create a DDIM Scheduler} \usage{ -ddim_scheduler_create( - num_train_timesteps = 1000, - num_inference_steps = 50, - eta = 0, - beta_schedule = c("linear", "scaled_linear", "cosine"), - beta_start = 0.00085, - beta_end = 0.012, - rescale_betas_zero_snr = FALSE, - dtype = torch::torch_float32(), - device = c(torch::torch_device("cpu"), torch::torch_device("cuda")) -) +ddim_scheduler_create(num_train_timesteps = 1000, num_inference_steps = 50, + eta = 0, + beta_schedule = c("linear", "scaled_linear", "cosine"), + beta_start = 0.00085, beta_end = 0.012, + rescale_betas_zero_snr = FALSE, + dtype = torch::torch_float32(), + device = c(torch::torch_device("cpu"), torch::torch_device("cuda"))) } \arguments{ \item{num_train_timesteps}{Integer. The number of diffusion steps used to diff --git a/man/ddim_scheduler_step.Rd b/man/ddim_scheduler_step.Rd index e9ac78b..9df3960 100644 --- a/man/ddim_scheduler_step.Rd +++ b/man/ddim_scheduler_step.Rd @@ -3,22 +3,12 @@ \alias{ddim_scheduler_step} \title{Perform a DDIM scheduler step} \usage{ -ddim_scheduler_step( - model_output, - timestep, - sample, - schedule, - eta = 0, - use_clipped_model_output = FALSE, - thresholding = FALSE, - generator = NULL, - variance_noise = NULL, - clip_sample = FALSE, - set_alpha_to_one = FALSE, - prediction_type = c("epsilon", "sample", "v_prediction"), - dtype = torch::torch_float32(), - device = "cpu" -) +ddim_scheduler_step(model_output, timestep, sample, schedule, eta = 0, + use_clipped_model_output = FALSE, thresholding = FALSE, + generator = NULL, variance_noise = NULL, + clip_sample = FALSE, set_alpha_to_one = FALSE, + prediction_type = c("epsilon", "sample", "v_prediction"), + dtype = torch::torch_float32(), device = "cpu") } \arguments{ \item{model_output}{Numeric array. The output from the diffusion model, typically diff --git a/man/dequantize_int4.Rd b/man/dequantize_int4.Rd deleted file mode 100644 index 15cd2a9..0000000 --- a/man/dequantize_int4.Rd +++ /dev/null @@ -1,26 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{dequantize_int4} -\alias{dequantize_int4} -\title{Dequantize INT4 Tensor} -\usage{ -dequantize_int4(q, dtype = torch::torch_float16(), device = "cpu") -} -\arguments{ -\item{q}{List. Quantized data from `quantize_int4()`.} - -\item{dtype}{Torch dtype. Output dtype (default float16).} - -\item{device}{Character. Target device.} -} -\value{ -Tensor with original shape and specified dtype. -} -\description{ -Reconstructs a float tensor from INT4-quantized data. -} -\examples{ -\dontrun{ -q <- quantize_int4(weights) -weights_approx <- dequantize_int4(q, dtype = torch_float16(), device = "cuda") -} -} diff --git a/man/diagonal_gaussian_distribution.Rd b/man/diagonal_gaussian_distribution.Rd deleted file mode 100644 index 486a788..0000000 --- a/man/diagonal_gaussian_distribution.Rd +++ /dev/null @@ -1,13 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{diagonal_gaussian_distribution} -\alias{diagonal_gaussian_distribution} -\title{Diagonal Gaussian Distribution} -\usage{ -diagonal_gaussian_distribution(parameters) -} -\arguments{ -\item{parameters}{Tensor of concatenated mean and log variance.} -} -\description{ -Represents a diagonal Gaussian distribution for VAE latents. -} diff --git a/man/dit_ltx23.Rd b/man/dit_ltx23.Rd new file mode 100644 index 0000000..d85b219 --- /dev/null +++ b/man/dit_ltx23.Rd @@ -0,0 +1,11 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_ltx23} +\alias{dit_ltx23} +\title{LTX-2.3 Audio-Video Diffusion Transformer} +\description{ +Fresh R port of the LTX-2 transformer from the diffusers reference +(Apache-2.0, src/diffusers/models/transformers/transformer_ltx2.py), +configured for LTX 2.3: gated attention, cross-attention modulation, +prompt AdaLN, split RoPE, and connector-projected text embeddings +(no in-model caption projection). +} diff --git a/man/dit_ltx23_modules.Rd b/man/dit_ltx23_modules.Rd new file mode 100644 index 0000000..abc9f9d --- /dev/null +++ b/man/dit_ltx23_modules.Rd @@ -0,0 +1,10 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_ltx23_modules} +\alias{dit_ltx23_modules} +\title{LTX-2.3 Transformer Building Blocks} +\description{ +Fresh R port of the LTX-2 transformer components from the diffusers +reference (Apache-2.0, src/diffusers/models/transformers/ +transformer_ltx2.py and shared modules). Field names mirror the +diffusers module tree so checkpoint keys map 1:1. +} diff --git a/man/dit_offloaded_forward.Rd b/man/dit_offloaded_forward.Rd deleted file mode 100644 index 09db2fd..0000000 --- a/man/dit_offloaded_forward.Rd +++ /dev/null @@ -1,53 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{dit_offloaded_forward} -\alias{dit_offloaded_forward} -\title{DiT Chunk-based Forward Pass} -\usage{ -dit_offloaded_forward( - hidden_states, - layers, - chunk_size = 1L, - device = "cuda", - verbose = FALSE, - ... -) -} -\arguments{ -\item{hidden_states}{Input tensor.} - -\item{layers}{List of transformer layers (on CPU).} - -\item{chunk_size}{Integer. Number of layers to load at once (default 1).} - -\item{device}{Target device for computation.} - -\item{verbose}{Logical. Print progress.} - -\item{...}{Additional arguments passed to each layer.} -} -\value{ -Output tensor (on CPU). -} -\description{ -Runs transformer layers in chunks, moving each chunk to GPU before -computation and back to CPU after. Balances memory usage with speed. -} -\examples{ -\dontrun{ -# Layer-by-layer for 8GB VRAM -output <- dit_offloaded_forward( - hidden_states, - model$transformer_blocks, - chunk_size = 1, - device = "cuda" -) - -# Chunk-based for 16GB VRAM -output <- dit_offloaded_forward( - hidden_states, - model$transformer_blocks, - chunk_size = 12, # 12 layers at a time - device = "cuda" -) -} -} diff --git a/man/dot-ltx23_jit_pack_block.Rd b/man/dot-ltx23_jit_pack_block.Rd new file mode 100644 index 0000000..3d11af7 --- /dev/null +++ b/man/dot-ltx23_jit_pack_block.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_jit_pack_block} +\alias{.ltx23_jit_pack_block} +\title{Pack a transformer block's weights for the JIT stack} +\usage{ +.ltx23_jit_pack_block(block) +} +\arguments{ +\item{block}{An NF4-quantized \code{ltx23_transformer_block}.} +} +\value{ +List of 114 tensors. +} +\description{ +Returns the block's tensors in the fixed 114-slot layout consumed by +the compiled \code{stack_nf4}/\code{block_nf4} TorchScript functions. +Tensor handles are borrowed, not copied. +} +\keyword{internal} diff --git a/man/dot-ltx23_jit_run_stack.Rd b/man/dot-ltx23_jit_run_stack.Rd new file mode 100644 index 0000000..4d7f3e1 --- /dev/null +++ b/man/dot-ltx23_jit_run_stack.Rd @@ -0,0 +1,24 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_jit_run_stack} +\alias{.ltx23_jit_run_stack} +\title{Run the block stack through the compiled TorchScript path} +\usage{ +.ltx23_jit_run_stack(blocks, hidden_states, audio_hidden_states, + encoder_hidden_states, audio_encoder_hidden_states, temb, + temb_audio, temb_ca_scale_shift, + temb_ca_audio_scale_shift, temb_ca_gate, + temb_ca_audio_gate, temb_prompt, temb_prompt_audio, + video_rotary_emb, audio_rotary_emb, ca_video_rotary_emb, + ca_audio_rotary_emb, encoder_attention_mask = NULL, + audio_encoder_attention_mask = NULL) +} +\value{ +list(hidden_states, audio_hidden_states) +} +\description{ +One R-to-libtorch crossing for all blocks: no per-op dispatch, no R +tensor garbage, fused SDPA. Masks must already be additive +\code{[B, 1, 1, S]} (or NULL); rope tensors are the \code{[.., r]} +cos/sin pairs used by the eager path. +} +\keyword{internal} diff --git a/man/dot-ltx23_pin_component.Rd b/man/dot-ltx23_pin_component.Rd new file mode 100644 index 0000000..048aa48 --- /dev/null +++ b/man/dot-ltx23_pin_component.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_pin_component} +\alias{.ltx23_pin_component} +\title{Pin a component's tensors for fast phase transfer} +\usage{ +.ltx23_pin_component(module) +} +\arguments{ +\item{module}{An nn_module on the CPU.} +} +\value{ +A list of \code{list(live, pinned)} tensor pairs, or NULL + if pinning is unavailable (no CUDA, or page-locking failed). +} +\description{ +Pin a component's tensors for fast phase transfer +} +\keyword{internal} diff --git a/man/dot-ltx23_release_attn_buffers.Rd b/man/dot-ltx23_release_attn_buffers.Rd new file mode 100644 index 0000000..f7e07ab --- /dev/null +++ b/man/dot-ltx23_release_attn_buffers.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_release_attn_buffers} +\alias{.ltx23_release_attn_buffers} +\title{Release the attention scratch buffers} +\usage{ +.ltx23_release_attn_buffers() +} +\value{ +Invisibly, NULL. +} +\description{ +Release the attention scratch buffers +} +\keyword{internal} diff --git a/man/dot-ltx23_release_vae_traces.Rd b/man/dot-ltx23_release_vae_traces.Rd new file mode 100644 index 0000000..a58a5b3 --- /dev/null +++ b/man/dot-ltx23_release_vae_traces.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_release_vae_traces} +\alias{.ltx23_release_vae_traces} +\title{Release all cached decode traces} +\usage{ +.ltx23_release_vae_traces() +} +\value{ +Invisibly, NULL. +} +\description{ +Traces hold references to the weight tensors they captured; drop +them when a component leaves the GPU so its memory actually frees. +} +\keyword{internal} diff --git a/man/dot-ltx23_staged_offload.Rd b/man/dot-ltx23_staged_offload.Rd new file mode 100644 index 0000000..9143965 --- /dev/null +++ b/man/dot-ltx23_staged_offload.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_staged_offload} +\alias{.ltx23_staged_offload} +\title{Return a pinned component to the CPU} +\usage{ +.ltx23_staged_offload(staging) +} +\description{ +Weights are immutable during inference, so the pinned host copies +are still current: offload is a pointer swap, no transfer. +} +\keyword{internal} diff --git a/man/dot-ltx23_staged_onload.Rd b/man/dot-ltx23_staged_onload.Rd new file mode 100644 index 0000000..222e554 --- /dev/null +++ b/man/dot-ltx23_staged_onload.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_staged_onload} +\alias{.ltx23_staged_onload} +\title{Move a pinned component onto the compute device} +\usage{ +.ltx23_staged_onload(staging, device) +} +\description{ +Non-blocking copies from pinned memory share the default stream, +so later kernels are ordered after them; no explicit sync needed. +} +\keyword{internal} diff --git a/man/dot-ltx23_traced_call.Rd b/man/dot-ltx23_traced_call.Rd new file mode 100644 index 0000000..35f59c7 --- /dev/null +++ b/man/dot-ltx23_traced_call.Rd @@ -0,0 +1,25 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.ltx23_traced_call} +\alias{.ltx23_traced_call} +\title{Run a module forward through a shape-specialized trace} +\usage{ +.ltx23_traced_call(module, x, forward = NULL, tag = "") +} +\arguments{ +\item{module}{The nn_module (identity for the cache key; also the +default callable).} + +\item{x}{Input tensor.} + +\item{forward}{Optional closure wrapping the call (for extra fixed +arguments like \code{causal}); must be pure in \code{x}.} + +\item{tag}{Character. Distinguishes call variants of one module.} +} +\value{ +The forward result. +} +\description{ +Run a module forward through a shape-specialized trace +} +\keyword{internal} diff --git a/man/download_component.Rd b/man/download_component.Rd index 684b0f5..6353746 100644 --- a/man/download_component.Rd +++ b/man/download_component.Rd @@ -3,13 +3,8 @@ \alias{download_component} \title{Download a single TorchScript model component} \usage{ -download_component( - model_name = "sd21", - component, - device = "cpu", - overwrite = FALSE, - show_progress = TRUE -) +download_component(model_name = "sd21", component, device = "cpu", + overwrite = FALSE, show_progress = TRUE) } \arguments{ \item{model_name}{Character string, the name of the model (e.g., \code{"sd21"}).} diff --git a/man/download_ltx2.Rd b/man/download_ltx2.Rd new file mode 100644 index 0000000..bddfd40 --- /dev/null +++ b/man/download_ltx2.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_ltx2} +\alias{download_ltx2} +\title{Download the LTX-2.3 checkpoint and build the fp8 artifact} +\usage{ +download_ltx2(quantize = TRUE, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), + text_encoder = TRUE, verbose = TRUE) +} +\arguments{ +\item{quantize}{Logical. Build the fp8 artifact after downloading.} + +\item{output_dir}{Directory for the fp8 artifact.} + +\item{text_encoder}{Logical. Also fetch the Gemma3 text encoder and +tokenizer (~25 GB, shared with LTX-2.0; from the Lightricks/LTX-2 +repo).} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, a list with \code{checkpoint} (source path or NULL), + \code{fp8_dir}, and \code{text_encoder_dir}. +} +\description{ +Skips work that is already done: a valid fp8 manifest short-circuits +everything; a cached 46 GB source skips the download. The source file +may be deleted after quantization (it is never removed automatically). +} diff --git a/man/download_ltx23.Rd b/man/download_ltx23.Rd new file mode 100644 index 0000000..2f6f47d --- /dev/null +++ b/man/download_ltx23.Rd @@ -0,0 +1,10 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_ltx23} +\alias{download_ltx23} +\title{Download and Prepare LTX-2.3 Model Weights} +\description{ +Downloads the LTX-2.3 distilled checkpoint (46 GB, LTX-2 Community +License) and the Gemma3 text encoder from HuggingFace with an explicit +consent prompt, then quantizes the transformer to the local fp8 +artifact (~26 GB) used by the GPU-poor pipeline. +} diff --git a/man/download_model.Rd b/man/download_model.Rd index abe915b..1dc4108 100644 --- a/man/download_model.Rd +++ b/man/download_model.Rd @@ -3,14 +3,10 @@ \alias{download_model} \title{Download TorchScript model files for Stable Diffusion} \usage{ -download_model( - model_name = "sd21", - devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), - unet_dtype_str = NULL, - overwrite = FALSE, - show_progress = TRUE, - download_models = FALSE -) +download_model(model_name = "sd21", + devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), + unet_dtype_str = NULL, overwrite = FALSE, show_progress = TRUE, + download_models = FALSE) } \arguments{ \item{model_name}{Name of the model (e.g., "sd21" for stable-diffusion-2-1)} @@ -31,8 +27,11 @@ A named list of full file paths, keyed by component name. \description{ Downloads the required model files (e.g., UNet, decoder, text encoder) for a given Stable Diffusion model using \code{hfhub::hub_download()}. +} +\details{ Files are cached by hfhub (typically \code{~/.cache/huggingface/hub/}). Legacy files in the old \code{R_user_dir()} location are also recognized. + } \examples{ \dontrun{ diff --git a/man/encode_bpe.Rd b/man/encode_bpe.Rd index f0ad760..d99ea9e 100644 --- a/man/encode_bpe.Rd +++ b/man/encode_bpe.Rd @@ -3,15 +3,8 @@ \alias{encode_bpe} \title{Encode text to token IDs} \usage{ -encode_bpe( - tokenizer, - text, - add_special_tokens = TRUE, - max_length = NULL, - padding = "none", - truncation = FALSE, - return_tensors = "list" -) +encode_bpe(tokenizer, text, add_special_tokens = TRUE, max_length = NULL, + padding = "none", truncation = FALSE, return_tensors = "list") } \arguments{ \item{tokenizer}{A bpe_tokenizer object.} diff --git a/man/encode_text_ltx2.Rd b/man/encode_text_ltx2.Rd deleted file mode 100644 index 3035579..0000000 --- a/man/encode_text_ltx2.Rd +++ /dev/null @@ -1,52 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{encode_text_ltx2} -\alias{encode_text_ltx2} -\title{Encode Text for LTX2} -\usage{ -encode_text_ltx2( - prompt, - backend = "random", - model_path = NULL, - tokenizer_path = NULL, - text_encoder = NULL, - embeddings_file = NULL, - api_url = NULL, - max_sequence_length = 1024L, - caption_channels = 3840L, - device = "cpu", - dtype = torch::torch_float32() -) -} -\arguments{ -\item{prompt}{Character vector of prompts.} - -\item{backend}{Character. Backend to use ("gemma3", "precomputed", "api", "random").} - -\item{model_path}{Character. Path to Gemma3 model directory (for "gemma3" backend).} - -\item{tokenizer_path}{Character. Path to tokenizer (for "gemma3" backend, defaults to model_path).} - -\item{text_encoder}{Pre-loaded Gemma3 text encoder module (for "gemma3" backend).} - -\item{embeddings_file}{Character. Path to pre-computed embeddings (for "precomputed" backend).} - -\item{api_url}{Character. URL of text encoding API (for "api" backend).} - -\item{max_sequence_length}{Integer. Maximum sequence length (default 1024).} - -\item{caption_channels}{Integer. Caption embedding dimension (default 3840).} - -\item{device}{Character. Device for tensors.} - -\item{dtype}{torch_dtype. Data type for tensors.} -} -\value{ -List with prompt_embeds and prompt_attention_mask tensors. -} -\description{ -Encodes text prompts for LTX2 video generation. Supports multiple backends: -- "gemma3": Native R torch Gemma3 text encoder -- "precomputed": Load pre-computed embeddings from file -- "api": Call an HTTP API for text encoding -- "random": Generate random embeddings (for testing only) -} diff --git a/man/encode_with_gemma3.Rd b/man/encode_with_gemma3.Rd index 1bb2e16..bd67717 100644 --- a/man/encode_with_gemma3.Rd +++ b/man/encode_with_gemma3.Rd @@ -3,16 +3,9 @@ \alias{encode_with_gemma3} \title{Encode text with Gemma3 for LTX-2} \usage{ -encode_with_gemma3( - prompts, - model = NULL, - tokenizer = NULL, - max_sequence_length = 1024L, - scale_factor = 8, - device = "cuda", - dtype = "float16", - verbose = TRUE -) +encode_with_gemma3(prompts, model = NULL, tokenizer = NULL, + max_sequence_length = 1024L, device = "cuda", + dtype = "float16", verbose = TRUE) } \arguments{ \item{prompts}{Character vector of prompts.} @@ -23,8 +16,6 @@ encode_with_gemma3( \item{max_sequence_length}{Integer. Maximum sequence length.} -\item{scale_factor}{Numeric. Scale factor for packing (default 8).} - \item{device}{Character. Device for computation.} \item{dtype}{Character. Data type.} @@ -32,9 +23,13 @@ encode_with_gemma3( \item{verbose}{Logical. Print progress.} } \value{ -List with prompt_embeds and prompt_attention_mask. +List with prompt_embeds (raw stacked hidden states, + shape \code{[batch, seq_len, hidden_size, num_layers + 1]}) and + prompt_attention_mask. } \description{ Full pipeline for encoding text prompts using Gemma3 text encoder. -Returns packed embeddings ready for LTX-2 connectors. +Returns the raw stacked per-layer hidden states (embedding layer plus all +transformer layers) for downstream connector modules, which handle +normalization and projection themselves. } diff --git a/man/feed_forward.Rd b/man/feed_forward.Rd deleted file mode 100644 index 668cb05..0000000 --- a/man/feed_forward.Rd +++ /dev/null @@ -1,19 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{feed_forward} -\alias{feed_forward} -\title{FeedForward module} -\usage{ -feed_forward( - dim, - dim_out = NULL, - mult = 4L, - dropout = 0, - activation_fn = "gelu-approximate", - inner_dim = NULL, - bias = TRUE -) -} -\description{ -FeedForward module -} -\keyword{internal} diff --git a/man/flowmatch_calculate_shift.Rd b/man/flowmatch_calculate_shift.Rd index 7433111..003673a 100644 --- a/man/flowmatch_calculate_shift.Rd +++ b/man/flowmatch_calculate_shift.Rd @@ -3,13 +3,8 @@ \alias{flowmatch_calculate_shift} \title{Calculate shift for dynamic shifting} \usage{ -flowmatch_calculate_shift( - seq_len, - base_seq_len = 256L, - max_seq_len = 4096L, - base_shift = 0.5, - max_shift = 1.15 -) +flowmatch_calculate_shift(seq_len, base_seq_len = 256L, max_seq_len = 4096L, + base_shift = 0.5, max_shift = 1.15) } \arguments{ \item{seq_len}{Integer. The sequence length (num_patches).} diff --git a/man/flowmatch_scheduler_create.Rd b/man/flowmatch_scheduler_create.Rd index e51e60d..0a2d899 100644 --- a/man/flowmatch_scheduler_create.Rd +++ b/man/flowmatch_scheduler_create.Rd @@ -3,18 +3,12 @@ \alias{flowmatch_scheduler_create} \title{Create a FlowMatch Euler Discrete Scheduler} \usage{ -flowmatch_scheduler_create( - num_train_timesteps = 1000L, - shift = 1, - use_dynamic_shifting = FALSE, - base_shift = 0.5, - max_shift = 1.15, - base_seq_len = 256L, - max_seq_len = 4096L, - invert_sigmas = FALSE, - shift_terminal = NULL, - time_shift_type = c("exponential", "linear") -) +flowmatch_scheduler_create(num_train_timesteps = 1000L, shift = 1, + use_dynamic_shifting = FALSE, base_shift = 0.5, + max_shift = 1.15, base_seq_len = 256L, + max_seq_len = 4096L, invert_sigmas = FALSE, + shift_terminal = NULL, + time_shift_type = c("exponential", "linear")) } \arguments{ \item{num_train_timesteps}{Integer. The number of diffusion steps used to diff --git a/man/flowmatch_scheduler_step.Rd b/man/flowmatch_scheduler_step.Rd index 4db2184..63f3eaa 100644 --- a/man/flowmatch_scheduler_step.Rd +++ b/man/flowmatch_scheduler_step.Rd @@ -3,13 +3,8 @@ \alias{flowmatch_scheduler_step} \title{Perform a FlowMatch scheduler step} \usage{ -flowmatch_scheduler_step( - model_output, - timestep, - sample, - schedule, - generator = NULL -) +flowmatch_scheduler_step(model_output, timestep, sample, schedule, + generator = NULL) } \arguments{ \item{model_output}{torch tensor. The output from the diffusion model diff --git a/man/flowmatch_set_timesteps.Rd b/man/flowmatch_set_timesteps.Rd index 89881cb..78c2f7f 100644 --- a/man/flowmatch_set_timesteps.Rd +++ b/man/flowmatch_set_timesteps.Rd @@ -3,14 +3,8 @@ \alias{flowmatch_set_timesteps} \title{Set timesteps for inference} \usage{ -flowmatch_set_timesteps( - schedule, - num_inference_steps = 50L, - device = "cpu", - mu = NULL, - sigmas = NULL, - timesteps = NULL -) +flowmatch_set_timesteps(schedule, num_inference_steps = 50L, device = "cpu", + mu = NULL, sigmas = NULL, timesteps = NULL) } \arguments{ \item{schedule}{List. The FlowMatch scheduler object.} diff --git a/man/fp8_ltx23.Rd b/man/fp8_ltx23.Rd new file mode 100644 index 0000000..bc77e7f --- /dev/null +++ b/man/fp8_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{fp8_ltx23} +\alias{fp8_ltx23} +\title{FP8 Weight Storage for the LTX-2.3 Transformer} +\description{ +GPU-poor weight handling: the large attention/FFN linears of the DiT +are stored as float8_e4m3fn with per-tensor scales (the official LTX +quantization policy), kept CPU-resident (optionally pinned), and +dequantized on the compute device inside each forward. Everything +else (norms, embeddings, modulation tables, biases) stays bfloat16. +Requires a safetensors build with F8 dtype support. +} diff --git a/man/gelu_activation.Rd b/man/gelu_activation.Rd deleted file mode 100644 index 7ff959a..0000000 --- a/man/gelu_activation.Rd +++ /dev/null @@ -1,11 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{gelu_activation} -\alias{gelu_activation} -\title{GELU activation with optional approximation} -\usage{ -gelu_activation(dim_in, dim_out, approximate = "none", bias = TRUE) -} -\description{ -GELU activation with optional approximation -} -\keyword{internal} diff --git a/man/gemma3_rotary_embedding.Rd b/man/gemma3_rotary_embedding.Rd index f840809..c80dea0 100644 --- a/man/gemma3_rotary_embedding.Rd +++ b/man/gemma3_rotary_embedding.Rd @@ -3,12 +3,8 @@ \alias{gemma3_rotary_embedding} \title{Gemma3 Rotary Position Embeddings} \usage{ -gemma3_rotary_embedding( - dim, - max_position_embeddings = 8192L, - base = 10000, - scaling_factor = 1 -) +gemma3_rotary_embedding(dim, max_position_embeddings = 8192L, base = 10000, + scaling_factor = 1) } \arguments{ \item{dim}{Integer. Head dimension.} diff --git a/man/get_timestep_embedding.Rd b/man/get_timestep_embedding.Rd deleted file mode 100644 index ce6b5e1..0000000 --- a/man/get_timestep_embedding.Rd +++ /dev/null @@ -1,18 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{get_timestep_embedding} -\alias{get_timestep_embedding} -\title{Get timestep embedding (sinusoidal)} -\usage{ -get_timestep_embedding( - timesteps, - embedding_dim, - flip_sin_to_cos = FALSE, - downscale_freq_shift = 1, - scale = 1, - max_period = 10000 -) -} -\description{ -Get timestep embedding (sinusoidal) -} -\keyword{internal} diff --git a/man/gpu_poor.Rd b/man/gpu_poor.Rd deleted file mode 100644 index 9d6ee20..0000000 --- a/man/gpu_poor.Rd +++ /dev/null @@ -1,8 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{gpu_poor} -\alias{gpu_poor} -\title{GPU-Poor Memory Management for LTX-2} -\description{ -wan2GP-style memory optimizations for running LTX-2 video generation -on limited VRAM (6-16GB). -} diff --git a/man/img2img.Rd b/man/img2img.Rd index f25797d..ad8cdbc 100644 --- a/man/img2img.Rd +++ b/man/img2img.Rd @@ -3,29 +3,13 @@ \alias{img2img} \title{Image-to-Image Generation with Stable Diffusion} \usage{ -img2img( - input_image, - prompt, - negative_prompt = NULL, - img_dim = 512, - model_name = c("sd21", "sdxl"), - pipeline = NULL, - devices = "auto", - unet_dtype_str = "float16", - download_models = FALSE, - scheduler = "ddim", - num_inference_steps = 50, - strength = 0.8, - guidance_scale = 7.5, - seed = NULL, - save_file = TRUE, - filename = NULL, - metadata_path = NULL, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, - ... -) +img2img(input_image, prompt, negative_prompt = NULL, img_dim = 512, + model_name = c("sd21", "sdxl"), pipeline = NULL, devices = "auto", + unet_dtype_str = "float16", download_models = FALSE, + scheduler = "ddim", num_inference_steps = 50, strength = 0.8, + guidance_scale = 7.5, seed = NULL, save_file = TRUE, filename = NULL, + metadata_path = NULL, use_native_decoder = FALSE, + use_native_text_encoder = FALSE, use_native_unet = FALSE, ...) } \arguments{ \item{input_image}{Path to the input image or a tensor representing the image.} diff --git a/man/int4_linear.Rd b/man/int4_linear.Rd deleted file mode 100644 index 5749828..0000000 --- a/man/int4_linear.Rd +++ /dev/null @@ -1,43 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{int4_linear} -\alias{int4_linear} -\title{INT4 Linear Layer} -\usage{ -int4_linear( - in_features, - out_features, - bias = TRUE, - device = "cuda", - dtype = torch::torch_float16() -) -} -\arguments{ -\item{in_features}{Integer. Size of each input sample.} - -\item{out_features}{Integer. Size of each output sample.} - -\item{bias}{Logical. If TRUE, adds a learnable bias (stored in float16).} - -\item{device}{Character. Device for the layer.} - -\item{dtype}{torch_dtype. Data type for dequantized operations.} -} -\value{ -nn_module with INT4 weight storage. -} -\description{ -A linear layer that stores weights in INT4 format and dequantizes on-the-fly -during forward pass. This enables running large models on limited VRAM by -keeping weights compressed on GPU. -} -\details{ -The layer stores: -- `weight_packed`: uint8 tensor with packed INT4 values -- `weight_scales`: float32 tensor with per-block scales -- `weight_shape`: original weight shape -- `bias`: optional float16 bias - -During forward(), weights are dequantized to the target dtype, the matmul -is performed, and the dequantized tensor is freed. This allows ~40GB models -to run with ~10GB of VRAM for weights. -} diff --git a/man/int4_linear_from_quantized.Rd b/man/int4_linear_from_quantized.Rd deleted file mode 100644 index a30ba7e..0000000 --- a/man/int4_linear_from_quantized.Rd +++ /dev/null @@ -1,30 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{int4_linear_from_quantized} -\alias{int4_linear_from_quantized} -\title{Create INT4 Linear from Pre-quantized Weights} -\usage{ -int4_linear_from_quantized( - q_weight, - q_bias = NULL, - bias_tensor = NULL, - device = "cuda", - dtype = torch::torch_float16() -) -} -\arguments{ -\item{q_weight}{List with packed, scales, orig_shape from load_int4_weights().} - -\item{q_bias}{Optional. Quantized bias (or NULL for no bias).} - -\item{bias_tensor}{Optional. Float tensor for bias (if not quantized).} - -\item{device}{Character. Target device.} - -\item{dtype}{torch_dtype. Target dtype for operations.} -} -\value{ -int4_linear module with loaded weights. -} -\description{ -Creates an INT4 linear layer from pre-quantized weight data. -} diff --git a/man/jit_ltx23.Rd b/man/jit_ltx23.Rd new file mode 100644 index 0000000..4055155 --- /dev/null +++ b/man/jit_ltx23.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{jit_ltx23} +\alias{jit_ltx23} +\title{LTX-2.3 JIT Block Stack} +\description{ +TorchScript compilation of the 48-block NF4 transformer step (cf. +the torch skill's JIT-decode pattern proven in whisper and +chatterbox). Eager execution crosses R -> lantern per op (~190 us +each) and leaves every intermediate as an R tensor handle that only +dies at gc(); at high resolution that forces a per-block gc() +costing the vast majority of step time. Compiled, the whole block +stack is one crossing: intermediates are freed eagerly by libtorch, +no R garbage accumulates, no per-block gc is needed, and attention +runs through the fused \code{scaled_dot_product_attention} kernel +instead of a materialized score matrix. +} +\details{ +Weights are passed per call as a flat \code{List[Tensor]} (borrowed +by reference, no copies) with a fixed per-block layout; the packer +and the TorchScript indices must stay in lockstep (parity-tested). + +} diff --git a/man/jit_vae_ltx23.Rd b/man/jit_vae_ltx23.Rd new file mode 100644 index 0000000..82f5f0d --- /dev/null +++ b/man/jit_vae_ltx23.Rd @@ -0,0 +1,33 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{jit_vae_ltx23} +\alias{jit_vae_ltx23} +\title{JIT-Traced Decode for the LTX-2.3 VAEs and Vocoder} +\description{ +The video/audio decoders and the vocoder are static feed-forward +graphs, so \code{torch::jit_trace} converts them wholesale: one +R-to-libtorch crossing per forward, intermediates freed eagerly by +libtorch instead of accumulating as R handles until gc. Traces are +shape-specialized (runtime sizes bake into the graph as constants; +a mismatched input errors), so they are cached per instance, input +shape, dtype, device, and call tag, and re-traced on a miss. +} +\details{ +A trace captures the module's weight tensors, which would pin them +on the GPU across phase offloads; the pipeline releases all traces +whenever a component offloads (\code{.ltx23_release_vae_traces}). + +Tracing hazard on this lantern build: if the allocator callback +runs R's gc \emph{during} trace recording (memory pressure), the +recorded graph can capture garbage argument values (observed as +corrupted \code{narrow} starts on the full-size decoder; verified +5/5 clean once gc cannot fire mid-trace). Defenses, in order: a +gc + cache flush right before each trace so pressure starts near +zero, \code{tryCatch} around trace and replay, and a one-time +validation of every fresh trace against the eager output — any +mismatch permanently blacklists that shape and runs eager. With +those in place the traced path cannot corrupt output — but per +render it measured slower than eager (traces are released on phase +offload, so every render re-pays trace + validation), so it stays +opt-in: \code{options(diffuseR.jit_vae = TRUE)}. + +} diff --git a/man/linear_to_int4.Rd b/man/linear_to_int4.Rd deleted file mode 100644 index 183ad5d..0000000 --- a/man/linear_to_int4.Rd +++ /dev/null @@ -1,20 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{linear_to_int4} -\alias{linear_to_int4} -\title{Create INT4 Linear from Standard Linear} -\usage{ -linear_to_int4(linear_module, device = "cuda", dtype = torch::torch_float16()) -} -\arguments{ -\item{linear_module}{nn_linear module to convert.} - -\item{device}{Character. Target device for INT4 weights.} - -\item{dtype}{torch_dtype. Target dtype for dequantized operations.} -} -\value{ -int4_linear module with quantized weights. -} -\description{ -Converts a standard nn_linear layer to an INT4 linear layer. -} diff --git a/man/load_gemma3_text_encoder.Rd b/man/load_gemma3_text_encoder.Rd index e7e07e6..bffba3d 100644 --- a/man/load_gemma3_text_encoder.Rd +++ b/man/load_gemma3_text_encoder.Rd @@ -3,12 +3,8 @@ \alias{load_gemma3_text_encoder} \title{Load Gemma3 Text Model from safetensors} \usage{ -load_gemma3_text_encoder( - model_path, - device = "cpu", - dtype = "float16", - verbose = TRUE -) +load_gemma3_text_encoder(model_path, device = "cpu", dtype = "float16", + verbose = TRUE) } \arguments{ \item{model_path}{Character. Path to directory containing model files.} diff --git a/man/load_int4_into_model.Rd b/man/load_int4_into_model.Rd deleted file mode 100644 index ebfecfa..0000000 --- a/man/load_int4_into_model.Rd +++ /dev/null @@ -1,41 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_int4_into_model} -\alias{load_int4_into_model} -\title{Load INT4 Weights into Model} -\usage{ -load_int4_into_model( - model, - int4_weights, - device = "cuda", - dtype = torch::torch_float16(), - verbose = TRUE -) -} -\arguments{ -\item{model}{nn_module. The model to convert.} - -\item{int4_weights}{List of quantized weights from `load_int4_weights()`.} - -\item{device}{Character. Target device for INT4 weights.} - -\item{dtype}{torch_dtype. Target dtype for dequantized operations.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -The model with linear layers replaced by INT4 versions. -} -\description{ -Replaces linear layers in a model with INT4 versions and loads quantized weights. -This is the main entry point for running large models with INT4 quantization. -} -\details{ -This function: -1. Identifies linear layers by matching parameter names ending in ".weight" -2. Creates INT4Linear layers with matching dimensions -3. Loads quantized weights and biases -4. Replaces the original layers in the model - -The INT4 weights stay compressed on GPU (~10GB for a 40GB model). -During forward(), each layer dequantizes on-the-fly, keeping memory usage low. -} diff --git a/man/load_int4_weights.Rd b/man/load_int4_weights.Rd deleted file mode 100644 index d74ae3a..0000000 --- a/man/load_int4_weights.Rd +++ /dev/null @@ -1,27 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_int4_weights} -\alias{load_int4_weights} -\title{Load INT4 Quantized Weights} -\usage{ -load_int4_weights(path, verbose = TRUE) -} -\arguments{ -\item{path}{Character. Path to safetensors file or base path for sharded files. -For sharded files, pass the base path (e.g., "model_int4.safetensors") and -the function will find all shards matching the pattern.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -List of quantized parameter structures ready for `dequantize_int4()`. -} -\description{ -Loads INT4 quantized weights from safetensors file(s). -} -\examples{ -\dontrun{ -q <- load_int4_weights("model_int4.safetensors") -# Dequantize specific parameter on GPU -weight <- dequantize_int4(q[["linear.weight"]], device = "cuda") -} -} diff --git a/man/load_int4_weights_into_model.Rd b/man/load_int4_weights_into_model.Rd deleted file mode 100644 index c274f75..0000000 --- a/man/load_int4_weights_into_model.Rd +++ /dev/null @@ -1,21 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_int4_weights_into_model} -\alias{load_int4_weights_into_model} -\title{Load INT4 Weights into INT4 Model} -\usage{ -load_int4_weights_into_model(model, int4_weights, verbose = TRUE) -} -\arguments{ -\item{model}{nn_module created with INT4 layers.} - -\item{int4_weights}{List from `load_int4_weights()`.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -Model with INT4 weights loaded (invisibly). -} -\description{ -Loads pre-quantized INT4 weights into a model created with `make_linear()` -when `diffuseR.use_int4 = TRUE`. -} diff --git a/man/load_ltx2_connector_weights.Rd b/man/load_ltx2_connector_weights.Rd deleted file mode 100644 index eb7dc35..0000000 --- a/man/load_ltx2_connector_weights.Rd +++ /dev/null @@ -1,18 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_connector_weights} -\alias{load_ltx2_connector_weights} -\title{Load weights into LTX2 connectors module} -\usage{ -load_ltx2_connector_weights(connectors, weights, verbose = TRUE) -} -\arguments{ -\item{connectors}{LTX2 connectors module} - -\item{weights}{Named list of weight tensors} - -\item{verbose}{Print progress} -} -\description{ -Load weights into LTX2 connectors module -} -\keyword{internal} diff --git a/man/load_ltx2_connectors.Rd b/man/load_ltx2_connectors.Rd deleted file mode 100644 index e4f6c6b..0000000 --- a/man/load_ltx2_connectors.Rd +++ /dev/null @@ -1,30 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_connectors} -\alias{load_ltx2_connectors} -\title{Load LTX2 Text Connectors from safetensors} -\usage{ -load_ltx2_connectors( - weights_path, - config_path = NULL, - device = "cpu", - dtype = "float32", - verbose = TRUE -) -} -\arguments{ -\item{weights_path}{Character. Path to safetensors file.} - -\item{config_path}{Character. Optional path to config.json.} - -\item{device}{Character. Device to load weights to. Default: "cpu"} - -\item{dtype}{Character. Data type ("float32", "float16"). Default: "float32"} - -\item{verbose}{Logical. Print loading progress. Default: TRUE} -} -\value{ -Initialized ltx2_text_connectors module -} -\description{ -Load pre-trained LTX2 connector weights from HuggingFace safetensors file. -} diff --git a/man/load_ltx2_transformer.Rd b/man/load_ltx2_transformer.Rd deleted file mode 100644 index 08ebd2d..0000000 --- a/man/load_ltx2_transformer.Rd +++ /dev/null @@ -1,31 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_transformer} -\alias{load_ltx2_transformer} -\title{Load LTX2 DiT Transformer from safetensors} -\usage{ -load_ltx2_transformer( - weights_dir, - config_path = NULL, - device = "cpu", - dtype = "float16", - verbose = TRUE -) -} -\arguments{ -\item{weights_dir}{Character. Directory containing safetensors files.} - -\item{config_path}{Character. Optional path to config.json. If NULL, uses default config.} - -\item{device}{Character. Device to load weights to. Default: "cpu"} - -\item{dtype}{Character. Data type ("float32", "float16", "bfloat16"). Default: "float16"} - -\item{verbose}{Logical. Print loading progress. Default: TRUE} -} -\value{ -Initialized ltx2_video_transformer3d module -} -\description{ -Load pre-trained LTX2 transformer weights from HuggingFace safetensors files. -Supports both single file and sharded multi-file loading. -} diff --git a/man/load_ltx2_transformer_sharded.Rd b/man/load_ltx2_transformer_sharded.Rd deleted file mode 100644 index b82e7d0..0000000 --- a/man/load_ltx2_transformer_sharded.Rd +++ /dev/null @@ -1,25 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_transformer_sharded} -\alias{load_ltx2_transformer_sharded} -\title{Load sharded transformer weights} -\usage{ -load_ltx2_transformer_sharded( - transformer, - weights_dir, - index_path, - verbose = TRUE -) -} -\arguments{ -\item{transformer}{LTX2 transformer module} - -\item{weights_dir}{Directory containing sharded safetensors} - -\item{index_path}{Path to index.json} - -\item{verbose}{Print progress} -} -\description{ -Load sharded transformer weights -} -\keyword{internal} diff --git a/man/load_ltx2_transformer_weights.Rd b/man/load_ltx2_transformer_weights.Rd deleted file mode 100644 index a03aec2..0000000 --- a/man/load_ltx2_transformer_weights.Rd +++ /dev/null @@ -1,18 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_transformer_weights} -\alias{load_ltx2_transformer_weights} -\title{Load weights into LTX2 transformer module} -\usage{ -load_ltx2_transformer_weights(transformer, weights, verbose = TRUE) -} -\arguments{ -\item{transformer}{LTX2 transformer module} - -\item{weights}{Named list of weight tensors} - -\item{verbose}{Print progress} -} -\description{ -Load weights into LTX2 transformer module -} -\keyword{internal} diff --git a/man/load_ltx2_vae.Rd b/man/load_ltx2_vae.Rd deleted file mode 100644 index 6228561..0000000 --- a/man/load_ltx2_vae.Rd +++ /dev/null @@ -1,31 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_vae} -\alias{load_ltx2_vae} -\title{Load LTX2 Video VAE from safetensors} -\usage{ -load_ltx2_vae( - weights_path, - config_path = NULL, - device = "cpu", - dtype = "float32", - verbose = TRUE -) -} -\arguments{ -\item{weights_path}{Character. Path to safetensors file or directory containing weights.} - -\item{config_path}{Character. Optional path to config.json. If NULL and weights_path -is a directory, looks for config.json in that directory. Otherwise uses default config.} - -\item{device}{Character. Device to load weights to. Default: "cpu"} - -\item{dtype}{Character or torch dtype. Data type. Default: "float32"} - -\item{verbose}{Logical. Print loading progress. Default: TRUE} -} -\value{ -Initialized ltx2_video_vae module -} -\description{ -Load pre-trained LTX2 VAE weights from a HuggingFace safetensors file. -} diff --git a/man/load_ltx2_vae_weights.Rd b/man/load_ltx2_vae_weights.Rd deleted file mode 100644 index 037c7a8..0000000 --- a/man/load_ltx2_vae_weights.Rd +++ /dev/null @@ -1,21 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{load_ltx2_vae_weights} -\alias{load_ltx2_vae_weights} -\title{Load weights into LTX2 VAE module} -\usage{ -load_ltx2_vae_weights(vae, weights, verbose = TRUE) -} -\arguments{ -\item{vae}{LTX2 VAE module} - -\item{weights}{Named list of weight tensors from safetensors} - -\item{verbose}{Logical. Print loading progress} -} -\value{ -The VAE with loaded weights (invisibly) -} -\description{ -Maps HuggingFace safetensors parameter names to R module parameters. -} -\keyword{internal} diff --git a/man/load_model_component.Rd b/man/load_model_component.Rd index edc91b6..d8efbcb 100644 --- a/man/load_model_component.Rd +++ b/man/load_model_component.Rd @@ -3,14 +3,8 @@ \alias{load_model_component} \title{Load a specific component of a diffusion model} \usage{ -load_model_component( - component, - model_name = "sd21", - device = "cpu", - unet_dtype_str = NULL, - download = TRUE, - use_native = FALSE -) +load_model_component(component, model_name = "sd21", device = "cpu", + unet_dtype_str = NULL, download = TRUE, use_native = FALSE) } \arguments{ \item{component}{Character string, the component to load: "unet", "decoder", or "text_encoder".} diff --git a/man/load_pipeline.Rd b/man/load_pipeline.Rd index 783d704..b805ca7 100644 --- a/man/load_pipeline.Rd +++ b/man/load_pipeline.Rd @@ -3,16 +3,9 @@ \alias{load_pipeline} \title{Load a diffusion model pipeline} \usage{ -load_pipeline( - model_name, - m2d, - i2i = FALSE, - unet_dtype_str, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, - ... -) +load_pipeline(model_name, m2d, i2i = FALSE, unet_dtype_str, + use_native_decoder = FALSE, use_native_text_encoder = FALSE, + use_native_unet = FALSE, ...) } \arguments{ \item{model_name}{The name of the model to load.} diff --git a/man/ltx23_ada_layer_norm_single.Rd b/man/ltx23_ada_layer_norm_single.Rd new file mode 100644 index 0000000..bd49a10 --- /dev/null +++ b/man/ltx23_ada_layer_norm_single.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_ada_layer_norm_single} +\alias{ltx23_ada_layer_norm_single} +\title{Adaptive layer norm single (adaLN-single)} +\usage{ +ltx23_ada_layer_norm_single(embedding_dim, num_mod_params = 6L) +} +\arguments{ +\item{embedding_dim}{Integer. Model dimension.} + +\item{num_mod_params}{Integer. Number of modulation parameter vectors.} +} +\value{ +Module whose forward returns + \code{list(mod_params [N, num_mod_params * dim], embedded_timestep [N, dim])}. +} +\description{ +Embeds a timestep/sigma and projects it to a configurable number of +modulation parameter vectors. +} diff --git a/man/ltx23_adain_filter_latent.Rd b/man/ltx23_adain_filter_latent.Rd new file mode 100644 index 0000000..16b1388 --- /dev/null +++ b/man/ltx23_adain_filter_latent.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_adain_filter_latent} +\alias{ltx23_adain_filter_latent} +\title{Adaptive instance normalization between latent tensors} +\usage{ +ltx23_adain_filter_latent(latents, reference_latents, factor = 1) +} +\arguments{ +\item{latents}{Tensor [B, C, F, H, W].} + +\item{reference_latents}{Tensor with the target statistics.} + +\item{factor}{Numeric blend in [-10, 10]; 0 is identity.} +} +\value{ +Filtered latents. +} +\description{ +Matches each (batch, channel) slice's mean/std to the reference +latents, blended by \code{factor} (cf. diffusers +\code{LTX2LatentUpsamplePipeline.adain_filter_latent}). +} diff --git a/man/ltx23_antialias_act1d.Rd b/man/ltx23_antialias_act1d.Rd new file mode 100644 index 0000000..de5f9af --- /dev/null +++ b/man/ltx23_antialias_act1d.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_antialias_act1d} +\alias{ltx23_antialias_act1d} +\title{Anti-aliased activation} +\usage{ +ltx23_antialias_act1d(channels, ratio = 2L, kernel_size = 12L) +} +\arguments{ +\item{channels}{Integer. Channels for the SnakeBeta activation.} + +\item{ratio,kernel_size}{Integers. Resampling config.} +} +\description{ +Upsample 2x, apply the activation, downsample 2x. +} diff --git a/man/ltx23_apply_interleaved_rotary_emb.Rd b/man/ltx23_apply_interleaved_rotary_emb.Rd new file mode 100644 index 0000000..d191c6d --- /dev/null +++ b/man/ltx23_apply_interleaved_rotary_emb.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_apply_interleaved_rotary_emb} +\alias{ltx23_apply_interleaved_rotary_emb} +\title{Apply interleaved rotary embeddings} +\usage{ +ltx23_apply_interleaved_rotary_emb(x, freqs) +} +\arguments{ +\item{x}{Tensor of shape [B, S, C].} + +\item{freqs}{List of two tensors (cos, sin), each [B, S, C].} +} +\value{ +Tensor with the same shape and dtype as \code{x}. +} +\description{ +Rotates adjacent element pairs of the last dimension: +\code{out = x * cos + rotate_half(x) * sin} with pairs interleaved +(elements 1,2 form the first complex pair). +} diff --git a/man/ltx23_apply_split_rotary_emb.Rd b/man/ltx23_apply_split_rotary_emb.Rd new file mode 100644 index 0000000..a1d9f57 --- /dev/null +++ b/man/ltx23_apply_split_rotary_emb.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_apply_split_rotary_emb} +\alias{ltx23_apply_split_rotary_emb} +\title{Apply split rotary embeddings} +\usage{ +ltx23_apply_split_rotary_emb(x, freqs) +} +\arguments{ +\item{x}{Tensor of shape [B, H, T, D] (per-head layout), or [B, T, H*D] +which is reshaped per-head when \code{freqs} is 4D.} + +\item{freqs}{List of two tensors (cos, sin), each [B, H, T, D/2].} +} +\value{ +Tensor with the same shape and dtype as \code{x}. +} +\description{ +Rotates element pairs formed by splitting the last dimension in half: +element i pairs with element i + d/2. The cos/sin tensors carry half +the head dimension. +} diff --git a/man/ltx23_attention.Rd b/man/ltx23_attention.Rd new file mode 100644 index 0000000..36dea3b --- /dev/null +++ b/man/ltx23_attention.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_attention} +\alias{ltx23_attention} +\title{LTX-2 attention layer} +\usage{ +ltx23_attention(query_dim, heads = 8L, kv_heads = NULL, dim_head = 64L, + bias = TRUE, cross_attention_dim = NULL, out_bias = TRUE, + norm_eps = 1e-06, norm_elementwise_affine = TRUE, + rope_type = "split", apply_gated_attention = FALSE) +} +\arguments{ +\item{query_dim}{Integer. Query feature dimension.} + +\item{dim_head}{Integer. Per-head dimension.} + +\item{cross_attention_dim}{Integer or NULL. Key/value input dimension +(NULL for self-attention).} + +\item{norm_eps}{Numeric. RMS norm epsilon.} + +\item{norm_elementwise_affine}{Logical. RMS norms carry weights.} + +\item{rope_type}{"split" or "interleaved".} + +\item{apply_gated_attention}{Logical. Add per-head sigmoid output gates.} + +\item{heads,kv_heads}{Integers. Attention head counts.} + +\item{bias,out_bias}{Logicals. Projection biases.} +} +\description{ +Attention with RMS q/k norms across heads, optional per-head output +gating (LTX-2.3), separate query/key RoPE (for a2v/v2a cross +attention), and optional STG perturbation (skip attention, use the +value projection). +} diff --git a/man/ltx23_audio_causal_conv2d.Rd b/man/ltx23_audio_causal_conv2d.Rd new file mode 100644 index 0000000..0b81bc3 --- /dev/null +++ b/man/ltx23_audio_causal_conv2d.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_causal_conv2d} +\alias{ltx23_audio_causal_conv2d} +\title{Causal 2D convolution for audio spectrograms} +\usage{ +ltx23_audio_causal_conv2d(in_channels, out_channels, kernel_size = 3L, + stride = 1L, causality_axis = "height") +} +\arguments{ +\item{kernel_size}{Integer or length-2 vector.} + +\item{stride}{Integer.} + +\item{causality_axis}{"height", "width", "width-compatibility", or "none".} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +Pads asymmetrically along the causal axis ("height" = time frames for +LTX audio) before an unpadded Conv2d. +} diff --git a/man/ltx23_audio_decoder.Rd b/man/ltx23_audio_decoder.Rd new file mode 100644 index 0000000..d29537e --- /dev/null +++ b/man/ltx23_audio_decoder.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_decoder} +\alias{ltx23_audio_decoder} +\title{LTX-2.3 audio VAE decoder} +\usage{ +ltx23_audio_decoder(base_channels = 128L, output_channels = 2L, + num_res_blocks = 2L, latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), causality_axis = "height", + mel_bins = 64L) +} +\arguments{ +\item{base_channels}{Integer.} + +\item{output_channels}{Integer. Audio channels (2 = stereo).} + +\item{num_res_blocks}{Integer. Per-level ResNet count (a stage runs +\code{num_res_blocks + 1} blocks).} + +\item{latent_channels}{Integer.} + +\item{ch_mult}{Integer vector. Channel multipliers per level.} + +\item{causality_axis}{Character.} + +\item{mel_bins}{Integer. Output mel bins (crop/pad target).} +} +\description{ +Latents [B, 8, L, 16] -> mel spectrogram [B, 2, 4L - 3, 64]. +} diff --git a/man/ltx23_audio_downsample.Rd b/man/ltx23_audio_downsample.Rd new file mode 100644 index 0000000..e86a97a --- /dev/null +++ b/man/ltx23_audio_downsample.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_downsample} +\alias{ltx23_audio_downsample} +\title{LTX audio downsampler} +\usage{ +ltx23_audio_downsample(in_channels, causality_axis = "height") +} +\arguments{ +\item{in_channels}{Integer.} + +\item{causality_axis}{Character.} +} +\description{ +Causal zero-pad followed by a plain stride-2 conv (reference +LTX2AudioDownsample; note the conv is unwrapped, so its checkpoint +key is \code{downsample.conv.*}). +} diff --git a/man/ltx23_audio_encoder.Rd b/man/ltx23_audio_encoder.Rd new file mode 100644 index 0000000..ebeaab8 --- /dev/null +++ b/man/ltx23_audio_encoder.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_encoder} +\alias{ltx23_audio_encoder} +\title{LTX-2.3 audio VAE encoder} +\usage{ +ltx23_audio_encoder(base_channels = 128L, in_channels = 2L, + num_res_blocks = 2L, latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), causality_axis = "height") +} +\arguments{ +\item{in_channels}{Integer. Mel channels (2 = stereo).} + +\item{base_channels,num_res_blocks,latent_channels,ch_mult,causality_axis}{See \code{\link{ltx23_audio_decoder}}.} +} +\description{ +Mel spectrogram [B, 2, T, 64] -> latent distribution moments +[B, 2 * latent_channels, ceil(T/4), 16]. Structure mirrors the +decoder: causal convs, parameterless pixel norms, ResNet stages with +stride-2 downsampling between levels (reference LTX2AudioEncoder). +} diff --git a/man/ltx23_audio_mel_frontend.Rd b/man/ltx23_audio_mel_frontend.Rd new file mode 100644 index 0000000..d9572ab --- /dev/null +++ b/man/ltx23_audio_mel_frontend.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_mel_frontend} +\alias{ltx23_audio_mel_frontend} +\title{Build the 16 kHz log-mel frontend for audio conditioning} +\usage{ +ltx23_audio_mel_frontend(filter_length = 1024L, hop_length = 160L, + n_mels = 64L, sample_rate = 16000L, fmin = 0, + fmax = 8000) +} +\arguments{ +\item{filter_length,hop_length,n_mels,sample_rate,fmin,fmax}{The +checkpoint preprocessing parameters (defaults are LTX-2.3's).} +} +\value{ +An \code{ltx23_mel_stft} module. +} +\description{ +An \code{\link{ltx23_mel_stft}} whose STFT and mel bases are +constructed (not checkpoint-loaded) with the audio VAE's +preprocessing spec. +} diff --git a/man/ltx23_audio_resnet_block.Rd b/man/ltx23_audio_resnet_block.Rd new file mode 100644 index 0000000..9b3c8a9 --- /dev/null +++ b/man/ltx23_audio_resnet_block.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_resnet_block} +\alias{ltx23_audio_resnet_block} +\title{LTX audio ResNet block} +\usage{ +ltx23_audio_resnet_block(in_channels, out_channels = NULL, + causality_axis = "height") +} +\arguments{ +\item{causality_axis}{Character.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +PixelNorm -> SiLU -> causal conv, twice, with a 1x1 causal conv +shortcut (\code{nin_shortcut}) on channel change. +} diff --git a/man/ltx23_audio_upsample.Rd b/man/ltx23_audio_upsample.Rd new file mode 100644 index 0000000..8388f7e --- /dev/null +++ b/man/ltx23_audio_upsample.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_upsample} +\alias{ltx23_audio_upsample} +\title{LTX audio upsampler} +\usage{ +ltx23_audio_upsample(in_channels, causality_axis = "height") +} +\arguments{ +\item{in_channels}{Integer.} + +\item{causality_axis}{Character.} +} +\description{ +Nearest 2x interpolation, causal conv, then a crop of the first +element along the causal axis. +} diff --git a/man/ltx23_audio_vae.Rd b/man/ltx23_audio_vae.Rd new file mode 100644 index 0000000..bf4cb34 --- /dev/null +++ b/man/ltx23_audio_vae.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_audio_vae} +\alias{ltx23_audio_vae} +\title{LTX-2.3 audio VAE} +\usage{ +ltx23_audio_vae(base_channels = 128L, output_channels = 2L, + num_res_blocks = 2L, latent_channels = 8L, + ch_mult = c(1L, 2L, 4L), causality_axis = "height", + mel_bins = 64L, in_channels = 2L) +} +\arguments{ +\item{in_channels}{Integer. Mel input channels (2 = stereo).} + +\item{base_channels,output_channels,num_res_blocks,latent_channels,ch_mult,causality_axis,mel_bins}{See \code{\link{ltx23_audio_decoder}}.} +} +\description{ +Encoder + decoder plus the per-channel latent statistics loaded from +the checkpoint. Encoding is used for audio-conditioned generation +(lip sync); decoding for generated audio. +} diff --git a/man/ltx23_causal_conv3d.Rd b/man/ltx23_causal_conv3d.Rd new file mode 100644 index 0000000..e7d74a0 --- /dev/null +++ b/man/ltx23_causal_conv3d.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_causal_conv3d} +\alias{ltx23_causal_conv3d} +\title{Causal 3D convolution} +\usage{ +ltx23_causal_conv3d(in_channels, out_channels, kernel_size = 3L, stride = 1L, + spatial_padding_mode = "zeros") +} +\arguments{ +\item{kernel_size}{Integer or length-3 vector (t, h, w).} + +\item{stride}{Integer or length-3 vector.} + +\item{spatial_padding_mode}{Character. Conv padding mode.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +Spatial padding is handled by the convolution; temporal padding +replicates the first frame (causal) or both edge frames (non-causal), +chosen at call time. +} diff --git a/man/ltx23_census.Rd b/man/ltx23_census.Rd new file mode 100644 index 0000000..762edcb --- /dev/null +++ b/man/ltx23_census.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_census} +\alias{ltx23_census} +\title{Summarize checkpoint key coverage} +\usage{ +ltx23_census(ckpt) +} +\arguments{ +\item{ckpt}{An \code{ltx23_checkpoint}.} +} +\value{ +A data.frame with one row per component group and its key count. +} +\description{ +Summarize checkpoint key coverage +} diff --git a/man/ltx23_connector_transformer_1d.Rd b/man/ltx23_connector_transformer_1d.Rd new file mode 100644 index 0000000..058e6f5 --- /dev/null +++ b/man/ltx23_connector_transformer_1d.Rd @@ -0,0 +1,30 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_connector_transformer_1d} +\alias{ltx23_connector_transformer_1d} +\title{1D connector transformer} +\usage{ +ltx23_connector_transformer_1d(num_attention_heads = 32L, + attention_head_dim = 128L, num_layers = 8L, + num_learnable_registers = 128L, + rope_base_seq_len = 4096L, rope_theta = 10000, + rope_double_precision = TRUE, eps = 1e-06, + rope_type = "split", gated_attention = TRUE) +} +\arguments{ +\item{num_learnable_registers}{Integer or NULL. Register count (the +sequence length must be divisible by it).} + +\item{eps}{Numeric. Norm epsilon.} + +\item{gated_attention}{Logical. Per-head attention output gates.} + +\item{num_attention_heads,attention_head_dim,num_layers}{Transformer shape.} + +\item{rope_base_seq_len,rope_theta,rope_double_precision,rope_type}{RoPE config.} +} +\description{ +Replaces padded positions with learnable registers (valid tokens are +front-aligned in their original order; the tail is filled with +registers indexed by absolute position, after which the attention mask +is cleared), then runs 1D transformer blocks with rotary embeddings. +} diff --git a/man/ltx23_denormalize_latents.Rd b/man/ltx23_denormalize_latents.Rd new file mode 100644 index 0000000..b07a773 --- /dev/null +++ b/man/ltx23_denormalize_latents.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_denormalize_latents} +\alias{ltx23_denormalize_latents} +\title{Denormalize latents with the VAE's per-channel statistics} +\usage{ +ltx23_denormalize_latents(latents, latents_mean, latents_std) +} +\arguments{ +\item{latents}{Tensor [B, C, F, H, W].} + +\item{latents_mean,latents_std}{Tensors [C].} +} +\value{ +Denormalized latents ready for the decoder. +} +\description{ +Denormalize latents with the VAE's per-channel statistics +} diff --git a/man/ltx23_distilled_sigmas.Rd b/man/ltx23_distilled_sigmas.Rd new file mode 100644 index 0000000..90d84a8 --- /dev/null +++ b/man/ltx23_distilled_sigmas.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_distilled_sigmas} +\alias{ltx23_distilled_sigmas} +\title{Official distilled sigma schedule} +\usage{ +ltx23_distilled_sigmas() +} +\value{ +Numeric vector of length 9. +} +\description{ +The distilled LTX sigma values (with terminal zero appended), as +published in the Apache-2.0 diffusers reference +(pipelines/ltx2/utils.py). +} diff --git a/man/ltx23_downsample1d.Rd b/man/ltx23_downsample1d.Rd new file mode 100644 index 0000000..ecc73cb --- /dev/null +++ b/man/ltx23_downsample1d.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_downsample1d} +\alias{ltx23_downsample1d} +\title{Anti-aliasing 1D downsampler (low-pass then stride)} +\usage{ +ltx23_downsample1d(ratio = 2L, kernel_size = NULL) +} +\arguments{ +\item{ratio}{Integer. Downsampling ratio.} + +\item{kernel_size}{Integer or NULL (default 6*ratio rounded even).} +} +\description{ +Anti-aliasing 1D downsampler (low-pass then stride) +} diff --git a/man/ltx23_encode_audio.Rd b/man/ltx23_encode_audio.Rd new file mode 100644 index 0000000..eee39ed --- /dev/null +++ b/man/ltx23_encode_audio.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_encode_audio} +\alias{ltx23_encode_audio} +\title{Encode audio into normalized, packed conditioning latents} +\usage{ +ltx23_encode_audio(audio_vae, wav, audio_num_frames, frontend = NULL) +} +\arguments{ +\item{audio_vae}{An \code{ltx23_audio_vae} (with encoder weights).} + +\item{wav}{Matrix [2, samples] in [-1, 1] at 16 kHz (see +\code{\link{ltx23_read_audio}}).} + +\item{audio_num_frames}{Integer. Target latent length.} + +\item{frontend}{Optional prebuilt \code{\link{ltx23_audio_mel_frontend}}.} +} +\value{ +Packed normalized latents [1, audio_num_frames, 128] (float32). +} +\description{ +Pads or trims the waveform so the latent length equals +\code{audio_num_frames} (mel frames \code{4L - 3}, mirroring the +decoder's \code{target_frames}), computes the log-mel, encodes in +argmax mode, packs, and normalizes with the checkpoint statistics. +} diff --git a/man/ltx23_encode_video_frames.Rd b/man/ltx23_encode_video_frames.Rd new file mode 100644 index 0000000..420c726 --- /dev/null +++ b/man/ltx23_encode_video_frames.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_encode_video_frames} +\alias{ltx23_encode_video_frames} +\title{Encode pixel frames to normalized video latents} +\usage{ +ltx23_encode_video_frames(vae, frames) +} +\arguments{ +\item{vae}{An \code{ltx23_video_vae}.} + +\item{frames}{Tensor [1, 3, F, H, W] in [-1, 1] (see +\code{\link{ltx23_preprocess_frames}}).} +} +\value{ +Normalized latents [1, 128, F', H/32, W/32] (float32). +} +\description{ +VAE encode in "argmax" mode (the distribution mean), then normalize +with the checkpoint's per-channel statistics — the exact inverse of +the decode path. +} diff --git a/man/ltx23_feed_forward.Rd b/man/ltx23_feed_forward.Rd new file mode 100644 index 0000000..e142ead --- /dev/null +++ b/man/ltx23_feed_forward.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_feed_forward} +\alias{ltx23_feed_forward} +\title{LTX feed-forward layer} +\usage{ +ltx23_feed_forward(dim, mult = 4L) +} +\arguments{ +\item{dim}{Integer. Input/output dimension.} + +\item{mult}{Integer. Hidden dimension multiplier.} +} +\description{ +Linear -> GELU (tanh approximation) -> Linear with 4x hidden dim, +matching diffusers \code{FeedForward(activation_fn="gelu-approximate")} +state-dict names (net.0.proj, net.2). +} diff --git a/man/ltx23_fp8_linear.Rd b/man/ltx23_fp8_linear.Rd new file mode 100644 index 0000000..dd87539 --- /dev/null +++ b/man/ltx23_fp8_linear.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_fp8_linear} +\alias{ltx23_fp8_linear} +\title{FP8 linear layer} +\usage{ +ltx23_fp8_linear(out_features, in_features, bias = TRUE) +} +\arguments{ +\item{bias}{Logical.} + +\item{out_features,in_features}{Integers.} +} +\description{ +Weight lives as float8_e4m3fn plus a float32 scale in plain module +fields (so \code{$to(device)} moves only the bias); the forward pass +ships 1 byte/param to the input's device, upcasts, rescales, and runs +\code{nnf_linear}. +} diff --git a/man/ltx23_get_timestep_embedding.Rd b/man/ltx23_get_timestep_embedding.Rd new file mode 100644 index 0000000..9b6e90d --- /dev/null +++ b/man/ltx23_get_timestep_embedding.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_get_timestep_embedding} +\alias{ltx23_get_timestep_embedding} +\title{Sinusoidal timestep embedding} +\usage{ +ltx23_get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos = TRUE, + downscale_freq_shift = 0, max_period = 10000) +} +\arguments{ +\item{timesteps}{1D tensor of timestep values.} + +\item{embedding_dim}{Integer. Output embedding size.} + +\item{flip_sin_to_cos}{Logical. Put cos before sin.} + +\item{downscale_freq_shift}{Numeric. Frequency delta control.} + +\item{max_period}{Numeric. Maximum embedding frequency period.} +} +\value{ +Tensor [N, embedding_dim]. +} +\description{ +DDPM-style sinusoidal embedding. LTX uses \code{flip_sin_to_cos=TRUE} +(cos first) and \code{downscale_freq_shift=0}. +} diff --git a/man/ltx23_is_fp8_cast_key.Rd b/man/ltx23_is_fp8_cast_key.Rd new file mode 100644 index 0000000..7cac9c4 --- /dev/null +++ b/man/ltx23_is_fp8_cast_key.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_is_fp8_cast_key} +\alias{ltx23_is_fp8_cast_key} +\title{Test whether a mapped DiT key is in the official fp8 cast set} +\usage{ +ltx23_is_fp8_cast_key(mapped_key) +} +\arguments{ +\item{mapped_key}{Character vector of mapped (diffusers-style) +parameter names.} +} +\value{ +Logical vector. +} +\description{ +Test whether a mapped DiT key is in the official fp8 cast set +} diff --git a/man/ltx23_kaiser_sinc_filter1d.Rd b/man/ltx23_kaiser_sinc_filter1d.Rd new file mode 100644 index 0000000..0a67b1f --- /dev/null +++ b/man/ltx23_kaiser_sinc_filter1d.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_kaiser_sinc_filter1d} +\alias{ltx23_kaiser_sinc_filter1d} +\title{Kaiser sinc low-pass filter kernel} +\usage{ +ltx23_kaiser_sinc_filter1d(cutoff, half_width, kernel_size) +} +\arguments{ +\item{cutoff}{Numeric. Normalized cutoff in (0, 0.5].} + +\item{half_width}{Numeric. Transition band half width.} + +\item{kernel_size}{Integer.} +} +\value{ +Tensor [kernel_size]. +} +\description{ +Kaiser sinc low-pass filter kernel +} diff --git a/man/ltx23_latent_upsampler.Rd b/man/ltx23_latent_upsampler.Rd new file mode 100644 index 0000000..fa4cae1 --- /dev/null +++ b/man/ltx23_latent_upsampler.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_latent_upsampler} +\alias{ltx23_latent_upsampler} +\title{LTX-2.3 latent upsampler model} +\usage{ +ltx23_latent_upsampler(in_channels = 128L, mid_channels = 1024L, + num_blocks_per_stage = 4L) +} +\arguments{ +\item{in_channels}{Integer. Latent channels.} + +\item{mid_channels}{Integer.} + +\item{num_blocks_per_stage}{Integer.} +} +\description{ +Latents [B, 128, F, H, W] -> [B, 128, F, 2H, 2W]. +} diff --git a/man/ltx23_load_group.Rd b/man/ltx23_load_group.Rd new file mode 100644 index 0000000..8a89164 --- /dev/null +++ b/man/ltx23_load_group.Rd @@ -0,0 +1,37 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_load_group} +\alias{ltx23_load_group} +\title{Stream a checkpoint key group into a module} +\usage{ +ltx23_load_group(ckpt, keys, module, map_key = identity, verbose = TRUE, + gc_every = 50L) +} +\arguments{ +\item{ckpt}{An \code{ltx23_checkpoint}.} + +\item{keys}{Character vector of checkpoint keys to load (one group +from \code{\link{ltx23_split_keys}}).} + +\item{module}{A torch nn_module to populate.} + +\item{map_key}{Function mapping a checkpoint key to the module's +parameter/buffer name, or NA to skip the key deliberately.} + +\item{verbose}{Logical. Report progress and coverage.} + +\item{gc_every}{Integer. Run \code{gc()} after this many tensors.} +} +\value{ +Invisibly, a list with \code{unmapped} (checkpoint keys that + found no destination), \code{skipped} (keys the mapper declined), + and \code{unfilled} (module parameters/buffers never written). + A perfectly loaded group has zero \code{unmapped} and zero + \code{unfilled}. +} +\description{ +Reads tensors one at a time from an open checkpoint and copies them +into the matching parameters/buffers of \code{module}. Destination +names are derived by \code{map_key}; \code{$copy_()} handles any +dtype/device conversion, so the module may already live on its target +device in its target dtype. +} diff --git a/man/ltx23_load_pipeline.Rd b/man/ltx23_load_pipeline.Rd new file mode 100644 index 0000000..ede7113 --- /dev/null +++ b/man/ltx23_load_pipeline.Rd @@ -0,0 +1,49 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_load_pipeline} +\alias{ltx23_load_pipeline} +\title{Load the LTX-2.3 generation components from a single-file checkpoint} +\usage{ +ltx23_load_pipeline(checkpoint_path, device = "cuda", dtype = "bfloat16", + transformer_device = "cpu", + components = c("dit", "connectors", "vae", "audio_vae", "vocoder"), + pin = TRUE, attn_chunk = NULL, phase_offload = TRUE, + verbose = TRUE) +} +\arguments{ +\item{checkpoint_path}{Path to the single-file checkpoint (e.g. +\code{ltx-2.3-22b-distilled-1.1.safetensors}) or to an fp8 artifact +directory produced by \code{\link{ltx23_quantize_fp8}}. With the fp8 +artifact, the transformer loads with CPU-resident fp8 weights that +stream to \code{device} during the forward pass.} + +\item{device}{Character. Device for the small components (VAEs, +vocoder, connectors) and, with fp8, the transformer residents.} + +\item{dtype}{Character. "bfloat16" (checkpoint native) or "float32".} + +\item{transformer_device}{Character. Device for the transformer +weights when loading the plain (non-fp8) checkpoint.} + +\item{components}{Character vector. Which components to load.} + +\item{pin}{Logical. Pin fp8 host memory (fp8 artifact only).} + +\item{attn_chunk}{Integer or NULL. Query-chunk size for attention +(see \code{\link{ltx23_set_attn_chunk}}).} + +\item{phase_offload}{Logical. Load the small components (connectors, +VAEs, vocoder) to the CPU; the pipeline moves each onto the compute +device only for its phase.} + +\item{verbose}{Logical.} +} +\value{ +A list with the loaded modules and the checkpoint config, + class \code{ltx23_pipeline}. +} +\description{ +Builds the transformer, connectors, video VAE, audio VAE, and vocoder +with the LTX 2.3 configuration and streams the checkpoint weights into +them. The Gemma3 text encoder ships separately (see +\code{\link{load_gemma3_text_encoder}}). +} diff --git a/man/ltx23_load_transformer_fp8.Rd b/man/ltx23_load_transformer_fp8.Rd new file mode 100644 index 0000000..767c8f4 --- /dev/null +++ b/man/ltx23_load_transformer_fp8.Rd @@ -0,0 +1,30 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_load_transformer_fp8} +\alias{ltx23_load_transformer_fp8} +\title{Load the LTX-2.3 transformer with FP8 weights} +\usage{ +ltx23_load_transformer_fp8(ckpt, device = "cuda", pin = TRUE, verbose = TRUE, + ...) +} +\arguments{ +\item{ckpt}{An fp8 \code{ltx23_checkpoint} +(\code{\link{ltx23_open_fp8_checkpoint}}).} + +\item{device}{Character. Device for the resident (non-fp8) weights.} + +\item{pin}{Logical. Pin the fp8 host memory for faster transfers.} + +\item{verbose}{Logical.} + +\item{...}{Passed to \code{\link{ltx23_transformer}} (tiny test configs).} +} +\value{ +The loaded \code{ltx23_transformer}. +} +\description{ +Builds the transformer, swaps the official cast-set linears for +\code{\link{ltx23_fp8_linear}}, loads fp8 weights CPU-side (optionally +pinned) and everything else as bfloat16 on \code{device}. Sets +\code{options(diffuseR.block_gc = TRUE)} so the transformer runs +per-block garbage collection over the dequantized temporaries. +} diff --git a/man/ltx23_load_transformer_nf4.Rd b/man/ltx23_load_transformer_nf4.Rd new file mode 100644 index 0000000..cdab665 --- /dev/null +++ b/man/ltx23_load_transformer_nf4.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_load_transformer_nf4} +\alias{ltx23_load_transformer_nf4} +\title{Load the LTX-2.3 transformer with resident NF4 weights} +\usage{ +ltx23_load_transformer_nf4(ckpt, device = "cuda", verbose = TRUE, ...) +} +\arguments{ +\item{ckpt}{An NF4 \code{ltx23_checkpoint} +(\code{\link{ltx23_open_fp8_checkpoint}} reads any shard artifact).} + +\item{device}{Character.} + +\item{verbose}{Logical.} + +\item{...}{Passed to \code{\link{ltx23_transformer}} (tiny test configs).} +} +\value{ +The loaded \code{ltx23_transformer}. +} +\description{ +Builds the transformer, swaps the cast-set linears for +\code{\link{ltx23_nf4_linear}}, and loads everything onto +\code{device}: at ~4.5 bits/parameter the whole 22B transformer stays +GPU-resident, avoiding per-step weight transfers. +} diff --git a/man/ltx23_load_upsampler.Rd b/man/ltx23_load_upsampler.Rd new file mode 100644 index 0000000..c54b1f3 --- /dev/null +++ b/man/ltx23_load_upsampler.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_load_upsampler} +\alias{ltx23_load_upsampler} +\title{Load the LTX-2.3 spatial upscaler weights} +\usage{ +ltx23_load_upsampler(path, device = "cuda", dtype = "bfloat16", verbose = TRUE) +} +\arguments{ +\item{path}{Path to e.g. \code{ltx-2.3-spatial-upscaler-x2-1.1.safetensors}.} + +\item{verbose}{Logical.} + +\item{device,dtype}{Placement for the loaded model.} +} +\value{ +The loaded \code{ltx23_latent_upsampler}. +} +\description{ +The checkpoint keys match this module tree directly. +} diff --git a/man/ltx23_map_audio_vae_key.Rd b/man/ltx23_map_audio_vae_key.Rd new file mode 100644 index 0000000..cb5ea99 --- /dev/null +++ b/man/ltx23_map_audio_vae_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_map_audio_vae_key} +\alias{ltx23_map_audio_vae_key} +\title{Map an official audio VAE checkpoint key to the R module name} +\usage{ +ltx23_map_audio_vae_key(key) +} +\arguments{ +\item{key}{Character. Checkpoint key.} +} +\value{ +Character. +} +\description{ +Map an official audio VAE checkpoint key to the R module name +} diff --git a/man/ltx23_map_connector_key.Rd b/man/ltx23_map_connector_key.Rd new file mode 100644 index 0000000..ad12720 --- /dev/null +++ b/man/ltx23_map_connector_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_map_connector_key} +\alias{ltx23_map_connector_key} +\title{Map an official connectors checkpoint key to the R module name} +\usage{ +ltx23_map_connector_key(key) +} +\arguments{ +\item{key}{Character. Checkpoint key.} +} +\value{ +Character. Module parameter name. +} +\description{ +Map an official connectors checkpoint key to the R module name +} diff --git a/man/ltx23_map_dit_key.Rd b/man/ltx23_map_dit_key.Rd new file mode 100644 index 0000000..80f8cd3 --- /dev/null +++ b/man/ltx23_map_dit_key.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_map_dit_key} +\alias{ltx23_map_dit_key} +\title{Map an official DiT checkpoint key to the R module name} +\usage{ +ltx23_map_dit_key(key) +} +\arguments{ +\item{key}{Character. Checkpoint key (with or without the +\code{model.diffusion_model.} prefix).} +} +\value{ +Character. Module parameter/buffer name. +} +\description{ +Applies the official-to-diffusers renames for the LTX-2.3 transformer +(cf. diffusers scripts/convert_ltx2_to_diffusers.py). Our module tree +matches the diffusers names, so this is the full mapping. +} diff --git a/man/ltx23_map_vae_key.Rd b/man/ltx23_map_vae_key.Rd new file mode 100644 index 0000000..b90195e --- /dev/null +++ b/man/ltx23_map_vae_key.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_map_vae_key} +\alias{ltx23_map_vae_key} +\title{Map an official VAE checkpoint key to the R module name} +\usage{ +ltx23_map_vae_key(key) +} +\arguments{ +\item{key}{Character. Checkpoint key (with or without "vae." prefix).} +} +\value{ +Character. Module parameter/buffer name. +} +\description{ +The official checkpoint stores the encoder/decoder as flat block lists +(down_blocks.0-8 / up_blocks.0-8) where downsamplers/upsamplers and +the mid block are separate entries; diffusers (and this port) nest +them. Index mapping per diffusers convert_ltx2_to_diffusers.py. +} diff --git a/man/ltx23_map_vocoder_key.Rd b/man/ltx23_map_vocoder_key.Rd new file mode 100644 index 0000000..5064077 --- /dev/null +++ b/man/ltx23_map_vocoder_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_map_vocoder_key} +\alias{ltx23_map_vocoder_key} +\title{Map an official vocoder checkpoint key to the R module name} +\usage{ +ltx23_map_vocoder_key(key) +} +\arguments{ +\item{key}{Character. Checkpoint key.} +} +\value{ +Character. Module parameter/buffer name. +} +\description{ +Map an official vocoder checkpoint key to the R module name +} diff --git a/man/ltx23_mel_stft.Rd b/man/ltx23_mel_stft.Rd new file mode 100644 index 0000000..5d36d8e --- /dev/null +++ b/man/ltx23_mel_stft.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_mel_stft} +\alias{ltx23_mel_stft} +\title{Causal log-mel spectrogram with checkpoint-loaded bases} +\usage{ +ltx23_mel_stft(filter_length = 512L, hop_length = 80L, window_length = 512L, + num_mel_channels = 64L) +} +\arguments{ +\item{filter_length,hop_length,window_length,num_mel_channels}{Integers.} +} +\description{ +Causal log-mel spectrogram with checkpoint-loaded bases +} diff --git a/man/ltx23_memory_profile.Rd b/man/ltx23_memory_profile.Rd new file mode 100644 index 0000000..cd2b2f8 --- /dev/null +++ b/man/ltx23_memory_profile.Rd @@ -0,0 +1,33 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_memory_profile} +\alias{ltx23_memory_profile} +\title{Get an LTX-2.3 memory profile} +\usage{ +ltx23_memory_profile(vram_gb = NULL) +} +\arguments{ +\item{vram_gb}{Numeric or NULL (auto-detect free VRAM).} +} +\value{ +Named list with device/dtype placement, \code{attn_chunk}, + \code{pin_weights}, and resolution caps. +} +\description{ +Selects transformer precision, component placement, and attention +chunking for the available VRAM. Measured on an RTX 5060 Ti (16 GB): +fp8 streaming peaks ~11.6 GB (without phase offloading) at +512x320x49; NF4 keeps the whole 22B transformer resident (~12.5 GB) +and removes the ~21 GB/step PCIe weight streaming. The NF4 profile +renders 1280x704x121 with audio in ~23 min at a 15.7 GB peak +(tiled VAE decode, in-place feed-forward GELU, and the default +\code{diffuseR.attn_budget} of 1.5e8 all required at that size). +} +\details{ +\describe{ +\item{precision "nf4"}{Weights resident on the GPU; fastest steps; +about 9 percent weight round-trip error.} +\item{precision "fp8"}{Weights CPU-resident, streamed per linear; +near-bf16 quality; each step pays the PCIe transfer.} +} + +} diff --git a/man/ltx23_nf4_dequantize.Rd b/man/ltx23_nf4_dequantize.Rd new file mode 100644 index 0000000..5ec7975 --- /dev/null +++ b/man/ltx23_nf4_dequantize.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_nf4_dequantize} +\alias{ltx23_nf4_dequantize} +\title{Dequantize NF4 data to a float tensor} +\usage{ +ltx23_nf4_dequantize(packed, absmax, shape, dtype = torch::torch_bfloat16(), + chunk_elements = 8388608L, out = NULL) +} +\arguments{ +\item{packed}{uint8 tensor of packed index pairs.} + +\item{absmax}{float32 tensor of per-block scales.} + +\item{shape}{Integer vector. Original tensor shape.} + +\item{dtype}{Target torch dtype.} + +\item{chunk_elements}{Integer. Elements dequantized per slice (bounds +the int64 index temporary).} + +\item{out}{Optional preallocated tensor of \code{shape} to write into +(avoids allocating a fresh weight tensor per call).} +} +\value{ +Tensor of \code{shape} in \code{dtype} on the input's device. +} +\description{ +Dequantize NF4 data to a float tensor +} diff --git a/man/ltx23_nf4_linear.Rd b/man/ltx23_nf4_linear.Rd new file mode 100644 index 0000000..2bc15c6 --- /dev/null +++ b/man/ltx23_nf4_linear.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_nf4_linear} +\alias{ltx23_nf4_linear} +\title{NF4 linear layer} +\usage{ +ltx23_nf4_linear(out_features, in_features, bias = TRUE) +} +\arguments{ +\item{bias}{Logical.} + +\item{out_features,in_features}{Integers.} +} +\description{ +Packed weights and per-block scales are registered as buffers, so +they move with the module (uint8 packs are untouched by dtype +conversions). The forward pass dequantizes on the weight's device. +} diff --git a/man/ltx23_nf4_quantize.Rd b/man/ltx23_nf4_quantize.Rd new file mode 100644 index 0000000..ec76592 --- /dev/null +++ b/man/ltx23_nf4_quantize.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_nf4_quantize} +\alias{ltx23_nf4_quantize} +\title{Quantize a tensor to NF4} +\usage{ +ltx23_nf4_quantize(x) +} +\arguments{ +\item{x}{Float tensor (any shape; total elements must be a multiple +of 128, i.e. two 64-element blocks - always true for the LTX +linears).} +} +\value{ +List with \code{packed} (uint8, two indices per byte) and + \code{absmax} (float32, one per 64-element block). +} +\description{ +Quantize a tensor to NF4 +} diff --git a/man/ltx23_normalize_latents.Rd b/man/ltx23_normalize_latents.Rd new file mode 100644 index 0000000..936d9d3 --- /dev/null +++ b/man/ltx23_normalize_latents.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_normalize_latents} +\alias{ltx23_normalize_latents} +\title{Normalize latents with the VAE's per-channel statistics} +\usage{ +ltx23_normalize_latents(latents, latents_mean, latents_std) +} +\arguments{ +\item{latents}{Tensor [B, C, F, H, W].} + +\item{latents_mean,latents_std}{Tensors [C].} +} +\value{ +Normalized latents. +} +\description{ +Normalize latents with the VAE's per-channel statistics +} diff --git a/man/ltx23_open_checkpoint.Rd b/man/ltx23_open_checkpoint.Rd new file mode 100644 index 0000000..980ada2 --- /dev/null +++ b/man/ltx23_open_checkpoint.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_open_checkpoint} +\alias{ltx23_open_checkpoint} +\title{Open an LTX-2.3 checkpoint} +\usage{ +ltx23_open_checkpoint(path, require_version = "2.3") +} +\arguments{ +\item{path}{Path to the checkpoint .safetensors file.} + +\item{require_version}{Character. Required \code{model_version} prefix +(default "2.3"). Set to NULL to skip the check.} +} +\value{ +An object of class \code{ltx23_checkpoint}: a list with + \code{handle} (safetensors reader), \code{keys}, \code{version}, + \code{config} (parsed component configs, or NULL), and \code{path}. +} +\description{ +Opens a single-file LTX checkpoint lazily (header only), validates the +\code{model_version} metadata, and parses the embedded component +configuration. +} +\examples{ +\dontrun{ +ckpt <- ltx23_open_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") +str(ltx23_split_keys(ckpt$keys), max.level = 1) +} +} diff --git a/man/ltx23_open_fp8_checkpoint.Rd b/man/ltx23_open_fp8_checkpoint.Rd new file mode 100644 index 0000000..2b3ada0 --- /dev/null +++ b/man/ltx23_open_fp8_checkpoint.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_open_fp8_checkpoint} +\alias{ltx23_open_fp8_checkpoint} +\title{Open an FP8 shard directory as a checkpoint} +\usage{ +ltx23_open_fp8_checkpoint(dir) +} +\arguments{ +\item{dir}{The fp8 artifact directory (with manifest.json).} +} +\value{ +An \code{ltx23_checkpoint}. +} +\description{ +Presents the sharded fp8 artifact through the same interface as +\code{\link{ltx23_open_checkpoint}} so the group loaders work +unchanged. +} diff --git a/man/ltx23_per_channel_rms_norm.Rd b/man/ltx23_per_channel_rms_norm.Rd new file mode 100644 index 0000000..74616f7 --- /dev/null +++ b/man/ltx23_per_channel_rms_norm.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_per_channel_rms_norm} +\alias{ltx23_per_channel_rms_norm} +\title{Per-channel RMS normalization} +\usage{ +ltx23_per_channel_rms_norm(eps = 1e-08) +} +\arguments{ +\item{eps}{Numeric. Stability epsilon.} +} +\description{ +Normalizes by the root-mean-square across the channel dimension +(dim 2 of [B, C, F, H, W]); no learned parameters. +} diff --git a/man/ltx23_per_token_rms_norm.Rd b/man/ltx23_per_token_rms_norm.Rd new file mode 100644 index 0000000..62c3a99 --- /dev/null +++ b/man/ltx23_per_token_rms_norm.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_per_token_rms_norm} +\alias{ltx23_per_token_rms_norm} +\title{Per-token RMS normalization over the channel axis} +\usage{ +ltx23_per_token_rms_norm(x, eps = 1e-06) +} +\arguments{ +\item{x}{Tensor [B, S, C, L] of stacked per-layer hidden states.} + +\item{eps}{Numeric. Stability epsilon.} +} +\value{ +Tensor of the same shape. +} +\description{ +Per-token RMS normalization over the channel axis +} diff --git a/man/ltx23_prepare_conditioned_latents.Rd b/man/ltx23_prepare_conditioned_latents.Rd new file mode 100644 index 0000000..6688cef --- /dev/null +++ b/man/ltx23_prepare_conditioned_latents.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_prepare_conditioned_latents} +\alias{ltx23_prepare_conditioned_latents} +\title{Build conditioned initial latents and the conditioning mask} +\usage{ +ltx23_prepare_conditioned_latents(cond_latents, latent_frames, latent_height, + latent_width, noise, cond_noise_scale = 0) +} +\arguments{ +\item{cond_latents}{Normalized condition latents +[1, 128, k, H', W'] from \code{\link{ltx23_encode_video_frames}}.} + +\item{noise}{Tensor [1, 128, F', H', W'] of standard noise (caller +provides so seeding stays in one place).} + +\item{cond_noise_scale}{Numeric. Optional partial noising of the +conditioned tokens (diffusers \code{noise_scale}, default 0).} + +\item{latent_frames,latent_height,latent_width}{Integers. Full +latent geometry of the generation.} +} +\value{ +list(latents [1, S, 128] float32 packed, + conditioning_mask [1, S] float32 packed). +} +\description{ +i2v (\code{cond_latents} has one latent frame): the encoded frame is +repeated across all latent frames and only latent frame 0 is marked +conditioned. Continuation (k latent frames): the prefix tokens are +replaced and marked. Unconditioned positions start as pure noise. +} diff --git a/man/ltx23_preprocess_frames.Rd b/man/ltx23_preprocess_frames.Rd new file mode 100644 index 0000000..1b88752 --- /dev/null +++ b/man/ltx23_preprocess_frames.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_preprocess_frames} +\alias{ltx23_preprocess_frames} +\title{Preprocess an image (or frame stack) for VAE encoding} +\usage{ +ltx23_preprocess_frames(x, height, width) +} +\arguments{ +\item{x}{Path to a PNG/JPEG, or an array [H, W, 3] (values in +[0, 1]), or a [F, H, W, 3] array of frames.} + +\item{height,width}{Integers. Target size (multiples of 32).} +} +\value{ +Float32 tensor [1, 3, F, height, width] in [-1, 1]. +} +\description{ +Mirrors the diffusers VideoProcessor: bilinear resize so the shorter +relative side matches, center-crop to the exact target, and scale to +[-1, 1]. +} diff --git a/man/ltx23_quantize_fp8.Rd b/man/ltx23_quantize_fp8.Rd new file mode 100644 index 0000000..3ed610c --- /dev/null +++ b/man/ltx23_quantize_fp8.Rd @@ -0,0 +1,30 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_quantize_fp8} +\alias{ltx23_quantize_fp8} +\title{Quantize an LTX-2.3 checkpoint to FP8 shards} +\usage{ +ltx23_quantize_fp8(checkpoint_path, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), + shard_bytes = 4e+09, force = FALSE, verbose = TRUE) +} +\arguments{ +\item{checkpoint_path}{Source .safetensors (46 GB bf16 single file).} + +\item{output_dir}{Output directory for shards + manifest.} + +\item{shard_bytes}{Numeric. Approximate shard size (default 4 GB).} + +\item{force}{Logical. Re-quantize even if a valid manifest exists.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the manifest list. +} +\description{ +Streams the single-file bf16 checkpoint tensor by tensor. DiT +attention/FFN linear weights are stored as float8_e4m3fn with a +float32 absmax/448 per-tensor scale (\code{_scale} sibling); +everything else is copied through unchanged. Output shards carry the +original key names plus a manifest for skip-if-exists. +} diff --git a/man/ltx23_quantize_nf4.Rd b/man/ltx23_quantize_nf4.Rd new file mode 100644 index 0000000..14c508b --- /dev/null +++ b/man/ltx23_quantize_nf4.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_quantize_nf4} +\alias{ltx23_quantize_nf4} +\title{Quantize an LTX-2.3 checkpoint to NF4 shards} +\usage{ +ltx23_quantize_nf4(checkpoint_path, + output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-nf4"), + shard_bytes = 4e+09, force = FALSE, verbose = TRUE) +} +\arguments{ +\item{checkpoint_path}{Source .safetensors (bf16 single file).} + +\item{output_dir}{Output directory for shards + manifest.} + +\item{shard_bytes}{Numeric. Approximate shard size.} + +\item{force}{Logical. Re-quantize even if a valid manifest exists.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the manifest list. +} +\description{ +Same streaming layout and cast policy as +\code{\link{ltx23_quantize_fp8}}, but cast-set weights are stored as +NF4 (\code{} packed uint8 + \code{_absmax} float32 blocks + +the original shape recovered from the model config at load time). +Non-cast tensors are copied through unchanged. The manifest carries +\code{format = "nf4"}. +} diff --git a/man/ltx23_read_audio.Rd b/man/ltx23_read_audio.Rd new file mode 100644 index 0000000..a048355 --- /dev/null +++ b/man/ltx23_read_audio.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_read_audio} +\alias{ltx23_read_audio} +\title{Read an audio file as 16 kHz stereo PCM} +\usage{ +ltx23_read_audio(path, sample_rate = 16000L) +} +\arguments{ +\item{path}{Audio file.} + +\item{sample_rate}{Integer.} +} +\value{ +Matrix [2, samples] in [-1, 1]. +} +\description{ +Decodes MP3/WAV/etc. via \code{av} to 16-bit PCM at the target rate +and parses the RIFF container in base R. +} diff --git a/man/ltx23_read_tail_frames.Rd b/man/ltx23_read_tail_frames.Rd new file mode 100644 index 0000000..2897411 --- /dev/null +++ b/man/ltx23_read_tail_frames.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_read_tail_frames} +\alias{ltx23_read_tail_frames} +\title{Read the trailing frames of a video file} +\usage{ +ltx23_read_tail_frames(path, n = 9L) +} +\arguments{ +\item{path}{Video file.} + +\item{n}{Integer. Trailing frame count.} +} +\value{ +Array [n, H, W, 3] in [0, 1]. +} +\description{ +Extracts the last \code{n} frames of an MP4 (via \code{av}) for use +as continuation conditioning. +} diff --git a/man/ltx23_release_dequant_buffers.Rd b/man/ltx23_release_dequant_buffers.Rd new file mode 100644 index 0000000..46972c3 --- /dev/null +++ b/man/ltx23_release_dequant_buffers.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_release_dequant_buffers} +\alias{ltx23_release_dequant_buffers} +\title{Release the NF4 dequantization buffers} +\usage{ +ltx23_release_dequant_buffers() +} +\value{ +Invisibly, NULL. +} +\description{ +Frees the cached per-shape weight buffers (e.g. before decoding at +high resolution). +} diff --git a/man/ltx23_rms_norm.Rd b/man/ltx23_rms_norm.Rd new file mode 100644 index 0000000..b7214ae --- /dev/null +++ b/man/ltx23_rms_norm.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_rms_norm} +\alias{ltx23_rms_norm} +\title{RMS normalization} +\usage{ +ltx23_rms_norm(dim, eps = 1e-06, elementwise_affine = TRUE) +} +\arguments{ +\item{dim}{Integer. Normalized dimension size.} + +\item{eps}{Numeric. Stability epsilon.} + +\item{elementwise_affine}{Logical. Learn a scale weight.} +} +\description{ +Variance is computed in float32; the result is cast back to the input +dtype (or the weight dtype when elementwise affine). +} diff --git a/man/ltx23_rotary_pos_embed.Rd b/man/ltx23_rotary_pos_embed.Rd new file mode 100644 index 0000000..8d6dcdd --- /dev/null +++ b/man/ltx23_rotary_pos_embed.Rd @@ -0,0 +1,47 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_rotary_pos_embed} +\alias{ltx23_rotary_pos_embed} +\title{LTX-2.3 audio/video rotary position embedder} +\usage{ +ltx23_rotary_pos_embed(dim, patch_size = 1L, patch_size_t = 1L, + base_num_frames = 20L, base_height = 2048L, + base_width = 2048L, sampling_rate = 16000L, + hop_length = 160L, scale_factors = c(8L, 32L, 32L), + theta = 10000, causal_offset = 1L, modality = "video", + double_precision = TRUE, rope_type = "split", + num_attention_heads = 32L) +} +\arguments{ +\item{dim}{Integer. Rotary dimension (attention head dim x heads for +split type at model level; see reference).} + +\item{scale_factors}{Integer vector. VAE (time, height, width) scale +factors.} + +\item{theta}{Numeric. RoPE theta.} + +\item{causal_offset}{Integer. Temporal offset for the causal VAE +(first frame has stride 1).} + +\item{modality}{"video" or "audio".} + +\item{double_precision}{Logical. Compute base frequencies in float64.} + +\item{rope_type}{"split" (LTX 2.3) or "interleaved".} + +\item{num_attention_heads}{Integer. Needed for the split layout.} + +\item{patch_size,patch_size_t}{Integers. Spatial/temporal patch sizes.} + +\item{base_num_frames,base_height,base_width}{Integers. Base grid the +coordinates are normalized against.} + +\item{sampling_rate,hop_length}{Integers. Audio spectrogram params.} +} +\description{ +Computes RoPE cos/sin frequency tensors from spatiotemporal patch +coordinates. Video coordinates are 3D (frames scaled to seconds via +fps, height, width in pixel space); audio coordinates are 1D +(seconds). Coordinates are patch boundaries [start, end); the midpoint +is used as the position. +} diff --git a/man/ltx23_rotary_pos_embed_1d.Rd b/man/ltx23_rotary_pos_embed_1d.Rd new file mode 100644 index 0000000..9324bd9 --- /dev/null +++ b/man/ltx23_rotary_pos_embed_1d.Rd @@ -0,0 +1,25 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_rotary_pos_embed_1d} +\alias{ltx23_rotary_pos_embed_1d} +\title{1D rotary embeddings for the text connectors} +\usage{ +ltx23_rotary_pos_embed_1d(dim, base_seq_len = 4096L, theta = 10000, + double_precision = TRUE, rope_type = "split", + num_attention_heads = 32L) +} +\arguments{ +\item{dim}{Integer. Rotary dimension (connector inner dim).} + +\item{base_seq_len}{Integer. Base sequence length for normalization.} + +\item{theta}{Numeric. RoPE theta.} + +\item{double_precision}{Logical. Compute base frequencies in float64.} + +\item{rope_type}{"split" (LTX-2.3) or "interleaved".} + +\item{num_attention_heads}{Integer. For the split per-head layout.} +} +\description{ +1D rotary embeddings for the text connectors +} diff --git a/man/ltx23_set_attn_chunk.Rd b/man/ltx23_set_attn_chunk.Rd new file mode 100644 index 0000000..0b9345a --- /dev/null +++ b/man/ltx23_set_attn_chunk.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_set_attn_chunk} +\alias{ltx23_set_attn_chunk} +\title{Set the attention query-chunk size across a transformer} +\usage{ +ltx23_set_attn_chunk(transformer, chunk) +} +\arguments{ +\item{transformer}{An \code{ltx23_transformer}.} + +\item{chunk}{Integer or NULL.} +} +\value{ +Invisibly, the transformer. +} +\description{ +R torch has no fused attention, so the [B, H, S, S] matrix +materializes; chunking queries bounds the peak. NULL disables +chunking. +} diff --git a/man/ltx23_snake_beta.Rd b/man/ltx23_snake_beta.Rd new file mode 100644 index 0000000..ec90b66 --- /dev/null +++ b/man/ltx23_snake_beta.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_snake_beta} +\alias{ltx23_snake_beta} +\title{SnakeBeta activation} +\usage{ +ltx23_snake_beta(channels, eps = 1e-09) +} +\arguments{ +\item{channels}{Integer.} + +\item{eps}{Numeric.} +} +\description{ +\code{x + (1 / (beta + eps)) * sin(x * alpha)^2} with per-channel +log-scale alpha/beta parameters. +} diff --git a/man/ltx23_split_keys.Rd b/man/ltx23_split_keys.Rd new file mode 100644 index 0000000..7fdbefe --- /dev/null +++ b/man/ltx23_split_keys.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_split_keys} +\alias{ltx23_split_keys} +\title{Split checkpoint keys by component} +\usage{ +ltx23_split_keys(keys) +} +\arguments{ +\item{keys}{Character vector of checkpoint tensor names.} +} +\value{ +Named list of character vectors: \code{dit}, + \code{connectors}, \code{vae}, \code{audio_vae}, \code{vocoder}, + and \code{other} (anything unrecognized; should be empty). +} +\description{ +Splits the flat key space of an LTX single-file checkpoint into its +component groups. Connector tensors live under the +\code{model.diffusion_model.} prefix alongside the transformer, plus a +top-level \code{text_embedding_projection.} group; both are routed to +the \code{connectors} component. +} diff --git a/man/ltx23_stage2_distilled_sigmas.Rd b/man/ltx23_stage2_distilled_sigmas.Rd new file mode 100644 index 0000000..2826575 --- /dev/null +++ b/man/ltx23_stage2_distilled_sigmas.Rd @@ -0,0 +1,13 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_stage2_distilled_sigmas} +\alias{ltx23_stage2_distilled_sigmas} +\title{Stage-2 distilled sigma schedule (two-stage refinement)} +\usage{ +ltx23_stage2_distilled_sigmas() +} +\value{ +Numeric vector of length 4. +} +\description{ +Stage-2 distilled sigma schedule (two-stage refinement) +} diff --git a/man/ltx23_text_connectors.Rd b/man/ltx23_text_connectors.Rd new file mode 100644 index 0000000..ab6722c --- /dev/null +++ b/man/ltx23_text_connectors.Rd @@ -0,0 +1,53 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_text_connectors} +\alias{ltx23_text_connectors} +\title{LTX-2.3 text connectors} +\usage{ +ltx23_text_connectors(caption_channels = 3840L, text_proj_in_factor = 49L, + video_connector_num_attention_heads = 32L, + video_connector_attention_head_dim = 128L, + video_connector_num_layers = 8L, + video_connector_num_learnable_registers = 128L, + video_gated_attn = TRUE, + audio_connector_num_attention_heads = 32L, + audio_connector_attention_head_dim = 64L, + audio_connector_num_layers = 8L, + audio_connector_num_learnable_registers = 128L, + audio_gated_attn = TRUE, + connector_rope_base_seq_len = 4096L, rope_theta = 10000, + rope_double_precision = TRUE, rope_type = "split", + video_hidden_dim = 4096L, audio_hidden_dim = 2048L, + proj_bias = TRUE) +} +\arguments{ +\item{caption_channels}{Integer. Text encoder hidden size (3840 for +Gemma3-12B).} + +\item{text_proj_in_factor}{Integer. Number of stacked hidden states +(num_layers + 1 = 49 for Gemma3-12B).} + +\item{video_connector_num_learnable_registers}{Integer or NULL.} + +\item{video_gated_attn}{Logical.} + +\item{audio_connector_num_learnable_registers}{Integer or NULL.} + +\item{audio_gated_attn}{Logical.} + +\item{proj_bias}{Logical. Projection bias (TRUE for LTX-2.3).} + +\item{video_connector_num_attention_heads,video_connector_attention_head_dim,video_connector_num_layers}{Video connector shape (LTX-2.3: 32 x 128, 8 layers).} + +\item{audio_connector_num_attention_heads,audio_connector_attention_head_dim,audio_connector_num_layers}{Audio connector shape (LTX-2.3: 32 x 64, 8 layers).} + +\item{connector_rope_base_seq_len,rope_theta,rope_double_precision,rope_type}{RoPE config.} + +\item{video_hidden_dim,audio_hidden_dim}{Integers. Projection targets +(DiT inner dims: 4096 / 2048).} +} +\description{ +Takes raw stacked per-layer text encoder hidden states and produces +the video and audio text embeddings for the DiT: per-token RMS norm, +per-modality sqrt(dim ratio) rescaling and projection, then a +per-modality 1D connector transformer. +} diff --git a/man/ltx23_tone_map_latents.Rd b/man/ltx23_tone_map_latents.Rd new file mode 100644 index 0000000..a238ce4 --- /dev/null +++ b/man/ltx23_tone_map_latents.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_tone_map_latents} +\alias{ltx23_tone_map_latents} +\title{Sigmoid tone mapping for latents} +\usage{ +ltx23_tone_map_latents(latents, compression) +} +\arguments{ +\item{latents}{Tensor.} + +\item{compression}{Numeric in [0, 1].} +} +\value{ +Tone-mapped latents. +} +\description{ +Compresses the latent dynamic range (cf. diffusers +\code{tone_map_latents}). \code{compression} 0 is identity, 1 is the +full effect. +} diff --git a/man/ltx23_transformer.Rd b/man/ltx23_transformer.Rd new file mode 100644 index 0000000..5a1837f --- /dev/null +++ b/man/ltx23_transformer.Rd @@ -0,0 +1,65 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_transformer} +\alias{ltx23_transformer} +\title{LTX-2.3 video transformer model} +\usage{ +ltx23_transformer(in_channels = 128L, out_channels = 128L, patch_size = 1L, + patch_size_t = 1L, num_attention_heads = 32L, + attention_head_dim = 128L, cross_attention_dim = 4096L, + vae_scale_factors = c(8L, 32L, 32L), pos_embed_max_pos = 20L, + base_height = 2048L, base_width = 2048L, gated_attn = TRUE, + cross_attn_mod = TRUE, audio_in_channels = 128L, + audio_out_channels = 128L, audio_patch_size = 1L, + audio_patch_size_t = 1L, audio_num_attention_heads = 32L, + audio_attention_head_dim = 64L, + audio_cross_attention_dim = 2048L, audio_scale_factor = 4L, + audio_pos_embed_max_pos = 20L, audio_sampling_rate = 16000L, + audio_hop_length = 160L, audio_gated_attn = TRUE, + audio_cross_attn_mod = TRUE, num_layers = 48L, + norm_eps = 1e-06, rope_theta = 10000, + rope_double_precision = TRUE, causal_offset = 1L, + timestep_scale_multiplier = 1000, + cross_attn_timestep_scale_multiplier = 1000, + rope_type = "split", perturbed_attn = TRUE) +} +\arguments{ +\item{cross_attention_dim}{Integer. Video text embedding dimension.} + +\item{vae_scale_factors}{Integer vector. VAE (time, height, width) scales.} + +\item{audio_cross_attention_dim}{Integer. Audio text embedding dimension.} + +\item{num_layers}{Integer. Transformer block count.} + +\item{norm_eps}{Numeric. Norm epsilon.} + +\item{rope_type}{"split" (LTX-2.3) or "interleaved".} + +\item{in_channels,out_channels}{Integers. Video latent channels.} + +\item{patch_size,patch_size_t}{Integers. Video patch sizes.} + +\item{num_attention_heads,attention_head_dim}{Video attention shape.} + +\item{pos_embed_max_pos,base_height,base_width}{RoPE base grid.} + +\item{audio_in_channels,audio_out_channels}{Integers. Audio latent channels.} + +\item{audio_patch_size,audio_patch_size_t}{Integers. Audio patch sizes.} + +\item{audio_num_attention_heads,audio_attention_head_dim}{Audio attention shape.} + +\item{audio_scale_factor,audio_pos_embed_max_pos,audio_sampling_rate,audio_hop_length}{Audio latent grid parameters.} + +\item{rope_theta,rope_double_precision,causal_offset}{RoPE parameters.} + +\item{timestep_scale_multiplier,cross_attn_timestep_scale_multiplier}{Timestep scaling (inputs arrive already scaled; the ratio modulates +the a2v/v2a gates).} + +\item{gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn}{LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints).} +} +\description{ +Dual-stream audio/video DiT. Text embeddings arrive already projected +to the video (\code{inner_dim}) and audio (\code{audio_inner_dim}) +dimensions by the connector modules. +} diff --git a/man/ltx23_transformer_block.Rd b/man/ltx23_transformer_block.Rd new file mode 100644 index 0000000..6391570 --- /dev/null +++ b/man/ltx23_transformer_block.Rd @@ -0,0 +1,44 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_transformer_block} +\alias{ltx23_transformer_block} +\title{LTX-2 transformer block} +\usage{ +ltx23_transformer_block(dim, num_attention_heads, attention_head_dim, + cross_attention_dim, audio_dim, + audio_num_attention_heads, audio_attention_head_dim, + audio_cross_attention_dim, video_gated_attn = TRUE, + video_cross_attn_adaln = TRUE, audio_gated_attn = TRUE, + audio_cross_attn_adaln = TRUE, eps = 1e-06, + elementwise_affine = FALSE, rope_type = "split", + perturbed_attn = TRUE) +} +\arguments{ +\item{cross_attention_dim}{Integer. Text embedding dim for video.} + +\item{audio_cross_attention_dim}{Integer. Text embedding dim for audio.} + +\item{eps}{Numeric. Norm epsilon.} + +\item{elementwise_affine}{Logical. Block norms carry weights (FALSE for LTX).} + +\item{rope_type}{"split" or "interleaved".} + +\item{perturbed_attn}{Logical. Enable the STG perturbation arguments.} + +\item{dim,audio_dim}{Integers. Video/audio stream dimensions.} + +\item{num_attention_heads,attention_head_dim}{Video attention shape.} + +\item{audio_num_attention_heads,audio_attention_head_dim}{Audio attention shape.} + +\item{video_gated_attn,audio_gated_attn}{Logicals. Per-head output gates.} + +\item{video_cross_attn_adaln,audio_cross_attn_adaln}{Logicals. LTX-2.3 +text cross-attention modulation (9 mod params instead of 6).} +} +\description{ +Dual-stream (video + audio) block: modulated self-attention per +modality, text cross-attention per modality (with LTX-2.3 query and +key/value modulation), bidirectional audio-video cross-attention with +global+per-block modulation, and modulated feed-forward. +} diff --git a/man/ltx23_tune_gc.Rd b/man/ltx23_tune_gc.Rd new file mode 100644 index 0000000..8c280bf --- /dev/null +++ b/man/ltx23_tune_gc.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_tune_gc} +\alias{ltx23_tune_gc} +\title{Tune the torch CUDA allocator for large-resident inference} +\usage{ +ltx23_tune_gc(footprint_gb = 12, total_gb = NULL) +} +\arguments{ +\item{footprint_gb}{Numeric. Expected resident GPU footprint in GB +(NF4 transformer: ~12).} + +\item{total_gb}{Numeric or NULL (auto-detect total VRAM).} +} +\value{ +Invisibly, the applied reserved rate (NULL if skipped). +} +\description{ +Stops the allocator GC storm (cf. ~/skills/torch +torch-jit-gc-performance.md): lantern proactively calls R's gc() +whenever reserved memory exceeds \code{torch.cuda_allocator_reserved_rate} +(default 0.20) of the card. With ~75\\% of VRAM occupied by resident +weights that fires on nearly every allocation. Raising the rate to the +actual footprint is safe here because the LTX hot loops compute into +persistent scratch buffers (near-zero per-step garbage). Also raises +the host-allocation GC threshold and defaults +\code{PYTORCH_CUDA_ALLOC_CONF} to expandable segments. Must run before +the first CUDA op; user-set options win. +} diff --git a/man/ltx23_upsample1d.Rd b/man/ltx23_upsample1d.Rd new file mode 100644 index 0000000..bdf5800 --- /dev/null +++ b/man/ltx23_upsample1d.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_upsample1d} +\alias{ltx23_upsample1d} +\title{Anti-aliasing 1D upsampler (transposed low-pass)} +\usage{ +ltx23_upsample1d(ratio = 2L, kernel_size = NULL, window_type = "kaiser", + persistent = TRUE) +} +\arguments{ +\item{ratio}{Integer. Upsampling ratio.} + +\item{kernel_size}{Integer or NULL.} + +\item{window_type}{"kaiser" (BigVGAN default) or "hann" (final resampler).} + +\item{persistent}{Logical. Register the filter as a buffer (present in +checkpoints); FALSE stores the computed filter as a plain field.} +} +\description{ +Anti-aliasing 1D upsampler (transposed low-pass) +} diff --git a/man/ltx23_video_decoder3d.Rd b/man/ltx23_video_decoder3d.Rd new file mode 100644 index 0000000..7b34f94 --- /dev/null +++ b/man/ltx23_video_decoder3d.Rd @@ -0,0 +1,45 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_decoder3d} +\alias{ltx23_video_decoder3d} +\title{LTX-2.3 video decoder} +\usage{ +ltx23_video_decoder3d(in_channels = 128L, out_channels = 3L, + block_out_channels = c(256L, 512L, 512L, 1024L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + upsample_type = c("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + patch_size = 4L, patch_size_t = 1L, + resnet_norm_eps = 1e-06, is_causal = FALSE, + upsample_residual = c(FALSE, FALSE, FALSE, FALSE), + upsample_factor = c(2L, 2L, 1L, 2L), + spatial_padding_mode = "zeros") +} +\arguments{ +\item{block_out_channels}{Integer vector (config order).} + +\item{spatio_temporal_scaling}{Logical vector per up block.} + +\item{layers_per_block}{Integer vector (config order; first entry is +the mid block after reversal).} + +\item{upsample_type}{Character vector per up block (not reversed).} + +\item{resnet_norm_eps}{Numeric.} + +\item{is_causal}{Logical. FALSE for LTX (symmetric temporal padding).} + +\item{upsample_residual}{Logical vector per up block.} + +\item{upsample_factor}{Integer vector per up block.} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers. Latent and pixel channels.} + +\item{patch_size,patch_size_t}{Integers.} +} +\description{ +Latents [B, 128, F, H, W] -> pixel video [B, 3, 8F - 7, 32H, 32W]. +Block channel lists are given encoder-side (as in the config) and +reversed internally; \code{upsample_type} is indexed directly. +} diff --git a/man/ltx23_video_down_block3d.Rd b/man/ltx23_video_down_block3d.Rd new file mode 100644 index 0000000..e8004be --- /dev/null +++ b/man/ltx23_video_down_block3d.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_down_block3d} +\alias{ltx23_video_down_block3d} +\title{LTX video down block} +\usage{ +ltx23_video_down_block3d(in_channels, out_channels = NULL, num_layers = 1L, + resnet_eps = 1e-06, spatio_temporal_scale = TRUE, + downsample_type = "spatiotemporal", + spatial_padding_mode = "zeros") +} +\arguments{ +\item{num_layers}{Integer. ResNet count.} + +\item{resnet_eps}{Numeric.} + +\item{spatio_temporal_scale}{Logical. Whether to downsample at all.} + +\item{downsample_type}{"spatial", "temporal", or "spatiotemporal".} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +ResNet stack (at the input channel count) followed by a +pixel-unshuffle downsampler that also changes the channel count. +} diff --git a/man/ltx23_video_downsampler3d.Rd b/man/ltx23_video_downsampler3d.Rd new file mode 100644 index 0000000..0afebb6 --- /dev/null +++ b/man/ltx23_video_downsampler3d.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_downsampler3d} +\alias{ltx23_video_downsampler3d} +\title{Pixel-unshuffle 3D downsampler} +\usage{ +ltx23_video_downsampler3d(in_channels, out_channels, stride = c(1L, 1L, 1L), + spatial_padding_mode = "zeros") +} +\arguments{ +\item{stride}{Length-3 integer vector (t, h, w).} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +Conv followed by space/time-to-channel rearrangement, plus a grouped +channel-mean residual of the same rearrangement. +} diff --git a/man/ltx23_video_encoder3d.Rd b/man/ltx23_video_encoder3d.Rd new file mode 100644 index 0000000..ca384a9 --- /dev/null +++ b/man/ltx23_video_encoder3d.Rd @@ -0,0 +1,38 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_encoder3d} +\alias{ltx23_video_encoder3d} +\title{LTX-2.3 video encoder} +\usage{ +ltx23_video_encoder3d(in_channels = 3L, out_channels = 128L, + block_out_channels = c(256L, 512L, 1024L, 1024L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + patch_size = 4L, patch_size_t = 1L, + resnet_norm_eps = 1e-06, is_causal = TRUE, + spatial_padding_mode = "zeros") +} +\arguments{ +\item{block_out_channels}{Integer vector. Per-block output channels.} + +\item{spatio_temporal_scaling}{Logical vector per block.} + +\item{layers_per_block}{Integer vector (blocks then mid).} + +\item{downsample_type}{Character vector per block.} + +\item{resnet_norm_eps}{Numeric.} + +\item{is_causal}{Logical.} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers. Pixel and latent channels.} + +\item{patch_size,patch_size_t}{Integers. Pixel patchification.} +} +\description{ +Pixel video [B, 3, F, H, W] -> latent statistics +[B, 2 * latent_channels, F/8, H/32, W/32] (mean and a uniform log-var +channel broadcast across the latent channels). +} diff --git a/man/ltx23_video_mid_block3d.Rd b/man/ltx23_video_mid_block3d.Rd new file mode 100644 index 0000000..ac45a31 --- /dev/null +++ b/man/ltx23_video_mid_block3d.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_mid_block3d} +\alias{ltx23_video_mid_block3d} +\title{LTX video mid block} +\usage{ +ltx23_video_mid_block3d(in_channels, num_layers = 1L, resnet_eps = 1e-06, + spatial_padding_mode = "zeros") +} +\arguments{ +\item{in_channels}{Integer.} + +\item{num_layers}{Integer.} + +\item{resnet_eps}{Numeric.} + +\item{spatial_padding_mode}{Character.} +} +\description{ +A plain ResNet stack at a fixed channel count. +} diff --git a/man/ltx23_video_resnet_block3d.Rd b/man/ltx23_video_resnet_block3d.Rd new file mode 100644 index 0000000..2506593 --- /dev/null +++ b/man/ltx23_video_resnet_block3d.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_resnet_block3d} +\alias{ltx23_video_resnet_block3d} +\title{LTX 3D ResNet block} +\usage{ +ltx23_video_resnet_block3d(in_channels, out_channels = NULL, eps = 1e-06, + spatial_padding_mode = "zeros") +} +\arguments{ +\item{eps}{Numeric. Shortcut LayerNorm epsilon.} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +PerChannelRMSNorm -> SiLU -> causal conv, twice, with a LayerNorm + +1x1 Conv3d shortcut when the channel count changes. +} diff --git a/man/ltx23_video_up_block3d.Rd b/man/ltx23_video_up_block3d.Rd new file mode 100644 index 0000000..f234c16 --- /dev/null +++ b/man/ltx23_video_up_block3d.Rd @@ -0,0 +1,32 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_up_block3d} +\alias{ltx23_video_up_block3d} +\title{LTX video up block} +\usage{ +ltx23_video_up_block3d(in_channels, out_channels = NULL, num_layers = 1L, + resnet_eps = 1e-06, spatio_temporal_scale = TRUE, + upsample_type = "spatiotemporal", + upsample_residual = FALSE, upscale_factor = 1L, + spatial_padding_mode = "zeros") +} +\arguments{ +\item{num_layers}{Integer.} + +\item{resnet_eps}{Numeric.} + +\item{spatio_temporal_scale}{Logical.} + +\item{upsample_type}{"spatial", "temporal", or "spatiotemporal".} + +\item{upsample_residual}{Logical.} + +\item{upscale_factor}{Integer.} + +\item{spatial_padding_mode}{Character.} + +\item{in_channels,out_channels}{Integers.} +} +\description{ +Optional channel-changing conv-in ResNet, pixel-shuffle upsampler, +then a ResNet stack at the output channel count. +} diff --git a/man/ltx23_video_upsampler3d.Rd b/man/ltx23_video_upsampler3d.Rd new file mode 100644 index 0000000..e2d25ca --- /dev/null +++ b/man/ltx23_video_upsampler3d.Rd @@ -0,0 +1,24 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_upsampler3d} +\alias{ltx23_video_upsampler3d} +\title{Pixel-shuffle 3D upsampler} +\usage{ +ltx23_video_upsampler3d(in_channels, stride = c(1L, 1L, 1L), residual = FALSE, + upscale_factor = 1L, spatial_padding_mode = "zeros") +} +\arguments{ +\item{in_channels}{Integer.} + +\item{stride}{Length-3 integer vector (t, h, w).} + +\item{residual}{Logical. Add the rearranged input as a residual.} + +\item{upscale_factor}{Integer.} + +\item{spatial_padding_mode}{Character.} +} +\description{ +Conv followed by channel-to-space/time rearrangement, with an optional +channel-repeat residual and an upscale factor that divides the conv +output channels. +} diff --git a/man/ltx23_video_vae.Rd b/man/ltx23_video_vae.Rd new file mode 100644 index 0000000..bc3226a --- /dev/null +++ b/man/ltx23_video_vae.Rd @@ -0,0 +1,44 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_video_vae} +\alias{ltx23_video_vae} +\title{LTX-2.3 video VAE} +\usage{ +ltx23_video_vae(in_channels = 3L, out_channels = 3L, latent_channels = 128L, + block_out_channels = c(256L, 512L, 1024L, 1024L), + decoder_block_out_channels = c(256L, 512L, 512L, 1024L), + layers_per_block = c(4L, 6L, 4L, 2L, 2L), + decoder_layers_per_block = c(4L, 6L, 4L, 2L, 2L), + spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + decoder_spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), + downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + upsample_type = c("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + upsample_residual = c(FALSE, FALSE, FALSE, FALSE), + upsample_factor = c(2L, 2L, 1L, 2L), patch_size = 4L, + patch_size_t = 1L, resnet_norm_eps = 1e-06, + encoder_causal = TRUE, decoder_causal = FALSE, + encoder_spatial_padding_mode = "zeros", + decoder_spatial_padding_mode = "zeros") +} +\arguments{ +\item{latent_channels}{Integer.} + +\item{resnet_norm_eps}{Numeric.} + +\item{in_channels,out_channels}{Integers. Pixel channels.} + +\item{block_out_channels,layers_per_block,spatio_temporal_scaling,downsample_type}{Encoder configuration (see \code{\link{ltx23_video_encoder3d}}).} + +\item{decoder_block_out_channels,decoder_layers_per_block,decoder_spatio_temporal_scaling,upsample_type,upsample_residual,upsample_factor}{Decoder configuration (see \code{\link{ltx23_video_decoder3d}}).} + +\item{patch_size,patch_size_t}{Integers. Pixel patchification.} + +\item{encoder_causal,decoder_causal}{Logicals. Temporal padding modes.} + +\item{encoder_spatial_padding_mode,decoder_spatial_padding_mode}{Characters.} +} +\description{ +Encoder + decoder + per-channel latent statistics (loaded from the +checkpoint's \code{per_channel_statistics}). The checkpoint's +\code{scaling_factor} is 1.0, so latent (de)normalization is purely +the per-channel affine map. +} diff --git a/man/ltx23_vocoder.Rd b/man/ltx23_vocoder.Rd new file mode 100644 index 0000000..e9ecaa1 --- /dev/null +++ b/man/ltx23_vocoder.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_vocoder} +\alias{ltx23_vocoder} +\title{LTX-2.3 vocoder stage} +\usage{ +ltx23_vocoder(in_channels = 128L, hidden_channels = 1536L, out_channels = 2L, + upsample_kernel_sizes = c(11L, 4L, 4L, 4L, 4L, 4L), + upsample_factors = c(5L, 2L, 2L, 2L, 2L, 2L), + resnet_kernel_sizes = c(3L, 7L, 11L), + resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), c(1L, 3L, 5L)), + antialias_ratio = 2L, antialias_kernel_size = 12L, + final_bias = FALSE) +} +\arguments{ +\item{in_channels}{Integer. Flattened input channels (C * mel bins / 1).} + +\item{hidden_channels}{Integer.} + +\item{out_channels}{Integer.} + +\item{resnet_kernel_sizes}{Integer vector.} + +\item{resnet_dilations}{List of integer vectors.} + +\item{final_bias}{Logical.} + +\item{upsample_kernel_sizes,upsample_factors}{Integer vectors.} + +\item{antialias_ratio,antialias_kernel_size}{Integers.} +} +\description{ +Mel spectrogram [B, C, T, M] -> waveform [B, out_channels, samples]. +Channel and mel dims are flattened into conv channels; each upsample +stage halves the channel count and averages three parallel ResNet +branches. +} diff --git a/man/ltx23_vocoder_resblock.Rd b/man/ltx23_vocoder_resblock.Rd new file mode 100644 index 0000000..8d5f7db --- /dev/null +++ b/man/ltx23_vocoder_resblock.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_vocoder_resblock} +\alias{ltx23_vocoder_resblock} +\title{Vocoder ResNet block (AMP)} +\usage{ +ltx23_vocoder_resblock(channels, kernel_size = 3L, dilations = c(1L, 3L, 5L), + antialias_ratio = 2L, antialias_kernel_size = 12L) +} +\arguments{ +\item{channels}{Integer.} + +\item{kernel_size}{Integer.} + +\item{dilations}{Integer vector.} + +\item{antialias_ratio,antialias_kernel_size}{Integers.} +} +\description{ +Dilated conv pairs, each preceded by an anti-aliased SnakeBeta +activation, with residual connections. +} diff --git a/man/ltx23_vocoder_with_bwe.Rd b/man/ltx23_vocoder_with_bwe.Rd new file mode 100644 index 0000000..81550b5 --- /dev/null +++ b/man/ltx23_vocoder_with_bwe.Rd @@ -0,0 +1,46 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_vocoder_with_bwe} +\alias{ltx23_vocoder_with_bwe} +\title{LTX-2.3 vocoder with bandwidth extension} +\usage{ +ltx23_vocoder_with_bwe(in_channels = 128L, hidden_channels = 1536L, + out_channels = 2L, + upsample_kernel_sizes = c(11L, 4L, 4L, 4L, 4L, 4L), + upsample_factors = c(5L, 2L, 2L, 2L, 2L, 2L), + resnet_kernel_sizes = c(3L, 7L, 11L), + resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), c(1L, 3L, 5L)), + bwe_in_channels = 128L, bwe_hidden_channels = 512L, + bwe_upsample_kernel_sizes = c(12L, 11L, 4L, 4L, 4L), + bwe_upsample_factors = c(6L, 5L, 2L, 2L, 2L), + bwe_resnet_kernel_sizes = c(3L, 7L, 11L), + bwe_resnet_dilations = list(c(1L, 3L, 5L), c(1L, 3L, 5L), c(1L, 3L, 5L)), + filter_length = 512L, hop_length = 80L, + window_length = 512L, num_mel_channels = 64L, + input_sampling_rate = 16000L, + output_sampling_rate = 48000L) +} +\arguments{ +\item{out_channels}{Integer. Audio channels.} + +\item{hop_length}{Integer. Mel analysis hop.} + +\item{in_channels,bwe_in_channels}{Integers. Flattened mel input channels.} + +\item{hidden_channels,bwe_hidden_channels}{Integers.} + +\item{upsample_kernel_sizes,upsample_factors,bwe_upsample_kernel_sizes,bwe_upsample_factors}{Integer vectors. Per-stage transposed-conv configs.} + +\item{resnet_kernel_sizes,bwe_resnet_kernel_sizes}{Integer vectors.} + +\item{resnet_dilations,bwe_resnet_dilations}{Lists of integer vectors.} + +\item{filter_length,window_length,num_mel_channels}{Integers. Mel +re-analysis configuration.} + +\item{input_sampling_rate,output_sampling_rate}{Integers.} +} +\description{ +Full mel [B, 2, T, 64] -> 48 kHz stereo waveform pipeline: 16 kHz +vocoder, causal mel re-analysis, BWE vocoder residual added to a +Hann-resampled skip path, clamped to [-1, 1]. +} diff --git a/man/ltx2_ada_layer_norm_single.Rd b/man/ltx2_ada_layer_norm_single.Rd deleted file mode 100644 index dd4e6b3..0000000 --- a/man/ltx2_ada_layer_norm_single.Rd +++ /dev/null @@ -1,15 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_ada_layer_norm_single} -\alias{ltx2_ada_layer_norm_single} -\title{LTX2 AdaLayerNorm Single} -\usage{ -ltx2_ada_layer_norm_single( - embedding_dim, - num_mod_params = 6L, - use_additional_conditions = FALSE -) -} -\description{ -LTX2 AdaLayerNorm Single -} -\keyword{internal} diff --git a/man/ltx2_attention.Rd b/man/ltx2_attention.Rd deleted file mode 100644 index 8d62fea..0000000 --- a/man/ltx2_attention.Rd +++ /dev/null @@ -1,24 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_attention} -\alias{ltx2_attention} -\title{LTX2 Attention module} -\usage{ -ltx2_attention( - query_dim, - heads = 8L, - kv_heads = 8L, - dim_head = 64L, - dropout = 0, - bias = TRUE, - cross_attention_dim = NULL, - out_bias = TRUE, - qk_norm = "rms_norm_across_heads", - norm_eps = 1e-06, - norm_elementwise_affine = TRUE, - rope_type = "interleaved" -) -} -\description{ -LTX2 Attention module -} -\keyword{internal} diff --git a/man/ltx2_audio_video_rotary_pos_embed.Rd b/man/ltx2_audio_video_rotary_pos_embed.Rd deleted file mode 100644 index 08c61bb..0000000 --- a/man/ltx2_audio_video_rotary_pos_embed.Rd +++ /dev/null @@ -1,27 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_audio_video_rotary_pos_embed} -\alias{ltx2_audio_video_rotary_pos_embed} -\title{LTX2 Audio-Video Rotary Positional Embeddings} -\usage{ -ltx2_audio_video_rotary_pos_embed( - dim, - patch_size = 1L, - patch_size_t = 1L, - base_num_frames = 20L, - base_height = 2048L, - base_width = 2048L, - sampling_rate = 16000L, - hop_length = 160L, - scale_factors = c(8L, 32L, 32L), - theta = 10000, - causal_offset = 1L, - modality = "video", - double_precision = TRUE, - rope_type = "interleaved", - num_attention_heads = 32L -) -} -\description{ -LTX2 Audio-Video Rotary Positional Embeddings -} -\keyword{internal} diff --git a/man/ltx2_connector_transformer_1d.Rd b/man/ltx2_connector_transformer_1d.Rd deleted file mode 100644 index b94e9d4..0000000 --- a/man/ltx2_connector_transformer_1d.Rd +++ /dev/null @@ -1,22 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_connector_transformer_1d} -\alias{ltx2_connector_transformer_1d} -\title{1D Connector Transformer for LTX2} -\usage{ -ltx2_connector_transformer_1d( - num_attention_heads = 30L, - attention_head_dim = 128L, - num_layers = 2L, - num_learnable_registers = 128L, - rope_base_seq_len = 4096L, - rope_theta = 10000, - rope_double_precision = TRUE, - eps = 1e-06, - causal_temporal_positioning = FALSE, - rope_type = "interleaved" -) -} -\description{ -1D Connector Transformer for LTX2 -} -\keyword{internal} diff --git a/man/ltx2_memory_profile.Rd b/man/ltx2_memory_profile.Rd deleted file mode 100644 index e183e9b..0000000 --- a/man/ltx2_memory_profile.Rd +++ /dev/null @@ -1,45 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_memory_profile} -\alias{ltx2_memory_profile} -\title{Get LTX-2 Memory Profile} -\usage{ -ltx2_memory_profile(vram_gb = NULL, model = "ltx2-19b-fp4") -} -\arguments{ -\item{vram_gb}{Numeric. Available VRAM in GB, or NULL for auto-detection.} - -\item{model}{Character. Model variant: "ltx2-19b-fp4" (default), "ltx2-19b-fp8", -or "ltx2-19b-distilled".} -} -\value{ -A list with memory profile settings. -} -\description{ -Determines optimal memory configuration based on available VRAM. -} -\details{ -LTX-2 is a 19B parameter model. Even with FP4 quantization (~10GB weights), -it requires careful memory management. The GPU-poor approach: - -1. Text encoding runs on CPU (cached) -2. DiT loaded in chunks, processed layer-by-layer, unloaded -3. VAE loaded after DiT unload, decode with tiling, unload - -Memory profiles: -\describe{ - \item{high}{16GB+ - FP4 DiT with chunk loading, VAE on GPU} - \item{medium}{12GB - FP4 DiT chunk loading, VAE tiled} - \item{low}{8GB - FP4 DiT layer-by-layer, VAE tiled small} - \item{very_low}{6GB - FP4 layer-by-layer, VAE on CPU} - \item{cpu_only}{All on CPU} -} -} -\examples{ -\dontrun{ -# Auto-detect profile -profile <- ltx2_memory_profile() - -# Specific VRAM -profile <- ltx2_memory_profile(vram_gb = 8) -} -} diff --git a/man/ltx2_rotary_pos_embed_1d.Rd b/man/ltx2_rotary_pos_embed_1d.Rd deleted file mode 100644 index 3ece718..0000000 --- a/man/ltx2_rotary_pos_embed_1d.Rd +++ /dev/null @@ -1,18 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_rotary_pos_embed_1d} -\alias{ltx2_rotary_pos_embed_1d} -\title{1D Rotary Position Embeddings for LTX2 Text Connectors} -\usage{ -ltx2_rotary_pos_embed_1d( - dim, - base_seq_len = 4096L, - theta = 10000, - double_precision = TRUE, - rope_type = "interleaved", - num_attention_heads = 32L -) -} -\description{ -1D Rotary Position Embeddings for LTX2 Text Connectors -} -\keyword{internal} diff --git a/man/ltx2_text_connectors.Rd b/man/ltx2_text_connectors.Rd deleted file mode 100644 index 8ae0ff4..0000000 --- a/man/ltx2_text_connectors.Rd +++ /dev/null @@ -1,60 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_text_connectors} -\alias{ltx2_text_connectors} -\title{LTX2 Text Connectors} -\usage{ -ltx2_text_connectors( - caption_channels = 3840L, - text_proj_in_factor = 49L, - video_connector_num_attention_heads = 30L, - video_connector_attention_head_dim = 128L, - video_connector_num_layers = 2L, - video_connector_num_learnable_registers = NULL, - audio_connector_num_attention_heads = 30L, - audio_connector_attention_head_dim = 128L, - audio_connector_num_layers = 2L, - audio_connector_num_learnable_registers = NULL, - connector_rope_base_seq_len = 4096L, - rope_theta = 10000, - rope_double_precision = TRUE, - causal_temporal_positioning = FALSE, - rope_type = "split" -) -} -\arguments{ -\item{caption_channels}{Integer. Dimension of caption embeddings (default 3840).} - -\item{text_proj_in_factor}{Integer. Factor for input projection (default 1).} - -\item{video_connector_num_attention_heads}{Integer. Number of attention heads for video connector.} - -\item{video_connector_attention_head_dim}{Integer. Attention head dimension for video.} - -\item{video_connector_num_layers}{Integer. Number of transformer layers for video.} - -\item{video_connector_num_learnable_registers}{Integer. Number of learnable registers for video.} - -\item{audio_connector_num_attention_heads}{Integer. Number of attention heads for audio connector.} - -\item{audio_connector_attention_head_dim}{Integer. Attention head dimension for audio.} - -\item{audio_connector_num_layers}{Integer. Number of transformer layers for audio.} - -\item{audio_connector_num_learnable_registers}{Integer. Number of learnable registers for audio.} - -\item{connector_rope_base_seq_len}{Integer. Base sequence length for RoPE.} - -\item{rope_theta}{Numeric. RoPE theta parameter.} - -\item{rope_double_precision}{Logical. Use double precision for RoPE.} - -\item{causal_temporal_positioning}{Logical. Use causal temporal positioning.} - -\item{rope_type}{Character. RoPE type ("interleaved" or "split").} -} -\value{ -nn_module for text connectors. -} -\description{ -Transforms packed text encoder hidden states for video and audio streams. -} diff --git a/man/ltx2_transformer_block_1d.Rd b/man/ltx2_transformer_block_1d.Rd deleted file mode 100644 index 2944e1e..0000000 --- a/man/ltx2_transformer_block_1d.Rd +++ /dev/null @@ -1,18 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_transformer_block_1d} -\alias{ltx2_transformer_block_1d} -\title{1D Transformer Block for LTX2 Text Connectors} -\usage{ -ltx2_transformer_block_1d( - dim, - num_attention_heads, - attention_head_dim, - activation_fn = "gelu-approximate", - eps = 1e-06, - rope_type = "interleaved" -) -} -\description{ -1D Transformer Block for LTX2 Text Connectors -} -\keyword{internal} diff --git a/man/ltx2_video_causal_conv3d.Rd b/man/ltx2_video_causal_conv3d.Rd deleted file mode 100644 index 5d28977..0000000 --- a/man/ltx2_video_causal_conv3d.Rd +++ /dev/null @@ -1,35 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_causal_conv3d} -\alias{ltx2_video_causal_conv3d} -\title{LTX2 Video Causal 3D Convolution} -\usage{ -ltx2_video_causal_conv3d( - in_channels, - out_channels, - kernel_size = 3L, - stride = 1L, - dilation = 1L, - groups = 1L, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer. Output channels.} - -\item{kernel_size}{Integer or vector of 3. Convolution kernel size.} - -\item{stride}{Integer or vector of 3. Stride.} - -\item{dilation}{Integer or vector of 3. Dilation.} - -\item{groups}{Integer. Convolution groups. Default: 1} - -\item{spatial_padding_mode}{Character. Padding mode for spatial dims. Default: "zeros"} -} -\description{ -3D convolution with runtime-selectable causal or non-causal padding. -Causal mode pads temporally by repeating first frame. -Non-causal mode pads temporally by repeating first and last frames. -} diff --git a/man/ltx2_video_decoder3d.Rd b/man/ltx2_video_decoder3d.Rd deleted file mode 100644 index 5b2c3a8..0000000 --- a/man/ltx2_video_decoder3d.Rd +++ /dev/null @@ -1,54 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_decoder3d} -\alias{ltx2_video_decoder3d} -\title{LTX2 Video Decoder} -\usage{ -ltx2_video_decoder3d( - in_channels = 128L, - out_channels = 3L, - block_out_channels = c(256L, 512L, 1024L), - spatio_temporal_scaling = c(TRUE, TRUE, TRUE), - layers_per_block = c(5L, 5L, 5L, 5L), - patch_size = 4L, - patch_size_t = 1L, - resnet_norm_eps = 1e-06, - is_causal = FALSE, - inject_noise = c(FALSE, FALSE, FALSE, FALSE), - timestep_conditioning = FALSE, - upsample_residual = c(TRUE, TRUE, TRUE), - upsample_factor = c(2L, 2L, 2L), - spatial_padding_mode = "reflect" -) -} -\arguments{ -\item{in_channels}{Integer. Latent channels.} - -\item{out_channels}{Integer. Output channels (typically 3 for RGB).} - -\item{block_out_channels}{Integer vector. Output channels per block.} - -\item{spatio_temporal_scaling}{Logical vector. Whether each block upscales.} - -\item{layers_per_block}{Integer vector. Number of layers per block.} - -\item{patch_size}{Integer. Spatial patch size.} - -\item{patch_size_t}{Integer. Temporal patch size.} - -\item{resnet_norm_eps}{Numeric. Epsilon for normalization.} - -\item{is_causal}{Logical. Whether to use causal convolutions.} - -\item{inject_noise}{Logical vector. Whether to inject noise per block.} - -\item{timestep_conditioning}{Logical. Whether to use timestep conditioning.} - -\item{upsample_residual}{Logical vector. Whether upsamplers use residual.} - -\item{upsample_factor}{Integer vector. Channel upscale factors.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Decodes latent representations back to video frames. -} diff --git a/man/ltx2_video_down_block3d.Rd b/man/ltx2_video_down_block3d.Rd deleted file mode 100644 index 6cbb8b0..0000000 --- a/man/ltx2_video_down_block3d.Rd +++ /dev/null @@ -1,39 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_down_block3d} -\alias{ltx2_video_down_block3d} -\title{LTX2 Video Down Block 3D} -\usage{ -ltx2_video_down_block3d( - in_channels, - out_channels = NULL, - num_layers = 1L, - dropout = 0, - resnet_eps = 1e-06, - resnet_act_fn = "swish", - spatio_temporal_scale = TRUE, - downsample_type = "conv", - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer or NULL. Output channels.} - -\item{num_layers}{Integer. Number of ResNet layers.} - -\item{dropout}{Numeric. Dropout rate.} - -\item{resnet_eps}{Numeric. Epsilon for normalization.} - -\item{resnet_act_fn}{Character. Activation function.} - -\item{spatio_temporal_scale}{Logical. Whether to use downsampling.} - -\item{downsample_type}{Character. Type of downsampling.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Encoder down block with multiple ResNet layers and optional downsampling. -} diff --git a/man/ltx2_video_encoder3d.Rd b/man/ltx2_video_encoder3d.Rd deleted file mode 100644 index a1eafc8..0000000 --- a/man/ltx2_video_encoder3d.Rd +++ /dev/null @@ -1,45 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_encoder3d} -\alias{ltx2_video_encoder3d} -\title{LTX2 Video Encoder} -\usage{ -ltx2_video_encoder3d( - in_channels = 3L, - out_channels = 128L, - block_out_channels = c(256L, 512L, 1024L, 2048L), - spatio_temporal_scaling = c(TRUE, TRUE, TRUE, TRUE), - layers_per_block = c(4L, 6L, 6L, 2L, 2L), - downsample_type = c("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - patch_size = 4L, - patch_size_t = 1L, - resnet_norm_eps = 1e-06, - is_causal = TRUE, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels (typically 3 for RGB).} - -\item{out_channels}{Integer. Latent channels.} - -\item{block_out_channels}{Integer vector. Output channels per block.} - -\item{spatio_temporal_scaling}{Logical vector. Whether each block downscales.} - -\item{layers_per_block}{Integer vector. Number of layers per block.} - -\item{downsample_type}{Character vector. Type of downsampling per block.} - -\item{patch_size}{Integer. Spatial patch size.} - -\item{patch_size_t}{Integer. Temporal patch size.} - -\item{resnet_norm_eps}{Numeric. Epsilon for normalization.} - -\item{is_causal}{Logical. Whether to use causal convolutions.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Encodes video frames into latent space with 3D causal convolutions. -} diff --git a/man/ltx2_video_mid_block3d.Rd b/man/ltx2_video_mid_block3d.Rd deleted file mode 100644 index 1d87390..0000000 --- a/man/ltx2_video_mid_block3d.Rd +++ /dev/null @@ -1,36 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_mid_block3d} -\alias{ltx2_video_mid_block3d} -\title{LTX2 Video Mid Block 3D} -\usage{ -ltx2_video_mid_block3d( - in_channels, - num_layers = 1L, - dropout = 0, - resnet_eps = 1e-06, - resnet_act_fn = "swish", - inject_noise = FALSE, - timestep_conditioning = FALSE, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{num_layers}{Integer. Number of ResNet layers.} - -\item{dropout}{Numeric. Dropout rate.} - -\item{resnet_eps}{Numeric. Epsilon for normalization.} - -\item{resnet_act_fn}{Character. Activation function.} - -\item{inject_noise}{Logical. Whether to inject noise.} - -\item{timestep_conditioning}{Logical. Whether to use timestep conditioning.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Middle block with ResNet layers and optional timestep conditioning. -} diff --git a/man/ltx2_video_resnet_block3d.Rd b/man/ltx2_video_resnet_block3d.Rd deleted file mode 100644 index 0d599e0..0000000 --- a/man/ltx2_video_resnet_block3d.Rd +++ /dev/null @@ -1,37 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_resnet_block3d} -\alias{ltx2_video_resnet_block3d} -\title{LTX2 Video ResNet Block 3D} -\usage{ -ltx2_video_resnet_block3d( - in_channels, - out_channels = NULL, - dropout = 0, - eps = 1e-06, - non_linearity = "silu", - inject_noise = FALSE, - timestep_conditioning = FALSE, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer or NULL. Output channels (defaults to in_channels).} - -\item{dropout}{Numeric. Dropout rate. Default: 0.0} - -\item{eps}{Numeric. Epsilon for normalization. Default: 1e-6} - -\item{non_linearity}{Character. Activation function. Default: "silu"} - -\item{inject_noise}{Logical. Whether to inject noise. Default: FALSE} - -\item{timestep_conditioning}{Logical. Whether to use timestep conditioning. Default: FALSE} - -\item{spatial_padding_mode}{Character. Padding mode. Default: "zeros"} -} -\description{ -3D ResNet block with per-channel RMS normalization and optional -noise injection and timestep conditioning. -} diff --git a/man/ltx2_video_transformer_3d_model.Rd b/man/ltx2_video_transformer_3d_model.Rd deleted file mode 100644 index 405a881..0000000 --- a/man/ltx2_video_transformer_3d_model.Rd +++ /dev/null @@ -1,107 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_transformer_3d_model} -\alias{ltx2_video_transformer_3d_model} -\title{LTX2 Video Transformer 3D Model (Audio-Video)} -\usage{ -ltx2_video_transformer_3d_model( - in_channels, - out_channels, - patch_size, - patch_size_t, - num_attention_heads, - attention_head_dim, - cross_attention_dim, - vae_scale_factors, - 32L, - 32L), - pos_embed_max_pos, - base_height, - base_width, - audio_in_channels, - audio_out_channels, - audio_patch_size, - audio_patch_size_t, - audio_num_attention_heads, - audio_attention_head_dim, - audio_cross_attention_dim -) -} -\arguments{ -\item{in_channels}{Integer. Video input channels (default: 128).} - -\item{out_channels}{Integer. Video output channels (default: 128).} - -\item{patch_size}{Integer. Spatial patch size (default: 1).} - -\item{patch_size_t}{Integer. Temporal patch size (default: 1).} - -\item{num_attention_heads}{Integer. Video attention heads (default: 32).} - -\item{attention_head_dim}{Integer. Video attention head dimension (default: 128).} - -\item{cross_attention_dim}{Integer. Video cross-attention dimension (default: 4096).} - -\item{vae_scale_factors}{Integer vector. VAE scale factors (default: c(8, 32, 32)).} - -\item{pos_embed_max_pos}{Integer. Max position for RoPE (default: 20).} - -\item{base_height}{Integer. Base height for RoPE (default: 2048).} - -\item{base_width}{Integer. Base width for RoPE (default: 2048).} - -\item{audio_in_channels}{Integer. Audio input channels (default: 128).} - -\item{audio_out_channels}{Integer. Audio output channels (default: 128).} - -\item{audio_patch_size}{Integer. Audio patch size (default: 1).} - -\item{audio_patch_size_t}{Integer. Audio temporal patch size (default: 1).} - -\item{audio_num_attention_heads}{Integer. Audio attention heads (default: 32).} - -\item{audio_attention_head_dim}{Integer. Audio head dimension (default: 64).} - -\item{audio_cross_attention_dim}{Integer. Audio cross-attention dim (default: 2048).} - -\item{audio_scale_factor}{Integer. Audio scale factor (default: 4).} - -\item{audio_pos_embed_max_pos}{Integer. Audio max position (default: 20).} - -\item{audio_sampling_rate}{Integer. Audio sampling rate (default: 16000).} - -\item{audio_hop_length}{Integer. Audio hop length (default: 160).} - -\item{num_layers}{Integer. Number of transformer layers (default: 48).} - -\item{activation_fn}{Character. Activation function (default: "gelu-approximate").} - -\item{qk_norm}{Character. QK normalization type (default: "rms_norm_across_heads").} - -\item{norm_elementwise_affine}{Logical. Use elementwise affine in norms (default: FALSE).} - -\item{norm_eps}{Numeric. Epsilon for normalization (default: 1e-6).} - -\item{caption_channels}{Integer. Caption embedding channels (default: 3840).} - -\item{attention_bias}{Logical. Use bias in attention (default: TRUE).} - -\item{attention_out_bias}{Logical. Use bias in attention output (default: TRUE).} - -\item{rope_theta}{Numeric. Theta for RoPE (default: 10000).} - -\item{rope_double_precision}{Logical. Use double precision for RoPE (default: TRUE).} - -\item{causal_offset}{Integer. Causal offset for RoPE (default: 1).} - -\item{timestep_scale_multiplier}{Numeric. Timestep scale (default: 1000).} - -\item{cross_attn_timestep_scale_multiplier}{Numeric. Cross-attn timestep scale (default: 1000).} - -\item{rope_type}{Character. RoPE type: "interleaved" or "split" (default: "interleaved").} -} -\value{ -An nn_module representing the LTX2 video transformer. -} -\description{ -Full audio-video transformer matching HuggingFace diffusers implementation. -} diff --git a/man/ltx2_video_transformer_block.Rd b/man/ltx2_video_transformer_block.Rd deleted file mode 100644 index a66683e..0000000 --- a/man/ltx2_video_transformer_block.Rd +++ /dev/null @@ -1,27 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_transformer_block} -\alias{ltx2_video_transformer_block} -\title{LTX2 Video Transformer Block (Audio-Video)} -\usage{ -ltx2_video_transformer_block( - dim, - num_attention_heads, - attention_head_dim, - cross_attention_dim, - audio_dim, - audio_num_attention_heads, - audio_attention_head_dim, - audio_cross_attention_dim, - qk_norm = "rms_norm_across_heads", - activation_fn = "gelu-approximate", - attention_bias = TRUE, - attention_out_bias = TRUE, - eps = 1e-06, - elementwise_affine = FALSE, - rope_type = "interleaved" -) -} -\description{ -LTX2 Video Transformer Block (Audio-Video) -} -\keyword{internal} diff --git a/man/ltx2_video_up_block3d.Rd b/man/ltx2_video_up_block3d.Rd deleted file mode 100644 index e3f7c70..0000000 --- a/man/ltx2_video_up_block3d.Rd +++ /dev/null @@ -1,48 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_up_block3d} -\alias{ltx2_video_up_block3d} -\title{LTX2 Video Up Block 3D} -\usage{ -ltx2_video_up_block3d( - in_channels, - out_channels = NULL, - num_layers = 1L, - dropout = 0, - resnet_eps = 1e-06, - resnet_act_fn = "swish", - spatio_temporal_scale = TRUE, - inject_noise = FALSE, - timestep_conditioning = FALSE, - upsample_residual = FALSE, - upscale_factor = 1L, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer or NULL. Output channels.} - -\item{num_layers}{Integer. Number of ResNet layers.} - -\item{dropout}{Numeric. Dropout rate.} - -\item{resnet_eps}{Numeric. Epsilon for normalization.} - -\item{resnet_act_fn}{Character. Activation function.} - -\item{spatio_temporal_scale}{Logical. Whether to use upsampling.} - -\item{inject_noise}{Logical. Whether to inject noise.} - -\item{timestep_conditioning}{Logical. Whether to use timestep conditioning.} - -\item{upsample_residual}{Logical. Whether upsampler uses residual.} - -\item{upscale_factor}{Integer. Channel upscale factor.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Decoder up block with upsampling and ResNet layers. -} diff --git a/man/ltx2_video_vae.Rd b/man/ltx2_video_vae.Rd deleted file mode 100644 index 0015e11..0000000 --- a/man/ltx2_video_vae.Rd +++ /dev/null @@ -1,101 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx2_video_vae} -\alias{ltx2_video_vae} -\title{LTX2 Video VAE} -\usage{ -ltx2_video_vae( - in_channels, - out_channels, - latent_channels, - block_out_channels, - 512L, - 1024L, - 2048L), - decoder_block_out_channels, - 512L, - 1024L), - layers_per_block, - 6L, - 6L, - 2L, - 2L), - decoder_layers_per_block, - 5L, - 5L, - 5L), - spatio_temporal_scaling, - TRUE, - TRUE, - TRUE), - decoder_spatio_temporal_scaling, - TRUE, - TRUE), - decoder_inject_noise, - FALSE, - FALSE, - FALSE), - downsample_type, - "temporal", - "spatiotemporal", - "spatiotemporal"), - upsample_residual, - TRUE, - TRUE), - upsample_factor, - 2L, - 2L), - timestep_conditioning, - patch_size, - patch_size_t, - resnet_norm_eps -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer. Output channels.} - -\item{latent_channels}{Integer. Latent space channels.} - -\item{block_out_channels}{Integer vector. Encoder block channels.} - -\item{decoder_block_out_channels}{Integer vector. Decoder block channels.} - -\item{layers_per_block}{Integer vector. Encoder layers per block.} - -\item{decoder_layers_per_block}{Integer vector. Decoder layers per block.} - -\item{spatio_temporal_scaling}{Logical vector. Encoder scaling.} - -\item{decoder_spatio_temporal_scaling}{Logical vector. Decoder scaling.} - -\item{decoder_inject_noise}{Logical vector. Noise injection per decoder block.} - -\item{downsample_type}{Character vector. Downsampling types.} - -\item{upsample_residual}{Logical vector. Upsampler residual flags.} - -\item{upsample_factor}{Integer vector. Upsampler factors.} - -\item{timestep_conditioning}{Logical. Whether to use timestep conditioning.} - -\item{patch_size}{Integer. Spatial patch size.} - -\item{patch_size_t}{Integer. Temporal patch size.} - -\item{resnet_norm_eps}{Numeric. Normalization epsilon.} - -\item{scaling_factor}{Numeric. Latent scaling factor.} - -\item{encoder_causal}{Logical. Encoder causality.} - -\item{decoder_causal}{Logical. Decoder causality.} - -\item{encoder_spatial_padding_mode}{Character. Encoder padding mode.} - -\item{decoder_spatial_padding_mode}{Character. Decoder padding mode.} -} -\description{ -Full VAE with encoder and decoder, supporting tiled encoding/decoding -for GPU-poor memory management. -} diff --git a/man/ltx_video_downsampler3d.Rd b/man/ltx_video_downsampler3d.Rd deleted file mode 100644 index c908f87..0000000 --- a/man/ltx_video_downsampler3d.Rd +++ /dev/null @@ -1,24 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx_video_downsampler3d} -\alias{ltx_video_downsampler3d} -\title{LTX Video Downsampler 3D} -\usage{ -ltx_video_downsampler3d( - in_channels, - out_channels, - stride = 1L, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{out_channels}{Integer. Output channels.} - -\item{stride}{Integer or vector of 3. Downsampling stride.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Spatiotemporal downsampling with strided pixel unshuffle + convolution. -} diff --git a/man/ltx_video_upsampler3d.Rd b/man/ltx_video_upsampler3d.Rd deleted file mode 100644 index 58f5176..0000000 --- a/man/ltx_video_upsampler3d.Rd +++ /dev/null @@ -1,27 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{ltx_video_upsampler3d} -\alias{ltx_video_upsampler3d} -\title{LTX Video Upsampler 3D} -\usage{ -ltx_video_upsampler3d( - in_channels, - stride = 1L, - residual = FALSE, - upscale_factor = 1L, - spatial_padding_mode = "zeros" -) -} -\arguments{ -\item{in_channels}{Integer. Input channels.} - -\item{stride}{Integer or vector of 3. Upsampling stride.} - -\item{residual}{Logical. Whether to use residual connection.} - -\item{upscale_factor}{Integer. Channel upscale factor.} - -\item{spatial_padding_mode}{Character. Padding mode.} -} -\description{ -Spatiotemporal upsampling with pixel shuffle + optional residual. -} diff --git a/man/make_linear.Rd b/man/make_linear.Rd deleted file mode 100644 index 0dc97d5..0000000 --- a/man/make_linear.Rd +++ /dev/null @@ -1,27 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{make_linear} -\alias{make_linear} -\title{Create Linear Layer (Standard or INT4)} -\usage{ -make_linear(in_features, out_features, bias = TRUE) -} -\arguments{ -\item{in_features}{Integer. Input dimension.} - -\item{out_features}{Integer. Output dimension.} - -\item{bias}{Logical. Include bias term.} -} -\value{ -nn_linear or int4_linear module. -} -\description{ -Factory function that creates either a standard nn_linear or INT4 linear layer -based on package options. Use this instead of torch::nn_linear() in model code. -} -\details{ -Behavior controlled by options: -- `diffuseR.use_int4`: If TRUE, create INT4 layer (default FALSE) -- `diffuseR.int4_device`: Device for INT4 layers (default "cuda") -- `diffuseR.int4_dtype`: Dtype for INT4 operations (default torch_float16()) -} diff --git a/man/memory_ltx23.Rd b/man/memory_ltx23.Rd new file mode 100644 index 0000000..21773c0 --- /dev/null +++ b/man/memory_ltx23.Rd @@ -0,0 +1,11 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{memory_ltx23} +\alias{memory_ltx23} +\title{LTX-2.3 Memory Profiles and CUDA GC Tuning} +\description{ +Memory management for running the 22B LTX-2.3 transformer on limited +VRAM, built on the patterns proven in the whisper and chatterbox +packages: torch allocator GC tuning before the first CUDA op, fp8 +CPU-resident streaming weights, query-chunked attention, and +phase-sequential component placement. +} diff --git a/man/models2devices.Rd b/man/models2devices.Rd index 4bee2f4..607e1b6 100644 --- a/man/models2devices.Rd +++ b/man/models2devices.Rd @@ -3,12 +3,8 @@ \alias{models2devices} \title{models2devices} \usage{ -models2devices( - model_name, - devices = "cpu", - unet_dtype_str = NULL, - download_models = FALSE -) +models2devices(model_name, devices = "cpu", unet_dtype_str = NULL, + download_models = FALSE) } \arguments{ \item{model_name}{A character string representing the name of the model to be used.} diff --git a/man/nf4_ltx23.Rd b/man/nf4_ltx23.Rd new file mode 100644 index 0000000..e2502f8 --- /dev/null +++ b/man/nf4_ltx23.Rd @@ -0,0 +1,13 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{nf4_ltx23} +\alias{nf4_ltx23} +\title{NF4 Weight Storage for the LTX-2.3 Transformer} +\description{ +4-bit NormalFloat quantization (the QLoRA scheme: per-block absmax +normalization against a 16-level quantile code, two indices packed +per byte). At ~4.5 bits/parameter the 22B transformer fits in about +12.5 GB, small enough to stay resident on a 16 GB GPU: no per-step +PCIe weight streaming, at a small quality cost relative to fp8. +Quantization and dequantization are pure torch ops (bucketize, +index_select) - no custom kernels. +} diff --git a/man/pack_text_embeds.Rd b/man/pack_text_embeds.Rd deleted file mode 100644 index aeacde2..0000000 --- a/man/pack_text_embeds.Rd +++ /dev/null @@ -1,34 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{pack_text_embeds} -\alias{pack_text_embeds} -\title{Pack Text Embeddings (Gemma-style)} -\usage{ -pack_text_embeds( - text_hidden_states, - sequence_lengths, - padding_side = "left", - scale_factor = 8, - eps = 1e-06, - device = "cpu" -) -} -\arguments{ -\item{text_hidden_states}{Tensor of shape [batch, seq_len, hidden_dim, num_layers].} - -\item{sequence_lengths}{Integer vector of valid sequence lengths per batch item.} - -\item{padding_side}{Character. "left" or "right".} - -\item{scale_factor}{Numeric. Scale factor for normalization (default 8).} - -\item{eps}{Numeric. Epsilon for numerical stability.} - -\item{device}{Character. Device for tensors.} -} -\value{ -Tensor of shape [batch, seq_len, hidden_dim * num_layers]. -} -\description{ -Normalizes and packs text encoder hidden states from multiple layers. -This is used when working with raw Gemma outputs. -} diff --git a/man/pack_video_latents.Rd b/man/pack_video_latents.Rd deleted file mode 100644 index 0669042..0000000 --- a/man/pack_video_latents.Rd +++ /dev/null @@ -1,21 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{pack_video_latents} -\alias{pack_video_latents} -\title{Pack Video Latents for DiT} -\usage{ -pack_video_latents(latents, patch_size = 1L, patch_size_t = 1L) -} -\arguments{ -\item{latents}{Tensor of shape \[B, C, F, H, W\].} - -\item{patch_size}{Integer. Spatial patch size (default 1).} - -\item{patch_size_t}{Integer. Temporal patch size (default 1).} -} -\value{ -Packed tensor of shape \[B, num_patches, C*patch_size_t*patch_size^2\]. -} -\description{ -Transforms 5D video latents \[B, C, F, H, W\] to 3D sequence \[B, F*H*W, C\] -for input to the DiT transformer. -} diff --git a/man/per_channel_rms_norm.Rd b/man/per_channel_rms_norm.Rd deleted file mode 100644 index 9c7b198..0000000 --- a/man/per_channel_rms_norm.Rd +++ /dev/null @@ -1,16 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{per_channel_rms_norm} -\alias{per_channel_rms_norm} -\title{Per-channel RMS normalization} -\usage{ -per_channel_rms_norm(channel_dim = 2L, eps = 1e-08) -} -\arguments{ -\item{channel_dim}{Integer. Dimension for RMS computation (1-indexed). Default: 2 (channels)} - -\item{eps}{Numeric. Small constant for numerical stability. Default: 1e-8} -} -\description{ -Normalizes tensor by root-mean-square along the channel dimension: -y = x / sqrt(mean(x^2, dim=channel_dim, keepdim=TRUE) + eps) -} diff --git a/man/pixart_alpha_combined_timestep_size_embeddings.Rd b/man/pixart_alpha_combined_timestep_size_embeddings.Rd deleted file mode 100644 index c665a34..0000000 --- a/man/pixart_alpha_combined_timestep_size_embeddings.Rd +++ /dev/null @@ -1,15 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{pixart_alpha_combined_timestep_size_embeddings} -\alias{pixart_alpha_combined_timestep_size_embeddings} -\title{PixArt Alpha Combined Timestep Size Embeddings} -\usage{ -pixart_alpha_combined_timestep_size_embeddings( - embedding_dim, - size_emb_dim, - use_additional_conditions = FALSE -) -} -\description{ -PixArt Alpha Combined Timestep Size Embeddings -} -\keyword{internal} diff --git a/man/pixart_alpha_text_projection.Rd b/man/pixart_alpha_text_projection.Rd deleted file mode 100644 index c9ec192..0000000 --- a/man/pixart_alpha_text_projection.Rd +++ /dev/null @@ -1,16 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{pixart_alpha_text_projection} -\alias{pixart_alpha_text_projection} -\title{PixArt Alpha Text Projection} -\usage{ -pixart_alpha_text_projection( - in_features, - hidden_size, - out_features = NULL, - act_fn = "gelu_tanh" -) -} -\description{ -PixArt Alpha Text Projection -} -\keyword{internal} diff --git a/man/quantize_int4.Rd b/man/quantize_int4.Rd deleted file mode 100644 index d9bbe43..0000000 --- a/man/quantize_int4.Rd +++ /dev/null @@ -1,36 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{quantize_int4} -\alias{quantize_int4} -\title{Quantize Tensor to INT4} -\usage{ -quantize_int4(x, block_size = 64L) -} -\arguments{ -\item{x}{Tensor. Input float tensor.} - -\item{block_size}{Integer. Number of values per scale factor (default 64).} -} -\value{ -A list with: - - `packed`: uint8 tensor with packed INT4 values - - `scales`: float tensor with per-block scale factors - - `orig_shape`: original tensor shape - - `orig_numel`: original number of elements - - `block_size`: block size used -} -\description{ -Quantizes a float tensor to 4-bit integers with block-wise scaling. -Two INT4 values are packed per byte for 7-8x compression. -} -\details{ -INT4 range is -8 to 7. Values are scaled per block, quantized, shifted to -unsigned (0-15), and packed two per byte. Compression is ~7x vs float32, -~3.5x vs float16. Typical reconstruction error is 10-12\% of std. -} -\examples{ -\dontrun{ -x <- torch_randn(c(4096, 4096)) * 0.02 -q <- quantize_int4(x) -x_back <- dequantize_int4(q) -} -} diff --git a/man/quantize_ltx2_transformer.Rd b/man/quantize_ltx2_transformer.Rd deleted file mode 100644 index 8ca3b4a..0000000 --- a/man/quantize_ltx2_transformer.Rd +++ /dev/null @@ -1,58 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{quantize_ltx2_transformer} -\alias{quantize_ltx2_transformer} -\title{Quantize LTX-2 Transformer to INT4} -\usage{ -quantize_ltx2_transformer( - model_id = "Lightricks/LTX-2", - output_dir = NULL, - block_size = 64L, - force = FALSE, - download = FALSE, - verbose = TRUE -) -} -\arguments{ -\item{model_id}{Character. HuggingFace model ID (default "Lightricks/LTX-2").} - -\item{output_dir}{Character. Directory to save quantized weights. -Default uses `tools::R_user_dir("diffuseR", "cache")`.} - -\item{block_size}{Integer. Block size for INT4 quantization (default 64).} - -\item{force}{Logical. Re-quantize even if cached file exists.} - -\item{download}{Logical. If TRUE, download model from HuggingFace if not cached.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -Character. Path to the quantized weights file. -} -\description{ -Downloads (if needed) and quantizes the LTX-2 19B transformer to INT4 format. -The quantized weights are cached for future use. -} -\details{ -LTX-2 is a 19B parameter model (~40GB in BF16). INT4 quantization reduces -this to ~5.7GB, fitting in 16GB VRAM with room for activations. - -The function: -1. Uses hfhub to locate/download the model from HuggingFace -2. Loads each safetensor shard -3. Quantizes all weights to INT4 (block-wise, 64 values per scale) -4. Saves as safetensors file -} -\examples{ -\dontrun{ -# Quantize and cache (first run takes ~10-20 minutes) -weights_path <- quantize_ltx2_transformer() - -# Load quantized weights -q <- load_int4_weights(weights_path) - -# Dequantize specific layer on GPU -layer_weight <- dequantize_int4(q[["transformer_blocks.0.attn1.to_q.weight"]], - device = "cuda", dtype = torch_float16()) -} -} diff --git a/man/quantize_ltx2_vae.Rd b/man/quantize_ltx2_vae.Rd deleted file mode 100644 index b02eb48..0000000 --- a/man/quantize_ltx2_vae.Rd +++ /dev/null @@ -1,34 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{quantize_ltx2_vae} -\alias{quantize_ltx2_vae} -\title{Quantize LTX-2 VAE to INT4} -\usage{ -quantize_ltx2_vae( - model_id = "Lightricks/LTX-2", - output_dir = NULL, - block_size = 64L, - force = FALSE, - download = FALSE, - verbose = TRUE -) -} -\arguments{ -\item{model_id}{Character. HuggingFace model ID (default "Lightricks/LTX-2").} - -\item{output_dir}{Character. Directory to save quantized weights. -Default uses `tools::R_user_dir("diffuseR", "cache")`.} - -\item{block_size}{Integer. Block size for INT4 quantization.} - -\item{force}{Logical. Re-quantize even if cached file exists.} - -\item{download}{Logical. If TRUE, download model from HuggingFace if not cached.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -Character. Path to the quantized weights file. -} -\description{ -Quantizes the LTX-2 video VAE to INT4 format. -} diff --git a/man/quantize_model_int4.Rd b/man/quantize_model_int4.Rd deleted file mode 100644 index 92c4c34..0000000 --- a/man/quantize_model_int4.Rd +++ /dev/null @@ -1,20 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{quantize_model_int4} -\alias{quantize_model_int4} -\title{Quantize Model Weights to INT4} -\usage{ -quantize_model_int4(module, block_size = 64L, verbose = TRUE) -} -\arguments{ -\item{module}{nn_module. The model to quantize.} - -\item{block_size}{Integer. Block size for quantization.} - -\item{verbose}{Logical. Print progress.} -} -\value{ -List of quantized parameters (does not modify original module). -} -\description{ -Quantizes all parameters in a torch module to INT4 format. -} diff --git a/man/quantize_safetensors_int4.Rd b/man/quantize_safetensors_int4.Rd deleted file mode 100644 index 68f842e..0000000 --- a/man/quantize_safetensors_int4.Rd +++ /dev/null @@ -1,31 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{quantize_safetensors_int4} -\alias{quantize_safetensors_int4} -\title{Quantize Safetensor Weights to INT4} -\usage{ -quantize_safetensors_int4(paths, output_path, block_size = 64L, verbose = TRUE) -} -\arguments{ -\item{paths}{Character vector. Paths to safetensor files.} - -\item{output_path}{Character. Path to save INT4 weights (.safetensors).} - -\item{block_size}{Integer. Block size for quantization (default 64).} - -\item{verbose}{Logical. Print progress.} -} -\value{ -Invisible NULL. Writes quantized weights to output_path. -} -\description{ -Loads weights from safetensors file(s) and quantizes to INT4. -Useful for quantizing large models without loading the full module. -} -\examples{ -\dontrun{ -# Quantize sharded transformer weights -paths <- list.files("~/.cache/huggingface/.../transformer", - pattern = "safetensors$", full.names = TRUE) -quantize_safetensors_int4(paths, "dit_int4.safetensors") -} -} diff --git a/man/rms_norm.Rd b/man/rms_norm.Rd deleted file mode 100644 index 13a1c29..0000000 --- a/man/rms_norm.Rd +++ /dev/null @@ -1,11 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{rms_norm} -\alias{rms_norm} -\title{RMS Normalization} -\usage{ -rms_norm(dim, eps = 1e-06, elementwise_affine = TRUE) -} -\description{ -RMS Normalization -} -\keyword{internal} diff --git a/man/rope.Rd b/man/rope.Rd deleted file mode 100644 index f9af843..0000000 --- a/man/rope.Rd +++ /dev/null @@ -1,14 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{rope} -\alias{rope} -\title{Rotary Position Embeddings (RoPE) for LTX-2 Video Models} -\description{ -Implementation of rotary positional embeddings for 3D video (spatiotemporal) -coordinates as used in LTX-2. RoPE encodes position information directly -into the attention queries and keys without adding position embeddings. -} -\references{ -Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). -"RoFormer: Enhanced Transformer with Rotary Position Embedding." -\url{https://arxiv.org/abs/2104.09864} -} diff --git a/man/rope_embedder_create.Rd b/man/rope_embedder_create.Rd deleted file mode 100644 index 1ac953a..0000000 --- a/man/rope_embedder_create.Rd +++ /dev/null @@ -1,54 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{rope_embedder_create} -\alias{rope_embedder_create} -\title{Create RoPE position embedder for video} -\usage{ -rope_embedder_create( - dim, - patch_size = 1L, - patch_size_t = 1L, - base_num_frames = 20L, - base_height = 2048L, - base_width = 2048L, - scale_factors = c(8, 32, 32), - theta = 10000, - causal_offset = 1L, - double_precision = TRUE, - rope_type = c("interleaved", "split"), - num_attention_heads = 32L -) -} -\arguments{ -\item{dim}{Integer. Dimension for RoPE (typically attention head dim * num_heads).} - -\item{patch_size}{Integer. Spatial patch size. Default: 1} - -\item{patch_size_t}{Integer. Temporal patch size. Default: 1} - -\item{base_num_frames}{Integer. Base number of frames for normalization. Default: 20} - -\item{base_height}{Integer. Base height for normalization. Default: 2048} - -\item{base_width}{Integer. Base width for normalization. Default: 2048} - -\item{scale_factors}{Numeric vector of length 3. VAE scale factors for -(temporal, height, width). Default: c(8, 32, 32)} - -\item{theta}{Numeric. Base frequency for RoPE. Default: 10000.0} - -\item{causal_offset}{Integer. Offset for causal VAE modeling. Default: 1} - -\item{double_precision}{Logical. Whether to use float64 for frequency -computation. Default: TRUE} - -\item{rope_type}{Character. Type of RoPE: "interleaved" or "split". Default: "interleaved"} - -\item{num_attention_heads}{Integer. Number of attention heads (for split RoPE). Default: 32} -} -\value{ -A list containing RoPE configuration and methods. -} -\description{ -Creates a RoPE embedder configured for video generation, handling -spatiotemporal coordinates (frames, height, width). -} diff --git a/man/rope_forward.Rd b/man/rope_forward.Rd deleted file mode 100644 index def96b5..0000000 --- a/man/rope_forward.Rd +++ /dev/null @@ -1,25 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{rope_forward} -\alias{rope_forward} -\title{Compute RoPE frequencies from coordinates} -\usage{ -rope_forward(embedder, coords, device = NULL) -} -\arguments{ -\item{embedder}{List. RoPE embedder configuration.} - -\item{coords}{torch tensor. Coordinate tensor from rope_prepare_video_coords().} - -\item{device}{Character or torch device. Device for output tensors.} -} -\value{ -A list with: - \describe{ - \item{cos_freqs}{Cosine frequencies tensor} - \item{sin_freqs}{Sine frequencies tensor} - } -} -\description{ -Converts spatiotemporal coordinates to (cos, sin) frequency tensors -for applying rotary embeddings. -} diff --git a/man/rope_ltx23.Rd b/man/rope_ltx23.Rd new file mode 100644 index 0000000..8f0479d --- /dev/null +++ b/man/rope_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{rope_ltx23} +\alias{rope_ltx23} +\title{LTX-2.3 Rotary Positional Embeddings} +\description{ +Fresh R port of the LTX rotary positional embedding scheme from the +diffusers reference implementation (Apache-2.0, +src/diffusers/models/transformers/transformer_ltx2.py). LTX 2.3 uses +the "split" RoPE layout everywhere; "interleaved" is kept for +completeness. Frequencies are computed in float64 per the checkpoint +config (frequencies_precision) and applied in float32. +} diff --git a/man/rope_prepare_video_coords.Rd b/man/rope_prepare_video_coords.Rd deleted file mode 100644 index 7765da9..0000000 --- a/man/rope_prepare_video_coords.Rd +++ /dev/null @@ -1,38 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{rope_prepare_video_coords} -\alias{rope_prepare_video_coords} -\title{Prepare video coordinates for RoPE} -\usage{ -rope_prepare_video_coords( - embedder, - batch_size, - num_frames, - height, - width, - device = "cpu", - fps = 24 -) -} -\arguments{ -\item{embedder}{List. RoPE embedder configuration.} - -\item{batch_size}{Integer. Batch size.} - -\item{num_frames}{Integer. Number of latent frames.} - -\item{height}{Integer. Latent height.} - -\item{width}{Integer. Latent width.} - -\item{device}{Character or torch device. Device for tensors.} - -\item{fps}{Numeric. Video frames per second. Default: 24.0} -} -\value{ -torch tensor of shape (batch_size, 3, num_patches, 2). -} -\description{ -Creates per-dimension patch boundaries for video coordinates in pixel space. -Returns tensor of shape (batch_size, 3, num_patches, 2) where dimension 1 -represents (frame, height, width) and dimension 3 represents (start, end). -} diff --git a/man/save_int4_weights.Rd b/man/save_int4_weights.Rd deleted file mode 100644 index 7457673..0000000 --- a/man/save_int4_weights.Rd +++ /dev/null @@ -1,46 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{save_int4_weights} -\alias{save_int4_weights} -\title{Save INT4 Quantized Weights} -\usage{ -save_int4_weights( - quantized_params, - path, - max_shard_size = 2e+09, - verbose = TRUE -) -} -\arguments{ -\item{quantized_params}{List of quantized parameters from `quantize_model_int4()`.} - -\item{path}{Character. Base path for safetensors files. If multiple shards needed, -files will be named `path-00001-of-NNNNN.safetensors`.} - -\item{max_shard_size}{Numeric. Maximum bytes per shard (default 2GB to avoid R -integer overflow issues).} - -\item{verbose}{Logical. Print progress.} -} -\value{ -Invisible character vector of saved file paths. -} -\description{ -Saves INT4 quantized weights to disk as sharded safetensors files. -} -\details{ -Weights are saved in safetensors format with the following structure: -\itemize{ - \item `{name}::packed` - uint8 tensor with packed INT4 values - \item `{name}::scales` - float32 tensor with per-block scales - \item `{name}::shape` - int64 tensor with original shape -} - -Large models are automatically sharded to avoid R's 2GB vector limit. -The block size is fixed at 64 (standard for INT4 quantization). -} -\examples{ -\dontrun{ -q <- quantize_model_int4(model) -save_int4_weights(q, "model_int4.safetensors") -} -} diff --git a/man/save_video.Rd b/man/save_video.Rd index 63c9f3b..8608d49 100644 --- a/man/save_video.Rd +++ b/man/save_video.Rd @@ -3,15 +3,8 @@ \alias{save_video} \title{Save Video to File} \usage{ -save_video( - video, - file, - fps = 24, - format = NULL, - backend = "auto", - quality = 85, - verbose = TRUE -) +save_video(video, file, fps = 24, format = NULL, backend = "auto", + quality = 85, verbose = TRUE) } \arguments{ \item{video}{Array of video frames with shape [T, H, W, C] where C is 3 (RGB). diff --git a/man/save_video_ffmpeg.Rd b/man/save_video_ffmpeg.Rd index 71acd55..03ceca9 100644 --- a/man/save_video_ffmpeg.Rd +++ b/man/save_video_ffmpeg.Rd @@ -3,14 +3,8 @@ \alias{save_video_ffmpeg} \title{Save Video using FFmpeg} \usage{ -save_video_ffmpeg( - video, - file, - fps = 24, - format = "mp4", - quality = 85, - verbose = TRUE -) +save_video_ffmpeg(video, file, fps = 24, format = "mp4", quality = 85, + verbose = TRUE) } \arguments{ \item{video}{Array of video frames [T, H, W, C].} diff --git a/man/save_video_frames.Rd b/man/save_video_frames.Rd deleted file mode 100644 index 1ae8832..0000000 --- a/man/save_video_frames.Rd +++ /dev/null @@ -1,24 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{save_video_frames} -\alias{save_video_frames} -\title{Save video frames to file} -\usage{ -save_video_frames(video_array, output_file, fps = 24, verbose = TRUE) -} -\arguments{ -\item{video_array}{Numeric array with dimensions [frames, height, width, channels]. -Values should be in range [0, 1].} - -\item{output_file}{Character. Path to output video file.} - -\item{fps}{Numeric. Frames per second. Default: 24.} - -\item{verbose}{Logical. Print progress messages. Default: TRUE.} -} -\value{ -Invisibly returns the output file path. -} -\description{ -Saves an array of video frames to an MP4 file using ffmpeg via the av package. -} -\keyword{internal} diff --git a/man/save_video_ltx23.Rd b/man/save_video_ltx23.Rd new file mode 100644 index 0000000..21eb7aa --- /dev/null +++ b/man/save_video_ltx23.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{save_video_ltx23} +\alias{save_video_ltx23} +\title{Save an LTX video (optionally with audio) to MP4} +\usage{ +save_video_ltx23(video, filename, fps = 24, audio = NULL, sample_rate = 48000L, + verbose = TRUE) +} +\arguments{ +\item{video}{Array [frames, height, width, 3] in [0, 1].} + +\item{filename}{Output path (.mp4).} + +\item{fps}{Numeric.} + +\item{audio}{Optional numeric matrix [channels, samples] in [-1, 1].} + +\item{sample_rate}{Integer.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the filename. +} +\description{ +Uses the av package (Suggests) to encode frames and mux the audio +track. +} diff --git a/man/sequential_cfg_forward.Rd b/man/sequential_cfg_forward.Rd deleted file mode 100644 index 72e2d3d..0000000 --- a/man/sequential_cfg_forward.Rd +++ /dev/null @@ -1,46 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{sequential_cfg_forward} -\alias{sequential_cfg_forward} -\title{Sequential CFG Forward Pass} -\usage{ -sequential_cfg_forward( - model, - latents, - timestep, - prompt_embeds, - negative_prompt_embeds, - guidance_scale, - ... -) -} -\arguments{ -\item{model}{The DiT model.} - -\item{latents}{Current latent tensor.} - -\item{timestep}{Current timestep tensor.} - -\item{prompt_embeds}{Conditional prompt embeddings.} - -\item{negative_prompt_embeds}{Unconditional prompt embeddings.} - -\item{guidance_scale}{CFG scale.} - -\item{...}{Additional arguments to model forward pass.} -} -\value{ -The CFG-combined noise prediction. -} -\description{ -Runs unconditional and conditional forward passes separately to halve -peak activation memory. For GPU-poor scenarios. -} -\examples{ -\dontrun{ -noise_pred <- sequential_cfg_forward( - model, latents, timestep, - prompt_embeds, negative_prompt_embeds, - guidance_scale = 4.0 -) -} -} diff --git a/man/staging_ltx23.Rd b/man/staging_ltx23.Rd new file mode 100644 index 0000000..7b2e245 --- /dev/null +++ b/man/staging_ltx23.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{staging_ltx23} +\alias{staging_ltx23} +\title{Pinned Staging for Phase-Sequential Components} +\description{ +Phase offloading moves each large component (transformer, +connectors, VAEs, vocoder) between CPU and GPU every render. From +pageable memory those copies run through the driver's bounce +buffer at a fraction of PCIe speed; page-locked (pinned) host +memory transfers by DMA at full rate. Each component's parameters +and buffers are pinned once at load; onload swaps every tensor to +a non-blocking GPU copy of its pinned source, and offload simply +re-points the tensors at the still-valid pinned copies — weights +are immutable during inference, so offload moves no bytes at all. +} +\details{ +Costs: the pinned copies are non-swappable host RAM for the life +of the pipeline (about the model's CPU footprint), and page-locking +~20GB adds ~30s to pipeline load. Measured per render the win is +small (~2.7s: warm pinned onload 0.54s vs pageable 0.99s, offload +0.06s vs 2.36s) because the real phase-transition cost was +allocator pool regrowth, fixed separately by the loader's pool +pre-warm and by not emptying the CUDA cache between phases. Off by +default; enable for long multi-render sessions with +\code{options(diffuseR.pin_staging = TRUE)} before +\code{ltx23_load_pipeline} (breaks even after ~11 renders). + +} diff --git a/man/text_encoder2_native.Rd b/man/text_encoder2_native.Rd index a07a7c7..9661799 100644 --- a/man/text_encoder2_native.Rd +++ b/man/text_encoder2_native.Rd @@ -3,14 +3,8 @@ \alias{text_encoder2_native} \title{Native CLIP Text Encoder 2 (OpenCLIP ViT-bigG for SDXL)} \usage{ -text_encoder2_native( - vocab_size = 49408, - context_length = 77, - embed_dim = 1280, - num_layers = 32, - num_heads = 20, - mlp_dim = 5120 -) +text_encoder2_native(vocab_size = 49408, context_length = 77, embed_dim = 1280, + num_layers = 32, num_heads = 20, mlp_dim = 5120) } \arguments{ \item{vocab_size}{Vocabulary size (default 49408)} diff --git a/man/text_encoder_native.Rd b/man/text_encoder_native.Rd index 0c0fec7..2279664 100644 --- a/man/text_encoder_native.Rd +++ b/man/text_encoder_native.Rd @@ -3,15 +3,9 @@ \alias{text_encoder_native} \title{Native CLIP Text Encoder} \usage{ -text_encoder_native( - vocab_size = 49408, - context_length = 77, - embed_dim = 768, - num_layers = 12, - num_heads = 12, - mlp_dim = 3072, - apply_final_ln = TRUE -) +text_encoder_native(vocab_size = 49408, context_length = 77, embed_dim = 768, + num_layers = 12, num_heads = 12, mlp_dim = 3072, + apply_final_ln = TRUE) } \arguments{ \item{vocab_size}{Vocabulary size (default 49408)} diff --git a/man/timestep_embedding.Rd b/man/timestep_embedding.Rd index cdc0bd8..f357184 100644 --- a/man/timestep_embedding.Rd +++ b/man/timestep_embedding.Rd @@ -3,12 +3,8 @@ \alias{timestep_embedding} \title{Sinusoidal Timestep Embedding} \usage{ -timestep_embedding( - timesteps, - dim, - flip_sin_to_cos = TRUE, - downscale_freq_shift = 0L -) +timestep_embedding(timesteps, dim, flip_sin_to_cos = TRUE, + downscale_freq_shift = 0L) } \arguments{ \item{timesteps}{Tensor of timesteps (batch_size,)} diff --git a/man/timestep_embedding_module.Rd b/man/timestep_embedding_module.Rd deleted file mode 100644 index 3fd7221..0000000 --- a/man/timestep_embedding_module.Rd +++ /dev/null @@ -1,16 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{timestep_embedding_module} -\alias{timestep_embedding_module} -\title{Timestep embedding MLP} -\usage{ -timestep_embedding_module( - in_channels, - time_embed_dim, - act_fn = "silu", - out_dim = NULL -) -} -\description{ -Timestep embedding MLP -} -\keyword{internal} diff --git a/man/timesteps_module.Rd b/man/timesteps_module.Rd deleted file mode 100644 index fdd9170..0000000 --- a/man/timesteps_module.Rd +++ /dev/null @@ -1,16 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{timesteps_module} -\alias{timesteps_module} -\title{Timesteps module} -\usage{ -timesteps_module( - num_channels, - flip_sin_to_cos = TRUE, - downscale_freq_shift = 0, - scale = 1L -) -} -\description{ -Timesteps module -} -\keyword{internal} diff --git a/man/tokenize_gemma3.Rd b/man/tokenize_gemma3.Rd index 98c5a34..c939dd4 100644 --- a/man/tokenize_gemma3.Rd +++ b/man/tokenize_gemma3.Rd @@ -3,13 +3,8 @@ \alias{tokenize_gemma3} \title{Tokenize text for Gemma3} \usage{ -tokenize_gemma3( - tokenizer, - text, - max_length = 1024L, - padding = "max_length", - return_tensors = "pt" -) +tokenize_gemma3(tokenizer, text, max_length = 1024L, padding = "max_length", + return_tensors = "pt") } \arguments{ \item{tokenizer}{Gemma3 tokenizer object.} diff --git a/man/txt2img_sd21.Rd b/man/txt2img_sd21.Rd index 55d3332..90362f1 100644 --- a/man/txt2img_sd21.Rd +++ b/man/txt2img_sd21.Rd @@ -3,28 +3,13 @@ \alias{txt2img_sd21} \title{Generate an image from a text prompt using a diffusion pipeline} \usage{ -txt2img_sd21( - prompt, - negative_prompt = NULL, - img_dim = 768, - pipeline = NULL, - devices = "auto", - unet_dtype_str = NULL, - download_models = FALSE, - scheduler = "ddim", - timesteps = NULL, - initial_latents = NULL, - num_inference_steps = 50, - guidance_scale = 7.5, - seed = NULL, - save_file = TRUE, - filename = NULL, - metadata_path = NULL, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, - ... -) +txt2img_sd21(prompt, negative_prompt = NULL, img_dim = 768, pipeline = NULL, + devices = "auto", unet_dtype_str = NULL, download_models = FALSE, + scheduler = "ddim", timesteps = NULL, initial_latents = NULL, + num_inference_steps = 50, guidance_scale = 7.5, seed = NULL, + save_file = TRUE, filename = NULL, metadata_path = NULL, + use_native_decoder = FALSE, use_native_text_encoder = FALSE, + use_native_unet = FALSE, ...) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} diff --git a/man/txt2img_sdxl.Rd b/man/txt2img_sdxl.Rd index c514f10..bc133d6 100644 --- a/man/txt2img_sdxl.Rd +++ b/man/txt2img_sdxl.Rd @@ -3,30 +3,14 @@ \alias{txt2img_sdxl} \title{Generate an image from a text prompt using SDXL} \usage{ -txt2img_sdxl( - prompt, - negative_prompt = NULL, - img_dim = 1024, - pipeline = NULL, - devices = "auto", - memory_profile = NULL, - unet_dtype_str = NULL, - download_models = FALSE, - scheduler = "ddim", - timesteps = NULL, - initial_latents = NULL, - num_inference_steps = 30, - guidance_scale = 7.5, - seed = NULL, - save_file = TRUE, - filename = NULL, - metadata_path = NULL, - use_native_decoder = FALSE, - use_native_text_encoder = FALSE, - use_native_unet = FALSE, - verbose = TRUE, - ... -) +txt2img_sdxl(prompt, negative_prompt = NULL, img_dim = 1024, pipeline = NULL, + devices = "auto", memory_profile = NULL, unet_dtype_str = NULL, + download_models = FALSE, scheduler = "ddim", timesteps = NULL, + initial_latents = NULL, num_inference_steps = 30, + guidance_scale = 7.5, seed = NULL, save_file = TRUE, + filename = NULL, metadata_path = NULL, use_native_decoder = FALSE, + use_native_text_encoder = FALSE, use_native_unet = FALSE, + verbose = TRUE, ...) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} diff --git a/man/txt2vid_ltx2.Rd b/man/txt2vid_ltx2.Rd index 83c0387..4f608ae 100644 --- a/man/txt2vid_ltx2.Rd +++ b/man/txt2vid_ltx2.Rd @@ -1,96 +1,113 @@ % tinyrox says don't edit this manually, but it can't stop you! \name{txt2vid_ltx2} \alias{txt2vid_ltx2} -\title{Generate Video from Text Prompt using LTX-2} +\title{Generate video (and audio) with LTX-2.3} \usage{ -txt2vid_ltx2( - prompt, - negative_prompt = NULL, - width = 768L, - height = 512L, - num_frames = 121L, - fps = 24, - num_inference_steps = 8L, - guidance_scale = 4, - memory_profile = "auto", - text_backend = "gemma3", - text_model_path = NULL, - text_api_url = NULL, - vae = NULL, - dit = NULL, - connectors = NULL, - seed = NULL, - output_file = NULL, - output_format = "mp4", - return_latents = FALSE, - verbose = TRUE -) +txt2vid_ltx2(prompt, pipeline, text_encoder = NULL, tokenizer = NULL, + prompt_embeds = NULL, width = 768L, height = 512L, + num_frames = 121L, frame_rate = 24, + sigmas = ltx23_distilled_sigmas(), guidance_scale = 1, + seed = NULL, device = "cuda", dtype = "bfloat16", filename = NULL, + max_sequence_length = 1024L, decode_video = TRUE, + decode_audio = TRUE, two_stage = FALSE, upsampler = NULL, + adain_factor = 1, tone_map_compression = 0, phase_offload = TRUE, + image = NULL, condition_video = NULL, conditioning_frames = 9L, + cond_noise_scale = 0, audio = NULL, verbose = TRUE) } \arguments{ -\item{prompt}{Character. Text prompt describing the video to generate.} +\item{prompt}{Character. The prompt.} -\item{negative_prompt}{Character. Optional negative prompt.} +\item{pipeline}{An \code{ltx23_pipeline} from +\code{\link{ltx23_load_pipeline}}.} -\item{width}{Integer. Video width in pixels (default 768).} +\item{prompt_embeds}{Optional precomputed list with +\code{prompt_embeds} (raw stacked Gemma3 states) and +\code{prompt_attention_mask}; bypasses the text encoder.} -\item{height}{Integer. Video height in pixels (default 512).} +\item{num_frames}{Integer. 8k + 1 frames (e.g. 121).} -\item{num_frames}{Integer. Number of frames to generate (default 121).} +\item{frame_rate}{Numeric. Frames per second.} -\item{fps}{Numeric. Frames per second (default 24).} +\item{sigmas}{Numeric vector. Denoising schedule (default: official +distilled schedule; must end in 0).} -\item{num_inference_steps}{Integer. Number of denoising steps (default 8 for distilled).} +\item{guidance_scale}{Numeric. Only 1 (no CFG) is supported; the +distilled checkpoints are trained for CFG-free sampling.} -\item{guidance_scale}{Numeric. CFG scale (default 4.0).} +\item{seed}{Integer or NULL.} -\item{memory_profile}{Character or list. Memory profile: "auto" for auto-detection, -or a profile from `ltx2_memory_profile()`.} +\item{device}{Character. Compute device for the denoising loop.} -\item{text_backend}{Character. Text encoding backend: "gemma3" (native), "api", "precomputed", or "random".} +\item{dtype}{Character. Model compute dtype ("bfloat16" or "float32").} -\item{text_model_path}{Character. Path to Gemma3 model (for "gemma3" backend). Supports glob patterns.} +\item{filename}{Character or NULL. Output video path (.mp4). Audio is +muxed in when the av package is available.} -\item{text_api_url}{Character. URL for text encoding API (if backend = "api").} +\item{max_sequence_length}{Integer. Text token length (multiple of 128).} -\item{vae}{Optional. Pre-loaded VAE module.} +\item{two_stage}{Logical. Generate at half resolution, upsample the +latents 2x spatially, and refine over the stage-2 schedule +(resolution must then be a multiple of 64; requires +\code{upsampler}).} -\item{dit}{Optional. Pre-loaded DiT transformer module.} +\item{upsampler}{An \code{\link{ltx23_latent_upsampler}} (see +\code{\link{ltx23_load_upsampler}}).} -\item{connectors}{Optional. Pre-loaded text connectors module.} +\item{adain_factor}{Numeric. AdaIN blend of the upsampled latents +toward the stage-1 statistics (0 disables).} -\item{seed}{Integer. Random seed for reproducibility.} +\item{tone_map_compression}{Numeric in [0, 1]. Optional latent tone +mapping before stage 2.} -\item{output_file}{Character. Path to save output video (NULL for no save).} +\item{phase_offload}{Logical. Move each small component to the compute +device only for its phase (text encoding, upsampling, decoding) and +back to the CPU afterwards, keeping the denoise phase as the sole +GPU tenant.} -\item{output_format}{Character. Output format: "mp4", "gif", or "frames".} +\item{image}{Optional start image for image-to-video: a PNG/JPEG path +or an [H, W, 3] array in [0, 1]. The image conditions the first +frame; the rest of the video is generated (reference i2v).} -\item{return_latents}{Logical. If TRUE, also return final latents.} +\item{condition_video}{Optional continuation source: a video path +(its trailing \code{conditioning_frames} frames are used) or an +[F, H, W, 3] array. The clip's tail becomes the frozen prefix of +the new video, so the output's first \code{conditioning_frames} +frames overlap the source (trim or crossfade when concatenating).} -\item{verbose}{Logical. Print progress messages.} +\item{conditioning_frames}{Integer. Trailing pixel frames taken from +\code{condition_video} (8k + 1, default 9 = 2 latent frames).} + +\item{cond_noise_scale}{Numeric in [0, 1]. Optional partial noising +of the conditioned tokens (0 = keep them exactly).} + +\item{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 +clean, frozen audio latents that the video attends to while +denoising, and the original samples are muxed into the output +(audio decoding is skipped).} + +\item{verbose}{Logical.} + +\item{text_encoder,tokenizer}{Gemma3 model and tokenizer (or paths; +see \code{\link{load_gemma3_text_encoder}} and +\code{\link{gemma3_tokenizer}}). Ignored when \code{prompt_embeds} +is supplied.} + +\item{width,height}{Integers. Output resolution (multiples of 32).} + +\item{decode_video,decode_audio}{Logicals. Decode the respective +latents (disable for latent-space work).} } \value{ -A list with: - - `video`: Array of video frames [frames, height, width, channels] - - `latents`: (if return_latents=TRUE) Final latent tensor - - `metadata`: Generation metadata +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. } \description{ -Generates video using the LTX-2 diffusion transformer model. -} -\examples{ -\dontrun{ -# Basic usage -result <- txt2vid_ltx2("A cat walking on a beach at sunset") - -# With specific settings -result <- txt2vid_ltx2( - prompt = "A timelapse of clouds moving over mountains", - width = 512, - height = 320, - num_frames = 61, - num_inference_steps = 8, - seed = 42, - output_file = "clouds.mp4" -) -} +Distilled text-to-video generation: encodes the prompt with Gemma3 + +connectors, denoises joint audio/video latents over the official +8-step distilled schedule (no classifier-free guidance), decodes the +video with the causal VAE and the audio with the audio VAE + BWE +vocoder, and optionally muxes both into an MP4. } diff --git a/man/txt2vid_ltx23.Rd b/man/txt2vid_ltx23.Rd new file mode 100644 index 0000000..3c20e62 --- /dev/null +++ b/man/txt2vid_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{txt2vid_ltx23} +\alias{txt2vid_ltx23} +\title{LTX-2.3 Text-to-Video Pipeline} +\description{ +Fresh R port of the LTX-2 text-to-video flow from the diffusers +reference (Apache-2.0, pipelines/ltx2/pipeline_ltx2.py), specialized +for the distilled LTX 2.3 checkpoints: 8-step official sigma schedule, +no classifier-free guidance, joint audio-video denoising with an Euler +velocity step, and audio decoding through the audio VAE and BWE +vocoder to 48 kHz stereo. +} diff --git a/man/unet_native.Rd b/man/unet_native.Rd index df9d8a0..d5f70f9 100644 --- a/man/unet_native.Rd +++ b/man/unet_native.Rd @@ -3,14 +3,10 @@ \alias{unet_native} \title{Native UNet for Stable Diffusion} \usage{ -unet_native( - in_channels = 4L, - out_channels = 4L, - block_out_channels = c(320L, 640L, 1280L, 1280L), - layers_per_block = 2L, - cross_attention_dim = 1024L, - attention_head_dim = 64L -) +unet_native(in_channels = 4L, out_channels = 4L, + block_out_channels = c(320L, 640L, 1280L, 1280L), + layers_per_block = 2L, cross_attention_dim = 1024L, + attention_head_dim = 64L) } \arguments{ \item{in_channels}{Input channels (default 4 for latent space)} diff --git a/man/unet_sdxl_native.Rd b/man/unet_sdxl_native.Rd index 2f01ca8..4e6668a 100644 --- a/man/unet_sdxl_native.Rd +++ b/man/unet_sdxl_native.Rd @@ -3,17 +3,12 @@ \alias{unet_sdxl_native} \title{Native SDXL UNet} \usage{ -unet_sdxl_native( - in_channels = 4L, - out_channels = 4L, - block_out_channels = c(320L, 640L, 1280L), - layers_per_block = 2L, - transformer_layers_per_block = c(0L, 2L, 10L), - cross_attention_dim = 2048L, - attention_head_dim = 64L, - addition_embed_dim = 1280L, - addition_time_embed_dim = 256L -) +unet_sdxl_native(in_channels = 4L, out_channels = 4L, + block_out_channels = c(320L, 640L, 1280L), + layers_per_block = 2L, + transformer_layers_per_block = c(0L, 2L, 10L), + cross_attention_dim = 2048L, attention_head_dim = 64L, + addition_embed_dim = 1280L, addition_time_embed_dim = 256L) } \arguments{ \item{in_channels}{Input channels (default 4 for latent space)} diff --git a/man/unpack_video_latents.Rd b/man/unpack_video_latents.Rd deleted file mode 100644 index 526a14f..0000000 --- a/man/unpack_video_latents.Rd +++ /dev/null @@ -1,33 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{unpack_video_latents} -\alias{unpack_video_latents} -\title{Unpack Video Latents from DiT} -\usage{ -unpack_video_latents( - latents, - num_frames, - height, - width, - patch_size = 1L, - patch_size_t = 1L -) -} -\arguments{ -\item{latents}{Tensor of shape \[B, num_patches, D\].} - -\item{num_frames}{Integer. Target number of latent frames.} - -\item{height}{Integer. Target latent height.} - -\item{width}{Integer. Target latent width.} - -\item{patch_size}{Integer. Spatial patch size (default 1).} - -\item{patch_size_t}{Integer. Temporal patch size (default 1).} -} -\value{ -Unpacked tensor of shape \[B, C, F, H, W\]. -} -\description{ -Transforms 3D sequence \[B, num_patches, D\] back to 5D video latents \[B, C, F, H, W\]. -} diff --git a/man/upsampler_ltx23.Rd b/man/upsampler_ltx23.Rd new file mode 100644 index 0000000..7f56376 --- /dev/null +++ b/man/upsampler_ltx23.Rd @@ -0,0 +1,11 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{upsampler_ltx23} +\alias{upsampler_ltx23} +\title{LTX-2.3 Spatial Latent Upsampler} +\description{ +Fresh R port of the LTX latent upsampler from the diffusers reference +(Apache-2.0, pipelines/ltx2/latent_upsampler.py and +pipeline_ltx2_latent_upsample.py), with the LTX 2.3 configuration: +Conv3d ResBlock stages around a per-frame 2x pixel-shuffle spatial +upsampler (no rational resampler). Operates on unnormalized latents. +} diff --git a/man/vae_ltx2.Rd b/man/vae_ltx2.Rd deleted file mode 100644 index 1436d72..0000000 --- a/man/vae_ltx2.Rd +++ /dev/null @@ -1,8 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{vae_ltx2} -\alias{vae_ltx2} -\title{LTX2 Video VAE} -\description{ -3D causal VAE for video encoding/decoding as used in LTX-2. -Supports GPU-poor tiling for memory-efficient processing. -} diff --git a/man/vae_ltx23.Rd b/man/vae_ltx23.Rd new file mode 100644 index 0000000..241fc82 --- /dev/null +++ b/man/vae_ltx23.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vae_ltx23} +\alias{vae_ltx23} +\title{LTX-2.3 Causal Video VAE} +\description{ +Fresh R port of the LTX-2 video autoencoder from the diffusers +reference (Apache-2.0, autoencoder_kl_ltx2.py), with LTX 2.3 defaults: +encoder blocks (256, 512, 1024, 1024), a 4-up-block decoder with mixed +(spatiotemporal, spatiotemporal, temporal, spatial) upsampling, no +upsample residuals, and zeros spatial padding throughout. The encoder +is causal; the decoder is not. +} diff --git a/man/vae_ltx23_modules.Rd b/man/vae_ltx23_modules.Rd new file mode 100644 index 0000000..4c8bbc7 --- /dev/null +++ b/man/vae_ltx23_modules.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vae_ltx23_modules} +\alias{vae_ltx23_modules} +\title{LTX-2.3 Video VAE Building Blocks} +\description{ +Fresh R port of the LTX-2 causal video autoencoder blocks from the +diffusers reference (Apache-2.0, +src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py). Training +and unused inference branches (noise injection, timestep conditioning, +plain-conv downsampling) are intentionally not ported; the 2.3 +checkpoints carry no such weights. +} diff --git a/man/vae_ltx2_modules.Rd b/man/vae_ltx2_modules.Rd deleted file mode 100644 index 0322f37..0000000 --- a/man/vae_ltx2_modules.Rd +++ /dev/null @@ -1,8 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{vae_ltx2_modules} -\alias{vae_ltx2_modules} -\title{LTX2 Video VAE Module Building Blocks} -\description{ -Low-level nn_module components for the LTX2 3D causal VAE. -Used by vae_ltx2.R encoder/decoder. -} diff --git a/man/validate_resolution.Rd b/man/validate_resolution.Rd deleted file mode 100644 index 14a2e77..0000000 --- a/man/validate_resolution.Rd +++ /dev/null @@ -1,28 +0,0 @@ -% tinyrox says don't edit this manually, but it can't stop you! -\name{validate_resolution} -\alias{validate_resolution} -\title{Validate Resolution Against Profile} -\usage{ -validate_resolution(height, width, num_frames, profile) -} -\arguments{ -\item{height}{Integer. Requested height.} - -\item{width}{Integer. Requested width.} - -\item{num_frames}{Integer. Requested number of frames.} - -\item{profile}{Memory profile from `ltx2_memory_profile()`.} -} -\value{ -List with adjusted height, width, num_frames and warning if adjusted. -} -\description{ -Checks if requested resolution fits within memory profile limits. -} -\examples{ -\dontrun{ -profile <- ltx2_memory_profile(vram_gb = 8) -validated <- validate_resolution(720, 1280, 60, profile) -} -} diff --git a/man/vocoder_ltx23.Rd b/man/vocoder_ltx23.Rd new file mode 100644 index 0000000..811a1e5 --- /dev/null +++ b/man/vocoder_ltx23.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vocoder_ltx23} +\alias{vocoder_ltx23} +\title{LTX-2.3 Vocoder with Bandwidth Extension} +\description{ +Fresh R port of the LTX-2 BigVGAN-style vocoder from the diffusers +reference (Apache-2.0, pipelines/ltx2/vocoder.py). The 2.3 vocoder +runs a 16 kHz stage (hidden 1536, snakebeta activations with +anti-aliased up/downsampling), re-analyzes its output into a causal +log-mel spectrogram, and feeds a bandwidth-extension vocoder whose +residual is added to a Hann-resampled skip path for 48 kHz output. +The Kaiser sinc / Hann filters and STFT bases are checkpoint buffers. +Runs in float32 (small model; snakebeta is precision-sensitive). +} diff --git a/man/vram.Rd b/man/vram.Rd new file mode 100644 index 0000000..6165a84 --- /dev/null +++ b/man/vram.Rd @@ -0,0 +1,8 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vram} +\alias{vram} +\title{VRAM Detection and Management Utilities} +\description{ +Device detection, VRAM reporting, and module offloading helpers shared +by the image and video pipelines. +} diff --git a/man/write_wav.Rd b/man/write_wav.Rd new file mode 100644 index 0000000..67ccac8 --- /dev/null +++ b/man/write_wav.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{write_wav} +\alias{write_wav} +\title{Write a 16-bit PCM WAV file} +\usage{ +write_wav(audio, path, sample_rate = 48000L) +} +\arguments{ +\item{audio}{Numeric matrix [channels, samples] in [-1, 1].} + +\item{path}{Output path.} + +\item{sample_rate}{Integer.} +} +\value{ +Invisibly, the path. +} +\description{ +Minimal RIFF writer in base R. +} diff --git a/ref/pipeline_ltx2.R b/ref/pipeline_ltx2.R deleted file mode 100644 index eeb4d30..0000000 --- a/ref/pipeline_ltx2.R +++ /dev/null @@ -1,1147 +0,0 @@ -# Converted from PyTorch by pyrotechnics -# Review: indexing (0->1 based), integer literals (add L), -# and block structure (braces may need adjustment) - -# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# import copy -# import inspect -# from typing import Any, Callable, Dict, List, Optional, Union - -# import numpy as np -# import torch -# from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast - -# from ...callbacks import MultiPipelineCallbacks, PipelineCallback -# from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin -# from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video -# from ...models.transformers import LTX2VideoTransformer3DModel -# from ...schedulers import FlowMatchEulerDiscreteScheduler -# from ...utils import is_torch_xla_available, logging, replace_example_docstring -# from ...utils.torch_utils import randn_tensor -# from ...video_processor import VideoProcessor -# from ..pipeline_utils import DiffusionPipeline -# from .connectors import LTX2TextConnectors -# from .pipeline_output import LTX2PipelineOutput -# from .vocoder import LTX2Vocoder - - -if (is_torch_xla_available()) { -# import torch_xla.core.xla_model as xm - - XLA_AVAILABLE <- TRUE -} else { - XLA_AVAILABLE <- FALSE - -logger <- logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING <- """ - Examples: - ```py - >>> import torch - >>> from diffusers import LTX2Pipeline - >>> from diffusers.pipelines.ltx2.export_utils import encode_video - - >>> pipe <- LTX2Pipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) - >>> pipe.enable_model_cpu_offload() - - >>> prompt <- "A woman with long brown hair && light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket && has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm && natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage" - >>> negative_prompt <- "worst quality, inconsistent motion, blurry, jittery, distorted" - - >>> frame_rate <- 24.0 - >>> video, audio <- pipe( - ... prompt <- prompt, - ... negative_prompt <- negative_prompt, - ... width <- 768, - ... height <- 512, - ... num_frames <- 121, - ... frame_rate <- frame_rate, - ... num_inference_steps <- 40, - ... guidance_scale <- 4.0, - ... output_type <- "np", - ... return_dict <- FALSE, - ... ) - >>> video <- (video * 255).round().astype("uint8") - >>> video <- torch_from_numpy(video) - - >>> encode_video( - ... video[0], - ... fps <- frame_rate, - ... audio <- audio[0].float()$cpu(), - ... audio_sample_rate <- pipe.vocoder.config.output_sampling_rate, # should be 24000 - ... output_path <- "video.mp4", - ... ) - ``` -""" - - -# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift -def calculate_shift( - image_seq_len, - base_seq_len <- 256, - max_seq_len <- 4096, - base_shift <- 0.5, - max_shift <- 1.15, -): - m <- (max_shift - base_shift) / (max_seq_len - base_seq_len) - b <- base_shift - m * base_seq_len - mu <- image_seq_len * m + b - return(mu) - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps <- NULL, - device <- NULL, - timesteps <- NULL, - sigmas <- NULL, - ^kwargs, -): - r""" - Calls the scheduler's `set_timesteps` method && retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` - must be `NULL`. - device (`str` || `torch$device`, *optional*): - The device to which the timesteps should be moved to. If `NULL`, the timesteps are ! moved. - timesteps (`List[int]`, *optional*): - Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, - `num_inference_steps` && `sigmas` must be `NULL`. - sigmas (`List[float]`, *optional*): - Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, - `num_inference_steps` && `timesteps` must be `NULL`. - - Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler && the - second element is the number of inference steps. - """ - if (!is.null(timesteps) && !is.null(sigmas)) { - raise ValueError("Only one of `timesteps` || `sigmas` can be passed. Please choose one to set custom values") - if (!is.null(timesteps)) { - accepts_timesteps <- "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if (! accepts_timesteps) { - raise ValueError( - sprintf("The current scheduler class {scheduler.__class__}'s `set_timesteps` does ! support custom") - sprintf(" timestep schedules. Please check whether you are using the correct scheduler.") - ) - scheduler.set_timesteps(timesteps=timesteps, device=device, ^kwargs) - timesteps <- scheduler.timesteps - num_inference_steps <- length(timesteps) - } else if (!is.null(sigmas)) { - accept_sigmas <- "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if (! accept_sigmas) { - raise ValueError( - sprintf("The current scheduler class {scheduler.__class__}'s `set_timesteps` does ! support custom") - sprintf(" sigmas schedules. Please check whether you are using the correct scheduler.") - ) - scheduler.set_timesteps(sigmas=sigmas, device=device, ^kwargs) - timesteps <- scheduler.timesteps - num_inference_steps <- length(timesteps) - } else { - scheduler.set_timesteps(num_inference_steps, device=device, ^kwargs) - timesteps <- scheduler.timesteps - return(list(timesteps, num_inference_steps)) - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg -def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): - r""" - Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality && fix overexposure. Based on - Section 3.4 from [Common Diffusion Noise Schedules && Sample Steps are - Flawed](https:%/%huggingface.co/papers/2305.08891). - - Args: - noise_cfg (`torch.Tensor`): - The predicted noise tensor for the guided diffusion process. - noise_pred_text (`torch.Tensor`): - The predicted noise tensor for the text-guided diffusion process. - guidance_rescale (`float`, *optional*, defaults to 0.0): - A rescale factor applied to the noise predictions. - - Returns: - noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. - """ - std_text <- noise_pred_text$std(dim=list(range(1, noise_pred_text.ndim)), keepdim=TRUE) - std_cfg <- noise_cfg$std(dim=list(range(1, noise_cfg.ndim)), keepdim=TRUE) - # rescale the results from guidance (fixes overexposure) - noise_pred_rescaled <- noise_cfg * (std_text / std_cfg) - # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images - noise_cfg <- guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg - return(noise_cfg) - - -class LTX2Pipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): - r""" - Pipeline for text-to-video generation. - - Reference: https:%/%github.com/Lightricks/LTX-Video - - Args: - transformer ([`LTXVideoTransformer3DModel`]): - Conditional Transformer architecture to denoise the encoded video latents. - scheduler ([`FlowMatchEulerDiscreteScheduler`]): - A scheduler to be used in combination with `transformer` to denoise the encoded image latents. - vae ([`AutoencoderKLLTXVideo`]): - Variational Auto-Encoder (VAE) Model to encode && decode images to && from latent representations. - text_encoder ([`T5EncoderModel`]): - [T5](https:%/%huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically - the [google/t5-v1_1-xxl](https:%/%huggingface.co/google/t5-v1_1-xxl) variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https:%/%huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer (`T5TokenizerFast`): - Second Tokenizer of class - [T5TokenizerFast](https:%/%huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast). - connectors ([`LTX2TextConnectors`]): - Text connector stack used to adapt text encoder hidden states for the video && audio branches. - """ - - model_cpu_offload_seq <- "text_encoder->connectors->transformer->vae->audio_vae->vocoder" - _optional_components <- [] - _callback_tensor_inputs <- ["latents", "prompt_embeds", "negative_prompt_embeds"] - - def __init__( - self, - scheduler: FlowMatchEulerDiscreteScheduler, - vae: AutoencoderKLLTX2Video, - audio_vae: AutoencoderKLLTX2Audio, - text_encoder: Gemma3ForConditionalGeneration, - tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast], - connectors: LTX2TextConnectors, - transformer: LTX2VideoTransformer3DModel, - vocoder: LTX2Vocoder, - ): - super().__init__() - - self$register_modules( - vae <- vae, - audio_vae <- audio_vae, - text_encoder <- text_encoder, - tokenizer <- tokenizer, - connectors <- connectors, - transformer <- transformer, - vocoder <- vocoder, - scheduler <- scheduler, - ) - - self$vae_spatial_compression_ratio <- ( - self$vae.spatial_compression_ratio if getattr(self, "vae", NULL) !is.null else 32 - ) - self$vae_temporal_compression_ratio <- ( - self$vae.temporal_compression_ratio if getattr(self, "vae", NULL) !is.null else 8 - ) - # TODO: check whether the MEL compression ratio logic here is corrct - self$audio_vae_mel_compression_ratio <- ( - self$audio_vae.mel_compression_ratio if getattr(self, "audio_vae", NULL) !is.null else 4 - ) - self$audio_vae_temporal_compression_ratio <- ( - self$audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", NULL) !is.null else 4 - ) - self$transformer_spatial_patch_size <- ( - self$transformer.config.patch_size if getattr(self, "transformer", NULL) !is.null else 1 - ) - self$transformer_temporal_patch_size <- ( - self$transformer.config.patch_size_t if getattr(self, "transformer") !is.null else 1 - ) - - self$audio_sampling_rate <- ( - self$audio_vae.config.sample_rate if getattr(self, "audio_vae", NULL) !is.null else 16000 - ) - self$audio_hop_length <- ( - self$audio_vae.config.mel_hop_length if getattr(self, "audio_vae", NULL) !is.null else 160 - ) - - self$video_processor <- VideoProcessor(vae_scale_factor=self$vae_spatial_compression_ratio) - self$tokenizer_max_length <- ( - self$tokenizer.model_max_length if getattr(self, "tokenizer", NULL) !is.null else 1024 - ) - - @staticmethod - def _pack_text_embeds( - text_hidden_states: torch.Tensor, - sequence_lengths: torch.Tensor, - device: Union[str, torch$device], - padding_side <- "left", - scale_factor <- 8, - eps <- 1e-6, - ) -> torch.Tensor: - """ - Packs && normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch && - per-layer in a masked fashion (only over non-padded positions). - - Args: - text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): - Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). - sequence_lengths (`torch.Tensor of shape `(batch_size,)`): - The number of valid (non-padded) tokens for each batch instance. - device: (`str` || `torch$device`, *optional*): - torch device to place the resulting embeddings on - padding_side: (`str`, *optional*, defaults to `"left"`): - Whether the text tokenizer performs padding on the `"left"` || `"right"`. - scale_factor (`int`, *optional*, defaults to `8`): - Scaling factor to multiply the normalized hidden states by. - eps (`float`, *optional*, defaults to `1e-6`): - A small positive value for numerical stability when performing normalization. - - Returns: - `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: - Normed && flattened text encoder hidden states. - """ - batch_size, seq_len, hidden_dim, num_layers <- text_hidden_states.shape - original_dtype <- text_hidden_states$dtype - - # Create padding mask - token_indices <- torch_arange(seq_len, device=device)$unsqueeze(0) - if (padding_side == "right") { - # For right padding, valid tokens are from 0 to sequence_length-1 - mask <- token_indices < sequence_lengths[:, NULL] # [batch_size, seq_len] - } else if (padding_side == "left") { - # For left padding, valid tokens are from (T - sequence_length) to T-1 - start_indices <- seq_len - sequence_lengths[:, NULL] # [batch_size, 1] - mask <- token_indices >= start_indices # [B, T] - } else { - raise ValueError(sprintf("padding_side must be 'left' || 'right', got {padding_side}")) - mask <- mask[:, :, NULL, NULL] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] - - # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) - masked_text_hidden_states <- text_hidden_states$masked_fill(~mask, 0.0) - num_valid_positions <- (sequence_lengths * hidden_dim)$view(batch_size, 1, 1, 1) - masked_mean <- masked_text_hidden_states$sum(dim=(1, 2), keepdim=TRUE) / (num_valid_positions + eps) - - # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) - x_min <- text_hidden_states$masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=TRUE) - x_max <- text_hidden_states$masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=TRUE) - - # Normalization - normalized_hidden_states <- (text_hidden_states - masked_mean) / (x_max - x_min + eps) - normalized_hidden_states <- normalized_hidden_states * scale_factor - - # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) - normalized_hidden_states <- normalized_hidden_states$flatten(2) - mask_flat <- mask$squeeze(-1)$expand(-1, -1, hidden_dim * num_layers) - normalized_hidden_states <- normalized_hidden_states$masked_fill(~mask_flat, 0.0) - normalized_hidden_states <- normalized_hidden_states$to(dtype=original_dtype) - return(normalized_hidden_states) - - def _get_gemma_prompt_embeds( - self, - prompt: Union[str, List[str]], - num_videos_per_prompt <- 1, - max_sequence_length <- 1024, - scale_factor <- 8, - device <- NULL, - dtype <- NULL, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` || `List[str]`, *optional*): - prompt to be encoded - device: (`str` || `torch$device`): - torch device to place the resulting embeddings on - dtype: (`torch$dtype`): - torch dtype to cast the prompt embeds to - max_sequence_length (`int`, defaults to 1024): Maximum sequence length to use for the prompt. - """ - device <- device || self$_execution_device - dtype <- dtype || self$text_encoder$dtype - - prompt <- [prompt] if inherits(prompt, "str") else prompt - batch_size <- length(prompt) - - if (getattr(self, "tokenizer", NULL) !is.null) { - # Gemma expects left padding for chat-style prompts - self$tokenizer.padding_side <- "left" - if (self$tokenizer.is.null(pad_token)) { - self$tokenizer.pad_token <- self$tokenizer.eos_token - - prompt <- [p.strip() for p in prompt] - text_inputs <- self$tokenizer( - prompt, - padding <- "max_length", - max_length <- max_sequence_length, - truncation <- TRUE, - add_special_tokens <- TRUE, - return_tensors <- "pt", - ) - text_input_ids <- text_inputs.input_ids - prompt_attention_mask <- text_inputs.attention_mask - text_input_ids <- text_input_ids$to(device) - prompt_attention_mask <- prompt_attention_mask$to(device) - - text_encoder_outputs <- self$text_encoder( - input_ids <- text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=TRUE - ) - text_encoder_hidden_states <- text_encoder_outputs.hidden_states - text_encoder_hidden_states <- torch_stack(text_encoder_hidden_states, dim=-1) - sequence_lengths <- prompt_attention_mask$sum(dim=-1) - - prompt_embeds <- self$_pack_text_embeds( - text_encoder_hidden_states, - sequence_lengths, - device <- device, - padding_side <- self$tokenizer.padding_side, - scale_factor <- scale_factor, - ) - prompt_embeds <- prompt_embeds$to(dtype=dtype) - - # duplicate text embeddings for each generation per prompt, using mps friendly method - _, seq_len, _ <- prompt_embeds.shape - prompt_embeds <- prompt_embeds$repeat(1, num_videos_per_prompt, 1) - prompt_embeds <- prompt_embeds$view(batch_size * num_videos_per_prompt, seq_len, -1) - - prompt_attention_mask <- prompt_attention_mask$view(batch_size, -1) - prompt_attention_mask <- prompt_attention_mask$repeat(num_videos_per_prompt, 1) - - return(list(prompt_embeds, prompt_attention_mask)) - - def encode_prompt( - self, - prompt: Union[str, List[str]], - negative_prompt <- NULL, - do_classifier_free_guidance <- TRUE, - num_videos_per_prompt <- 1, - prompt_embeds <- NULL, - negative_prompt_embeds <- NULL, - prompt_attention_mask <- NULL, - negative_prompt_attention_mask <- NULL, - max_sequence_length <- 1024, - scale_factor <- 8, - device <- NULL, - dtype <- NULL, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` || `List[str]`, *optional*): - prompt to be encoded - negative_prompt (`str` || `List[str]`, *optional*): - The prompt || prompts ! to guide the image generation. If ! defined, one has to pass - `negative_prompt_embeds` instead. Ignored when ! using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - do_classifier_free_guidance (`bool`, *optional*, defaults to `TRUE`): - Whether to use classifier free guidance || !. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - Number of videos that should be generated per prompt. torch device to place the resulting embeddings on - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If ! - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If ! provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - device: (`torch$device`, *optional*): - torch device - dtype: (`torch$dtype`, *optional*): - torch dtype - """ - device <- device || self$_execution_device - - prompt <- [prompt] if inherits(prompt, "str") else prompt - if (!is.null(prompt)) { - batch_size <- length(prompt) - } else { - batch_size <- prompt_embeds.shape[0] - - if (is.null(prompt_embeds)) { - prompt_embeds, prompt_attention_mask <- self$_get_gemma_prompt_embeds( - prompt <- prompt, - num_videos_per_prompt <- num_videos_per_prompt, - max_sequence_length <- max_sequence_length, - scale_factor <- scale_factor, - device <- device, - dtype <- dtype, - ) - - if (do_classifier_free_guidance && is.null(negative_prompt_embeds)) { - negative_prompt <- negative_prompt || "" - negative_prompt <- batch_size * [negative_prompt] if inherits(negative_prompt, "str") else negative_prompt - - if (!is.null(prompt) && type(prompt) is ! type(negative_prompt)) { - raise TypeError( - sprintf("`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=") - sprintf(" {type(prompt)}.") - ) - } else if (batch_size != length(negative_prompt)) { - raise ValueError( - sprintf("`negative_prompt`: {negative_prompt} has batch size {length(negative_prompt)}, but `prompt`:") - sprintf(" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches") - " the batch size of `prompt`." - ) - - negative_prompt_embeds, negative_prompt_attention_mask <- self$_get_gemma_prompt_embeds( - prompt <- negative_prompt, - num_videos_per_prompt <- num_videos_per_prompt, - max_sequence_length <- max_sequence_length, - scale_factor <- scale_factor, - device <- device, - dtype <- dtype, - ) - - return(list(prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask)) - - def check_inputs( - self, - prompt, - height, - width, - callback_on_step_end_tensor_inputs <- NULL, - prompt_embeds <- NULL, - negative_prompt_embeds <- NULL, - prompt_attention_mask <- NULL, - negative_prompt_attention_mask <- NULL, - ): - if (height % 32 != 0 || width % 32 != 0) { - raise ValueError(sprintf("`height` && `width` have to be divisible by 32 but are {height} && {width}.")) - - if !is.null(callback_on_step_end_tensor_inputs) && ! all( - k in self$_callback_tensor_inputs for k in callback_on_step_end_tensor_inputs - ): - raise ValueError( - sprintf("`callback_on_step_end_tensor_inputs` has to be in {self$_callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k ! in self$_callback_tensor_inputs]}") - ) - - if (!is.null(prompt) && !is.null(prompt_embeds)) { - raise ValueError( - sprintf("Cannot forward both `prompt`: {prompt} && `prompt_embeds`: {prompt_embeds}. Please make sure to") - " only forward one of the two." - ) - } else if (is.null(prompt) && is.null(prompt_embeds)) { - raise ValueError( - "Provide either `prompt` || `prompt_embeds`. Cannot leave both `prompt` && `prompt_embeds` undefined." - ) - } else if (!is.null(prompt) && (! inherits(prompt, "str") && ! inherits(prompt, "list"))) { - raise ValueError(sprintf("`prompt` has to be of type `str` || `list` but is {type(prompt)}")) - - if (!is.null(prompt_embeds) && is.null(prompt_attention_mask)) { - raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") - - if (!is.null(negative_prompt_embeds) && is.null(negative_prompt_attention_mask)) { - raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") - - if (!is.null(prompt_embeds) && !is.null(negative_prompt_embeds)) { - if (prompt_embeds.shape != negative_prompt_embeds.shape) { - raise ValueError( - "`prompt_embeds` && `negative_prompt_embeds` must have the same shape when passed directly, but" - sprintf(" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`") - sprintf(" {negative_prompt_embeds.shape}.") - ) - if (prompt_attention_mask.shape != negative_prompt_attention_mask.shape) { - raise ValueError( - "`prompt_attention_mask` && `negative_prompt_attention_mask` must have the same shape when passed directly, but" - sprintf(" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`") - sprintf(" {negative_prompt_attention_mask.shape}.") - ) - - @staticmethod - def _pack_latents(latents: torch.Tensor, patch_size= 1, patch_size_t= 1) -> torch.Tensor: - # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. - # The patch dimensions are then permuted and collapsed into the channel dimension of shape: - # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). - # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features - batch_size, num_channels, num_frames, height, width <- latents.shape - post_patch_num_frames <- num_frames %/% patch_size_t - post_patch_height <- height %/% patch_size - post_patch_width <- width %/% patch_size - latents <- latents$reshape( - batch_size, - -1, - post_patch_num_frames, - patch_size_t, - post_patch_height, - patch_size, - post_patch_width, - patch_size, - ) - latents <- latents$permute(0, 2, 4, 6, 1, 3, 5, 7)$flatten(4, 7)$flatten(1, 3) - return(latents) - - @staticmethod - def _unpack_latents( - latents: torch.Tensor, num_frames, height, width, patch_size <- 1, patch_size_t= 1 - ) -> torch.Tensor: - # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) - # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of - # what happens in the `_pack_latents` method. - batch_size <- latents$size(0) - latents <- latents$reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) - latents <- latents$permute(0, 4, 1, 5, 2, 6, 3, 7)$flatten(6, 7)$flatten(4, 5)$flatten(2, 3) - return(latents) - - @staticmethod - def _denormalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor <- 1.0 - ) -> torch.Tensor: - # Denormalize latents across the channel dimension [B, C, F, H, W] - latents_mean <- latents_mean$view(1, -1, 1, 1, 1)$to(latents$device, latents$dtype) - latents_std <- latents_std$view(1, -1, 1, 1, 1)$to(latents$device, latents$dtype) - latents <- latents * latents_std / scaling_factor + latents_mean - return(latents) - - @staticmethod - def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): - latents_mean <- latents_mean$to(latents$device, latents$dtype) - latents_std <- latents_std$to(latents$device, latents$dtype) - return((latents * latents_std) + latents_mean) - - @staticmethod - def _pack_audio_latents( - latents: torch.Tensor, patch_size <- NULL, patch_size_t= NULL - ) -> torch.Tensor: - # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins - if (!is.null(patch_size) && !is.null(patch_size_t)) { - # Packs the latents into a patch sequence of shape [B, L // p_t * M // p, C * p_t * p] (a ndim=3 tnesor). - # dim=1 is the effective audio sequence length and dim=2 is the effective audio input feature size. - batch_size, num_channels, latent_length, latent_mel_bins <- latents.shape - post_patch_latent_length <- latent_length / patch_size_t - post_patch_mel_bins <- latent_mel_bins / patch_size - latents <- latents$reshape( - batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size - ) - latents <- latents$permute(0, 2, 4, 1, 3, 5)$flatten(3, 5)$flatten(1, 2) - } else { - # Packs the latents into a patch sequence of shape [B, L, C * M]. This implicitly assumes a (mel) - # patch_size of M (all mel bins constitutes a single patch) and a patch_size_t of 1. - latents <- latents$transpose(1, 2)$flatten(2, 3) # [B, C, L, M] --> [B, L, C * M] - return(latents) - - @staticmethod - def _unpack_audio_latents( - latents: torch.Tensor, - latent_length, - num_mel_bins, - patch_size <- NULL, - patch_size_t <- NULL, - ) -> torch.Tensor: - # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], - # where L is the latent audio length and M is the number of mel bins. - if (!is.null(patch_size) && !is.null(patch_size_t)) { - batch_size <- latents$size(0) - latents <- latents$reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) - latents <- latents$permute(0, 3, 1, 4, 2, 5)$flatten(4, 5)$flatten(2, 3) - } else { - # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. - latents <- latents.unflatten(2, (-1, num_mel_bins))$transpose(1, 2) - return(latents) - - def prepare_latents( - self, - batch_size <- 1, - num_channels_latents <- 128, - height <- 512, - width <- 768, - num_frames <- 121, - dtype <- NULL, - device <- NULL, - generator <- NULL, - latents <- NULL, - ) -> torch.Tensor: - if (!is.null(latents)) { - return(latents$to(device=device, dtype=dtype)) - - height <- height %/% self$vae_spatial_compression_ratio - width <- width %/% self$vae_spatial_compression_ratio - num_frames <- (num_frames - 1) %/% self$vae_temporal_compression_ratio + 1 - - shape <- (batch_size, num_channels_latents, num_frames, height, width) - - if (inherits(generator, "list") && length(generator) != batch_size) { - raise ValueError( - sprintf("You have passed a list of generators of length {length(generator)}, but requested an effective batch") - sprintf(" size of {batch_size}. Make sure the batch size matches the length of the generators.") - ) - - latents <- randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents <- self$_pack_latents( - latents, self$transformer_spatial_patch_size, self$transformer_temporal_patch_size - ) - return(latents) - - def prepare_audio_latents( - self, - batch_size <- 1, - num_channels_latents <- 8, - num_mel_bins <- 64, - num_frames <- 121, - frame_rate <- 25.0, - sampling_rate <- 16000, - hop_length <- 160, - dtype <- NULL, - device <- NULL, - generator <- NULL, - latents <- NULL, - ) -> torch.Tensor: - duration_s <- num_frames / frame_rate - latents_per_second <- ( - float(sampling_rate) / float(hop_length) / float(self$audio_vae_temporal_compression_ratio) - ) - latent_length <- round(duration_s * latents_per_second) - - if (!is.null(latents)) { - return(latents$to(device=device, dtype=dtype), latent_length) - - # TODO: confirm whether this logic is correct - latent_mel_bins <- num_mel_bins %/% self$audio_vae_mel_compression_ratio - - shape <- (batch_size, num_channels_latents, latent_length, latent_mel_bins) - - if (inherits(generator, "list") && length(generator) != batch_size) { - raise ValueError( - sprintf("You have passed a list of generators of length {length(generator)}, but requested an effective batch") - sprintf(" size of {batch_size}. Make sure the batch size matches the length of the generators.") - ) - - latents <- randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents <- self$_pack_audio_latents(latents) - return(list(latents, latent_length)) - - @property - def guidance_scale(self): - return(self$_guidance_scale) - - @property - def guidance_rescale(self): - return(self$_guidance_rescale) - - @property - def do_classifier_free_guidance(self): - return(self$_guidance_scale > 1.0) - - @property - def num_timesteps(self): - return(self$_num_timesteps) - - @property - def current_timestep(self): - return(self$_current_timestep) - - @property - def attention_kwargs(self): - return(self$_attention_kwargs) - - @property - def interrupt(self): - return(self$_interrupt) - -# NOTE: wrap body in with_no_grad({ ... }) - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = NULL, - negative_prompt <- NULL, - height <- 512, - width <- 768, - num_frames <- 121, - frame_rate <- 24.0, - num_inference_steps <- 40, - timesteps <- NULL, - guidance_scale <- 4.0, - guidance_rescale <- 0.0, - num_videos_per_prompt <- 1, - generator <- NULL, - latents <- NULL, - audio_latents <- NULL, - prompt_embeds <- NULL, - prompt_attention_mask <- NULL, - negative_prompt_embeds <- NULL, - negative_prompt_attention_mask <- NULL, - decode_timestep: Union[float, List[float]] = 0.0, - decode_noise_scale <- NULL, - output_type <- "pil", - return_dict <- TRUE, - attention_kwargs <- NULL, - callback_on_step_end, NULL]] = NULL, - callback_on_step_end_tensor_inputs <- ["latents"], - max_sequence_length <- 1024, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` || `List[str]`, *optional*): - The prompt || prompts to guide the image generation. If ! defined, one has to pass `prompt_embeds`. - instead. - height (`int`, *optional*, defaults to `512`): - The height in pixels of the generated image. This is set to 480 by default for the best results. - width (`int`, *optional*, defaults to `768`): - The width in pixels of the generated image. This is set to 848 by default for the best results. - num_frames (`int`, *optional*, defaults to `121`): - The number of video frames to generate - frame_rate (`float`, *optional*, defaults to `24.0`): - The frames per second (FPS) of the generated video. - num_inference_steps (`int`, *optional*, defaults to 40): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - timesteps (`List[int]`, *optional*): - Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument - in their `set_timesteps` method. If ! defined, the default behavior when `num_inference_steps` is - passed will be used. Must be in descending order. - guidance_scale (`float`, *optional*, defaults to `4.0`): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https:%/%huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https:%/%huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor proposed by [Common Diffusion Noise Schedules && Sample Steps are - Flawed](https:%/%huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of - [Common Diffusion Noise Schedules && Sample Steps are - Flawed](https:%/%huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when - using zero terminal SNR. - num_videos_per_prompt (`int`, *optional*, defaults to 1): - The number of videos to generate per prompt. - generator (`torch.Generator` || `List[torch.Generator]`, *optional*): - One || a list of [torch generator(s)](https:%/%pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video - generation. Can be used to tweak the same generation with different prompts. If ! provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - audio_latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for audio - generation. Can be used to tweak the same generation with different prompts. If ! provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If ! - provided, text embeddings will be generated from `prompt` input argument. - prompt_attention_mask (`torch.Tensor`, *optional*): - Pre-generated attention mask for text embeddings. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If ! - provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. - negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): - Pre-generated attention mask for negative text embeddings. - decode_timestep (`float`, defaults to `0.0`): - The timestep at which generated video is decoded. - decode_noise_scale (`float`, defaults to `NULL`): - The interpolation factor between random noise && denoised latents at the decode timestep. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https:%/%pillow.readthedocs.io/en/stable/): `PIL.Image.Image` || `np.array`. - return_dict (`bool`, *optional*, defaults to `TRUE`): - Whether || ! to return a [`~pipelines.ltx.LTX2PipelineOutput`] instead of a plain tuple. - attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self$processor` in - [diffusers.models.attention_processor](https:%/%github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step, timestep, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*, defaults to `["latents"]`): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int`, *optional*, defaults to `1024`): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.ltx.LTX2PipelineOutput`] || `tuple`: - If `return_dict` is `TRUE`, [`~pipelines.ltx.LTX2PipelineOutput`] is returned, otherwise a `tuple` is - returned where the first element is a list with the generated images. - """ - - if (isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks))) { - callback_on_step_end_tensor_inputs <- callback_on_step_end.tensor_inputs - - # 1. Check inputs. Raise error if not correct - self$check_inputs( - prompt <- prompt, - height <- height, - width <- width, - callback_on_step_end_tensor_inputs <- callback_on_step_end_tensor_inputs, - prompt_embeds <- prompt_embeds, - negative_prompt_embeds <- negative_prompt_embeds, - prompt_attention_mask <- prompt_attention_mask, - negative_prompt_attention_mask <- negative_prompt_attention_mask, - ) - - self$_guidance_scale <- guidance_scale - self$_guidance_rescale <- guidance_rescale - self$_attention_kwargs <- attention_kwargs - self$_interrupt <- FALSE - self$_current_timestep <- NULL - - # 2. Define call parameters - if (!is.null(prompt) && inherits(prompt, "str")) { - batch_size <- 1 - } else if (!is.null(prompt) && inherits(prompt, "list")) { - batch_size <- length(prompt) - } else { - batch_size <- prompt_embeds.shape[0] - - device <- self$_execution_device - - # 3. Prepare text embeddings - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self$encode_prompt( - prompt <- prompt, - negative_prompt <- negative_prompt, - do_classifier_free_guidance <- self$do_classifier_free_guidance, - num_videos_per_prompt <- num_videos_per_prompt, - prompt_embeds <- prompt_embeds, - negative_prompt_embeds <- negative_prompt_embeds, - prompt_attention_mask <- prompt_attention_mask, - negative_prompt_attention_mask <- negative_prompt_attention_mask, - max_sequence_length <- max_sequence_length, - device <- device, - ) - if (self$do_classifier_free_guidance) { - prompt_embeds <- torch_cat([negative_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask <- torch_cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - - additive_attention_mask <- (1 - prompt_attention_mask$to(prompt_embeds$dtype)) * -1000000.0 - connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask <- self$connectors( - prompt_embeds, additive_attention_mask, additive_mask <- TRUE - ) - - # 4. Prepare latent variables - latent_num_frames <- (num_frames - 1) %/% self$vae_temporal_compression_ratio + 1 - latent_height <- height %/% self$vae_spatial_compression_ratio - latent_width <- width %/% self$vae_spatial_compression_ratio - video_sequence_length <- latent_num_frames * latent_height * latent_width - - num_channels_latents <- self$transformer.config.in_channels - latents <- self$prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - torch.float32, - device, - generator, - latents, - ) - - num_mel_bins <- self$audio_vae.config.mel_bins if getattr(self, "audio_vae", NULL) !is.null else 64 - latent_mel_bins <- num_mel_bins %/% self$audio_vae_mel_compression_ratio - - num_channels_latents_audio <- ( - self$audio_vae.config.latent_channels if getattr(self, "audio_vae", NULL) !is.null else 8 - ) - audio_latents, audio_num_frames <- self$prepare_audio_latents( - batch_size * num_videos_per_prompt, - num_channels_latents <- num_channels_latents_audio, - num_mel_bins <- num_mel_bins, - num_frames <- num_frames, # Video frames, audio frames will be calculated from this - frame_rate <- frame_rate, - sampling_rate <- self$audio_sampling_rate, - hop_length <- self$audio_hop_length, - dtype <- torch.float32, - device <- device, - generator <- generator, - latents <- audio_latents, - ) - - # 5. Prepare timesteps - sigmas <- np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - mu <- calculate_shift( - video_sequence_length, - self$scheduler.config.get("base_image_seq_len", 1024), - self$scheduler.config.get("max_image_seq_len", 4096), - self$scheduler.config.get("base_shift", 0.95), - self$scheduler.config.get("max_shift", 2.05), - ) - # For now, duplicate the scheduler for use with the audio latents - audio_scheduler <- copy.deepcopy(self$scheduler) - _, _ <- retrieve_timesteps( - audio_scheduler, - num_inference_steps, - device, - timesteps, - sigmas <- sigmas, - mu <- mu, - ) - timesteps, num_inference_steps <- retrieve_timesteps( - self$scheduler, - num_inference_steps, - device, - timesteps, - sigmas <- sigmas, - mu <- mu, - ) - num_warmup_steps <- max(length(timesteps) - num_inference_steps * self$scheduler.order, 0) - self$_num_timesteps <- length(timesteps) - - # 6. Prepare micro-conditions - rope_interpolation_scale <- ( - self$vae_temporal_compression_ratio / frame_rate, - self$vae_spatial_compression_ratio, - self$vae_spatial_compression_ratio, - ) - # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop - video_coords <- self$transformer.rope.prepare_video_coords( - latents.shape[0], latent_num_frames, latent_height, latent_width, latents$device, fps <- frame_rate - ) - audio_coords <- self$transformer.audio_rope.prepare_audio_coords( - audio_latents.shape[0], audio_num_frames, audio_latents$device - ) - - # 7. Denoising loop - with self$progress_bar(total=num_inference_steps) as progress_bar: - # TODO: tuple unpacking in for loop - for ( for i, t in enumerate(timesteps): in for i, t in seq_along(timesteps):) { - if (self$interrupt) { - continue - - self$_current_timestep <- t - - latent_model_input <- torch_cat(c(latents) * 2) if self$do_classifier_free_guidance else latents - latent_model_input <- latent_model_input$to(prompt_embeds$dtype) - audio_latent_model_input <- ( - torch_cat(c(audio_latents) * 2) if self$do_classifier_free_guidance else audio_latents - ) - audio_latent_model_input <- audio_latent_model_input$to(prompt_embeds$dtype) - - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep <- t$expand(latent_model_input.shape[0]) - - with self$transformer.cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio <- self$transformer( - hidden_states <- latent_model_input, - audio_hidden_states <- audio_latent_model_input, - encoder_hidden_states <- connector_prompt_embeds, - audio_encoder_hidden_states <- connector_audio_prompt_embeds, - timestep <- timestep, - encoder_attention_mask <- connector_attention_mask, - audio_encoder_attention_mask <- connector_attention_mask, - num_frames <- latent_num_frames, - height <- latent_height, - width <- latent_width, - fps <- frame_rate, - audio_num_frames <- audio_num_frames, - video_coords <- video_coords, - audio_coords <- audio_coords, - # rope_interpolation_scale=rope_interpolation_scale, - attention_kwargs <- attention_kwargs, - return_dict <- FALSE, - ) - noise_pred_video <- noise_pred_video$float() - noise_pred_audio <- noise_pred_audio$float() - - if (self$do_classifier_free_guidance) { - noise_pred_video_uncond, noise_pred_video_text <- noise_pred_video$chunk(2) - noise_pred_video <- noise_pred_video_uncond + self$guidance_scale * ( - noise_pred_video_text - noise_pred_video_uncond - ) - - noise_pred_audio_uncond, noise_pred_audio_text <- noise_pred_audio$chunk(2) - noise_pred_audio <- noise_pred_audio_uncond + self$guidance_scale * ( - noise_pred_audio_text - noise_pred_audio_uncond - ) - - if (self$guidance_rescale > 0) { - # Based on 3.4. in https://huggingface.co/papers/2305.08891 - noise_pred_video <- rescale_noise_cfg( - noise_pred_video, noise_pred_video_text, guidance_rescale <- self$guidance_rescale - ) - noise_pred_audio <- rescale_noise_cfg( - noise_pred_audio, noise_pred_audio_text, guidance_rescale <- self$guidance_rescale - ) - - # compute the previous noisy sample x_t -> x_t-1 - latents <- self$scheduler.step(noise_pred_video, t, latents, return_dict=FALSE)[0] - # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in - # the step method (such as _step_index) - audio_latents <- audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=FALSE)[0] - - if (!is.null(callback_on_step_end)) { - callback_kwargs <- {} - for (k in callback_on_step_end_tensor_inputs) { - callback_kwargs[k] = locals()[k] - callback_outputs <- callback_on_step_end(self, i, t, callback_kwargs) - - latents <- callback_outputs.pop("latents", latents) - prompt_embeds <- callback_outputs.pop("prompt_embeds", prompt_embeds) - - # call the callback, if provided - if (i == length(timesteps) - 1 || ((i + 1) > num_warmup_steps && (i + 1) % self$scheduler.order == 0)) { - progress_bar.update() - - if (XLA_AVAILABLE) { - xm.mark_step() - - latents <- self$_unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self$transformer_spatial_patch_size, - self$transformer_temporal_patch_size, - ) - latents <- self$_denormalize_latents( - latents, self$vae.latents_mean, self$vae.latents_std, self$vae.config.scaling_factor - ) - - audio_latents <- self$_denormalize_audio_latents( - audio_latents, self$audio_vae.latents_mean, self$audio_vae.latents_std - ) - audio_latents <- self$_unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) - - if (output_type == "latent") { - video <- latents - audio <- audio_latents - } else { - latents <- latents$to(prompt_embeds$dtype) - - if (! self$vae.config.timestep_conditioning) { - timestep <- NULL - } else { - noise <- randn_tensor(latents.shape, generator=generator, device=device, dtype=latents$dtype) - if (! inherits(decode_timestep, "list")) { - decode_timestep <- [decode_timestep] * batch_size - if (is.null(decode_noise_scale)) { - decode_noise_scale <- decode_timestep - } else if (! inherits(decode_noise_scale, "list")) { - decode_noise_scale <- [decode_noise_scale] * batch_size - - timestep <- torch_tensor(decode_timestep, device=device, dtype=latents$dtype) - decode_noise_scale <- torch_tensor(decode_noise_scale, device=device, dtype=latents$dtype)[ - :, NULL, NULL, NULL, NULL - ] - latents <- (1 - decode_noise_scale) * latents + decode_noise_scale * noise - - latents <- latents$to(self$vae$dtype) - video <- self$vae.decode(latents, timestep, return_dict=FALSE)[0] - video <- self$video_processor.postprocess_video(video, output_type=output_type) - - audio_latents <- audio_latents$to(self$audio_vae$dtype) - generated_mel_spectrograms <- self$audio_vae.decode(audio_latents, return_dict=FALSE)[0] - audio <- self$vocoder(generated_mel_spectrograms) - - # Offload all models - self$maybe_free_model_hooks() - - if (! return_dict) { - return(list(video, audio)) - - return(LTX2PipelineOutput(frames=video, audio=audio)) - diff --git a/ref/scheduling_flow_match_euler_discrete.R b/ref/scheduling_flow_match_euler_discrete.R deleted file mode 100644 index ce4a57d..0000000 --- a/ref/scheduling_flow_match_euler_discrete.R +++ /dev/null @@ -1,610 +0,0 @@ -# Converted from PyTorch by pyrotechnics -# Review: indexing (0->1 based), integer literals (add L), -# and block structure (braces may need adjustment) - -# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# import math -# from dataclasses import dataclass -# from typing import List, Optional, Tuple, Union - -# import numpy as np -# import torch - -# from ..configuration_utils import ConfigMixin, register_to_config -# from ..utils import BaseOutput, is_scipy_available, logging -# from .scheduling_utils import SchedulerMixin - - -if (is_scipy_available()) { -# import scipy.stats - -logger <- logging.get_logger(__name__) # pylint: disable=invalid-name - - -@dataclass -class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): - """ - Output class for the scheduler's `step` function output. - - Args: - prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): - Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the - denoising loop. - """ - - prev_sample: torch.FloatTensor - - -class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): - """ - Euler scheduler. - - This model inherits from [`SchedulerMixin`] && [`ConfigMixin`]. Check the superclass documentation for the generic - methods the library implements for all schedulers such as loading && saving. - - Args: - num_train_timesteps (`int`, defaults to 1000): - The number of diffusion steps to train the model. - shift (`float`, defaults to 1.0): - The shift value for the timestep schedule. - use_dynamic_shifting (`bool`, defaults to FALSE): - Whether to apply timestep shifting on-the-fly based on the image resolution. - base_shift (`float`, defaults to 0.5): - Value to stabilize image generation. Increasing `base_shift` reduces variation && image is more consistent - with desired output. - max_shift (`float`, defaults to 1.15): - Value change allowed to latent vectors. Increasing `max_shift` encourages more variation && image may be - more exaggerated || stylized. - base_image_seq_len (`int`, defaults to 256): - The base image sequence length. - max_image_seq_len (`int`, defaults to 4096): - The maximum image sequence length. - invert_sigmas (`bool`, defaults to FALSE): - Whether to invert the sigmas. - shift_terminal (`float`, defaults to NULL): - The end value of the shifted timestep schedule. - use_karras_sigmas (`bool`, defaults to FALSE): - Whether to use Karras sigmas for step sizes in the noise schedule during sampling. - use_exponential_sigmas (`bool`, defaults to FALSE): - Whether to use exponential sigmas for step sizes in the noise schedule during sampling. - use_beta_sigmas (`bool`, defaults to FALSE): - Whether to use beta sigmas for step sizes in the noise schedule during sampling. - time_shift_type (`str`, defaults to "exponential"): - The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" || "linear". - stochastic_sampling (`bool`, defaults to FALSE): - Whether to use stochastic sampling. - """ - - _compatibles <- [] - order <- 1 - - @register_to_config - def __init__( - self, - num_train_timesteps <- 1000, - shift <- 1.0, - use_dynamic_shifting <- FALSE, - base_shift <- 0.5, - max_shift <- 1.15, - base_image_seq_len <- 256, - max_image_seq_len <- 4096, - invert_sigmas <- FALSE, - shift_terminal <- NULL, - use_karras_sigmas <- FALSE, - use_exponential_sigmas <- FALSE, - use_beta_sigmas <- FALSE, - time_shift_type <- "exponential", - stochastic_sampling <- FALSE, - ): - if (self$config.use_beta_sigmas && ! is_scipy_available()) { - raise ImportError("Make sure to install scipy if you want to use beta sigmas.") - if (sum([self$config.use_beta_sigmas, self$config.use_exponential_sigmas, self$config.use_karras_sigmas]) > 1) { - raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." - ) - if (time_shift_type ! in {"exponential", "linear"}) { - raise ValueError("`time_shift_type` must either be 'exponential' || 'linear'.") - - timesteps <- np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() - timesteps <- torch_from_numpy(timesteps)$to(dtype=torch.float32) - - sigmas <- timesteps / num_train_timesteps - if (! use_dynamic_shifting) { - # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution - sigmas <- shift * sigmas / (1 + (shift - 1) * sigmas) - - self$timesteps <- sigmas * num_train_timesteps - - self$_step_index <- NULL - self$_begin_index <- NULL - - self$_shift <- shift - - self$sigmas <- sigmas$to("cpu") # to avoid too much CPU/GPU communication - self$sigma_min <- self$sigmas[-1].item() - self$sigma_max <- self$sigmas[0].item() - - @property - def shift(self): - """ - The value used for shifting. - """ - return(self$_shift) - - @property - def step_index(self): - """ - The index counter for current timestep. It will increase 1 after each scheduler step. - """ - return(self$_step_index) - - @property - def begin_index(self): - """ - The index for the first timestep. It should be set from pipeline with `set_begin_index` method. - """ - return(self$_begin_index) - - # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index= 0): - """ - Sets the begin index for the scheduler. This function should be run from pipeline before the inference. - - Args: - begin_index (`int`, defaults to `0`): - The begin index for the scheduler. - """ - self$_begin_index <- begin_index - - def set_shift(self, shift): - self$_shift <- shift - - def scale_noise( - self, - sample: torch.FloatTensor, - timestep: torch.FloatTensor, - noise: torch.FloatTensor, - ) -> torch.FloatTensor: - """ - Forward process in flow-matching - - Args: - sample (`torch.FloatTensor`): - The input sample. - timestep (`torch.FloatTensor`): - The current timestep in the diffusion chain. - noise (`torch.FloatTensor`): - The noise tensor. - - Returns: - `torch.FloatTensor`: - A scaled input sample. - """ - # Make sure sigmas and timesteps have the same device and dtype as original_samples - sigmas <- self$sigmas$to(device=sample$device, dtype=sample$dtype) - - if (sample$device.type == "mps" && torch_is_floating_point(timestep)) { - # mps does not support float64 - schedule_timesteps <- self$timesteps$to(sample$device, dtype=torch.float32) - timestep <- timestep$to(sample$device, dtype=torch.float32) - } else { - schedule_timesteps <- self$timesteps$to(sample$device) - timestep <- timestep$to(sample$device) - - # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index - if (self$is.null(begin_index)) { - step_indices <- [self$index_for_timestep(t, schedule_timesteps) for t in timestep] - } else if (self$!is.null(step_index)) { - # add_noise is called after first denoising step (for inpainting) - step_indices <- [self$step_index] * timestep.shape[0] - } else { - # add noise is called before first denoising step to create initial latent(img2img) - step_indices <- [self$begin_index] * timestep.shape[0] - - sigma <- sigmas[step_indices].flatten() - while length(sigma.shape) < length(sample.shape): - sigma <- sigma$unsqueeze(-1) - - sample <- sigma * noise + (1.0 - sigma) * sample - - return(sample) - - def _sigma_to_t(self, sigma): - return(sigma * self$config.num_train_timesteps) - - def time_shift(self, mu, sigma, t: torch.Tensor): - if (self$config.time_shift_type == "exponential") { - return(self$_time_shift_exponential(mu, sigma, t)) - } else if (self$config.time_shift_type == "linear") { - return(self$_time_shift_linear(mu, sigma, t)) - - def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: - r""" - Stretches && shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config - value. - - Reference: - https:%/%github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 - - Args: - t (`torch.Tensor`): - A tensor of timesteps to be stretched && shifted. - - Returns: - `torch.Tensor`: - A tensor of adjusted timesteps such that the final value equals `self$config.shift_terminal`. - """ - one_minus_z <- 1 - t - scale_factor <- one_minus_z[-1] / (1 - self$config.shift_terminal) - stretched_t <- 1 - (one_minus_z / scale_factor) - return(stretched_t) - - def set_timesteps( - self, - num_inference_steps <- NULL, - device: Union[str, torch$device] = NULL, - sigmas <- NULL, - mu <- NULL, - timesteps <- NULL, - ): - """ - Sets the discrete timesteps used for the diffusion chain (to be run before inference). - - Args: - num_inference_steps (`int`, *optional*): - The number of diffusion steps used when generating samples with a pre-trained model. - device (`str` || `torch$device`, *optional*): - The device to which the timesteps should be moved to. If `NULL`, the timesteps are ! moved. - sigmas (`List[float]`, *optional*): - Custom values for sigmas to be used for each diffusion step. If `NULL`, the sigmas are computed - automatically. - mu (`float`, *optional*): - Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep - shifting. - timesteps (`List[float]`, *optional*): - Custom values for timesteps to be used for each diffusion step. If `NULL`, the timesteps are computed - automatically. - """ - if (self$config.use_dynamic_shifting && is.null(mu)) { - raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `TRUE`") - - if (!is.null(sigmas) && !is.null(timesteps)) { - if (length(sigmas) != length(timesteps)) { - raise ValueError("`sigmas` && `timesteps` should have the same length") - - if (!is.null(num_inference_steps)) { - if (!is.null(sigmas) && length(sigmas) != num_inference_steps) || ( - !is.null(timesteps) && length(timesteps) != num_inference_steps - ): - raise ValueError( - "`sigmas` && `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided" - ) - } else { - num_inference_steps <- length(sigmas) if !is.null(sigmas) else length(timesteps) - - self$num_inference_steps <- num_inference_steps - - # 1. Prepare default sigmas - is_timesteps_provided <- !is.null(timesteps) - - if (is_timesteps_provided) { - timesteps <- np.array(timesteps).astype(np.float32) - - if (is.null(sigmas)) { - if (is.null(timesteps)) { - timesteps <- np.linspace( - self$_sigma_to_t(self$sigma_max), self$_sigma_to_t(self$sigma_min), num_inference_steps - ) - sigmas <- timesteps / self$config.num_train_timesteps - } else { - sigmas <- np.array(sigmas).astype(np.float32) - num_inference_steps <- length(sigmas) - - # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of - # "exponential" or "linear" type is applied - if (self$config.use_dynamic_shifting) { - sigmas <- self$time_shift(mu, 1.0, sigmas) - } else { - sigmas <- self$shift * sigmas / (1 + (self$shift - 1) * sigmas) - - # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value - if (self$config.shift_terminal) { - sigmas <- self$stretch_shift_to_terminal(sigmas) - - # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules - if (self$config.use_karras_sigmas) { - sigmas <- self$_convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - } else if (self$config.use_exponential_sigmas) { - sigmas <- self$_convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - } else if (self$config.use_beta_sigmas) { - sigmas <- self$_convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) - - # 5. Convert sigmas and timesteps to tensors and move to specified device - sigmas <- torch_from_numpy(sigmas)$to(dtype=torch.float32, device=device) - if (! is_timesteps_provided) { - timesteps <- sigmas * self$config.num_train_timesteps - } else { - timesteps <- torch_from_numpy(timesteps)$to(dtype=torch.float32, device=device) - - # 6. Append the terminal sigma value. - # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the - # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi - if (self$config.invert_sigmas) { - sigmas <- 1.0 - sigmas - timesteps <- sigmas * self$config.num_train_timesteps - sigmas <- torch_cat([sigmas, torch_ones(1, device=sigmas$device)]) - } else { - sigmas <- torch_cat([sigmas, torch_zeros(1, device=sigmas$device)]) - - self$timesteps <- timesteps - self$sigmas <- sigmas - self$_step_index <- NULL - self$_begin_index <- NULL - - def index_for_timestep(self, timestep, schedule_timesteps=NULL): - if (is.null(schedule_timesteps)) { - schedule_timesteps <- self$timesteps - - indices <- (schedule_timesteps == timestep).nonzero() - - # The sigma index that is taken for the **very** first `step` - # is always the second index (or the last index if there is only 1) - # This way we can ensure we don't accidentally skip a sigma in - # case we start in the middle of the denoising schedule (e.g. for image-to-image) - pos <- 1 if length(indices) > 1 else 0 - - return(indices[pos].item()) - - def _init_step_index(self, timestep): - if (self$is.null(begin_index)) { - if (isinstance(timestep, torch.Tensor)) { - timestep <- timestep$to(self$timesteps$device) - self$_step_index <- self$index_for_timestep(timestep) - } else { - self$_step_index <- self$_begin_index - - def step( - self, - model_output: torch.FloatTensor, - timestep: Union[float, torch.FloatTensor], - sample: torch.FloatTensor, - s_churn <- 0.0, - s_tmin <- 0.0, - s_tmax <- float("inf"), - s_noise <- 1.0, - generator <- NULL, - per_token_timesteps <- NULL, - return_dict <- TRUE, - ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion - process from the learned model outputs (most often the predicted noise). - - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.FloatTensor`): - A current instance of a sample created by the diffusion process. - s_churn (`float`): - s_tmin (`float`): - s_tmax (`float`): - s_noise (`float`, defaults to 1.0): - Scaling factor for noise added to the sample. - generator (`torch.Generator`, *optional*): - A random number generator. - per_token_timesteps (`torch.Tensor`, *optional*): - The timesteps for each token in the sample. - return_dict (`bool`): - Whether || ! to return a - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] || tuple. - - Returns: - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] || `tuple`: - If return_dict is `TRUE`, - [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, - otherwise a tuple is returned where the first element is the sample tensor. - """ - - if ( - inherits(timestep, "int") - || isinstance(timestep, torch.IntTensor) - || isinstance(timestep, torch.LongTensor) - ): - raise ValueError( - ( - "Passing integer indices (e.g. from `seq_along(timesteps)`) as timesteps to" - " `FlowMatchEulerDiscreteScheduler.step()` is ! supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - - if (self$is.null(step_index)) { - self$_init_step_index(timestep) - - # Upcast to avoid precision issues when computing prev_sample - sample <- sample$to(torch.float32) - - if (!is.null(per_token_timesteps)) { - per_token_sigmas <- per_token_timesteps / self$config.num_train_timesteps - - sigmas <- self$sigmas[:, NULL, NULL] - lower_mask <- sigmas < per_token_sigmas[NULL] - 1e-6 - lower_sigmas <- lower_mask * sigmas - lower_sigmas, _ <- lower_sigmas$max(dim=0) - - current_sigma <- per_token_sigmas[..., NULL] - next_sigma <- lower_sigmas[..., NULL] - dt <- current_sigma - next_sigma - } else { - sigma_idx <- self$step_index - sigma <- self$sigmas[sigma_idx] - sigma_next <- self$sigmas[sigma_idx + 1] - - current_sigma <- sigma - next_sigma <- sigma_next - dt <- sigma_next - sigma - - if (self$config.stochastic_sampling) { - x0 <- sample - current_sigma * model_output - noise <- torch_randn_like(sample) - prev_sample <- (1.0 - next_sigma) * x0 + next_sigma * noise - } else { - prev_sample <- sample + dt * model_output - - # upon completion increase step index by one - self$_step_index += 1 - if (is.null(per_token_timesteps)) { - # Cast sample back to model compatible dtype - prev_sample <- prev_sample$to(model_output$dtype) - - if (! return_dict) { - return(list(prev_sample,)) - - return(FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: - """ - Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative - Models](https:%/%huggingface.co/papers/2206.00364). - - Args: - in_sigmas (`torch.Tensor`): - The input sigma values to be converted. - num_inference_steps (`int`): - The number of inference steps to generate the noise schedule for. - - Returns: - `torch.Tensor`: - The converted sigma values following the Karras noise schedule. - """ - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if (hasattr(self$config, "sigma_min")) { - sigma_min <- self$config.sigma_min - } else { - sigma_min <- NULL - - if (hasattr(self$config, "sigma_max")) { - sigma_max <- self$config.sigma_max - } else { - sigma_max <- NULL - - sigma_min <- sigma_min if !is.null(sigma_min) else in_sigmas[-1].item() - sigma_max <- sigma_max if !is.null(sigma_max) else in_sigmas[0].item() - - rho <- 7.0 # 7.0 is the value used in the paper - ramp <- np.linspace(0, 1, num_inference_steps) - min_inv_rho <- sigma_min ^ (1 / rho) - max_inv_rho <- sigma_max ^ (1 / rho) - sigmas <- (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ^ rho - return(sigmas) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential - def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: - """ - Construct an exponential noise schedule. - - Args: - in_sigmas (`torch.Tensor`): - The input sigma values to be converted. - num_inference_steps (`int`): - The number of inference steps to generate the noise schedule for. - - Returns: - `torch.Tensor`: - The converted sigma values following an exponential schedule. - """ - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if (hasattr(self$config, "sigma_min")) { - sigma_min <- self$config.sigma_min - } else { - sigma_min <- NULL - - if (hasattr(self$config, "sigma_max")) { - sigma_max <- self$config.sigma_max - } else { - sigma_max <- NULL - - sigma_min <- sigma_min if !is.null(sigma_min) else in_sigmas[-1].item() - sigma_max <- sigma_max if !is.null(sigma_max) else in_sigmas[0].item() - - sigmas <- np$exp(np.linspace(math$log(sigma_max), math$log(sigma_min), num_inference_steps)) - return(sigmas) - - # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta - def _convert_to_beta( - self, in_sigmas: torch.Tensor, num_inference_steps, alpha <- 0.6, beta= 0.6 - ) -> torch.Tensor: - """ - Construct a beta noise schedule as proposed in [Beta Sampling is All You - Need](https:%/%huggingface.co/papers/2407.12173). - - Args: - in_sigmas (`torch.Tensor`): - The input sigma values to be converted. - num_inference_steps (`int`): - The number of inference steps to generate the noise schedule for. - alpha (`float`, *optional*, defaults to `0.6`): - The alpha parameter for the beta distribution. - beta (`float`, *optional*, defaults to `0.6`): - The beta parameter for the beta distribution. - - Returns: - `torch.Tensor`: - The converted sigma values following a beta distribution schedule. - """ - - # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers - if (hasattr(self$config, "sigma_min")) { - sigma_min <- self$config.sigma_min - } else { - sigma_min <- NULL - - if (hasattr(self$config, "sigma_max")) { - sigma_max <- self$config.sigma_max - } else { - sigma_max <- NULL - - sigma_min <- sigma_min if !is.null(sigma_min) else in_sigmas[-1].item() - sigma_max <- sigma_max if !is.null(sigma_max) else in_sigmas[0].item() - - sigmas <- np.array( - [ - sigma_min + (ppf * (sigma_max - sigma_min)) - for ppf in [ - scipy.stats.beta.ppf(timestep, alpha, beta) - for timestep in 1 - np.linspace(0, 1, num_inference_steps) - ] - ] - ) - return(sigmas) - - def _time_shift_exponential(self, mu, sigma, t): - return(math$exp(mu) / (math$exp(mu) + (1 / t - 1) ^ sigma)) - - def _time_shift_linear(self, mu, sigma, t): - return(mu / (mu + (1 / t - 1) ^ sigma)) - - def __len__(self): - return(self$config.num_train_timesteps) - diff --git a/tools/gen_fixtures.sh b/tools/gen_fixtures.sh new file mode 100755 index 0000000..e3f2f73 --- /dev/null +++ b/tools/gen_fixtures.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Run LTX-2.3 fixture generators against the diffusers reference checkout. +# Usage: tools/gen_fixtures.sh [script.py ...] (default: all gen_fixtures_*.py) +# +# Uses uv with a CPU-only torch index; never touches the system Python. +set -euo pipefail +cd "$(dirname "$0")/.." + +scripts=("$@") +if [ ${#scripts[@]} -eq 0 ]; then + scripts=(tools/gen_fixtures_*.py) +fi + +for s in "${scripts[@]}"; do + echo "== $s" + uv run --no-project \ + --index https://download.pytorch.org/whl/cpu \ + --index-strategy unsafe-best-match \ + --with torch \ + --with numpy \ + --with safetensors \ + --with huggingface_hub \ + --with packaging \ + --with filelock \ + --with regex \ + --with requests \ + --with tqdm \ + --with pillow \ + python "$s" +done diff --git a/tools/gen_fixtures_audio.py b/tools/gen_fixtures_audio.py new file mode 100644 index 0000000..d741f38 --- /dev/null +++ b/tools/gen_fixtures_audio.py @@ -0,0 +1,142 @@ +# Generate audio VAE + vocoder parity fixtures for the LTX-2.3 R port. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.autoencoders.autoencoder_kl_ltx2_audio import ( # noqa: E402 + LTX2AudioDecoder, +) +from diffusers.pipelines.ltx2.vocoder import ( # noqa: E402 + DownSample1d, + LTX2Vocoder, + LTX2VocoderWithBWE, + SnakeBeta, + UpSample1d, + kaiser_sinc_filter1d, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(31) +fx = {} + + +def add_module(prefix, module): + for name, p in module.state_dict().items(): + fx[f"{prefix}.{name}"] = p + + +# --- Audio VAE decoder (tiny, real structure) ----------------------------------- +adec = LTX2AudioDecoder( + base_channels=8, + output_channels=2, + num_res_blocks=1, + in_channels=2, + latent_channels=4, + ch_mult=(1, 2), + norm_type="pixel", + causality_axis="height", + mid_block_add_attention=False, + mel_bins=8, +).eval() +add_module("adec", adec) +za = torch.randn(1, 4, 5, 4) +fx["adec_x"] = za +with torch.no_grad(): + fx["adec_out"] = adec(za) +print("audio decoder out:", tuple(fx["adec_out"].shape)) # expect (1, 2, 17, 8) + +# --- Kaiser sinc filter + window parity ------------------------------------------- +fx["kaiser_filt_12"] = kaiser_sinc_filter1d(0.25, 0.3, 12) +fx["kaiser_filt_13"] = kaiser_sinc_filter1d(0.1, 0.12, 13) +fx["kaiser_window_12"] = torch.kaiser_window(12, beta=4.7, periodic=False) + +# --- Up/DownSample1d --------------------------------------------------------------- +ds1 = DownSample1d(ratio=2, kernel_size=12).eval() +us1 = UpSample1d(ratio=2, kernel_size=12).eval() +ush = UpSample1d(ratio=4, window_type="hann").eval() +xw = torch.randn(2, 3, 40) +fx["w_x"] = xw +with torch.no_grad(): + fx["w_down"] = ds1(xw) + fx["w_up"] = us1(xw) + fx["w_up_hann"] = ush(xw) + +# --- SnakeBeta ----------------------------------------------------------------------- +sb = SnakeBeta(channels=3) +with torch.no_grad(): + sb.alpha.copy_(torch.randn(3) * 0.1) + sb.beta.copy_(torch.randn(3) * 0.1) +add_module("sb", sb) +with torch.no_grad(): + fx["sb_out"] = sb(xw) + +# --- Tiny vocoder stage ---------------------------------------------------------------- +voc = LTX2Vocoder( + in_channels=8, # 2 channels x 4 mel bins + hidden_channels=16, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3]], + act_fn="snakebeta", + antialias=True, + final_act_fn=None, + final_bias=False, +).eval() +add_module("voc", voc) +mel_in = torch.randn(1, 2, 6, 4) # [B, C, T, M] +fx["voc_x"] = mel_in +with torch.no_grad(): + fx["voc_out"] = voc(mel_in) +print("vocoder out:", tuple(fx["voc_out"].shape)) # expect (1, 2, 24) + +# --- Tiny BWE wrapper -------------------------------------------------------------------- +bwe = LTX2VocoderWithBWE( + in_channels=8, + hidden_channels=16, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3]], + act_fn="snakebeta", + antialias=True, + final_act_fn=None, + final_bias=False, + bwe_in_channels=16, # 2 channels x 8 mel channels + bwe_hidden_channels=8, + bwe_upsample_kernel_sizes=[8, 4], + bwe_upsample_factors=[4, 2], + bwe_resnet_kernel_sizes=[3], + bwe_resnet_dilations=[[1, 3]], + bwe_act_fn="snakebeta", + bwe_antialias=True, + bwe_final_act_fn=None, + bwe_final_bias=False, + filter_length=8, + hop_length=2, + window_length=8, + num_mel_channels=8, + input_sampling_rate=100, + output_sampling_rate=400, +).eval() +# Fill STFT/mel bases with reproducible values (checkpoints carry real ones) +with torch.no_grad(): + bwe.mel_stft.stft_fn.forward_basis.copy_(torch.randn_like(bwe.mel_stft.stft_fn.forward_basis) * 0.1) + bwe.mel_stft.mel_basis.copy_(torch.rand_like(bwe.mel_stft.mel_basis)) +add_module("bwe", bwe) +with torch.no_grad(): + fx["bwe_out"] = bwe(mel_in) +print("bwe out:", tuple(fx["bwe_out"].shape)) # expect (1, 2, 96) + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "audio_ltx23.safetensors")) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/audio_ltx23.safetensors") diff --git a/tools/gen_fixtures_connectors.py b/tools/gen_fixtures_connectors.py new file mode 100644 index 0000000..e33abf6 --- /dev/null +++ b/tools/gen_fixtures_connectors.py @@ -0,0 +1,64 @@ +# Generate connector parity fixtures for the LTX-2.3 R port. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors # noqa: E402 + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(11) +fx = {} + +conn = LTX2TextConnectors( + caption_channels=8, + text_proj_in_factor=3, + video_connector_num_attention_heads=2, + video_connector_attention_head_dim=8, + video_connector_num_layers=2, + video_connector_num_learnable_registers=4, + video_gated_attn=True, + audio_connector_num_attention_heads=2, + audio_connector_attention_head_dim=4, + audio_connector_num_layers=2, + audio_connector_num_learnable_registers=4, + audio_gated_attn=True, + rope_type="split", + per_modality_projections=True, + video_hidden_dim=16, + audio_hidden_dim=8, + proj_bias=True, +).eval() + +for name, p in conn.state_dict().items(): + fx[f"conn.{name}"] = p + +B, S, C, L = 2, 8, 8, 3 +states = torch.randn(B, S, C, L) +mask = torch.ones(B, S) +mask[0, :3] = 0 # left padding on first batch element +mask[1, :1] = 0 + +fx["c_states"] = states +fx["c_mask"] = mask + +with torch.no_grad(): + video_emb, audio_emb, out_mask = conn(states, mask) +fx["c_video_emb"] = video_emb +fx["c_audio_emb"] = audio_emb +fx["c_out_mask"] = out_mask.to(torch.float32) + +# 3D (flattened) input variant should give identical output +with torch.no_grad(): + video_emb3, audio_emb3, _ = conn(states.flatten(2), mask) +fx["c_video_emb3"] = video_emb3 + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "connectors_ltx23.safetensors")) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/connectors_ltx23.safetensors") diff --git a/tools/gen_fixtures_dit.py b/tools/gen_fixtures_dit.py new file mode 100644 index 0000000..7c5798e --- /dev/null +++ b/tools/gen_fixtures_dit.py @@ -0,0 +1,167 @@ +# Generate DiT parity fixtures for the LTX-2.3 R port. +# Tiny LTX-2.3-configured reference transformer: weights + inputs + outputs. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.transformers.transformer_ltx2 import ( # noqa: E402 + LTX2AdaLayerNormSingle, + LTX2Attention, + LTX2VideoTransformer3DModel, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(7) +fx = {} + +# --- Gated attention with split RoPE ------------------------------------------ +attn = LTX2Attention( + query_dim=16, heads=2, kv_heads=2, dim_head=8, + rope_type="split", apply_gated_attention=True, +).eval() +for name, p in attn.state_dict().items(): + fx[f"attn.{name}"] = p + +B, S = 2, 6 +x = torch.randn(B, S, 16) +cos = torch.randn(B, 2, S, 4) +sin = torch.randn(B, 2, S, 4) +mask = torch.zeros(B, 1, S) +mask[:, :, -2:] = -10000.0 +fx["attn_x"] = x +fx["attn_cos"] = cos +fx["attn_sin"] = sin +fx["attn_mask"] = mask +with torch.no_grad(): + fx["attn_out"] = attn(x, attention_mask=mask, query_rotary_emb=(cos, sin)) + fx["attn_out_nomask"] = attn(x, query_rotary_emb=(cos, sin)) + +# --- AdaLN single with 9 mod params -------------------------------------------- +ada = LTX2AdaLayerNormSingle(16, num_mod_params=9).eval() +for name, p in ada.state_dict().items(): + fx[f"ada.{name}"] = p +t = torch.tensor([700.0, 300.0]) +with torch.no_grad(): + ada_out, ada_emb = ada(t, batch_size=2, hidden_dtype=torch.float32) +fx["ada_t"] = t +fx["ada_out"] = ada_out +fx["ada_emb"] = ada_emb + +# --- Tiny full model (LTX-2.3 flags) -------------------------------------------- +model = LTX2VideoTransformer3DModel( + in_channels=4, + out_channels=4, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + num_layers=2, + gated_attn=True, + cross_attn_mod=True, + audio_gated_attn=True, + audio_cross_attn_mod=True, + use_prompt_embeddings=False, + perturbed_attn=True, + rope_type="split", +).eval() +sd = model.state_dict() +for name, p in sd.items(): + fx[f"model.{name}"] = p +print(f"tiny model params: {len(sd)}") + +num_frames, height, width = 2, 3, 4 +n_video = num_frames * height * width +n_audio = 5 +n_text = 7 +hidden = torch.randn(1, n_video, 4) +audio_hidden = torch.randn(1, n_audio, 4) +enc = torch.randn(1, n_text, 16) +audio_enc = torch.randn(1, n_text, 8) +enc_mask = torch.ones(1, n_text) +enc_mask[:, -2:] = 0 +timestep = torch.tensor([700.0]) +sigma = torch.tensor([0.7]) + +fx["m_hidden"] = hidden +fx["m_audio_hidden"] = audio_hidden +fx["m_enc"] = enc +fx["m_audio_enc"] = audio_enc +fx["m_enc_mask"] = enc_mask +fx["m_timestep"] = timestep +fx["m_sigma"] = sigma + +with torch.no_grad(): + out = model( + hidden_states=hidden, + audio_hidden_states=audio_hidden, + encoder_hidden_states=enc, + audio_encoder_hidden_states=audio_enc, + timestep=timestep, + sigma=sigma, + encoder_attention_mask=enc_mask, + audio_encoder_attention_mask=enc_mask, + num_frames=num_frames, + height=height, + width=width, + fps=24.0, + audio_num_frames=n_audio, + use_cross_timestep=True, + return_dict=False, + ) +fx["m_out_video"] = out[0] +fx["m_out_audio"] = out[1] + +# Same forward with modalities isolated (no a2v/v2a cross attention) +with torch.no_grad(): + out_iso = model( + hidden_states=hidden, + audio_hidden_states=audio_hidden, + encoder_hidden_states=enc, + audio_encoder_hidden_states=audio_enc, + timestep=timestep, + sigma=sigma, + num_frames=num_frames, + height=height, + width=width, + audio_num_frames=n_audio, + isolate_modalities=True, + use_cross_timestep=True, + return_dict=False, + ) +fx["m_out_video_iso"] = out_iso[0] +fx["m_out_audio_iso"] = out_iso[1] + +# STG: perturb all batch elements at block 1 +with torch.no_grad(): + out_stg = model( + hidden_states=hidden, + audio_hidden_states=audio_hidden, + encoder_hidden_states=enc, + audio_encoder_hidden_states=audio_enc, + timestep=timestep, + sigma=sigma, + num_frames=num_frames, + height=height, + width=width, + audio_num_frames=n_audio, + spatio_temporal_guidance_blocks=[1], + use_cross_timestep=True, + return_dict=False, + ) +fx["m_out_video_stg"] = out_stg[0] +fx["m_out_audio_stg"] = out_stg[1] + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "dit_ltx23.safetensors")) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/dit_ltx23.safetensors") diff --git a/tools/gen_fixtures_rope.py b/tools/gen_fixtures_rope.py new file mode 100644 index 0000000..c6d164f --- /dev/null +++ b/tools/gen_fixtures_rope.py @@ -0,0 +1,109 @@ +# Generate RoPE parity fixtures for the LTX-2.3 R port. +# +# Runs the diffusers reference implementation (Apache-2.0) on small fixed +# inputs and saves {inputs, expected outputs} as safetensors fixtures that +# the R tinytest suite compares against. Run via tools/gen_fixtures.sh; +# never executed at package test/run time. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.transformers.transformer_ltx2 import ( # noqa: E402 + LTX2AudioVideoRotaryPosEmbed, + apply_interleaved_rotary_emb, + apply_split_rotary_emb, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(42) +fx = {} + +# --- apply_split_rotary_emb: 4D per-head input ------------------------------- +B, H, T, D = 2, 4, 6, 16 +x4 = torch.randn(B, H, T, D) +cos4 = torch.randn(B, H, T, D // 2) +sin4 = torch.randn(B, H, T, D // 2) +fx["split4_x"] = x4 +fx["split4_cos"] = cos4 +fx["split4_sin"] = sin4 +fx["split4_out"] = apply_split_rotary_emb(x4, (cos4, sin4)) + +# --- apply_split_rotary_emb: 3D input reshaped per-head ---------------------- +x3 = torch.randn(B, T, H * D) +fx["split3_x"] = x3 +fx["split3_out"] = apply_split_rotary_emb(x3, (cos4, sin4)) + +# --- apply_interleaved_rotary_emb --------------------------------------------- +S, C = 6, 16 +xi = torch.randn(B, S, C) +cosi = torch.randn(B, S, C) +sini = torch.randn(B, S, C) +fx["inter_x"] = xi +fx["inter_cos"] = cosi +fx["inter_sin"] = sini +fx["inter_out"] = apply_interleaved_rotary_emb(xi, (cosi, sini)) + +# --- embedder: video coords + split freqs ------------------------------------- +rope_v = LTX2AudioVideoRotaryPosEmbed( + dim=64, + base_num_frames=20, + base_height=2048, + base_width=2048, + scale_factors=(8, 32, 32), + theta=10000.0, + modality="video", + double_precision=True, + rope_type="split", + num_attention_heads=4, +) +vc = rope_v.prepare_video_coords(batch_size=2, num_frames=3, height=4, width=6, device="cpu", fps=24.0) +fx["video_coords"] = vc +vcos, vsin = rope_v(vc) +fx["video_cos"] = vcos +fx["video_sin"] = vsin + +# fp16 dtype preservation through apply (values compared in the R test) +xv = torch.randn(2, 4, 3 * 4 * 6, 16, dtype=torch.float16) +fx["video_x_f16"] = xv +fx["video_out_f16"] = apply_split_rotary_emb(xv, (vcos, vsin)) + +# --- embedder: audio coords + split freqs -------------------------------------- +rope_a = LTX2AudioVideoRotaryPosEmbed( + dim=32, + base_num_frames=20, + sampling_rate=16000, + hop_length=160, + scale_factors=(8, 32, 32), + modality="audio", + double_precision=True, + rope_type="split", + num_attention_heads=4, +) +ac = rope_a.prepare_audio_coords(batch_size=2, num_frames=5, device="cpu", shift=0) +fx["audio_coords"] = ac +acos, asin = rope_a(ac) +fx["audio_cos"] = acos +fx["audio_sin"] = asin + +# --- embedder: interleaved variant --------------------------------------------- +rope_i = LTX2AudioVideoRotaryPosEmbed( + dim=64, + modality="video", + double_precision=True, + rope_type="interleaved", + num_attention_heads=4, +) +icos, isin = rope_i(vc) +fx["video_inter_cos"] = icos +fx["video_inter_sin"] = isin + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "rope_ltx23.safetensors")) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/rope_ltx23.safetensors") diff --git a/tools/gen_fixtures_vae.py b/tools/gen_fixtures_vae.py new file mode 100644 index 0000000..439ad2b --- /dev/null +++ b/tools/gen_fixtures_vae.py @@ -0,0 +1,109 @@ +# Generate video VAE parity fixtures for the LTX-2.3 R port. +# Tiny encoder/decoder with the real 2.3 block structure (mixed +# down/upsample types, factors (2,2,1,2), no residual, zeros padding). + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.autoencoders.autoencoder_kl_ltx2 import ( # noqa: E402 + LTX2VideoCausalConv3d, + LTX2VideoDecoder3d, + LTX2VideoDownsampler3d, + LTX2VideoEncoder3d, + LTX2VideoResnetBlock3d, + LTX2VideoUpsampler3d, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(23) +fx = {} + + +def add_module(prefix, module): + for name, p in module.state_dict().items(): + fx[f"{prefix}.{name}"] = p + + +# --- Causal conv: causal and non-causal paths ---------------------------------- +cc = LTX2VideoCausalConv3d(4, 6, kernel_size=3).eval() +add_module("cc", cc) +x = torch.randn(2, 4, 5, 6, 6) +fx["cc_x"] = x +with torch.no_grad(): + fx["cc_out_causal"] = cc(x, causal=True) + fx["cc_out_noncausal"] = cc(x, causal=False) + +# --- Resnet block with channel change (LayerNorm + conv shortcut) -------------- +rb = LTX2VideoResnetBlock3d(in_channels=4, out_channels=8).eval() +add_module("rb", rb) +with torch.no_grad(): + fx["rb_out"] = rb(x, causal=True) + +# --- Downsampler (temporal) and upsampler (spatiotemporal, residual) ------------ +ds = LTX2VideoDownsampler3d(in_channels=8, out_channels=16, stride=(2, 1, 1)).eval() +add_module("ds", ds) +xd = torch.randn(1, 8, 5, 4, 4) +fx["ds_x"] = xd +with torch.no_grad(): + fx["ds_out"] = ds(xd, causal=True) + +us = LTX2VideoUpsampler3d(in_channels=16, stride=(2, 2, 2), residual=True, upscale_factor=2).eval() +add_module("us", us) +xu = torch.randn(1, 16, 3, 4, 4) +fx["us_x"] = xu +with torch.no_grad(): + fx["us_out"] = us(xu, causal=False) + +# --- Tiny encoder (2.3 structure) ------------------------------------------------ +enc = LTX2VideoEncoder3d( + in_channels=3, + out_channels=4, + block_out_channels=(8, 16, 32, 32), + spatio_temporal_scaling=(True, True, True, True), + layers_per_block=(1, 1, 1, 1, 1), + downsample_type=("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + patch_size=4, + patch_size_t=1, + is_causal=True, + spatial_padding_mode="zeros", +).eval() +add_module("enc", enc) +xe = torch.randn(1, 3, 9, 32, 32) +fx["enc_x"] = xe +with torch.no_grad(): + fx["enc_out"] = enc(xe) + +# --- Tiny decoder (2.3 structure: 4 up blocks, mixed types) ---------------------- +dec = LTX2VideoDecoder3d( + in_channels=4, + out_channels=3, + block_out_channels=(16, 32, 32, 64), + spatio_temporal_scaling=(True, True, True, True), + layers_per_block=(1, 1, 1, 1, 1), + upsample_type=("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + patch_size=4, + patch_size_t=1, + is_causal=False, + inject_noise=(False,) * 5, + timestep_conditioning=False, + upsample_residual=(False,) * 4, + upsample_factor=(2, 2, 1, 2), + spatial_padding_mode="zeros", +).eval() +add_module("dec", dec) +xz = torch.randn(1, 4, 3, 2, 2) +fx["dec_x"] = xz +with torch.no_grad(): + fx["dec_out"] = dec(xz) +print("decoder out shape:", tuple(fx["dec_out"].shape)) + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "vae_ltx23.safetensors")) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/vae_ltx23.safetensors") diff --git a/tools/ltx23_compare.R b/tools/ltx23_compare.R new file mode 100644 index 0000000..46371e6 --- /dev/null +++ b/tools/ltx23_compare.R @@ -0,0 +1,156 @@ +# treesitR coverage report: diffusers LTX-2 reference vs the diffuseR port. +# +# Parses the Python reference files and the R port with tree-sitter, +# extracts class/function inventories, and prints a coverage table using +# a curated mapping (which also documents deliberate omissions). +# +# Usage: r tools/ltx23_compare.R + +library(treesitR) + +`%||%` <- function(x, y) if (is.null(x)) y else x + +py_entities <- function(path) { + src <- paste(readLines(path, warn = FALSE), collapse = "\n") + parser <- ts_parser_new() + ts_parser_set_language(parser, ts_language_python()) + tree <- ts_parse(parser, src) + root <- ts_tree_root_node(tree) + + out <- character(0) + # Top-level classes and functions only (methods are per-class detail) + for (i in seq_len(ts_node_named_child_count(root))) { + node <- ts_node_named_child(root, i - 1L) + if (ts_node_type(node) == "decorated_definition") { + node <- ts_node_child_by_field(node, "definition") + } + type <- ts_node_type(node) + if (type %in% c("class_definition", "function_definition")) { + name <- ts_node_text(ts_node_child_by_field(node, "name")) + prefix <- if (type == "class_definition") "class" else "def" + out[[length(out) + 1L]] <- paste(prefix, name) + } + } + unique(out) +} + +r_entities <- function(path) { + src <- paste(readLines(path, warn = FALSE), collapse = "\n") + parser <- ts_parser_new() + ts_parser_set_language(parser, ts_language_r()) + tree <- ts_parse(parser, src) + root <- ts_tree_root_node(tree) + + out <- character(0) + for (i in seq_len(ts_node_named_child_count(root))) { + node <- ts_node_named_child(root, i - 1L) + if (ts_node_type(node) == "binary_operator") { + lhs <- ts_node_child_by_field(node, "lhs") + if (!is.null(lhs) && ts_node_type(lhs) == "identifier") { + out[[length(out) + 1L]] <- ts_node_text(lhs) + } + } + } + unique(out) +} + +ref_dir <- "ref/upstream/diffusers/src/diffusers" +py_files <- c( + transformer = file.path(ref_dir, "models/transformers/transformer_ltx2.py"), + video_vae = file.path(ref_dir, "models/autoencoders/autoencoder_kl_ltx2.py"), + audio_vae = file.path(ref_dir, "models/autoencoders/autoencoder_kl_ltx2_audio.py"), + connectors = file.path(ref_dir, "pipelines/ltx2/connectors.py"), + vocoder = file.path(ref_dir, "pipelines/ltx2/vocoder.py"), + upsampler = file.path(ref_dir, "pipelines/ltx2/latent_upsampler.py") +) + +# Curated mapping: reference entity -> R counterpart or a documented skip +MAPPING <- c( + # transformer_ltx2.py + "def apply_interleaved_rotary_emb" = "ltx23_apply_interleaved_rotary_emb", + "def apply_split_rotary_emb" = "ltx23_apply_split_rotary_emb", + "class LTX2AdaLayerNormSingle" = "ltx23_ada_layer_norm_single", + "class LTX2AudioVideoAttnProcessor" = "ltx23_attention (inlined)", + "class LTX2PerturbedAttnProcessor" = "ltx23_attention (inlined)", + "class LTX2Attention" = "ltx23_attention", + "class LTX2VideoTransformerBlock" = "ltx23_transformer_block", + "class LTX2AudioVideoRotaryPosEmbed" = "ltx23_rotary_pos_embed", + "class LTX2VideoTransformer3DModel" = "ltx23_transformer", + "class AudioVisualModelOutput" = "SKIP: plain list return", + # autoencoder_kl_ltx2.py + "class PerChannelRMSNorm" = "ltx23_per_channel_rms_norm", + "class LTX2VideoCausalConv3d" = "ltx23_causal_conv3d", + "class LTX2VideoResnetBlock3d" = "ltx23_video_resnet_block3d", + "class LTX2VideoDownsampler3d" = "ltx23_video_downsampler3d", + "class LTX2VideoUpsampler3d" = "ltx23_video_upsampler3d", + "class LTX2VideoDownBlock3D" = "ltx23_video_down_block3d", + "class LTX2VideoMidBlock3d" = "ltx23_video_mid_block3d", + "class LTX2VideoUpBlock3d" = "ltx23_video_up_block3d", + "class LTX2VideoEncoder3d" = "ltx23_video_encoder3d", + "class LTX2VideoDecoder3d" = "ltx23_video_decoder3d", + "class AutoencoderKLLTX2Video" = "ltx23_video_vae", + # autoencoder_kl_ltx2_audio.py + "class LTX2AudioCausalConv2d" = "ltx23_audio_causal_conv2d", + "class LTX2AudioPixelNorm" = "ltx23_per_channel_rms_norm (shared)", + "class LTX2AudioAttnBlock" = "SKIP: mid_block_add_attention FALSE in 2.3", + "class LTX2AudioResnetBlock" = "ltx23_audio_resnet_block", + "class LTX2AudioDownsample" = "SKIP: encoder-only (t2v never encodes audio)", + "class LTX2AudioUpsample" = "ltx23_audio_upsample", + "class LTX2AudioAudioPatchifier" = "SKIP: pipeline packs latents directly", + "class LTX2AudioEncoder" = "SKIP: encoder-only (t2v never encodes audio)", + "class LTX2AudioDecoder" = "ltx23_audio_decoder", + "class AutoencoderKLLTX2Audio" = "ltx23_audio_vae", + # connectors.py + "def per_layer_masked_mean_norm" = "SKIP: LTX-2.0 path (per_modality_projections)", + "def per_token_rms_norm" = "ltx23_per_token_rms_norm", + "class LTX2RotaryPosEmbed1d" = "ltx23_rotary_pos_embed_1d", + "class LTX2TransformerBlock1d" = "ltx23_transformer_block_1d", + "class LTX2ConnectorTransformer1d" = "ltx23_connector_transformer_1d", + "class LTX2TextConnectors" = "ltx23_text_connectors", + # vocoder.py + "def kaiser_sinc_filter1d" = "ltx23_kaiser_sinc_filter1d", + "class DownSample1d" = "ltx23_downsample1d", + "class UpSample1d" = "ltx23_upsample1d", + "class AntiAliasAct1d" = "ltx23_antialias_act1d", + "class SnakeBeta" = "ltx23_snake_beta", + "class ResBlock" = "ltx23_vocoder_resblock / ltx23_upsampler_res_block", + "class LTX2Vocoder" = "ltx23_vocoder", + "class CausalSTFT" = "ltx23_causal_stft", + "class MelSTFT" = "ltx23_mel_stft", + "class LTX2VocoderWithBWE" = "ltx23_vocoder_with_bwe", + # latent_upsampler.py + "class PixelShuffleND" = "nnf_pixel_shuffle (2D case only in 2.3)", + "class BlurDownsample" = "SKIP: rational resampler off in 2.3", + "class SpatialRationalResampler" = "SKIP: rational resampler off in 2.3", + "class LTX2LatentUpsamplerModel" = "ltx23_latent_upsampler" +) + +r_files <- list.files("R", pattern = "ltx23|txt2vid_ltx23", full.names = TRUE) +r_defined <- unlist(lapply(r_files, r_entities)) + +missing <- character(0) +cat(sprintf("%-45s %-50s %s\n", "REFERENCE", "R PORT", "STATUS")) +cat(strrep("-", 110), "\n") +for (section in names(py_files)) { + cat("##", section, "\n") + for (entity in py_entities(py_files[[section]])) { + mapped <- if (entity %in% names(MAPPING)) MAPPING[[entity]] else NULL + status <- if (is.null(mapped)) { + missing <- c(missing, paste0(section, ": ", entity)) + "MISSING FROM MAPPING" + } else if (startsWith(mapped, "SKIP")) { + "skipped (documented)" + } else { + r_name <- strsplit(mapped, " ")[[1]][1] + if (r_name %in% r_defined || grepl("nnf_", r_name)) "ported" else "MAPPED BUT NOT DEFINED" + } + cat(sprintf("%-45s %-50s %s\n", entity, mapped %||% "-", status)) + } +} + +if (length(missing)) { + cat("\nUNMAPPED reference entities:\n") + for (m in missing) cat(" -", m, "\n") + quit(status = 1) +} +cat("\nAll reference entities accounted for.\n")