Update README.md

README.md

@@ -19,11 +19,11 @@ tags:
 
 # Usage with Diffusers
 
-To use this quantized FLUX.1 [dev] checkpoint, you need to install the 🧨 diffusers and bitsandbytes library:
+To use this quantized FLUX.1 [dev] checkpoint, you need to install the 🧨 diffusers and torchao library:
 
 ```
 pip install -U diffusers
-pip install -U bitsandbytes
+pip install -U torchao
 ```
 
 After installing the required library, you can run the following script: 
@@ -65,13 +65,13 @@ This checkpoint was created with the following script using "black-forest-labs/F
 import torch
 from diffusers import FluxPipeline
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
-from diffusers.quantizers import PipelineQuantizationConfig
-from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
+from diffusers import TorchAoConfig as DiffusersTorchAoConfig
+from transformers import TorchAoConfig as TransformersTorchAoConfig
 
 pipeline_quant_config = PipelineQuantizationConfig(
     quant_mapping={
-        "transformer": DiffusersBitsAndBytesConfig(load_in_8bit=True),
-        "text_encoder_2": TransformersBitsAndBytesConfig(load_in_8bit=True),
+        "transformer": DiffusersTorchAoConfig("int8_weight_only"),
+        "text_encoder_2": TransformersTorchAoConfig("int8_weight_only"),
     }
 )