Fix get_allow_patterns to include non-indexed safetensors files

src/exo/worker/download/huggingface_utils.py

@@ -136,8 +136,17 @@ def get_allow_patterns(weight_map: dict[str, str], shard: ShardMetadata) -> list
                         shard_specific_patterns.add(filename)
                 else:
                     shard_specific_patterns.add(filename)
+
         else:
             shard_specific_patterns = set(["*.safetensors"])
+
+        # Include safetensors files from non-shardable components that have no index file
+        for component in shard.model_meta.components:
+            if not component.can_shard and component.safetensors_index_filename is None:
+                component_pattern = (
+                    f"{component.component_path.rstrip('/')}/*.safetensors"
+                )
+                shard_specific_patterns.add(component_pattern)
     else:
         if weight_map:
             for tensor_name, filename in weight_map.items():

src/exo/worker/engines/mflux/shard_mflux.py

@@ -152,7 +152,7 @@ class FluxJointPipelineLastBlock(CustomMlxJointBlock):
 class FluxJointToSingleTransition(CustomMlxJointBlock):
     def __init__(
         self,
-        original_block: nn.Module,
+        original_block: _JointBlock,
         rank: int,
         world_size: int,
         group: mx.distributed.Group,

src/exo/worker/engines/mflux/utils_mflux.py

@@ -9,21 +9,9 @@ from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init
 
 
 def initialize_mflux(bound_instance: BoundInstance) -> Flux1:
-    """Initialize Flux1 model for single-node or distributed inference.
-
-    Automatically detects multi-node setup and applies pipeline parallelism
-    to the transformer component while keeping other components replicated.
-
-    Args:
-        bound_instance: Instance binding information with shard metadata
-
-    Returns:
-        Flux1 model, with transformer sharded if running in distributed mode
-    """
     model_id = bound_instance.bound_shard.model_meta.model_id
     model_path = build_model_path(model_id)
 
-    # Check if this is a distributed setup
     is_distributed = len(bound_instance.instance.shard_assignments.node_to_runner) > 1
 
     if not is_distributed: