diff --git a/bulkgen/cli.py b/bulkgen/cli.py index 550cc06..aa6c6a8 100644 --- a/bulkgen/cli.py +++ b/bulkgen/cli.py @@ -13,6 +13,7 @@ import typer from bulkgen.builder import BuildEvent, BuildResult, run_build from bulkgen.config import ProjectConfig, load_config from bulkgen.graph import build_graph, get_build_order +from bulkgen.models import ALL_MODELS app = typer.Typer(name="bulkgen", help="AI artifact build tool.") @@ -176,3 +177,53 @@ def graph() -> None: if preds: arrow = click.style(" <- ", dim=True) click.echo(f" {node}{arrow}{', '.join(preds)}") + + +@app.command() +def models() -> None: + """List available models and their capabilities.""" + name_width = max(len(m.name) for m in ALL_MODELS) + provider_width = max(len(m.provider) for m in ALL_MODELS) + type_width = max(len(m.type) for m in ALL_MODELS) + + header_name = "Model".ljust(name_width) + header_provider = "Provider".ljust(provider_width) + header_type = "Type".ljust(type_width) + header_caps = "Capabilities" + + click.echo( + click.style(header_name, bold=True) + + " " + + click.style(header_provider, bold=True) + + " " + + click.style(header_type, bold=True) + + " " + + click.style(header_caps, bold=True) + ) + click.echo( + "─" * name_width + + " " + + "─" * provider_width + + " " + + "─" * type_width + + " " + + "─" * len(header_caps) + ) + + for model in ALL_MODELS: + name_col = model.name.ljust(name_width) + provider_col = model.provider.ljust(provider_width) + type_col = model.type.ljust(type_width) + caps_col = ", ".join(model.capabilities) + + type_color = "green" if model.type == "image" else "cyan" + + click.echo( + click.style(name_col, fg="white", bold=True) + + " " + + provider_col + + " " + + click.style(type_col, fg=type_color) + + " " + + click.style(caps_col, dim=True) + ) diff --git a/bulkgen/models.py b/bulkgen/models.py new file mode 100644 index 0000000..667c7fd --- /dev/null +++ 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