Refactor filestructure

src/exo/worker/engines/image/__init__.py

@@ -1,8 +1,14 @@
 from exo.worker.engines.image.base import ImageGenerator
+from exo.worker.engines.image.distributed_model import (
+    DistributedImageModel,
+    initialize_image_model,
+)
 from exo.worker.engines.image.generate import generate_image, warmup_image_generator
 
 __all__ = [
+    "DistributedImageModel",
     "ImageGenerator",
     "generate_image",
+    "initialize_image_model",
     "warmup_image_generator",
 ]

src/exo/worker/engines/image/distributed_model.py

@@ -15,11 +15,9 @@ from PIL import Image
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.download.download_utils import build_model_path
-from exo.worker.engines.mflux.config import get_config_for_model
-from exo.worker.engines.mflux.config.model_config import ImageModelConfig
-from exo.worker.engines.mflux.pipefusion import create_adapter_for_model
-from exo.worker.engines.mflux.pipefusion.adapter import ModelAdapter
-from exo.worker.engines.mflux.pipefusion.diffusion_runner import DiffusionRunner
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models import create_adapter_for_model, get_config_for_model
+from exo.worker.engines.image.pipeline import DiffusionRunner, ModelAdapter
 from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier
 from exo.worker.runner.bootstrap import logger
 
@@ -256,3 +254,8 @@ class DistributedImageModel:
             image_strength=runtime_config.image_strength,
             generation_time=0,  # TODO: Track time in runner if needed
         )
+
+
+def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel:
+    """Initialize DistributedImageModel from a BoundInstance."""
+    return DistributedImageModel.from_bound_instance(bound_instance)

src/exo/worker/engines/image/models/__init__.py

@@ -0,0 +1,80 @@
+from pathlib import Path
+from typing import Callable
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.flux import (
+    FLUX_DEV_CONFIG,
+    FLUX_SCHNELL_CONFIG,
+    FluxModelAdapter,
+)
+from exo.worker.engines.image.pipeline.adapter import ModelAdapter
+
+__all__ = [
+    "create_adapter_for_model",
+    "get_config_for_model",
+    "FluxModelAdapter",
+    "FLUX_DEV_CONFIG",
+    "FLUX_SCHNELL_CONFIG",
+]
+
+# Type alias for adapter factory functions
+# Factory takes (config, model_id, local_path, quantize) and returns a ModelAdapter
+AdapterFactory = Callable[[ImageModelConfig, str, Path, int | None], ModelAdapter]
+
+# Registry maps model_family string to adapter factory
+_ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
+    "flux": FluxModelAdapter,
+}
+
+# Config registry: maps model ID patterns to configs
+_CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
+    "flux.1-schnell": FLUX_SCHNELL_CONFIG,
+    "flux.1-dev": FLUX_DEV_CONFIG,
+}
+
+
+def get_config_for_model(model_id: str) -> ImageModelConfig:
+    """Get configuration for a model ID.
+
+    Args:
+        model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell")
+
+    Returns:
+        The model configuration
+
+    Raises:
+        ValueError: If no configuration found for model ID
+    """
+    model_id_lower = model_id.lower()
+
+    for pattern, config in _CONFIG_REGISTRY.items():
+        if pattern in model_id_lower:
+            return config
+
+    raise ValueError(f"No configuration found for model: {model_id}")
+
+
+def create_adapter_for_model(
+    config: ImageModelConfig,
+    model_id: str,
+    local_path: Path,
+    quantize: int | None = None,
+) -> ModelAdapter:
+    """Create a model adapter for the given configuration.
+
+    Args:
+        config: The model configuration
+        model_id: The model identifier
+        local_path: Path to the model weights
+        quantize: Optional quantization bits
+
+    Returns:
+        A ModelAdapter instance
+
+    Raises:
+        ValueError: If no adapter found for model family
+    """
+    factory = _ADAPTER_REGISTRY.get(config.model_family)
+    if factory is None:
+        raise ValueError(f"No adapter found for model family: {config.model_family}")
+    return factory(config, model_id, local_path, quantize)

src/exo/worker/engines/image/models/flux/__init__.py

@@ -0,0 +1,11 @@
+from exo.worker.engines.image.models.flux.adapter import FluxModelAdapter
+from exo.worker.engines.image.models.flux.config import (
+    FLUX_DEV_CONFIG,
+    FLUX_SCHNELL_CONFIG,
+)
+
+__all__ = [
+    "FluxModelAdapter",
+    "FLUX_DEV_CONFIG",
+    "FLUX_SCHNELL_CONFIG",
+]

src/exo/worker/engines/image/models/flux/adapter.py

@@ -13,13 +13,13 @@ from mflux.models.flux.model.flux_transformer.joint_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 BlockType, ImageModelConfig
-from exo.worker.engines.mflux.pipefusion.adapter import (
+from exo.worker.engines.image.config import BlockType, ImageModelConfig
+from exo.worker.engines.image.pipeline.adapter import (
     BlockWrapperMode,
     JointBlockInterface,
     SingleBlockInterface,
 )
-from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
 
 
 class FluxModelAdapter:

src/exo/worker/engines/image/models/flux/config.py

@@ -1,4 +1,4 @@
-from exo.worker.engines.mflux.config.model_config import (
+from exo.worker.engines.image.config import (
     BlockType,
     ImageModelConfig,
     TransformerBlockConfig,

src/exo/worker/engines/image/pipeline/__init__.py

@@ -0,0 +1,25 @@
+from exo.worker.engines.image.pipeline.adapter import (
+    BlockWrapperMode,
+    JointBlockInterface,
+    ModelAdapter,
+    SingleBlockInterface,
+)
+from exo.worker.engines.image.pipeline.block_wrapper import (
+    BlockWrapper,
+    JointBlockWrapper,
+    SingleBlockWrapper,
+)
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.pipeline.runner import DiffusionRunner
+
+__all__ = [
+    "BlockWrapperMode",
+    "BlockWrapper",
+    "DiffusionRunner",
+    "ImagePatchKVCache",
+    "JointBlockInterface",
+    "JointBlockWrapper",
+    "ModelAdapter",
+    "SingleBlockInterface",
+    "SingleBlockWrapper",
+]

src/exo/worker/engines/image/pipeline/adapter.py

@@ -4,8 +4,8 @@ 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 BlockType, ImageModelConfig
-from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.config import BlockType, ImageModelConfig
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
 
 
 class AttentionInterface(Protocol):

src/exo/worker/engines/image/pipeline/block_wrapper.py

@@ -2,14 +2,14 @@ from typing import Any
 
 import mlx.core as mx
 
-from exo.worker.engines.mflux.config.model_config import BlockType
-from exo.worker.engines.mflux.pipefusion.adapter import (
+from exo.worker.engines.image.config import BlockType
+from exo.worker.engines.image.pipeline.adapter import (
     BlockWrapperMode,
     JointBlockInterface,
     ModelAdapter,
     SingleBlockInterface,
 )
-from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
 
 
 class JointBlockWrapper:

src/exo/worker/engines/image/pipeline/runner.py

@@ -8,13 +8,13 @@ from mflux.utils.exceptions import StopImageGenerationException
 from tqdm import tqdm
 
 from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.engines.mflux.config.model_config import ImageModelConfig
-from exo.worker.engines.mflux.pipefusion.adapter import BlockWrapperMode, ModelAdapter
-from exo.worker.engines.mflux.pipefusion.block_wrapper import (
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.pipeline.adapter import BlockWrapperMode, ModelAdapter
+from exo.worker.engines.image.pipeline.block_wrapper import (
     JointBlockWrapper,
     SingleBlockWrapper,
 )
-from exo.worker.engines.mflux.pipefusion.kv_cache import ImagePatchKVCache
+from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache
 
 
 def calculate_patch_heights(latent_height: int, num_patches: int, patch_size: int):

src/exo/worker/engines/mflux/config/__init__.py

@@ -1,47 +0,0 @@
-from exo.worker.engines.mflux.config.flux import FLUX_DEV_CONFIG, FLUX_SCHNELL_CONFIG
-from exo.worker.engines.mflux.config.model_config import (
-    BlockType,
-    ImageModelConfig,
-    TransformerBlockConfig,
-)
-
-__all__ = [
-    "BlockType",
-    "ImageModelConfig",
-    "TransformerBlockConfig",
-    "FLUX_SCHNELL_CONFIG",
-    "FLUX_DEV_CONFIG",
-    "get_config_for_model",
-]
-
-
-# Config registry: maps model ID patterns to configs
-_CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
-    "flux.1-schnell": FLUX_SCHNELL_CONFIG,
-    "flux1-schnell": FLUX_SCHNELL_CONFIG,
-    "schnell": FLUX_SCHNELL_CONFIG,
-    "flux.1-dev": FLUX_DEV_CONFIG,
-    "flux1-dev": FLUX_DEV_CONFIG,
-    "dev": FLUX_DEV_CONFIG,
-}
-
-
-def get_config_for_model(model_id: str) -> ImageModelConfig:
-    """Get configuration for a model ID.
-
-    Args:
-        model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell")
-
-    Returns:
-        The model configuration
-
-    Raises:
-        ValueError: If no configuration found for model ID
-    """
-    model_id_lower = model_id.lower()
-
-    for pattern, config in _CONFIG_REGISTRY.items():
-        if pattern in model_id_lower:
-            return config
-
-    raise ValueError(f"No configuration found for model: {model_id}")

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

@@ -1,35 +0,0 @@
-from pathlib import Path
-from typing import Callable
-
-from exo.worker.engines.mflux.config.model_config import ImageModelConfig
-from exo.worker.engines.mflux.pipefusion.adapter import ModelAdapter
-from exo.worker.engines.mflux.pipefusion.diffusion_runner import DiffusionRunner
-from exo.worker.engines.mflux.pipefusion.flux_adapter import FluxModelAdapter
-
-__all__ = [
-    "create_adapter_for_model",
-    "DiffusionRunner",
-    "ModelAdapter",
-    "FluxModelAdapter",
-]
-
-# Type alias for adapter factory functions
-# Factory takes (config, model_id, local_path, quantize) and returns a ModelAdapter
-AdapterFactory = Callable[[ImageModelConfig, str, Path, int | None], ModelAdapter]
-
-# Registry maps model_family string to adapter factory
-_ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
-    "flux": FluxModelAdapter,
-}
-
-
-def create_adapter_for_model(
-    config: ImageModelConfig,
-    model_id: str,
-    local_path: Path,
-    quantize: int | None = None,
-) -> ModelAdapter:
-    factory = _ADAPTER_REGISTRY.get(config.model_family)
-    if factory is None:
-        raise ValueError(f"No adapter found for model family: {config.model_family}")
-    return factory(config, model_id, local_path, quantize)

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

@@ -1,7 +0,0 @@
-from exo.shared.types.worker.instances import BoundInstance
-from exo.worker.engines.mflux.distributed_model import DistributedImageModel
-
-
-def initialize_mflux(bound_instance: BoundInstance) -> DistributedImageModel:
-    """Initialize DistributedImageModel from a BoundInstance."""
-    return DistributedImageModel.from_bound_instance(bound_instance)

src/exo/worker/runner/runner.py

@@ -49,7 +49,7 @@ from exo.worker.engines.image import (
     generate_image,
     warmup_image_generator,
 )
-from exo.worker.engines.mflux.utils_mflux import initialize_mflux
+from exo.worker.engines.image import initialize_image_model
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
@@ -138,7 +138,7 @@ def main(
                             ModelTask.TextToImage in model_tasks
                             or ModelTask.ImageToImage in model_tasks
                         ):
-                            model = initialize_mflux(bound_instance)
+                            model = initialize_image_model(bound_instance)
                         else:
                             raise ValueError(
                                 f"Unknown model task(s): {model_card.tasks}"