109 lines
2.7 KiB
Python
109 lines
2.7 KiB
Python
"""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
|