Create safetensors_converter

safetensors_converter

@@ -0,0 +1,72 @@
+# Adapted from https://github.com/DiffusionDalmation/pt_to_safetensors_converter_notebook
+
+import os
+from typing import Any, Dict
+import torch
+from safetensors.torch import save_file
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Supported file extensions
+SUPPORTED_EXTENSIONS = ['.pth', '.pt', '.bin', '.ckpt']
+
+def is_supported_file(input_path):
+    """Check if file has a supported extension"""
+    _, ext = os.path.splitext(input_path.lower())
+    return ext in SUPPORTED_EXTENSIONS
+
+def convert_file(input_path, output_path):
+    # Load the PyTorch model
+    model = torch.load(input_path, map_location=device)
+
+    if "string_to_param" in model: # Embeddings are a bit of a weird dict.
+        (s_model, dmeta) = process_embedding_file(model)
+    elif "state_dict" in model: # Ckpts or vaes are a standard list of layers.
+        (s_model, dmeta) = process_ckpt_file(model)
+    else: # No clue, try simple conversion.
+        s_model = model
+        dmeta = {"ckpt": None, "step": None, "dim": None}
+    
+    try:
+        save_file (s_model, output_path)
+    except Exception as e:
+        raise ValueError(f"Unknown filetype: {input_path} | {str(e)}")
+
+    return dmeta
+
+def process_embedding_file(model):
+    # Extract the embedding tensors
+    model_tensors = model.get('string_to_param').get('*')
+    s_model = {
+          'emb_params': model_tensors
+            }
+
+    # Metadata extraction.
+    dmeta = {"ckpt": None, "step": None, "dim": None}
+    if ('sd_checkpoint_name' in model) and (model['sd_checkpoint_name'] is not None):
+        dmeta["ckpt"] = model['sd_checkpoint_name']
+    if ('step' in model) and (model['step'] is not None):
+        dmeta["step"] = model["step"]
+    dmeta["dim"] = model_tensors.shape
+
+    return s_model, dmeta
+
+def process_ckpt_file(model):
+    # Extract the state dictionary
+    s_model = model["state_dict"]
+
+    # Metadata extraction.
+    dmeta = {"ckpt": None, "step": None, "dim": None}
+    dmeta["step"] = model.get('step', model.get('global_step'))
+
+    return s_model, dmeta
+    
+    if verbose:
+        # Print the requested training information, if it exists
+        step = model.get('step', model.get('global_step'))
+        if step is not None:
+            print(f"Trained for {step} steps.")
+        else:
+            print("Step not found in the model.")
+        print()
+    return s_model