Update app.py

app.py

@@ -13,22 +13,25 @@ from torchao.quantization import (
     Int8DynamicActivationInt8WeightConfig,
     Float8WeightOnlyConfig,
     Float8DynamicActivationFloat8WeightConfig,
+    GemliteUIntXWeightOnlyConfig,
 )
 
 MAP_QUANT_TYPE_TO_NAME = {
-    "int4_weight_only": "int4wo",
-    "int8_weight_only": "int8wo",
-    "int8_dynamic_activation_int8_weight": "int8da8w8",
-    "float8_weight_only": "float8wo",
-    "float8_dynamic_activation_float8_weight": "float8da8w8",
+    "Int4WeightOnly": "int4wo",
+    "GemliteUIntXWeightOnly": "intxwo-gemlite"
+    "Int8WeightOnly": "int8wo",
+    "Int8DynamicActivationInt8Weight": "int8da8w8",
+    "Float8WeightOnly": "float8wo",
+    "Float8DynamicActivationFloat8Weight": "float8da8w8",
     "autoquant": "autoquant",
 }
 MAP_QUANT_TYPE_TO_CONFIG = {
-    "int4_weight_only": Int4WeightOnlyConfig,
-    "int8_weight_only": Int8WeightOnlyConfig,
-    "int8_dynamic_activation_int8_weight": Int8DynamicActivationInt8WeightConfig,
-    "float8_weight_only": Float8WeightOnlyConfig,
-    "float8_dynamic_activation_float8_weight": Float8DynamicActivationFloat8WeightConfig,
+    "Int4WeightOnly": Int4WeightOnlyConfig,
+    "GemliteUIntXWeightOnly": GemliteUIntXWeightOnlyConfig,
+    "Int8WeightOnly": Int8WeightOnlyConfig,
+    "Int8DynamicActivationInt8Weight": Int8DynamicActivationInt8WeightConfig,
+    "Float8WeightOnly": Float8WeightOnlyConfig,
+    "Float8DynamicActivationFloat8Weight": Float8DynamicActivationFloat8WeightConfig,
 }
 
 
@@ -56,8 +59,7 @@ def check_model_exists(
             repo_name = f"{username}/{quantized_model_name}"
         else:
             if (
-                quantization_type == "int4_weight_only"
-                or quantization_type == "int8_weight_only"
+                quantization_type in ["Int4WeightOnly", "GemliteUIntXWeightOnly"]
             ) and (group_size is not None):
                 repo_name = f"{username}/{model_name.split('/')[-1]}-ao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}-gs{group_size}"
             else:
@@ -173,13 +175,13 @@ def quantize_model(
     print(f"Quantizing model: {quantization_type}")
     progress(0, desc="Preparing Quantization")
     if (
-        quantization_type == "int8_weight_only"
+        quantization_type == "GemliteUIntXWeightOnly"
     ):
         quant_config = MAP_QUANT_TYPE_TO_CONFIG[quantization_type](
             group_size=group_size
         )
         quantization_config = TorchAoConfig(quant_config)
-    elif quantization_type == "int4_weight_only":
+    elif quantization_type == "Int4WeightOnly":
         from torchao.dtypes import Int4CPULayout
         
         quant_config = MAP_QUANT_TYPE_TO_CONFIG[quantization_type](
@@ -233,8 +235,7 @@ def save_model(
             repo_name = f"{username}/{quantized_model_name}"
         else:
             if (
-                quantization_type == "int4_weight_only"
-                or quantization_type == "int8_weight_only"
+                quantization_type in ["Int4WeightOnly", "GemliteUIntXWeightOnly"]
             ) and (group_size is not None):
                 repo_name = f"{username}/{model_name.split('/')[-1]}-ao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}-gs{group_size}"
             else:
@@ -318,7 +319,7 @@ def quantize_and_save(
             return """
             <div class="error-box">
                 <h3>❌ Group Size Error</h3>
-                <p>Group Size is a number for int4_weight_only and int8_weight_only or empty for int8_weight_only</p>
+                <p>Group Size is a parameter for Int4WeightOnly or GemliteUIntXWeightOnly</p>
             </div>
             """
 
@@ -492,11 +493,12 @@ with gr.Blocks(css=css) as demo:
                     quantization_type = gr.Dropdown(
                         info="Select the Quantization method",
                         choices=[
-                            "int4_weight_only",
-                            "int8_weight_only",
-                            "int8_dynamic_activation_int8_weight",
-                            "float8_weight_only",
-                            "float8_dynamic_activation_float8_weight",
+                            "Int4WeightOnly",
+                            "GemliteUIntXWeightOnly"
+                            "Int8WeightOnly",
+                            "Int8DynamicActivationInt8Weight",
+                            "Float8WeightOnly",
+                            "Float8DynamicActivationFloat8Weight",
                             "autoquant",
                         ],
                         value="int8_weight_only",
@@ -549,11 +551,18 @@ with gr.Blocks(css=css) as demo:
             ## 📝 Quantization Options
             
             ### Quantization Types
-            - **int4_weight_only**: 4-bit weight-only quantization 
-            - **int8_weight_only**: 8-bit weight-only quantization
-            - **int8_dynamic_activation_int8_weight**: 8-bit quantization for both weights and activations
-            - **float8_weight_only**: float8-bit weight-only quantization
-            - **float8_dynamic_activation_float8_weight**: float8-bit quantization for both weights and activations
+                            "Int4WeightOnly",
+                            "GemliteUIntXWeightOnly"
+                            "Int8WeightOnly",
+                            "Int8DynamicActivationInt8Weight",
+                            "Float8WeightOnly",
+                            "Float8DynamicActivationFloat8Weight",
+            - **Int4WeightOnly**: 4-bit weight-only quantization
+            - **GemliteUIntXWeightOnly**: uintx gemlite quantization (default to 4 bit only for now)
+            - **Int8WeightOnly**: 8-bit weight-only quantization
+            - **Int8DynamicActivationInt8Weight**: 8-bit quantization for both weights and activations
+            - **Float8WeightOnly**: float8-bit weight-only quantization
+            - **Float8DynamicActivationFloat8Weight**: float8-bit quantization for both weights and activations
             - **autoquant**: automatic quantization (uses the best quantization method for the model)
 
             ### Group Size