update log message + assertion

Dashboard fixes (TypeScript errors from `npm run check`):
- TopologyGraph.svelte: remove reference to deleted sendBackMultiaddr
  property, fix type inference for debug edge labels
- ModelCard.svelte: add missing topoWidth/topoHeight to early return
- +page.svelte: fix nested property access for deviceRank

Backend fix:
- info_gatherer.py: send initial MiscData on startup so friendly name
  appears immediately instead of showing "Unknown" until it changes

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

.github/workflows/build-app.yml

@@ -1,16 +1,5 @@
 name: Build EXO macOS DMG
 
-# Release workflow:
-# 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown
-# 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0
-# 3. This workflow builds, signs, and notarizes the DMG
-# 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown)
-# 5. DMG and appcast.xml are uploaded to S3
-# 6. The draft GitHub Release is published with the DMG attached
-#
-# For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional.
-# If no draft exists, a release is auto-created with generated notes.
-
 on:
   workflow_dispatch:
   push:
@@ -22,10 +11,8 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
-    permissions:
-      contents: write
     env:
-      SPARKLE_VERSION: 2.9.0-beta.1
+      SPARKLE_VERSION: 2.8.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
       SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }}
       SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }}
@@ -100,52 +87,6 @@ jobs:
             exit 1
           fi
 
-      - name: Fetch and validate release notes
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          # Find draft release by name using gh release list (more reliable with default token)
-          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
-          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
-
-          if [[ -z "$DRAFT_EXISTS" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: No draft release found for tag $GITHUB_REF_NAME"
-            echo "Please create a draft release with release notes before pushing the tag."
-            exit 1
-          fi
-
-          # Fetch full release details via API to get body and ID
-          echo "Found draft release, fetching details..."
-          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
-
-          # Extract release notes
-          NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
-          if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "Draft release has no notes (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: Draft release exists but has no release notes"
-            echo "Please add release notes to the draft release before pushing the tag."
-            exit 1
-          fi
-
-          # Save release ID for later publishing
-          RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id')
-          echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV
-          echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV
-
-          echo "Found draft release (ID: $RELEASE_ID), saving release notes..."
-          echo "$NOTES" > /tmp/release_notes.md
-          echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV
-
       # ============================================================
       # Install dependencies
       # ============================================================
@@ -172,22 +113,11 @@ jobs:
           uv python install
           uv sync --locked
 
-      - name: Install Nix
-        uses: cachix/install-nix-action@v31
-        with:
-          nix_path: nixpkgs=channel:nixos-unstable
-
-      - name: Configure Cachix
-        uses: cachix/cachix-action@v14
-        with:
-          name: exo
-          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
-
       - name: Build dashboard
         run: |
-          DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths)
+          cd dashboard
-          mkdir -p dashboard/build
+          npm ci
-          cp -r "$DASHBOARD_OUT"/* dashboard/build/
+          npm run build
 
       - name: Install Sparkle CLI
         run: |
@@ -363,28 +293,6 @@ jobs:
             $CHANNEL_FLAG \
             .
 
-      - name: Inject release notes into appcast
-        if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true'
-        env:
-          RELEASE_VERSION: ${{ env.RELEASE_VERSION }}
-        run: |
-          # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+)
-          export NOTES=$(cat "$RELEASE_NOTES_FILE")
-
-          # Insert description after the enclosure tag for this version
-          awk '
-            /<enclosure[^>]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) {
-              print
-              print "            <description sparkle:format=\"markdown\"><![CDATA["
-              print ENVIRON["NOTES"]
-              print "            ]]></description>"
-              next
-            }
-            { print }
-          ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml
-
-          echo "Injected markdown release notes for version $RELEASE_VERSION"
-
       # ============================================================
       # Upload artifacts
       # ============================================================
@@ -417,26 +325,3 @@ jobs:
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
             aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache
           fi
-
-      - name: Publish GitHub Release
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg"
-
-          if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then
-            # Update the draft release with the tag and upload DMG
-            gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \
-              -f tag_name="$GITHUB_REF_NAME" \
-              -F draft=false
-            gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber
-            echo "Published release $GITHUB_REF_NAME with DMG attached"
-          else
-            # Alpha without draft release - create one with auto-generated notes
-            gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \
-              --title "$GITHUB_REF_NAME" \
-              --generate-notes \
-              --prerelease
-            echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes"
-          fi

.github/workflows/pipeline.yml

@@ -20,12 +20,6 @@ jobs:
         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"
@@ -94,19 +88,9 @@ jobs:
 
       - uses: ./.github/actions/typecheck
 
-  nix:
+  nix-flake-check:
-    name: Build and check (${{ matrix.system }})
+    name: Check Nix flake
-    runs-on: ${{ matrix.runner }}
+    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - runner: macos-26
-            system: aarch64-darwin
-          - runner: ubuntu-latest
-            system: x86_64-linux
-          - runner: ubuntu-24.04-arm
-            system: aarch64-linux
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -117,20 +101,83 @@ jobs:
         with:
           nix_path: nixpkgs=channel:nixos-unstable
 
-      - uses: cachix/cachix-action@v14
-        name: Configure Cachix
-        with:
-          name: exo
-          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
-
-      - name: Build all Nix outputs
-        run: |
-          nix flake show --json | jq -r '
-            [
-              (.packages."${{ matrix.system }}" // {} | keys[] | ".#packages.${{ matrix.system }}.\(.)"),
-              (.devShells."${{ matrix.system }}" // {} | keys[] | ".#devShells.${{ matrix.system }}.\(.)")
-            ] | .[]
-          ' | xargs nix build
-
       - name: Run nix flake check
-        run: nix flake check
+        run: |
+          nix flake check
+        shell: bash
+
+#  ci:
+#    needs: typecheck
+#    runs-on: ubuntu-latest
+#    permissions:
+#      contents: read
+#    env:
+#      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+#    steps:
+#      - name: Checkout repository
+#        uses: actions/checkout@v4
+#        with:
+#          fetch-depth: 0
+#          token: ${{ secrets.GITHUB_TOKEN }}
+#          lfs: true
+#
+#      - 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 EXO_HOME and API_PORT
+#        run: |
+#          EXO_HOME=$(mktemp -d -t exo-ci-XXXXXXXX)
+#          # Generate random port (macOS compatible method)
+#          API_PORT=$((49152 + RANDOM % (65535 - 49152 + 1)))
+#          echo "EXO_HOME=$EXO_HOME" >> $GITHUB_ENV
+#          echo "API_PORT=$API_PORT" >> $GITHUB_ENV
+#          echo "Created EXO_HOME: $EXO_HOME"
+#          echo "Generated API_PORT: $API_PORT"
+#        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
+#
+#      - uses: ./.github/actions/lint-check
+#
+#      - uses: ./.github/actions/unit-test
+#
+#      - name: Cleanup EXO_HOME
+#        run: |
+#          echo "Cleaning up EXO_HOME: $EXO_HOME"
+#          rm -rf "$EXO_HOME"
+#        shell: bash
+#        if: always()

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

.mlx_typings/mflux/callbacks/__init__.pyi

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

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

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

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

.mlx_typings/mflux/cli/__init__.pyi

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

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

.mlx_typings/mflux/models/__init__.pyi

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

.mlx_typings/mflux/models/common/__init__.pyi

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

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

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

.mlx_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"]

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

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

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

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

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

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

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

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

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

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

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

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

.mlx_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",
-]

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

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

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

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

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

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

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

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

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

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

.mlx_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",
-]

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

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

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

.mlx_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"]

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

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

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

.mlx_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",
-]

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

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

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

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

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

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

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

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

.mlx_typings/mflux/models/depth_pro/depth_pro_initializer.pyi

@@ -1,9 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.depth_pro.model.depth_pro_model import DepthProModel
-
-class DepthProInitializer:
-    @staticmethod
-    def init(model: DepthProModel, quantize: int | None = ...) -> None: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class FeatureFusionBlock2d(nn.Module):
-    def __init__(self, num_features: int, deconv: bool = ...) -> None: ...
-    def __call__(self, x0: mx.array, x1: mx.array | None = ...) -> mx.array: ...

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

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class MultiresConvDecoder(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(
-        self,
-        x0_latent: mx.array,
-        x1_latent: mx.array,
-        x0_features: mx.array,
-        x1_features: mx.array,
-        x_global_features: mx.array,
-    ) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class ResidualBlock(nn.Module):
-    def __init__(self, num_features: int) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,20 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from dataclasses import dataclass
-from pathlib import Path
-from PIL import Image
-
-@dataclass
-class DepthResult:
-    depth_image: Image.Image
-    depth_array: mx.array
-    min_depth: float
-    max_depth: float
-    ...
-
-class DepthPro:
-    def __init__(self, quantize: int | None = ...) -> None: ...
-    def create_depth_map(self, image_path: str | Path) -> DepthResult: ...

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

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class DepthProModel(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(
-        self, x0: mx.array, x1: mx.array, x2: mx.array
-    ) -> tuple[mx.array, mx.array]: ...

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

@@ -1,15 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class DepthProUtil:
-    @staticmethod
-    def split(x: mx.array, overlap_ratio: float = ...) -> mx.array: ...
-    @staticmethod
-    def interpolate(x: mx.array, size=..., scale_factor=...):  # -> array:
-        ...
-    @staticmethod
-    def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array: ...

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

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class Attention(nn.Module):
-    def __init__(
-        self, dim: int = ..., head_dim: int = ..., num_heads: int = ...
-    ) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class DinoVisionTransformer(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array, mx.array]: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class LayerScale(nn.Module):
-    def __init__(self, dims: int, init_values: float = ...) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class MLP(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class PatchEmbed(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class TransformerBlock(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class DepthProEncoder(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(
-        self, x0: mx.array, x1: mx.array, x2: mx.array
-    ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: ...

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

@@ -1,16 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class UpSampleBlock(nn.Module):
-    def __init__(
-        self,
-        dim_in: int = ...,
-        dim_int: int = ...,
-        dim_out: int = ...,
-        upsample_layers: int = ...,
-    ) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-import mlx.nn as nn
-
-class FOVHead(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, x: mx.array) -> mx.array: ...

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

@@ -1,23 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.loading.weight_definition import (
-    ComponentDefinition,
-    TokenizerDefinition,
-)
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class DepthProWeightDefinition:
-    @staticmethod
-    def get_components() -> List[ComponentDefinition]: ...
-    @staticmethod
-    def get_tokenizers() -> List[TokenizerDefinition]: ...
-    @staticmethod
-    def get_download_patterns() -> List[str]: ...
-    @staticmethod
-    def quantization_predicate(path: str, module) -> bool: ...

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

@@ -1,13 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.mapping.weight_mapping import (
-    WeightMapping,
-    WeightTarget,
-)
-
-class DepthProWeightMapping(WeightMapping):
-    @staticmethod
-    def get_mapping() -> List[WeightTarget]: ...

.mlx_typings/mflux/models/fibo/latent_creator/fibo_latent_creator.pyi

@@ -1,13 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-
-class FiboLatentCreator:
-    @staticmethod
-    def create_noise(seed: int, height: int, width: int) -> mx.array: ...
-    @staticmethod
-    def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: ...
-    @staticmethod
-    def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ...

.mlx_typings/mflux/models/fibo/weights/fibo_weight_definition.pyi

@@ -1,23 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.loading.weight_definition import (
-    ComponentDefinition,
-    TokenizerDefinition,
-)
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class FIBOWeightDefinition:
-    @staticmethod
-    def get_components() -> List[ComponentDefinition]: ...
-    @staticmethod
-    def get_tokenizers() -> List[TokenizerDefinition]: ...
-    @staticmethod
-    def get_download_patterns() -> List[str]: ...
-    @staticmethod
-    def quantization_predicate(path: str, module) -> bool: ...

.mlx_typings/mflux/models/fibo/weights/fibo_weight_mapping.pyi

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.mapping.weight_mapping import (
-    WeightMapping,
-    WeightTarget,
-)
-
-class FIBOWeightMapping(WeightMapping):
-    @staticmethod
-    def get_transformer_mapping() -> List[WeightTarget]: ...
-    @staticmethod
-    def get_text_encoder_mapping() -> List[WeightTarget]: ...
-    @staticmethod
-    def get_vae_mapping() -> List[WeightTarget]: ...

.mlx_typings/mflux/models/fibo_vlm/tokenizer/qwen2vl_image_processor.pyi

@@ -1,8 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.qwen.tokenizer.qwen_image_processor import QwenImageProcessor
-
-class Qwen2VLImageProcessor(QwenImageProcessor):
-    def __init__(self) -> None: ...

.mlx_typings/mflux/models/fibo_vlm/tokenizer/qwen2vl_processor.pyi

@@ -1,28 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import Optional, Union
-from PIL import Image
-
-class Qwen2VLProcessor:
-    def __init__(self, tokenizer) -> None: ...
-    def apply_chat_template(
-        self,
-        messages,
-        tokenize: bool = ...,
-        add_generation_prompt: bool = ...,
-        return_tensors: Optional[str] = ...,
-        return_dict: bool = ...,
-        **kwargs,
-    ):  # -> dict[Any, Any]:
-        ...
-    def __call__(
-        self,
-        text: Optional[Union[str, list[str]]] = ...,
-        images: Optional[Union[Image.Image, list[Image.Image]]] = ...,
-        padding: bool = ...,
-        return_tensors: Optional[str] = ...,
-        **kwargs,
-    ):  # -> dict[Any, Any]:
-        ...

.mlx_typings/mflux/models/fibo_vlm/weights/fibo_vlm_weight_definition.pyi

@@ -1,24 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.loading.weight_definition import (
-    ComponentDefinition,
-    TokenizerDefinition,
-)
-
-"""
-This type stub file was generated by pyright.
-"""
-QWEN2VL_CHAT_TEMPLATE = ...
-
-class FIBOVLMWeightDefinition:
-    @staticmethod
-    def get_components() -> List[ComponentDefinition]: ...
-    @staticmethod
-    def get_tokenizers() -> List[TokenizerDefinition]: ...
-    @staticmethod
-    def get_download_patterns() -> List[str]: ...
-    @staticmethod
-    def quantization_predicate(path: str, module) -> bool: ...

.mlx_typings/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.pyi

@@ -1,15 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import List
-from mflux.models.common.weights.mapping.weight_mapping import (
-    WeightMapping,
-    WeightTarget,
-)
-
-class FIBOVLMWeightMapping(WeightMapping):
-    @staticmethod
-    def get_vlm_decoder_mapping(num_layers: int = ...) -> List[WeightTarget]: ...
-    @staticmethod
-    def get_vlm_visual_mapping(depth: int = ...) -> List[WeightTarget]: ...

.mlx_typings/mflux/models/flux/__init__.pyi

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

.mlx_typings/mflux/models/flux/cli/__init__.pyi

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

.mlx_typings/mflux/models/flux/flux_initializer.pyi

@@ -1,53 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.config import ModelConfig
-
-class FluxInitializer:
-    @staticmethod
-    def init(
-        model,
-        model_config: ModelConfig,
-        quantize: int | None,
-        model_path: str | None = ...,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-        custom_transformer=...,
-    ) -> None: ...
-    @staticmethod
-    def init_depth(
-        model,
-        model_config: ModelConfig,
-        quantize: int | None,
-        model_path: str | None = ...,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> None: ...
-    @staticmethod
-    def init_redux(
-        model,
-        model_config: ModelConfig,
-        quantize: int | None,
-        model_path: str | None = ...,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> None: ...
-    @staticmethod
-    def init_controlnet(
-        model,
-        model_config: ModelConfig,
-        quantize: int | None,
-        model_path: str | None = ...,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> None: ...
-    @staticmethod
-    def init_concept(
-        model,
-        model_config: ModelConfig,
-        quantize: int | None,
-        model_path: str | None = ...,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> None: ...

.mlx_typings/mflux/models/flux/latent_creator/__init__.pyi

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

.mlx_typings/mflux/models/flux/latent_creator/flux_latent_creator.pyi

@@ -1,19 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class FluxLatentCreator:
-    @staticmethod
-    def create_noise(seed: int, height: int, width: int) -> mx.array: ...
-    @staticmethod
-    def pack_latents(
-        latents: mx.array, height: int, width: int, num_channels_latents: int = ...
-    ) -> mx.array: ...
-    @staticmethod
-    def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/__init__.pyi

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

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class CLIPEmbeddings(nn.Module):
-    def __init__(self, dims: int) -> None: ...
-    def __call__(self, tokens: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder.pyi

@@ -1,14 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class CLIPEncoder(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, tokens: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder_layer.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class CLIPEncoderLayer(nn.Module):
-    def __init__(self, layer: int) -> None: ...
-    def __call__(
-        self, hidden_states: mx.array, causal_attention_mask: mx.array
-    ) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_mlp.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class CLIPMLP(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...
-    @staticmethod
-    def quick_gelu(input_array: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.pyi

@@ -1,18 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class CLIPSdpaAttention(nn.Module):
-    head_dimension = ...
-    batch_size = ...
-    num_heads = ...
-    def __init__(self) -> None: ...
-    def __call__(
-        self, hidden_states: mx.array, causal_attention_mask: mx.array
-    ) -> mx.array: ...
-    @staticmethod
-    def reshape_and_transpose(x, batch_size, num_heads, head_dim):  # -> array:
-        ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_text_model.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class CLIPTextModel(nn.Module):
-    def __init__(self, dims: int, num_encoder_layers: int) -> None: ...
-    def __call__(self, tokens: mx.array) -> tuple[mx.array, mx.array]: ...
-    @staticmethod
-    def create_causal_attention_mask(input_shape: tuple) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/encoder_clip.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class EncoderCLIP(nn.Module):
-    def __init__(self, num_encoder_layers: int) -> None: ...
-    def __call__(
-        self, tokens: mx.array, causal_attention_mask: mx.array
-    ) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/prompt_encoder.pyi

@@ -1,25 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mflux.models.common.tokenizer import Tokenizer
-from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import (
-    CLIPEncoder,
-)
-from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class PromptEncoder:
-    @staticmethod
-    def encode_prompt(
-        prompt: str,
-        prompt_cache: dict[str, tuple[mx.array, mx.array]],
-        t5_tokenizer: Tokenizer,
-        clip_tokenizer: Tokenizer,
-        t5_text_encoder: T5Encoder,
-        clip_text_encoder: CLIPEncoder,
-    ) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_attention.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5Attention(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_block.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5Block(nn.Module):
-    def __init__(self, layer: int) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_dense_relu_dense.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5DenseReluDense(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...
-    @staticmethod
-    def new_gelu(input_array: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_encoder.pyi

@@ -1,14 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-"""
-This type stub file was generated by pyright.
-"""
-
-class T5Encoder(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, tokens: mx.array): ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_feed_forward.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5FeedForward(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_layer_norm.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5LayerNorm(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_self_attention.pyi

@@ -1,16 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class T5SelfAttention(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(self, hidden_states: mx.array) -> mx.array: ...
-    @staticmethod
-    def shape(states):  # -> array:
-        ...
-    @staticmethod
-    def un_shape(states):  # -> array:
-        ...

.mlx_typings/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.pyi

@@ -1,10 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class AdaLayerNormContinuous(nn.Module):
-    def __init__(self, embedding_dim: int, conditioning_embedding_dim: int) -> None: ...
-    def __call__(self, x: mx.array, text_embeddings: mx.array) -> mx.array: ...

.mlx_typings/mflux/models/flux/model/flux_transformer/ada_layer_norm_zero.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-from mlx import nn
-
-class AdaLayerNormZero(nn.Module):
-    def __init__(self) -> None: ...
-    def __call__(
-        self, hidden_states: mx.array, text_embeddings: mx.array
-    ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: ...