diff --git a/README.md b/README.md index c175241..7e271b7 100644 --- a/README.md +++ b/README.md @@ -97,10 +97,13 @@ import sys sys.path.append("/path/to/mflux/src") from mflux.config.config import Config -from mflux.flux import Flux1 +from mflux.flux.flux import Flux1 # Load the model -flux = Flux1.from_alias(alias="schnell") # "schnell" or "dev" +flux = Flux1.from_alias( + alias="schnell", # "schnell" or "dev" + quantize=8, # 4 or 8 +) # Generate an image image = flux.generate_image( diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py index 35d9356..96964c9 100644 --- a/src/mflux/flux/flux.py +++ b/src/mflux/flux/flux.py @@ -126,8 +126,11 @@ class Flux1: return latents @staticmethod - def from_alias(alias: str) -> "Flux1": - return Flux1(ModelConfig.from_alias(alias)) + def from_alias(alias: str, quantize: int | None = None) -> "Flux1": + return Flux1( + model_config=ModelConfig.from_alias(alias), + quantize=quantize, + ) def _set_model_weights(self, weights): self.vae.update(weights.vae)