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()