updating to use AoBaseConfig

app.py

@@ -12,12 +12,15 @@ from torchao.quantization import (
     Int8WeightOnlyConfig,
     Int8DynamicActivationInt8WeightConfig,
     Float8WeightOnlyConfig,
+    Float8DynamicActivationFloat8WeightConfig,
 )
 
 MAP_QUANT_TYPE_TO_NAME = {
     "int4_weight_only": "int4wo",
     "int8_weight_only": "int8wo",
-    "int8_dynamic_activation_int8_weight": "int8da8w",
+    "int8_dynamic_activation_int8_weight": "int8da8w8",
+    "float8_weight_only": "float8wo",
+    "float8_dynamic_activation_float8_weight": "float8da8w8",
     "autoquant": "autoquant",
 }
 MAP_QUANT_TYPE_TO_CONFIG = {
@@ -25,6 +28,7 @@ MAP_QUANT_TYPE_TO_CONFIG = {
     "int8_weight_only": Int8WeightOnlyConfig,
     "int8_dynamic_activation_int8_weight": Int8DynamicActivationInt8WeightConfig,
     "float8_weight_only": Float8WeightOnlyConfig,
+    "float8_dynamic_activation_float8_weight": Float8DynamicActivationFloat8WeightConfig,
 }
 
 
@@ -164,16 +168,30 @@ It's quantized using the TorchAO library using the [torchao-my-repo](https://hug
 
 
 def quantize_model(
-    model_name, quantization_type, group_size=128, auth_token=None, username=None
+    model_name, quantization_type, group_size=128, auth_token=None, username=None, progress=gr.Progress()
 ):
     print(f"Quantizing model: {quantization_type}")
+    progress(0, desc="Preparing Quantization")
     if (
-        quantization_type == "int4_weight_only"
-        or quantization_type == "int8_weight_only"
+        quantization_type == "int8_weight_only"
     ):
-        quantization_config = TorchAoConfig(quantization_type, group_size=group_size)
-    else:
+        quant_config = MAP_QUANT_TYPE_TO_CONFIG[quantization_type](
+            group_size=group_size
+        )
+        quantization_config = TorchAoConfig(quant_config)
+    elif quantization_type == "int4_weight_only":
+        from torchao.dtypes import Int4CPULayout
+        
+        quant_config = MAP_QUANT_TYPE_TO_CONFIG[quantization_type](
+            group_size=group_size, layout=Int4CPULayout()
+        )
+        quantization_config = TorchAoConfig(quant_config)
+    elif quantization_type == "autoquant":
         quantization_config = TorchAoConfig(quantization_type)
+    else:
+        quant_config = MAP_QUANT_TYPE_TO_CONFIG[quantization_type]()
+        quantization_config = TorchAoConfig(quant_config)
+    progress(0.10, desc="Quantizing model")
     model = AutoModel.from_pretrained(
         model_name,
         torch_dtype="auto",
@@ -181,7 +199,7 @@ def quantize_model(
         device_map="cpu",
         use_auth_token=auth_token.token,
     )
-
+    progress(0.45, desc="Quantization completed")
     return model
 
 
@@ -193,7 +211,10 @@ def save_model(
     username=None,
     auth_token=None,
     quantized_model_name=None,
+    public=True,
+    progress=gr.Progress(),
 ):
+    progress(0.50, desc="Preparing to push")
     print("Saving quantized model")
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Load and save the tokenizer
@@ -203,10 +224,11 @@ def save_model(
         tokenizer.save_pretrained(tmpdirname, use_auth_token=auth_token.token)
 
         # Save the model
+        progress(0.60, desc="Saving model")
         model.save_pretrained(
             tmpdirname, safe_serialization=False, use_auth_token=auth_token.token
         )
-
+        
         if quantized_model_name:
             repo_name = f"{username}/{quantized_model_name}"
         else:
@@ -217,19 +239,21 @@ def save_model(
                 repo_name = f"{username}/{model_name.split('/')[-1]}-ao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}-gs{group_size}"
             else:
                 repo_name = f"{username}/{model_name.split('/')[-1]}-ao-{MAP_QUANT_TYPE_TO_NAME[quantization_type.lower()]}"
-
+        progress(0.70, desc="Creating model card")
         model_card = create_model_card(model_name, quantization_type, group_size)
         with open(os.path.join(tmpdirname, "README.md"), "w") as f:
             f.write(model_card)
         # Push to Hub
         api = HfApi(token=auth_token.token)
-        api.create_repo(repo_name, exist_ok=True)
+        api.create_repo(repo_name, exist_ok=True, private=not public)
+        progress(0.80, desc="Pushing to Hub")
         api.upload_folder(
             folder_path=tmpdirname,
             repo_id=repo_name,
             repo_type="model",
         )
-
+        progress(1.00, desc="Pushing to Hub completed")
+    
     import io
     from contextlib import redirect_stdout
     import html
@@ -273,6 +297,7 @@ def quantize_and_save(
     quantization_type,
     group_size,
     quantized_model_name,
+    public,
 ):
     if oauth_token is None:
         return """
@@ -332,8 +357,10 @@ def quantize_and_save(
             profile.username,
             oauth_token,
             quantized_model_name,
+            public,
         )
     except Exception as e:
+        # raise e
         return str(e)
 
 
@@ -464,24 +491,44 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
                             "int4_weight_only",
                             "int8_weight_only",
                             "int8_dynamic_activation_int8_weight",
+                            "float8_weight_only",
+                            "float8_dynamic_activation_float8_weight",
                             "autoquant",
                         ],
                         value="int8_weight_only",
                         filterable=False,
                         show_label=False,
                     )
+
                     group_size = gr.Textbox(
                         info="Group Size (only for int4_weight_only and int8_weight_only)",
                         value="128",
-                        interactive=True,
+                        interactive=(quantization_type.value == "int4_weight_only" or quantization_type.value == "int8_weight_only"),
                         show_label=False,
                     )
-                    quantized_model_name = gr.Textbox(
-                        info="Custom name for your quantized model (optional)",
-                        value="",
-                        interactive=True,
-                        show_label=False,
+
+            gr.Markdown(
+                        """
+                        ### 💾 Saving Settings
+                        """
                     )
+            with gr.Row():
+                quantized_model_name = gr.Textbox(
+                    label="✏️ Model Name",
+                    info="Model Name (optional : to override default)",
+                    value="",
+                    interactive=True,
+                    elem_classes="model-name-textbox",
+                    show_label=False,
+                )
+            with gr.Row():
+                public = gr.Checkbox(
+                    label="🌐 Make model public",
+                    info="If checked, the model will be publicly accessible",
+                    value=True,
+                    interactive=True,
+                    show_label=True,
+                )
 
         with gr.Column():
             quantize_button = gr.Button(
@@ -517,11 +564,10 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
             - int8_weight_only typically reduces size by about 50%
             """
         )
-
     # Keep existing click handler
     quantize_button.click(
         fn=quantize_and_save,
-        inputs=[model_name, quantization_type, group_size, quantized_model_name],
+        inputs=[model_name, quantization_type, group_size, quantized_model_name, public],
         outputs=[output_link],
     )