feat: refine dit modules.

Author: xiao-yu-chen

xllm/api_service/api_service_impl.h

@@ -19,8 +19,8 @@ limitations under the License.
 #include <memory>
 
 #include "call.h"
-#include "core/runtime/dit_master.h"
 #include "core/runtime/llm_master.h"
+
 namespace xllm {
 
 template <typename T>

xllm/api_service/image_generation_service_impl.cpp

@@ -44,8 +44,8 @@ bool send_result_to_client_brpc(std::shared_ptr<ImageGenerationCall> call,
   proto_output->mutable_results()->Reserve(outputs.size());
   for (const auto& output : outputs) {
     auto* proto_result = proto_output->add_results();
-    // proto_result->set_base64(output.image_tensor); // TODO proto tensor to
-    // base64
+
+    // proto_result->set_image(output.image);
     proto_result->set_width(output.width);
     proto_result->set_height(output.height);
     proto_result->set_seed(output.seed);
@@ -73,12 +73,11 @@ void ImageGenerationServiceImpl::process_async(
     call->finish_with_error(StatusCode::UNKNOWN, "Model not supported");
     return;
   }
-  // create RequestParams for image generation request
-  // set is_image_generation and max_tokens = 1 to control engine step once.
+
+  // create DiTRequestParams for image generation request
   DiTRequestParams request_params(
       rpc_request, call->get_x_request_id(), call->get_x_request_time());
-  // TODO only support input_str for now
-  auto& input = rpc_request.input().prompt();
+
   // schedule the request
   master_->handle_request(
       std::move(request_params),
@@ -88,16 +87,13 @@ void ImageGenerationServiceImpl::process_async(
        request_id = request_params.request_id,
        created_time = absl::ToUnixSeconds(absl::Now())](
           const DiTRequestOutput& req_output) -> bool {
-        LOG(INFO) << "into callback before request finished";
-        LOG(INFO) << req_output.outputs.size();
-        LOG(INFO) << req_output.outputs[0].image_tensor;
         if (req_output.status.has_value()) {
           const auto& status = req_output.status.value();
           if (!status.ok()) {
             return call->finish_with_error(status.code(), status.message());
           }
         }
-        LOG(INFO) << "into callback after request finished";
+
         return send_result_to_client_brpc(
             call, request_id, created_time, model, req_output);
       });

xllm/api_service/image_generation_service_impl.h

@@ -16,8 +16,6 @@ limitations under the License.
 #pragma once
 #include <absl/container/flat_hash_set.h>
 
-#include "api_service/api_service_impl.h"
-#include "api_service/call.h"
 #include "api_service/non_stream_call.h"
 #include "image_generation.pb.h"
 

xllm/core/framework/batch/batch.h

@@ -21,7 +21,6 @@ limitations under the License.
 #include <limits>
 #include <vector>
 
-#include "framework/request/dit_request_params.h"
 #include "framework/request/mm_data.h"
 #include "framework/request/request.h"
 #include "framework/request/sequence.h"

xllm/core/framework/batch/dit_batch.cpp

@@ -24,13 +24,14 @@ limitations under the License.
 namespace xllm {
 
 DiTForwardInput xllm::DiTBatch::prepare_forward_input() {
+  CHECK(!dit_request_vec_.empty());
+
   DiTForwardInput forward_input;
-  if (dit_request_data_vec_.empty()) {
-    return forward_input;
-  }
-  forward_input.input_params = dit_request_data_vec_[0].input_params;
-  forward_input.generation_params = dit_request_data_vec_[0].generation_params;
+  forward_input.input_params = dit_request_vec_[0]->state().input_params();
+  forward_input.generation_params =
+      dit_request_vec_[0]->state().generation_params();
+
   return forward_input;
 }
 
-}  // namespace xllm
+}  // namespace xllm

xllm/core/framework/batch/dit_batch.h

@@ -21,25 +21,25 @@ limitations under the License.
 #include <limits>
 #include <vector>
 
-#include "framework/request/dit_request_params.h"
+#include "framework/request/dit_request.h"
 #include "runtime/dit_forward_params.h"
 
 namespace xllm {
 
 struct DiTBatch {
  public:
   DiTBatch() = default;
-  void add(const DiTRequestParams& dit_request_state) {
-    dit_request_data_vec_.emplace_back(dit_request_state);
+  void add(const std::shared_ptr<DiTRequest>& request) {
+    dit_request_vec_.emplace_back(request);
   }
-  size_t size() const { return dit_request_data_vec_.size(); }
-  bool empty() const { return dit_request_data_vec_.empty(); }
+  size_t size() const { return dit_request_vec_.size(); }
+  bool empty() const { return dit_request_vec_.empty(); }
 
   // prepare forward input
   DiTForwardInput prepare_forward_input();
 
  private:
-  std::vector<DiTRequestParams> dit_request_data_vec_;
+  std::vector<std::shared_ptr<DiTRequest>> dit_request_vec_;
 };
 
-}  // namespace xllm
+}  // namespace xllm

xllm/core/framework/dit_model_loader.cpp

@@ -237,13 +237,13 @@ DiTModelLoader::DiTModelLoader(const std::string& model_root_path)
 
   const nlohmann::json root_json = model_index_reader.data();
   if (!root_json.is_object()) {
-    LOG(FATAL) << "DITModelLoader: model_index.json root is not an object!";
+    LOG(FATAL) << "DiTModelLoader: model_index.json root is not an object!";
   }
 
   // parse model_index.json & initialize model_loader
   for (const auto& [json_key, json_value] : root_json.items()) {
     if (!json_value.is_array() || json_value.size() != 2) {
-      LOG(WARNING) << "DITModelLoader: Invalid format for component! "
+      LOG(WARNING) << "DiTModelLoader: Invalid format for component! "
                    << "JsonKey=" << json_key
                    << ", Expected [library, class_name] array";
       continue;
@@ -254,13 +254,13 @@ DiTModelLoader::DiTModelLoader(const std::string& model_root_path)
         std::filesystem::path(model_root_path_) / json_key;
     const std::string component_folder = component_folder_path.string();
     if (!std::filesystem::exists(component_folder)) {
-      LOG(WARNING) << "DITModelLoader: Component folder not found! "
+      LOG(WARNING) << "DiTModelLoader: Component folder not found! "
                    << "ComponentName=" << component_name
                    << ", Folder=" << component_folder;
       continue;
     }
     if (!std::filesystem::is_directory(component_folder)) {
-      LOG(WARNING) << "DITModelLoader: Component path is not a directory! "
+      LOG(WARNING) << "DiTModelLoader: Component path is not a directory! "
                    << "ComponentName=" << component_name
                    << ", Path=" << component_folder;
       continue;

xllm/core/framework/model/dit_model.h

@@ -29,8 +29,8 @@ class DiTModel : public torch::nn::Module {
  public:
   ~DiTModel() override = default;
 
-  virtual torch::Tensor forward(const InputParams& input_params,
-                                const GenerationParams& gen_params) = 0;
+  virtual torch::Tensor forward(const DiTInputParams& input_params,
+                                const DiTGenerationParams& gen_params) = 0;
   virtual torch::Device device() const = 0;
   virtual const torch::TensorOptions& options() const = 0;
   virtual void load_model(std::unique_ptr<DiTModelLoader> loader) = 0;
@@ -43,8 +43,8 @@ class DiTModelImpl : public DiTModel {
       : model_(std::move(model)), options_(options) {
     LOG(INFO) << "DiTModelImpl created.";
   }
-  torch::Tensor forward(const InputParams& input_params,
-                        const GenerationParams& gen_params) override {
+  torch::Tensor forward(const DiTInputParams& input_params,
+                        const DiTGenerationParams& gen_params) override {
     return model_->forward(input_params, gen_params);
   }
   torch::Device device() const override { return options_.device(); }

xllm/core/framework/request/CMakeLists.txt

@@ -5,10 +5,13 @@ cc_library(
   NAME 
     request
   HDRS
+    dit_request.h
+    dit_request_params.h
     finish_reason.h
     incremental_decoder.h
     mm_data.h
     mm_input_helper.h
+    request_base.h
     request.h
     dit_request.h
     request_output.h
@@ -23,6 +26,7 @@ cc_library(
     stopping_checker.h
     priority_comparator.h
   SRCS
+    dit_request.cpp
     finish_reason.cpp
     incremental_decoder.cpp
     mm_data.cpp

xllm/core/framework/request/dit_request.cpp

@@ -31,19 +31,9 @@ DiTRequest::DiTRequest(const std::string& request_id,
                        const std::string& x_request_id,
                        const std::string& x_request_time,
                        const DiTRequestState& state,
-                       const std::string& service_request_id,
-                       bool offline,
-                       int32_t slo_ms,
-                       RequestPriority priority)
-    : created_time_(absl::Now()),
-      request_id_(request_id),
-      service_request_id_(service_request_id),
-      x_request_id_(x_request_id),
-      x_request_time_(x_request_time),
-      state_(state),
-      offline_(offline),
-      slo_ms_(slo_ms),
-      priority_(priority) {}
+                       const std::string& service_request_id)
+    : RequestBase(request_id, x_request_id, x_request_time, service_request_id),
+      state_(state) {}
 
 bool DiTRequest::finished() const { return true; }
 
@@ -54,16 +44,23 @@ void DiTRequest::log_statistic(double total_latency) {
             << "total_latency: " << total_latency * 1000 << "ms";
 }
 
-DiTRequestOutput DiTRequest::generate_dit_output(DiTForwardOutput dit_output) {
+const DiTRequestOutput DiTRequest::generate_output(
+    DiTForwardOutput dit_output) {
   DiTRequestOutput output;
   output.request_id = request_id_;
   output.service_request_id = service_request_id_;
   output.status = Status(StatusCode::OK);
   output.finished = finished();
   output.cancelled = false;
+
   DiTGenerationOutput result;
-  result.image_tensor = dit_output.image;
+  result.image = dit_output.image;
+  result.height = state_.generation_params().height;
+  result.width = state_.generation_params().width;
+  result.seed = state_.generation_params().seed.value();
   output.outputs.push_back(result);
+
   return output;
 }
-}  // namespace xllm
+
+}  // namespace xllm