diff --git a/runtime/executor/BUILD b/runtime/executor/BUILD
index 08d41abe5..0e5c25411 100644
--- a/runtime/executor/BUILD
+++ b/runtime/executor/BUILD
@@ -386,6 +386,7 @@ cc_library(
         "@litert//litert/cc:litert_model",
         "@litert//litert/cc:litert_ranked_tensor_type",
         "//runtime/components:model_resources",
+        "//runtime/proto:llm_model_type_cc_proto",
         "//runtime/util:litert_status_util",
     ] + select({
         "@platforms//os:ios": [],
@@ -404,6 +405,7 @@ cc_test(
         ":llm_executor_settings",
         ":llm_litert_compiled_model_executor_factory",
         "@com_google_googletest//:gtest_main",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@litert//litert/cc:litert_environment",
@@ -412,7 +414,6 @@ cc_test(
         "//runtime/components:model_resources_litert_lm",
         "//runtime/util:litert_lm_loader",
         "//runtime/util:litert_status_util",
-        "//runtime/util:model_asset_bundle_resources",
         "//runtime/util:scoped_file",
         "//runtime/util:test_utils",
     ],
diff --git a/runtime/executor/llm_litert_compiled_model_executor_factory.cc b/runtime/executor/llm_litert_compiled_model_executor_factory.cc
index 3872b9c33..6b9bc71ac 100644
--- a/runtime/executor/llm_litert_compiled_model_executor_factory.cc
+++ b/runtime/executor/llm_litert_compiled_model_executor_factory.cc
@@ -34,6 +34,7 @@
 #include "runtime/executor/llm_executor.h"
 #include "runtime/executor/llm_executor_settings.h"
 #include "runtime/executor/llm_litert_compiled_model_executor.h"
+#include "runtime/proto/llm_model_type.pb.h"
 #include "runtime/util/status_macros.h"
 
 #if !defined(LITERT_DISABLE_NPU)
@@ -131,10 +132,11 @@ absl::StatusOr<std::unique_ptr<LlmExecutor>>
 CreateCpuOrGpuLlmLiteRtCompiledModelExecutor(
     LlmExecutorSettings executor_settings, Environment& lrt_env,
     ModelResources& resources) {
+  std::unique_ptr<LlmExecutor> executor;
+
   ASSIGN_OR_RETURN(const litert::Model* litert_model,
                    resources.GetTFLiteModel(ModelType::kTfLitePrefillDecode));
 
-  std::unique_ptr<LlmExecutor> executor;
   ASSIGN_OR_RETURN(bool is_dynamic_model, IsDynamicModel(*litert_model));
   if (is_dynamic_model) {
     ASSIGN_OR_RETURN(executor, LlmLiteRtCompiledModelExecutorDynamic::Create(
diff --git a/runtime/executor/llm_litert_compiled_model_executor_factory_test.cc b/runtime/executor/llm_litert_compiled_model_executor_factory_test.cc
index 4bdd1df2e..fcb021397 100644
--- a/runtime/executor/llm_litert_compiled_model_executor_factory_test.cc
+++ b/runtime/executor/llm_litert_compiled_model_executor_factory_test.cc
@@ -21,6 +21,7 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"  // from @com_google_absl
 #include "absl/status/statusor.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "litert/cc/litert_environment.h"  // from @litert
@@ -31,7 +32,6 @@
 #include "runtime/executor/llm_executor_settings.h"
 #include "runtime/util/litert_lm_loader.h"
 #include "runtime/util/scoped_file.h"
-#include "runtime/util/status_macros.h"
 #include "runtime/util/test_utils.h"  // IWYU pragma: keep
 
 namespace litert::lm {
diff --git a/runtime/framework/resource_management/serial_execution_manager.cc b/runtime/framework/resource_management/serial_execution_manager.cc
index a63fdc1f8..61905f97c 100644
--- a/runtime/framework/resource_management/serial_execution_manager.cc
+++ b/runtime/framework/resource_management/serial_execution_manager.cc
@@ -786,7 +786,7 @@ absl::Status SerialExecutionManager::AddDecodeTask(
     session_info->stop_token_detector->ResetBatch(num_output_candidates);
     std::optional<Sampler*> optional_sampler = std::nullopt;
     std::optional<litert::TensorBuffer> decoded_ids_buffer = std::nullopt;
-    if (session_info->sampler != nullptr) {
+    if (session_info->sampler != nullptr && !is_text_diffusion) {
       optional_sampler = session_info->sampler.get();
       std::vector<int> decoded_ids(num_output_candidates,
                                    session_info->last_prefill_token_id);
diff --git a/runtime/framework/resource_management/threaded_execution_manager.cc b/runtime/framework/resource_management/threaded_execution_manager.cc
index 786a2a858..1a2199c61 100644
--- a/runtime/framework/resource_management/threaded_execution_manager.cc
+++ b/runtime/framework/resource_management/threaded_execution_manager.cc
@@ -908,7 +908,8 @@ absl::Status ThreadedExecutionManager::AddDecodeTask(
     session_info->stop_token_detector->ResetBatch(num_output_candidates);
     std::optional<Sampler*> optional_sampler = std::nullopt;
     std::optional<litert::TensorBuffer> decoded_ids_buffer = std::nullopt;
-    if (session_info->sampler != nullptr) {
+    bool is_text_diffusion = false;
+    if (session_info->sampler != nullptr && !is_text_diffusion) {
       optional_sampler = session_info->sampler.get();
       std::vector<int> decoded_ids(num_output_candidates,
                                    session_info->last_prefill_token_id);