Skip to content

Commit 651666d

Browse files
committed
update docs
1 parent 2bd9302 commit 651666d

File tree

1 file changed

+33
-0
lines changed

1 file changed

+33
-0
lines changed

docs/source/en/quantization/torchao.md

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,39 @@ Some quantization methods are aliases (for example, `int8wo` is the commonly use
8888

8989
Refer to the official torchao documentation for a better understanding of the available quantization methods and the exhaustive list of configuration options available.
9090

91+
## Serializing and Deserializing quantized models
92+
93+
To serialize a quantized model in a given dtype, first load the model with the desired quantization dtype and then save it using the [`~ModelMixin.save_pretrained`] method.
94+
95+
```python
96+
import torch
97+
from diffusers import FluxTransformer2DModel, TorchAoConfig
98+
99+
quantization_config = TorchAoConfig("int8wo")
100+
transformer = FluxTransformer2DModel.from_pretrained(
101+
"black-forest-labs/Flux.1-Dev",
102+
subfolder="transformer",
103+
quantization_config=quantization_config,
104+
torch_dtype=torch.bfloat16,
105+
)
106+
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)
107+
```
108+
109+
To load a serialized quantized model, use the [`~ModelMixin.from_pretrained`] method.
110+
111+
```python
112+
import torch
113+
from diffusers import FluxPipeline, FluxTransformer2DModel
114+
115+
transformer = FluxTransformer2DModel.from_pretrained("/path/to/flux_int8wo", torch_dtype=torch.bfloat16, use_safetensors=False)
116+
pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)
117+
pipe.to("cuda")
118+
119+
prompt = "A cat holding a sign that says hello world"
120+
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
121+
image.save("output.png")
122+
```
123+
91124
## Resources
92125

93126
- [TorchAO Quantization API](https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md)

0 commit comments

Comments
 (0)