Better llm model type check

src/exo/master/main.py

@@ -156,7 +156,6 @@ class Master:
 
                             self.command_task_mapping[command.command_id] = task_id
                         case ImageGeneration():
-                            # TODO(ciaran): refactor with ChatCompletion
                             instance_task_counts: dict[InstanceId, int] = {}
                             for instance in self.state.instances.values():
                                 if (
@@ -200,7 +199,6 @@ class Master:
 
                             self.command_task_mapping[command.command_id] = task_id
                         case ImageEdits():
-                            # TODO(ciaran): refactor with ChatCompletion
                             instance_task_counts: dict[InstanceId, int] = {}
                             for instance in self.state.instances.values():
                                 if (

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

@@ -42,7 +42,7 @@ def generate_image(
     # Parse parameters
     width, height = parse_size(task.size)
     quality: Literal["low", "medium", "high"] = task.quality or "medium"
-    seed = 2  # TODO(ciaran): Consider adding seed to ImageGenerationTaskParams
+    seed = 2  # TODO(ciaran): Randomise when not testing anymore
 
     # Generate using the model's generate method
     image = model.generate(

src/exo/worker/runner/runner.py

@@ -50,6 +50,7 @@ from exo.worker.engines.image import (
     initialize_image_model,
     warmup_image_generator,
 )
+from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
@@ -159,8 +160,7 @@ def main(
 
                         logger.info(f"warming up inference for instance: {instance}")
                         if ModelTask.TextGeneration in model_tasks:
-                            # assert isinstance(model, Model) TODO(ciaran): not actually Model
-                            assert model and not isinstance(model, ImageGenerator)
+                            assert model and isinstance(model, Model)
                             assert tokenizer
                             assert sampler
 
@@ -192,8 +192,7 @@ def main(
                     case ChatCompletion(
                         task_params=task_params, command_id=command_id
                     ) if isinstance(current_status, RunnerReady):
-                        # assert isinstance(model, Model) TODO(ciaran): not actually Model
-                        assert model and not isinstance(model, ImageGenerator)
+                        assert model and isinstance(model, Model)
                         assert tokenizer
                         assert sampler
                         logger.info(f"received chat request: {str(task)[:500]}")