add size

app.py

@@ -8,6 +8,7 @@ from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from bitsandbytes.nn import Linear4bit
 from packaging import version
 import os
+from tqdm import tqdm
 
 
 def hello(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
@@ -108,6 +109,9 @@ def quantize_model(
     progress(0.33, desc="Quantizing")
 
     # Quantize model
+    # Calculate original model sizeo
+    original_size_gb = get_model_size(model)
+
     modules = list(model.named_modules())
     for idx, (_, module) in enumerate(modules):
         if isinstance(module, Linear4bit):
@@ -116,12 +120,13 @@ def quantize_model(
         progress(0.33 + (0.33 * idx / len(modules)), desc="Quantizing")
 
     progress(0.66, desc="Quantized successfully")
-    return model
+    return model, original_size_gb
 
 
 def save_model(
     model,
     model_name,
+    original_size_gb,
     quant_type_4,
     double_quant_4,
     compute_type_4,
@@ -189,7 +194,7 @@ def save_model(
         <div style="line-height: 1.2; font-size: 0.75em;">{model_architecture_str_html}</div>
     </div>
     """
-    return f'🔗 Quantized Model <br/><h1> 🤗 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a><br/><br/>📊 Model Architecture<br/>{model_architecture_info}'
+    return f'🎉 Quantized Model <br/><h1> 🚀 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a><br/><br/>📋 Model Architecture<br/>{model_architecture_info}<br/><br/>📦 Model Size <br/><br/> Original (bf16)≈ {original_size_gb} GB → Quantized ≈ {get_model_size(model)} GB'
 
 
 def quantize_and_save(
@@ -231,7 +236,7 @@ def quantize_and_save(
     try:
         # Download phase
         progress(0, desc="Starting quantization process")
-        quantized_model = quantize_model(
+        quantized_model, original_size_gb = quantize_model(
             model_name,
             quant_type_4,
             double_quant_4,
@@ -243,6 +248,7 @@ def quantize_and_save(
         final_message = save_model(
             quantized_model,
             model_name,
+            original_size_gb,
             quant_type_4,
             double_quant_4,
             compute_type_4,
@@ -264,7 +270,30 @@ def quantize_and_save(
             <p>{error_message}</p>
         </div>
         """
-
+def get_model_size(model):
+    """
+    Calculate the size of a PyTorch model in gigabytes.
+    
+    Args:
+        model: PyTorch model
+        
+    Returns:
+        float: Size of the model in GB
+    """
+    # Get model state dict
+    state_dict = model.state_dict()
+    
+    # Calculate total size in bytes
+    total_size = 0
+    for param in state_dict.values():
+        # Calculate bytes for each parameter
+        total_size += param.nelement() * param.element_size()
+    
+    # Convert bytes to gigabytes (1 GB = 1,073,741,824 bytes)
+    size_gb = total_size / (1024 ** 3)
+    size_gb = round(size_gb, 2)
+    
+    return size_gb
 
 css = """/* Custom CSS to allow scrolling */
 .gradio-container {overflow-y: auto;}
@@ -477,7 +506,8 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
     gr.Markdown(
         """
         # 🤗 BitsAndBytes Quantizer : Create your own BNB Quants ! ✨
-        
+        <br/>
+        <br/>
         """
     )