diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ca64ad19..15045dcc 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 2442865b..3a65ee1b 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 1352f41f..e35f5fdb 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. +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. diff --git a/docs/seedvr2.md b/docs/seedvr2.md new file mode 100644 index 00000000..24f55032 --- /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 bcf73595..2bff0507 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 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 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 a3496f4f..70107a39 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 a1d82e78..53c6b705 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 00000000..87e542e9 --- /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 order; + sd::Tensor reverse; + sd::Tensor positions; + sd::Tensor text_positions; + std::vector 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{"out"} : std::vector{"attn", "mlp"}) { + for (const auto& mode : output_ ? std::vector{"shift", "scale"} : std::vector{"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(config.dim, config.hidden, !swiglu_); + blocks["proj_out"] = std::make_shared(config.hidden, config.dim, !swiglu_); + if (swiglu_) + blocks["proj_in_gate"] = std::make_shared(config.dim, config.hidden, false); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto up = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, x); + auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["proj_in_gate"])->forward(ctx, x), up) + : ggml_ext_gelu(ctx->ggml_ctx, up); + return std::dynamic_pointer_cast(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{"all"} : std::vector{"vid", "txt"}) { + blocks["attn.proj_qkv." + branch] = std::make_shared(config.dim, 3 * config.dim, false); + blocks["attn.proj_out." + branch] = std::make_shared(config.dim, config.dim); + blocks["attn.norm_q." + branch] = std::make_shared(128, 1e-5f); + blocks["attn.norm_k." + branch] = std::make_shared(128, 1e-5f); + blocks["ada." + branch] = std::make_shared(config.dim); + blocks["mlp." + branch] = std::make_shared(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 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(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(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(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 video_windows; + ggml_tensor* tout = nullptr; + int start = 0; + for (int length : plan.lengths) { + std::vector 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(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(132, config.dim); + blocks["txt_in"] = std::make_shared(config.text_dim, config.dim); + blocks["emb_in.proj_in"] = std::make_shared(256, config.dim); + blocks["emb_in.proj_hid"] = std::make_shared(config.dim, config.dim); + blocks["emb_in.proj_out"] = std::make_shared(config.dim, 6 * config.dim); + for (int i = 0; i < config.layers; ++i) + blocks["blocks." + std::to_string(i)] = std::make_shared(config, i); + blocks["vid_out.proj"] = std::make_shared(config.dim, 64); + if (config.swiglu) { + blocks["vid_out_norm"] = std::make_shared(config.dim, 1e-5f); + blocks["vid_out_ada"] = std::make_shared(config.dim, true); + } + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* timestep, const std::vector& plans, VSR::GraphInputs& inputs) { + auto linear = [&](const std::string& name, ggml_tensor* x) { return std::dynamic_pointer_cast(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(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(blocks["vid_out_norm"])->forward(ctx, vid); + vid = std::dynamic_pointer_cast(blocks["vid_out_ada"])->forward(ctx, vid, emb, "out"); + } + return linear("vid_out.proj", vid); + } + }; + + struct SeedVR2Runner : GGMLRunner { + SeedVR2Config config; + NaDiT model; + sd::Tensor timestep; + + SeedVR2Runner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr 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& tensors) { model.get_param_tensors(tensors, "diffusion"); } + + sd::Tensor compute(int n_threads, const sd::Tensor& input, int width, int height, int frames) { + std::vector 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(); + } + }; +} + +#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 00000000..5f38aebc --- /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 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({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 kernel = {3, 3, 3}, std::tuple stride = {1, 1, 1}, std::tuple 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(1, max_patch_bytes / row_bytes); + struct RestorePadding { + std::tuple& target; + std::tuple 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(0, top), std::min(x->ne[1], bottom), false); + strip = ggml_ext_pad_ext(c, ctx->backend, strip, 0, 0, int(std::max(0, -top)), + int(std::max(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& 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(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(in); + blocks["norm2"] = std::make_shared(out); + blocks["conv1"] = std::make_shared(in, out); + blocks["conv2"] = std::make_shared(out, out); + if (in != out) + blocks["conv_shortcut"] = std::make_shared(in, out, std::tuple{1, 1, 1}, std::tuple{1, 1, 1}, std::tuple{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(blocks["conv1"])->forward(ctx, h); + h = ggml_silu(ctx->ggml_ctx, causal_norm(ctx, h, blocks["norm2"])); + h = std::dynamic_pointer_cast(blocks["conv2"])->forward(ctx, h); + if (blocks.count("conv_shortcut")) + x = std::dynamic_pointer_cast(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(channels); + for (const auto& name : {"to_q", "to_k", "to_v", "to_out.0"}) + blocks[name] = std::make_shared(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(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(channels, channels * 4 * temporal_, std::tuple{1, 1, 1}, std::tuple{1, 1, 1}, std::tuple{0, 0, 0}); + blocks["conv"] = std::make_shared(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(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(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(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(previous, channels[i]); + previous = channels[i]; + } + if (i < 3) { + if (decode) + blocks[stage + "upsamplers.0"] = std::make_shared(previous, i < 2); + else + blocks[stage + "downsamplers.0.conv"] = std::make_shared(previous, previous, + std::tuple{i ? 3 : 1, 3, 3}, std::tuple{i ? 2 : 1, 2, 2}, std::tuple{i ? 1 : 0, 0, 0}); + } + } + blocks[root + "mid_block.resnets.0"] = std::make_shared(512, 512); + blocks[root + "mid_block.attentions.0"] = std::make_shared(512); + blocks[root + "mid_block.resnets.1"] = std::make_shared(512, 512); + blocks[root + "conv_norm_out"] = std::make_shared(previous); + blocks[root + "conv_out"] = std::make_shared(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(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 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& tensors) { model.get_param_tensors(tensors, "vae"); } + sd::Tensor compute(int n_threads, const sd::Tensor& 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 78a5e106..738723fa 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> vae_conversion_map_layer; - if (vae_conversion_map_layer.empty()) { + static const auto vae_conversion_map_layer = [] { + std::vector> 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> vae_conversion_map_attn = { {"norm.", "group_norm."}, diff --git a/src/pipeline/relay_vsr.cpp b/src/pipeline/relay_vsr.cpp index bc5cad71..c8b0d72f 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 00000000..5fe2eaa4 --- /dev/null +++ b/src/pipeline/seedvr2.cpp @@ -0,0 +1,167 @@ +#include "pipeline/video_upscale.h" + +#include + +#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 dit_; + std::unique_ptr 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(runtime.backends.runtime_backend(SDBackendModule::DIFFUSION), weights, runtime.manager); + vae_ = std::make_unique(runtime.backends.runtime_backend(SDBackendModule::VAE), weights, runtime.manager); + std::set used; + auto register_model = [&](auto& runner, ModelComponent component, SDBackendModule module) { + std::map 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> 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({pw - width, height, 3, 1}), 0); + if (ph != height) + image = sd::ops::concat(image, sd::zeros({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 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({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(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 create_seedvr2_upscaler() { + return std::make_unique(); + } + } +} diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp index 5b34a928..f9b631e6 100644 --- a/src/pipeline/video_upscale.cpp +++ b/src/pipeline/video_upscale.cpp @@ -2,11 +2,62 @@ #include +#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(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 62d0af10..ec5815c4 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 #include #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 manager; int n_threads = 1; + SDBackendAssignment flash_attn_assignment; + SDBackendAssignment conv_direct_assignment; + std::map flash_attn_by_backend; + std::map conv2d_direct_by_backend; + std::map 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 create_relay_vsr_upscaler(); + bool is_seedvr2_model(const String2TensorStorage& weights); + std::unique_ptr create_seedvr2_upscaler(); class VideoUpscaler { VideoUpscaleRuntime runtime_;