refactor: use StrEnum for model capabilities instead of plain strings

bulkgen/providers/models.py

@@ -3,9 +3,25 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from enum import StrEnum
 from typing import Literal
 
 
+class Capability(StrEnum):
+    """Capabilities a model may support."""
+
+    TEXT_GENERATION = "text generation"
+    VISION = "vision"
+    TEXT_TO_IMAGE = "text-to-image"
+    HIGH_RESOLUTION = "high resolution"
+    REFERENCE_IMAGES = "reference images"
+    CONTROL_IMAGES = "control images"
+    EDGE_DETECTION = "edge detection"
+    DEPTH_MAP = "depth map"
+    INPAINTING = "inpainting"
+    OUTPAINTING = "outpainting"
+
+
 @dataclass(frozen=True)
 class ModelInfo:
     """Describes a supported model and its capabilities."""
@@ -13,7 +29,7 @@ class ModelInfo:
     name: str
     provider: str
     type: Literal["text", "image"]
-    capabilities: list[str]
+    capabilities: list[Capability]
 
 
 TEXT_MODELS: list[ModelInfo] = [
@@ -21,25 +37,25 @@ TEXT_MODELS: list[ModelInfo] = [
         name="mistral-large-latest",
         provider="Mistral",
         type="text",
-        capabilities=["text generation"],
+        capabilities=[Capability.TEXT_GENERATION],
     ),
     ModelInfo(
         name="mistral-small-latest",
         provider="Mistral",
         type="text",
-        capabilities=["text generation"],
+        capabilities=[Capability.TEXT_GENERATION],
     ),
     ModelInfo(
         name="pixtral-large-latest",
         provider="Mistral",
         type="text",
-        capabilities=["text generation", "vision"],
+        capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
     ),
     ModelInfo(
         name="pixtral-12b-latest",
         provider="Mistral",
         type="text",
-        capabilities=["text generation", "vision"],
+        capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
     ),
 ]
 
@@ -48,61 +64,69 @@ IMAGE_MODELS: list[ModelInfo] = [
         name="flux-dev",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image"],
+        capabilities=[Capability.TEXT_TO_IMAGE],
     ),
     ModelInfo(
         name="flux-pro",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image"],
+        capabilities=[Capability.TEXT_TO_IMAGE],
     ),
     ModelInfo(
         name="flux-pro-1.1",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image"],
+        capabilities=[Capability.TEXT_TO_IMAGE],
     ),
     ModelInfo(
         name="flux-pro-1.1-ultra",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "high resolution"],
+        capabilities=[Capability.TEXT_TO_IMAGE, Capability.HIGH_RESOLUTION],
     ),
     ModelInfo(
         name="flux-2-pro",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "reference images"],
+        capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
     ),
     ModelInfo(
         name="flux-kontext-pro",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "reference images"],
+        capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
     ),
     ModelInfo(
         name="flux-pro-1.0-canny",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "control images", "edge detection"],
+        capabilities=[
+            Capability.TEXT_TO_IMAGE,
+            Capability.CONTROL_IMAGES,
+            Capability.EDGE_DETECTION,
+        ],
     ),
     ModelInfo(
         name="flux-pro-1.0-depth",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "control images", "depth map"],
+        capabilities=[
+            Capability.TEXT_TO_IMAGE,
+            Capability.CONTROL_IMAGES,
+            Capability.DEPTH_MAP,
+        ],
     ),
     ModelInfo(
         name="flux-pro-1.0-fill",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "inpainting"],
+        capabilities=[Capability.TEXT_TO_IMAGE, Capability.INPAINTING],
     ),
     ModelInfo(
         name="flux-pro-1.0-expand",
         provider="BlackForestLabs",
         type="image",
-        capabilities=["text-to-image", "outpainting"],
+        capabilities=[Capability.TEXT_TO_IMAGE, Capability.OUTPAINTING],
     ),
 ]