From d926df8f956f7e4e94e250fa61861e1f72ee6198 Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Mon, 5 Jan 2026 15:13:07 +0000 Subject: [PATCH] Fix comms issue for models without single blocks --- src/exo/worker/engines/image/pipeline/runner.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index ffc93d93..171dddab 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/src/exo/worker/engines/image/pipeline/runner.py @@ -583,8 +583,8 @@ class DiffusionRunner: hidden_states, encoder_hidden_states ) - if self.has_single_blocks: - # We continue with single blocks on this stage + if self.has_single_blocks or self.is_last_stage: + # Keep locally: either for single blocks or final projection hidden_states = concatenated else: # Send concatenated state to next stage (which has single blocks) @@ -773,7 +773,8 @@ class DiffusionRunner: if self.owns_concat_stage: patch_concat = self.adapter.merge_streams(patch, encoder_hidden_states) - if self.has_single_blocks: + if self.has_single_blocks or self.is_last_stage: + # Keep locally: either for single blocks or final projection patch = patch_concat else: mx.eval(