Update download_model.py

download_model.py

@@ -1,4 +1,5 @@
 from transformers import VisionEncoderDecoderModel, TrOCRProcessor
+import torch
 import os
 
 MODEL_ID = 'anuashok/ocr-captcha-v3'
@@ -10,8 +11,15 @@ os.makedirs(SAVE_PATH, exist_ok=True)
 processor = TrOCRProcessor.from_pretrained(MODEL_ID)
 model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID)
 
-print("Saving to ./model...")
+print("Quantizing to INT8...")
+quantized_model = torch.quantization.quantize_dynamic(
+    model,
+    {torch.nn.Linear},
+    dtype=torch.qint8
+)
+
+print("Saving quantized model to ./model...")
 processor.save_pretrained(SAVE_PATH)
-model.save_pretrained(SAVE_PATH)
+quantized_model.save_pretrained(SAVE_PATH)
 
-print("✅ Model saved successfully")
+print("✅ Quantized model saved successfully")