feat: add bulkgen models command listing available models with capabilities

2026-02-14 23:06:17 +01:00 · 2026-02-14 23:06:17 +01:00 · b536ff9d79
commit b536ff9d79
parent 47b53db760
2 changed files with 160 additions and 0 deletions
--- a/bulkgen/models.py
+++ b/bulkgen/models.py
@ -0,0 +1,109 @@
+"""Registry of supported models and their capabilities."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class ModelInfo:
+    """Describes a supported model and its capabilities."""
+
+    name: str
+    provider: str
+    type: Literal["text", "image"]
+    capabilities: list[str]
+
+
+TEXT_MODELS: list[ModelInfo] = [
+    ModelInfo(
+        name="mistral-large-latest",
+        provider="Mistral",
+        type="text",
+        capabilities=["text generation"],
+    ),
+    ModelInfo(
+        name="mistral-small-latest",
+        provider="Mistral",
+        type="text",
+        capabilities=["text generation"],
+    ),
+    ModelInfo(
+        name="pixtral-large-latest",
+        provider="Mistral",
+        type="text",
+        capabilities=["text generation", "vision"],
+    ),
+    ModelInfo(
+        name="pixtral-12b-latest",
+        provider="Mistral",
+        type="text",
+        capabilities=["text generation", "vision"],
+    ),
+]
+
+IMAGE_MODELS: list[ModelInfo] = [
+    ModelInfo(
+        name="flux-dev",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image"],
+    ),
+    ModelInfo(
+        name="flux-pro",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.1",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.1-ultra",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "high resolution"],
+    ),
+    ModelInfo(
+        name="flux-2-pro",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "reference images"],
+    ),
+    ModelInfo(
+        name="flux-kontext-pro",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "reference images"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.0-canny",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "control images", "edge detection"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.0-depth",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "control images", "depth map"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.0-fill",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "inpainting"],
+    ),
+    ModelInfo(
+        name="flux-pro-1.0-expand",
+        provider="BlackForestLabs",
+        type="image",
+        capabilities=["text-to-image", "outpainting"],
+    ),
+]
+
+ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS