Update src/pipeline.py

src/pipeline.py

@@ -1,4 +1,4 @@
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+# from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 
@@ -12,11 +12,58 @@ from pipelines.models import TextToImageRequest
 from torch import Generator
 import time
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
-from torchao.quantization import quantize_,  float8_dynamic_activation_float8_weight #PerRow,
+# from torchao.quantization import quantize_,  float8_dynamic_activation_float8_weight #PerRow,
 import os
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01"
 Pipeline = None
 
+class W8A16LinearLayer(nn.Module):
+    def __init__(self, in_features, out_features, bias=True, dtype=torch.float32):
+        super().__init__()
+        
+        self.register_buffer(
+            "int8_weights",
+            torch.randint(-128, 127, (out_features, in_features), dtype=torch.int8))
+        
+        self.register_buffer("scales", torch.randn((out_features), dtype=dtype))
+        
+        if bias:
+            self.register_buffer("bias", torch.randn((1, out_features), dtype=dtype))
+        else:
+    
+    def quantize(self, weights):
+        w_fp32 = weights.clone().to(torch.float32)
+
+        scales = w_fp32.abs().max(dim=-1).values / 127
+        scales = scales.to(weights.dtype)
+
+        int8_weights = torch.round(weights/scales.unsqueeze(1)).to(torch.int8)
+
+        self.int8_weights = int8_weights
+        self.scales = scales            self.bias = None
+
+    def forward(self, input):
+        return w8_a16_forward(self.int8_weights, input, self.scales, self.bias)
+
+def replace_linear_with_target_and_quantize(module, target_class, module_name_to_exclude):
+    for name, child in module.named_children():
+        if isinstance(child, nn.Linear) and not \
+        any([x == name for x in module_name_to_exclude]):
+            old_bias = child.bias
+            old_weight = child.weight
+
+            new_module = target_class(child.in_features, child.out_features, old_bias is not None, child.weight.dtype)
+            setattr(module, name, new_module)
+
+            getattr(module, name).quantize(old_weight)
+            
+            if old_bias is not None:
+              getattr(module, name).bias = old_bias
+        else:
+            # Recursively call the function for nested modules
+            replace_linear_with_target_and_quantize(child, target_class, module_name_to_exclude)
+
+
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 def empty_cache():
     start = time.time()
@@ -39,12 +86,14 @@ def load_pipeline() -> Pipeline:
         text_encoder_2 = text_encoder_2,
         torch_dtype=dtype,
         )
-    quantize_(pipeline.transformer, float8_dynamic_activation_float8_weight())
+    # quantize_(pipeline.transformer, float8_dynamic_activation_float8_weight())
+    
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.cuda.set_per_process_memory_fraction(0.99)
     pipeline.text_encoder.to(memory_format=torch.channels_last)
     pipeline.transformer.to(memory_format=torch.channels_last)
+    replace_linear_with_target_and_quantize(pipeline.transformer, W8A16LinearLayer, [])
     
 
     pipeline.vae.to(memory_format=torch.channels_last)