Shows how to run Flux schnell under 17GBs without bells and whistles. It additionally shows how to serialize the quantized checkpoint and load it back.

inference_with_torchao_serialized.py

import torch
from huggingface_hub import hf_hub_download
from diffusers import FluxTransformer2DModel, DiffusionPipeline

dtype, device = torch.bfloat16, "cuda"
ckpt_id = "black-forest-labs/FLUX.1-schnell"

with torch.device("meta"):
    config = FluxTransformer2DModel.load_config(ckpt_id, subfolder="transformer")
    model = FluxTransformer2DModel.from_config(config).to(dtype)

ckpt_path = hf_hub_download(repo_id="sayakpaul/flux.1-schell-int8wo", filename="flux_schnell_int8wo.pt")
state_dict = torch.load(ckpt_path, map_location="cpu")
model.load_state_dict(state_dict, assign=True)

pipeline = DiffusionPipeline.from_pretrained(ckpt_id, transformer=model, torch_dtype=dtype).to("cuda")
image = pipeline(
	"cat", guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256
).images[0]
image.save("flux_schnell_int8.png")

run_flux_schnell_torchao.py

# Install `torchao` from source: https://github.com/pytorch/ao
# Install PyTorch nightly

from diffusers import DiffusionPipeline, FluxTransformer2DModel, AutoencoderKL
from transformers import T5EncoderModel, CLIPTextModel
from torchao.quantization import quantize_, int8_weight_only
import torch 

ckpt_id = "black-forest-labs/FLUX.1-schnell"

# Quantize the components individually.
# If quality is taking a hit then don't quantize all components.
# Mix and match.

############ Diffusion Transformer ############
transformer = FluxTransformer2DModel.from_pretrained(
    ckpt_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
quantize_(transformer, int8_weight_only())

############ Text Encoder ############
text_encoder = CLIPTextModel.from_pretrained(
    ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
)
quantize_(text_encoder, int8_weight_only())

############ Text Encoder 2 ############
text_encoder_2 = T5EncoderModel.from_pretrained(
    ckpt_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16
)
quantize_(text_encoder_2, int8_weight_only())

############ VAE ############
vae = AutoencoderKL.from_pretrained(
    ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16
)
quantize_(vae, int8_weight_only())

# Initialize the pipeline now.
pipeline = DiffusionPipeline.from_pretrained(
	ckpt_id, 
    transformer=transformer, 
    vae=vae,
    text_encoder=text_encoder, 
    text_encoder_2=text_encoder_2, 
    torch_dtype=torch.bfloat16
).to("cuda")
image = pipeline(
	"cat", guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256
).images[0]

torch.cuda.empty_cache()
memory = (torch.cuda.memory_allocated() / 1024 / 1024 / 1024) 
print(f"{memory=:.3f} GB")

image.save("quantized_image.png")

serialize_with_torchao.py

from diffusers import FluxTransformer2DModel
from torchao.quantization import quantize_, int8_weight_only
import torch 

ckpt_id = "black-forest-labs/FLUX.1-schnell"

transformer = FluxTransformer2DModel.from_pretrained(
    ckpt_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
quantize_(transformer, int8_weight_only())

# should ideally be possible with safetensors but
# https://github.com/huggingface/safetensors/issues/515
# this checkpoint is 12GB instead of 23GB.
torch.save(transformer.state_dict(), "flux_schnell_int8wo.pt")

/home/abdallah/Desktop/webui/qunata-6h/runquanta.py:15: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  transformer_state_dict = torch.load(os.path.join(save_directory, "transformer.pt"), map_location="cpu")
/home/abdallah/Desktop/webui/qunata-6h/runquanta.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  vae_state_dict = torch.load(os.path.join(save_directory, "vae.pt"), map_location="cpu")
Traceback (most recent call last):
  File "/home/abdallah/Desktop/webui/qunata-6h/runquanta.py", line 27, in <module>
    text_encoder = CLIPTextModel(text_encoder_config)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abdallah/Desktop/webui/qunata-6h/.venv/lib/python3.12/site-packages/transformers/models/clip/modeling_clip.py", line 946, in __init__
    self.text_model = CLIPTextTransformer(config)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abdallah/Desktop/webui/qunata-6h/.venv/lib/python3.12/site-packages/transformers/models/clip/modeling_clip.py", line 840, in __init__
    self.embeddings = CLIPTextEmbeddings(config)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abdallah/Desktop/webui/qunata-6h/.venv/lib/python3.12/site-packages/transformers/models/clip/modeling_clip.py", line 205, in __init__
    self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/abdallah/Desktop/webui/qunata-6h/.venv/lib/python3.12/site-packages/torch/nn/modules/sparse.py", line 167, in __init__
    torch.empty((num_embeddings, embedding_dim), **factory_kwargs),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: empty() received an invalid combination of arguments - got (tuple, dtype=NoneType, device=NoneType), but expected one of:
 * (tuple of ints size, *, tuple of names names, torch.memory_format memory_format = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (tuple of ints size, *, torch.memory_format memory_format = None, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)

@tin2tin we don't support Windows really well simply because torch.compile() doesnt support Windows very well. Feel free to open an issue on our repo though we might choose to prioritze this depending on demand.

For now try installing ao from source using USE_CPP=0 pip install .

This works on windows, in case anyone wants to know
The commands on cmd prompt are actually:

set USE_CPP=0
pip install <path to downloaded torchao folder>

Edit 12-01-2025:
as of now, even this is throwing an error:

AttributeError: Can't get attribute 'PlainAQTLayout' on <module 'torchao.dtypes.affine_quantized_tensor' from '...Python\\Python311\\site-packages\\torchao\\dtypes\\affine_quantized_tensor.py'>

For those trying to use the one under 17GBs the link is dead, you can you use this instead:

import torch
from diffusers import FluxTransformer2DModel, DiffusionPipeline

dtype, device = torch.bfloat16, "cuda"
ckpt_id = "black-forest-labs/FLUX.1-schnell"

model = FluxTransformer2DModel.from_pretrained(
    "sayakpaul/flux.1-schell-int8wo-improved", torch_dtype=dtype, use_safetensors=False
)
pipeline = DiffusionPipeline.from_pretrained(ckpt_id, transformer=model, torch_dtype=dtype).to("cuda")
image = pipeline(
    "cat", guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256
).images[0]
image.save("flux_schnell_int8.png")

@MostHumble
Getting an error running that code:
AttributeError: Can't get attribute 'PlainAQTLayout' on <module 'torchao.dtypes.affine_quantized_tensor' from '...Python\\Python311\\site-packages\\torchao\\dtypes\\affine_quantized_tensor.py'>
On torchao == 0.7.0

@tin2tin
Ended up having the same issue, reported it here: https://huggingface.co/sayakpaul/flux.1-schell-int8wo-improved/discussions/2#6783e9e49c52f42c530ff7c1

Please take the issue with torchao. Until it's resolved, either:

1. Use torchao integration from Diffusers.
2. Downgrade torchao installation.

Other than that, I can't provide additional suggestions.