Update app.py

app.py

@@ -1,10 +1,12 @@
-import gradio as gr
-from math import exp
+import io
 import re
 import struct
-import requests
-import io
 from enum import IntEnum
+from math import floor
+
+import requests
+
+import gradio as gr
 
 
 class GGUFValueType(IntEnum):
@@ -274,28 +276,27 @@ def estimate_vram(metadata, gpu_layers, ctx_size, cache_type):
             gpu_layers = n_layers
 
         # Convert cache_type to numeric
-        cache_type_map = {'fp16': 16, 'q8_0': 8, 'q4_0': 4}
-        cache_type_numeric = cache_type_map.get(cache_type, 16)
+        if cache_type == 'q4_0':
+            cache_type = 4
+        elif cache_type == 'q8_0':
+            cache_type = 8
+        else:
+            cache_type = 16
 
         # Derived features
         size_per_layer = size_in_mb / max(n_layers, 1e-6)
-        context_per_layer = context_length / max(n_layers, 1e-6)
-        ffn_per_embedding = feed_forward_dim / max(embedding_dim, 1e-6)
-        kv_cache_factor = n_kv_heads * cache_type_numeric * ctx_size
-
-        # Helper function for smaller
-        def smaller(x, y):
-            return 1 if x < y else 0
+        kv_cache_factor = n_kv_heads * cache_type * ctx_size
+        embedding_per_context = embedding_dim / ctx_size
 
         # Calculate VRAM using the model
+        # Details: https://oobabooga.github.io/blog/posts/gguf-vram-formula/
         vram = (
-            (size_per_layer - 21.19195204848197)
-            * exp(0.0001047328491557063 * size_in_mb * smaller(ffn_per_embedding, 2.671096993407845))
-            + 0.0006621544775632052 * context_per_layer
-            + 3.34664386576376e-05 * kv_cache_factor
-        ) * (1.363306170123392 + gpu_layers) + 1255.163594536052
+            (size_per_layer - 17.99552795246051 + 3.148552680382576e-05 * kv_cache_factor)
+            * (gpu_layers + max(0.9690636483914102, cache_type - (floor(50.77817218646521 * embedding_per_context) + 9.987899908205632)))
+            + 1516.522943869404
+        )
 
-        return max(0, vram)  # Ensure non-negative result
+        return vram
 
     except Exception as e:
         print(f"Error in VRAM calculation: {e}")
@@ -310,7 +311,7 @@ def estimate_vram_wrapper(model_metadata, gpu_layers, ctx_size, cache_type):
     # Use cache_type directly (it's already a string from the radio button)
     try:
         result = estimate_vram(model_metadata, gpu_layers, ctx_size, cache_type)
-        conservative = result + 906
+        conservative = result + 577
         return f"""<div id="vram-info">
         <div>Expected VRAM usage: <span class="value">{result:.0f} MiB</span></div>
         <div>Safe estimate: <span class="value">{conservative:.0f} MiB</span> - 95% chance the VRAM is at most this.</div>