Update app.py

app.py

@@ -4,6 +4,8 @@ from huggingface_hub import HfApi, upload_folder
 import gradio as gr
 import hf_utils
 import utils
+from safetensors import safe_open
+from safetensors.torch import save_file
 
 subprocess.run(["git", "clone", "https://github.com/huggingface/diffusers", "diffs"])
 
@@ -40,7 +42,7 @@ 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, sd_version, token, path_in_repo, ema):
+def convert_and_push(radio_model_names, input_model, ckpt_name, sd_version, token, path_in_repo, ema, safetensors):
     extract_ema = ema == "ema"
     
     if sd_version == None:
@@ -54,6 +56,16 @@ def convert_and_push(radio_model_names, input_model, ckpt_name, sd_version, toke
         # 1. Download the checkpoint file
         ckpt_path, revision = hf_utils.download_file(repo_id=model_id, filename=ckpt_name, token=token)
 
+        if safetensors == "yes":
+            tensors = {}
+            with safe_open(ckpt_path, framework="pt", device="cpu") as f:
+               for key in f.keys():
+                   tensors[key] = f.get_tensor(key)
+
+            new_checkpoint_path = "/".join(ckpt_path.split("/")[:-1] + ["model_safe.ckpt"])
+            torch.save(tensors, new_checkpoint_path)
+            ckpt_path = new_checkpoint_path
+
         # 2. Run the conversion script
         os.makedirs(model_id, exist_ok=True)
         run_command = [
@@ -133,6 +145,7 @@ with gr.Blocks() as demo:
                 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")
                 ema = gr.Radio(label="Extract EMA or non-EMA?", choices=["ema", "non-ema"])
+                safetensors = gr.Radio(label="Extract from safetensors", choices=["yes", "no"], value="no")
                 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")
@@ -163,7 +176,7 @@ with gr.Blocks() as demo:
 
     btn_convert.click(
         fn=convert_and_push,
-        inputs=[radio_model_names, input_model, radio_ckpts, radio_sd_version, input_token, path_in_repo, ema],
+        inputs=[radio_model_names, input_model, radio_ckpts, radio_sd_version, input_token, path_in_repo, ema, safetensors],
         outputs=error_output,
         scroll_to_output=True
     )