Update src/pipeline.py

src/pipeline.py

@@ -26,6 +26,36 @@ torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.cuda.set_per_process_memory_fraction(0.95)
 
+
+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)
+        qtensor = torch.clamp(qtensor, self.qmin, self.qmax)
+        return (qtensor - zero_point) * scale, scale, zero_point
+
+class ModelQuantization:
+    def __init__(self, model, bits=9):
+        self.model = model
+        self.quant = BasicQuantization(bits)
+
+    def quantize_model(self):
+        for name, module in self.model.named_modules():
+            if isinstance(module, torch.nn.Linear):
+                if hasattr(module, 'weightML'):
+                    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)
+
+
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 def empty_cache():
     start = time.time()
@@ -65,6 +95,8 @@ def load_pipeline() -> Pipeline:
     # vae.to(dtype=DTYPE)
     
     # quantize_(vae, int8_weight_only())
+    quantizer = ModelQuantization(vae)
+    quantizer.quantize_model()