Update app.py

app.py

@@ -40,8 +40,11 @@ def get_ckpt_names(token, radio_model_names, input_model):
     except Exception as e:
         return error_str(e), gr.update(choices=[]), None
 
-def convert_and_push(radio_model_names, input_model, ckpt_name, token, path_in_repo):
+def convert_and_push(radio_model_names, input_model, ckpt_name, sd_version, token, path_in_repo):
     
+    if sd_version == None:
+        return error_str("You must select a stable diffusion version.", title="Invalid input")
+
     model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names
 
     try:
@@ -60,6 +63,8 @@ def convert_and_push(radio_model_names, input_model, ckpt_name, token, path_in_r
                 ckpt_path,
                 "--dump_path" ,
                 model_id,
+                "--sd_version",
+                sd_version
             ]
         )
 
@@ -84,6 +89,8 @@ def convert_and_push(radio_model_names, input_model, ckpt_name, token, path_in_r
 
         return f"""Successfully converted the checkpoint and opened a PR to add the weights to the model repo.
                 You can view and merge the PR [here]({hf_utils.get_pr_url(HfApi(token=token), model_id, commit_message)})."""
+
+        return "Done"
     
     except Exception as e:
         return error_str(e)
@@ -106,6 +113,7 @@ with gr.Blocks() as demo:
                     max_lines=1,
                     label="Enter your Hugging Face token",
                     placeholder="READ permission is enough",
+                    type='password'
                 )
                 gr.Markdown("You can get a token [here](https://huggingface.co/settings/tokens)")
                 with gr.Group(visible=False) as group_model:
@@ -124,6 +132,7 @@ with gr.Blocks() as demo:
                 gr.Markdown("## 2. Convert to Diffusers🧨")
                 radio_ckpts = gr.Radio(label="Choose the checkpoint to convert", visible=False)
                 path_in_repo = gr.Textbox(label="Path where the weights will be saved", placeholder="Leave empty for root folder")
+                radio_sd_version = gr.Radio(label="Choose the model version", choices=["v1", "v2", "v2.1"])
                 gr.Markdown("Conversion may take a few minutes.")
                 btn_convert = gr.Button("Convert & Push")
 
@@ -153,7 +162,7 @@ with gr.Blocks() as demo:
 
     btn_convert.click(
         fn=convert_and_push,
-        inputs=[radio_model_names, input_model, radio_ckpts, input_token, path_in_repo],
+        inputs=[radio_model_names, input_model, radio_ckpts, radio_sd_version, input_token, path_in_repo],
         outputs=error_output,
         scroll_to_output=True
     )
@@ -169,4 +178,4 @@ with gr.Blocks() as demo:
     """)
     
 demo.queue()
-demo.launch(share=utils.is_google_colab())
+demo.launch(debug=True, share=utils.is_google_colab())