`.typings/mflux/models/common/weights/loading/weight_definition.pyi`

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

import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, TYPE_CHECKING, TypeAlias
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.depth_pro.weights.depth_pro_weight_definition import (
    DepthProWeightDefinition,
)
from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition
from mflux.models.fibo_vlm.weights.fibo_vlm_weight_definition import (
    FIBOVLMWeightDefinition,
)
from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition
from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition
from mflux.models.seedvr2.weights.seedvr2_weight_definition import (
    SeedVR2WeightDefinition,
)
from mflux.models.z_image.weights.z_image_weight_definition import (
    ZImageWeightDefinition,
)

"""
This type stub file was generated by pyright.
"""
if TYPE_CHECKING:
    WeightDefinitionType: TypeAlias = type[
        FluxWeightDefinition
        | FIBOWeightDefinition
        | FIBOVLMWeightDefinition
        | QwenWeightDefinition
        | ZImageWeightDefinition
        | SeedVR2WeightDefinition
        | DepthProWeightDefinition
    ]

@dataclass
class ComponentDefinition:
    name: str
    hf_subdir: str
    mapping_getter: Callable[[], List[WeightTarget]] | None = ...
    model_attr: str | None = ...
    num_blocks: int | None = ...
    num_layers: int | None = ...
    loading_mode: str = ...
    precision: mx.Dtype | None = ...
    skip_quantization: bool = ...
    bulk_transform: Callable[[mx.array], mx.array] | None = ...
    weight_subkey: str | None = ...
    download_url: str | None = ...
    weight_prefix_filters: List[str] | None = ...
    weight_files: List[str] | None = ...

@dataclass
class TokenizerDefinition:
    name: str
    hf_subdir: str
    tokenizer_class: str = ...
    fallback_subdirs: List[str] | None = ...
    download_patterns: List[str] | None = ...
    encoder_class: type[BaseTokenizer] | None = ...
    max_length: int = ...
    padding: str = ...
    template: str | None = ...
    use_chat_template: bool = ...
    chat_template_kwargs: dict | None = ...
    add_special_tokens: bool = ...
    processor_class: type | None = ...
    image_token: str = ...
    chat_template: str | None = ...