Update app.py

app.py

@@ -1,38 +1,21 @@
 import gradio as gr
 import torch
-from safetensors.torch import save_file
+from safetensors.torch import save_model
 import requests
 import os
 
-def resolve_shared_tensors(state_dict):
- storage_map = {}
- for k, v in state_dict.items():
- if not isinstance(v, torch.Tensor):
- continue
- storage_data_ptr = v.storage().data_ptr()
- if storage_data_ptr in storage_map:
- # Clone the tensor to create independent storage
- state_dict[k] = v.clone()
- else:
- storage_map[storage_data_ptr] = k
- return state_dict
-
 def convert_ckpt_to_safetensors(input_path, output_path):
  # Load the .ckpt file
  state_dict = torch.load(input_path, map_location='cpu')
 
- # Check for nested 'state_dict' key
+ # If the checkpoint has a 'state_dict' key, extract it
  if 'state_dict' in state_dict:
  state_dict = state_dict['state_dict']
+ elif 'model' in state_dict:
+ state_dict = state_dict['model']
 
- # Filter out non-tensor entries
- tensor_state_dict = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)}
-
- # Resolve shared tensors
- tensor_state_dict = resolve_shared_tensors(tensor_state_dict)
-
- # Save as .safetensors
- save_file(tensor_state_dict, output_path)
+ # Save the entire state dictionary, including non-tensor entries
+ save_model(state_dict, output_path)
 
 def process(url, uploaded_file):
  if url: