converting cptk to safe tensors

app.py

@@ -0,0 +1,89 @@
+import gradio as gr
+import torch
+from safetensors.torch import save_file
+import requests
+import os
+
+def download_ckpt_file(ckpt_url):
+    """
+    Downloads the .ckpt file from the provided Hugging Face URL.
+    """
+    try:
+        # Get the filename from the URL
+        filename = ckpt_url.split("/")[-1]
+        response = requests.get(ckpt_url)
+        
+        # Save the file locally
+        with open(filename, 'wb') as f:
+            f.write(response.content)
+        
+        return filename
+    except Exception as e:
+        return None, f"Error downloading the file: {str(e)}"
+
+def convert_ckpt_to_safetensors(ckpt_file):
+    """
+    Converts a .ckpt file to safetensors format.
+    """
+    try:
+        # Load the checkpoint
+        checkpoint = torch.load(ckpt_file, map_location='cpu')
+        
+        # Ensure the checkpoint contains a 'state_dict'
+        if 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+
+        # Remove any prefixes if necessary (e.g., 'module.')
+        new_state_dict = {}
+        for key, value in state_dict.items():
+            if key.startswith('module.'):
+                new_key = key[len('module.'):]
+            else:
+                new_key = key
+            new_state_dict[new_key] = value
+
+        # Save to safetensors format
+        output_file = ckpt_file.replace(".ckpt", ".safetensors")
+        save_file(new_state_dict, output_file)
+        return output_file
+    except Exception as e:
+        return f"Error converting to safetensors: {str(e)}"
+
+def handle_conversion(ckpt_url):
+    """
+    Handles the entire process of downloading the .ckpt file from the link,
+    converting it to safetensors, and providing the user with the output file.
+    """
+    # Download the .ckpt file
+    filename = download_ckpt_file(ckpt_url)
+    if not filename:
+        return None, "Failed to download the file."
+    
+    # Convert the .ckpt file to safetensors
+    safetensors_file = convert_ckpt_to_safetensors(filename)
+    
+    # If the conversion is successful, return the safetensors file
+    if safetensors_file.endswith(".safetensors"):
+        return safetensors_file
+    else:
+        return None, safetensors_file
+
+# Gradio Interface
+def convert_and_download(ckpt_url):
+    safetensors_file, message = handle_conversion(ckpt_url)
+    
+    if safetensors_file:
+        return safetensors_file  # Provide the converted file for download
+    else:
+        return message  # Provide the error message
+
+# Define the Gradio interface
+gr.Interface(
+    fn=convert_and_download,
+    inputs=gr.Textbox(label="Hugging Face CKPT File URL", placeholder="Enter the link to a .ckpt file"),
+    outputs=gr.File(label="Download .safetensors file"),
+    title="CKPT to Safetensors Converter",
+    description="Enter the Hugging Face URL for a .ckpt file and convert it to safetensors format."
+).launch()