From 892ab4e295bdde6fa495b348774ddef0e809e4bd Mon Sep 17 00:00:00 2001 From: chuanhaoyu Date: Tue, 18 Aug 2026 18:34:45 +0800 Subject: [PATCH] perf: keep raw dit image tensors on cpu until vae preprocess. --- xllm/core/runtime/dit_forward_params.h | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/xllm/core/runtime/dit_forward_params.h b/xllm/core/runtime/dit_forward_params.h index d5e8e72d75..bc926bcce5 100644 --- a/xllm/core/runtime/dit_forward_params.h +++ b/xllm/core/runtime/dit_forward_params.h @@ -205,28 +205,30 @@ struct DiTForwardInput { } if (images.defined()) { - input.images = images.to(device, /*dtype=*/torch::kUInt8); + input.images = images.to(torch::kCPU, /*dtype=*/torch::kUInt8); } if (mask_images.defined()) { - input.mask_images = mask_images.to(device, /*dtype=*/torch::kUInt8); + input.mask_images = mask_images.to(torch::kCPU, /*dtype=*/torch::kUInt8); } for (auto& img : input.images_list) { - img = img.to(device, /*dtype=*/torch::kUInt8); + img = img.to(torch::kCPU, /*dtype=*/torch::kUInt8); } if (control_image.defined()) { - input.control_image = control_image.to(device, /*dtype=*/torch::kUInt8); + input.control_image = + control_image.to(torch::kCPU, /*dtype=*/torch::kUInt8); } if (last_images.defined()) { - input.last_images = last_images.to(device, /*dtype=*/torch::kUInt8); + input.last_images = last_images.to(torch::kCPU, /*dtype=*/torch::kUInt8); } if (prompt_audio.defined()) { input.prompt_audio = prompt_audio.to(device, torch::kFloat32); } + return input; }