add new kill bridge script

.envrc

@@ -1,8 +1 @@
 use flake
-
-# creates .venv if doesn't exist and loads its environment
-export VIRTUAL_ENV=".venv"
-if ! [ -d "./$VIRTUAL_ENV" ]; then
-  uv venv
-fi
-layout python

.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

@@ -32,6 +32,7 @@ jobs:
       SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
       SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
       SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
+      EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
       AWS_REGION: ${{ secrets.AWS_REGION }}
       EXO_BUILD_NUMBER: ${{ github.run_number }}
       EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -158,7 +159,7 @@ jobs:
           fi
 
       - name: Install Homebrew packages
-        run: brew install just awscli
+        run: brew install just awscli macmon
 
       - name: Install UV
         uses: astral-sh/setup-uv@v6
@@ -238,92 +239,10 @@ jobs:
           # Export keychain path for other steps
           echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV
 
-      # ============================================================
-      # Pre-flight credential / profile validation
-      # Runs BEFORE the ~16 min build so auth/expiry failures surface in <1 min.
-      # ============================================================
-
-      - name: Validate Apple notarization credentials
-        env:
-          APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }}
-          APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
-          APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
-        run: |
-          # All-or-nothing: either all three creds are set, or none are.
-          CRED_COUNT=0
-          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
-            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
-          done
-          if [[ "$CRED_COUNT" -eq 0 ]]; then
-            echo "No notarization credentials configured — skipping notarization for this build."
-            exit 0
-          fi
-          if [[ "$CRED_COUNT" -ne 3 ]]; then
-            echo "ERROR: partial notarization credentials set ($CRED_COUNT/3). Aborting before build."
-            exit 1
-          fi
-          # Cheap, ~5s, auth-only call. Fails instantly with a clear message if
-          # the app-specific password is stale, wrong team-id, etc.
-          echo "Verifying Apple notarization credentials via notarytool history..."
-          if ! xcrun notarytool history \
-            --apple-id "$APPLE_NOTARIZATION_USERNAME" \
-            --password "$APPLE_NOTARIZATION_PASSWORD" \
-            --team-id "$APPLE_NOTARIZATION_TEAM" >/dev/null; then
-            echo "ERROR: notarytool rejected the provided credentials. Fix before rerunning."
-            echo "Common causes: app-specific password expired/revoked, wrong team-id,"
-            echo "Apple ID not on the team, or 2FA not configured for this Apple ID."
-            exit 1
-          fi
-          echo "Apple notarization credentials OK."
-
-      - name: Validate provisioning profile expiry
-        run: |
-          PROFILE="$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile"
-          if [[ ! -f "$PROFILE" ]]; then
-            echo "ERROR: provisioning profile not found at $PROFILE"
-            exit 1
-          fi
-          EXPIRY=$(security cms -D -i "$PROFILE" | plutil -extract ExpirationDate raw -o - - 2>/dev/null || true)
-          if [[ -z "$EXPIRY" ]]; then
-            echo "WARNING: could not read ExpirationDate from provisioning profile; skipping expiry check."
-            exit 0
-          fi
-          # Try a couple of known plutil date formats. If none parse, skip the check rather
-          # than risk a false-positive "expired" block on a format we didn't anticipate.
-          EXPIRY_EPOCH=""
-          for fmt in "%Y-%m-%dT%H:%M:%SZ" "%Y-%m-%d %H:%M:%S %z" "%Y-%m-%d %H:%M:%S +0000"; do
-            if parsed=$(date -j -f "$fmt" "$EXPIRY" +%s 2>/dev/null); then
-              EXPIRY_EPOCH="$parsed"
-              break
-            fi
-          done
-          if [[ -z "$EXPIRY_EPOCH" ]]; then
-            echo "WARNING: could not parse ExpirationDate '$EXPIRY'; skipping expiry check."
-            exit 0
-          fi
-          NOW_EPOCH=$(date +%s)
-          if [[ "$EXPIRY_EPOCH" -le "$NOW_EPOCH" ]]; then
-            echo "ERROR: provisioning profile expired on $EXPIRY. Regenerate it before rerunning."
-            exit 1
-          fi
-          DAYS_LEFT=$(( (EXPIRY_EPOCH - NOW_EPOCH) / 86400 ))
-          echo "Provisioning profile valid until $EXPIRY ($DAYS_LEFT days remaining)."
-          if [[ "$DAYS_LEFT" -lt 14 ]]; then
-            echo "WARNING: profile expires in under 14 days — regenerate soon."
-          fi
-
       # ============================================================
       # Build the bundle
       # ============================================================
 
-      - name: Add pinned macmon to PATH
-        run: |
-          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
-          echo "Using macmon from: $MACMON_DIR"
-          echo "$MACMON_DIR" >> $GITHUB_PATH
-          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
-          brew uninstall macmon 2>/dev/null || true
-
       - name: Build PyInstaller bundle
         run: uv run pyinstaller packaging/pyinstaller/exo.spec
 
@@ -346,6 +265,7 @@ jobs:
             EXO_BUILD_COMMIT="$GITHUB_SHA" \
             SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
             SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
+            EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
             CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
             CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
           mkdir -p ../../output
@@ -378,41 +298,11 @@ jobs:
           APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
           APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
         run: |
-          set -o pipefail
           cd output
           security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH"
           SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}')
-
-          # Fail fast if notarization creds are partial. All-or-nothing.
-          CRED_COUNT=0
-          for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
-            [[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
-          done
-          if [[ "$CRED_COUNT" -ne 0 && "$CRED_COUNT" -ne 3 ]]; then
-            echo "ERROR: partial Apple notarization credentials set ($CRED_COUNT/3). Aborting."
-            exit 1
-          fi
-
           /usr/bin/codesign --deep --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" EXO.app
-
-          # Pre-flight: verify the signed app BEFORE building DMG and submitting to Apple.
-          # If this fails, notarization will fail too — cheap way to fail in seconds, not 15 minutes.
-          echo "===== codesign --verify EXO.app ====="
-          if ! /usr/bin/codesign --verify --deep --strict --verbose=2 EXO.app; then
-            echo "ERROR: EXO.app failed codesign verification. Dumping signing status of every executable:"
-            find EXO.app -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" -o -name "*.framework" \) -print0 |
-              while IFS= read -r -d '' f; do
-                printf -- '--- %s\n' "$f"
-                /usr/bin/codesign -dv --verbose=2 "$f" 2>&1 | sed 's/^/    /' || true
-              done
-            exit 1
-          fi
-
-          # Gatekeeper assessment. A failure here strongly predicts notarization rejection.
-          echo "===== spctl assessment (predicts notarization outcome) ====="
-          /usr/bin/spctl -a -vvv -t install EXO.app || echo "WARNING: spctl assessment failed — notarization is likely to fail too."
-
           mkdir -p dmg-root
           cp -R EXO.app dmg-root/
           ln -s /Applications dmg-root/Applications
@@ -420,22 +310,12 @@ jobs:
           hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME"
           /usr/bin/codesign --force --timestamp --options runtime \
             --sign "$SIGNING_IDENTITY" "$DMG_NAME"
-
-          echo "===== codesign --verify DMG ====="
-          if ! /usr/bin/codesign --verify --verbose=2 "$DMG_NAME"; then
-            echo "ERROR: DMG failed codesign verification."
-            exit 1
-          fi
-
           if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then
-            echo "===== notarytool submit ====="
-            # `|| true` so set -e doesn't abort before we can echo output / fetch the log.
-            # We rely on the parsed STATUS below to decide pass/fail.
             SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \
               --apple-id "$APPLE_NOTARIZATION_USERNAME" \
               --password "$APPLE_NOTARIZATION_PASSWORD" \
               --team-id "$APPLE_NOTARIZATION_TEAM" \
-              --wait --timeout 15m 2>&1) || true
+              --wait --timeout 15m 2>&1)
             echo "$SUBMISSION_OUTPUT"
 
             SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}')
@@ -516,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 }}
@@ -532,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,60 +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
-
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '
@@ -102,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

@@ -18,6 +18,7 @@ digest.txt
 app/EXO/build/
 dist/
 
+
 # rust
 target/
 **/*.rs.bk
@@ -27,17 +28,3 @@ target/
 dashboard/build/
 dashboard/node_modules/
 dashboard/.svelte-kit/
-
-# host config snapshots
-hosts_*.json
-.swp
-
-# bench files
-bench/**/*.json
-
-# tmp
-tmp/models
-/build/exo
-/.claude/skills
-/.claude
-/.codex

.idea/exo-v2.iml

@@ -0,0 +1,31 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="EMPTY_MODULE" version="4">
+  <component name="FacetManager">
+    <facet type="Python" name="Python facet">
+      <configuration sdkName="Python 3.13 virtualenv at ~/Desktop/exo/.venv" />
+    </facet>
+  </component>
+  <component name="Go" enabled="true" />
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/scripts/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/tests" isTestSource="true" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/util/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/networking/examples" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/networking/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/networking/tests" isTestSource="true" />
+      <sourceFolder url="file://$MODULE_DIR$/rust/system_custodian/src" isTestSource="false" />
+      <excludeFolder url="file://$MODULE_DIR$/.venv" />
+      <excludeFolder url="file://$MODULE_DIR$/.direnv" />
+      <excludeFolder url="file://$MODULE_DIR$/build" />
+      <excludeFolder url="file://$MODULE_DIR$/dist" />
+      <excludeFolder url="file://$MODULE_DIR$/.go_cache" />
+      <excludeFolder url="file://$MODULE_DIR$/rust/target" />
+    </content>
+    <orderEntry type="jdk" jdkName="Python 3.13 (exo)" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+    <orderEntry type="library" name="Python 3.13 virtualenv at ~/Desktop/exo/.venv interpreter library" level="application" />
+  </component>
+</module>

.idea/externalDependencies.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalDependencies">
-    <plugin id="al.aoli.intellijdirenv" />
+    <plugin id="systems.fehn.intellijdirenv" />
   </component>
 </project>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,14 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ourVersions">
+        <value>
+          <list size="1">
+            <item index="0" class="java.lang.String" itemvalue="3.14" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/misc.xml

@@ -4,7 +4,7 @@
     <option name="sdkName" value="Python 3.13 (exo)" />
   </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
-  <component name="RuffConfiguration">
-    <option name="enabled" value="true" />
+  <component name="PythonCompatibilityInspectionAdvertiser">
+    <option name="version" value="3" />
   </component>
 </project>

.mlx_typings/mlx/core/__init__.pyi

@@ -215,22 +215,6 @@ class StreamContext:
         traceback: object | None = ...,
     ) -> None: ...
 
-def device_info() -> dict[str, str | int]:
-    """
-    Get information about the GPU device and system settings.
-
-    Currently returns:
-
-    * ``architecture``
-    * ``max_buffer_size``
-    * ``max_recommended_working_set_size``
-    * ``memory_size``
-    * ``resource_limit``
-
-    Returns:
-        dict: A dictionary with string keys and string or integer values.
-    """
-
 def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
     """
     Element-wise absolute value.
@@ -1155,7 +1139,7 @@ class array:
     ) -> array:
         """See :func:`flatten`."""
 
-    def reshape(self, *shape: int, stream: Stream | Device | None = ...) -> array:
+    def reshape(self, *shape, stream: Stream | Device | None = ...) -> array:
         """
         Equivalent to :func:`reshape` but the shape can be passed either as a
         :obj:`tuple` or as separate arguments.
@@ -1238,7 +1222,7 @@ class array:
     ) -> array:
         """See :func:`swapaxes`."""
 
-    def transpose(self, *axes: int, stream: Stream | Device | None = ...) -> array:
+    def transpose(self, *axes, stream: Stream | Device | None = ...) -> array:
         """
         Equivalent to :func:`transpose` but the axes can be passed either as
         a tuple or as separate arguments.
@@ -1767,12 +1751,12 @@ def clip(
         array: The clipped array.
     """
 
-def compile[F: Callable[..., object]](
-    fun: F,
+def compile(
+    fun: Callable,
     inputs: object | None = ...,
     outputs: object | None = ...,
     shapeless: bool = ...,
-) -> F:
+) -> Callable:
     """
     Returns a compiled function which produces the same output as ``fun``.
 
@@ -2382,7 +2366,7 @@ class custom_function:
 def default_device() -> Device:
     """Get the default device."""
 
-def default_stream(device: Device | DeviceType) -> Stream:
+def default_stream(device: Device) -> Stream:
     """Get the device's default stream."""
 
 def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
@@ -2396,7 +2380,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
+def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.
@@ -2915,8 +2899,8 @@ def gather_mm(
     a: array,
     b: array,
     /,
-    lhs_indices: array | None = ...,
-    rhs_indices: array | None = ...,
+    lhs_indices: array,
+    rhs_indices: array,
     *,
     sorted_indices: bool = ...,
     stream: Stream | Device | None = ...,
@@ -4707,7 +4691,6 @@ def softmax(
     /,
     axis: int | Sequence[int] | None = ...,
     *,
-    precise: bool = ...,
     stream: Stream | Device | None = ...,
 ) -> array:
     """

.mlx_typings/mlx/nn/__init__.pyi

@@ -0,0 +1,9 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from layers import *
+from utils import *
+
+from . import init as init
+from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -0,0 +1,20 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from activations import *
+from base import *
+from containers import *
+from convolution import *
+from convolution_transpose import *
+from distributed import *
+from dropout import *
+from embedding import *
+from linear import *
+from normalization import *
+from pooling import *
+from positional_encoding import *
+from quantized import *
+from recurrent import *
+from transformer import *
+from upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,14 +53,10 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    __call__: Callable
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 
-    def __getitem__(self, key: str) -> mx.array | Module: ...
-    def get(
-        self, key: str, default: mx.array | Module | None = ...
-    ) -> mx.array | Module | None: ...
     @property
     def training(self):  # -> bool:
         """Boolean indicating if the model is in training mode."""
@@ -204,7 +200,7 @@ class Module(dict):
     ) -> mx.MX_ARRAY_TREE:  # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
         """Return the submodules that do not contain other modules."""
 
-    def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module:
+    def update(self, parameters: dict, strict: bool = ...) -> Module:
         """Replace the parameters of this Module with the provided ones in the
         dict of dicts and lists.
 

.mlx_typings/mlx/nn/layers/convolution.pyi

@@ -30,10 +30,6 @@ class Conv1d(Module):
         bias (bool, optional): If ``True`` add a learnable bias to the output.
             Default: ``True``
     """
-
-    weight: mx.array
-    bias: mx.array | None
-    groups: int
     def __init__(
         self,
         in_channels: int,

.mlx_typings/mlx/nn/layers/linear.pyi

@@ -40,10 +40,6 @@ class Linear(Module):
         bias (bool, optional): If set to ``False`` then the layer will
           not use a bias. Default is ``True``.
     """
-
-    weight: mx.array
-    bias: mx.array | None
-
     def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...
     def to_quantized(

.mlx_typings/mlx/nn/layers/normalization.pyi

@@ -88,9 +88,6 @@ class RMSNorm(Module):
         dims (int): The feature dimension of the input to normalize over
         eps (float): A small additive constant for numerical stability
     """
-
-    weight: mx.array
-
     def __init__(self, dims: int, eps: float = ...) -> None: ...
     def __call__(self, x) -> mx.array: ...
 

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -2,7 +2,7 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Any, Callable, Optional, Union
+from typing import Callable, Optional, Union
 
 import mlx.core as mx
 from base import Module
@@ -13,10 +13,8 @@ def quantize(
     bits: int = ...,
     *,
     mode: str = ...,
-    class_predicate: Optional[
-        Callable[[str, Module], Union[bool, dict[str, Any]]]
-    ] = ...,
-) -> None:
+    class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
+):  # -> None:
     """Quantize the sub-modules of a module according to a predicate.
 
     By default all layers that define a ``to_quantized(group_size, bits)``

.mlx_typings/mlx/utils.pyi

@@ -7,10 +7,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from mlx.core import MX_ARRAY_TREE
 
 def tree_map(
-    fn: Callable[..., Any],
-    tree: Any,
-    *rest: Any,
-    is_leaf: Callable[..., bool] | None = ...,
+    fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
 ) -> Any:
     """Applies ``fn`` to the leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -47,11 +44,11 @@ def tree_map(
     """
 
 def tree_map_with_path(
-    fn: Callable[..., Any],
+    fn: Callable,
     tree: Any,
     *rest: Any,
-    is_leaf: Callable[..., bool] | None = ...,
-    path: str | None = ...,
+    is_leaf: Optional[Callable] = ...,
+    path: Optional[Any] = ...,
 ) -> Any:
     """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -83,9 +80,9 @@ def tree_map_with_path(
 def tree_flatten(
     tree: Any,
     prefix: str = ...,
-    is_leaf: Callable[..., bool] | None = ...,
-    destination: list[tuple[str, Any]] | dict[str, Any] | None = ...,
-) -> list[tuple[str, Any]] | dict[str, Any]:
+    is_leaf: Optional[Callable] = ...,
+    destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ...,
+) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
     """Flattens a Python tree to a list of key, value tuples.
 
     The keys are using the dot notation to define trees of arbitrary depth and
@@ -121,7 +118,7 @@ def tree_flatten(
             the Python tree.
     """
 
-def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any:
+def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
     """Recreate a Python tree from its flat representation.
 
     .. code-block:: python

.mlx_typings/mlx_lm/generate.pyi

@@ -3,12 +3,13 @@ This type stub file was generated by pyright.
 """
 
 import contextlib
+from dataclasses import dataclass
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
+
 import mlx.core as mx
 import mlx.nn as nn
-from dataclasses import dataclass
-from collections import deque
-from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
 from transformers import PreTrainedTokenizer
+
 from .tokenizer_utils import TokenizerWrapper
 
 DEFAULT_PROMPT = ...
@@ -28,9 +29,8 @@ def str2bool(string):  # -> bool:
     ...
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
-    ...
 
-generation_stream: mx.Stream
+generation_stream = ...
 
 @contextlib.contextmanager
 def wired_limit(
@@ -43,7 +43,6 @@ def wired_limit(
     async eval could be running pass in the streams to synchronize with prior
     to exiting the context manager.
     """
-    ...
 @dataclass
 class GenerationResponse:
     """
@@ -74,11 +73,9 @@ class GenerationResponse:
     finish_reason: Optional[str] = ...
 
 def maybe_quantize_kv_cache(
-    prompt_cache: Any,
-    quantized_kv_start: int | None,
-    kv_group_size: int | None,
-    kv_bits: int | None,
-) -> None: ...
+    prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+):  # -> None:
+    ...
 def generate_step(
     prompt: mx.array,
     model: nn.Module,
@@ -92,7 +89,7 @@ def generate_step(
     kv_bits: Optional[int] = ...,
     kv_group_size: int = ...,
     quantized_kv_start: int = ...,
-    prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
+    prompt_progress_callback: Optional[Callable[[int], int]] = ...,
     input_embeddings: Optional[mx.array] = ...,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -118,7 +115,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
+        prompt_progress_callback (Callable[[int], int]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -126,7 +123,6 @@ def generate_step(
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
-    ...
 
 def speculative_generate_step(
     prompt: mx.array,
@@ -172,7 +168,6 @@ def speculative_generate_step(
         Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
           and a bool indicating if the token was generated by the draft model
     """
-    ...
 
 def stream_generate(
     model: nn.Module,
@@ -180,7 +175,7 @@ def stream_generate(
     prompt: Union[str, mx.array, List[int]],
     max_tokens: int = ...,
     draft_model: Optional[nn.Module] = ...,
-    **kwargs: Any,
+    **kwargs: object,
 ) -> Generator[GenerationResponse, None, None]:
     """
     A generator producing text based on the given prompt from the model.
@@ -202,7 +197,6 @@ def stream_generate(
         GenerationResponse: An instance containing the generated text segment and
             associated metadata. See :class:`GenerationResponse` for details.
     """
-    ...
 
 def generate(
     model: nn.Module,
@@ -223,9 +217,6 @@ def generate(
        kwargs: The remaining options get passed to :func:`stream_generate`.
           See :func:`stream_generate` for more details.
     """
-    ...
-
-def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
 @dataclass
 class BatchStats:
     """
@@ -249,263 +240,10 @@ class BatchStats:
     generation_time: float = ...
     peak_memory: float = ...
 
-class SequenceStateMachine:
-    """A state machine that uses one Aho-Corasick trie per state to efficiently
-    track state across a generated sequence.
-
-    The transitions are provided as state -> [(sequence, new_state)].
-
-    Example:
-
-        sm = SequenceStateMachine(
-            transitions={
-                "normal": [
-                    (think_start_tokens, "reasoning"),
-                    (tool_start_tokens, "tool"),
-                    (eos, None),
-                ],
-                "reasoning": [
-                    (think_end_tokens, "normal"),
-                    (eos, None),
-                ],
-                "tool": [
-                    (tool_end_tokens, None),
-                    (eos, None)
-                ],
-            },
-            initial="normal"
-        )
-    """
-    def __init__(self, transitions=..., initial=...) -> None: ...
-    def __deepcopy__(self, memo):  # -> SequenceStateMachine:
-        ...
-    def make_state(self):  # -> tuple[str, Any, dict[Any, Any]]:
-        ...
-    @staticmethod
-    def match(state, x):  # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
-        ...
-
-class PromptProcessingBatch:
-    """
-    A batch processor for prompt tokens with support for incremental processing.
-
-    This class handles batched prompt processing, managing KV caches and preparing
-    tokens for generation. It supports extending, filtering, and splitting batches.
-    """
-    @dataclass
-    class Response:
-        uid: int
-        progress: tuple
-        end_of_segment: bool
-        end_of_prompt: bool
-        ...
-
-    def __init__(
-        self,
-        model: nn.Module,
-        uids: List[int],
-        caches: List[List[Any]],
-        tokens: Optional[List[List[int]]] = ...,
-        prefill_step_size: int = ...,
-        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
-        fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        logits_processors: Optional[
-            List[List[Callable[[mx.array, mx.array], mx.array]]]
-        ] = ...,
-        state_machines: Optional[List[SequenceStateMachine]] = ...,
-        max_tokens: Optional[List[int]] = ...,
-    ) -> None: ...
-    def __len__(self):  # -> int:
-        ...
-    def extract_cache(self, idx: int) -> List[Any]: ...
-    def extend(self, batch):  # -> None:
-        ...
-    def split(self, indices: List[int]):  # -> Self:
-        ...
-    def filter(self, keep: List[int]):  # -> None:
-        ...
-    def prompt(self, tokens: List[List[int]]):  # -> None:
-        """
-        Process prompt tokens through the model.
-
-        Args:
-            tokens: List of token sequences to process.
-        """
-        ...
-
-    def generate(self, tokens: List[List[int]]):  # -> GenerationBatch:
-        """
-        Transition from prompt processing to generation.
-
-        Args:
-            tokens: Final tokens for each sequence to start generation.
-
-        Returns:
-            A GenerationBatch ready for token generation.
-        """
-        ...
-
-    @classmethod
-    def empty(
-        cls,
-        model: nn.Module,
-        fallback_sampler: Callable[[mx.array], mx.array],
-        prefill_step_size: int = ...,
-    ):  # -> Self:
-        ...
-
-class GenerationBatch:
-    """
-    A batched token generator that manages multiple sequences in parallel.
-
-    This class handles the generation phase after prompt processing, managing
-    KV caches, sampling, and stop sequence detection for multiple sequences.
-    """
-    @dataclass
-    class Response:
-        uid: int
-        token: int
-        logprobs: mx.array
-        finish_reason: Optional[str]
-        current_state: Optional[str]
-        match_sequence: Optional[List[int]]
-        prompt_cache: Optional[List[Any]]
-        all_tokens: Optional[List[int]]
-        ...
-
-    model: nn.Module
-    uids: List[int]
-    prompt_cache: List[Any]
-    tokens: List[List[int]]
-    samplers: Optional[List[Callable[[mx.array], mx.array]]]
-    fallback_sampler: Callable[[mx.array], mx.array]
-    logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
-    state_machines: List[SequenceStateMachine]
-    max_tokens: List[int]
-    _current_tokens: Optional[mx.array]
-    _current_logprobs: mx.array | List[mx.array]
-    _next_tokens: Optional[mx.array]
-    _next_logprobs: mx.array | List[mx.array]
-    _token_context: List[Any]
-    _num_tokens: List[int]
-    _matcher_states: List[Any]
-
-    def __init__(
-        self,
-        model: nn.Module,
-        uids: List[int],
-        inputs: mx.array,
-        prompt_cache: List[Any],
-        tokens: List[List[int]],
-        samplers: Optional[List[Callable[[mx.array], mx.array]]],
-        fallback_sampler: Callable[[mx.array], mx.array],
-        logits_processors: Optional[
-            List[List[Callable[[mx.array, mx.array], mx.array]]]
-        ],
-        state_machines: List[SequenceStateMachine],
-        max_tokens: List[int],
-    ) -> None: ...
-    def __len__(self) -> int: ...
-    def extend(self, batch: GenerationBatch) -> None: ...
-    def extract_cache(self, idx: int) -> List[Any]: ...
-    def filter(self, keep: List[int]) -> None: ...
-    def _step(self) -> Tuple[List[int], List[mx.array]]: ...
-    def next(self) -> List[Response]:
-        """
-        Generate the next batch of tokens.
-
-        Returns:
-            List of Response objects for each sequence in the batch.
-        """
-        ...
-
-    @classmethod
-    def empty(
-        cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
-    ):  # -> Self:
-        ...
-
-class BatchGenerator:
-    """
-    A batch generator implements continuous batching.
-
-    This class provides automatic management of prompt processing and generation
-    batches, handling the transition between the two.
-
-    It also allows for segmented prompt processing which guarantees that the
-    generator will stop at these boundaries when processing an input.
-    """
-    def __init__(
-        self,
-        model: nn.Module,
-        max_tokens: int = ...,
-        stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
-        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        logits_processors: Optional[
-            List[Callable[[mx.array, mx.array], mx.array]]
-        ] = ...,
-        completion_batch_size: int = ...,
-        prefill_batch_size: int = ...,
-        prefill_step_size: int = ...,
-    ) -> None: ...
-    def close(self) -> None: ...
-    def __del__(self):  # -> None:
-        ...
-    @contextlib.contextmanager
-    def stats(self, stats=...):  # -> Generator[Any | BatchStats, Any, None]:
-        ...
-    _unprocessed_sequences: deque[tuple[Any, ...]]
-    _prompt_batch: PromptProcessingBatch
-    _generation_batch: GenerationBatch
-    _currently_processing: list[Any]
-    _gen_tokens_counter: int
-    _steps_counter: int
-    def _next(
-        self,
-    ) -> tuple[
-        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
-    ]: ...
-    def insert(
-        self,
-        prompts: List[List[int]],
-        max_tokens: Optional[List[int]] = ...,
-        caches: Optional[List[List[Any]]] = ...,
-        all_tokens: Optional[List[List[int]]] = ...,
-        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
-        logits_processors: Optional[
-            List[List[Callable[[mx.array, mx.array], mx.array]]]
-        ] = ...,
-        state_machines: Optional[List[SequenceStateMachine]] = ...,
-    ) -> List[int]: ...
-    def insert_segments(
-        self,
-        segments: List[List[List[int]]],
-        max_tokens: Optional[List[int]] = ...,
-        caches: Optional[List[List[Any]]] = ...,
-        all_tokens: Optional[List[List[int]]] = ...,
-        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
-        logits_processors: Optional[
-            List[List[Callable[[mx.array, mx.array], mx.array]]]
-        ] = ...,
-        state_machines: Optional[List[SequenceStateMachine]] = ...,
-    ) -> List[int]: ...
-    def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
-    def remove(
-        self, uids: List[int], return_prompt_caches: bool = ...
-    ) -> dict[int, Any]: ...
-    @property
-    def prompt_cache_nbytes(self) -> int: ...
-    def next(
-        self,
-    ) -> tuple[
-        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
-    ]: ...
-    def next_generated(self) -> List[GenerationBatch.Response]: ...
-
 @dataclass
 class BatchResponse:
     """
-    A data object to hold a batch generation response.
+    An data object to hold a batch generation response.
 
     Args:
         texts: (List[str]): The generated text for each prompt.
@@ -514,18 +252,55 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
-    caches: Optional[List[List[Any]]]
-    ...
+
+@dataclass
+class Batch:
+    uids: List[int]
+    y: mx.array
+    logprobs: mx.array
+    max_tokens: List[int]
+    num_tokens: List[int]
+    cache: List[Any]
+    def __len__(self):  # -> int:
+        ...
+    def filter(self, keep_idx: List[int]):  # -> None:
+        ...
+    def extend(self, other):  # -> None:
+        ...
+
+class BatchGenerator:
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+
+    def __init__(
+        self,
+        model,
+        max_tokens: int = ...,
+        stop_tokens: Optional[set] = ...,
+        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        completion_batch_size: int = ...,
+        prefill_batch_size: int = ...,
+        prefill_step_size: int = ...,
+    ) -> None: ...
+    def insert(
+        self, prompts, max_tokens: Union[List[int], int, None] = ...
+    ):  # -> list[Any]:
+        ...
+    def stats(self):  # -> BatchStats:
+        ...
+    def next(self):  # -> list[Any]:
+        ...
 
 def batch_generate(
     model,
     tokenizer,
-    prompts: List[List[int]],
-    prompt_caches: Optional[List[List[Any]]] = ...,
+    prompts: List[int],
     max_tokens: Union[int, List[int]] = ...,
     verbose: bool = ...,
-    return_prompt_caches: bool = ...,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -534,22 +309,14 @@ def batch_generate(
     Args:
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompts (List[List[int]]): The input prompts.
-       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
-          for each input prompt. Note, unlike ``generate_step``, the caches
-          won't be updated in-place.
+       prompt (List[List[int]]): The input prompts.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
-       return_prompt_caches (bool): Return the prompt caches in the batch
-          responses. Default: ``False``.
-       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
-    ...
 
 def main():  # -> None:
     ...

.mlx_typings/mlx_lm/models/base.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Optional
 
 import mlx.core as mx
 
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
     bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries: mx.array,
-    keys: mx.array,
-    values: mx.array,
-    cache: Optional[Any],
+    queries,
+    keys,
+    values,
+    cache,
     scale: float,
     mask: Optional[mx.array],
     sinks: Optional[mx.array] = ...,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -11,12 +11,9 @@ import mlx.core as mx
 class Cache(Protocol):
     keys: mx.array
     values: mx.array
-    offset: int
-    def update_and_fetch(
-        self, keys: mx.array, values: mx.array
-    ) -> tuple[mx.array, mx.array]: ...
+    def update_and_fetch(self, keys: mx.array, values: mx.array) -> None: ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
 
@@ -88,18 +85,16 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array | None
-    values: mx.array | None
-    offset: int
+    keys: mx.array
+    values: mx.array
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
     @property
     def meta_state(self) -> Literal[""]: ...
     @meta_state.setter
     def meta_state(self, v) -> None: ...
-    def trim(self, n: int) -> int: ...
     def is_trimmable(self) -> Literal[False]: ...
     @classmethod
     def from_state(cls, state, meta_state) -> Self: ...
@@ -115,13 +110,15 @@ class ConcatenateKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -131,7 +128,10 @@ class QuantizedKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> Any:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(
+        self,
+    ):  # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -143,7 +143,8 @@ class QuantizedKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -155,30 +156,22 @@ class KVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ) -> tuple[array, array]: ...
     @state.setter
     def state(self, v) -> None: ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(
-        self, *args: Any, **kwargs: Any
-    ) -> mx.array | Literal["causal"] | None: ...
+    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
+        ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
-    keys: mx.array | None
-    values: mx.array | None
-    keep: int
-    max_size: int
-    _idx: int
     def __init__(self, max_size, keep=...) -> None: ...
-    def _trim(
-        self, trim_size: int, v: mx.array, append: mx.array | None = ...
-    ) -> mx.array: ...
     def update_and_fetch(
         self, keys, values
     ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
@@ -186,7 +179,8 @@ class RotatingKVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ):  # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -198,7 +192,8 @@ class RotatingKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
@@ -213,7 +208,8 @@ class ArraysCache(_BaseCache):
         ...
     def __getitem__(self, idx): ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> list[Any | array] | list[array]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -227,7 +223,8 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int) -> mx.array | None: ...
+    def make_mask(self, N: int):  # -> array | None:
+        ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -238,7 +235,8 @@ class ChunkedKVCache(KVCache):
         ...
     def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     @property
     def meta_state(self):  # -> tuple[str, ...]:
         ...
@@ -251,9 +249,10 @@ class CacheList(_BaseCache):
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n): ...
     @property
-    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
+    def state(self):  # -> list[Any]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -268,14 +267,29 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step: int
-    keys: array | None
-    values: array | None
-    offset: array
-    left_padding: array
-    _idx: int
-    def __init__(self, left_padding: List[int]) -> None: ...
-    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
+    step = ...
+    def __init__(self, left_padding: List[int]) -> None:
+        """
+        The BatchKV cache expects inputs to be left-padded.
+
+        E.g. the following prompts:
+
+            [1, 3, 5]
+            [7]
+            [2, 6, 8, 9]
+
+        Should be padded like so:
+
+            [0, 1, 3, 5]
+            [0, 0, 0, 7]
+            [2, 6, 8, 9]
+
+        And ``left_padding`` specifies the amount of padding for each.
+        In this case, ``left_padding = [1, 3, 0]``.
+        """
+
+    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
+        ...
     @property
     def state(
         self,
@@ -301,21 +315,12 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step: int
-    keys: array | None
-    values: array | None
-    offset: array
-    left_padding: array
-    max_size: int
-    _idx: int
-    _offset: int
-    rotated: bool
-    _lengths: array | None
-    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
-    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
-    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
-    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
-    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
+    step = ...
+    def __init__(self, max_size, left_padding: List[int]) -> None: ...
+    def update_and_fetch(
+        self, keys, values
+    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
+        ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/deepseek_v3.pyi

@@ -5,7 +5,6 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx_lm.models.mla import MultiLinear
 
 from .base import BaseModelArgs
 from .switch_layers import SwitchGLU
@@ -61,10 +60,7 @@ class DeepseekV3Attention(nn.Module):
     q_b_proj: nn.Linear
     kv_a_proj_with_mqa: nn.Linear
     kv_a_layernorm: nn.RMSNorm
-    # kv_b_proj: nn.Linear
-    embed_q: MultiLinear
-    unembed_out: MultiLinear
-
+    kv_b_proj: nn.Linear
     o_proj: nn.Linear
     rope: Any
 

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -73,9 +73,6 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array: ...
 
 class SwitchMLP(nn.Module):
-    fc1: SwitchLinear
-    fc2: SwitchLinear
-
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/sample_utils.pyi

@@ -48,11 +48,7 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
-    presence_penalty: Optional[float] = ...,
-    presence_context_size: Optional[int] = ...,
-    frequency_penalty: Optional[float] = ...,
-    frequency_context_size: Optional[int] = ...,
-) -> list[Callable[[mx.array, mx.array], mx.array]]:
+):  # -> list[Any]:
     """
     Make logits processors for use with ``generate_step``.
 

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -39,11 +39,11 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    def reset(self) -> None: ...
-    def add_token(self, token: int) -> None: ...
-    def finalize(self) -> None: ...
+    def reset(self): ...
+    def add_token(self, token): ...
+    def finalize(self): ...
     @property
-    def last_segment(self) -> str:
+    def last_segment(self):
         """Return the last segment of readable text since last time this property was accessed."""
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
@@ -108,23 +108,16 @@ class TokenizerWrapper:
     _tokenizer: PreTrainedTokenizerFast
     eos_token_id: int | None
     eos_token: str | None
-    eos_token_ids: list[int] | set[int] | None
     bos_token_id: int | None
     bos_token: str | None
     vocab_size: int
     all_special_tokens: list[str]
-    think_start: str | None
-    think_end: str | None
-    think_start_id: int | None
-    think_end_id: int | None
-    think_start_tokens: list[int] | None
-    think_end_tokens: list[int] | None
 
     def __init__(
         self,
         tokenizer: Any,
         detokenizer_class: Any = ...,
-        eos_token_ids: list[int] | set[int] | None = ...,
+        eos_token_ids: list[int] | None = ...,
         chat_template: Any = ...,
         tool_parser: Any = ...,
         tool_call_start: str | None = ...,

.prettierrc

@@ -0,0 +1,3 @@
+{
+  "useTabs": true
+}

.typings/mflux/__init__.pyi

@@ -1,7 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import os
-
-if "TOKENIZERS_PARALLELISM" not in os.environ: ...

.typings/mflux/callbacks/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/callbacks/callback.pyi

@@ -1,47 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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):
-    def call_before_loop(
-        self,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        canny_image: PIL.Image.Image | None = ...,
-        depth_image: PIL.Image.Image | None = ...,
-    ) -> None: ...
-
-class InLoopCallback(Protocol):
-    def call_in_loop(
-        self,
-        t: int,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        time_steps: tqdm,
-    ) -> None: ...
-
-class AfterLoopCallback(Protocol):
-    def call_after_loop(
-        self, seed: int, prompt: str, latents: mx.array, config: Config
-    ) -> None: ...
-
-class InterruptCallback(Protocol):
-    def call_interrupt(
-        self,
-        t: int,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        time_steps: tqdm,
-    ) -> None: ...

.typings/mflux/callbacks/callback_registry.pyi

@@ -1,24 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import TYPE_CHECKING
-from mflux.callbacks.callback import (
-    AfterLoopCallback,
-    BeforeLoopCallback,
-    InLoopCallback,
-    InterruptCallback,
-)
-from mflux.callbacks.generation_context import GenerationContext
-from mflux.models.common.config.config import Config
-
-if TYPE_CHECKING: ...
-
-class CallbackRegistry:
-    def __init__(self) -> None: ...
-    def register(self, callback) -> None: ...
-    def start(self, seed: int, prompt: str, config: Config) -> GenerationContext: ...
-    def before_loop_callbacks(self) -> list[BeforeLoopCallback]: ...
-    def in_loop_callbacks(self) -> list[InLoopCallback]: ...
-    def after_loop_callbacks(self) -> list[AfterLoopCallback]: ...
-    def interrupt_callbacks(self) -> list[InterruptCallback]: ...

.typings/mflux/callbacks/generation_context.pyi

@@ -1,29 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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
-
-if TYPE_CHECKING: ...
-
-class GenerationContext:
-    def __init__(
-        self, registry: CallbackRegistry, seed: int, prompt: str, config: Config
-    ) -> None: ...
-    def before_loop(
-        self,
-        latents: mx.array,
-        *,
-        canny_image: PIL.Image.Image | None = ...,
-        depth_image: PIL.Image.Image | None = ...,
-    ) -> None: ...
-    def in_loop(self, t: int, latents: mx.array, time_steps: tqdm = ...) -> None: ...
-    def after_loop(self, latents: mx.array) -> None: ...
-    def interruption(
-        self, t: int, latents: mx.array, time_steps: tqdm = ...
-    ) -> None: ...

.typings/mflux/cli/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/cli/defaults/defaults.pyi

@@ -1,22 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import os
-
-BATTERY_PERCENTAGE_STOP_LIMIT = ...
-CONTROLNET_STRENGTH = ...
-DEFAULT_DEV_FILL_GUIDANCE = ...
-DEFAULT_DEPTH_GUIDANCE = ...
-DIMENSION_STEP_PIXELS = ...
-GUIDANCE_SCALE = ...
-GUIDANCE_SCALE_KONTEXT = ...
-IMAGE_STRENGTH = ...
-MODEL_CHOICES = ...
-MODEL_INFERENCE_STEPS = ...
-QUANTIZE_CHOICES = ...
-if os.environ.get("MFLUX_CACHE_DIR"):
-    MFLUX_CACHE_DIR = ...
-else:
-    MFLUX_CACHE_DIR = ...
-MFLUX_LORA_CACHE_DIR = ...

.typings/mflux/models/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/cli/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/config/__init__.pyi

@@ -1,8 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.config.config import Config
-from mflux.models.common.config.model_config import ModelConfig
-
-__all__ = ["Config", "ModelConfig"]

.typings/mflux/models/common/config/config.pyi

@@ -1,66 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from pathlib import Path
-from typing import Any
-from tqdm import tqdm
-from mflux.models.common.config.model_config import ModelConfig
-
-logger = ...
-
-class Config:
-    def __init__(
-        self,
-        model_config: ModelConfig,
-        num_inference_steps: int = ...,
-        height: int = ...,
-        width: int = ...,
-        guidance: float = ...,
-        image_path: Path | str | None = ...,
-        image_strength: float | None = ...,
-        depth_image_path: Path | str | None = ...,
-        redux_image_paths: list[Path | str] | None = ...,
-        redux_image_strengths: list[float] | None = ...,
-        masked_image_path: Path | str | None = ...,
-        controlnet_strength: float | None = ...,
-        scheduler: str = ...,
-    ) -> None: ...
-    @property
-    def height(self) -> int: ...
-    @property
-    def width(self) -> int: ...
-    @width.setter
-    def width(self, value):  # -> None:
-        ...
-    @property
-    def image_seq_len(self) -> int: ...
-    @property
-    def guidance(self) -> float: ...
-    @property
-    def num_inference_steps(self) -> int: ...
-    @property
-    def precision(self) -> mx.Dtype: ...
-    @property
-    def num_train_steps(self) -> int: ...
-    @property
-    def image_path(self) -> Path | None: ...
-    @property
-    def image_strength(self) -> float | None: ...
-    @property
-    def depth_image_path(self) -> Path | None: ...
-    @property
-    def redux_image_paths(self) -> list[Path] | None: ...
-    @property
-    def redux_image_strengths(self) -> list[float] | None: ...
-    @property
-    def masked_image_path(self) -> Path | None: ...
-    @property
-    def init_time_step(self) -> int: ...
-    @property
-    def time_steps(self) -> tqdm: ...
-    @property
-    def controlnet_strength(self) -> float | None: ...
-    @property
-    def scheduler(self) -> Any: ...

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

@@ -1,86 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from functools import lru_cache
-from typing import Literal
-
-class ModelConfig:
-    precision: mx.Dtype = ...
-    def __init__(
-        self,
-        priority: int,
-        aliases: list[str],
-        model_name: str,
-        base_model: str | None,
-        controlnet_model: str | None,
-        custom_transformer_model: str | None,
-        num_train_steps: int | None,
-        max_sequence_length: int | None,
-        supports_guidance: bool | None,
-        requires_sigma_shift: bool | None,
-        transformer_overrides: dict | None = ...,
-    ) -> None: ...
-    @staticmethod
-    @lru_cache
-    def dev() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def schnell() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_kontext() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_fill() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_redux() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_depth() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_controlnet_canny() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def schnell_controlnet_canny() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_controlnet_upscaler() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_fill_catvton() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def krea_dev() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def flux2_klein_4b() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def flux2_klein_9b() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def qwen_image() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def qwen_image_edit() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def fibo() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def z_image_turbo() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def seedvr2_3b() -> ModelConfig: ...
-    def x_embedder_input_dim(self) -> int: ...
-    def is_canny(self) -> bool: ...
-    @staticmethod
-    def from_name(
-        model_name: str, base_model: Literal["dev", "schnell", "krea-dev"] | None = ...
-    ) -> ModelConfig: ...
-
-AVAILABLE_MODELS = ...

.typings/mflux/models/common/latent_creator/__init__.pyi

@@ -1,7 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/latent_creator/latent_creator.pyi

@@ -1,49 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from pathlib import Path
-from typing import TYPE_CHECKING, TypeAlias
-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
-from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator
-from mflux.models.z_image.latent_creator.z_image_latent_creator import (
-    ZImageLatentCreator,
-)
-
-if TYPE_CHECKING:
-    LatentCreatorType: TypeAlias = type[
-        FiboLatentCreator | FluxLatentCreator | QwenLatentCreator | ZImageLatentCreator
-    ]
-
-class Img2Img:
-    def __init__(
-        self,
-        vae: nn.Module,
-        latent_creator: LatentCreatorType,
-        sigmas: mx.array,
-        init_time_step: int,
-        image_path: str | Path | None,
-        tiling_config: TilingConfig | None = ...,
-    ) -> None: ...
-
-class LatentCreator:
-    @staticmethod
-    def create_for_txt2img_or_img2img(
-        seed: int, height: int, width: int, img2img: Img2Img
-    ) -> mx.array: ...
-    @staticmethod
-    def encode_image(
-        vae: nn.Module,
-        image_path: str | Path,
-        height: int,
-        width: int,
-        tiling_config: TilingConfig | None = ...,
-    ) -> mx.array: ...
-    @staticmethod
-    def add_noise_by_interpolation(
-        clean: mx.array, noise: mx.array, sigma: float
-    ) -> mx.array: ...

.typings/mflux/models/common/lora/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

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

@@ -1,13 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mlx import nn
-from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
-
-class FusedLoRALinear(nn.Module):
-    def __init__(
-        self, base_linear: nn.Linear | nn.QuantizedLinear, loras: list[LoRALinear]
-    ) -> None: ...
-    def __call__(self, x):  # -> array:
-        ...

.typings/mflux/models/common/lora/layer/linear_lora_layer.pyi

@@ -1,22 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mlx import nn
-
-class LoRALinear(nn.Module):
-    @staticmethod
-    def from_linear(
-        linear: nn.Linear | nn.QuantizedLinear, r: int = ..., scale: float = ...
-    ):  # -> LoRALinear:
-        ...
-    def __init__(
-        self,
-        input_dims: int,
-        output_dims: int,
-        r: int = ...,
-        scale: float = ...,
-        bias: bool = ...,
-    ) -> None: ...
-    def __call__(self, x):  # -> array:
-        ...

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

@@ -1,26 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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
-class PatternMatch:
-    source_pattern: str
-    target_path: str
-    matrix_name: str
-    transpose: bool
-    transform: Callable[[mx.array], mx.array] | None = ...
-
-class LoRALoader:
-    @staticmethod
-    def load_and_apply_lora(
-        lora_mapping: list[LoRATarget],
-        transformer: nn.Module,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> tuple[list[str], list[float]]: ...

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

@@ -1,21 +0,0 @@
-"""
-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
-
-@dataclass
-class LoRATarget:
-    model_path: str
-    possible_up_patterns: List[str]
-    possible_down_patterns: List[str]
-    possible_alpha_patterns: List[str] = ...
-    up_transform: Callable[[mx.array], mx.array] | None = ...
-    down_transform: Callable[[mx.array], mx.array] | None = ...
-
-class LoRAMapping(Protocol):
-    @staticmethod
-    def get_mapping() -> List[LoRATarget]: ...

.typings/mflux/models/common/lora/mapping/lora_saver.pyi

@@ -1,9 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.nn as nn
-
-class LoRASaver:
-    @staticmethod
-    def bake_and_strip_lora(module: nn.Module) -> nn.Module: ...

.typings/mflux/models/common/lora/mapping/lora_transforms.pyi

@@ -1,35 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-
-class LoraTransforms:
-    @staticmethod
-    def split_q_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_k_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_v_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_q_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_k_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_v_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_q_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_k_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_v_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_mlp_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_q_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_k_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_v_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_mlp_down(tensor: mx.array) -> mx.array: ...

.typings/mflux/models/common/resolution/__init__.pyi

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.resolution.config_resolution import ConfigResolution
-from mflux.models.common.resolution.lora_resolution import LoraResolution
-from mflux.models.common.resolution.path_resolution import PathResolution
-from mflux.models.common.resolution.quantization_resolution import (
-    QuantizationResolution,
-)
-
-__all__ = [
-    "ConfigResolution",
-    "LoraResolution",
-    "PathResolution",
-    "QuantizationResolution",
-]

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

@@ -1,39 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from enum import Enum
-from typing import NamedTuple
-
-class QuantizationAction(Enum):
-    NONE = ...
-    STORED = ...
-    REQUESTED = ...
-
-class PathAction(Enum):
-    LOCAL = ...
-    HUGGINGFACE_CACHED = ...
-    HUGGINGFACE = ...
-    ERROR = ...
-
-class LoraAction(Enum):
-    LOCAL = ...
-    REGISTRY = ...
-    HUGGINGFACE_COLLECTION_CACHED = ...
-    HUGGINGFACE_COLLECTION = ...
-    HUGGINGFACE_REPO_CACHED = ...
-    HUGGINGFACE_REPO = ...
-    ERROR = ...
-
-class ConfigAction(Enum):
-    EXACT_MATCH = ...
-    EXPLICIT_BASE = ...
-    INFER_SUBSTRING = ...
-    ERROR = ...
-
-class Rule(NamedTuple):
-    priority: int
-    name: str
-    check: str
-    action: QuantizationAction | PathAction | LoraAction | ConfigAction
-    ...

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

@@ -1,14 +0,0 @@
-"""
-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: ...
-logger = ...
-
-class ConfigResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(model_name: str, base_model: str | None = ...) -> ModelConfig: ...

.typings/mflux/models/common/resolution/lora_resolution.pyi

@@ -1,21 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-
-logger = ...
-
-class LoraResolution:
-    RULES = ...
-    _registry: dict[str, Path] = ...
-    @staticmethod
-    def resolve(path: str) -> str: ...
-    @staticmethod
-    def resolve_paths(paths: list[str] | None) -> list[str]: ...
-    @staticmethod
-    def resolve_scales(scales: list[float] | None, num_paths: int) -> list[float]: ...
-    @staticmethod
-    def get_registry() -> dict[str, Path]: ...
-    @staticmethod
-    def discover_files(library_paths: list[Path]) -> dict[str, Path]: ...

.typings/mflux/models/common/resolution/path_resolution.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-
-logger = ...
-
-class PathResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(path: str | None, patterns: list[str] | None = ...) -> Path | None: ...

.typings/mflux/models/common/resolution/quantization_resolution.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-logger = ...
-
-class QuantizationResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(
-        stored: int | None, requested: int | None
-    ) -> tuple[int | None, str | None]: ...

.typings/mflux/models/common/schedulers/__init__.pyi

@@ -1,26 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from .flow_match_euler_discrete_scheduler import FlowMatchEulerDiscreteScheduler
-from .linear_scheduler import LinearScheduler
-from .seedvr2_euler_scheduler import SeedVR2EulerScheduler
-
-__all__ = [
-    "LinearScheduler",
-    "FlowMatchEulerDiscreteScheduler",
-    "SeedVR2EulerScheduler",
-]
-
-class SchedulerModuleNotFound(ValueError): ...
-class SchedulerClassNotFound(ValueError): ...
-class InvalidSchedulerType(TypeError): ...
-
-SCHEDULER_REGISTRY = ...
-
-def register_contrib(scheduler_object, scheduler_name=...):  # -> None:
-    ...
-def try_import_external_scheduler(
-    scheduler_object_path: str,
-):  # -> type[BaseScheduler]:
-    ...

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

@@ -1,16 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from abc import ABC, abstractmethod
-
-class BaseScheduler(ABC):
-    @property
-    @abstractmethod
-    def sigmas(self) -> mx.array: ...
-    @abstractmethod
-    def step(
-        self, noise: mx.array, timestep: int, latents: mx.array, **kwargs
-    ) -> mx.array: ...
-    def scale_model_input(self, latents: mx.array, t: int) -> mx.array: ...

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

@@ -1,26 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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
-
-if TYPE_CHECKING: ...
-
-class FlowMatchEulerDiscreteScheduler(BaseScheduler):
-    def __init__(self, config: Config) -> None: ...
-    @property
-    def sigmas(self) -> mx.array: ...
-    @property
-    def timesteps(self) -> mx.array: ...
-    def set_image_seq_len(self, image_seq_len: int) -> None: ...
-    @staticmethod
-    def get_timesteps_and_sigmas(
-        image_seq_len: int, num_inference_steps: int, num_train_timesteps: int = ...
-    ) -> tuple[mx.array, mx.array]: ...
-    def step(
-        self, noise: mx.array, timestep: int, latents: mx.array, **kwargs
-    ) -> mx.array: ...
-    def scale_model_input(self, latents: mx.array, t: int) -> mx.array: ...

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

@@ -1,20 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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
-
-if TYPE_CHECKING: ...
-
-class LinearScheduler(BaseScheduler):
-    def __init__(self, config: Config) -> None: ...
-    @property
-    def sigmas(self) -> mx.array: ...
-    @property
-    def timesteps(self) -> mx.array: ...
-    def step(
-        self, noise: mx.array, timestep: int, latents: mx.array, **kwargs
-    ) -> mx.array: ...

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

@@ -1,20 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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
-
-if TYPE_CHECKING: ...
-
-class SeedVR2EulerScheduler(BaseScheduler):
-    def __init__(self, config: Config) -> None: ...
-    @property
-    def timesteps(self) -> mx.array: ...
-    @property
-    def sigmas(self) -> mx.array: ...
-    def step(
-        self, noise: mx.array, timestep: int, latents: mx.array, **kwargs
-    ) -> mx.array: ...

.typings/mflux/models/common/tokenizer/__init__.pyi

@@ -1,24 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.tokenizer.tokenizer import (
-    BaseTokenizer,
-    LanguageTokenizer,
-    Tokenizer,
-    VisionLanguageTokenizer,
-)
-from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader
-from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
-
-"""
-This type stub file was generated by pyright.
-"""
-__all__ = [
-    "Tokenizer",
-    "BaseTokenizer",
-    "LanguageTokenizer",
-    "VisionLanguageTokenizer",
-    "TokenizerLoader",
-    "TokenizerOutput",
-]

.typings/mflux/models/common/tokenizer/tokenizer.pyi

@@ -1,74 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from abc import ABC, abstractmethod
-from typing import Protocol, runtime_checkable
-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.
-"""
-
-@runtime_checkable
-class Tokenizer(Protocol):
-    tokenizer: PreTrainedTokenizer
-    def tokenize(
-        self,
-        prompt: str | list[str],
-        images: list[Image.Image] | None = ...,
-        max_length: int | None = ...,
-        **kwargs,
-    ) -> TokenizerOutput: ...
-
-class BaseTokenizer(ABC):
-    def __init__(
-        self, tokenizer: PreTrainedTokenizer, max_length: int = ...
-    ) -> None: ...
-    @abstractmethod
-    def tokenize(
-        self,
-        prompt: str | list[str],
-        images: list[Image.Image] | None = ...,
-        max_length: int | None = ...,
-        **kwargs,
-    ) -> TokenizerOutput: ...
-
-class LanguageTokenizer(BaseTokenizer):
-    def __init__(
-        self,
-        tokenizer: PreTrainedTokenizer,
-        max_length: int = ...,
-        padding: str = ...,
-        return_attention_mask: bool = ...,
-        template: str | None = ...,
-        use_chat_template: bool = ...,
-        chat_template_kwargs: dict | None = ...,
-        add_special_tokens: bool = ...,
-    ) -> None: ...
-    def tokenize(
-        self,
-        prompt: str | list[str],
-        images: list[Image.Image] | None = ...,
-        max_length: int | None = ...,
-        **kwargs,
-    ) -> TokenizerOutput: ...
-
-class VisionLanguageTokenizer(BaseTokenizer):
-    def __init__(
-        self,
-        tokenizer: PreTrainedTokenizer,
-        processor,
-        max_length: int = ...,
-        template: str | None = ...,
-        image_token: str = ...,
-    ) -> None: ...
-    def tokenize(
-        self,
-        prompt: str | list[str],
-        images: list[Image.Image] | None = ...,
-        max_length: int | None = ...,
-        **kwargs,
-    ) -> TokenizerOutput: ...

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

@@ -1,22 +0,0 @@
-"""
-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
-
-"""
-This type stub file was generated by pyright.
-"""
-if TYPE_CHECKING: ...
-
-class TokenizerLoader:
-    @staticmethod
-    def load(definition: TokenizerDefinition, model_path: str) -> BaseTokenizer: ...
-    @staticmethod
-    def load_all(
-        definitions: list[TokenizerDefinition],
-        model_path: str,
-        max_length_overrides: dict[str, int] | None = ...,
-    ) -> dict[str, BaseTokenizer]: ...

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

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from dataclasses import dataclass
-
-"""
-This type stub file was generated by pyright.
-"""
-
-@dataclass
-class TokenizerOutput:
-    input_ids: mx.array
-    attention_mask: mx.array
-    pixel_values: mx.array | None = ...
-    image_grid_thw: mx.array | None = ...

.typings/mflux/models/common/vae/__init__.pyi

@@ -1,8 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.vae.tiling_config import TilingConfig
-from mflux.models.common.vae.vae_tiler import VAETiler
-
-__all__ = ["TilingConfig", "VAETiler"]

.typings/mflux/models/common/vae/tiling_config.pyi

@@ -1,13 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from dataclasses import dataclass
-
-@dataclass(frozen=True, slots=True)
-class TilingConfig:
-    vae_decode_tiles_per_dim: int | None = ...
-    vae_decode_overlap: int = ...
-    vae_encode_tiled: bool = ...
-    vae_encode_tile_size: int = ...
-    vae_encode_tile_overlap: int = ...

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

@@ -1,27 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from typing import Callable
-
-class VAETiler:
-    @staticmethod
-    def encode_image_tiled(
-        *,
-        image: mx.array,
-        encode_fn: Callable[[mx.array], mx.array],
-        latent_channels: int,
-        tile_size: tuple[int, int] = ...,
-        tile_overlap: tuple[int, int] = ...,
-        spatial_scale: int = ...,
-    ) -> mx.array: ...
-    @staticmethod
-    def decode_image_tiled(
-        *,
-        latent: mx.array,
-        decode_fn: Callable[[mx.array], mx.array],
-        tile_size: tuple[int, int] = ...,
-        tile_overlap: tuple[int, int] = ...,
-        spatial_scale: int = ...,
-    ) -> mx.array: ...

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

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-from mflux.models.common.vae.tiling_config import TilingConfig
-
-class VAEUtil:
-    @staticmethod
-    def encode(
-        vae: nn.Module, image: mx.array, tiling_config: TilingConfig | None = ...
-    ) -> mx.array: ...
-    @staticmethod
-    def decode(
-        vae: nn.Module, latent: mx.array, tiling_config: TilingConfig | None = ...
-    ) -> mx.array: ...

.typings/mflux/models/common/weights/__init__.pyi

@@ -1,18 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.weights.loading.loaded_weights import LoadedWeights, MetaData
-from mflux.models.common.weights.loading.weight_applier import WeightApplier
-from mflux.models.common.weights.loading.weight_definition import ComponentDefinition
-from mflux.models.common.weights.loading.weight_loader import WeightLoader
-from mflux.models.common.weights.saving.model_saver import ModelSaver
-
-__all__ = [
-    "ComponentDefinition",
-    "LoadedWeights",
-    "MetaData",
-    "ModelSaver",
-    "WeightApplier",
-    "WeightLoader",
-]

.typings/mflux/models/common/weights/loading/loaded_weights.pyi

@@ -1,18 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from dataclasses import dataclass
-
-@dataclass
-class MetaData:
-    quantization_level: int | None = ...
-    mflux_version: str | None = ...
-
-@dataclass
-class LoadedWeights:
-    components: dict[str, dict]
-    meta_data: MetaData
-    def __getattr__(self, name: str) -> dict | None: ...
-    def num_transformer_blocks(self, component_name: str = ...) -> int: ...
-    def num_single_transformer_blocks(self, component_name: str = ...) -> int: ...

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

@@ -1,30 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-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,
-    WeightDefinitionType,
-)
-
-if TYPE_CHECKING: ...
-
-class WeightApplier:
-    @staticmethod
-    def apply_and_quantize_single(
-        weights: LoadedWeights,
-        model: nn.Module,
-        component: ComponentDefinition,
-        quantize_arg: int | None,
-        quantization_predicate=...,
-    ) -> int | None: ...
-    @staticmethod
-    def apply_and_quantize(
-        weights: LoadedWeights,
-        models: dict[str, nn.Module],
-        quantize_arg: int | None,
-        weight_definition: WeightDefinitionType,
-    ) -> int | None: ...

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

@@ -1,73 +0,0 @@
-"""
-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 = ...

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

@@ -1,23 +0,0 @@
-"""
-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,
-    WeightDefinitionType,
-)
-
-if TYPE_CHECKING: ...
-logger = ...
-
-class WeightLoader:
-    @staticmethod
-    def load_single(
-        component: ComponentDefinition, repo_id: str, file_pattern: str = ...
-    ) -> LoadedWeights: ...
-    @staticmethod
-    def load(
-        weight_definition: WeightDefinitionType, model_path: str | None = ...
-    ) -> LoadedWeights: ...

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

@@ -1,16 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from typing import Dict, List, Optional
-from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
-
-class WeightMapper:
-    @staticmethod
-    def apply_mapping(
-        hf_weights: Dict[str, mx.array],
-        mapping: List[WeightTarget],
-        num_blocks: Optional[int] = ...,
-        num_layers: Optional[int] = ...,
-    ) -> Dict: ...

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

@@ -1,23 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from dataclasses import dataclass
-from typing import Callable, List, Optional, Protocol
-
-"""
-This type stub file was generated by pyright.
-"""
-
-@dataclass
-class WeightTarget:
-    to_pattern: str
-    from_pattern: List[str]
-    transform: Optional[Callable[[mx.array], mx.array]] = ...
-    required: bool = ...
-    max_blocks: Optional[int] = ...
-
-class WeightMapping(Protocol):
-    @staticmethod
-    def get_mapping() -> List[WeightTarget]: ...

.typings/mflux/models/common/weights/mapping/weight_transforms.pyi

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-
-class WeightTransforms:
-    @staticmethod
-    def reshape_gamma_to_1d(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def transpose_patch_embed(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def transpose_conv3d_weight(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def transpose_conv2d_weight(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def transpose_conv_transpose2d_weight(tensor: mx.array) -> mx.array: ...

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

@@ -1,14 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import Any, TYPE_CHECKING
-from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType
-
-if TYPE_CHECKING: ...
-
-class ModelSaver:
-    @staticmethod
-    def save_model(
-        model: Any, bits: int, base_path: str, weight_definition: WeightDefinitionType
-    ) -> None: ...