"""
This type stub file was generated by pyright.
"""

import mlx.core as mx
from functools import lru_cache
from typing import Literal

class ModelConfig:
    precision: mx.Dtype = ...
    def __init__(
        self,
        priority: int,
        aliases: list[str],
        model_name: str,
        base_model: str | None,
        controlnet_model: str | None,
        custom_transformer_model: str | None,
        num_train_steps: int | None,
        max_sequence_length: int | None,
        supports_guidance: bool | None,
        requires_sigma_shift: bool | None,
        transformer_overrides: dict | None = ...,
    ) -> None: ...
    @staticmethod
    @lru_cache
    def dev() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def schnell() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_kontext() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_fill() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_redux() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_depth() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_controlnet_canny() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def schnell_controlnet_canny() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_controlnet_upscaler() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def dev_fill_catvton() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def krea_dev() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def flux2_klein_4b() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def flux2_klein_9b() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def qwen_image() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def qwen_image_edit() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def fibo() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def z_image_turbo() -> ModelConfig: ...
    @staticmethod
    @lru_cache
    def seedvr2_3b() -> ModelConfig: ...
    def x_embedder_input_dim(self) -> int: ...
    def is_canny(self) -> bool: ...
    @staticmethod
    def from_name(
        model_name: str, base_model: Literal["dev", "schnell", "krea-dev"] | None = ...
    ) -> ModelConfig: ...

AVAILABLE_MODELS = ...
