Update app.py

app.py

@@ -1,9 +1,8 @@
 import os
 import subprocess
-from huggingface_hub import HfApi, upload_folder
+from huggingface_hub import HfApi, upload_folder, whoami, list_models, hf_hub_download, upload_file
 import gradio as gr
 import requests
-from huggingface_hub import whoami, list_models
 
 
 def error_str(error, title="Error"):
@@ -69,7 +68,33 @@ 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, hardware, private_space, other_model_name, radio_model_names, model_id, title, description, prefix, update, token):
+def add_space_badge_to_model_card(model_id, token):
+
+    readme_file = 'README.md'
+    model_card = hf_hub_download(repo_id=model_id, filename=readme_file, token=token)
+
+    with open(model_card, "r") as f:
+        content = f.read()
+
+    content = content.split("---\n")
+    content[2] = "[![Open In Spaces](https://camo.githubusercontent.com/00380c35e60d6b04be65d3d94a58332be5cc93779f630bcdfc18ab9a3a7d3388/68747470733a2f2f696d672e736869656c64732e696f2f62616467652f25463025394625413425393725323048756767696e67253230466163652d5370616365732d626c7565)](https://huggingface.co/spaces/" + model_id + ")\n" + content[2]
+    content = "---\n".join(content)
+
+    with open(readme_file, "w") as f:
+        f.write(content)
+
+    upload_file(
+        path_or_fileobj=readme_file,
+        path_in_repo=readme_file,
+        repo_id=model_id,
+        token=token,
+        create_pr=True,
+        commit_message="Add Space badge to model card",
+    )
+
+    os.remove(readme_file)
+
+def create_and_push(space_type, hardware, private_space, add_badge, other_model_name, radio_model_names, model_id, title, description, prefix, update, token):
 
     try:
 
@@ -124,25 +149,34 @@ def create_and_push(space_type, hardware, private_space, other_model_name, radio
         # 5. Delete the app.py file
         os.remove("app.py")
 
-        return f"""Successfully created space at: <a href="{repo_url}" target="_blank">{repo_url}</a>"""
+        # 6. Add the Space badge to the model card
+        if add_badge:
+            add_space_badge_to_model_card(model_id, token)
+
+        return f"""
+                    Successfully created space at: <a href="{repo_url}" target="_blank">{repo_url}</a> <br>
+                    Opened a PR to add the space badge at: <a href="https://huggingface.co/{model_id}" target="_blank">https://huggingface.co/{model_id}</a>
+                """
 
     except Exception as e:
         return error_str(e)
 
 
 DESCRIPTION = """### Create a gradio space for your Diffusers🧨 model
-                With this space, you can easily create a gradio demo for your Diffusers model and share it with the community.<br>
-                1️⃣ Make sure you have created your hugging face account<br>
-                2️⃣ Generate a token here with write access<br>
-                3️⃣ Choose a stable diffusion base model, there are thousands of them here<br>
-                4️⃣ Choose Space type<br>
-                5️⃣ Choose the new Space Hardware<br>
-                It is done.
+                With this space, you can easily create a gradio demo for your Diffusers model and share it with the community.
             """
+            # <br>
+            #     1️⃣ Make sure you have created your hugging face account<br>
+            #     2️⃣ Generate a token here with write access<br>
+            #     3️⃣ Choose a stable diffusion base model, there are thousands of them here<br>
+            #     4️⃣ Choose Space type<br>
+            #     5️⃣ Choose the new Space Hardware<br>
+            #     It is done.
+            # """
 
 with gr.Blocks() as demo:
 
-    # gr.Markdown(DESCRIPTION)
+    gr.Markdown(DESCRIPTION)
     with gr.Row():
 
         with gr.Column(scale=11):
@@ -167,15 +201,20 @@ with gr.Blocks() as demo:
                 title = gr.Textbox(label="Title", placeholder="e.g. Diffusers Demo")
                 description = gr.Textbox(label="Description", placeholder="e.g. Demo for my awesome Diffusers model", lines=5)
                 prefix = gr.Textbox(label="Prefix tokens", placeholder="Tokens that are required to be present in the prompt, e.g. `rick and morty style`")
+                
                 gr.Markdown("""#### Choose space type
                 - **Simple** - Runs on GPU using Hugging Face inference API, but you cannot control image generation parameters.
                 - **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")
+                add_badge = gr.Checkbox(label="Add Space badge to the model card (will open a PR)")
+                
                 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")
+                
+                btn_create = gr.Button("Create the space")
 
     error_output = gr.Markdown(label="Output")
 
@@ -194,9 +233,9 @@ with gr.Blocks() as demo:
         queue=False,
         scroll_to_output=True)
 
-    brn_create.click(
+    btn_create.click(
         fn=create_and_push,
-        inputs=[space_type, hardware, private_space, other_model_name, radio_model_names, name, title, description, prefix, update, input_token],
+        inputs=[space_type, hardware, private_space, add_badge, other_model_name, radio_model_names, name, title, description, prefix, update, input_token],
         outputs=[error_output],
         scroll_to_output=True
     )
@@ -212,4 +251,4 @@ with gr.Blocks() as demo:
     """)
     
 demo.queue()
-demo.launch(debug=True)
+demo.launch(debug=True)