Compare commits

..

17 Commits

Author SHA1 Message Date
dmcc73 37ad1fb3ed Call warmup_speculative at startup to pre-compile LpB kernels
The warmup_speculative() function was defined but never called.
Custom Metal kernels (LpB) require first-call compilation (~200ms).
Without warmup, the first speculative cycle is slow, dragging down
average TPS by 10-20% on short generations.

In mlx_bench testing: cold 48 TPS → warm 60 TPS for DFlash,
cold 39 TPS → warm 44 TPS for MTP.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-01 23:15:05 +01:00
dmcc73 b47a287f3e Add EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction
For profiling: extract_top_logprobs() does 11 .item() calls +
argpartition on 248K vocab per token. Testing if this is the
source of speculative overhead vs mlx_bench.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 23:41:47 +01:00
dmcc73 e1cf376e45 Add speculative warmup: compile MTP + verify kernels at startup
The standard warmup only runs S=1 generation, leaving speculative
kernels (S>1 verify, speculative GDN kernel, MTP draft) uncompiled.
First real speculative cycle had compilation overhead.

New warmup_speculative(): prefills a short prompt, runs 3 speculative
cycles to compile all kernels before real requests arrive.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 22:39:34 +01:00
dmcc73 75932cbcca Fix stop token dropping valid tokens before it
When <|im_end|> appeared in accepted drafts, all preceding tokens in
the cycle were returned with finish_reason="stop", causing exo to
drop them (exo skips adding tokens with finish_reason="stop").

Symptom: γ=0 outputs "20", γ=1 outputs "2", γ=2 outputs nothing —
losing exactly γ tokens at the end.

Fix: yield tokens before the stop normally (no finish_reason), buffer
the stop token, let _yield_buffered return it with finish_reason="stop".

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:44:01 +01:00
dmcc73 199a4ab7e0 EXO_SPECULATIVE_TEMP overrides model sampling temperature globally
When set, overrides the request's temperature for both the model's
sampler AND the speculative acceptance. This allows testing greedy
baseline (γ=0) and greedy speculative (γ=2) with the same T=0.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:24:25 +01:00
dmcc73 4818b9a3db EXO_SPECULATIVE_TEMP overrides request temp when set
If EXO_SPECULATIVE_TEMP is explicitly set, use it (for testing greedy).
If not set, use the request's temperature (production behavior).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:23:05 +01:00
dmcc73 f0433505a8 Fix speculative temp: use request temperature, not global env var
The speculative cycle was using EXO_SPECULATIVE_TEMP (global) instead
of the request's actual temperature. This caused greedy decoding in
speculative while the model sampled at T=0.7, producing different
(shorter) output and missing responses after </think>.

Now passes task_params.temperature from submit() to MTPBatchGenerator
per-request via _request_temp[uid]. Falls back to self.temp (env var)
if not set.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:21:57 +01:00
dmcc73 88bc1656a2 Fix MTP prefill to use all captured positions
Was using prompt_pre_norm[:, :-1, :] (missing last position).
Now uses full prompt_pre_norm paired with all_prompt_tokens[1:S_pre+1],
matching the mlx_bench MTPBatchGenerator's prefill behavior.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 18:06:53 +01:00
dmcc73 7bd1ba6605 Fix MTP prefill: do it in submit() with correct prompt tokens
Bug: _CapturingEmbed was overwritten by BatchGenerator's 2-token insert,
causing MTP prefill to silently skip (len check failed: 2 < N-1). MTP
drafted without any prompt context → low acceptance → low TPS.

Fix: Do MTP prefill in ExoBatchGenerator.submit() right after main model
prefill, using all_prompt_tokens (available as local variable). Remove
_CapturingEmbed entirely. Simplify _first_step_and_prefill to just
capture decode pre_norm.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 17:42:20 +01:00
dmcc73 e7c5d56e83 Add LpB kernel patches for Qwen3.5 dense models (27B, 9B)
Loop-over-B custom GEMV kernels for expanding projections (N > K):
gate_proj, up_proj, down_proj, in_proj_qkv, in_proj_z, out_proj, q_proj.

These reduce S>1 verification cost from ~7ms/token to ~3ms/token,
critical for speculative decoding speedup.

Auto-detected for model_type=qwen3_5 (dense models like 27B, 9B).
MoE models (qwen3_5_moe) use the existing batched fused patches instead.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 17:00:02 +01:00
dmcc73 dd71182457 Fix MTP prefill for exo: capture prompt tokens via embed_tokens wrapper
Exo does its own prefill outside BatchGenerator, so batch.tokens only
has the last 2 tokens. Added _CapturingEmbed wrapper on embed_tokens to
capture the full prompt token ids during prefill. MTP prefill now uses
these captured tokens instead of batch.tokens.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 15:57:25 +01:00
dmcc73 09012d3799 Auto-extract MTP weights from HuggingFace model repo
When EXO_SPECULATIVE=1, MTP weights are resolved in order:
1. EXO_MTP_WEIGHTS=/path/to/file (explicit path)
2. EXO_MTP_MODEL=Qwen/Qwen3.5-27B (explicit HF repo)
3. Auto-detect: if model has mtp_num_hidden_layers > 0 and is
   Qwen3.5, defaults to Qwen/Qwen3.5-27B

Downloads safetensors from HF, extracts model.mtp.* tensors,
caches to ~/.cache/exo/mtp_weights/ for future use.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 15:19:48 +01:00
dmcc73 ce19267d2d Pass temperature and alpha to MTP speculative decoding
Default temp=0.7 (matching exo's default) so probabilistic acceptance
runs correctly. Configurable via EXO_SPECULATIVE_TEMP and
EXO_SPECULATIVE_ALPHA env vars.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 15:12:59 +01:00
dmcc73 8a65a51569 Add MTP speculative decoding for Qwen3.5 models
Integrates MTP-based speculative decoding into exo's BatchGenerator.
When enabled via EXO_SPECULATIVE=1 and EXO_MTP_WEIGHTS=<path>,
MTPBatchGenerator replaces the standard MlxBatchGenerator for BS=1
inference, drafting γ tokens with the model's built-in MTP head and
verifying at S=γ+1.

New files in speculative/:
- mtp_module.py: MTPPredictor + speculative_forward (kernel swap for
  GDN rollback) + draft_tokens (lazy MTP chaining)
- mtp_batch_generator.py: MTPBatchGenerator subclassing mlx_lm's
  BatchGenerator with token buffering and BS>1 fallback
- speculative_cache.py: SpeculativeArraysCache for GDN state rollback
- speculative_gdn_kernel.py: Metal kernel with per-step state output

Environment variables:
  EXO_SPECULATIVE=1              Enable speculative decoding
  EXO_MTP_WEIGHTS=/path/to/file  Path to MTP weights safetensors
  EXO_SPECULATIVE_GAMMA=2        Draft tokens per cycle (default: 2)

MTP weights must be extracted from the original HF model (e.g.
Qwen/Qwen3.5-27B) as they are stripped during MLX quantization.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-30 15:10:11 +01:00
dmcc73 a2de281c67 Replace GDN projections with register-sharing batched kernel
Old kernel used grid z=B, loading weights B times independently.
New kernel loads weights once into registers and computes B dot products.
11-14% faster at B=2-4 in full model benchmarks (194 vs 174 TPS at B=2).
B=1 generates identical code, no regression.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-19 15:13:39 +00:00
dmcc73 9394d04f5f Add LCB TPS benchmark script
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-18 17:37:36 +00:00
dmcc73 92c04b0aa5 Add batched fused Metal kernel patches for Qwen3.5 MoE decode
Custom Metal kernels with register-level weight sharing for batch sizes 1-8.
Fuses o_proj + RMSNorm + gate GEMV + softmax + topk + SwiGLU + down_proj + epilogue
into 4 dispatches per MoE layer, plus fused GDN and GQA attention projections.
Falls back to vanilla for B>8 or S>1 (prefill). Controlled by EXO_FUSED_KERNELS env var.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-18 17:36:28 +00:00
499 changed files with 9593 additions and 10302 deletions
-20
View File
@@ -1,20 +0,0 @@
from enum import Enum
class HarmonyEncodingName(Enum):
HARMONY_GPT_OSS = ...
class HarmonyEncoding: ...
class HarmonyError(Exception): ...
class Role(Enum):
ASSISTANT = ...
class StreamableParser:
last_content_delta: str
current_channel: str | None
current_recipient: str | None
def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
def process(self, token_id: int) -> None: ...
def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...
-17
View File
@@ -1,17 +0,0 @@
class NvmlMemoryInfo:
used: int
total: int
free: int
class NvmlUtilizationRates:
gpu: int
memory: int
def nvmlInit() -> None: ...
def nvmlShutdown() -> None: ...
def nvmlDeviceGetCount() -> int: ...
def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...
-61
View File
@@ -1,61 +0,0 @@
from typing import Any, Sequence
from torch import backends as backends
from torch import cuda as cuda
from torch import distributed as distributed
__version__: str
class version:
cuda: str
class dtype: ...
bfloat16: dtype
float16: dtype
float32: dtype
int8: dtype
int32: dtype
int64: dtype
long: dtype
float8_e4m3fn: dtype
class Tensor:
shape: Sequence[int]
dtype: dtype
def __getitem__(self, key: Any) -> Tensor: ...
def __setitem__(self, key: Any, value: Any) -> None: ...
def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
def cpu(self) -> Tensor: ...
def detach(self) -> Tensor: ...
def clone(self) -> Tensor: ...
def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
def view(self, *shape: Any) -> Tensor: ...
def squeeze(self, dim: int = ...) -> Tensor: ...
def unsqueeze(self, dim: int) -> Tensor: ...
def permute(self, *dims: int) -> Tensor: ...
def float(self) -> Tensor: ...
def numpy(self) -> Any: ...
def numel(self) -> int: ...
def nelement(self) -> int: ...
@property
def is_cuda(self) -> bool: ...
@property
def device(self) -> device: ...
def __len__(self) -> int: ...
def data_ptr(self) -> int: ...
def tolist(self) -> Any: ...
def abs(self) -> Tensor: ...
def max(self) -> Tensor: ...
def mean(self) -> Tensor: ...
def sum(self, dim: int = ...) -> Tensor: ...
def item(self) -> float: ...
def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def from_numpy(ndarray: Any) -> Tensor: ...
def inference_mode() -> Any: ...
class device:
def __init__(self, type: str, index: int = ...) -> None: ...
@@ -1 +0,0 @@
from torch.backends import cuda as cuda
@@ -1 +0,0 @@
def is_built() -> bool: ...
-10
View File
@@ -1,10 +0,0 @@
class _DeviceProperties:
total_memory: int
def is_available() -> bool: ...
def get_device_name(device: int) -> str: ...
def get_device_properties(device: int) -> _DeviceProperties: ...
def empty_cache() -> None: ...
def mem_get_info() -> tuple[int, int]: ...
def synchronize() -> None: ...
def max_memory_allocated() -> int: ...
@@ -1,2 +0,0 @@
def is_initialized() -> bool: ...
def destroy_process_group() -> None: ...
-1
View File
@@ -1 +0,0 @@
__version__: str
-2
View File
@@ -1,2 +0,0 @@
class ModelConfig:
max_model_len: int
-18
View File
@@ -1,18 +0,0 @@
from dataclasses import dataclass
@dataclass
class EngineArgs:
model: str = ...
served_model_name: str | list[str] | None = ...
tokenizer: str | None = ...
trust_remote_code: bool = ...
dtype: str = ...
seed: int = ...
max_model_len: int | None = ...
gpu_memory_utilization: float = ...
enforce_eager: bool = ...
tensor_parallel_size: int = ...
pipeline_parallel_size: int = ...
quantization: str | None = ...
load_format: str = ...
enable_sleep_mode: bool = ...
-17
View File
@@ -1,17 +0,0 @@
class CompletionOutput:
index: int
text: str
token_ids: list[int]
cumulative_logprob: float | None
logprobs: object | None
finish_reason: str | None
stop_reason: int | str | None
def finished(self) -> bool: ...
class RequestOutput:
request_id: str
prompt: str | None
prompt_token_ids: list[int] | None
outputs: list[CompletionOutput]
finished: bool
-11
View File
@@ -1,11 +0,0 @@
class SamplingParams:
n: int
temperature: float
top_p: float
top_k: int
min_p: float
seed: int | None
stop: str | list[str] | None
max_tokens: int | None
logprobs: int | None
repetition_penalty: float
@@ -1,3 +0,0 @@
from vllm.tokenizers.protocol import TokenizerLike
__all__ = ["TokenizerLike"]
@@ -1,15 +0,0 @@
from typing import Protocol
class TokenizerLike(Protocol):
@property
def eos_token_id(self) -> int: ...
@property
def vocab_size(self) -> int: ...
def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
def apply_chat_template(
self,
messages: list[dict[str, str]],
tools: list[dict[str, object]] | None = ...,
**kwargs: object,
) -> str | list[int]: ...
-1
View File
@@ -1 +0,0 @@
-1
View File
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
from collections.abc import Sequence
from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
from vllm.v1.kv_cache_interface import KVCacheConfig
class KVCacheBlocks:
blocks: tuple[Sequence[KVCacheBlock], ...]
def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
def get_block_ids(self) -> tuple[list[int], ...]: ...
class KVCacheManager:
block_pool: BlockPool
kv_cache_config: KVCacheConfig
enable_caching: bool
num_kv_cache_groups: int
coordinator: object
def __init__(self, *args: object, **kwargs: object) -> None: ...
def allocate_slots(
self, request: object, num_new_tokens: int, *args: object, **kwargs: object
) -> KVCacheBlocks | None: ...
def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks: ...
@@ -1,16 +0,0 @@
class KVCacheBlock:
block_id: int
ref_cnt: int
def __init__(self, block_id: int) -> None: ...
class FreeKVCacheBlockQueue:
def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
class BlockPool:
blocks: list[KVCacheBlock]
free_block_queue: FreeKVCacheBlockQueue
num_gpu_blocks: int
enable_caching: bool
def get_num_free_blocks(self) -> int: ...
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...
@@ -1,22 +0,0 @@
from vllm.config import ModelConfig
from vllm.engine.arg_utils import EngineArgs
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.tokenizers import TokenizerLike
class LLMEngine:
tokenizer: TokenizerLike | None
model_config: ModelConfig
@classmethod
def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
def add_request(
self,
request_id: str,
prompt: str,
params: SamplingParams,
arrival_time: float | None = ...,
) -> None: ...
def step(self) -> list[RequestOutput]: ...
def has_unfinished_requests(self) -> bool: ...
def get_tokenizer(self) -> TokenizerLike: ...
@@ -1,23 +0,0 @@
from dataclasses import dataclass
@dataclass
class KVCacheSpec:
block_size: int
num_kv_heads: int
head_size: int
@dataclass
class KVCacheGroupSpec:
layer_names: list[str]
kv_cache_spec: KVCacheSpec
@dataclass
class KVCacheTensorSpec:
shared_by: list[str]
size: int
@dataclass
class KVCacheConfig:
num_blocks: int
kv_cache_groups: list[KVCacheGroupSpec]
kv_cache_tensors: list[KVCacheTensorSpec]
-6
View File
@@ -1,6 +0,0 @@
class Request:
request_id: str
prompt_token_ids: list[int] | None
num_prompt_tokens: int
num_computed_tokens: int
num_tokens: int
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
import torch
class _CompilationConfig:
static_forward_context: dict[str, object]
class _ModelConfig:
hf_config: object
class GPUModelRunner:
kv_caches: list[torch.Tensor]
compilation_config: _CompilationConfig
model_config: _ModelConfig | None
def _allocate_kv_cache_tensors(
self, kv_cache_config: object
) -> dict[str, torch.Tensor]: ...
def initialize_kv_cache_tensors(
self, kv_cache_config: object, kernel_block_sizes: list[int]
) -> dict[str, torch.Tensor]: ...
def _reshape_kv_cache_tensors(
self,
kv_cache_config: object,
raw_tensors: dict[str, torch.Tensor],
kernel_block_sizes: list[int],
) -> dict[str, torch.Tensor]: ...
@@ -1,6 +0,0 @@
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
class Worker:
model_runner: GPUModelRunner
def determine_available_memory(self) -> int: ...
def initialize_from_config(self, kv_cache_config: object) -> None: ...
-1
View File
@@ -1 +0,0 @@
def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from typing import Protocol
import mlx.core as mx
import PIL.Image
import tqdm
from typing import Protocol
from mflux.models.common.config.config import Config
class BeforeLoopCallback(Protocol):
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.callbacks.callback import (
AfterLoopCallback,
BeforeLoopCallback,
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
import PIL.Image
import tqdm
from typing import TYPE_CHECKING
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.models.common.config.config import Config
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import Any
import mlx.core as mx
from mflux.models.common.config.model_config import ModelConfig
from tqdm import tqdm
from mflux.models.common.config.model_config import ModelConfig
logger = ...
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from functools import lru_cache
from typing import Literal
import mlx.core as mx
class ModelConfig:
precision: mx.Dtype = ...
def __init__(
@@ -2,10 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import TYPE_CHECKING, TypeAlias
import mlx.core as mx
from mlx import nn
from mflux.models.common.vae.tiling_config import TilingConfig
from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
@@ -13,7 +13,6 @@ from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreat
from mflux.models.z_image.latent_creator.z_image_latent_creator import (
ZImageLatentCreator,
)
from mlx import nn
if TYPE_CHECKING:
LatentCreatorType: TypeAlias = type[
@@ -2,8 +2,8 @@
This type stub file was generated by pyright.
"""
from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
from mlx import nn
from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
class FusedLoRALinear(nn.Module):
def __init__(
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from collections.abc import Callable
from dataclasses import dataclass
import mlx.core as mx
import mlx.nn as nn
from collections.abc import Callable
from dataclasses import dataclass
from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
@dataclass
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from collections.abc import Callable
from dataclasses import dataclass
from typing import List, Protocol
import mlx.core as mx
@dataclass
class LoRATarget:
model_path: str
@@ -36,3 +36,4 @@ class Rule(NamedTuple):
name: str
check: str
action: QuantizationAction | PathAction | LoraAction | ConfigAction
...
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.config.model_config import ModelConfig
if TYPE_CHECKING: ...
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from abc import ABC, abstractmethod
import mlx.core as mx
from abc import ABC, abstractmethod
class BaseScheduler(ABC):
@property
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -4,10 +4,9 @@ This type stub file was generated by pyright.
from abc import ABC, abstractmethod
from typing import Protocol, runtime_checkable
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
from PIL import Image
from transformers import PreTrainedTokenizer
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
"""
This type stub file was generated by pyright.
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.common.weights.loading.weight_definition import TokenizerDefinition
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from dataclasses import dataclass
import mlx.core as mx
from dataclasses import dataclass
"""
This type stub file was generated by pyright.
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import Callable
import mlx.core as mx
from typing import Callable
class VAETiler:
@staticmethod
@@ -3,8 +3,8 @@ This type stub file was generated by pyright.
"""
import mlx.core as mx
from mflux.models.common.vae.tiling_config import TilingConfig
from mlx import nn
from mflux.models.common.vae.tiling_config import TilingConfig
class VAEUtil:
@staticmethod
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.nn as nn
from typing import TYPE_CHECKING
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, List, TypeAlias
import mlx.core as mx
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
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,
)
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import Dict, List, Optional
import mlx.core as mx
from typing import Dict, List, Optional
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
class WeightMapper:
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, Optional, Protocol
import mlx.core as mx
"""
This type stub file was generated by pyright.
"""
@@ -2,8 +2,7 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING, Any
from typing import Any, TYPE_CHECKING
from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType
if TYPE_CHECKING: ...
@@ -2,10 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from pathlib import Path
import mlx.core as mx
from PIL import Image
@dataclass
@@ -14,6 +13,7 @@ class DepthResult:
depth_array: mx.array
min_depth: float
max_depth: float
...
class DepthPro:
def __init__(self, quantize: int | None = ...) -> None: ...
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,

Some files were not shown because too many files have changed in this diff Show More