Update app.py

app.py

@@ -1,55 +1,85 @@
+import os
 import subprocess
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, upload_folder
 import gradio as gr
+import hf_utils
 
 subprocess.run(["git", "clone", "https://github.com/huggingface/diffusers.git", "diffs"])
 
 def error_str(error, title="Error"):
     return f"""#### {title}
-            {error}"""
+            {error}"""  if error else ""
+
+def on_token_change(token):
+    model_names, error = hf_utils.get_my_model_names(token)
+    if model_names:
+        model_names.append("Other")
+
+    return gr.update(visible=bool(model_names)), gr.update(choices=model_names, value=model_names[0] if model_names else None), gr.update(value=error_str(error))
 
 def url_to_model_id(model_id_str):
     return model_id_str.split("/")[-2] + "/" + model_id_str.split("/")[-1] if model_id_str.startswith("https://huggingface.co/") else model_id_str
     
-def get_ckpt_names(model_id = "nitrosocke/mo-di-diffusion"):
+def get_ckpt_names(token, radio_model_names, input_model):
     
-    if model_id == "":
-        return error_str("Please enter a model name.", title="Invalid input"), None, None
+    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names
+
+    if token == "" or model_id == "":
+        return error_str("Please enter both a token and a model name.", title="Invalid input"), gr.update(choices=[]), gr.update(visible=False)
 
     try:
-        api = HfApi()
-        ckpt_files = [f for f in api.list_repo_files(url_to_model_id(model_id)) if f.endswith(".ckpt")]
+        api = HfApi(token=token)
+        ckpt_files = [f for f in api.list_repo_files(repo_id=model_id) if f.endswith(".ckpt")]
         
-        if len(ckpt_files) == 0:
-            return error_str("No checkpoint files found in the model repo."), None, None
+        if not ckpt_files:
+            return error_str("No checkpoint files found in the model repo."), gr.update(choices=[]), gr.update(visible=False)
         
-        return None, gr.update(choices=ckpt_files, visible=True), gr.update(visible=True)
+        return None, gr.update(choices=ckpt_files, value=ckpt_files[0], visible=True), gr.update(visible=True)
         
     except Exception as e:
-        return error_str(e), None, None
+        return error_str(e), gr.update(choices=[]), None
 
-def convert(model_id, ckpt_name, token = "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"):
+def convert_and_push(radio_model_names, input_model, ckpt_name, token):
     
-    model_id = url_to_model_id(model_id)
-
-    # 1. Download the checkpoint file
-    ckpt_path = hf_hub_download(repo_id=model_id, filename=ckpt_name)
-
-    # 2. Run the conversion script
-    subprocess.run(
-        [
-            "python3",
-            "./diffs/scripts/convert_original_stable_diffusion_to_diffusers.py",
-            "--checkpoint_path",
-            ckpt_path,
-            "--dump_path" ,
-            model_id,
-        ]
-    )
+    model_id = url_to_model_id(input_model) if radio_model_names == "Other" else radio_model_names
+
+    try:
+        model_id = url_to_model_id(model_id)
+
+        # 1. Download the checkpoint file
+        ckpt_path, revision = hf_utils.download_file(repo_id=model_id, filename=ckpt_name, token=token)
+
+        # 2. Run the conversion script
+        subprocess.run(
+            [
+                "python3",
+                "./diffs/scripts/convert_original_stable_diffusion_to_diffusers.py",
+                "--checkpoint_path",
+                ckpt_path,
+                "--dump_path" ,
+                model_id,
+            ]
+        )
+
+        # 3. Push to the model repo
+        upload_folder(
+            folder_path=model_id,
+            repo_id=model_id,
+            token=token,
+            create_pr=True,
+        )
+
+        # # 4. Delete the downloaded checkpoint file, yaml files, and the converted model folder
+        hf_utils.delete_file(revision)
+        subprocess.run(["rm", "-rf", model_id.split('/')[0]])
+        import glob
+        for f in glob.glob("*.yaml*"):
+            subprocess.run(["rm", "-rf", f])
+
+        return "Success"
     
-    import os
-    return f"""files in current directory:
-    {[f for f in os.listdir(".")]}"""
+    except Exception as e:
+        return error_str(e)
 
     
 with gr.Blocks() as demo:
@@ -57,40 +87,61 @@ with gr.Blocks() as demo:
     with gr.Row():
 
         with gr.Column(scale=11):
-            with gr.Group():
+            with gr.Column():
                 gr.Markdown("## 1. Load model info")
                 input_token = gr.Textbox(
                     max_lines=1,
-                    label="Hugging Face token",
+                    label="Enter your Hugging Face token",
                     placeholder="hf_...",
                 )
-                gr.Markdown("Get your token [here](https://huggingface.co/settings/tokens).")
-                input_model = gr.Textbox(
-                    max_lines=1,
-                    label="Model name or URL",
-                    placeholder="username/model_name",
-                )
+                gr.Markdown("You can get a token [here](https://huggingface.co/settings/tokens).")
+                with gr.Group(visible=False) as group_model:
+                    radio_model_names = gr.Radio(label="Choose a model")
+                    input_model = gr.Textbox(
+                        max_lines=1,
+                        label="Model name or URL",
+                        placeholder="username/model_name",
+                        visible=False,
+                    )
 
             btn_get_ckpts = gr.Button("Load")
 
-        with gr.Column(scale=10, visible=False) as col_convert:
-            gr.Markdown("## 2. Convert to Diffusers🧨")
-            radio_ckpts = gr.Radio(label="Choose a checkpoint to convert", visible=False)
-            btn_convert = gr.Button("Convert")
+        with gr.Column(scale=10):
+            with gr.Column(visible=False) as group_convert:
+                gr.Markdown("## 2. Convert to Diffusers🧨")
+                radio_ckpts = gr.Radio(label="Choose the checkpoint to convert", visible=False)
+                gr.Markdown("Conversion may take a few minutes.")
+                btn_convert = gr.Button("Convert & Push")
 
     error_output = gr.Markdown(label="Output")
+
+    input_token.change(
+        fn=on_token_change,
+        inputs=input_token,
+        outputs=[group_model, radio_model_names, error_output],
+        queue=False,
+        scroll_to_output=True)
+    
+    radio_model_names.change(
+        lambda x: gr.update(visible=x == "Other"),
+        inputs=radio_model_names,
+        outputs=input_model,
+        queue=False,
+        scroll_to_output=True)
+    
     btn_get_ckpts.click(
         fn=get_ckpt_names,
-        inputs=[input_model],
-        outputs=[error_output, radio_ckpts, col_convert],
+        inputs=[input_token, radio_model_names, input_model],
+        outputs=[error_output, radio_ckpts, group_convert],
         scroll_to_output=True
     )
 
     btn_convert.click(
-        fn=convert,
-        inputs=[input_model, radio_ckpts, input_token],
+        fn=convert_and_push,
+        inputs=[radio_model_names, input_model, radio_ckpts, input_token],
         outputs=error_output,
         scroll_to_output=True
     )
 
-demo.launch()
+demo.queue()
+demo.launch()