130 lines
3.2 KiB
Python
130 lines
3.2 KiB
Python
"""Registry of supported models and their capabilities."""
|
|
|
|
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"
|
|
REFERENCE_IMAGES = "reference images"
|
|
CONTROL_IMAGES = "control images"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelInfo:
|
|
"""Describes a supported model and its capabilities."""
|
|
|
|
name: str
|
|
provider: str
|
|
type: Literal["text", "image"]
|
|
capabilities: list[Capability]
|
|
|
|
|
|
TEXT_MODELS: list[ModelInfo] = [
|
|
ModelInfo(
|
|
name="mistral-large-latest",
|
|
provider="Mistral",
|
|
type="text",
|
|
capabilities=[Capability.TEXT_GENERATION],
|
|
),
|
|
ModelInfo(
|
|
name="mistral-small-latest",
|
|
provider="Mistral",
|
|
type="text",
|
|
capabilities=[Capability.TEXT_GENERATION],
|
|
),
|
|
ModelInfo(
|
|
name="pixtral-large-latest",
|
|
provider="Mistral",
|
|
type="text",
|
|
capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
|
|
),
|
|
ModelInfo(
|
|
name="pixtral-12b-latest",
|
|
provider="Mistral",
|
|
type="text",
|
|
capabilities=[Capability.TEXT_GENERATION, Capability.VISION],
|
|
),
|
|
]
|
|
|
|
IMAGE_MODELS: list[ModelInfo] = [
|
|
ModelInfo(
|
|
name="flux-dev",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.1",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.1-ultra",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE],
|
|
),
|
|
ModelInfo(
|
|
name="flux-2-pro",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
|
|
),
|
|
ModelInfo(
|
|
name="flux-kontext-pro",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[Capability.TEXT_TO_IMAGE, Capability.REFERENCE_IMAGES],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.0-canny",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[
|
|
Capability.TEXT_TO_IMAGE,
|
|
Capability.CONTROL_IMAGES,
|
|
],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.0-depth",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[
|
|
Capability.TEXT_TO_IMAGE,
|
|
Capability.CONTROL_IMAGES,
|
|
],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.0-fill",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[
|
|
Capability.TEXT_TO_IMAGE,
|
|
],
|
|
),
|
|
ModelInfo(
|
|
name="flux-pro-1.0-expand",
|
|
provider="BlackForestLabs",
|
|
type="image",
|
|
capabilities=[
|
|
Capability.TEXT_TO_IMAGE,
|
|
],
|
|
),
|
|
]
|
|
|
|
ALL_MODELS: list[ModelInfo] = TEXT_MODELS + IMAGE_MODELS
|