Point mlx-lm to davidmcc73 fork with context parallelism support

Store pre-shard head count and distributed group on each attention
layer during sharding, enabling automatic TP→CP switching at runtime
when context length exceeds a threshold.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.cuda_typings/openai_harmony/__init__.pyi

@@ -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: ...

.cuda_typings/pynvml/__init__.pyi

@@ -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: ...

.cuda_typings/torch/__init__.pyi

@@ -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: ...

.cuda_typings/torch/backends/__init__.pyi

@@ -1 +0,0 @@
-from torch.backends import cuda as cuda

.cuda_typings/torch/backends/cuda/__init__.pyi

@@ -1 +0,0 @@
-def is_built() -> bool: ...

.cuda_typings/torch/cuda/__init__.pyi

@@ -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: ...

.cuda_typings/torch/distributed/__init__.pyi

@@ -1,2 +0,0 @@
-def is_initialized() -> bool: ...
-def destroy_process_group() -> None: ...

.cuda_typings/vllm/__init__.pyi

@@ -1 +0,0 @@
-__version__: str

.cuda_typings/vllm/config.pyi

@@ -1,2 +0,0 @@
-class ModelConfig:
-    max_model_len: int

.cuda_typings/vllm/engine/arg_utils.pyi

@@ -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 = ...

.cuda_typings/vllm/outputs.pyi

@@ -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

.cuda_typings/vllm/sampling_params.pyi

@@ -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

.cuda_typings/vllm/tokenizers/__init__.pyi

@@ -1,3 +0,0 @@
-from vllm.tokenizers.protocol import TokenizerLike
-
-__all__ = ["TokenizerLike"]

.cuda_typings/vllm/tokenizers/protocol.pyi

@@ -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]: ...

.cuda_typings/vllm/v1/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/core/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/core/kv_cache_manager.pyi

@@ -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: ...

.cuda_typings/vllm/v1/core/kv_cache_utils.pyi

@@ -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]: ...

.cuda_typings/vllm/v1/engine/llm_engine.pyi

@@ -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: ...

.cuda_typings/vllm/v1/kv_cache_interface.pyi

@@ -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]

.cuda_typings/vllm/v1/request.pyi

@@ -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

.cuda_typings/vllm/v1/worker/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/worker/gpu_model_runner.pyi

@@ -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]: ...

.cuda_typings/vllm/v1/worker/gpu_worker.pyi

@@ -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: ...

.cuda_typings/vllm/v1/worker/utils.pyi

@@ -1 +0,0 @@
-def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...

.github/actions/typecheck/action.yml

@@ -0,0 +1,12 @@
+name: Type Check
+
+description: "Run type checker"
+
+runs:
+  using: "composite"
+  steps:
+    - name: Run type checker
+      run: |
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
+      shell: bash

.github/workflows/build-app.yml

@@ -396,7 +396,7 @@ jobs:
           path: output/EXO-${{ env.RELEASE_VERSION }}.dmg
 
       - name: Upload to S3
-        if: env.SPARKLE_S3_BUCKET != ''
+        if: env.SPARKLE_S3_BUCKET != '' && github.ref_type == 'tag'
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -412,12 +412,6 @@ jobs:
             PREFIX="${PREFIX}/"
           fi
           DMG_NAME="EXO-${RELEASE_VERSION}.dmg"
-
-          if [[ "${{ github.ref_type }}" != "tag" ]]; then
-            aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-${GITHUB_SHA}.dmg"
-            exit 0
-          fi
-
           aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}${DMG_NAME}"
           if [[ "$IS_ALPHA" != "true" ]]; then
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"

.github/workflows/pipeline.yml

@@ -8,6 +8,92 @@ 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: Configure git user
+        run: |
+          git config --local user.email "github-actions@users.noreply.github.com"
+          git config --local user.name  "github-actions bot"
+        shell: bash
+
+      - name: Pull LFS files
+        run: |
+          echo "Pulling Git LFS files..."
+          git lfs pull
+        shell: bash
+
+      - name: Setup Nix Environment
+        run: |
+          echo "Checking for nix installation..."
+          
+          # Check if nix binary exists directly
+          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
+            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
+            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
+            echo "PATH=$PATH" >> $GITHUB_ENV
+            nix --version
+          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
+            echo "Found nix profile script, sourcing..."
+            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
+            nix --version
+          elif command -v nix >/dev/null 2>&1; then
+            echo "Nix already in PATH"
+            nix --version
+          else
+            echo "Nix not found. Debugging info:"
+            echo "Contents of /nix/var/nix/profiles/default/:"
+            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
+            echo "Contents of /nix/var/nix/profiles/default/bin/:"
+            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
+            exit 1
+          fi
+        shell: bash
+
+      - name: Configure basedpyright include for local MLX
+        run: |
+          RUNNER_LABELS='${{ toJSON(runner.labels) }}'
+          if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
+            if [ -d "/Users/Shared/mlx" ]; then
+              echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
+              awk '
+                BEGIN { in=0 }
+                /^\[tool\.basedpyright\]/ { in=1; print; next }
+                in && /^\[/ { in=0 }  # next section
+                in && /^[ \t]*include[ \t]*=/ {
+                  print "include = [\"/Users/Shared/mlx\"]"
+                  next
+                }
+                { print }
+              ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
+
+              echo "New [tool.basedpyright] section:"
+              sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
+            else
+              echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
+            fi
+          else
+            echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
+          fi
+        shell: bash
+
+      - uses: ./.github/actions/typecheck
+
   nix:
     name: Build and check (${{ matrix.system }})
     runs-on: ${{ matrix.runner }}
@@ -37,63 +123,6 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Build Metal packages (macOS only)
-        if: runner.os == 'macOS'
-        run: |
-          # Try to build metal-toolchain first (may succeed via cachix cache hit)
-          if nix build .#metal-toolchain 2>/dev/null; then
-            echo "metal-toolchain built successfully (likely cache hit)"
-          else
-            echo "metal-toolchain build failed, extracting from Xcode..."
-
-            NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="
-            NAR_NAME="metal-toolchain-17C48.nar"
-
-            # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS
-            WORK_DIR="${RUNNER_TEMP}/metal-work"
-            mkdir -p "$WORK_DIR"
-
-            # Download the Metal toolchain component
-            xcodebuild -downloadComponent MetalToolchain
-
-            # Find and mount the DMG
-            DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1)
-            if [ -z "$DMG_PATH" ]; then
-              echo "Error: Could not find Metal toolchain DMG"
-              exit 1
-            fi
-
-            echo "Found DMG at: $DMG_PATH"
-            hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg"
-
-            # Copy the toolchain
-            cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export"
-            hdiutil detach "${WORK_DIR}/metal-dmg"
-
-            # Create NAR and add to store
-            nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}"
-            STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}")
-            echo "Added NAR to store: $STORE_PATH"
-
-            # Verify the hash matches
-            ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}")
-            if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then
-              echo "Warning: NAR hash mismatch!"
-              echo "Expected: $NAR_HASH"
-              echo "Actual:   $ACTUAL_HASH"
-              echo "The metal-toolchain.nix may need updating"
-            fi
-
-            # Clean up
-            rm -rf "$WORK_DIR"
-
-            # Retry the build now that NAR is in store
-            nix build .#metal-toolchain
-          fi
-
-          # Build mlx (depends on metal-toolchain)
-          nix build .#mlx
-
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '
@@ -105,16 +134,3 @@ jobs:
 
       - name: Run nix flake check
         run: nix flake check
-
-      - name: Run pytest (macOS only)
-        if: runner.os == 'macOS'
-        run: |
-          # Build the test environment (requires relaxed sandbox for uv2nix on macOS)
-          TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths)
-
-          # Run pytest outside sandbox (needs GPU access for MLX)
-          export HOME="$RUNNER_TEMP"
-          export EXO_TESTS=1
-          export EXO_DASHBOARD_DIR="$PWD/dashboard/" 
-          export EXO_RESOURCES_DIR="$PWD/resources" 
-          $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib

.gitignore

@@ -28,13 +28,3 @@ target/
 dashboard/build/
 dashboard/node_modules/
 dashboard/.svelte-kit/
-
-# host config snapshots
-hosts_*.json
-.swp
-
-# bench files
-bench/**/*.json
-
-# tmp
-tmp/models

.mlx_typings/mflux/callbacks/callback.pyi

@@ -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):

.mlx_typings/mflux/callbacks/callback_registry.pyi

@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
 """
 
 from typing import TYPE_CHECKING
-
 from mflux.callbacks.callback import (
     AfterLoopCallback,
     BeforeLoopCallback,

.mlx_typings/mflux/callbacks/generation_context.pyi

@@ -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
 

.mlx_typings/mflux/models/common/config/config.pyi

@@ -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 = ...
 

.mlx_typings/mflux/models/common/config/model_config.pyi

@@ -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__(

.mlx_typings/mflux/models/common/latent_creator/latent_creator.pyi

@@ -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[

.mlx_typings/mflux/models/common/lora/layer/fused_linear_lora_layer.pyi

@@ -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__(

.mlx_typings/mflux/models/common/lora/mapping/lora_loader.pyi

@@ -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

.mlx_typings/mflux/models/common/lora/mapping/lora_mapping.pyi

@@ -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

.mlx_typings/mflux/models/common/resolution/actions.pyi

@@ -36,3 +36,4 @@ class Rule(NamedTuple):
     name: str
     check: str
     action: QuantizationAction | PathAction | LoraAction | ConfigAction
+    ...

.mlx_typings/mflux/models/common/resolution/config_resolution.pyi

@@ -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: ...

.mlx_typings/mflux/models/common/schedulers/base_scheduler.pyi

@@ -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

.mlx_typings/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.pyi

@@ -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
 

.mlx_typings/mflux/models/common/schedulers/linear_scheduler.pyi

@@ -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
 

.mlx_typings/mflux/models/common/schedulers/seedvr2_euler_scheduler.pyi

@@ -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
 

.mlx_typings/mflux/models/common/tokenizer/tokenizer.pyi

@@ -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.

.mlx_typings/mflux/models/common/tokenizer/tokenizer_loader.pyi

@@ -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
 

.mlx_typings/mflux/models/common/tokenizer/tokenizer_output.pyi

@@ -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.

.mlx_typings/mflux/models/common/vae/vae_tiler.pyi

@@ -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

.mlx_typings/mflux/models/common/vae/vae_util.pyi

@@ -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

.mlx_typings/mflux/models/common/weights/loading/weight_applier.pyi

@@ -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,

.mlx_typings/mflux/models/common/weights/loading/weight_definition.pyi

@@ -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,
 )

.mlx_typings/mflux/models/common/weights/loading/weight_loader.pyi

@@ -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,

.mlx_typings/mflux/models/common/weights/mapping/weight_mapper.pyi

@@ -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:

.mlx_typings/mflux/models/common/weights/mapping/weight_mapping.pyi

@@ -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.
 """

.mlx_typings/mflux/models/common/weights/saving/model_saver.pyi

@@ -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: ...

.mlx_typings/mflux/models/depth_pro/model/depth_pro.pyi

@@ -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: ...

.mlx_typings/mflux/models/depth_pro/weights/depth_pro_weight_definition.pyi

@@ -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,

.mlx_typings/mflux/models/depth_pro/weights/depth_pro_weight_mapping.pyi

@@ -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,