Fix embedding reparser regression

tools/torch_tools.py

@@ -16,16 +16,29 @@ def get_target_dtype_ref(target_dtype: str) -> torch.dtype:
 		raise ValueError(f"Invalid target_dtype: {target_dtype}")
 
 def convert_ckpt_to_safetensors(ckpt_upload: io.BytesIO, target_dtype) -> dict:
+	if isinstance(ckpt_upload, bytes):
+		ckpt_upload = io.BytesIO(ckpt_upload)
+	
 	target_dtype = get_target_dtype_ref(target_dtype)
-	ckpt_data = ckpt_upload.getvalue()
 	
 	# Load the checkpoint
-	checkpoint = torch.load(ckpt_data, map_location="cpu")
+	loaded_dict = torch.load(ckpt_upload, map_location="cpu")
 	
-	# Convert the checkpoint to a dictionary of tensors
 	tensor_dict = {}
-	for key, val in checkpoint.items():
-		tensor_dict[key] = val.to(dtype=target_dtype)
+	
+	is_embedding = 'string_to_param' in loaded_dict
+	if is_embedding:
+		emb_tensor = loaded_dict.get('string_to_param', {}).get('*', None)
+		if emb_tensor is not None:
+			emb_tensor = emb_tensor.to(dtype=target_dtype)
+			tensor_dict = {
+				'emb_params': emb_tensor
+			}
+	else:
+		# Convert weights in a checkpoint to a dictionary of tensors
+		for key, val in loaded_dict.items():
+			if isinstance(val, torch.Tensor):
+				tensor_dict[key] = val.to(dtype=target_dtype)
 	
 	return tensor_dict