change dtype

app.py

@@ -25,7 +25,7 @@ def check_model_exists(oauth_token: gr.OAuthToken | None, username, model_name,
         if quantized_model_name : 
             repo_name = f"{username}/{quantized_model_name}"
         else : 
-            repo_name = f"{username}/{model_name.split('/')[-1]}-BNB-INT4"
+            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"
 
         if repo_name in model_names:
             return f"Model '{repo_name}' already exists in your repository."
@@ -83,7 +83,7 @@ def quantize_model(model_name, quant_type_4, double_quant_4, compute_type_4, qua
         bnb_4bit_compute_dtype=DTYPE_MAPPING[compute_type_4],
     )
 
-    model = AutoModel.from_pretrained(model_name, quantization_config=quantization_config, device_map="cpu", use_auth_token=auth_token.token)
+    model = AutoModel.from_pretrained(model_name, quantization_config=quantization_config, device_map="cpu", use_auth_token=auth_token.token, torch_dtype=torch.bfloat16)
     for _ , module in model.named_modules():
         if isinstance(module, Linear4bit):
             module.to("cuda")
@@ -99,7 +99,7 @@ def save_model(model, model_name, quant_type_4, double_quant_4, compute_type_4,
         if quantized_model_name : 
             repo_name = f"{username}/{quantized_model_name}"
         else :   
-            repo_name = f"{username}/{model_name.split('/')[-1]}-BNB-INT4"
+            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"
             
 
         model_card = create_model_card(repo_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4)
@@ -425,5 +425,4 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
 
 if __name__ == "__main__":
     demo.launch(share=True)
-# Launch the app
-# demo.launch(share=True, debug=True)
+