8 bit onxx (#5)

- 8 bit onxx (a35a0a82d6202fd081cef55fdf17fadf23a43395)


Co-authored-by: Aditi Baheti <[email protected]>

app.py

@@ -5,6 +5,10 @@ from diffusers import DiffusionPipeline
 import torch
 from huggingface_hub import login
 import os
+import bitsandbytes as bnb
+import onnx
+import onnxruntime as ort
+from onnxruntime.quantization import quantize_dynamic, QuantType
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -16,17 +20,32 @@ login(token=HUGGINGFACE_TOKEN)
 base_model_repo = "stabilityai/stable-diffusion-3-medium-diffusers"
 lora_weights_path = "./pytorch_lora_weights.safetensors"
 
-# Load the base model
+# Load the base model with 8-bit precision
 pipeline = DiffusionPipeline.from_pretrained(
     base_model_repo,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     use_auth_token=HUGGINGFACE_TOKEN
 )
+bnb.optim.load_int8_model(pipeline.model, device=device)
+
 pipeline.load_lora_weights(lora_weights_path)
 pipeline.enable_sequential_cpu_offload()  # Efficient memory usage
 pipeline.enable_xformers_memory_efficient_attention()  # Enable xformers memory efficient attention
 pipeline = pipeline.to(device)
 
+# Export to ONNX
+onnx_model_path = "model.onnx"
+pipeline.model.eval()
+dummy_input = torch.randn(1, 3, 512, 512, device=device)
+torch.onnx.export(pipeline.model, dummy_input, onnx_model_path, export_params=True, opset_version=11, do_constant_folding=True, input_names=['input'], output_names=['output'])
+
+# Quantize ONNX model to 8-bit
+quantized_model_path = "model_quantized.onnx"
+quantize_dynamic(onnx_model_path, quantized_model_path, weight_type=QuantType.QUInt8)
+
+# Load quantized ONNX model
+session = ort.InferenceSession(quantized_model_path)
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 768  # Reduce max image size to fit within memory constraints