diff --git a/DESCRIPTION b/DESCRIPTION index 0a973bf..ccba33d 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.24 +Version: 0.2.2.25 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), diff --git a/NAMESPACE b/NAMESPACE index 4ccbece..baccd7e 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -1,5 +1,6 @@ # tinyrox says don't edit this manually, but it can't stop you! +export(animatediff_load_pipeline) export(auto_devices) export(bpe_tokenizer) export(clear_vram) @@ -9,12 +10,14 @@ export(convert_sd21_pt_to_diffusers) export(ddim_scheduler_create) export(ddim_scheduler_step) export(decode_bpe) +export(download_animatediff) export(download_component) export(download_flux1) export(download_flux2_klein) export(download_flux3_klein) export(download_ltx2) export(download_model) +export(download_sd15) export(download_sd21) export(download_sdxl) export(download_zimage_turbo) @@ -105,6 +108,7 @@ export(load_gemma3_text_encoder) export(load_gemma4_nf4) export(load_gemma4_text_encoder) export(load_model_component) +export(load_motion_adapter) export(load_pipeline) export(load_qwen3_text_encoder) export(load_t5_text_encoder) @@ -202,6 +206,7 @@ export(ltx23_vocoder_resblock) export(ltx23_vocoder_with_bwe) export(ltx25_open_checkpoint) export(models2devices) +export(motion_adapter) export(offload_to_cpu) export(post_quant_conv) export(preprocess_image) @@ -241,7 +246,9 @@ export(txt2img_flux3) export(txt2img_sd21) export(txt2img_sdxl) export(txt2img_zimage) +export(txt2vid_animatediff) export(txt2vid_ltx2) +export(unet_motion) export(unet_native) export(unet_native_from_safetensors) export(unet_native_from_torchscript) diff --git a/NEWS.md b/NEWS.md index 6ccfa17..4914837 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,41 @@ +# diffuseR 0.2.2.25 + +* AnimateDiff text-to-video on Stable Diffusion 1.5: + `animatediff_load_pipeline()` and `txt2vid_animatediff()`, with the + v3 or v2 motion module (Apache-2.0). `download_animatediff()` + converts the original `.ckpt` in R; `download_sd15()` fetches SD 1.5. + - Prompt keyframes (`prompt = c("1" = ..., "16" = ...)`) blend the + prompt embeddings frame by frame, so the subject changes over the + clip. + - Clips longer than 16 frames use FreeNoise: the motion modules + attend within overlapping 16-frame windows (`context_length`, + `context_stride`) and later windows start from shuffled copies of + earlier noise. The modules were trained on 16-frame clips; in one + window of 20, 24 or 32 frames both diffusers and diffuseR give + noise, so `context_length` above 16 warns. The last partial window + is averaged with per-frame weights; diffusers counts it with one + scalar weight, which at 26 frames turned the last frame into + swirls. + - SD 1.5 support in the native UNet (`num_attention_heads`, 1x1 conv + projections) and `sd_pipeline_from_safetensors(model_name = "sd15")`. + - The SD UNet spatial transformers use fused SDPA and GroupNorm eps + 1e-6, as in diffusers. + - Verified against diffusers on the real weights (fp32 CPU): text + encoder 3.7e-6, image UNet 7.7e-6, motion UNet 6.7e-6 (v3) and + 7.9e-6 (v2) max abs, FreeNoise forward at 24 frames 6.8e-6; + end to end at 256x256x8 88.5 dB PSNR. GPU fp16 at 512x512x32 + against diffusers on the GPU: 43.6 dB. + - Memory: the VAE decodes one frame at a time (`decode_chunk = 1`; + 4 frames set a 13.9 GiB peak at 512x512x16). The loader sizes the + allocator gc gates from the GPU memory free at load and turns on + expandable segments, and the motion UNet runs a minor gc after each + layer (`diffuseR.block_gc`): R's gc had taken 69-78 s of a 95-105 s + render. + - Measured on the RTX 5060 Ti 16 GB, 25 steps, with a second process + holding the card to 10.3 GiB free: 512x512x16 37 s (allocated peak + 9.3 GiB), 512x512x32 112 s (9.6 GiB), 768x768x16 106 s (9.9 GiB); + the same speed as on the free card. Pipeline load about 13 s. + # diffuseR 0.2.2.24 * Image models in a ~10.5 GiB GPU budget (a 16 GB card that keeps a diff --git a/R/animatediff.R b/R/animatediff.R new file mode 100644 index 0000000..1a2eebe --- /dev/null +++ b/R/animatediff.R @@ -0,0 +1,418 @@ +#' AnimateDiff text-to-video on Stable Diffusion 1.5 +#' +#' SD 1.5 (CreativeML OpenRAIL-M) with the AnimateDiff v2 or v3 motion +#' module (Apache-2.0, \code{guoyww/animatediff}). The pipeline follows +#' diffusers' \code{AnimateDiffPipeline}: DDIM with a linear beta schedule +#' and linspace timesteps, classifier-free guidance, frames decoded one +#' chunk at a time. +#' +#' @name animatediff +NULL + +.sd15_repo <- "stable-diffusion-v1-5/stable-diffusion-v1-5" + +# hub path -> name the native loaders read +.sd15_files <- c( + "unet/config.json" = "unet/config.json", + "unet/diffusion_pytorch_model.fp16.safetensors" = + "unet/diffusion_pytorch_model.safetensors", + "vae/config.json" = "vae/config.json", + "vae/diffusion_pytorch_model.fp16.safetensors" = + "vae/diffusion_pytorch_model.safetensors", + "text_encoder/config.json" = "text_encoder/config.json", + "text_encoder/model.fp16.safetensors" = "text_encoder/model.safetensors") + +.animatediff_motion_files <- c(v3 = "v3_sd15_mm.ckpt", v2 = "mm_sd_v15_v2.ckpt") + +.animatediff_dir <- function() { + file.path(tools::R_user_dir("diffuseR", "data"), "animatediff") +} + +#' Download Stable Diffusion 1.5 +#' +#' Fetches the fp16 UNet, VAE and CLIP text encoder (about 2.1 GB) from +#' \code{stable-diffusion-v1-5/stable-diffusion-v1-5} and links them under +#' the file names the native loaders read. +#' +#' @param verbose Logical. +#' +#' @return Invisibly, the diffusers-layout directory. +#' @export +download_sd15 <- function(verbose = TRUE) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to download model weights.") + } + if (!.hub_all_cached(.sd15_repo, names(.sd15_files)) && + !.ltx23_consent("Stable Diffusion 1.5 UNet + VAE + CLIP text encoder (~2.1 GB)")) { + stop("Download cancelled.", call. = FALSE) + } + out <- file.path(.animatediff_dir(), "sd15") + for (f in names(.sd15_files)) { + src <- hfhub::hub_download(.sd15_repo, f) + dest <- file.path(out, .sd15_files[[f]]) + dir.create(dirname(dest), recursive = TRUE, showWarnings = FALSE) + if (!file.exists(dest) && !file.symlink(src, dest)) { + file.copy(src, dest) + } + } + if (verbose) { + message("SD 1.5 ready: ", out) + } + invisible(out) +} + +#' Download an AnimateDiff motion module +#' +#' Fetches the original Apache-2.0 checkpoint from +#' \code{guoyww/animatediff} and converts it once to an fp16 safetensors +#' file (the pickle is read by \code{torch::load_state_dict()}, which only +#' rebuilds tensors). Keys keep their original names; +#' \code{\link{load_motion_adapter}} maps them. +#' +#' @param version "v3" (\code{v3_sd15_mm.ckpt}, no mid-block module) or +#' "v2" (\code{mm_sd_v15_v2.ckpt}). +#' @param verbose Logical. +#' +#' @return Invisibly, the path of the converted safetensors file. +#' @export +download_animatediff <- function(version = c("v3", "v2"), verbose = TRUE) { + version <- match.arg(version) + out <- file.path(.animatediff_dir(), paste0("motion-sd15-", version, + ".safetensors")) + if (file.exists(out)) { + return(invisible(out)) + } + ckpt_name <- .animatediff_motion_files[[version]] + if (!.hub_all_cached("guoyww/animatediff", ckpt_name) && + !.ltx23_consent(sprintf("AnimateDiff %s motion module (~1.7 GB)", version))) { + stop("Download cancelled.", call. = FALSE) + } + ckpt <- hfhub::hub_download("guoyww/animatediff", ckpt_name) + if (verbose) { + message("Converting ", ckpt_name, " to fp16 safetensors...") + } + .animatediff_convert(ckpt, out) + invisible(out) +} + +# Original motion checkpoint -> fp16 safetensors with the same keys. +.animatediff_convert <- function(ckpt, out) { + sd <- torch::load_state_dict(ckpt) + if (!is.null(sd$state_dict)) { + sd <- sd$state_dict + } + sd <- lapply(sd, function(t) { + if (t$is_floating_point()) t$to(dtype = torch::torch_float16()) else t + }) + dir.create(dirname(out), recursive = TRUE, showWarnings = FALSE) + tmp <- paste0(out, ".part") + safetensors::safe_save_file(sd, tmp) + file.rename(tmp, out) + invisible(out) +} + +#' Load the AnimateDiff SD 1.5 pipeline +#' +#' @param motion Motion module: "v3", "v2", or a path to a motion +#' safetensors file. +#' @param model_dir SD 1.5 diffusers directory; NULL uses +#' \code{\link{download_sd15}}. +#' @param device Compute device for the UNet and VAE. +#' @param dtype UNet and motion module dtype; NULL means float16 on CUDA, +#' float32 on CPU. +#' @param verbose Logical. +#' +#' @return A list with \code{unet} (a \code{\link{unet_motion}}), +#' \code{decoder}, \code{text_encoder}, \code{device}, \code{dtype} and +#' \code{train_frames} (the longest motion window that works, 16). +#' @export +animatediff_load_pipeline <- function(motion = "v3", model_dir = NULL, + device = "cuda", dtype = NULL, + verbose = TRUE) { + if (is.null(model_dir)) { + model_dir <- download_sd15(verbose = verbose) + } + path <- if (motion %in% names(.animatediff_motion_files)) { + download_animatediff(motion, verbose = verbose) + } else { + motion + } + if (is.null(dtype)) { + dtype <- if (identical(device, "cuda")) torch::torch_float16() else torch::torch_float32() + } + if (identical(device, "cuda")) { + .animatediff_gc_gates() + } + pipe <- sd_pipeline_from_safetensors(model_dir, model_name = "sd15", + devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), + unet_dtype = torch::torch_float32(), verbose = verbose) + adapter <- load_motion_adapter(path, verbose = verbose) + unet <- unet_motion(pipe$unet, adapter) + unet$to(device = device, dtype = dtype) + pipe$decoder$to(device = device) + list(unet = unet, decoder = pipe$decoder, text_encoder = pipe$text_encoder, + device = device, dtype = dtype, train_frames = .animatediff_train_frames) +} + +# Allocator gc gates sized from the memory free at load, not the card +# total: R's gc took 78 of 105 s at 512x512x16 under torch's default +# gates. A full gc runs only above allocated_rate x total; on a card +# shared down to 10.3 GiB free, the image models' 0.65 x 15.5 GiB was +# out of reach, misses ran only minor gcs, and the garbage promoted by +# the per-layer gcs filled the budget. Called before the first CUDA +# allocation so expandable segments take effect (768x768x16 ran out of +# memory on 5.5 GiB of fragments without). +.animatediff_gc_gates <- function(free_gb = .detect_vram(use_free = TRUE), + total_gb = .detect_vram(use_free = FALSE)) { + if (!isTRUE(free_gb > 0 && total_gb > 0)) { + return(invisible(NULL)) + } + .flux_gc_gates(footprint_gb = min(8, 0.7 * free_gb), + allocated_rate = 0.8 * free_gb / total_gb) +} + +# Clip length the SD 1.5 v2/v3 modules were trained on. Measured: at 20, +# 24 and 32 frames in one window both diffusers and diffuseR give noise. +.animatediff_train_frames <- 16L + +# FreeNoise "shuffle_context" noise (diffusers free_noise_utils.py): past +# the first window, each run of `stride` frames reuses the noise of the +# frames `len` earlier, shuffled, so later windows start from noise the +# motion module has already tied together. +.freenoise_latents <- function(latents, len, stride) { + n <- as.integer(latents$shape[3]) + if (n <= len) { + return(latents) + } + latents <- latents$clone() + for (i in seq.int(len, n - 1L, by = stride)) { + from <- seq.int(i - len, min(n, i - len + stride) - 1L) + to <- seq.int(i, min(n, i + length(from)) - 1L) + # torch's generator, so `seed` fixes the shuffle too + perm <- as.integer(torch::torch_randperm(length(from))) + 1L + src <- from[perm][seq_along(to)] + for (k in seq_along(to)) { + latents[, , to[k] + 1L, , ] <- latents[, , src[k] + 1L, , ] + } + } + latents +} + +# diffusers DDIMScheduler timesteps. "linspace" (AnimateDiff): 0..999 +# evenly, rounded half to even, no offset. "leading" (SD 1.5's own +# config): k * 1000 %/% n + steps_offset. +.ddim_timesteps <- function(n, spacing = c("linspace", "leading"), + steps_offset = 1L, train_steps = 1000L) { + spacing <- match.arg(spacing) + t <- if (spacing == "linspace") { + round(seq(0, train_steps - 1, length.out = n)) + } else { + seq.int(0L, n - 1L) * (train_steps %/% n) + steps_offset + } + as.integer(rev(t)) +} + +.ddim_alphas_cumprod <- function(beta_schedule = c("linear", "scaled_linear"), + beta_start = 0.00085, beta_end = 0.012, + train_steps = 1000L) { + betas <- switch(match.arg(beta_schedule), + linear = seq(beta_start, beta_end, length.out = train_steps), + scaled_linear = seq(sqrt(beta_start), sqrt(beta_end), + length.out = train_steps) ^ 2) + cumprod(1 - betas) +} + +# One deterministic DDIM step (eta 0, epsilon prediction). The previous +# timestep is t - 1000 / n whatever the spacing, as in diffusers; past 0 +# it uses alphas_cumprod[0] (set_alpha_to_one = FALSE). +.ddim_step <- function(sample, eps, t, n, alphas_cumprod) { + a_t <- alphas_cumprod[t + 1L] + prev <- t - length(alphas_cumprod) %/% n + a_prev <- if (prev >= 0) alphas_cumprod[prev + 1L] else alphas_cumprod[1] + x0 <- (sample - eps * sqrt(1 - a_t)) / sqrt(a_t) + x0 * sqrt(a_prev) + eps * sqrt(1 - a_prev) +} + +# Prompt schedule -> one embedding per frame [F, 77, D]. A single prompt +# is one row (the UNet repeats it). Keyframes are named by 1-based frame; +# frames between two keyframes blend their embeddings linearly, frames +# before the first or after the last hold it. +.animatediff_prompt_embeds <- function(prompt, num_frames, encode) { + if (length(prompt) == 1L && is.null(names(prompt))) { + return(encode(prompt)) + } + frames <- as.integer(names(prompt)) + if (anyNA(frames) || any(frames < 1L | frames > num_frames)) { + stop("Prompt keyframes must be named by frame number, 1 to ", + num_frames, call. = FALSE) + } + o <- order(frames) + frames <- frames[o] + embeds <- lapply(unname(as.character(unlist(prompt)))[o], encode) + rows <- lapply(seq_len(num_frames), function(f) { + k <- findInterval(f, frames) + if (k == 0L) { + return(embeds[[1]]) + } + if (k == length(frames)) { + return(embeds[[k]]) + } + w <- (f - frames[k]) / (frames[k + 1L] - frames[k]) + embeds[[k]] * (1 - w) + embeds[[k + 1L]] * w + }) + torch::torch_cat(rows, dim = 1L) +} + +#' Generate a video with AnimateDiff (SD 1.5) +#' +#' Text-to-video from pure noise. The image model draws every frame and the +#' motion module ties them together loosely, so objects drift and morph; +#' \code{prompt} keyframes push that further by changing the subject over +#' the clip. +#' +#' @param prompt A prompt, or keyframes: a character vector or list named +#' by 1-based frame number, e.g. \code{c("1" = "a red fox", "16" = "a +#' wolf made of smoke")}. Frames between keyframes blend the two prompt +#' embeddings. +#' @param negative_prompt Negative prompt (one for the whole clip). +#' @param pipeline From \code{\link{animatediff_load_pipeline}}; NULL loads +#' the default (v3, CUDA). +#' @param width Output width, a multiple of 8 (SD 1.5 is trained at 512). +#' @param height Output height, a multiple of 8. +#' @param num_frames Frames. Past \code{context_length} the motion +#' modules attend within overlapping windows (FreeNoise), so clips can +#' be longer than the 16 frames the modules were trained on. +#' @param context_length Frames per motion window. Longer than the +#' pipeline's \code{train_frames} (16) warns: the v2 and v3 modules +#' gave noise at 20, 24 and 32 frames, though their position table +#' covers 32. +#' @param context_stride Frames between window starts. +#' @param num_inference_steps DDIM steps. +#' @param guidance_scale Classifier-free guidance. Higher values (10-15) +#' push harder towards the prompt and can add artifacts. +#' @param seed Integer seed, or NULL. +#' @param latents Optional starting noise \code{[1, 4, F, H/8, W/8]}. +#' @param decode_chunk Frames per VAE decode. The fp32 decoder needs about +#' 2 GiB per 512x512 frame, so 4 at once set a 13.9 GiB peak where 1 +#' costs 5 GiB and about 2 s. +#' @param fps Frame rate of the saved file. +#' @param filename Output video path; NULL with \code{save_file = TRUE} +#' derives one from the prompt. +#' @param save_file Logical; write the video with \code{\link{save_video}}. +#' @param verbose Logical. +#' +#' @return A list: \code{video} (array \code{[F, H, W, 3]} in 0..1), +#' \code{file} (or NULL) and \code{metadata}. +#' +#' @examples +#' \dontrun{ +#' pipe <- animatediff_load_pipeline("v3") +#' txt2vid_animatediff("a lighthouse in a storm, oil painting", +#' pipeline = pipe, seed = 7, filename = "storm.mp4") +#' # prompt travel: the subject melts from one prompt into the next +#' txt2vid_animatediff(c("1" = "a red fox in snow", "16" = "a wolf made of smoke"), +#' pipeline = pipe, guidance_scale = 12, filename = "fox.mp4") +#' } +#' @export +txt2vid_animatediff <- function(prompt, negative_prompt = NULL, pipeline = NULL, + width = 512L, height = 512L, num_frames = 16L, + context_length = 16L, context_stride = 4L, + num_inference_steps = 25L, guidance_scale = 7.5, + seed = NULL, latents = NULL, decode_chunk = 1L, + fps = 8L, filename = NULL, save_file = TRUE, + verbose = TRUE) { + if (is.null(pipeline)) { + pipeline <- animatediff_load_pipeline(verbose = verbose) + } + if (width %% 8 != 0 || height %% 8 != 0) { + stop("width and height must be multiples of 8", call. = FALSE) + } + if (context_length > pipeline$train_frames) { + warning(sprintf(paste("context_length %d exceeds the %d frames the motion", + "module was trained on; at 20 frames and up the", + "output tends towards noise"), + context_length, pipeline$train_frames), call. = FALSE) + } + device <- pipeline$device + dtype <- pipeline$dtype + if (identical(device, "cuda") && is.null(getOption("diffuseR.block_gc"))) { + op <- options(diffuseR.block_gc = TRUE) + on.exit(options(op), add = TRUE) + } + start <- proc.time()[[3]] + + torch::with_no_grad({ + encode <- function(p) { + pipeline$text_encoder(CLIPTokenizer(p, pad_token = 49407L)) + } + cond <- .animatediff_prompt_embeds(prompt, num_frames, encode) + uncond <- encode(if (is.null(negative_prompt)) "" else negative_prompt) + cond <- cond$to(device = device, dtype = dtype) + uncond <- uncond$to(device = device, dtype = dtype) + + if (is.null(latents)) { + if (!is.null(seed)) { + torch::torch_manual_seed(seed) + } + latents <- torch::torch_randn(c(1L, 4L, num_frames, height %/% 8L, + width %/% 8L)) + latents <- .freenoise_latents(latents, context_length, context_stride) + } + latents <- latents$to(device = device, dtype = torch::torch_float32()) + ac <- .ddim_alphas_cumprod("linear") + steps <- .ddim_timesteps(num_inference_steps, "linspace") + for (i in seq_along(steps)) { + t <- torch::torch_tensor(steps[i], device = device) + x <- latents$to(dtype = dtype) + eps_u <- pipeline$unet(x, t, uncond, context_length, context_stride) + eps_c <- pipeline$unet(x, t, cond, context_length, context_stride) + eps <- (eps_u + (eps_c - eps_u) * guidance_scale)$to(dtype = torch::torch_float32()) + rm(eps_u, eps_c, x) + latents <- .ddim_step(latents, eps, steps[i], num_inference_steps, ac) + if (verbose) { + message(sprintf("step %d/%d", i, length(steps))) + } + } + video <- .animatediff_decode(pipeline$decoder, latents, decode_chunk) + }) + + file <- NULL + if (save_file) { + first <- if (is.list(prompt)) prompt[[1]] else prompt[1] + file <- filename %||% sub("\\.png$", ".mp4", filename_from_prompt(first, datetime = TRUE)) + save_video(video, file, fps = fps, verbose = verbose) + } + elapsed <- proc.time()[[3]] - start + if (verbose) { + message(sprintf("Video generated in %.1f s", elapsed)) + } + list(video = video, file = file, + metadata = list(prompt = prompt, negative_prompt = negative_prompt, + width = width, height = height, num_frames = num_frames, + context_length = context_length, + context_stride = context_stride, + num_inference_steps = num_inference_steps, + guidance_scale = guidance_scale, seed = seed, + fps = fps, seconds = elapsed, model = "animatediff-sd15")) +} + +# Latents [1, 4, F, h, w] -> R array [F, H, W, 3] in 0..1, a few frames +# per VAE call. +.animatediff_decode <- function(decoder, latents, chunk = 4L) { + f <- as.integer(latents$shape[3]) + frames <- (latents / 0.18215)[1, , , , ]$permute(c(2L, 1L, 3L, 4L)) + out <- NULL + for (s in seq.int(1L, f, by = chunk)) { + idx <- s:min(f, s + chunk - 1L) + x <- decoder(frames[idx, , , , drop = FALSE]$contiguous()) + x <- ((x / 2 + 0.5)$clamp(0, 1))$permute(c(1L, 3L, 4L, 2L))$cpu() + a <- as.array(x) + if (is.null(out)) { + out <- array(0, dim = c(f, dim(a)[2:4])) + } + out[idx, , , ] <- a + rm(x) + if (isTRUE(getOption("diffuseR.block_gc"))) gc(full = FALSE) + } + out +} diff --git a/R/motion_module.R b/R/motion_module.R new file mode 100644 index 0000000..bf89ecc --- /dev/null +++ b/R/motion_module.R @@ -0,0 +1,403 @@ +#' AnimateDiff motion modules +#' +#' Temporal transformers that turn the SD 1.5 image UNet into a video +#' model: after every resnet (and spatial attention) the UNet runs one +#' motion module, which attends across frames at each latent position. +#' Ported from diffusers \code{models/unets/unet_motion_model.py} +#' (\code{AnimateDiffTransformer3D}, \code{MotionModules}, +#' \code{MotionAdapter}, \code{UNetMotionModel}) and the double +#' self-attention \code{BasicTransformerBlock} in +#' \code{models/attention.py}, Apache-2.0. +#' +#' @name motion_module +#' @keywords internal +NULL + +# Sinusoidal frame positions [1, max_len, dim]: sin on even channels, cos +# on odd (diffusers SinusoidalPositionalEmbedding). Recomputed rather than +# loaded; the checkpoints' pos_encoder buffers hold the same table. +.motion_positions <- function(dim, max_len) { + pos <- seq.int(0L, max_len - 1L) + div <- exp(seq(0, dim - 1, by = 2) * (-log(10000) / dim)) + ang <- outer(pos, div) + pe <- matrix(0, max_len, dim) + pe[, seq(1L, dim, by = 2L)] <- sin(ang) + pe[, seq(2L, dim, by = 2L)] <- cos(ang) + torch::torch_tensor(pe, dtype = torch::torch_float32())$unsqueeze(1L) +} + +# One temporal transformer layer: two self-attentions over the frame axis +# (positions added to each normed input) and a GEGLU feed-forward. +motion_transformer_block <- torch::nn_module( + "MotionTransformerBlock", + initialize = function(dim, heads, max_len) { + self$norm1 <- torch::nn_layer_norm(dim) + self$attn1 <- UNetCrossAttention(dim, heads = heads, + dim_head = dim %/% heads) + self$norm2 <- torch::nn_layer_norm(dim) + self$attn2 <- UNetCrossAttention(dim, heads = heads, + dim_head = dim %/% heads) + self$norm3 <- torch::nn_layer_norm(dim) + self$ff <- FeedForward(dim) + self$pe <- torch::nn_buffer(.motion_positions(dim, max_len), + persistent = FALSE) + }, + attend = function(x) { + pe <- self$pe[, seq_len(x$shape[2]), ]$to(dtype = x$dtype) + x <- x + self$attn1(self$norm1(x) + pe) + x + self$attn2(self$norm2(x) + pe) + }, + forward = function(x, windows = NULL) { + if (is.null(windows)) { + x <- self$attend(x) + } else { + x <- .motion_windowed(x, windows, self$attend) + } + x + self$ff(self$norm3(x)) + } +) + +# FreeNoise (diffusers FreeNoiseTransformerBlock, Apache-2.0): the two +# attentions run on overlapping frame windows, each with positions from +# 0, and each frame takes the weighted mean of its windows' outputs. The +# feed-forward is per frame, so it runs once on the merged result. +.motion_windowed <- function(x, windows, attend) { + n <- x$shape[2] + acc <- torch::torch_zeros(c(x$shape[1], n, x$shape[3]), device = x$device) + cnt <- torch::torch_zeros(c(1L, n, 1L), device = x$device) + for (w in windows$frames) { + len <- w$end - w$start + 1L + h <- attend(x$narrow(2L, w$start, len))$to(dtype = torch::torch_float32()) + k1 <- w$keep[1] + nk <- length(w$keep) + wt <- torch::torch_tensor(windows$weights[w$keep], device = x$device)$ + view(c(1L, -1L, 1L)) + # narrow() views: indexing with a stored range copies, and add_ on + # the copy would leave acc untouched + acc$narrow(2L, w$start - 1L + k1, nk)$add_(h$narrow(2L, k1, nk) * wt) + cnt$narrow(2L, w$start - 1L + k1, nk)$add_(wt) + } + (acc / cnt)$to(dtype = x$dtype) +} + +# Frame windows for FreeNoise: length `len` every `stride` frames, plus +# one window ending on the last frame when the stride does not land +# there; that window contributes only the frames no earlier window +# covered. diffusers counts those frames with a single scalar weight +# (`weights[:, -n]`) while summing them with per-frame weights; here +# both use the per-frame weights, so a proper weighted mean. +.motion_windows <- function(num_frames, len = 16L, stride = 4L, + weighting = c("pyramid", "flat")) { + weighting <- match.arg(weighting) + if (num_frames < len) { + stop(sprintf("%d frames are fewer than the %d-frame context window", + num_frames, len), call. = FALSE) + } + starts <- seq.int(1L, num_frames - len + 1L, by = stride) + frames <- lapply(starts, function(s) { + list(start = s, end = s + len - 1L, keep = seq_len(len)) + }) + last <- starts[length(starts)] + len - 1L + if (last < num_frames) { + extra <- num_frames - last + frames <- c(frames, list(list(start = num_frames - len + 1L, + end = num_frames, + keep = seq.int(len - extra + 1L, len)))) + } + half <- len %/% 2L + weights <- if (weighting == "flat") { + rep(1, len) + } else if (len %% 2L == 0L) { + c(seq_len(half), rev(seq_len(half))) + } else { + c(seq_len(half), half + 1L, rev(seq_len(half))) + } + list(frames = frames, weights = as.numeric(weights)) +} + +# AnimateDiffTransformer3D: [B*F, C, H, W] in and out. The group norm runs +# over [B, C, F, H, W], so its statistics span all frames. +motion_transformer_3d <- torch::nn_module( + "MotionTransformer3D", + initialize = function(channels, heads = 8L, max_len = 32L, layers = 1L) { + self$norm <- torch::nn_group_norm(32L, channels, eps = 1e-6) + self$proj_in <- torch::nn_linear(channels, channels) + self$transformer_blocks <- torch::nn_module_list() + for (i in seq_len(layers)) { + self$transformer_blocks$append( + motion_transformer_block(channels, heads, max_len)) + } + self$proj_out <- torch::nn_linear(channels, channels) + }, + forward = function(x, num_frames, windows = NULL) { + d <- as.integer(x$shape) + b <- d[1] %/% num_frames + residual <- x + x <- x$reshape(c(b, num_frames, d[2], d[3], d[4]))$permute(c(1L, 3L, 2L, 4L, 5L)) + x <- self$norm(x) + x <- x$permute(c(1L, 4L, 5L, 3L, 2L))$reshape(c(b * d[3] * d[4], num_frames, d[2])) + x <- self$proj_in(x) + for (i in seq_along(self$transformer_blocks)) { + x <- self$transformer_blocks[[i]](x, windows) + } + x <- self$proj_out(x) + x <- x$reshape(c(b, d[3], d[4], num_frames, d[2]))$ + permute(c(1L, 4L, 5L, 2L, 3L))$contiguous()$reshape(d) + x + residual + } +) + +motion_modules <- torch::nn_module( + "MotionModules", + initialize = function(channels, n, heads, max_len) { + self$motion_modules <- torch::nn_module_list() + for (i in seq_len(n)) { + self$motion_modules$append( + motion_transformer_3d(channels, heads, max_len)) + } + } +) + +#' AnimateDiff motion adapter +#' +#' The motion modules for every block of an SD UNet, laid out as +#' diffusers' \code{MotionAdapter} so its parameter names match the +#' converted checkpoints: \code{down_blocks..motion_modules.}, +#' \code{mid_block.motion_modules.0} and \code{up_blocks..motion_modules.}. +#' Up blocks carry one more module than down blocks, as the UNet's up +#' blocks carry one more resnet. +#' +#' @param block_out_channels UNet block widths. +#' @param layers_per_block Motion modules per down block. +#' @param use_mid_block Whether the mid block has a motion module (v2 +#' does, v3 does not). +#' @param num_heads Attention heads per motion module. +#' @param max_len Longest frame window the positions cover. +#' +#' @return An \code{nn_module}. +#' @export +motion_adapter <- torch::nn_module( + "MotionAdapter", + initialize = function(block_out_channels = c(320L, 640L, 1280L, 1280L), + layers_per_block = 2L, use_mid_block = TRUE, + num_heads = 8L, max_len = 32L) { + self$max_len <- max_len + self$down_blocks <- torch::nn_module_list() + for (ch in block_out_channels) { + self$down_blocks$append( + motion_modules(ch, layers_per_block, num_heads, max_len)) + } + if (use_mid_block) { + self$mid_block <- motion_modules( + block_out_channels[length(block_out_channels)], 1L, + num_heads, max_len) + } + self$up_blocks <- torch::nn_module_list() + for (ch in rev(block_out_channels)) { + self$up_blocks$append( + motion_modules(ch, layers_per_block + 1L, num_heads, max_len)) + } + } +) + +# Original AnimateDiff key -> diffusers MotionAdapter key (the mapping in +# diffusers' scripts/convert_animatediff_motion_module_to_diffusers.py). +# Already-converted keys pass through unchanged. +.motion_remap <- function(key) { + key <- sub("^module\\.", "", key) + for (r in list(c(".temporal_transformer", ""), c(".norms.0", ".norm1"), + c(".norms.1", ".norm2"), c(".ff_norm", ".norm3"), + c(".attention_blocks.0", ".attn1"), + c(".attention_blocks.1", ".attn2"))) { + key <- gsub(r[1], r[2], key, fixed = TRUE) + } + key +} + +# Adapter shape from its keys: widths from each down block's first norm, +# the mid block's presence, and the frame window from a stored position +# table (32 when none is stored). +.motion_adapter_config <- function(handle, keys) { + keys <- vapply(keys, .motion_remap, "") + norm <- grep("^down_blocks\\.[0-9]+\\.motion_modules\\.0\\.norm\\.weight$", + keys, value = TRUE) + if (!length(norm)) { + stop("Not an AnimateDiff motion module: no down_blocks.*.motion_modules", + call. = FALSE) + } + idx <- as.integer(sub("^down_blocks\\.([0-9]+)\\..*", "\\1", norm)) + norm <- norm[order(idx)] + raw <- names(keys)[match(norm, keys)] + widths <- vapply(raw, function(k) as.integer(handle$get_tensor(k)$shape[1]), + integer(1), USE.NAMES = FALSE) + per_block <- length(unique(sub("^down_blocks\\.0\\.motion_modules\\.([0-9]+)\\..*", + "\\1", grep("^down_blocks\\.0\\.motion_modules\\.", keys, value = TRUE)))) + pe <- grep("pos_encoder\\.pe$|pos_embed\\.pe$", names(keys), value = TRUE) + max_len <- if (length(pe)) { + as.integer(handle$get_tensor(pe[1])$shape[2]) + } else { + 32L + } + list(block_out_channels = widths, layers_per_block = per_block, + use_mid_block = any(startsWith(keys, "mid_block.")), + max_len = max_len) +} + +#' Load an AnimateDiff motion adapter from safetensors +#' +#' Reads a motion module checkpoint in either the original AnimateDiff key +#' layout (as converted by \code{\link{download_animatediff}}) or diffusers' +#' \code{MotionAdapter} layout, builds a matching +#' \code{\link{motion_adapter}} and fills every parameter. The stored +#' position tables are skipped (the module computes the same sinusoids); +#' any other unmatched key, shape mismatch or unfilled parameter is an +#' error. +#' +#' @param path A \code{.safetensors} file. +#' @param verbose Logical. +#' +#' @return A \code{\link{motion_adapter}} in eval mode (float32, CPU). +#' @export +load_motion_adapter <- function(path, verbose = TRUE) { + handle <- safetensors::safetensors$new(path.expand(path), framework = "torch") + keys <- handle$keys() + cfg <- .motion_adapter_config(handle, keys) + adapter <- do.call(motion_adapter, cfg) + dests <- adapter$named_parameters() + filled <- character(0) + torch::with_no_grad({ + for (key in keys) { + if (grepl("pos_encoder\\.pe$|pos_embed\\.pe$", key)) { + next + } + name <- .motion_remap(key) + dest <- dests[[name]] + if (is.null(dest)) { + stop("Motion adapter load: unmapped key ", key, call. = FALSE) + } + src <- handle$get_tensor(key) + if (!identical(as.integer(src$shape), as.integer(dest$shape))) { + stop("Motion adapter load: shape mismatch at ", name, + call. = FALSE) + } + dest$copy_(src) + filled <- c(filled, name) + } + }) + unfilled <- setdiff(names(dests), filled) + if (length(unfilled)) { + stop("Motion adapter load: ", length(unfilled), " unfilled, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", "), call. = FALSE) + } + if (verbose) { + message(sprintf("Loaded %d motion adapter tensors (%s mid block, %d-frame window)", + length(filled), if (cfg$use_mid_block) "with" else "no", + cfg$max_len)) + } + adapter$eval() + adapter +} + +#' SD UNet with AnimateDiff motion modules +#' +#' Wraps a native SD 1.5 \code{\link{unet_native}} and a +#' \code{\link{motion_adapter}}. The forward takes latents +#' \code{[B, 4, F, H, W]}; frames run through the image layers as a batch +#' of \code{B * F} and through each motion module along the frame axis. +#' The text context is either \code{[B, 77, D]} (repeated per frame) or +#' \code{[B * F, 77, D]}, one row per frame, which is how prompt travel +#' conditions each frame on its own prompt. +#' +#' The forward also takes \code{context_length} and \code{context_stride}: +#' with more frames than \code{context_length}, the motion modules attend +#' within overlapping windows of that length (FreeNoise) instead of over +#' the whole clip. +#' +#' @param unet A \code{\link{unet_native}}. +#' @param adapter A \code{\link{motion_adapter}} with the same block widths. +#' +#' @return An \code{nn_module}. +#' @export +unet_motion <- torch::nn_module( + "UNetMotion", + initialize = function(unet, adapter) { + self$unet <- unet + self$adapter <- adapter + }, + forward = function(sample, timestep, encoder_hidden_states, + context_length = NULL, context_stride = 4L) { + u <- self$unet + a <- self$adapter + d <- as.integer(sample$shape) + f <- d[3] + windows <- NULL + if (!is.null(context_length) && f > context_length) { + windows <- .motion_windows(f, context_length, context_stride) + } else if (f > a$max_len) { + stop(sprintf("%d frames exceed the motion module's %d-frame window", + f, a$max_len), call. = FALSE) + } + # Minor gc per layer (diffuseR.block_gc, as the DiTs): a 16-frame + # forward leaves GBs of dead activations behind otherwise + block_gc <- isTRUE(getOption("diffuseR.block_gc")) + model_dtype <- u$time_embedding_linear_1$weight$dtype + t_emb <- timestep_embedding(timestep, u$block_out_channels[1], + flip_sin_to_cos = TRUE, + downscale_freq_shift = 0L) + t_emb <- u$time_embedding_linear_2(torch::nnf_silu( + u$time_embedding_linear_1(t_emb$to(dtype = model_dtype)))) + ctx <- encoder_hidden_states + if (ctx$shape[1] == d[1] && f > 1L) { + ctx <- ctx$repeat_interleave(f, dim = 1L) + } + + x <- sample$permute(c(1L, 3L, 2L, 4L, 5L))$reshape(c(d[1] * f, d[2], d[4], d[5])) + x <- u$conv_in(x) + res <- list(x) + for (i in seq_along(u$down_blocks)) { + blk <- u$down_blocks[[i]] + mm <- a$down_blocks[[i]]$motion_modules + for (j in seq_along(blk$resnets)) { + x <- blk$resnets[[j]](x, t_emb) + if (blk$has_attentions) { + x <- blk$attentions[[j]](x, ctx) + } + x <- mm[[j]](x, f, windows) + res <- c(res, list(x)) + if (block_gc) gc(full = FALSE) + } + if (blk$has_downsamplers) { + x <- blk$downsamplers[[1]](x) + res <- c(res, list(x)) + } + } + + x <- u$mid_block$resnets[[1]](x, t_emb) + x <- u$mid_block$attentions[[1]](x, ctx) + if (!is.null(a$mid_block)) { + x <- a$mid_block$motion_modules[[1]](x, f, windows) + } + x <- u$mid_block$resnets[[2]](x, t_emb) + + for (i in seq_along(u$up_blocks)) { + blk <- u$up_blocks[[i]] + mm <- a$up_blocks[[i]]$motion_modules + for (j in seq_along(blk$resnets)) { + skip <- res[[length(res)]] + res <- res[-length(res)] + x <- blk$resnets[[j]](torch::torch_cat(list(x, skip), dim = 2L), t_emb) + if (blk$has_attentions) { + x <- blk$attentions[[j]](x, ctx) + } + x <- mm[[j]](x, f, windows) + if (block_gc) gc(full = FALSE) + } + if (blk$has_upsamplers) { + x <- blk$upsamplers[[1]](x) + } + } + + x <- u$conv_out(torch::nnf_silu(u$conv_norm_out(x))) + x$reshape(c(d[1], f, d[2], d[4], d[5]))$permute(c(1L, 3L, 2L, 4L, 5L)) + } +) diff --git a/R/sd_pipeline_safetensors.R b/R/sd_pipeline_safetensors.R index 4e767e1..b81d4bc 100644 --- a/R/sd_pipeline_safetensors.R +++ b/R/sd_pipeline_safetensors.R @@ -107,7 +107,7 @@ download_sd21 <- function(verbose = TRUE) { #' #' @param diffusers_dir Directory with \code{unet/}, \code{vae/}, #' \code{text_encoder/} subdirectories. -#' @param model_name Currently "sd21" (SDXL pending its second encoder). +#' @param model_name "sd21" or "sd15" (see \code{\link{download_sd15}}). #' @param devices Named list of component devices (\code{unet}, #' \code{decoder}, \code{text_encoder}); defaults to all-CPU. #' @param unet_dtype A torch dtype for the UNet (default float16 on CUDA, @@ -120,11 +120,11 @@ download_sd21 <- function(verbose = TRUE) { sd_pipeline_from_safetensors <- function(diffusers_dir, model_name = "sd21", devices = NULL, unet_dtype = NULL, verbose = TRUE) { - if (!identical(model_name, "sd21")) { - stop("sd_pipeline_from_safetensors currently supports \"sd21\"; ", - "SDXL needs its second text encoder and added-conditioning ", - "embeddings (see tasks/todo.md).") + if (!model_name %in% c("sd21", "sd15")) { + stop("sd_pipeline_from_safetensors supports \"sd21\" and \"sd15\"; ", + "SDXL has sdxl_pipeline_from_safetensors().") } + sd15 <- identical(model_name, "sd15") diffusers_dir <- path.expand(diffusers_dir) if (is.null(devices)) { devices <- list(unet = "cpu", decoder = "cpu", text_encoder = "cpu") @@ -145,17 +145,24 @@ sd_pipeline_from_safetensors <- function(diffusers_dir, model_name = "sd21", # casts prompt embeds to unet_dtype and the latent to float32 before # decode); only the UNet runs at unet_dtype. Up-casting a float16 # artifact to float32 for these is lossless. + # SD 1.5's CLIP ViT-L uses quick_gelu and the final-LN last hidden + # state too. text_encoder <- text_encoder_native_from_safetensors( file.path(diffusers_dir, "text_encoder"), apply_final_ln = TRUE, - verbose = FALSE) + gelu_type = if (sd15) "quick" else "tanh", verbose = FALSE) text_encoder$to(device = torch::torch_device(devices$text_encoder), dtype = torch::torch_float32()) if (verbose) { message("Building UNet...") } - unet <- unet_native_from_safetensors(file.path(diffusers_dir, "unet"), - verbose = FALSE) + unet_args <- list(file.path(diffusers_dir, "unet"), verbose = FALSE) + if (sd15) { + # 768-wide CLIP context, 8 heads in every block + unet_args <- c(unet_args, list(cross_attention_dim = 768L, + num_attention_heads = 8L)) + } + unet <- do.call(unet_native_from_safetensors, unet_args) unet$to(device = torch::torch_device(devices$unet), dtype = unet_dtype) if (verbose) { diff --git a/R/unet.R b/R/unet.R index 5ed911c..330c8b6 100644 --- a/R/unet.R +++ b/R/unet.R @@ -1,3 +1,12 @@ +# Heads and head size for a SpatialTransformer at `channels`: SD 2.1 +# fixes the head size, SD 1.5 fixes the head count. +.unet_heads <- function(channels, attention_head_dim, num_heads = NULL) { + if (is.null(num_heads)) { + return(c(channels %/% attention_head_dim, attention_head_dim)) + } + c(num_heads, channels %/% num_heads) +} + #' Native UNet for Stable Diffusion #' #' Native R torch implementation of UNet2DConditionModel. @@ -9,6 +18,9 @@ #' @param layers_per_block Number of ResBlocks per down/up block #' @param cross_attention_dim Context dimension from text encoder #' @param attention_head_dim Dimension per attention head +#' @param num_attention_heads NULL (SD 2.1: heads = channels / +#' \code{attention_head_dim}) or a fixed head count for every block +#' (SD 1.5: 8, so the head size grows with the channels). #' #' @return An nn_module representing the UNet #' @export @@ -21,8 +33,10 @@ unet_native <- torch::nn_module( block_out_channels = c(320L, 640L, 1280L, 1280L), layers_per_block = 2L, cross_attention_dim = 1024L, - attention_head_dim = 64L + attention_head_dim = 64L, + num_attention_heads = NULL ) { + self$num_attention_heads <- num_attention_heads self$in_channels <- in_channels self$out_channels <- out_channels self$block_out_channels <- block_out_channels @@ -133,6 +147,7 @@ unet_native <- torch::nn_module( add_downsample = TRUE, add_attention = TRUE ) { + num_heads <- self$num_attention_heads # Create block as a simple nn_module DownBlock <- torch::nn_module( "DownBlock", @@ -152,10 +167,10 @@ unet_native <- torch::nn_module( self$has_attentions <- add_attention if (add_attention) { self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim + hd <- .unet_heads(out_channels, attention_head_dim, num_heads) for (i in seq_len(num_layers)) { self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, + SpatialTransformer(out_channels, hd[1], hd[2], depth = 1L, context_dim = cross_attention_dim) ) } @@ -178,6 +193,7 @@ unet_native <- torch::nn_module( cross_attention_dim, attention_head_dim ) { + num_heads <- self$num_attention_heads MidBlock <- torch::nn_module( "MidBlock", initialize = function() { @@ -187,10 +203,10 @@ unet_native <- torch::nn_module( self$resnets$append(UNetResBlock(channels, channels, time_embed_dim)) # Attention - n_heads <- channels %/% attention_head_dim + hd <- .unet_heads(channels, attention_head_dim, num_heads) self$attentions <- torch::nn_module_list() self$attentions$append( - SpatialTransformer(channels, n_heads, attention_head_dim, + SpatialTransformer(channels, hd[1], hd[2], depth = 1L, context_dim = cross_attention_dim) ) } @@ -210,6 +226,7 @@ unet_native <- torch::nn_module( ) { # resnet_skip_channels is a vector with one entry per resnet num_resnets <- length(resnet_skip_channels) + num_heads <- self$num_attention_heads UpBlock <- torch::nn_module( "UpBlock", @@ -230,10 +247,10 @@ unet_native <- torch::nn_module( self$has_attentions <- add_attention if (add_attention) { self$attentions <- torch::nn_module_list() - n_heads <- out_channels %/% attention_head_dim + hd <- .unet_heads(out_channels, attention_head_dim, num_heads) for (i in seq_len(num_resnets)) { self$attentions$append( - SpatialTransformer(out_channels, n_heads, attention_head_dim, + SpatialTransformer(out_channels, hd[1], hd[2], depth = 1L, context_dim = cross_attention_dim) ) } diff --git a/R/unet_modules.R b/R/unet_modules.R index 7f4819a..6f35759 100644 --- a/R/unet_modules.R +++ b/R/unet_modules.R @@ -201,12 +201,9 @@ UNetCrossAttention <- torch::nn_module( k <- k$view(c(b, context_len, h, self$dim_head))$transpose(2L, 3L) v <- v$view(c(b, context_len, h, self$dim_head))$transpose(2L, 3L) - # Scaled dot-product attention - attn <- torch::torch_matmul(q, k$transpose(3L, 4L)) * self$scale - attn <- torch::nnf_softmax(attn, dim = -1L) - - # Apply attention to values - out <- torch::torch_matmul(attn, v) + # Scaled dot-product attention (fused when available: the manual form + # holds an [B, H, S, S] score matrix, 4.3 GB for 16 frames at 512 px) + out <- .ltx23_sdpa(q, k, v) # Reshape back: (B, H, S, D) -> (B, S, H*D) out <- out$transpose(2L, 3L)$contiguous()$view(c(b, seq_len, -1L)) @@ -314,7 +311,8 @@ SpatialTransformer <- torch::nn_module( self$in_channels <- in_channels # Normalization - self$norm <- group_norm_32(in_channels) + # eps 1e-6 as in diffusers Transformer2DModel (the resnets use 1e-5) + self$norm <- torch::nn_group_norm(32, in_channels, eps = 1e-6) # Projection in/out self$proj_in <- torch::nn_linear(in_channels, inner_dim) diff --git a/R/unet_safetensors.R b/R/unet_safetensors.R index 7fa18ed..c5f9740 100644 --- a/R/unet_safetensors.R +++ b/R/unet_safetensors.R @@ -64,7 +64,13 @@ NULL next } src <- opened$handle$get_tensor(key) - if (!all(dest$shape == src$shape)) { + # SD 1.5's proj_in/proj_out are 1x1 convs; the native + # SpatialTransformer uses the equivalent linear + if (src$dim() == 4L && dest$dim() == 2L && + all(as.integer(src$shape[3:4]) == 1L)) { + src <- src$flatten(start_dim = 2L) + } + if (src$dim() != dest$dim() || !all(dest$shape == src$shape)) { mismatch <- c(mismatch, sprintf("%s (%s vs %s)", native_name, paste(as.integer(src$shape), collapse = "x"), paste(as.integer(dest$shape), collapse = "x"))) diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md index 7133fd7..e872229 100644 --- a/inst/REFERENCES.md +++ b/inst/REFERENCES.md @@ -70,6 +70,17 @@ module; this file documents the actual lineage, idea by idea. | Qwen3-4B encoder + Qwen2 byte-level BPE tokenizer | Shared with the FLUX.2 klein port (see above); the tokenizer files are byte-identical between the two checkpoints | | Weights | Tongyi-MAI/Z-Image-Turbo (Apache-2.0, ungated; downloaded by the user, never redistributed) | +## AnimateDiff (SD 1.5) + +| What | Source | +|---|---| +| Motion modules (temporal transformer, sinusoidal frame positions), motion adapter layout, motion UNet forward | Ported from HuggingFace **diffusers** (Apache-2.0): `models/unets/unet_motion_model.py`, `models/transformers/transformer_temporal.py`, `models/attention.py` | +| Original checkpoint key mapping | diffusers `scripts/convert_animatediff_motion_module_to_diffusers.py` (Apache-2.0); conversion runs in R via `torch::load_state_dict()` | +| Pipeline: DDIM (linear betas, linspace timesteps), CFG, per-frame prompt embeddings | diffusers `pipelines/animatediff/pipeline_animatediff.py` and `schedulers/scheduling_ddim.py` (Apache-2.0) | +| FreeNoise: overlapping motion windows with pyramid weights, shuffle-context noise | diffusers `FreeNoiseTransformerBlock` (`models/attention.py`) and `pipelines/free_noise_utils.py` (Apache-2.0); method from Qiu et al. (2023), arXiv:2310.15169. The last partial window is counted with per-frame weights, where diffusers uses one scalar | +| SD 1.5 UNet, CLIP ViT-L text encoder, VAE | The native SD 2.1/SDXL ports with SD 1.5 config (8 fixed heads, 768 context, 1x1 conv projections loaded as linears, quick GELU) | +| Weights | stable-diffusion-v1-5/stable-diffusion-v1-5 (CreativeML OpenRAIL-M) and guoyww/animatediff `v3_sd15_mm.ckpt` / `mm_sd_v15_v2.ckpt` (Apache-2.0); downloaded by the user, never redistributed | + ## Quantization | What | Source | diff --git a/inst/tinytest/test_animatediff.R b/inst/tinytest/test_animatediff.R new file mode 100644 index 0000000..aa40321 --- /dev/null +++ b/inst/tinytest/test_animatediff.R @@ -0,0 +1,156 @@ +# AnimateDiff motion modules, SD 1.5 head layout, DDIM timesteps and prompt +# keyframes. Tiny random modules; no weights needed. + +library(tinytest) +library(diffuseR) + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +# SD 2.1 fixes the head size, SD 1.5 the head count +expect_equal(diffuseR:::.unet_heads(640L, 64L), c(10L, 64L)) +expect_equal(diffuseR:::.unet_heads(640L, 64L, 8L), c(8L, 80L)) + +# Original AnimateDiff keys map onto the diffusers adapter layout +remap <- diffuseR:::.motion_remap +expect_equal(remap("down_blocks.1.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.1.to_q.weight"), + "down_blocks.1.motion_modules.0.transformer_blocks.0.attn2.to_q.weight") +expect_equal(remap("up_blocks.0.motion_modules.2.temporal_transformer.transformer_blocks.0.ff_norm.bias"), + "up_blocks.0.motion_modules.2.transformer_blocks.0.norm3.bias") +expect_equal(remap("mid_block.motion_modules.0.temporal_transformer.transformer_blocks.0.norms.0.weight"), + "mid_block.motion_modules.0.transformer_blocks.0.norm1.weight") +expect_equal(remap("down_blocks.0.motion_modules.0.proj_in.weight"), + "down_blocks.0.motion_modules.0.proj_in.weight") + +# Sinusoidal frame positions: sin on even channels, cos on odd +pe <- as.array(diffuseR:::.motion_positions(8L, 4L)) +expect_equal(dim(pe), c(1L, 4L, 8L)) +expect_equal(pe[1, 3, 1], sin(2), tolerance = 1e-6) +expect_equal(pe[1, 3, 2], cos(2), tolerance = 1e-6) +expect_equal(pe[1, 2, 3], sin(1 * exp(2 * -log(10000) / 8)), tolerance = 1e-6) + +# Adapter layout: up blocks carry one more module than down blocks +ad <- motion_adapter(c(32L, 64L), layers_per_block = 1L, use_mid_block = TRUE, + num_heads = 2L, max_len = 8L) +nm <- names(ad$named_parameters()) +expect_true("down_blocks.0.motion_modules.0.norm.weight" %in% nm) +expect_true("down_blocks.1.motion_modules.0.transformer_blocks.0.attn1.to_out.0.weight" %in% nm) +expect_true("up_blocks.1.motion_modules.1.transformer_blocks.0.ff.net.2.bias" %in% nm) +expect_false(any(grepl("^up_blocks\\.0\\.motion_modules\\.2\\.", nm))) +expect_true(any(startsWith(nm, "mid_block."))) +expect_false(any(grepl("\\.pe$", names(ad$state_dict())))) +ad_nomid <- motion_adapter(c(32L, 64L), layers_per_block = 1L, use_mid_block = FALSE, + num_heads = 2L, max_len = 8L) +expect_false(any(startsWith(names(ad_nomid$named_parameters()), "mid_block."))) + +# A motion module keeps [B*F, C, H, W] and mixes frames +torch::torch_manual_seed(1) +mm <- diffuseR:::motion_transformer_3d(32L, heads = 2L, max_len = 8L) +x <- torch::torch_randn(2L * 3L, 32L, 4L, 4L) +y <- torch::with_no_grad(mm(x, 3L)) +expect_equal(as.integer(y$shape), c(6L, 32L, 4L, 4L)) +x2 <- x$clone() +torch::with_no_grad(x2[2, , , ]$add_(1)) +y2 <- torch::with_no_grad(mm(x2, 3L)) +# frame 2 of clip 1 changed: frames 1 and 3 of clip 1 move, clip 2 does not +expect_true((y2[1, ..] - y[1, ..])$abs()$max()$item() > 1e-4) +expect_true((y2[3, ..] - y[3, ..])$abs()$max()$item() > 1e-4) +expect_true((y2[4:6, ..] - y[4:6, ..])$abs()$max()$item() < 1e-5) + +# Tiny SD 1.5-shaped UNet + adapter. With every motion proj_out zeroed the +# modules are identities, so each frame equals the image UNet's output. +unet <- unet_native(block_out_channels = c(32L, 32L, 64L, 64L), layers_per_block = 1L, + cross_attention_dim = 16L, num_attention_heads = 2L) +ad <- motion_adapter(c(32L, 32L, 64L, 64L), layers_per_block = 1L, + use_mid_block = TRUE, num_heads = 2L, max_len = 8L) +um <- unet_motion(unet, ad) +um$eval() +lat <- torch::torch_randn(1L, 4L, 3L, 16L, 16L) +ctx <- torch::torch_randn(1L, 77L, 16L) +t <- torch::torch_tensor(500L) +out <- torch::with_no_grad(um(lat, t, ctx)) +expect_equal(as.integer(out$shape), c(1L, 4L, 3L, 16L, 16L)) +torch::with_no_grad({ + for (p in names(ad$named_parameters())) { + if (grepl("\\.proj_out\\.", p)) ad$named_parameters()[[p]]$zero_() + } + out0 <- um(lat, t, ctx) + img <- unet(lat[, , 2, , ], t, ctx) +}) +expect_true((out0[, , 2, , ] - img)$abs()$max()$item() < 1e-5) +# per-frame context: [F, 77, D] rows condition each frame +ctx_f <- torch::torch_randn(3L, 77L, 16L) +out_f <- torch::with_no_grad(um(lat, t, ctx_f)) +img3 <- torch::with_no_grad(unet(lat[, , 3, , ], t, ctx_f[3:3, , ])) +expect_true((out_f[, , 3, , ] - img3)$abs()$max()$item() < 1e-5) +expect_error(um(torch::torch_randn(1L, 4L, 9L, 16L, 16L), t, ctx), "8-frame window") + +# FreeNoise windows: stride 4 from frame 1, plus one ending on the last +# frame that adds only the uncovered frame (diffusers: (0,16) (4,20) +# (8,24) (9,25) 0-based) +w <- diffuseR:::.motion_windows(25L, 16L, 4L) +expect_equal(vapply(w$frames, `[[`, 1L, "start"), c(1L, 5L, 9L, 10L)) +expect_equal(w$frames[[4]]$keep, 16L) +expect_equal(w$weights, c(1:8, 8:1)) +expect_equal(length(diffuseR:::.motion_windows(32L, 16L, 4L)$frames), 5L) +expect_error(diffuseR:::.motion_windows(8L, 16L, 4L), "fewer than") + +# One window covering the clip is the plain motion module; frames only one +# window covers get that window's attention output +torch::torch_manual_seed(2) +blk <- diffuseR:::motion_transformer_block(32L, 2L, 16L) +xs <- torch::torch_randn(5L, 12L, 32L) +torch::with_no_grad({ + plain <- blk(xs) + one <- blk(xs, diffuseR:::.motion_windows(12L, 12L, 4L)) + win <- blk(xs, diffuseR:::.motion_windows(12L, 8L, 4L)) + first <- blk(xs[, 1:8, ]) +}) +expect_true((one - plain)$abs()$max()$item() < 1e-5) +expect_true((win[, 1:4, ] - first[, 1:4, ])$abs()$max()$item() < 1e-5) +expect_true((win[, 5:8, ] - first[, 5:8, ])$abs()$max()$item() > 1e-4) +# the UNet forward windows past context_length instead of erroring +out_w <- torch::with_no_grad(um(torch::torch_randn(1L, 4L, 10L, 16L, 16L), t, ctx, + context_length = 8L, context_stride = 2L)) +expect_equal(as.integer(out_w$shape), c(1L, 4L, 10L, 16L, 16L)) + +# FreeNoise noise: frames past the first window reuse earlier noise, shuffled +torch::torch_manual_seed(3) +lat0 <- torch::torch_randn(1L, 1L, 22L, 2L, 2L) +lat1 <- diffuseR:::.freenoise_latents(lat0, 16L, 4L) +key <- function(l, i) round(as.numeric(l[1, 1, i, 1, 1]), 5) +expect_equal(sort(sapply(17:20, key, l = lat1)), sort(sapply(1:4, key, l = lat0))) +expect_equal(sort(sapply(21:22, key, l = lat1)) %in% sapply(5:8, key, l = lat1), + c(TRUE, TRUE)) +expect_true((lat1[, , 1:16, , ] - lat0[, , 1:16, , ])$abs()$max()$item() == 0) + +# Windows longer than the 16 training frames warn (the fake pipeline then +# fails at its missing text encoder) +expect_warning(try(txt2vid_animatediff("a fox", pipeline = list(train_frames = 16L), + context_length = 24L, save_file = FALSE), + silent = TRUE), + "trained on") + +# DDIM timesteps as diffusers: linspace (rounded half to even) and leading +expect_equal(diffuseR:::.ddim_timesteps(6L, "linspace"), c(999L, 799L, 599L, 400L, 200L, 0L)) +lead <- diffuseR:::.ddim_timesteps(50L, "leading") +expect_equal(c(lead[1], lead[50]), c(981L, 1L)) +ac <- diffuseR:::.ddim_alphas_cumprod("linear") +expect_equal(length(ac), 1000L) +s <- torch::torch_ones(2) +e <- torch::torch_zeros(2) +# zero noise prediction: x0 = x / sqrt(a_t); last step lands on alphas_cumprod[1] +last <- diffuseR:::.ddim_step(s, e, 0L, 6L, ac) +expect_equal(as.numeric(last[1]), sqrt(ac[1]) / sqrt(ac[1]), tolerance = 1e-6) +mid <- diffuseR:::.ddim_step(s, e, 599L, 6L, ac) +# previous timestep is t - 1000 %/% n = 433, i.e. alphas_cumprod[434] +expect_equal(as.numeric(mid[1]), sqrt(ac[434]) / sqrt(ac[600]), tolerance = 1e-6) + +# Prompt keyframes blend linearly and hold outside the keyframes +enc <- function(p) torch::torch_full(c(1L, 2L, 3L), switch(p, a = 0, b = 10)) +pe <- diffuseR:::.animatediff_prompt_embeds(c("2" = "a", "6" = "b"), 8L, enc) +v <- as.numeric(pe[, 1, 1]) +expect_equal(v, c(0, 0, 2.5, 5, 7.5, 10, 10, 10)) +expect_equal(as.integer(diffuseR:::.animatediff_prompt_embeds("a", 8L, enc)$shape[1]), 1L) +expect_error(diffuseR:::.animatediff_prompt_embeds(c("0" = "a"), 8L, enc), "frame number") diff --git a/man/animatediff.Rd b/man/animatediff.Rd new file mode 100644 index 0000000..a95ac30 --- /dev/null +++ b/man/animatediff.Rd @@ -0,0 +1,11 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{animatediff} +\alias{animatediff} +\title{AnimateDiff text-to-video on Stable Diffusion 1.5} +\description{ +SD 1.5 (CreativeML OpenRAIL-M) with the AnimateDiff v2 or v3 motion +module (Apache-2.0, \code{guoyww/animatediff}). The pipeline follows +diffusers' \code{AnimateDiffPipeline}: DDIM with a linear beta schedule +and linspace timesteps, classifier-free guidance, frames decoded one +chunk at a time. +} diff --git a/man/animatediff_load_pipeline.Rd b/man/animatediff_load_pipeline.Rd new file mode 100644 index 0000000..cda68ed --- /dev/null +++ b/man/animatediff_load_pipeline.Rd @@ -0,0 +1,35 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{animatediff_load_pipeline} +\alias{animatediff_load_pipeline} +\title{Load the AnimateDiff SD 1.5 pipeline} +\usage{ +animatediff_load_pipeline( + motion = "v3", + model_dir = NULL, + device = "cuda", + dtype = NULL, + verbose = TRUE +) +} +\arguments{ +\item{motion}{Motion module: "v3", "v2", or a path to a motion +safetensors file.} + +\item{model_dir}{SD 1.5 diffusers directory; NULL uses +\code{\link{download_sd15}}.} + +\item{device}{Compute device for the UNet and VAE.} + +\item{dtype}{UNet and motion module dtype; NULL means float16 on CUDA, +float32 on CPU.} + +\item{verbose}{Logical.} +} +\value{ +A list with \code{unet} (a \code{\link{unet_motion}}), + \code{decoder}, \code{text_encoder}, \code{device}, \code{dtype} and + \code{train_frames} (the longest motion window that works, 16). +} +\description{ +Load the AnimateDiff SD 1.5 pipeline +} diff --git a/man/download_animatediff.Rd b/man/download_animatediff.Rd new file mode 100644 index 0000000..84817a2 --- /dev/null +++ b/man/download_animatediff.Rd @@ -0,0 +1,23 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_animatediff} +\alias{download_animatediff} +\title{Download an AnimateDiff motion module} +\usage{ +download_animatediff(version = c("v3", "v2"), verbose = TRUE) +} +\arguments{ +\item{version}{"v3" (\code{v3_sd15_mm.ckpt}, no mid-block module) or +"v2" (\code{mm_sd_v15_v2.ckpt}).} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the path of the converted safetensors file. +} +\description{ +Fetches the original Apache-2.0 checkpoint from +\code{guoyww/animatediff} and converts it once to an fp16 safetensors +file (the pickle is read by \code{torch::load_state_dict()}, which only +rebuilds tensors). Keys keep their original names; +\code{\link{load_motion_adapter}} maps them. +} diff --git a/man/download_sd15.Rd b/man/download_sd15.Rd new file mode 100644 index 0000000..8c3e332 --- /dev/null +++ b/man/download_sd15.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_sd15} +\alias{download_sd15} +\title{Download Stable Diffusion 1.5} +\usage{ +download_sd15(verbose = TRUE) +} +\arguments{ +\item{verbose}{Logical.} +} +\value{ +Invisibly, the diffusers-layout directory. +} +\description{ +Fetches the fp16 UNet, VAE and CLIP text encoder (about 2.1 GB) from +\code{stable-diffusion-v1-5/stable-diffusion-v1-5} and links them under +the file names the native loaders read. +} diff --git a/man/load_motion_adapter.Rd b/man/load_motion_adapter.Rd new file mode 100644 index 0000000..90f5ec4 --- /dev/null +++ b/man/load_motion_adapter.Rd @@ -0,0 +1,24 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_motion_adapter} +\alias{load_motion_adapter} +\title{Load an AnimateDiff motion adapter from safetensors} +\usage{ +load_motion_adapter(path, verbose = TRUE) +} +\arguments{ +\item{path}{A \code{.safetensors} file.} + +\item{verbose}{Logical.} +} +\value{ +A \code{\link{motion_adapter}} in eval mode (float32, CPU). +} +\description{ +Reads a motion module checkpoint in either the original AnimateDiff key +layout (as converted by \code{\link{download_animatediff}}) or diffusers' +\code{MotionAdapter} layout, builds a matching +\code{\link{motion_adapter}} and fills every parameter. The stored +position tables are skipped (the module computes the same sinusoids); +any other unmatched key, shape mismatch or unfilled parameter is an +error. +} diff --git a/man/motion_adapter.Rd b/man/motion_adapter.Rd new file mode 100644 index 0000000..2492cd1 --- /dev/null +++ b/man/motion_adapter.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{motion_adapter} +\alias{motion_adapter} +\title{AnimateDiff motion adapter} +\usage{ +motion_adapter( + block_out_channels = c(320L, 640L, 1280L, 1280L), + layers_per_block = 2L, + use_mid_block = TRUE, + num_heads = 8L, + max_len = 32L +) +} +\arguments{ +\item{block_out_channels}{UNet block widths.} + +\item{layers_per_block}{Motion modules per down block.} + +\item{use_mid_block}{Whether the mid block has a motion module (v2 +does, v3 does not).} + +\item{num_heads}{Attention heads per motion module.} + +\item{max_len}{Longest frame window the positions cover.} +} +\value{ +An \code{nn_module}. +} +\description{ +The motion modules for every block of an SD UNet, laid out as +diffusers' \code{MotionAdapter} so its parameter names match the +converted checkpoints: \code{down_blocks..motion_modules.}, +\code{mid_block.motion_modules.0} and \code{up_blocks..motion_modules.}. +Up blocks carry one more module than down blocks, as the UNet's up +blocks carry one more resnet. +} diff --git a/man/motion_module.Rd b/man/motion_module.Rd new file mode 100644 index 0000000..4aecd15 --- /dev/null +++ b/man/motion_module.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{motion_module} +\alias{motion_module} +\title{AnimateDiff motion modules} +\description{ +Temporal transformers that turn the SD 1.5 image UNet into a video +model: after every resnet (and spatial attention) the UNet runs one +motion module, which attends across frames at each latent position. +Ported from diffusers \code{models/unets/unet_motion_model.py} +(\code{AnimateDiffTransformer3D}, \code{MotionModules}, +\code{MotionAdapter}, \code{UNetMotionModel}) and the double +self-attention \code{BasicTransformerBlock} in +\code{models/attention.py}, Apache-2.0. +} +\keyword{internal} diff --git a/man/sd_pipeline_from_safetensors.Rd b/man/sd_pipeline_from_safetensors.Rd index 63dec7c..2e6cc36 100644 --- a/man/sd_pipeline_from_safetensors.Rd +++ b/man/sd_pipeline_from_safetensors.Rd @@ -15,7 +15,7 @@ sd_pipeline_from_safetensors( \item{diffusers_dir}{Directory with \code{unet/}, \code{vae/}, \code{text_encoder/} subdirectories.} -\item{model_name}{Currently "sd21" (SDXL pending its second encoder).} +\item{model_name}{"sd21" or "sd15" (see \code{\link{download_sd15}}).} \item{devices}{Named list of component devices (\code{unet}, \code{decoder}, \code{text_encoder}); defaults to all-CPU.} diff --git a/man/txt2vid_animatediff.Rd b/man/txt2vid_animatediff.Rd new file mode 100644 index 0000000..ff9e694 --- /dev/null +++ b/man/txt2vid_animatediff.Rd @@ -0,0 +1,93 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{txt2vid_animatediff} +\alias{txt2vid_animatediff} +\title{Generate a video with AnimateDiff (SD 1.5)} +\usage{ +txt2vid_animatediff( + prompt, + negative_prompt = NULL, + pipeline = NULL, + width = 512L, + height = 512L, + num_frames = 16L, + context_length = 16L, + context_stride = 4L, + num_inference_steps = 25L, + guidance_scale = 7.5, + seed = NULL, + latents = NULL, + decode_chunk = 1L, + fps = 8L, + filename = NULL, + save_file = TRUE, + verbose = TRUE +) +} +\arguments{ +\item{prompt}{A prompt, or keyframes: a character vector or list named +by 1-based frame number, e.g. \code{c("1" = "a red fox", "16" = "a +wolf made of smoke")}. Frames between keyframes blend the two prompt +embeddings.} + +\item{negative_prompt}{Negative prompt (one for the whole clip).} + +\item{pipeline}{From \code{\link{animatediff_load_pipeline}}; NULL loads +the default (v3, CUDA).} + +\item{width}{Output width, a multiple of 8 (SD 1.5 is trained at 512).} + +\item{height}{Output height, a multiple of 8.} + +\item{num_frames}{Frames. Past \code{context_length} the motion +modules attend within overlapping windows (FreeNoise), so clips can +be longer than the 16 frames the modules were trained on.} + +\item{context_length}{Frames per motion window. Longer than the +pipeline's \code{train_frames} (16) warns: the v2 and v3 modules +gave noise at 20, 24 and 32 frames, though their position table +covers 32.} + +\item{context_stride}{Frames between window starts.} + +\item{num_inference_steps}{DDIM steps.} + +\item{guidance_scale}{Classifier-free guidance. Higher values (10-15) +push harder towards the prompt and can add artifacts.} + +\item{seed}{Integer seed, or NULL.} + +\item{latents}{Optional starting noise \code{[1, 4, F, H/8, W/8]}.} + +\item{decode_chunk}{Frames per VAE decode. The fp32 decoder needs about +2 GiB per 512x512 frame, so 4 at once set a 13.9 GiB peak where 1 +costs 5 GiB and about 2 s.} + +\item{fps}{Frame rate of the saved file.} + +\item{filename}{Output video path; NULL with \code{save_file = TRUE} +derives one from the prompt.} + +\item{save_file}{Logical; write the video with \code{\link{save_video}}.} + +\item{verbose}{Logical.} +} +\value{ +A list: \code{video} (array \code{[F, H, W, 3]} in 0..1), + \code{file} (or NULL) and \code{metadata}. +} +\description{ +Text-to-video from pure noise. The image model draws every frame and the +motion module ties them together loosely, so objects drift and morph; +\code{prompt} keyframes push that further by changing the subject over +the clip. +} +\examples{ +\dontrun{ +pipe <- animatediff_load_pipeline("v3") +txt2vid_animatediff("a lighthouse in a storm, oil painting", + pipeline = pipe, seed = 7, filename = "storm.mp4") +# prompt travel: the subject melts from one prompt into the next +txt2vid_animatediff(c("1" = "a red fox in snow", "16" = "a wolf made of smoke"), + pipeline = pipe, guidance_scale = 12, filename = "fox.mp4") +} +} diff --git a/man/unet_motion.Rd b/man/unet_motion.Rd new file mode 100644 index 0000000..350ab66 --- /dev/null +++ b/man/unet_motion.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{unet_motion} +\alias{unet_motion} +\title{SD UNet with AnimateDiff motion modules} +\usage{ +unet_motion(unet, adapter) +} +\arguments{ +\item{unet}{A \code{\link{unet_native}}.} + +\item{adapter}{A \code{\link{motion_adapter}} with the same block widths.} +} +\value{ +An \code{nn_module}. +} +\description{ +Wraps a native SD 1.5 \code{\link{unet_native}} and a +\code{\link{motion_adapter}}. The forward takes latents +\code{[B, 4, F, H, W]}; frames run through the image layers as a batch +of \code{B * F} and through each motion module along the frame axis. +The text context is either \code{[B, 77, D]} (repeated per frame) or +\code{[B * F, 77, D]}, one row per frame, which is how prompt travel +conditions each frame on its own prompt. +} +\details{ +The forward also takes \code{context_length} and \code{context_stride}: +with more frames than \code{context_length}, the motion modules attend +within overlapping windows of that length (FreeNoise) instead of over +the whole clip. + +} diff --git a/man/unet_native.Rd b/man/unet_native.Rd index df9d8a0..c3289d5 100644 --- a/man/unet_native.Rd +++ b/man/unet_native.Rd @@ -9,7 +9,8 @@ unet_native( block_out_channels = c(320L, 640L, 1280L, 1280L), layers_per_block = 2L, cross_attention_dim = 1024L, - attention_head_dim = 64L + attention_head_dim = 64L, + num_attention_heads = NULL ) } \arguments{ @@ -24,6 +25,10 @@ unet_native( \item{cross_attention_dim}{Context dimension from text encoder} \item{attention_head_dim}{Dimension per attention head} + +\item{num_attention_heads}{NULL (SD 2.1: heads = channels / +\code{attention_head_dim}) or a fixed head count for every block +(SD 1.5: 8, so the head size grows with the channels).} } \value{ An nn_module representing the UNet