Improvements for Hardware Selection (#4)

- Improvements for Hardware Selection (63e70686c0557a2acf134691d0431f972ca878e3)


Co-authored-by: Radamés Ajna <[email protected]>

app.py

@@ -2,7 +2,7 @@ import os
 import subprocess
 from huggingface_hub import HfApi, upload_folder
 import gradio as gr
-
+import requests
 from huggingface_hub import whoami, list_models
 
 
@@ -69,7 +69,7 @@ def on_load_model(user_model_id, other_model_id, token):
     except Exception as e:
         return None, None, None, None, gr.update(value=error_str(e))
 
-def create_and_push(space_type, other_model_name, radio_model_names, model_id, title, description, prefix, update, token):
+def create_and_push(space_type, hardware, private_space, other_model_name, radio_model_names, model_id, title, description, prefix, update, token):
 
     try:
 
@@ -80,7 +80,14 @@ def create_and_push(space_type, other_model_name, radio_model_names, model_id, t
             exist_ok=update,
             repo_type="space",
             space_sdk="gradio",
+            private=private_space
         )
+        api_url = f'https://huggingface.co/api/spaces/{model_id}'
+        headers = { "Authorization" : f"Bearer {token}"}
+        # add HUGGING_FACE_HUB_TOKEN secret to new space
+        requests.post(f'{api_url}/secrets', json={"key":"HUGGING_FACE_HUB_TOKEN","value":token}, headers=headers)
+        # set new Space Hardware flavor
+        requests.post(f'{api_url}/hardware', json={'flavor': hardware}, headers=headers)
 
         # 2. Replace the name, title, and description in the template
         with open("template/app_simple.py" if space_type == "Simple" else "template/app_advanced.py", "r") as f:
@@ -137,6 +144,7 @@ with gr.Blocks() as demo:
                 gr.Markdown("#### 1. Choose a model")
                 input_token = gr.Textbox(
                     max_lines=1,
+                    type="password",
                     label="Enter your Hugging Face token",
                     placeholder="WRITE permission is required!",
                 )
@@ -158,8 +166,9 @@ with gr.Blocks() as demo:
                 - **Advanced** - Runs on CPU by default, with the option to upgrade to GPU. You can control image generation parameters: guidance, number of steps, image size, etc. Also supports **image-to-image** generation.""")
                 space_type =gr.Radio(label="Space type", choices=["Simple", "Advanced"], value="Simple")
                 update = gr.Checkbox(label="Update the space if it already exists?")
-                
-
+                private_space = gr.Checkbox(label="Private Space")
+                gr.Markdown("Choose the new Space Hardware <small>[check pricing page](https://huggingface.co/pricing#spaces), you need payment method to upgrade your Space hardware</small>")
+                hardware = gr.Dropdown(["cpu-basic","cpu-upgrade","t4-small","t4-medium","a10g-small","a10g-large"],value = "cpu-basic", label="Space Hardware")
                 brn_create = gr.Button("Create the space")
 
     error_output = gr.Markdown(label="Output")
@@ -181,7 +190,7 @@ with gr.Blocks() as demo:
 
     brn_create.click(
         fn=create_and_push,
-        inputs=[space_type, other_model_name, radio_model_names, name, title, description, prefix, update, input_token],
+        inputs=[space_type, hardware, private_space, other_model_name, radio_model_names, name, title, description, prefix, update, input_token],
         outputs=[error_output],
         scroll_to_output=True
     )
@@ -197,4 +206,4 @@ with gr.Blocks() as demo:
     """)
     
 demo.queue()
-demo.launch(debug=True)
+demo.launch(debug=True)