Move some more logic to adaptor

src/exo/worker/engines/mflux/pipefusion/adapter.py

@@ -5,7 +5,7 @@ from typing import Any, Protocol
 import mlx.core as mx
 from mflux.config.runtime_config import RuntimeConfig
 
-from exo.worker.engines.mflux.config.model_config import ImageModelConfig
+from exo.worker.engines.mflux.config.model_config import BlockType, ImageModelConfig
 from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
 
 
@@ -281,3 +281,78 @@ class ModelAdapter(Protocol):
             List of single transformer blocks
         """
         ...
+
+    def get_blocks(self, transformer: Any) -> list[tuple[Any, BlockType]]:
+        """Get all transformer blocks in execution order with their types.
+
+        This method provides a unified view of all blocks, regardless of their
+        specific type (joint, single, unified). New model adapters can override
+        this to return blocks in their native arrangement.
+
+        Args:
+            transformer: The transformer model
+
+        Returns:
+            List of (block, block_type) tuples in execution order
+        """
+        ...
+
+    def apply_block(
+        self,
+        block: Any,
+        block_type: BlockType,
+        hidden_states: mx.array,
+        encoder_hidden_states: mx.array | None,
+        text_embeddings: mx.array,
+        rotary_embeddings: mx.array,
+        kv_cache: ImagePatchKVCache | None,
+        mode: "BlockWrapperMode",
+        text_seq_len: int,
+        patch_start: int | None = None,
+        patch_end: int | None = None,
+    ) -> tuple[mx.array, mx.array | None]:
+        """Apply any transformer block type.
+
+        This unified method dispatches to the appropriate block-specific logic
+        based on block_type. New model adapters can implement this directly
+        without needing separate apply_joint_block/apply_single_block methods.
+
+        Args:
+            block: The transformer block
+            block_type: Type of block (JOINT, SINGLE, or UNIFIED)
+            hidden_states: Image hidden states (or concatenated for SINGLE/UNIFIED)
+            encoder_hidden_states: Text hidden states (None for SINGLE/UNIFIED)
+            text_embeddings: Conditioning embeddings
+            rotary_embeddings: Rotary position embeddings
+            kv_cache: KV cache (None if not using cache)
+            mode: CACHING or PATCHED mode
+            text_seq_len: Text sequence length
+            patch_start: Start index for patched mode
+            patch_end: End index for patched mode
+
+        Returns:
+            Tuple of (hidden_states, encoder_hidden_states or None)
+            - For JOINT blocks: (image_hidden, text_hidden)
+            - For SINGLE/UNIFIED blocks: (concatenated_hidden, None)
+        """
+        ...
+
+    def merge_streams(
+        self,
+        hidden_states: mx.array,
+        encoder_hidden_states: mx.array,
+    ) -> mx.array:
+        """Merge image and text streams for transition to single/unified blocks.
+
+        This is called at the transition point from joint blocks (which process
+        image and text separately) to single/unified blocks (which process them
+        together). Override to customize the merge strategy.
+
+        Args:
+            hidden_states: Image hidden states
+            encoder_hidden_states: Text hidden states
+
+        Returns:
+            Merged hidden states (default: concatenate [text, image])
+        """
+        ...

src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py

@@ -276,9 +276,9 @@ class DistributedDenoising:
 
         # === PHASE 3: Joint→Single Transition ===
         if self.owns_concat_stage:
-            # Concatenate encoder and hidden states
-            concatenated = mx.concatenate(
-                [encoder_hidden_states, hidden_states], axis=1
+            # Merge encoder and hidden states using adapter hook
+            concatenated = self.adapter.merge_streams(
+                hidden_states, encoder_hidden_states
             )
 
             if self.has_single_blocks:
@@ -435,7 +435,7 @@ class DistributedDenoising:
                     )
 
             if self.owns_concat_stage:
-                patch_concat = mx.concatenate([encoder_hidden_states, patch], axis=1)
+                patch_concat = self.adapter.merge_streams(patch, encoder_hidden_states)
 
                 if self.has_single_blocks:
                     patch = patch_concat

src/exo/worker/engines/mflux/pipefusion/flux_adapter.py

@@ -16,7 +16,7 @@ from mflux.models.flux.model.flux_transformer.single_transformer_block import (
 from mflux.models.flux.model.flux_transformer.transformer import Transformer
 from mflux.models.flux.variants.txt2img.flux import Flux1
 
-from exo.worker.engines.mflux.config.model_config import ImageModelConfig
+from exo.worker.engines.mflux.config.model_config import BlockType, ImageModelConfig
 from exo.worker.engines.mflux.pipefusion.adapter import BlockWrapperMode
 from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
 
@@ -177,6 +177,84 @@ class FluxModelAdapter:
         """Get the list of single transformer blocks from the Flux model."""
         return list(transformer.single_transformer_blocks)
 
+    def get_blocks(
+        self, transformer: Transformer
+    ) -> list[tuple[JointTransformerBlock | SingleTransformerBlock, BlockType]]:
+        """Get all transformer blocks in execution order with their types.
+
+        For Flux models, this returns joint blocks first, then single blocks.
+        """
+        blocks: list[tuple[JointTransformerBlock | SingleTransformerBlock, BlockType]] = []
+        for block in self.get_joint_blocks(transformer):
+            blocks.append((block, BlockType.JOINT))
+        for block in self.get_single_blocks(transformer):
+            blocks.append((block, BlockType.SINGLE))
+        return blocks
+
+    def apply_block(
+        self,
+        block: JointTransformerBlock | SingleTransformerBlock,
+        block_type: BlockType,
+        hidden_states: mx.array,
+        encoder_hidden_states: mx.array | None,
+        text_embeddings: mx.array,
+        rotary_embeddings: mx.array,
+        kv_cache: ImagePatchKVCache | None,
+        mode: BlockWrapperMode,
+        text_seq_len: int,
+        patch_start: int | None = None,
+        patch_end: int | None = None,
+    ) -> tuple[mx.array, mx.array | None]:
+        """Apply any transformer block type.
+
+        Delegates to the appropriate block-specific method based on block_type.
+        """
+        if block_type == BlockType.JOINT:
+            assert isinstance(block, JointTransformerBlock)
+            assert encoder_hidden_states is not None
+            enc_out, hidden_out = self.apply_joint_block(
+                block=block,
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                text_embeddings=text_embeddings,
+                rotary_embeddings=rotary_embeddings,
+                kv_cache=kv_cache,
+                mode=mode,
+                text_seq_len=text_seq_len,
+                patch_start=patch_start,
+                patch_end=patch_end,
+            )
+            return hidden_out, enc_out
+        elif block_type == BlockType.SINGLE:
+            assert isinstance(block, SingleTransformerBlock)
+            hidden_out = self.apply_single_block(
+                block=block,
+                hidden_states=hidden_states,
+                text_embeddings=text_embeddings,
+                rotary_embeddings=rotary_embeddings,
+                kv_cache=kv_cache,
+                mode=mode,
+                text_seq_len=text_seq_len,
+                patch_start=patch_start,
+                patch_end=patch_end,
+            )
+            return hidden_out, None
+        else:
+            raise NotImplementedError(
+                f"Block type {block_type} not supported by FluxModelAdapter"
+            )
+
+    def merge_streams(
+        self,
+        hidden_states: mx.array,
+        encoder_hidden_states: mx.array,
+    ) -> mx.array:
+        """Merge image and text streams for transition to single blocks.
+
+        For Flux models, this concatenates [text, image] along the sequence dimension.
+        """
+        return mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
+
     # -------------------------------------------------------------------------
     # Joint block implementations
     # -------------------------------------------------------------------------