kfickel/hokusai
1
0
Fork 0
Code Issues 1 Pull requests 1 Projects Releases Packages Wiki Activity Actions

refactor: use StrEnum for model capabilities instead of plain strings

Browse source
This commit is contained in:
Konstantin Fickel 2026-02-15 08:04:28 +01:00
parent 6a80cfb78e
commit 8e3ed7010f
Signed by: kfickel
GPG key ID: A793722F9933C1A5
1 changed files with 39 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

54
bulkgen/providers/models.py
View file

@ -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],
),
]