Feat: support multiple dtypes per checkpoint

.github/workflows/pre-commit.yaml

@@ -0,0 +1,18 @@
+name: Lint (pre-commit)
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [assigned, opened, synchronize, reopened]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v3
+      with:
+        python-version: "3.11"
+    - uses: pre-commit/action@v3.0.0

.pre-commit-config.yaml

@@ -0,0 +1,37 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.0.1
+    hooks:
+      - id: check-docstring-first
+      - id: check-toml
+      - id: check-yaml
+        exclude: packaging/.*
+        args:
+          - --allow-multiple-documents
+      - id: mixed-line-ending
+        args: [--fix=lf]
+      - id: end-of-file-fixer
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: 'v0.3.4'
+    hooks:
+      - id: ruff
+        name: lint with ruff
+      - id: ruff
+        name: sort imports with ruff
+        args: [--select, I, --fix]
+      - id: ruff-format
+        name: format with ruff
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-merge-conflict
+      - id: check-vcs-permalinks
+      - id: debug-statements
+
+  - repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: v14.0.6
+    hooks:
+      - id: clang-format 

README.md

@@ -1,91 +1,187 @@
-<div align="center">
-<picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-white.png?raw=true">
-  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-black.png?raw=true">
-  <img alt="FlashPack Logo" src="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-black.png?raw=true">
-</picture>
-<h2>Disk-to-GPU Tensor loading at up to 25Gbps without GDS</h2>
-</div>
-
-<div align="center">
-<picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-white.png?raw=true">
-  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-black.png?raw=true">
-  <img alt="Benchmark Results" src="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-black.png?raw=true">
-</picture>
-<em>Run this benchmark in `scripts/run_benchmark.py`</em>
-</div>
-
-<div align="center">
-<picture>
-  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-white.png?raw=true">
-  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-black.png?raw=true">
-  <img alt="Benchmark Results" src="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-black.png?raw=true">
-</picture>
-<em>Run this benchmark in `tests/test_speed_comparison.py`</em>
-</div>
-
-# Integration Guide
-## Mixins
-### Diffusers/Transformers
-
-```py
-# Integration classes
-from flashpack.integrations.diffusers import FlashPackDiffusersModelMixin, FlashPackDiffusionPipeline
-from flashpack.integrations.transformers import FlashPackTransformersModelMixin
-
-# Base classes
-from diffusers.models import MyModel, SomeOtherModel
-from diffusers.pipelines import MyPipeline
-
-# Define mixed classes
-class FlashPackMyModel(MyModel, FlashPackDiffusersModelMixin):
-    pass
-
-class FlashPackMyPipeline(MyPipeline, FlashPackDiffusionPipine):
-    def __init__(
-        self,
-        my_model: FlashPackMyModel,
-        other_model: SomeOtherModel,
-    ) -> None:
-        super().__init__()
-
-# Load base pipeline
-pipeline = FlashPackMyPipeline.from_pretrained("some/repository")
-
-# Save flashpack pipeline
-pipeline.save_pretrained_flashpack(
-    "some_directory",
-    push_to_hub=False,  # pass repo_id when using this
-)
-
-# Load directly from flashpack directory or repository
-pipeline = FlashPackMyPipeline.from_pretrained_flashpack("my/flashpack-repository")
-```
-
-### Vanilla PyTorch
-
-```py
-from flashpack import FlashPackMixin
-
-class MyModule(nn.Module, FlashPackMixin):
-    def __init__(self, some_arg: int = 4) -> None:
-        ...
-
-module = MyModule(some_arg = 4)
-module.save_flashpack("model.flashpack")
-
-loaded_module = module.from_flashpack("model.flashpack", some_arg=4)
-```
-
-## Direct Integration
-
-```py
-from flashpack import pack_to_file, assign_from_file
-
-flashpack_path = "/path/to/model.flashpack"
-model = nn.Module(...)
-
-pack_to_file(model, flashpack_path)  # write state dict to file
-assign_from_file(model, flashpack_path)  # load state dict from file
-```
+<div align="center">
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-white.png?raw=true">
+  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-black.png?raw=true">
+  <img alt="FlashPack Logo" src="https://github.com/fal-ai/flashpack/blob/main/media/flashpack-logo-black.png?raw=true">
+</picture>
+<h2>Disk-to-GPU Tensor loading at up to 25Gbps without GDS</h2>
+</div>
+
+## Updates
+
+- **2025-11-25**: Now supports **multiple data types per checkpoint** with no regressions in speed!
+
+<div align="center">
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-white.png?raw=true">
+  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-black.png?raw=true">
+  <img alt="Benchmark Results" src="https://github.com/fal-ai/flashpack/blob/main/media/benchmark-black.png?raw=true">
+</picture>
+<em>Run this benchmark in `scripts/run_benchmark.py`</em>
+</div>
+
+<div align="center">
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-white.png?raw=true">
+  <source media="(prefers-color-scheme: light)" srcset="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-black.png?raw=true">
+  <img alt="Benchmark Results" src="https://github.com/fal-ai/flashpack/blob/main/media/load-state-dict-comparison-black.png?raw=true">
+</picture>
+<em>Run this benchmark in `tests/test_speed_comparison.py`</em>
+</div>
+
+# Integration Guide
+## Mixins
+### Diffusers/Transformers
+
+```py
+# Integration classes
+from flashpack.integrations.diffusers import FlashPackDiffusersModelMixin, FlashPackDiffusionPipeline
+from flashpack.integrations.transformers import FlashPackTransformersModelMixin
+
+# Base classes
+from diffusers.models import MyModel, SomeOtherModel
+from diffusers.pipelines import MyPipeline
+
+# Define mixed classes
+class FlashPackMyModel(MyModel, FlashPackDiffusersModelMixin):
+    pass
+
+class FlashPackMyPipeline(MyPipeline, FlashPackDiffusionPipine):
+    def __init__(
+        self,
+        my_model: FlashPackMyModel,
+        other_model: SomeOtherModel,
+    ) -> None:
+        super().__init__()
+
+# Load base pipeline
+pipeline = FlashPackMyPipeline.from_pretrained("some/repository")
+
+# Save flashpack pipeline
+pipeline.save_pretrained_flashpack(
+    "some_directory",
+    push_to_hub=False,  # pass repo_id when using this
+)
+
+# Load directly from flashpack directory or repository
+pipeline = FlashPackMyPipeline.from_pretrained_flashpack("my/flashpack-repository")
+```
+
+### Vanilla PyTorch
+
+```py
+from flashpack import FlashPackMixin
+
+class MyModule(nn.Module, FlashPackMixin):
+    def __init__(self, some_arg: int = 4) -> None:
+        ...
+
+module = MyModule(some_arg = 4)
+module.save_flashpack("model.flashpack")
+
+loaded_module = module.from_flashpack("model.flashpack", some_arg=4)
+```
+
+## Direct Integration
+
+```py
+from flashpack import pack_to_file, assign_from_file
+
+flashpack_path = "/path/to/model.flashpack"
+model = nn.Module(...)
+
+pack_to_file(model, flashpack_path)  # write state dict to file
+assign_from_file(model, flashpack_path)  # load state dict from file
+```
+
+# CLI Commands
+
+FlashPack provides a command-line interface for converting, inspecting, and reverting flashpack files.
+
+## `flashpack convert`
+
+Convert a model to a flashpack file.
+
+```bash
+flashpack convert <path_or_repo_id> [destination_path] [options]
+```
+
+**Arguments:**
+- `path_or_repo_id` - Local path or Hugging Face repository ID
+- `destination_path` - (Optional) Output path for the flashpack file
+
+**Options:**
+| Option | Description |
+|--------|-------------|
+| `--subfolder` | Subfolder of the model (for repo_id) |
+| `--variant` | Model variant (for repo_id) |
+| `--dtype` | Target dtype for the flashpack file. When omitted, no type changes are made |
+| `--ignore-names` | Tensor names to ignore (can be specified multiple times) |
+| `--ignore-prefixes` | Tensor prefixes to ignore (can be specified multiple times) |
+| `--ignore-suffixes` | Tensor suffixes to ignore (can be specified multiple times) |
+| `--use-transformers` | Load the path as a transformers model |
+| `--use-diffusers` | Load the path as a diffusers model |
+| `-v, --verbose` | Enable verbose output |
+
+**Examples:**
+```bash
+# Convert a local model
+flashpack convert ./my_model ./my_model.flashpack
+
+# Convert from Hugging Face
+flashpack convert stabilityai/stable-diffusion-xl-base-1.0 --subfolder unet --use-diffusers
+
+# Convert with specific dtype
+flashpack convert ./my_model ./my_model.flashpack --dtype float16
+```
+
+## `flashpack revert`
+
+Revert a flashpack file back to safetensors or torch format.
+
+```bash
+flashpack revert <path> [destination_path] [options]
+```
+
+**Arguments:**
+- `path` - Path to the flashpack file
+- `destination_path` - (Optional) Output path for the reverted file
+
+**Options:**
+| Option | Description |
+|--------|-------------|
+| `-v, --verbose` | Enable verbose output |
+
+**Example:**
+```bash
+flashpack revert ./my_model.flashpack ./my_model.safetensors
+```
+
+## `flashpack metadata`
+
+Print the metadata of a flashpack file.
+
+```bash
+flashpack metadata <path> [options]
+```
+
+**Arguments:**
+- `path` - Path to the flashpack file
+
+**Options:**
+| Option | Description |
+|--------|-------------|
+| `-i, --show-index` | Show the tensor index |
+| `-j, --json` | Output metadata in JSON format |
+
+**Examples:**
+```bash
+# View basic metadata
+flashpack metadata ./my_model.flashpack
+
+# View metadata with tensor index
+flashpack metadata ./my_model.flashpack --show-index
+
+# Output as JSON
+flashpack metadata ./my_model.flashpack --json
+```

pyproject.toml

@@ -46,4 +46,4 @@ dev = [
 ]
 
 [tool.setuptools_scm]
-write_to = "src/flashpack/version.py"
+write_to = "src/flashpack/version.py"

scripts/plot_benchmark.py

@@ -10,7 +10,7 @@
 # configuration
 accelerate_color = "#0f5ef3"
 flashpack_color = "#adff02"
-label_color = "#111111"
+label_color = "#eeeeee"
 model_labels = {
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": "Wan2.1 1.3B DiT",
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": "Wan2.1 14B DiT",

scripts/run_benchmark.py

@@ -8,13 +8,13 @@
 import torch
 from flashpack.integrations.diffusers import patch_diffusers_auto_model
 from flashpack.integrations.transformers import patch_transformers_auto_model
+from huggingface_hub import snapshot_download
 
 patch_diffusers_auto_model()
 patch_transformers_auto_model()
 
-from diffusers.models import AutoModel as DiffusersAutoModel
-from huggingface_hub import snapshot_download
-from transformers import AutoModel as TransformersAutoModel
+from diffusers.models import AutoModel as DiffusersAutoModel  # noqa: E402
+from transformers import AutoModel as TransformersAutoModel  # noqa: E402
 
 
 def test_model(
@@ -201,4 +201,4 @@ def sync_and_flush() -> None:
                 )
                 print_test_result(
                     test_model_name, accelerate_time, flashpack_time, total_bytes
-                )
+                )