DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-09 10:26:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add SeedVR2 image and video upscaling (#2300)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-09 01:13:57 +08:00
1 parent abc6718455
commit cd5594e9f9
14 files changed
+1035 -89
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CONTRIBUTING.md
+12
View File
@@ -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.
README.md
+3 -1
View File
@@ -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
docs/relay_vsr.md
+31 -51
View File
@@ -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.
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`.
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.
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.
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.
## Settings and limits
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.
- `--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.
docs/seedvr2.md
+69
View File
@@ -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).
examples/cli/main.cpp
+49 -23
View File
@@ -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(&params);
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(&params),
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;
examples/common/common.cpp
+1 -1
View File
@@ -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},
{"",
include/stable-diffusion.h
+5
View File
@@ -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);
src/model/diffusion/seedvr2.hpp
+325
View File
@@ -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__
src/model/vae/seedvr2_vae.hpp
+286
View File
@@ -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__
src/name_conversion.cpp
+9 -8
View File
@@ -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."},
src/pipeline/relay_vsr.cpp
+1 -1
View File
@@ -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);
src/pipeline/seedvr2.cpp
+167
View File
@@ -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>();
}
}
}
src/pipeline/video_upscale.cpp
+63 -4
View File
@@ -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) {
src/pipeline/video_upscale.h
+14
View File
@@ -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_;