feat(short_drama): turn a novel into vertical drama episodes on TRTMC

ASSET_LICENSES.md

@@ -218,4 +218,17 @@ with the rest of the project:
 - `families/openfold3/tests/data/openfold3_features.npz`
 - `families/openfold3/tests/data/openfold3_structure.json`
 
+## Short-drama sample novel
+
+`apps/short_drama/samples/ah_q_zheng_zhuan.txt` is the complete text of
+《阿Q正传》 (The True Story of Ah Q) by Lu Xun (1881-1936), first serialized in
+1921-1922. The author died in 1936, so the work is in the public domain in
+mainland China, Taiwan and countries with copyright terms of life plus 50 to 80
+years, and in the United States. The text was taken from the Chinese Wikisource
+edition based on the 1938 《鲁迅全集》, which Wikisource marks public domain. The
+file keeps the nine chapters and the closing date line; the corner brackets
+(『』「」) were changed to curly quotation marks and nothing else was edited.
+
+- Source: https://zh.wikisource.org/wiki/阿Q正傳 (simplified-Chinese view)
+
 <!-- Collaborative review anchor: batch 2. -->

CMakeLists.txt

@@ -478,6 +478,24 @@ if(TRTMC_BUILD_TESTS)
   target_compile_options(test_bundle_format_v1 PRIVATE -Wall -Wextra -Wpedantic)
   add_test(NAME bundle_format_v1 COMMAND test_bundle_format_v1)
 
+  add_executable(test_bundle_stream_reader core/runtime/tests/test_bundle_format_v1.cpp)
+  target_compile_definitions(test_bundle_stream_reader PRIVATE TRTMC_TEST_BUNDLE_STREAM_READER)
+  target_include_directories(test_bundle_stream_reader PRIVATE ${PROJECT_SOURCE_DIR}/core)
+  target_link_libraries(test_bundle_stream_reader PRIVATE trtmc_core)
+  target_compile_options(test_bundle_stream_reader PRIVATE -Wall -Wextra -Wpedantic)
+  add_test(NAME bundle_stream_reader COMMAND test_bundle_stream_reader)
+
+  # This executable supplies delayed CPU-only CUDA transfer stubs. It tests
+  # callback buffer ownership without allocating or running on a GPU.
+  add_executable(test_trt_plan_stream_reader core/runtime/tests/test_trt_module_dynamic_input.cpp)
+  target_compile_definitions(test_trt_plan_stream_reader PRIVATE TRTMC_TEST_PLAN_STREAM_READER)
+  target_include_directories(test_trt_plan_stream_reader PRIVATE
+    ${PROJECT_SOURCE_DIR}/core ${PROJECT_SOURCE_DIR}/core/runtime/include
+    ${TRTMC_TRT_INCLUDE_DIR} ${CUDAToolkit_INCLUDE_DIRS})
+  target_link_libraries(test_trt_plan_stream_reader PRIVATE Threads::Threads)
+  target_compile_options(test_trt_plan_stream_reader PRIVATE -Wall -Wextra -Wpedantic)
+  add_test(NAME trt_plan_stream_reader_cpu COMMAND test_trt_plan_stream_reader)
+
   add_executable(test_task_api core/runtime/tests/test_task_api.cpp)
   target_include_directories(test_task_api PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include)
   target_compile_options(test_task_api PRIVATE -Wall -Wextra -Wpedantic)
@@ -677,6 +695,37 @@ if(TRTMC_BUILD_TESTS)
   add_dependencies(test_cli trtmc_test_default_runtime_root)
   add_test(NAME family_loader COMMAND test_family_loader "${_trtmc_test_runtime_root}")
 
+  set(_trtmc_stream_test_root "${CMAKE_BINARY_DIR}/tests/runtime-stream")
+  add_library(trtmc_test_backend_stream SHARED core/runtime/tests/fake_backend.cpp)
+  target_include_directories(trtmc_test_backend_stream PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include)
+  target_link_libraries(trtmc_test_backend_stream PRIVATE CUDA::cudart)
+  target_compile_definitions(trtmc_test_backend_stream PRIVATE
+    TRTMC_FAKE_BACKEND_NAME="trt_rtx" TRTMC_TEST_STREAM_BACKEND)
+  set_target_properties(trtmc_test_backend_stream PROPERTIES
+    OUTPUT_NAME trtmc_backend_trt_rtx LIBRARY_OUTPUT_DIRECTORY "${_trtmc_stream_test_root}")
+
+  add_library(trtmc_test_family_stream SHARED core/runtime/tests/fake_family.cpp)
+  target_include_directories(trtmc_test_family_stream PRIVATE ${PROJECT_SOURCE_DIR}/core/runtime/include)
+  target_link_libraries(trtmc_test_family_stream PRIVATE trtmc_core)
+  target_compile_definitions(trtmc_test_family_stream PRIVATE TRTMC_TEST_STREAM_FAMILY)
+  set_target_properties(trtmc_test_family_stream PROPERTIES
+    OUTPUT_NAME trtmc_model_fake LIBRARY_OUTPUT_DIRECTORY "${_trtmc_stream_test_root}"
+    BUILD_RPATH "\$ORIGIN;\$ORIGIN/../..")
+
+  add_custom_target(trtmc_test_stream_runtime
+    COMMAND ${CMAKE_COMMAND} -E make_directory "${_trtmc_stream_test_root}"
+    COMMAND ${CMAKE_COMMAND} -E copy_if_different
+      $<TARGET_FILE:trtmc_test_backend_fake> "${_trtmc_stream_test_root}"
+    DEPENDS trtmc_test_backend_fake trtmc_test_backend_stream trtmc_test_family_stream)
+  add_executable(test_stream_backend_loader core/runtime/tests/test_family_loader.cpp)
+  target_include_directories(test_stream_backend_loader PRIVATE
+    ${PROJECT_SOURCE_DIR}/core/runtime/include ${PROJECT_SOURCE_DIR}/core)
+  target_link_libraries(test_stream_backend_loader PRIVATE trtmc_runtime)
+  target_compile_definitions(test_stream_backend_loader PRIVATE TRTMC_TEST_STREAM_BACKEND)
+  target_compile_options(test_stream_backend_loader PRIVATE -Wall -Wextra -Wpedantic)
+  add_dependencies(test_stream_backend_loader trtmc_test_stream_runtime trtmc_test_default_runtime_root)
+  add_test(NAME stream_backend_loader COMMAND test_stream_backend_loader "${_trtmc_stream_test_root}")
+
   set(_trtmc_api_test_runtime_root "${CMAKE_BINARY_DIR}/tests/api-runtime")
   add_library(trtmc_api_test_backend SHARED core/runtime/tests/fake_backend.cpp)
   target_include_directories(trtmc_api_test_backend PRIVATE

apps/perf_serving/trtmc_perf_serving/video_media.py

@@ -41,7 +41,7 @@ def decoded_identity(path: Path, width: int, height: int, count: int, *,
     if fps is not None and Fraction(streams[0]["avg_frame_rate"]) != fps:
         raise ValueError("downloaded video FPS differs from the requested presentation")
     digest = hashlib.sha256()
-    process = subprocess.Popen(["ffmpeg", "-nostdin", "-v", "error", "-i", str(path),
+    process = subprocess.Popen(["ffmpeg", "-nostdin", "-v", "error", "-i", str(path), "-map", "0:v:0",
                                 "-threads", "1", "-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1"],
                                stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
     try:
@@ -65,11 +65,14 @@ def decoded_identity(path: Path, width: int, height: int, count: int, *,
 
 
 def encode(directory: Path, path: Path, width: int, height: int, count: int,
-           fps: int, max_bytes: int) -> dict:
+           fps: int, max_bytes: int, audio: Path | None = None) -> dict:
+    """Encode the PNG frame sequence losslessly; mux a WAV track as AAC when given."""
     pixels = frame_identity(directory, width, height, count)
+    sound = (["-i", str(audio), "-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac", "-b:a", "192k"]
+             if audio is not None else ["-an"])
     subprocess.run(["ffmpeg", "-nostdin", "-v", "error", "-framerate", str(fps),
-                    "-start_number", "0", "-i", str(directory / "frame-%06d.png"),
-                    "-frames:v", str(count), "-an", "-c:v", "libx264rgb", "-preset", "ultrafast",
+                    "-start_number", "0", "-i", str(directory / "frame-%06d.png"), *sound,
+                    "-frames:v", str(count), "-c:v", "libx264rgb", "-preset", "ultrafast",
                     "-threads", "1", "-crf", "0", "-pix_fmt", "rgb24", "-movflags", "+faststart",
                     "-fs", str(max_bytes), "-y", str(path)],
                    check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, timeout=600)
@@ -79,4 +82,5 @@ def encode(directory: Path, path: Path, width: int, height: int, count: int,
         encoded_hash = hashlib.file_digest(file, "sha256").hexdigest()
     return {"output_width": width, "output_height": height, "output_frames": count,
             "output_pixels_sha256": pixels, "output_sha256": encoded_hash,
-            "output_bytes": path.stat().st_size, "fps": fps, "codec": "libx264rgb-crf0"}
+            "output_bytes": path.stat().st_size, "fps": fps, "codec": "libx264rgb-crf0",
+            "audio": audio is not None}

apps/server/main.cpp

@@ -4,6 +4,7 @@
  */
 
 #include "server/native_worker.h"
+#include "trtmc/control.hpp"
 #include "trtmc/core.hpp"
 
 #include <array>
@@ -64,6 +65,9 @@ int worker_main(int argc, char** argv) {
     std::string runtime_cache;
     std::uint64_t kv_cache_size = 0;
     bool cuda_graphs = false;
+    std::string byok_library;
+    std::string byok_function;
+    std::string byok_name;
     for (int index = 3; index < argc; ++index) {
         const std::string option = argv[index];
         if (option == "--runtime-root")
@@ -74,9 +78,21 @@ int worker_main(int argc, char** argv) {
             kv_cache_size = positive_integer(value_after(argc, argv, index, option), option);
         else if (option == "--cuda-graphs")
             cuda_graphs = true;
+        else if (option == "--byok-library")
+            byok_library = value_after(argc, argv, index, option);
+        else if (option == "--byok-function")
+            byok_function = value_after(argc, argv, index, option);
+        else if (option == "--byok-name")
+            byok_name = value_after(argc, argv, index, option);
         else
             throw std::invalid_argument("unknown _serve-worker option: " + option);
     }
+    const bool byok = !byok_library.empty() || !byok_function.empty() || !byok_name.empty();
+    if (byok && (byok_library.empty() || byok_function.empty() || byok_name.empty()))
+        throw std::invalid_argument("--byok-library, --byok-function and --byok-name go together");
+    // A bring-your-own kernel must be registered before its bundle deserializes.
+    if (byok)
+        trtmc::load_byok_kernel(byok_library, byok_function, byok_name, runtime_root);
     const trtmc::LoadOptions options{runtime_root, kv_cache_size, runtime_cache, cuda_graphs};
     return trtmc::server::run_bundle_worker(bundle, options, std::cin, std::cout);
 }

apps/server/native_worker.cpp

@@ -19,6 +19,7 @@
 #include "trtmc/stream.hpp"
 #include "trtmc/task.h"
 #include "trtmc/text.hpp"
+#include "trtmc/video.hpp"
 
 #include <algorithm>
 #include <chrono>
@@ -768,8 +769,144 @@ int run_existing_worker(ITask& task, std::istream& input, std::ostream& output,
     return run_text_worker(task, input, output);
 }
 
+constexpr std::size_t kMaxVideoReferences = 16;
+
+void write_le(std::ofstream& file, std::uint32_t value, int bytes) {
+    for (int index = 0; index < bytes; ++index)
+        file.put(static_cast<char>((value >> (8 * index)) & 0xFFU));
+}
+
+// 16-bit PCM WAV from the Task's borrowed interleaved float samples.
+Json write_wav(const std::filesystem::path& path, const AudioView& audio) {
+    if (!audio.sample_rate || *audio.sample_rate == 0 || audio.channels == 0 ||
+        audio.channels > 8 || audio.samples.size() % audio.channels != 0)
+        throw std::runtime_error("audio result lacks a sample rate or whole sample frames");
+    const std::uint64_t data_bytes = audio.samples.size() * 2U;
+    if (data_bytes > 0xFFFFFFFFULL - 36U)
+        throw std::runtime_error("audio result exceeds the WAV size limit");
+    const std::uint32_t rate = *audio.sample_rate;
+    const std::uint32_t channels = audio.channels;
+    std::ofstream file(path, std::ios::binary);
+    file.exceptions(std::ios::badbit | std::ios::failbit);
+    file.write("RIFF", 4);
+    write_le(file, static_cast<std::uint32_t>(36U + data_bytes), 4);
+    file.write("WAVEfmt ", 8);
+    write_le(file, 16, 4);
+    write_le(file, 1, 2);
+    write_le(file, channels, 2);
+    write_le(file, rate, 4);
+    write_le(file, rate * channels * 2U, 4);
+    write_le(file, channels * 2U, 2);
+    write_le(file, 16, 2);
+    file.write("data", 4);
+    write_le(file, static_cast<std::uint32_t>(data_bytes), 4);
+    for (const float sample : audio.samples) {
+        if (!std::isfinite(sample))
+            throw std::runtime_error("audio result contains non-finite samples");
+        const auto value =
+            static_cast<std::int16_t>(std::lround(std::clamp(sample, -1.0F, 1.0F) * 32767.0F));
+        write_le(file, static_cast<std::uint16_t>(value), 2);
+    }
+    return {{"audio_sample_rate", rate},
+            {"audio_channels", channels},
+            {"audio_samples", audio.samples.size() / channels}};
+}
+
+Json write_audio_video(const AudioVideoGenerationResult& result, const std::string& frame_directory,
+                       const std::string& audio_path, double wall_ms) {
+    const auto frames = result.frames();
+    if (frames.size() == 0 || frames[0].channels != 3 || frames[0].height == 0 ||
+        frames[0].width == 0)
+        throw std::runtime_error("reference video result has no RGB frames");
+    const auto height = frames[0].height;
+    const auto width = frames[0].width;
+    std::filesystem::create_directories(frame_directory);
+    for (std::size_t index = 0; index < frames.size(); ++index) {
+        const auto& frame = frames[index];
+        if (frame.height != height || frame.width != width || frame.channels != 3 ||
+            frame.pixel_count != std::uint64_t{height} * width * 3U)
+            throw std::runtime_error("reference video frames differ in geometry");
+        std::ostringstream name;
+        name << "frame-" << std::setw(6) << std::setfill('0') << index << ".png";
+        cli::io::save_png(
+            (std::filesystem::path(frame_directory) / name.str()).string(),
+            Span<const float>{frame.pixels, static_cast<std::size_t>(frame.pixel_count)},
+            static_cast<int>(width), static_cast<int>(height), 3);
+    }
+    const auto times = result.timestamps_seconds();
+    double fps = 0.0;
+    if (times.size() >= 2 && times[times.size() - 1] > times[0])
+        fps = static_cast<double>(times.size() - 1) / (times[times.size() - 1] - times[0]);
+    auto summary = write_wav(audio_path, result.audio());
+    summary.update(Json{{"width", width},
+                        {"height", height},
+                        {"channels", 3},
+                        {"num_frames", frames.size()},
+                        {"fps", fps},
+                        {"audio_start_seconds", result.audio_start_seconds()},
+                        {"model_call_ms", wall_ms},
+                        {"timing_scope", "public_task_call_wall"}});
+    return summary;
+}
+
+int run_reference_video_worker(const Model& model, std::istream& input, std::ostream& output) {
+    const auto task = model.task<ReferencesTextToAudioVideo>();
+    const auto fields = task.config_fields();
+    const Json ready{{"event", "ready"},
+                     {"protocol_version", 1},
+                     {"capabilities", Json::array({"reference_video_generation"})}};
+    return run_protocol(
+        ready,
+        {{"generate_reference_video",
+          [&](const Json& request) {
+              const auto prompt = string_field(request, "prompt");
+              const auto frame_directory = string_field(request, "output_path");
+              const auto audio_path = string_field(request, "audio_path");
+              const auto paths = request.value("reference_paths", Json::array());
+              if (!paths.is_array() || paths.empty() || paths.size() > kMaxVideoReferences)
+                  throw ProtocolError("reference_paths must list 1 to 16 images");
+              // The Task reads borrowed RGB uint8 pixels; keep every buffer alive until run
+              // returns.
+              std::vector<std::vector<std::uint8_t>> pixels;
+              pixels.reserve(paths.size());
+              ReferencesTextToAudioVideoRequest references{{}, prompt};
+              for (const auto& path : paths) {
+                  if (!path.is_string())
+                      throw ProtocolError("reference_paths must contain strings");
+                  cli::io::LoadedImage image;
+                  try {
+                      image = cli::io::read_image(path.get<std::string>());
+                  } catch (const std::exception&) {
+                      throw ProtocolError("reference must be a readable image file");
+                  }
+                  if (image.height <= 0 || image.width <= 0 || image.pixels.empty())
+                      throw ProtocolError("reference must contain RGB pixels");
+                  auto& bytes = pixels.emplace_back(image.pixels.size());
+                  std::transform(image.pixels.begin(), image.pixels.end(), bytes.begin(),
+                                 [](float value) {
+                                     return static_cast<std::uint8_t>(
+                                         std::lround(std::clamp(value, 0.0F, 1.0F) * 255.0F));
+                                 });
+                  references.references.emplace_back(
+                      ImageInput{Span<const std::uint8_t>{bytes.data(), bytes.size()},
+                                 static_cast<std::uint32_t>(image.height),
+                                 static_cast<std::uint32_t>(image.width)});
+              }
+              const auto config = parse_sdk_config(request, fields);
+              const auto started = std::chrono::steady_clock::now();
+              const auto result = task.run(references, config);
+              const auto wall_ms = std::chrono::duration<double, std::milli>(
+                                       std::chrono::steady_clock::now() - started)
+                                       .count();
+              return write_audio_video(result, frame_directory, audio_path, wall_ms);
+          }}},
+        input, output);
+}
+
 int run_model_worker(const Model& model, std::istream& input, std::ostream& output) {
     const auto primary = model.info().bundle_task;
+    if (primary == ReferencesTextToAudioVideo::kTask)
+        return run_reference_video_worker(model, input, output);
     if (primary == ImagesTextToImageEdit::kTask || primary == IImageEditing::kTask) {
         const auto task = model.task<ImagesTextToImageEdit>();
         const auto fields = task.config_fields();