diff --git a/src/exo/worker/engines/mflux/distributed_model.py b/src/exo/worker/engines/mflux/distributed_model.py index 5cd9f13a..5c90f50f 100644 --- a/src/exo/worker/engines/mflux/distributed_model.py +++ b/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 diff --git a/src/exo/worker/engines/mflux/pipefusion/__init__.py b/src/exo/worker/engines/mflux/pipefusion/__init__.py index c47b6b0c..53e6cbdf 100644 --- a/src/exo/worker/engines/mflux/pipefusion/__init__.py +++ b/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