.png)
1 day ago
1
Run this benchmark in `scripts/run_benchmark.py`
Run this benchmark in `tests/test_speed_comparison.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")
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)
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
