Update src/pipeline.py

src/pipeline.py

@@ -7,7 +7,33 @@ from PIL.Image import Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
 from diffusers import DiffusionPipeline
-from torchao.quantization import quantize_, int8_weight_only
+
+class BasicQuantization:
+    def __init__(self, bits=1):
+        self.bits = bits
+        self.qmin = -(2**(bits-1))
+        self.qmax = 2**(bits-1) - 1
+    def quantize_tensor(self, tensor):
+        scale = (tensor.max() - tensor.min()) / (self.qmax - self.qmin)
+        zero_point = self.qmin - torch.round(tensor.min() / scale)
+        qtensor = torch.round(tensor / scale + zero_point) #q
+        qtensor = torch.clamp(qtensor, self.qmin, self.qmax)
+        tensor_q = (qtensor - zero_point) * scale #d
+        return tensor_q, scale, zero_point
+class SDXLQuantization:
+    def __init__(self, model, bit_number=16):
+        self.model = model
+        self.quant = BasicQuantization(bit_number) 
+    def quantize_model(self, save_name=None):
+        quantized_layers_state = {}
+        for name, module in self.model.named_modules():
+            if isinstance(module, (torch.nn.Linear)): # nn.Conv2d
+                if hasattr(module, 'weight'):
+                    quantized_weight, _, _ = self.quant.quantize_tensor(module.weight)
+                    module.weight = torch.nn.Parameter(quantized_weight)
+                if hasattr(module, 'bias') and module.bias is not None:
+                    quantized_bias, _, _ = self.quant.quantize_tensor(module.bias)
+                    module.bias = torch.nn.Parameter(quantized_bias)
 
 Pipeline = None
 MODEL_ID = "black-forest-labs/FLUX.1-schnell"
@@ -23,7 +49,9 @@ def load_pipeline() -> Pipeline:
     vae = AutoencoderKL.from_pretrained(
         MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16
     )
-    quantize_(vae, int8_weight_only(), device="cuda")
+    instance = SDXLQuantization(vae, 8)
+    instance.quantize_model()
+    
     pipeline = DiffusionPipeline.from_pretrained(
         MODEL_ID,
         vae=vae,