feat: add RelayVSR video upscaling with CUDA window attention (#2285)

examples/cli/main.cpp

@@ -3,6 +3,7 @@
 #include <time.h>
 #include <algorithm>
 #include <cctype>
+#include <climits>
 #include <filesystem>
 #include <functional>
 #include <iostream>
@@ -33,10 +34,13 @@ const char* previews_str[] = {
 std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
 
 struct SDCliParams {
-    SDMode mode             = IMG_GEN;
-    std::string output_path = "output.png";
-    int output_begin_idx    = -1;
-    int compression_quality = 90;
+    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;
     std::string image_path;
     std::string metadata_format = "text";
 
@@ -62,6 +66,7 @@ struct SDCliParams {
         ArgOptions options;
 
         options.string_options = {
+            {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 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",
@@ -90,6 +95,8 @@ struct SDCliParams {
         };
 
         options.int_options = {
+            {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor},
+            {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_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",
@@ -190,7 +197,7 @@ struct SDCliParams {
         options.manual_options = {
             {"-M",
              "--mode",
-             "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
+             "run mode, one of [img_gen, adetailer, vid_gen, vid_upscale, upscale, convert, metadata], default: img_gen",
              on_mode_arg},
             {"",
              "--preview",
@@ -220,6 +227,10 @@ 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");
+            return false;
+        }
         if (mode != METADATA) {
             if (output_path.length() == 0) {
                 LOG_ERROR("error: the following arguments are required: output_path");
@@ -529,7 +540,7 @@ bool save_results(const SDCliParams& cli_params,
         return sucessful_reults != 0;
     }
 
-    if (cli_params.mode == VID_GEN && num_results > 1) {
+    if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) {
         if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
         std::string params = gen_params.embed_image_metadata
@@ -664,6 +675,50 @@ int main(int argc, const char* argv[]) {
         return 0;
     }
 
+    if (cli_params.mode == VID_UPSCALE) {
+        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()) {
+            LOG_ERROR("vid_upscale input contains no readable frames");
+            return 1;
+        }
+        if (owned_frames.size() > size_t(INT_MAX))
+            return 1;
+        std::vector<sd_image_t> frames;
+        for (const auto& frame : owned_frames)
+            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();
+        std::unique_ptr<video_upscaler_ctx_t, decltype(&free_video_upscaler_ctx)> upscaler(
+            new_video_upscaler_ctx(&params),
+            free_video_upscaler_ctx);
+        if (!upscaler)
+            return 1;
+        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;
+        if (!upscale_video(upscaler.get(), &upscale_params, &output))
+            return 1;
+        std::string video_info = output->generation_info;
+        SDImageVec results;
+        results.adopt_video(output);
+        return save_results(cli_params, ctx_params, gen_params, results, int(results.size()), nullptr, video_info) ? 0 : 1;
+    }
+
     if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
         LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
         return 1;