Update app.py

app.py

@@ -48,14 +48,14 @@ def on_token_change(token):
 def on_load_model(user_model_id, other_model_id, token):
 
     if not user_model_id and not other_model_id:
-        return None, None, None, None, gr.update(value=error_str("Please enter a model ID."))
+        return None, None, None, None, gr.update(value=error_str("Please enter a model ID.")), None
 
     try:
         model_id = url_to_model_id(other_model_id) if other_model_id else user_model_id
         original_model_id = model_id
 
         if not has_diffusion_model(model_id, token):
-            return None, None, None, None, gr.update(value=error_str("There are no diffusion weights in the model you selected."))
+            return None, None, None, None, gr.update(value=error_str("There are no diffusion weights in the model you selected.")), None
 
         user = whoami(token=token)
         model_id = user["name"] + "/" + model_id.split("/")[-1]
@@ -63,10 +63,10 @@ def on_load_model(user_model_id, other_model_id, token):
         
         description = f"""Demo for <a href="https://huggingface.co/{original_model_id}">{title}</a> Stable Diffusion model."""
 
-        return gr.update(visible=True), gr.update(value=model_id), gr.update(value=title), gr.update(value=description), None
+        return gr.update(visible=True), gr.update(value=model_id), gr.update(value=title), gr.update(value=description), None, original_model_id
     
     except Exception as e:
-        return None, None, None, None, gr.update(value=error_str(e))
+        return None, None, None, None, gr.update(value=error_str(e)), None
 
 def add_space_badge_to_model_card(model_id, token):
 
@@ -94,7 +94,7 @@ def add_space_badge_to_model_card(model_id, token):
 
     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):
+def create_and_push(space_type, hardware, private_space, add_badge, other_model_name, radio_model_names, model_id, title, description, prefix, update, token, original_model_id):
 
     try:
 
@@ -151,7 +151,7 @@ def create_and_push(space_type, hardware, private_space, add_badge, other_model_
 
         # 6. Add the Space badge to the model card
         if add_badge:
-            add_space_badge_to_model_card(model_id, token)
+            add_space_badge_to_model_card(original_model_id, token)
 
         return f"""
                     Successfully created space at: <a href="{repo_url}" target="_blank">{repo_url}</a> <br>
@@ -200,6 +200,7 @@ with gr.Blocks() as demo:
                 name = gr.Textbox(label="Name", placeholder="e.g. diffusers-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)
+                original_model_id = gr.Textbox(visible=True)
                 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
@@ -229,13 +230,13 @@ with gr.Blocks() as demo:
     btn_load.click(
         fn=on_load_model,
         inputs=[radio_model_names, other_model_name, input_token],
-        outputs=[group_create, name, title, description, error_output],
+        outputs=[group_create, name, title, description, error_output, original_model_id],
         queue=False,
         scroll_to_output=True)
 
     btn_create.click(
         fn=create_and_push,
-        inputs=[space_type, hardware, private_space, add_badge, 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, original_model_id],
         outputs=[error_output],
         scroll_to_output=True
     )