Commit cd5594e for stable-diffusion.cpp
commit cd5594e9f9e5964388b716036b319f1a81dd7c7e
Author: leejet <leejet714@gmail.com>
Date: Mon Nov 9 01:13:57 2026 +0800
feat: add SeedVR2 image and video upscaling (#2300)
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index ca64ad1..15045dc 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -29,6 +29,18 @@ PRs should include:
- Linked issue/PR context when applicable.
- Screenshots or sample outputs for UI/visual behavior changes.
+## Documentation
+
+Write model usage guides in `docs/` for users. Focus on supported weights,
+setup, runnable commands, parameter meanings, output behavior, and practical
+limits such as memory requirements. Include technical details only when they
+help users choose settings or troubleshoot a problem.
+
+Keep tensor layouts, internal class names, kernel implementation details,
+algorithm derivations, and development history out of usage guides. Keep task-specific
+profiling logs, numerical comparison procedures, and test artifacts local, and
+summarize relevant verification results in the PR description.
+
## README UI Listings
The [UIs list in the README](README.md#uis) only includes open-source projects that use `stable-diffusion.cpp` as a backend. Do not submit PRs to add closed-source UIs.
diff --git a/README.md b/README.md
index 2442865..3a65ee1 100644
--- a/README.md
+++ b/README.md
@@ -88,6 +88,9 @@ API and command-line option may change frequently.***
- [LTX-2.3/LTX-2.5](./docs/ltx2.md)
- [HunyuanVideo 1.5](./docs/hunyuan_video.md)
- [LingBot-Video](./docs/lingbot_video.md)
+ - Video Upscale Models
+ - [RelayVSR](./docs/relay_vsr.md)
+ - [SeedVR2](./docs/seedvr2.md)
- [PhotoMaker](./docs/photo_maker.md) support.
- [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
- Control Net support with SD 1.5
@@ -96,7 +99,6 @@ API and command-line option may change frequently.***
- Latent Consistency Models support (LCM/LCM-LoRA)
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
- Upscale images generated with [ESRGAN](./docs/esrgan.md)
- - Upscale videos with [RelayVSR](./docs/relay_vsr.md) (experimental, 2x and 4x)
- Supported backends
- CPU (AVX, AVX2 and AVX512 support for x86 architectures)
- CUDA
diff --git a/docs/relay_vsr.md b/docs/relay_vsr.md
index 1352f41..e35f5fd 100644
--- a/docs/relay_vsr.md
+++ b/docs/relay_vsr.md
@@ -1,24 +1,24 @@
-# RelayVSR video super-resolution
+# RelayVSR image and video super-resolution
-RelayVSR uses a Wan model to sample sparse reference latents, followed by the
-Flash conditional decoder to reconstruct every input frame. The model code
-follows [RelayVSR](https://github.com/kopperx/RelayVSR); the loader maps
-[ComfyUI repackaged weights](https://huggingface.co/Comfy-Org/RelayVSR) to the
-original module names.
+[RelayVSR](https://github.com/kopperx/RelayVSR) upscales images and videos by
+2x or 4x with `-M vid_upscale`.
-Use `-M vid_upscale` with an ordered directory of RGB images. Filenames are
-sorted lexicographically. Supply the source frame rate with `--fps`. Output
-keeps the input frame count. Use a video output supported by the CLI, such as
-`.avi` or `.webm`, or a numbered PNG output such as `output/%04d.png`.
+## Weights
-The scale is an explicit, required argument. It is not read from checkpoint
-metadata. Use matching Wan and Flash checkpoints:
+[ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/RelayVSR) load
+directly. Select the matching pair for the requested scale:
-| Scale | `--diffusion-model` | `--vae` |
+| `--upscale-factor` | `--diffusion-model` | `--vae` |
| --- | --- | --- |
| 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` |
| 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` |
+No separate text encoder or ordinary Wan VAE is needed.
+
+## Usage
+
+For a video, extract numbered frames into a directory, then run:
+
```sh
sd-cli -M vid_upscale \
--diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \
@@ -27,44 +27,24 @@ sd-cli -M vid_upscale \
--vsr-window-size 16 -s 42 -o output.avi
```
-`--vsr-window-size` includes both reference endpoints. The default 16 samples
-references at frame indices 0, 15, 30, and so on, always including the last
-frame. Shared endpoints update Flash history in the following window but are
-emitted only once. Single-frame input is also supported.
-
-Wan performs one flow-matching Euler step at timestep 1000 with guidance 1,
-using the prompt embedding stored in the checkpoint. No text encoder or
-ordinary Wan VAE is needed. Generic generation sampling settings do not
-configure this fixed sampling path.
-
-The Flash scale cannot be verified from its tensor shapes: the x2 and x4 files
-have identical shapes. The caller must select the correct pair. Wan's input
-projection is checked against the requested scale. Inputs are replicate-padded
-internally and outputs are cropped to exactly the input size times the scale.
-
-INT8 convrot uses the existing [quantized-weight execution path](int8_convrot.md).
-Backend and parameter placement use `--backend` and `--params-backend`, with
-Wan assigned to `diffusion` and Flash to `vae`. CUDA builds with CUTLASS use
-fused window attention on compute capability 8.0 or newer. It loads spatial
-neighborhoods directly, with FP16 Q/K/V and FP32 accumulation, and preserves
-Wan's boundary-shifted windows and Flash's boundary-truncated windows.
-Other backends use bounded spatial tiles with ordinary attention operators.
-Set `GGML_CUDA_DISABLE_WINDOW_ATTENTION=1` to use that fallback on CUDA for
-comparison. Floating-point results can differ between these implementations.
-
-Library callers initialize `sd_video_upscaler_ctx_params_t` with
-`sd_video_upscaler_ctx_params_init()`, set the model paths, explicit scale and
-backend options, and create a `video_upscaler_ctx_t` with
-`new_video_upscaler_ctx()`. Model architecture is detected from
-weight names and shapes; the scale remains caller-supplied.
-
-Initialize `sd_video_upscale_params_t` with `sd_video_upscale_params_init()`,
-then set the RGB frames, frame count, FPS, seed and RelayVSR reference window.
-Call `upscale_video()` and free the result with
-`free_sd_video_generation_result()`. Defaults are 24 FPS, seed 42 and a
-16-frame reference window. No sampler or scheduler selection is required.
-
-Each call starts fresh
-temporal state. Model caches are bounded by the temporal windows; the public
-API currently holds the complete input and output frame arrays in host memory.
-Audio is not copied by this path.
+Frames are read in filename order. Set `--fps` to the source video's frame rate.
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o output/%04d.png`.
+
+For a single image, replace `--input-frames input_frames` with `-i input.png`
+and use `-o output.png` (or JPEG/WebP). No FPS setting is needed for an image.
+`-i` / `--init-img` and `--input-frames` are mutually exclusive.
+
+## Settings and limits
+
+- `--upscale-factor` is required. Use 2 or 4 with the corresponding weights.
+ Output width and height are exactly the input dimensions times the scale.
+- `--vsr-window-size` sets the reference-frame window, including both endpoints.
+ The default 16 selects reference frames 0, 15, 30, and so on, and always
+ includes the final frame. Every input frame appears once in the output.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+ [backend selection](backend.md) to choose a device or weight placement.
+- Sampler, scheduler, and step settings do not configure RelayVSR.
+ `--vsr-temporal-chunk-size` is a SeedVR2 option.
+- Long videos require enough system memory to hold the input and output frames.
+ Audio is not copied to the output.
diff --git a/docs/seedvr2.md b/docs/seedvr2.md
new file mode 100644
index 0000000..24f5503
--- /dev/null
+++ b/docs/seedvr2.md
@@ -0,0 +1,69 @@
+# SeedVR2 image and video super-resolution
+
+SeedVR2 upscales images and videos with `-M vid_upscale`. Support is experimental.
+
+## Weights
+
+These [ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/SeedVR2)
+load directly, without conversion or a separate text encoder:
+
+| Model | `--diffusion-model` | `--vae` |
+| --- | --- | --- |
+| 3B | `seedvr2_3b_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+| 7B Sharp | `seedvr2_7b_sharp_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+
+Use the matching SeedVR2 VAE for either model. The 7B model requires more memory.
+
+## Upscale an image
+
+```sh
+sd-cli -M vid_upscale \
+ --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+ --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+ -i input.png --upscale-factor 2 -s 42 -o restored.png
+```
+
+`-i` / `--init-img` accepts a single image. Outputs can be PNG, JPEG, or WebP.
+No FPS or temporal chunk setting is needed for an image.
+
+## Upscale a video
+
+Extract the video into numbered frames, then pass the frame directory:
+
+```sh
+mkdir frames
+ffmpeg -i input.mp4 -fps_mode passthrough frames/%06d.png
+sd-cli -M vid_upscale \
+ --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+ --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+ --input-frames frames --upscale-factor 2 --fps 24 \
+ --vsr-temporal-chunk-size 1 -s 42 -o restored.avi
+```
+
+Set `--fps` to the source video's frame rate; 24 is only the example above.
+Frames are read in filename order. Use either `-i` or `--input-frames`, not both.
+
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o restored/%06d.png`.
+
+## Settings and limits
+
+- `--upscale-factor` is required and must be a positive integer. Output width
+ and height are exactly the input dimensions multiplied by this value.
+- `--vsr-temporal-chunk-size` controls how many video frames are processed
+ together. Valid values are 1, 5, 9, 13, and so on; the default is 5. A value of
+ 1 processes each frame independently and uses less memory. Larger values use
+ neighboring frames but require more memory. Chunk boundaries may be visible.
+- Frame count and order are preserved, including a shorter final chunk. Audio
+ is not copied to the output.
+- If you run out of GPU memory, try `--vsr-temporal-chunk-size 1`. A five-frame
+ chunk at 1728x960 output resolution exceeded 24 GB of GPU memory in testing.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+ [backend selection](backend.md) to choose a device or weight placement.
+- SeedVR2 does not use sampler, scheduler, or step settings.
+ `--vsr-window-size` is a RelayVSR option and does not affect SeedVR2.
+
+3B has been tested on CUDA and CPU; 7B has been tested on CUDA with small inputs.
+Other backends have not been tested with this integration.
+
+Model reference: [ByteDance-Seed/SeedVR](https://github.com/ByteDance-Seed/SeedVR).
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
index bcf7359..2bff050 100644
--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -35,12 +35,13 @@ std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
struct SDCliParams {
std::string input_frames;
- int video_upscale_factor = 0;
- int vsr_window_size = 16;
- SDMode mode = IMG_GEN;
- std::string output_path = "output.png";
- int output_begin_idx = -1;
- int compression_quality = 90;
+ int video_upscale_factor = 0;
+ int vsr_window_size = 16;
+ int vsr_temporal_chunk_size = 5;
+ SDMode mode = IMG_GEN;
+ std::string output_path = "output.png";
+ int output_begin_idx = -1;
+ int compression_quality = 90;
std::string image_path;
std::string metadata_format = "text";
@@ -66,7 +67,7 @@ struct SDCliParams {
ArgOptions options;
options.string_options = {
- {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 0, &input_frames},
+ {"", "--input-frames", "directory of ordered RGB frames for vid_upscale (use -i for a single image)", 0, &input_frames},
{"-o",
"--output",
"path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
@@ -95,8 +96,9 @@ struct SDCliParams {
};
options.int_options = {
- {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor},
+ {"", "--upscale-factor", "required positive video upscale factor (RelayVSR: 2 or 4; SeedVR2: integer)", &video_upscale_factor},
{"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size},
+ {"", "--vsr-temporal-chunk-size", "SeedVR2 input frames per independent chunk: 1 or 4n+1 (default: 5)", &vsr_temporal_chunk_size},
{"",
"--preview-interval",
"preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
@@ -227,8 +229,8 @@ struct SDCliParams {
}
bool validate() {
- if (mode == VID_UPSCALE && (input_frames.empty() || video_upscale_factor <= 0 || vsr_window_size < 2)) {
- LOG_ERROR("vid_upscale requires --input-frames, positive --upscale-factor, and --vsr-window-size >= 2");
+ if (mode == VID_UPSCALE && video_upscale_factor <= 0) {
+ LOG_ERROR("vid_upscale requires a positive --upscale-factor");
return false;
}
if (mode != METADATA) {
@@ -540,7 +542,9 @@ bool save_results(const SDCliParams& cli_params,
return sucessful_reults != 0;
}
- if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) {
+ const bool video_upscale_output = cli_params.mode == VID_UPSCALE &&
+ (gen_params.init_image_path.empty() || ext_lower == ".avi" || ext_lower == ".webm");
+ if ((cli_params.mode == VID_GEN && num_results > 1) || video_upscale_output) {
if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
ext = ".avi";
std::string params = gen_params.embed_image_metadata
@@ -676,12 +680,23 @@ int main(int argc, const char* argv[]) {
}
if (cli_params.mode == VID_UPSCALE) {
+ if (cli_params.input_frames.empty() == gen_params.init_image_path.empty()) {
+ LOG_ERROR("vid_upscale requires exactly one input: -i/--init-img or --input-frames");
+ return 1;
+ }
+ ctx_params.prepare_backend_assignments();
if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) {
LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps");
return 1;
}
std::vector<SDImageOwner> owned_frames;
- if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
+ if (!gen_params.init_image_path.empty()) {
+ owned_frames.emplace_back();
+ if (!load_sd_image_from_file(owned_frames.back().put(), gen_params.init_image_path.c_str(), 0, 0, 3)) {
+ LOG_ERROR("load image from '%s' failed", gen_params.init_image_path.c_str());
+ return 1;
+ }
+ } else if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
LOG_ERROR("vid_upscale input contains no readable frames");
return 1;
}
@@ -692,12 +707,22 @@ int main(int argc, const char* argv[]) {
frames.push_back(frame.get());
sd_video_upscaler_ctx_params_t params;
sd_video_upscaler_ctx_params_init(¶ms);
- params.diffusion_model_path = ctx_params.diffusion_model_path.c_str();
- params.vae_path = ctx_params.vae_path.c_str();
- params.upscale_factor = cli_params.video_upscale_factor;
- params.n_threads = ctx_params.n_threads;
- params.backend = ctx_params.effective_backend.c_str();
- params.params_backend = ctx_params.effective_params_backend.c_str();
+ params.diffusion_model_path = ctx_params.diffusion_model_path.c_str();
+ params.vae_path = ctx_params.vae_path.c_str();
+ params.upscale_factor = cli_params.video_upscale_factor;
+ params.n_threads = ctx_params.n_threads;
+ params.backend = ctx_params.effective_backend.c_str();
+ params.params_backend = ctx_params.effective_params_backend.c_str();
+ params.flash_attn_mode = ctx_params.flash_attn_mode.c_str();
+ std::string conv_direct_mode = ctx_params.conv_direct_mode;
+ if (conv_direct_mode.empty()) {
+ conv_direct_mode = "auto";
+ if (ctx_params.diffusion_conv_direct)
+ conv_direct_mode += ",diffusion=on";
+ if (ctx_params.vae_conv_direct)
+ conv_direct_mode += ",vae=on";
+ }
+ params.conv_direct_mode = conv_direct_mode.c_str();
std::unique_ptr<video_upscaler_ctx_t, decltype(&free_video_upscaler_ctx)> upscaler(
new_video_upscaler_ctx(¶ms),
free_video_upscaler_ctx);
@@ -706,11 +731,12 @@ int main(int argc, const char* argv[]) {
sd_video_generation_result_t* output = nullptr;
sd_video_upscale_params_t upscale_params;
sd_video_upscale_params_init(&upscale_params);
- upscale_params.frames = frames.data();
- upscale_params.frame_count = int(frames.size());
- upscale_params.fps = gen_params.fps;
- upscale_params.window_size = cli_params.vsr_window_size;
- upscale_params.seed = gen_params.seed;
+ upscale_params.frames = frames.data();
+ upscale_params.frame_count = int(frames.size());
+ upscale_params.fps = gen_params.fps;
+ upscale_params.window_size = cli_params.vsr_window_size;
+ upscale_params.temporal_chunk_size = cli_params.vsr_temporal_chunk_size;
+ upscale_params.seed = gen_params.seed;
if (!upscale_video(upscaler.get(), &upscale_params, &output))
return 1;
std::string video_info = output->generation_info;
diff --git a/examples/common/common.cpp b/examples/common/common.cpp
index a3496f4..70107a3 100644
--- a/examples/common/common.cpp
+++ b/examples/common/common.cpp
@@ -1136,7 +1136,7 @@ ArgOptions SDGenerationParams::get_options() {
&extra_ad_args},
{"-i",
"--init-img",
- "path to the init image",
+ "path to the init image, or a single input image for vid_upscale",
0,
&init_image_path},
{"",
diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h
index a1d82e7..53c6b70 100644
--- a/include/stable-diffusion.h
+++ b/include/stable-diffusion.h
@@ -625,6 +625,9 @@ typedef struct {
int n_threads;
const char* backend;
const char* params_backend;
+ // auto/on/off or diffusion/vae assignments; null selects auto.
+ const char* flash_attn_mode;
+ const char* conv_direct_mode;
} sd_video_upscaler_ctx_params_t;
typedef struct {
@@ -634,6 +637,8 @@ typedef struct {
int64_t seed;
// RelayVSR reference window, including both endpoints; defaults to 16.
int window_size;
+ // SeedVR2 independent temporal chunks in input frames (1 or 4n+1); defaults to 5.
+ int temporal_chunk_size;
} sd_video_upscale_params_t;
SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params);
diff --git a/src/model/diffusion/seedvr2.hpp b/src/model/diffusion/seedvr2.hpp
new file mode 100644
index 0000000..87e542e
--- /dev/null
+++ b/src/model/diffusion/seedvr2.hpp
@@ -0,0 +1,325 @@
+#ifndef __SD_MODEL_DIFFUSION_SEEDVR2_HPP__
+#define __SD_MODEL_DIFFUSION_SEEDVR2_HPP__
+
+#include "model/common/vsr.hpp"
+
+namespace SeedVR2 {
+
+ struct SeedVR2Config {
+ int dim = 0;
+ int layers = 0;
+ int mm_layers = 0;
+ int hidden = 0;
+ int text_dim = 0;
+ int text_len = 0;
+ int rope_freqs = 0;
+ bool swiglu = false;
+
+ static SeedVR2Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+ SeedVR2Config config;
+ auto at = [&](const std::string& key) -> const TensorStorage& {
+ auto found = weights.find(prefix + "." + key);
+ if (found == weights.end())
+ throw std::runtime_error("Missing SeedVR2 tensor: " + prefix + "." + key);
+ return found->second;
+ };
+ config.dim = int(at("vid_in.proj.weight").ne[1]);
+ config.hidden = int(at("blocks.0.mlp.vid.proj_in.weight").ne[1]);
+ config.text_dim = int(at("positive_conditioning").ne[0]);
+ config.text_len = int(at("positive_conditioning").ne[1]);
+ config.rope_freqs = int(at("blocks.0.attn.rope.rope.freqs").ne[0]);
+ config.swiglu = weights.count(prefix + ".blocks.0.mlp.vid.proj_in_gate.weight") != 0;
+ while (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.rope.rope.freqs")) {
+ if (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.proj_qkv.vid.weight"))
+ ++config.mm_layers;
+ ++config.layers;
+ }
+ if (at("vid_in.proj.weight").logical_ne0() != 132 || at("vid_out.proj.weight").ne[1] != 64 ||
+ config.text_dim != 5120 ||
+ !(config.swiglu ? config.dim == 2560 && config.layers == 32 && config.mm_layers == 10 && config.rope_freqs == 21
+ : config.dim == 3072 && config.layers == 36 && config.mm_layers == 36 && config.rope_freqs == 10))
+ throw std::runtime_error("Unsupported SeedVR2 architecture");
+ LOG_VERBOSE("seedvr2: dim = %d, layers = %d, mm_layers = %d, hidden = %d, rope_freqs = %d", config.dim, config.layers, config.mm_layers, config.hidden, config.rope_freqs);
+ return config;
+ }
+ };
+
+ struct WindowPlan {
+ sd::Tensor<int32_t> order;
+ sd::Tensor<int32_t> reverse;
+ sd::Tensor<float> positions;
+ sd::Tensor<float> text_positions;
+ std::vector<int> lengths;
+
+ WindowPlan(int width, int height, int frames, int text_len, bool shifted, bool multimodal)
+ : order({width * height * frames}), reverse({width * height * frames}), positions({1, 3, width * height * frames}), text_positions({1, 3, text_len}) {
+ const double scale = std::sqrt(720.0 * 1280.0 / 256.0 / (width * height));
+ const int wh = int(std::ceil(std::nearbyint(height * scale) / 3));
+ const int ww = int(std::ceil(std::nearbyint(width * scale) / 3));
+ const int wt = (std::min(frames, 30) + 3) / 4;
+ const double st = shifted && wt < frames ? 0.5 : 0;
+ const double sh = shifted && wh < height ? 0.5 : 0;
+ const double sw = shifted && ww < width ? 0.5 : 0;
+ const int nt = st ? int(std::ceil((frames - st) / wt)) + 1 : (frames + wt - 1) / wt;
+ const int nh = sh ? int(std::ceil((height - sh) / wh)) + 1 : (height + wh - 1) / wh;
+ const int nw = sw ? int(std::ceil((width - sw) / ww)) + 1 : (width + ww - 1) / ww;
+ int index = 0;
+ for (int iw = 0; iw < nw; ++iw)
+ for (int ih = 0; ih < nh; ++ih)
+ for (int it = 0; it < nt; ++it) {
+ int x0 = std::max(0, int((iw - sw) * ww)), x1 = std::min(width, int((iw - sw + 1) * ww));
+ int y0 = std::max(0, int((ih - sh) * wh)), y1 = std::min(height, int((ih - sh + 1) * wh));
+ int t0 = std::max(0, int((it - st) * wt)), t1 = std::min(frames, int((it - st + 1) * wt));
+ if (x1 <= x0 || y1 <= y0 || t1 <= t0)
+ continue;
+ lengths.push_back((x1 - x0) * (y1 - y0) * (t1 - t0));
+ for (int t = t0; t < t1; ++t)
+ for (int y = y0; y < y1; ++y)
+ for (int x = x0; x < x1; ++x, ++index) {
+ int original = (t * height + y) * width + x;
+ order.data()[index] = original;
+ reverse.data()[original] = index;
+ const int pos[] = {t - t0, y - y0, x - x0};
+ const int size[] = {t1 - t0, y1 - y0, x1 - x0};
+ for (int a = 0; a < 3; ++a)
+ positions.data()[3 * index + a] = multimodal ? float(pos[a] + (a == 0 ? text_len : 0))
+ : (size[a] == 1 ? -1.f : -1.f + 2.f * pos[a] / (size[a] - 1));
+ }
+ }
+ if (index != width * height * frames)
+ throw std::runtime_error("SeedVR2 window partition is incomplete");
+ for (int i = 0; i < text_len; ++i)
+ for (int a = 0; a < 3; ++a)
+ text_positions.data()[3 * i + a] = multimodal ? float(i) : 0.f;
+ }
+ };
+
+ class AdaSingle : public GGMLBlock {
+ int dim_;
+ bool output_;
+
+ protected:
+ void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+ for (const auto& layer : output_ ? std::vector<std::string>{"out"} : std::vector<std::string>{"attn", "mlp"}) {
+ for (const auto& mode : output_ ? std::vector<std::string>{"shift", "scale"} : std::vector<std::string>{"shift", "scale", "gate"})
+ params[layer + "_" + mode] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim_);
+ }
+ }
+
+ public:
+ AdaSingle(int dim, bool output = false)
+ : dim_(dim), output_(output) {}
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb, const std::string& layer, bool gate = false) {
+ auto c = ctx->ggml_ctx;
+ // The reference reshapes embedding to [D, 2, 3], not [6, D].
+ const int offset = layer == "mlp" ? 3 : 0;
+ auto mod = [&](int i, const char* name) {
+ auto value = ggml_cont(c, ggml_view_2d(c, emb, 1, dim_, 6 * emb->nb[0], (offset + i) * emb->nb[0]));
+ return ggml_add(c, ggml_reshape_1d(c, value, dim_), params[layer + "_" + name]);
+ };
+ if (gate)
+ return ggml_mul(c, x, mod(2, "gate"));
+ // AdaSingle uses scale directly; ggml_ext_modulate adds an identity scale.
+ return ggml_add(c, ggml_mul(c, x, mod(1, "scale")), mod(0, "shift"));
+ }
+ };
+
+ class SRMLP : public UnaryBlock {
+ bool swiglu_;
+
+ public:
+ explicit SRMLP(const SeedVR2Config& config)
+ : swiglu_(config.swiglu) {
+ blocks["proj_in"] = std::make_shared<Linear>(config.dim, config.hidden, !swiglu_);
+ blocks["proj_out"] = std::make_shared<Linear>(config.hidden, config.dim, !swiglu_);
+ if (swiglu_)
+ blocks["proj_in_gate"] = std::make_shared<Linear>(config.dim, config.hidden, false);
+ }
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+ auto up = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, x);
+ auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast<Linear>(blocks["proj_in_gate"])->forward(ctx, x), up)
+ : ggml_ext_gelu(ctx->ggml_ctx, up);
+ return std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, hidden);
+ }
+ };
+
+ class NaMMSRTransformerBlock : public GGMLBlock {
+ SeedVR2Config config_;
+ bool shared_;
+ bool last_;
+
+ protected:
+ void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+ params["attn.rope.rope.freqs"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config_.rope_freqs);
+ }
+
+ public:
+ NaMMSRTransformerBlock(const SeedVR2Config& config, int index)
+ : config_(config), shared_(index >= config.mm_layers), last_(config.swiglu && index == config.layers - 1) {
+ for (const auto& branch : shared_ ? std::vector<std::string>{"all"} : std::vector<std::string>{"vid", "txt"}) {
+ blocks["attn.proj_qkv." + branch] = std::make_shared<Linear>(config.dim, 3 * config.dim, false);
+ blocks["attn.proj_out." + branch] = std::make_shared<Linear>(config.dim, config.dim);
+ blocks["attn.norm_q." + branch] = std::make_shared<RMSNorm>(128, 1e-5f);
+ blocks["attn.norm_k." + branch] = std::make_shared<RMSNorm>(128, 1e-5f);
+ blocks["ada." + branch] = std::make_shared<AdaSingle>(config.dim);
+ blocks["mlp." + branch] = std::make_shared<SRMLP>(config);
+ }
+ }
+
+ ggml_tensor* rotation(GGMLRunnerContext* ctx, ggml_tensor* positions) {
+ auto c = ctx->ggml_ctx;
+ const int n = config_.rope_freqs;
+ auto p = ggml_repeat(c, positions, ggml_new_tensor_3d(c, GGML_TYPE_F32, n, 3, positions->ne[2]));
+ auto angle = ggml_mul(c, p, params["attn.rope.rope.freqs"]);
+ angle = ggml_reshape_2d(c, angle, 3 * n, positions->ne[2]);
+ angle = ggml_pad(c, angle, 64 - 3 * n, 0, 0, 0);
+ // Keep tokens out of ne[3]: CUDA concat otherwise launches once per token.
+ auto cs = ggml_reshape_3d(c, ggml_cos(c, angle), 1, 1, 64 * positions->ne[2]);
+ auto sn = ggml_reshape_3d(c, ggml_sin(c, angle), 1, 1, 64 * positions->ne[2]);
+ auto matrix = ggml_concat(c, ggml_concat(c, cs, ggml_scale(c, sn, -1.f), 0), ggml_concat(c, sn, cs, 0), 1);
+ return ggml_reshape_4d(c, matrix, 2, 2, 64, positions->ne[2]);
+ }
+
+ std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vid, ggml_tensor* txt, ggml_tensor* emb, const WindowPlan& plan, VSR::GraphInputs& inputs) {
+ auto c = ctx->ggml_ctx;
+ const int dim = config_.dim, heads = dim / 128;
+ auto branch = [&](bool text) -> std::string { return shared_ ? "all" : text ? "txt"
+ : "vid"; };
+ auto linear = [&](const std::string& name, bool text, ggml_tensor* x) {
+ return std::dynamic_pointer_cast<Linear>(blocks[name + "." + branch(text)])->forward(ctx, x);
+ };
+ auto ada = [&](ggml_tensor* x, bool text, const std::string& layer, bool gate = false) {
+ if (text && last_)
+ return x;
+ return std::dynamic_pointer_cast<AdaSingle>(blocks["ada." + branch(text)])->forward(ctx, x, emb, layer, gate);
+ };
+ auto vqkv = linear("attn.proj_qkv", false, ada(ggml_rms_norm(c, vid, 1e-5f), false, "attn"));
+ auto tqkv = linear("attn.proj_qkv", true, ada(ggml_rms_norm(c, txt, 1e-5f), true, "attn"));
+ vqkv = ggml_get_rows(c, vqkv, inputs.input(ctx, plan.order));
+ auto vpe = rotation(ctx, inputs.input(ctx, plan.positions));
+ auto tpe = rotation(ctx, inputs.input(ctx, plan.text_positions));
+ auto qkv = [&](ggml_tensor* x, bool text, ggml_tensor* pe) {
+ auto parts = ggml_ext_chunk(c, x, 3, 0);
+ for (int i = 0; i < 2; ++i) {
+ auto h = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]);
+ h = std::dynamic_pointer_cast<RMSNorm>(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h);
+ h = ggml_reshape_2d(c, h, dim, x->ne[1]);
+ parts[i] = VSR::apply_rope(ctx, h, pe, heads);
+ }
+ return parts;
+ };
+ auto vp = qkv(vqkv, false, vpe), tp = qkv(tqkv, true, tpe);
+ std::vector<ggml_tensor*> video_windows;
+ ggml_tensor* tout = nullptr;
+ int start = 0;
+ for (int length : plan.lengths) {
+ std::vector<ggml_tensor*> joined;
+ for (int i = 0; i < 3; ++i)
+ joined.push_back(ggml_concat(c, ggml_ext_slice(c, vp[i], 1, start, start + length), tp[i], 1));
+ auto out = ggml_ext_attention_ext(c, ctx->backend, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled);
+ auto vo = ggml_ext_slice(c, out, 1, 0, length);
+ auto to = ggml_ext_slice(c, out, 1, length, length + txt->ne[1]);
+ video_windows.push_back(vo);
+ tout = tout ? ggml_add(c, tout, to) : to;
+ start += length;
+ }
+ auto vout = ggml_ext_vec_concat(c, video_windows, 1);
+ vout = ggml_get_rows(c, vout, inputs.input(ctx, plan.reverse));
+ tout = ggml_scale(c, tout, 1.f / plan.lengths.size());
+ vid = ggml_add(c, vid, ada(linear("attn.proj_out", false, vout), false, "attn", true));
+ txt = ggml_add(c, txt, ada(linear("attn.proj_out", true, tout), true, "attn", true));
+ auto mlp = [&](ggml_tensor* x, bool text) {
+ auto h = ada(ggml_rms_norm(c, x, 1e-5f), text, "mlp");
+ h = std::dynamic_pointer_cast<SRMLP>(blocks["mlp." + branch(text)])->forward(ctx, h);
+ return ggml_add(c, x, ada(h, text, "mlp", true));
+ };
+ vid = mlp(vid, false);
+ if (!last_)
+ txt = mlp(txt, true);
+ return {vid, txt};
+ }
+ };
+
+ class NaDiT : public GGMLBlock {
+ SeedVR2Config config_;
+
+ protected:
+ void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+ params["positive_conditioning"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config_.text_dim, config_.text_len);
+ }
+
+ public:
+ explicit NaDiT(const SeedVR2Config& config)
+ : config_(config) {
+ blocks["vid_in.proj"] = std::make_shared<Linear>(132, config.dim);
+ blocks["txt_in"] = std::make_shared<Linear>(config.text_dim, config.dim);
+ blocks["emb_in.proj_in"] = std::make_shared<Linear>(256, config.dim);
+ blocks["emb_in.proj_hid"] = std::make_shared<Linear>(config.dim, config.dim);
+ blocks["emb_in.proj_out"] = std::make_shared<Linear>(config.dim, 6 * config.dim);
+ for (int i = 0; i < config.layers; ++i)
+ blocks["blocks." + std::to_string(i)] = std::make_shared<NaMMSRTransformerBlock>(config, i);
+ blocks["vid_out.proj"] = std::make_shared<Linear>(config.dim, 64);
+ if (config.swiglu) {
+ blocks["vid_out_norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+ blocks["vid_out_ada"] = std::make_shared<AdaSingle>(config.dim, true);
+ }
+ }
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* timestep, const std::vector<WindowPlan>& plans, VSR::GraphInputs& inputs) {
+ auto linear = [&](const std::string& name, ggml_tensor* x) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x); };
+ auto vid = linear("vid_in.proj", input);
+ auto txt = linear("txt_in", params["positive_conditioning"]);
+ auto emb = linear("emb_in.proj_in", timestep);
+ emb = linear("emb_in.proj_hid", ggml_silu(ctx->ggml_ctx, emb));
+ emb = linear("emb_in.proj_out", ggml_silu(ctx->ggml_ctx, emb));
+ for (int i = 0; i < config_.layers; ++i) {
+ auto name = "blocks." + std::to_string(i);
+ std::tie(vid, txt) = std::dynamic_pointer_cast<NaMMSRTransformerBlock>(blocks[name])->forward(ctx, vid, txt, emb, plans[i % 2], inputs);
+ sd::ggml_graph_cut::mark_graph_cut(vid, "seedvr2." + name, "vid");
+ sd::ggml_graph_cut::mark_graph_cut(txt, "seedvr2." + name, "txt");
+ }
+ if (config_.swiglu) {
+ vid = std::dynamic_pointer_cast<RMSNorm>(blocks["vid_out_norm"])->forward(ctx, vid);
+ vid = std::dynamic_pointer_cast<AdaSingle>(blocks["vid_out_ada"])->forward(ctx, vid, emb, "out");
+ }
+ return linear("vid_out.proj", vid);
+ }
+ };
+
+ struct SeedVR2Runner : GGMLRunner {
+ SeedVR2Config config;
+ NaDiT model;
+ sd::Tensor<float> timestep;
+
+ SeedVR2Runner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr<DeviceResidencyManager> manager)
+ : GGMLRunner(backend, manager), config(SeedVR2Config::detect_from_weights(weights, "diffusion")), model(config), timestep({256}) {
+ model.init(params_ctx, weights, "diffusion");
+ for (int i = 0; i < 128; ++i) {
+ const double angle = 1000.0 * std::pow(10000.0, -double(i) / 128);
+ timestep.data()[i] = float(std::sin(angle));
+ timestep.data()[i + 128] = float(std::cos(angle));
+ }
+ }
+ std::string get_desc() override { return "SeedVR2 DiT"; }
+ void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) { model.get_param_tensors(tensors, "diffusion"); }
+
+ sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input, int width, int height, int frames) {
+ std::vector<WindowPlan> plans;
+ for (bool shifted : {false, true})
+ plans.emplace_back(width, height, frames, config.text_len, shifted, config.swiglu);
+ VSR::GraphInputs inputs;
+ auto graph = [&]() {
+ inputs.begin_graph();
+ auto gf = new_graph_custom(MAX_GRAPH_SIZE);
+ auto ctx = get_context(gf);
+ auto out = model.forward(&ctx, make_input(input), make_input(timestep), plans, inputs);
+ ggml_build_forward_expand(gf, out);
+ return gf;
+ };
+ auto output = GGMLRunner::compute(graph, n_threads);
+ return output ? *output : sd::Tensor<float>();
+ }
+ };
+}
+
+#endif // __SD_MODEL_DIFFUSION_SEEDVR2_HPP__
diff --git a/src/model/vae/seedvr2_vae.hpp b/src/model/vae/seedvr2_vae.hpp
new file mode 100644
index 0000000..5f38aeb
--- /dev/null
+++ b/src/model/vae/seedvr2_vae.hpp
@@ -0,0 +1,286 @@
+#ifndef __SD_MODEL_VAE_SEEDVR2_VAE_HPP__
+#define __SD_MODEL_VAE_SEEDVR2_VAE_HPP__
+
+#include "model/common/ggml_block.hpp"
+
+namespace SeedVR2 {
+
+ struct SeedVR2VAEConfig {
+ std::vector<int> channels;
+ int latent_channels = 16;
+ int layers = 2;
+
+ static SeedVR2VAEConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+ SeedVR2VAEConfig config;
+ auto at = [&](const std::string& key) -> const TensorStorage& {
+ auto found = weights.find(prefix + "." + key);
+ if (found == weights.end())
+ throw std::runtime_error("Missing SeedVR2 VAE tensor: " + prefix + "." + key);
+ return found->second;
+ };
+ for (int i = 0; i < 4; ++i)
+ config.channels.push_back(int(at("encoder.down_blocks." + std::to_string(i) + ".resnets.0.norm2.weight").ne[0]));
+ if (config.channels != std::vector<int>({128, 256, 512, 512}) || at("encoder.conv_out.bias").ne[0] != 32 ||
+ at("decoder.conv_in.weight").ne[3] != 16 * 512 || at("decoder.conv_out.bias").ne[0] != 3 ||
+ !weights.count(prefix + ".decoder.up_blocks.0.upsamplers.0.upscale_conv.weight"))
+ throw std::runtime_error("SeedVR2 requires its matching 16-channel causal VAE");
+ LOG_VERBOSE("seedvr2 vae: latent_channels = %d, layers = %d, spatial factor = 8, temporal factor = 4", config.latent_channels, config.layers);
+ return config;
+ }
+ };
+
+ class InflatedCausalConv3d : public Conv3d {
+ int temporal_padding_;
+
+ public:
+ InflatedCausalConv3d(int in, int out, std::tuple<int, int, int> kernel = {3, 3, 3}, std::tuple<int, int, int> stride = {1, 1, 1}, std::tuple<int, int, int> padding = {1, 1, 1})
+ : Conv3d(in, out, kernel, stride, {0, std::get<1>(padding), std::get<2>(padding)}), temporal_padding_(2 * std::get<0>(padding)) {}
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+ auto c = ctx->ggml_ctx;
+ const int kt = std::get<0>(kernel_size), kh = std::get<1>(kernel_size), kw = std::get<2>(kernel_size);
+ if (ctx->conv3d_direct_enabled && !ctx->weight_adapter && x->ne[2] == 1 && kt > 1 &&
+ temporal_padding_ == kt - 1 && std::get<0>(dilation) == 1) {
+ // All causal taps read the same frame. Sum in F32 before rounding the
+ // equivalent spatial filter to F16, instead of copying and convolving the halo.
+ auto w = ggml_cast(c, params["weight"], GGML_TYPE_F32);
+ w = ggml_reshape_3d(c, w, kw * kh, kt, in_channels * out_channels);
+ auto spatial = ggml_ext_slice(c, w, 1, 0, 1, false);
+ for (int i = 1; i < kt; ++i)
+ spatial = ggml_add(c, spatial, ggml_ext_slice(c, w, 1, i, i + 1, false));
+ w = ggml_cast(c, ggml_reshape_4d(c, spatial, kw, kh, 1, in_channels * out_channels), GGML_TYPE_F16);
+ if (ctx->conv2d_direct_enabled) {
+ w = ggml_reshape_4d(c, w, kw, kh, in_channels, out_channels);
+ x = ggml_reshape_4d(c, ggml_ext_cont(c, x), x->ne[0], x->ne[1], in_channels, 1);
+ x = ggml_ext_conv_2d(c, x, w, bias ? params["bias"] : nullptr,
+ std::get<2>(stride), std::get<1>(stride),
+ std::get<2>(padding), std::get<1>(padding),
+ std::get<2>(dilation), std::get<1>(dilation), true);
+ return ggml_reshape_4d(c, x, x->ne[0], x->ne[1], 1, out_channels);
+ }
+ return ggml_ext_conv_3d(c, ctx->backend, x, w, bias ? params["bias"] : nullptr, in_channels,
+ std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
+ std::get<2>(padding), std::get<1>(padding), 0,
+ std::get<2>(dilation), std::get<1>(dilation), 1, force_prec_f32, true);
+ }
+ if (temporal_padding_) {
+ const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3];
+ // Merge spatial axes so concat does not launch separately for every channel.
+ x = ggml_reshape_3d(c, ggml_ext_cont(c, x), w * h, t, ch);
+ auto head = ggml_ext_slice(c, x, 1, 0, 1);
+ auto shape = ggml_new_tensor_3d(c, x->type, w * h, temporal_padding_, ch);
+ x = ggml_concat(c, ggml_repeat(c, head, shape), x, 1);
+ x = ggml_reshape_4d(c, x, w, h, t + temporal_padding_, ch);
+ }
+ if (ctx->conv3d_direct_enabled)
+ return Conv3d::forward(ctx, x);
+ const int st = std::get<0>(stride), sh = std::get<1>(stride), sw = std::get<2>(stride);
+ const int ph = std::get<1>(padding), pw = std::get<2>(padding);
+ const int64_t oh = (x->ne[1] + 2 * ph - kh) / sh + 1;
+ const int64_t ow = (x->ne[0] + 2 * pw - kw) / sw + 1;
+ const int64_t ot = (x->ne[2] - kt) / st + 1;
+ const int64_t row_bytes = ow * in_channels * kt * kh * kw * 2;
+ const int64_t max_patch_bytes = 256LL * 1024 * 1024;
+ if (row_bytes * oh * ot <= max_patch_bytes)
+ return Conv3d::forward(ctx, x);
+ // Bound im2col storage without spatial tiling: every strip keeps the exact
+ // convolution halo, and normalization still sees the complete frame.
+ const int64_t rows = std::max<int64_t>(1, max_patch_bytes / row_bytes);
+ struct RestorePadding {
+ std::tuple<int, int, int>& target;
+ std::tuple<int, int, int> value;
+ ~RestorePadding() { target = value; }
+ } restore{padding, padding};
+ padding = {0, 0, pw};
+ auto output = ggml_ext_zeros(c, ow, oh, ot, out_channels);
+ for (int64_t t = 0; t < ot; ++t) {
+ auto frame = ggml_ext_slice(c, x, 2, t * st, t * st + kt, false);
+ for (int64_t y = 0; y < oh; y += rows) {
+ const int64_t end = std::min(oh, y + rows);
+ const int64_t top = y * sh - ph, bottom = (end - 1) * sh - ph + kh;
+ auto strip = ggml_ext_slice(c, frame, 1, std::max<int64_t>(0, top), std::min(x->ne[1], bottom), false);
+ strip = ggml_ext_pad_ext(c, ctx->backend, strip, 0, 0, int(std::max<int64_t>(0, -top)),
+ int(std::max<int64_t>(0, bottom - x->ne[1])), 0, 0, 0, 0);
+ strip = ggml_ext_cont(c, strip);
+ strip = Conv3d::forward(ctx, strip);
+ output = ggml_set_inplace(c, output, strip, output->nb[1], output->nb[2], output->nb[3], y * output->nb[1] + t * output->nb[2]);
+ }
+ }
+ return output;
+ }
+ };
+
+ static ggml_tensor* causal_norm(GGMLRunnerContext* ctx, ggml_tensor* x, const std::shared_ptr<GGMLBlock>& norm) {
+ auto c = ctx->ggml_ctx;
+ // [B, C, T, H, W] -> [B*T, C, H, W]: GroupNorm never mixes frames.
+ x = ggml_cont(c, ggml_permute(c, x, 0, 1, 3, 2));
+ x = std::dynamic_pointer_cast<GroupNorm>(norm)->forward(ctx, x);
+ return ggml_cont(c, ggml_permute(c, x, 0, 1, 3, 2));
+ }
+
+ class ResnetBlock3D : public UnaryBlock {
+ public:
+ ResnetBlock3D(int in, int out) {
+ blocks["norm1"] = std::make_shared<GroupNorm32>(in);
+ blocks["norm2"] = std::make_shared<GroupNorm32>(out);
+ blocks["conv1"] = std::make_shared<InflatedCausalConv3d>(in, out);
+ blocks["conv2"] = std::make_shared<InflatedCausalConv3d>(out, out);
+ if (in != out)
+ blocks["conv_shortcut"] = std::make_shared<InflatedCausalConv3d>(in, out, std::tuple<int, int, int>{1, 1, 1}, std::tuple<int, int, int>{1, 1, 1}, std::tuple<int, int, int>{0, 0, 0});
+ }
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+ auto h = ggml_silu(ctx->ggml_ctx, causal_norm(ctx, x, blocks["norm1"]));
+ h = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv1"])->forward(ctx, h);
+ h = ggml_silu(ctx->ggml_ctx, causal_norm(ctx, h, blocks["norm2"]));
+ h = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv2"])->forward(ctx, h);
+ if (blocks.count("conv_shortcut"))
+ x = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv_shortcut"])->forward(ctx, x);
+ return ggml_add(ctx->ggml_ctx, h, x);
+ }
+ };
+
+ class SpatialAttention : public UnaryBlock {
+ public:
+ explicit SpatialAttention(int channels) {
+ blocks["group_norm"] = std::make_shared<GroupNorm32>(channels);
+ for (const auto& name : {"to_q", "to_k", "to_v", "to_out.0"})
+ blocks[name] = std::make_shared<Linear>(channels, channels);
+ }
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+ auto c = ctx->ggml_ctx;
+ const int64_t width = x->ne[0], height = x->ne[1], frames = x->ne[2], channels = x->ne[3];
+ auto h = causal_norm(ctx, x, blocks["group_norm"]);
+ h = ggml_reshape_3d(c, h, width * height, frames, channels);
+ h = ggml_cont(c, ggml_permute(c, h, 1, 2, 0, 3));
+ auto linear = [&](const char* name, ggml_tensor* value) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, value); };
+ h = ggml_ext_attention_ext(c, ctx->backend, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled);
+ h = linear("to_out.0", h);
+ h = ggml_cont(c, ggml_permute(c, h, 2, 0, 1, 3));
+ return ggml_add(c, x, ggml_reshape_4d(c, h, width, height, frames, channels));
+ }
+ };
+
+ class Upsample3D : public UnaryBlock {
+ int temporal_;
+
+ public:
+ Upsample3D(int channels, bool temporal)
+ : temporal_(temporal ? 2 : 1) {
+ blocks["upscale_conv"] = std::make_shared<InflatedCausalConv3d>(channels, channels * 4 * temporal_, std::tuple<int, int, int>{1, 1, 1}, std::tuple<int, int, int>{1, 1, 1}, std::tuple<int, int, int>{0, 0, 0});
+ blocks["conv"] = std::make_shared<InflatedCausalConv3d>(channels, channels);
+ }
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+ auto c = ctx->ggml_ctx;
+ const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3];
+ x = std::dynamic_pointer_cast<UnaryBlock>(blocks["upscale_conv"])->forward(ctx, x);
+ // [B, (y x z C), T, H, W] -> [B, C, T*z, H*y, W*x].
+ x = ggml_reshape_4d(c, x, w * h * t, ch, temporal_, 4);
+ x = ggml_cont(c, ggml_permute(c, x, 1, 2, 0, 3));
+ x = ggml_reshape_4d(c, x, temporal_, w * h, t, ch * 4);
+ x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3));
+ x = ggml_reshape_4d(c, x, w, h * t * temporal_ * ch, 2, 2);
+ x = ggml_cont(c, ggml_permute(c, x, 1, 2, 0, 3));
+ x = ggml_reshape_4d(c, x, w * 2, h, t * temporal_ * ch, 2);
+ x = ggml_cont(c, ggml_permute(c, x, 0, 2, 3, 1));
+ x = ggml_reshape_4d(c, x, w * 2, h * 2, t * temporal_, ch);
+ if (temporal_ == 2) {
+ auto first = ggml_ext_slice(c, x, 2, 0, 1);
+ x = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(c, x, 2, 2, t * 2), 2);
+ }
+ return std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"])->forward(ctx, x);
+ }
+ };
+
+ class VideoAutoencoderKL : public GGMLBlock {
+ SeedVR2VAEConfig config_;
+
+ public:
+ explicit VideoAutoencoderKL(const SeedVR2VAEConfig& config)
+ : config_(config) {
+ for (bool decode : {false, true}) {
+ std::string root = decode ? "decoder." : "encoder.";
+ auto channels = config.channels;
+ if (decode)
+ std::reverse(channels.begin(), channels.end());
+ int previous = channels.front();
+ blocks[root + "conv_in"] = std::make_shared<InflatedCausalConv3d>(decode ? 16 : 3, previous);
+ for (int i = 0; i < 4; ++i) {
+ auto stage = root + (decode ? "up_blocks." : "down_blocks.") + std::to_string(i) + ".";
+ for (int j = 0; j < config.layers + int(decode); ++j) {
+ blocks[stage + "resnets." + std::to_string(j)] = std::make_shared<ResnetBlock3D>(previous, channels[i]);
+ previous = channels[i];
+ }
+ if (i < 3) {
+ if (decode)
+ blocks[stage + "upsamplers.0"] = std::make_shared<Upsample3D>(previous, i < 2);
+ else
+ blocks[stage + "downsamplers.0.conv"] = std::make_shared<InflatedCausalConv3d>(previous, previous,
+ std::tuple<int, int, int>{i ? 3 : 1, 3, 3}, std::tuple<int, int, int>{i ? 2 : 1, 2, 2}, std::tuple<int, int, int>{i ? 1 : 0, 0, 0});
+ }
+ }
+ blocks[root + "mid_block.resnets.0"] = std::make_shared<ResnetBlock3D>(512, 512);
+ blocks[root + "mid_block.attentions.0"] = std::make_shared<SpatialAttention>(512);
+ blocks[root + "mid_block.resnets.1"] = std::make_shared<ResnetBlock3D>(512, 512);
+ blocks[root + "conv_norm_out"] = std::make_shared<GroupNorm32>(previous);
+ blocks[root + "conv_out"] = std::make_shared<InflatedCausalConv3d>(previous, decode ? 3 : 32);
+ }
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, bool decode) {
+ auto c = ctx->ggml_ctx;
+ std::string root = decode ? "decoder." : "encoder.";
+ auto apply = [&](const std::string& name) {
+ x = std::dynamic_pointer_cast<UnaryBlock>(blocks[root + name])->forward(ctx, x);
+ sd::ggml_graph_cut::mark_graph_cut(x, "seedvr2.vae." + root + name, "hidden");
+ };
+ auto mid = [&]() {
+ apply("mid_block.resnets.0");
+ apply("mid_block.attentions.0");
+ apply("mid_block.resnets.1");
+ };
+ apply("conv_in");
+ if (decode)
+ mid();
+ for (int i = 0; i < 4; ++i) {
+ auto stage = std::string(decode ? "up_blocks." : "down_blocks.") + std::to_string(i) + ".";
+ for (int j = 0; j < config_.layers + int(decode); ++j)
+ apply(stage + "resnets." + std::to_string(j));
+ if (i < 3) {
+ if (!decode)
+ x = ggml_pad(c, x, 1, 1, 0, 0);
+ apply(stage + (decode ? "upsamplers.0" : "downsamplers.0.conv"));
+ }
+ }
+ if (!decode)
+ mid();
+ x = ggml_silu(c, causal_norm(ctx, x, blocks[root + "conv_norm_out"]));
+ apply("conv_out");
+ return decode ? x : ggml_ext_slice(c, x, 3, 0, 16);
+ }
+ };
+
+ struct SeedVR2VAERunner : GGMLRunner {
+ SeedVR2VAEConfig config;
+ VideoAutoencoderKL model;
+
+ SeedVR2VAERunner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr<DeviceResidencyManager> manager)
+ : GGMLRunner(backend, manager), config(SeedVR2VAEConfig::detect_from_weights(weights, "vae")), model(config) {
+ model.init(params_ctx, weights, "vae");
+ }
+ std::string get_desc() override { return "SeedVR2 VAE"; }
+ void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) { model.get_param_tensors(tensors, "vae"); }
+ sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input, bool decode) {
+ auto graph = [&]() {
+ auto gf = new_graph_custom(MAX_GRAPH_SIZE * 4);
+ auto ctx = get_context(gf);
+ auto out = model.forward(&ctx, make_input(input), decode);
+ ggml_build_forward_expand(gf, out);
+ return gf;
+ };
+ auto output = GGMLRunner::compute(graph, n_threads);
+ return restore_trailing_singleton_dims(output, 4);
+ }
+ };
+}
+
+#endif // __SD_MODEL_VAE_SEEDVR2_VAE_HPP__
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index 78a5e10..738723f 100644
--- a/src/name_conversion.cpp
+++ b/src/name_conversion.cpp
@@ -1161,31 +1161,31 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
{"mid.attn_1.", "mid_block.attentions.0."},
};
- static std::vector<std::pair<std::string, std::string>> vae_conversion_map_layer;
- if (vae_conversion_map_layer.empty()) {
+ static const auto vae_conversion_map_layer = [] {
+ std::vector<std::pair<std::string, std::string>> mappings;
for (int i = 0; i < 4; ++i) {
// --- encoder down blocks ---
for (int j = 0; j < 2; ++j) {
std::string hf_down_prefix = "encoder.down_blocks." + std::to_string(i) + ".resnets." + std::to_string(j) + ".";
std::string sd_down_prefix = "encoder.down." + std::to_string(i) + ".block." + std::to_string(j) + ".";
- vae_conversion_map_layer.emplace_back(sd_down_prefix, hf_down_prefix);
+ mappings.emplace_back(sd_down_prefix, hf_down_prefix);
}
if (i < 3) {
std::string hf_downsample_prefix = "down_blocks." + std::to_string(i) + ".downsamplers.0.";
std::string sd_downsample_prefix = "down." + std::to_string(i) + ".downsample.";
- vae_conversion_map_layer.emplace_back(sd_downsample_prefix, hf_downsample_prefix);
+ mappings.emplace_back(sd_downsample_prefix, hf_downsample_prefix);
std::string hf_upsample_prefix = "up_blocks." + std::to_string(i) + ".upsamplers.0.";
std::string sd_upsample_prefix = "up." + std::to_string(3 - i) + ".upsample.";
- vae_conversion_map_layer.emplace_back(sd_upsample_prefix, hf_upsample_prefix);
+ mappings.emplace_back(sd_upsample_prefix, hf_upsample_prefix);
}
// --- decoder up blocks (reverse) ---
for (int j = 0; j < 3; ++j) {
std::string hf_up_prefix = "decoder.up_blocks." + std::to_string(i) + ".resnets." + std::to_string(j) + ".";
std::string sd_up_prefix = "decoder.up." + std::to_string(3 - i) + ".block." + std::to_string(j) + ".";
- vae_conversion_map_layer.emplace_back(sd_up_prefix, hf_up_prefix);
+ mappings.emplace_back(sd_up_prefix, hf_up_prefix);
}
}
@@ -1193,9 +1193,10 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
for (int i = 0; i < 2; ++i) {
std::string hf_mid_res_prefix = "mid_block.resnets." + std::to_string(i) + ".";
std::string sd_mid_res_prefix = "mid.block_" + std::to_string(i + 1) + ".";
- vae_conversion_map_layer.emplace_back(sd_mid_res_prefix, hf_mid_res_prefix);
+ mappings.emplace_back(sd_mid_res_prefix, hf_mid_res_prefix);
}
- }
+ return mappings;
+ }();
static const std::vector<std::pair<std::string, std::string>> vae_conversion_map_attn = {
{"norm.", "group_norm."},
diff --git a/src/pipeline/relay_vsr.cpp b/src/pipeline/relay_vsr.cpp
index bc5cad7..c8b0d72 100644
--- a/src/pipeline/relay_vsr.cpp
+++ b/src/pipeline/relay_vsr.cpp
@@ -84,7 +84,7 @@ namespace sd {
sources.emplace(source->second, tensor);
ggml_set_name(tensor, source->second.c_str());
}
- runner->set_flash_attention_enabled(true);
+ runtime.configure_runner(*runner, module, true, module == SDBackendModule::DIFFUSION);
return runtime.manager->register_param_tensors(component, std::move(sources),
runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
compute, params);
diff --git a/src/pipeline/seedvr2.cpp b/src/pipeline/seedvr2.cpp
new file mode 100644
index 0000000..5fe2eaa
--- /dev/null
+++ b/src/pipeline/seedvr2.cpp
@@ -0,0 +1,167 @@
+#include "pipeline/video_upscale.h"
+
+#include <climits>
+
+#include "core/rng_philox.hpp"
+#include "generation_info.h"
+#include "model/diffusion/seedvr2.hpp"
+#include "model/vae/seedvr2_vae.hpp"
+
+namespace sd {
+ namespace pipeline {
+
+ class SeedVR2Upscaler : public VideoUpscaleModel {
+ std::unique_ptr<SeedVR2::SeedVR2Runner> dit_;
+ std::unique_ptr<SeedVR2::SeedVR2VAERunner> vae_;
+ int scale_ = 0;
+ int n_threads_ = 1;
+
+ public:
+ bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& weights, int upscale_factor) override {
+ scale_ = upscale_factor;
+ n_threads_ = runtime.n_threads;
+ dit_ = std::make_unique<SeedVR2::SeedVR2Runner>(runtime.backends.runtime_backend(SDBackendModule::DIFFUSION), weights, runtime.manager);
+ vae_ = std::make_unique<SeedVR2::SeedVR2VAERunner>(runtime.backends.runtime_backend(SDBackendModule::VAE), weights, runtime.manager);
+ std::set<std::string> used;
+ auto register_model = [&](auto& runner, ModelComponent component, SDBackendModule module) {
+ std::map<std::string, ggml_tensor*> tensors;
+ runner->get_param_tensors(tensors);
+ for (const auto& [name, tensor] : tensors) {
+ if (!weights.count(name))
+ throw std::runtime_error("Missing SeedVR2 tensor: " + name);
+ used.insert(name);
+ }
+ runtime.configure_runner(*runner, module, module == SDBackendModule::VAE, module == SDBackendModule::VAE);
+ return runtime.manager->register_param_tensors(component, std::move(tensors),
+ runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
+ runtime.backends.runtime_backend(module), runtime.backends.params_backend(module));
+ };
+ if (!register_model(dit_, ModelComponent::Diffusion, SDBackendModule::DIFFUSION) ||
+ !register_model(vae_, ModelComponent::VAE, SDBackendModule::VAE) || !runtime.manager->validate_registered_tensors())
+ return false;
+ for (const auto& [name, storage] : weights) {
+ if (!used.count(name) && name != "diffusion.negative_conditioning" && !ends_with(name, ".comfy_quant"))
+ throw std::runtime_error("Unused SeedVR2 tensor: " + name);
+ }
+ LOG_INFO("SeedVR2 validated %zu parameter tensors; negative conditioning is unused at guidance 1", used.size());
+ return true;
+ }
+
+ void end() override {
+ if (dit_)
+ dit_->runner_end();
+ if (vae_)
+ vae_->runner_end();
+ }
+
+ bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) override {
+ const int chunk = params.temporal_chunk_size;
+ if (chunk < 1 || chunk > INT_MAX - 3 || (chunk - 1) % 4) {
+ LOG_ERROR("SeedVR2 temporal chunk size must be 1 or 4n+1 input frames");
+ return false;
+ }
+ struct Reset {
+ SeedVR2Upscaler* model;
+ ~Reset() { model->end(); }
+ } reset{this};
+ end();
+ const int width = int(params.frames[0].width) * scale_, height = int(params.frames[0].height) * scale_;
+ if (width > INT_MAX - 15 || height > INT_MAX - 15)
+ return false;
+ const int pw = (width + 15) / 16 * 16, ph = (height + 15) / 16 * 16;
+ PhiloxRNG rng;
+ rng.manual_seed(uint64_t(params.seed));
+ for (int start = 0; start < params.frame_count;) {
+ const int count = std::min(chunk, params.frame_count - start);
+ const int padded_count = count == 1 ? 1 : ((count - 1 + 3) / 4) * 4 + 1;
+ LOG_INFO("SeedVR2 chunk %d..%d/%d", start + 1, start + count, params.frame_count);
+ const int64_t begin = ggml_time_us();
+ std::vector<sd::Tensor<float>> images;
+ for (int i = 0; i < count; ++i) {
+ auto image = sd_image_to_tensor(params.frames[start + i], -1, -1, false);
+ image /= 255.f;
+ image = sd::ops::interpolate(image, {width, height, 3, 1}, sd::ops::InterpolateMode::Bicubic);
+ image = sd::ops::clamp(image, 0.f, 1.f);
+ if (pw != width)
+ image = sd::ops::concat(image, sd::zeros<float>({pw - width, height, 3, 1}), 0);
+ if (ph != height)
+ image = sd::ops::concat(image, sd::zeros<float>({pw, ph - height, 3, 1}), 1);
+ image *= 2.f;
+ image -= 1.f;
+ images.push_back(image);
+ }
+ auto pixels = images.front();
+ for (int i = 1; i < padded_count; ++i)
+ pixels = sd::ops::concat(pixels, images[std::min(i, count - 1)], 3);
+ pixels.permute_({0, 1, 3, 2});
+ const int64_t prepared = ggml_time_us();
+ auto latent = vae_->compute(n_threads_, pixels, false);
+ if (latent.empty())
+ return false;
+ latent *= 0.9152f;
+ const int64_t encoded = ggml_time_us();
+ const int lw = pw / 8, lh = ph / 8, lt = (padded_count - 1) / 4 + 1;
+ sd::Tensor<float> noise({lw, lh, lt, 16});
+ auto values = rng.randn(uint32_t(noise.numel()));
+ std::copy(values.begin(), values.end(), noise.data());
+ auto mask = sd::zeros<float>({lw, lh, lt, 1});
+ mask += 1.f;
+ auto input = sd::ops::concat(sd::ops::concat(noise, latent, 3), mask, 3);
+ input.reshape_({2, lw / 2, 2, lh / 2, lt, 33});
+ input.permute_({5, 0, 2, 1, 3, 4});
+ input.reshape_({132, (lw / 2) * (lh / 2) * lt});
+ auto prediction = dit_->compute(n_threads_, input, lw / 2, lh / 2, lt);
+ if (prediction.empty())
+ return false;
+ prediction.reshape_({16, 2, 2, lw / 2, lh / 2, lt});
+ prediction.permute_({1, 3, 2, 4, 5, 0});
+ prediction.reshape_({lw, lh, lt, 16});
+ latent = (noise - prediction) / 0.9152f;
+ const int64_t denoised = ggml_time_us();
+ auto decoded = vae_->compute(n_threads_, latent, true);
+ if (decoded.empty())
+ return false;
+ const int64_t decoded_at = ggml_time_us();
+ if (!std::all_of(decoded.values().begin(), decoded.values().end(), [](float x) { return std::isfinite(x); }))
+ throw std::runtime_error("SeedVR2 produced non-finite pixels");
+ decoded = sd::ops::clamp((decoded + 1.f) * 127.5f, 0.f, 255.f);
+ for (int i = 0; i < count; ++i) {
+ auto output = sd::ops::slice(sd::ops::slice(sd::ops::slice(decoded, 2, i, i + 1), 0, 0, width), 1, 0, height);
+ output.permute_({3, 0, 1, 2});
+ auto& image = result->frames[start + i];
+ image.width = width;
+ image.height = height;
+ image.channel = 3;
+ image.data = static_cast<uint8_t*>(malloc(size_t(width) * height * 3));
+ if (!image.data)
+ throw std::bad_alloc();
+ std::transform(output.values().begin(), output.values().end(), image.data, [](float x) { return uint8_t(x); });
+ }
+ if (auto progress = sd_get_progress_callback())
+ progress(start + count, params.frame_count, 0.f, sd_get_progress_callback_data());
+ LOG_INFO("SeedVR2 seconds: preprocess %.3f, VAE encode %.3f, DiT and latent transforms %.3f, VAE decode %.3f, postprocess %.3f",
+ (prepared - begin) / 1e6, (encoded - prepared) / 1e6, (denoised - encoded) / 1e6, (decoded_at - denoised) / 1e6, (ggml_time_us() - decoded_at) / 1e6);
+ start += count;
+ }
+ auto info = sd::upscale_generation_info(scale_);
+ info["kind"] = "video";
+ info["upscale"]["model"] = "seedvr2";
+ info["upscale"]["temporal_chunk_size"] = chunk;
+ info["seed"] = params.seed;
+ info["fps"] = params.fps;
+ info["video_frames"] = params.frame_count;
+ result->generation_info = sd::copy_generation_info(info);
+ return result->generation_info != nullptr;
+ }
+ };
+
+ bool is_seedvr2_model(const String2TensorStorage& weights) {
+ return weights.count("diffusion.vid_in.proj.weight") && weights.count("diffusion.blocks.0.attn.rope.rope.freqs") &&
+ weights.count("diffusion.positive_conditioning");
+ }
+
+ std::unique_ptr<VideoUpscaleModel> create_seedvr2_upscaler() {
+ return std::make_unique<SeedVR2Upscaler>();
+ }
+ }
+}
diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp
index 5b34a92..f9b631e 100644
--- a/src/pipeline/video_upscale.cpp
+++ b/src/pipeline/video_upscale.cpp
@@ -2,11 +2,62 @@
#include <climits>
+#include "core/ggml_runner.h"
+#include "core/probe/conv_direct.h"
+#include "core/probe/flash_attn.h"
#include "core/util.h"
namespace sd {
namespace pipeline {
+ bool VideoUpscaleRuntime::flash_attention_for(SDBackendModule module) {
+ const auto mode = flash_attn_assignment.get(module);
+ if (mode == "on" || mode == "off")
+ return mode == "on";
+ const auto devices = backends.runtime_backends(module);
+ bool enabled = !devices.empty(), any_enabled = false;
+ for (auto backend : devices) {
+ auto it = flash_attn_by_backend.find(backend);
+ if (it == flash_attn_by_backend.end())
+ it = flash_attn_by_backend.emplace(backend, sd_probe_flash_attention(backend, n_threads)).first;
+ enabled = enabled && it->second;
+ any_enabled = any_enabled || it->second;
+ }
+ if (any_enabled && !enabled)
+ LOG_INFO("flash attention auto: %s -> off (compute devices disagree)", sd_backend_module_name(module));
+ return enabled;
+ }
+
+ bool VideoUpscaleRuntime::conv_direct_for(SDBackendModule module, bool conv3d) {
+ const auto mode = conv_direct_assignment.get(module);
+ if (mode == "on" || mode == "off")
+ return mode == "on";
+ auto& cache = conv3d ? conv3d_direct_by_backend : conv2d_direct_by_backend;
+ const auto devices = backends.runtime_backends(module);
+ bool enabled = !devices.empty(), any_enabled = false;
+ for (auto backend : devices) {
+ auto it = cache.find(backend);
+ if (it == cache.end())
+ it = cache.emplace(backend, sd_probe_conv_direct(backend, conv3d)).first;
+ enabled = enabled && it->second;
+ any_enabled = any_enabled || it->second;
+ }
+ if (any_enabled && !enabled)
+ LOG_INFO("conv%dd direct auto: %s -> off (compute devices disagree)", conv3d ? 3 : 2, sd_backend_module_name(module));
+ return enabled;
+ }
+
+ void VideoUpscaleRuntime::configure_runner(GGMLRunner& runner, SDBackendModule module, bool conv2d, bool conv3d) {
+ const bool flash = flash_attention_for(module);
+ const bool direct2d = conv2d && conv_direct_for(module, false);
+ const bool direct3d = conv3d && conv_direct_for(module, true);
+ runner.set_flash_attention_enabled(flash);
+ runner.set_conv2d_direct_enabled(direct2d);
+ runner.set_conv3d_direct_enabled(direct3d);
+ LOG_INFO("video upscale %s: flash attention %s, conv2d direct %s, conv3d direct %s",
+ sd_backend_module_name(module), flash ? "on" : "off", direct2d ? "on" : "off", direct3d ? "on" : "off");
+ }
+
VideoUpscaler::~VideoUpscaler() {
if (model_)
model_->end();
@@ -26,6 +77,11 @@ namespace sd {
executor_ = std::make_unique<sd::ParallelExecutor>(runtime_.n_threads);
sd::ParallelScope scope(executor_.get());
std::string error;
+ if (!sd_parse_flash_attention_assignment(params.flash_attn_mode ? params.flash_attn_mode : "", &runtime_.flash_attn_assignment, &error) ||
+ !sd_parse_conv_direct_assignment(params.conv_direct_mode ? params.conv_direct_mode : "", &runtime_.conv_direct_assignment, &error)) {
+ LOG_ERROR("video upscale operator configuration: %s", error.c_str());
+ return false;
+ }
if (!runtime_.backends.init(params.backend, params.params_backend, nullptr, &error)) {
LOG_ERROR("video upscaler backend: %s", error.c_str());
return false;
@@ -41,6 +97,8 @@ namespace sd {
const auto& weights = loader.get_tensor_storage_map();
if (is_relay_vsr_model(weights))
model_ = create_relay_vsr_upscaler();
+ else if (is_seedvr2_model(weights))
+ model_ = create_seedvr2_upscaler();
else {
LOG_ERROR("unsupported video upscale model architecture");
return false;
@@ -104,10 +162,11 @@ void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params) {
void sd_video_upscale_params_init(sd_video_upscale_params_t* params) {
if (!params)
return;
- *params = {};
- params->fps = 24;
- params->seed = 42;
- params->window_size = 16;
+ *params = {};
+ params->fps = 24;
+ params->seed = 42;
+ params->window_size = 16;
+ params->temporal_chunk_size = 5;
}
video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params) {
diff --git a/src/pipeline/video_upscale.h b/src/pipeline/video_upscale.h
index 62d0af1..ec5815c 100644
--- a/src/pipeline/video_upscale.h
+++ b/src/pipeline/video_upscale.h
@@ -1,6 +1,7 @@
#ifndef __SD_PIPELINE_VIDEO_UPSCALE_H__
#define __SD_PIPELINE_VIDEO_UPSCALE_H__
+#include <map>
#include <memory>
#include "core/ggml_extend_backend.h"
@@ -8,6 +9,8 @@
#include "model_manager.h"
#include "stable-diffusion.h"
+struct GGMLRunner;
+
namespace sd {
namespace pipeline {
@@ -15,6 +18,15 @@ namespace sd {
SDBackendManager backends;
std::shared_ptr<ModelManager> manager;
int n_threads = 1;
+ SDBackendAssignment flash_attn_assignment;
+ SDBackendAssignment conv_direct_assignment;
+ std::map<ggml_backend_t, bool> flash_attn_by_backend;
+ std::map<ggml_backend_t, bool> conv2d_direct_by_backend;
+ std::map<ggml_backend_t, bool> conv3d_direct_by_backend;
+
+ bool flash_attention_for(SDBackendModule module);
+ bool conv_direct_for(SDBackendModule module, bool conv3d);
+ void configure_runner(GGMLRunner& runner, SDBackendModule module, bool conv2d, bool conv3d);
};
class VideoUpscaleModel {
@@ -27,6 +39,8 @@ namespace sd {
bool is_relay_vsr_model(const String2TensorStorage& weights);
std::unique_ptr<VideoUpscaleModel> create_relay_vsr_upscaler();
+ bool is_seedvr2_model(const String2TensorStorage& weights);
+ std::unique_ptr<VideoUpscaleModel> create_seedvr2_upscaler();
class VideoUpscaler {
VideoUpscaleRuntime runtime_;