Update README.md

README.md

@@ -68,14 +68,16 @@ pip install optimum-quanto
 2. Run the script:
 ```python
 import torch
-from safetensors.torch import save_file
+from safetensors.torch import save_file, load_file
 
-from diffusers import FluxTransformer2DModel, FluxPipeline
-from optimum.quanto import freeze, qfloat8, quantize
+from diffusers import FluxTransformer2DModel
+from optimum.quanto import freeze, qfloat8, quantize, QuantizedDiffusersModel
+
+class QuantizedFlux2DModel(QuantizedDiffusersModel):
+    base_class = FluxTransformer2DModel
 
 transformer = FluxTransformer2DModel.from_single_file("modern-anime.safetensors", torch_dtype=torch.bfloat16)
-quantize(transformer, weights=qfloat8)
-freeze(transformer)
+transformer = QuantizedFlux2DModel.quantize(transformer, weights=qfloat8)
 
 transformer.save_pretrained("transformer")
 ```