update

Author: a-r-r-o-w

src/diffusers/quantizers/torchao/torchao_quantizer.py

@@ -125,7 +125,7 @@ def validate_environment(self, *args, **kwargs):
     def update_torch_dtype(self, torch_dtype):
         quant_type = self.quantization_config.quant_type
 
-        if quant_type.startswith("int"):
+        if quant_type.startswith("int") or quant_type.startswith("uint"):
             if torch_dtype is not None and torch_dtype != torch.bfloat16:
                 logger.warning(
                     f"You are trying to set torch_dtype to {torch_dtype} for int4/int8/uintx quantization, but "

tests/quantization/torchao/test_torchao.py

@@ -654,3 +654,25 @@ def test_quantization(self):
             gc.collect()
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
+    
+    def test_serialization(self):
+        quantization_config = TorchAoConfig("int8wo")
+        components = self.get_dummy_components(quantization_config)
+        pipe = FluxPipeline(**components)
+        pipe.enable_model_cpu_offload()
+
+        weight = pipe.transformer.x_embedder.weight
+        self.assertTrue(isinstance(weight, AffineQuantizedTensor))
+
+        inputs = self.get_dummy_inputs(torch_device)
+        output = pipe(**inputs)[0].flatten()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            pipe.save_pretrained(tmp_dir, safe_serialization=False)
+            loaded_pipe = FluxPipeline.from_pretrained(tmp_dir, use_safetensors=False).to(torch_device)
+
+        weight = loaded_pipe.transformer.x_embedder.weight
+        self.assertTrue(isinstance(weight, AffineQuantizedTensor))
+        
+        loaded_output = loaded_pipe(**inputs)[0].flatten()
+        self.assertTrue(np.allclose(output, loaded_output, atol=1e-3, rtol=1e-3))