Update app.py

app.py

@@ -1,13 +1,22 @@
 import torch
-from safetensors.torch import save_file
+from safetensors.torch import save_file, load_file
 import gradio as gr
 import os
 
-def convert_embedding(sd15_embedding):
+def convert_embedding(uploaded_file):
     output_path = "embedding.safetensors"
+    file_extension = os.path.splitext(uploaded_file.name)[1]
 
-    sd15_embedding = torch.load(sd15_embedding.name, map_location=torch.device('cpu'), weights_only=True)
-    sd15_tensor = sd15_embedding['string_to_param']['*']
+    #The sample files are probably structured in these ways because the pt files were probably all created with automatic1111, and the safetensors files were probably created with kohya_ss
+    #If we learn of other programs that structure the embedding file differently, we'll have to adjust the logic.
+    if file_extension == '.pt':
+        sd15_embedding = torch.load(uploaded_file.name, map_location=torch.device('cpu'), weights_only=True)
+        sd15_tensor = sd15_embedding['string_to_param']['*']
+    elif file_extension == '.safetensors':
+        loaded_tensors = load_file(uploaded_file.name)
+        sd15_tensor = loaded_tensors['emb_params']
+    else:
+        raise ValueError("Unsupported file format")
     
     num_vectors = sd15_tensor.shape[0]
     clip_g_shape = (num_vectors, 1280)
@@ -30,3 +39,4 @@ iface = gr.Interface(
 
 if __name__ == "__main__":
     iface.launch()
+