Compare commits

..

17 Commits

Author SHA1 Message Date
Alex Cheema 7469f44e58 fix: clean up stale runners from state when instance is deleted
apply_instance_deleted() previously only removed the instance from
state.instances, leaving its runner entries orphaned in state.runners
with their last known status (e.g. RunnerReady). After a node kill and
rejoin, readiness checks would see these stale entries and attempt
inference against dead runner processes, causing post-recovery failures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 18:41:22 -08:00
Alex Cheema 5d26d2dcd6 fix: eliminate serialization bottlenecks in continuous batching pipeline
The batch prefill deferral (35973b86) was insufficient alone because
multiple other serialization points prevented true concurrent request
processing. This fixes four bottlenecks:

- Drain all available tasks per plan_step cycle instead of one-per-100ms
- Keep tasks in supervisor pending set after ACK to prevent re-dispatch
- Process TextGeneration tasks inline during decode loop (no break/restart)
- Pre-tokenize in queue_request so sync_and_insert_pending is lightweight
- Reuse prompt from queue_request for thinking detection (no double apply_chat_template)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 18:36:56 -08:00
Alex Cheema 35973b8698 fix: defer batch prefill for true continuous batching
Move sync_and_insert_pending() out of the per-task loop so all
concurrently-arrived requests share a single batched prefill pass.
Previously each request was prefilled individually as it arrived,
serializing what should be a parallel operation.

Also break early from the generation TimeBudget loop when new tasks
are waiting, so they get inserted sooner rather than blocking for
the full 0.5s budget.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 17:52:04 -08:00
Alex Cheema 41d9d2a61f fix: add has_thinking to mock tokenizers in edge case tests
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 17:02:07 -08:00
Alex Cheema efbf9850eb style: apply nix fmt to new edge case test file
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 16:55:38 -08:00
Alex Cheema 4a22c4b512 fix: restore per-request sampling, model-specific parsers, error handling, and tracing
- Use per-request temperature/top_p/top_k from TextGenerationTaskParams
  instead of hardcoded sampler defaults in BatchGenerationEngine
- Restore model-specific tokenizer patches (Kimi, GLM) at load time
- Add GptOssTracker for per-request GPT-OSS stream parsing with
  thinking channels and tool call routing
- Filter Kimi section boundary tokens from batch output
- Detect thinking prompt suffix and prepend think_start on first token
- Wrap batch_engine.step() in try/except, send ErrorChunks on failure
- Call _send_traces_if_enabled() when text generation tasks complete

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 16:55:18 -08:00
Alex Cheema 4d0fe5d17b test: add edge-case tests for continuous batching
Cover concurrent tool calls, length/stop finish reasons, multiple
completions per step, staggered draining, and batches of 5-10 requests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 16:27:45 -08:00
Alex Cheema 9fe7251796 fix: add generation loop, deferred task completion, and tool call tracking
The batch engine integration had three critical issues:
1. No generation loop - batch_engine.step() only ran during shutdown drain
2. Tasks marked complete before any tokens were generated
3. Tool calls dropped - parse_tool_calls pipeline was disconnected

Restructure runner main() into a two-phase while loop that alternates
between TimeBudget-based generation steps and task polling. Add
ToolCallTracker for per-request tool call state in the batch path,
and defer TaskStatusUpdated(Complete) until finish_reason is set.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 15:22:09 -08:00
Alex Cheema 1c8f69ce00 fix: address PR review comments for continuous batching
- Use get_args(FinishReason) instead of hardcoded finish reason checks
- Use new Python type syntax (def share_object[T]) instead of TypeVar
- Assert obj is not None for rank 0 with message to use mx_barrier()
- Raise RuntimeError on size=0 instead of silently returning None
- Simplify share_object callers since return type is now T (not T | None)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 14:29:13 -08:00
Alex Cheema f19166617a fix: use EventCollector instead of mp_channel to fix flaky test
Replace mp_channel event receiver with direct EventCollector in
test_continuous_batching.py to eliminate multiprocessing pipe race
condition that caused test_runner_status_reflects_active_requests
to intermittently miss RunnerRunning events.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 14:00:28 -08:00
Alex Cheema 51e959c979 feat: integrate BatchGenerationEngine into runner for continuous batching
Replace synchronous per-request text generation with BatchGenerationEngine,
enabling continuous batching of multiple concurrent inference requests.

- Runner accepts TextGeneration in both RunnerReady and RunnerRunning states
- Requests are queued and sync-inserted into the batch engine
- Batch engine is drained during shutdown to complete in-flight requests
- Only rank 0 emits ChunkGenerated events (distributed-safe)
- Enable previously-skipped continuous batching tests
- Update event ordering tests for the new batch-based flow

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 13:47:04 -08:00
Alex Cheema cd43588a04 Merge remote-tracking branch 'origin/main' into alexcheema/continuous-batching 2026-02-13 10:09:40 -08:00
Alex Cheema 7b879593bb fix: update continuous batching types after main merge
Replace ChatCompletionTaskParams with TextGenerationTaskParams and
ChatCompletion with TextGeneration to match the refactored type
hierarchy from main. Add missing usage parameter to GenerationResponse
constructors and add type annotations to StreamingDetokenizer stubs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-13 07:29:13 -08:00
Alex Cheema e4e895d7a8 Merge remote-tracking branch 'origin/main' into alexcheema/continuous-batching
# Conflicts:
#	AGENTS.md
2026-02-13 05:57:15 -08:00
Alex Cheema db400dbb75 skip continuous batching tests pending type migration
The continuous batching runner architecture references old types
(ChatCompletion, ChatCompletionTaskParams) that were renamed on main.
Skip the test module until the batch engine code is updated.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-02-05 06:18:40 -08:00
Alex Cheema 15fad9c632 Merge remote-tracking branch 'origin/main' into alexcheema/continuous-batching
# Conflicts:
#	.mlx_typings/mlx_lm/tokenizer_utils.pyi
#	src/exo/worker/runner/runner.py
#	src/exo/worker/runner/runner_supervisor.py
#	src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
2026-02-05 06:12:49 -08:00
Alex Cheema 842beefac0 feat: add continuous batching for distributed inference
Implements continuous batching using mlx_lm's BatchGenerator for efficient
multi-request handling in distributed mode.

Key changes:
- Add BatchGenerationEngine that wraps mlx_lm's BatchGenerator for continuous
  batching with prefill batching (up to 8 requests) and decode batching
- Add TimeBudget pattern for controlling generation loop timing with periodic
  distributed sync
- Add distributed_sync utilities for broadcasting objects across ranks using
  mx.distributed.all_sum()
- Stream tokens immediately as generated for smooth streaming (not in batches)
- Fix distributed correctness: deferred shutdown handling, sync_completions
  always syncs in distributed mode to prevent deadlocks

Performance results on Kimi K2 Thinking (658GB) with Tensor RDMA:
- Batch 1:  10.7 tok/s (baseline)
- Batch 4:  34.6 tok/s (3.2x speedup)
- Batch 16: 41.8 tok/s (3.9x speedup)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-22 12:27:15 +00:00
647 changed files with 11356 additions and 36584 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: ...
+27
View File
@@ -8,6 +8,33 @@ on:
- main
jobs:
typecheck:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v4
with:
lfs: false
- uses: cachix/install-nix-action@v31
with:
nix_path: nixpkgs=channel:nixos-unstable
- uses: cachix/cachix-action@v14
name: Configure Cachix
with:
name: exo
authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
- name: Load nix develop environment
run: nix run github:nicknovitski/nix-develop/v1
- name: Sync dependencies
run: uv sync --all-packages
- name: Run type checker
run: uv run basedpyright --project pyproject.toml
nix:
name: Build and check (${{ matrix.system }})
runs-on: ${{ matrix.runner }}
@@ -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