Implement model factory

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

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, Literal, Optional
 import mlx.core as mx
 from mflux.callbacks.callbacks import Callbacks
 from mflux.config.config import Config
-from mflux.config.model_config import ModelConfig
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator
 from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
@@ -21,7 +20,7 @@ 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 get_adapter_for_model
+from exo.worker.engines.mflux.pipefusion import create_model, get_adapter_for_model
 from exo.worker.engines.mflux.pipefusion.adapter import ModelAdapter
 from exo.worker.engines.mflux.pipefusion.distributed_denoising import (
     DistributedDenoising,
@@ -69,17 +68,8 @@ class DistributedImageModel:
         config = get_config_for_model(model_id)
         adapter = get_adapter_for_model(config)
 
-        # Create the appropriate mflux model based on family
-        if config.model_family == "flux":
-            model = Flux1(
-                model_config=ModelConfig.from_name(
-                    model_name=model_id, base_model=None
-                ),
-                local_path=str(local_path),
-                quantize=quantize,
-            )
-        else:
-            raise ValueError(f"Unsupported model family: {config.model_family}")
+        # Create the model using the factory registry
+        model = create_model(config, model_id, local_path, quantize)
 
         if group is not None:
             # Apply pipeline parallelism by wrapping the transformer

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

@@ -1,11 +1,16 @@
 """
-Adapter registry for model-specific operations.
+Adapter and model factory registries.
 
-This module provides a registry pattern for managing model adapters,
-allowing new model families to be added without modifying core code.
+This module provides registry patterns for managing model adapters and
+model factories, allowing new model families to be added without modifying
+core code.
 """
 
-from typing import Callable
+from pathlib import Path
+from typing import Any, Callable
+
+from mflux.config.model_config import ModelConfig
+from mflux.models.flux.variants.txt2img.flux import Flux1
 
 from exo.worker.engines.mflux.config.model_config import ImageModelConfig
 from exo.worker.engines.mflux.pipefusion.adapter import ModelAdapter
@@ -46,3 +51,63 @@ def register_adapter(model_family: str, factory: AdapterFactory) -> None:
         factory: A callable that takes an ImageModelConfig and returns a ModelAdapter
     """
     _ADAPTER_REGISTRY[model_family] = factory
+
+
+# =============================================================================
+# Model Factory Registry
+# =============================================================================
+
+# Type alias for model factory functions
+# Takes (model_id, local_path, quantize) and returns a model instance
+ModelFactory = Callable[[str, Path, int | None], Any]
+
+
+def _create_flux_model(model_id: str, local_path: Path, quantize: int | None) -> Flux1:
+    """Create a Flux1 model instance."""
+    return Flux1(
+        model_config=ModelConfig.from_name(model_name=model_id, base_model=None),
+        local_path=str(local_path),
+        quantize=quantize,
+    )
+
+
+# Registry maps model_family string to model factory
+_MODEL_REGISTRY: dict[str, ModelFactory] = {
+    "flux": _create_flux_model,
+}
+
+
+def create_model(
+    config: ImageModelConfig,
+    model_id: str,
+    local_path: Path,
+    quantize: int | None = None,
+) -> Any:
+    """Create a model instance for a model configuration.
+
+    Args:
+        config: The model configuration
+        model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell")
+        local_path: Path to the local model weights
+        quantize: Optional quantization bit width
+
+    Returns:
+        A model instance for the model family
+
+    Raises:
+        ValueError: If no factory is registered for the model family
+    """
+    factory = _MODEL_REGISTRY.get(config.model_family)
+    if factory is None:
+        raise ValueError(f"No model factory found for model family: {config.model_family}")
+    return factory(model_id, local_path, quantize)
+
+
+def register_model_factory(model_family: str, factory: ModelFactory) -> None:
+    """Register a new model factory for a model family.
+
+    Args:
+        model_family: The model family identifier (e.g., "flux", "fibo", "qwen")
+        factory: A callable that takes (model_id, local_path, quantize) and returns a model
+    """
+    _MODEL_REGISTRY[model_family] = factory