60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Callable, List, TypeAlias
|
|
|
|
import mlx.core as mx
|
|
|
|
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
|
|
|
|
if TYPE_CHECKING:
|
|
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.z_image.weights.z_image_weight_definition import ZImageWeightDefinition
|
|
|
|
WeightDefinitionType: TypeAlias = type[
|
|
FluxWeightDefinition
|
|
| FIBOWeightDefinition
|
|
| FIBOVLMWeightDefinition
|
|
| QwenWeightDefinition
|
|
| ZImageWeightDefinition
|
|
| DepthProWeightDefinition
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class ComponentDefinition:
|
|
name: str
|
|
hf_subdir: str
|
|
mapping_getter: Callable[[], List[WeightTarget]] | None = None
|
|
model_attr: str | None = None
|
|
num_blocks: int | None = None
|
|
num_layers: int | None = None
|
|
loading_mode: str = "mlx_native"
|
|
precision: mx.Dtype | None = None
|
|
skip_quantization: bool = False
|
|
bulk_transform: Callable[[mx.array], mx.array] | None = None
|
|
weight_subkey: str | None = None
|
|
download_url: str | None = None
|
|
weight_prefix_filters: List[str] | None = None
|
|
|
|
|
|
@dataclass
|
|
class TokenizerDefinition:
|
|
name: str
|
|
hf_subdir: str
|
|
tokenizer_class: str = "AutoTokenizer"
|
|
fallback_subdirs: List[str] | None = None
|
|
download_patterns: List[str] | None = None
|
|
encoder_class: type["BaseTokenizer"] | None = None
|
|
max_length: int = 512
|
|
padding: str = "max_length"
|
|
template: str | None = None
|
|
use_chat_template: bool = False
|
|
chat_template_kwargs: dict | None = field(default_factory=dict)
|
|
add_special_tokens: bool = True
|
|
processor_class: type | None = None
|
|
image_token: str = "<|image_pad|>"
|