add to org option

app.py

@@ -18,19 +18,26 @@ def hello(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) ->
 
 
 def check_model_exists(
-    oauth_token: gr.OAuthToken | None, username, model_name, quantized_model_name
+    oauth_token: gr.OAuthToken | None, username, model_name, quantized_model_name, upload_to_community
 ):
     """Check if a model exists in the user's Hugging Face repository."""
     try:
         models = list_models(author=username, token=oauth_token.token)
+        community_models = list_models(author="bnb-community", token=oauth_token.token)
         model_names = [model.id for model in models]
-        if quantized_model_name:
-            repo_name = f"{username}/{quantized_model_name}"
+        community_model_names = [model.id for model in community_models]
+        if upload_to_community:
+            repo_name = f"bnb-community/{model_name.split('/')[-1]}-bnb-4bit"
         else:
-            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"
+            if quantized_model_name:
+                repo_name = f"{username}/{quantized_model_name}"
+            else:
+                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."
+        elif repo_name in community_model_names:
+            return f"Model '{repo_name}' already exists in the bnb-community organization."
         else:
             return None  # Model does not exist
     except Exception as e:
@@ -200,6 +207,7 @@ def save_model(
     auth_token=None,
     quantized_model_name=None,
     public=False,
+    upload_to_community=False,
     progress=gr.Progress(),
 ):
     progress(0.67, desc="Preparing to push")
@@ -214,10 +222,13 @@ def save_model(
         progress(0.75, desc="Preparing to push")
 
         # Prepare repo name and model card
-        if quantized_model_name:
-            repo_name = f"{username}/{quantized_model_name}"
+        if upload_to_community:
+            repo_name = f"bnb-community/{model_name.split('/')[-1]}-bnb-4bit"
         else:
-            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"
+            if quantized_model_name:
+                repo_name = f"{username}/{quantized_model_name}"
+            else:
+                repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"
 
         model_card = create_model_card(
             model_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4
@@ -291,6 +302,7 @@ def quantize_and_save(
     quant_storage_4,
     quantized_model_name,
     public,
+    upload_to_community,
     progress=gr.Progress(),
 ):
     if oauth_token is None:
@@ -308,7 +320,7 @@ def quantize_and_save(
         </div>
         """
     exists_message = check_model_exists(
-        oauth_token, profile.username, model_name, quantized_model_name
+        oauth_token, profile.username, model_name, quantized_model_name, upload_to_community
     )
     if exists_message:
         return f"""
@@ -341,6 +353,7 @@ def quantize_and_save(
             oauth_token,
             quantized_model_name,
             public,
+            upload_to_community,
             progress,
         )
         # Clean up the model to free memory
@@ -685,6 +698,28 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
                             interactive=True,
                             show_label=True,
                         )
+                        
+                    with gr.Row():
+                        upload_to_community = gr.Checkbox(
+                            label="🤗 Upload to bnb-community",
+                            info="If checked, the model will be uploaded to the bnb-community organization",
+                            value=False,
+                            interactive=True,
+                            show_label=True,
+                        )
+                        
+                    # Add event handler to disable and clear model name when uploading to community
+                    def toggle_model_name(upload_to_community_checked):
+                        return gr.update(
+                            interactive=not upload_to_community_checked,
+                            value="Can't change model name when uploading to community" if upload_to_community_checked else quantized_model_name.value
+                        )
+                    
+                    upload_to_community.change(
+                        fn=toggle_model_name,
+                        inputs=[upload_to_community],
+                        outputs=quantized_model_name
+                    )
 
         with gr.Column():
             quantize_button = gr.Button(
@@ -704,6 +739,7 @@ with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
             quant_storage_4,
             quantized_model_name,
             public,
+            upload_to_community,
         ],
         outputs=[output_link],
         show_progress="full",