Create app.py

app.py

@@ -4,86 +4,47 @@ 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(input_path, output_path):
+    # Load the .ckpt file
+    state_dict = torch.load(input_path, map_location='cpu')
+    # Save as .safetensors
+    save_file(state_dict, output_path)
 
-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
+def process(url, uploaded_file):
+    if url:
+        # Download the .ckpt file
+        local_filename = 'model.ckpt'
+        with requests.get(url, stream=True) as r:
+            r.raise_for_status()
+            with open(local_filename, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=8192):
+                    f.write(chunk)
+    elif uploaded_file is not None:
+        # Save uploaded file
+        local_filename = uploaded_file.name
+        with open(local_filename, 'wb') as f:
+            f.write(uploaded_file.read())
+    else:
+        return "Please provide a URL or upload a .ckpt file."
 
-        # 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
+    output_filename = local_filename.replace('.ckpt', '.safetensors')
+    # Convert the .ckpt to .safetensors
+    convert_ckpt_to_safetensors(local_filename, output_filename)
 
-        # 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)}"
+    # Clean up the input file
+    os.remove(local_filename)
 
-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
+    return output_filename
 
-# 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
+iface = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Textbox(label="URL of .ckpt file", placeholder="Enter the URL here"),
+        gr.File(label="Or upload a .ckpt file", file_types=['.ckpt'])
+    ],
+    outputs=gr.File(label="Converted .safetensors file"),
+    title="CKPT to SafeTensors Converter",
+    description="Convert .ckpt files to .safetensors format. Provide a URL or upload your .ckpt file."
+)
 
-# 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()
+iface.launch()