Skip to content

fal-ai/flashpack

Repository files navigation

FlashPack Logo

Disk-to-GPU Tensor loading at up to 25Gbps without GDS

Benchmark Results Run this benchmark in `scripts/run_benchmark.py`
Benchmark Results Run this benchmark in `tests/test_speed_comparison.py`

Updates

  • 2025-11-25: Now supports multiple data types per checkpoint with no regressions in speed!

Integration Guide

Mixins

Diffusers/Transformers

# 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

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

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.

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:

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

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:

flashpack revert ./my_model.flashpack ./my_model.safetensors

flashpack metadata

Print the metadata of a flashpack file.

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:

# 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

About

High-throughput tensor loading for PyTorch

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Contributors 5

Languages