WireCell
/

dnn-roi

@@ -1,180 +0,0 @@
-#!/usr/bin/env python
-# filepath: convert_pth_to_safetensors.py
-import os
-import argparse
-import torch
-from safetensors.torch import save_file
-from pathlib import Path
-def convert_pth_to_safetensors(pth_path, output_dir=None, output_name=None):
-    """
-    Convert a PyTorch pickle (.pth) model to safetensors format.
-    Args:
-        pth_path (str): Path to the PyTorch .pth file
-        output_dir (str, optional): Directory to save the converted model. Defaults to same directory as input.
-        output_name (str, optional): Name for the output file. Defaults to input filename with .safetensors extension.
-    Returns:
-        str: Path to the saved safetensors file
-    """
-    print(f"Loading PyTorch model from: {pth_path}")
-    # Load PyTorch model
-    try:
-        state_dict = torch.load(pth_path, map_location="cpu")
-    except Exception as e:
-        raise RuntimeError(f"Failed to load PyTorch model: {e}")
-    # Handle different types of saved objects
-    if isinstance(state_dict, dict):
-        # Check if this is a state_dict or a full model save
-        if all(isinstance(v, torch.Tensor) for v in state_dict.values()):
-            # It's already a state_dict
-            print(f"Loaded state_dict with {len(state_dict)} parameters")
-        elif 'state_dict' in state_dict:
-            # It's a checkpoint with 'state_dict' key
-            state_dict = state_dict['state_dict']
-            print(f"Extracted state_dict from checkpoint with {len(state_dict)} parameters")
-        elif 'model_state_dict' in state_dict:
-            # It's a checkpoint with 'model_state_dict' key
-            state_dict = state_dict['model_state_dict']
-            print(f"Extracted model_state_dict from checkpoint with {len(state_dict)} parameters")
-        else:
-            # Try to find a key that contains tensors
-            tensor_keys = [k for k, v in state_dict.items() if isinstance(v, dict) and
-                          any(isinstance(item, torch.Tensor) for item in v.values())]
-            if tensor_keys:
-                state_dict = state_dict[tensor_keys[0]]
-                print(f"Extracted state_dict from key '{tensor_keys[0]}' with {len(state_dict)} parameters")
-            else:
-                raise ValueError("Could not find state_dict in the loaded file")
-    elif hasattr(state_dict, 'state_dict'):
-        # It's a full model object
-        state_dict = state_dict.state_dict()
-        print(f"Extracted state_dict from model object with {len(state_dict)} parameters")
-    else:
-        raise ValueError("Unsupported format: loaded object is not a state_dict or model")
-    # Ensure all values are tensors
-    for k, v in list(state_dict.items()):
-        if not isinstance(v, torch.Tensor):
-            print(f"Warning: Removing non-tensor value for key '{k}' of type {type(v)}")
-            state_dict.pop(k)
-    # Determine output path
-    if output_dir is None:
-        output_dir = os.path.dirname(pth_path)
-    if output_name is None:
-        base_name = os.path.basename(pth_path)
-        output_name = os.path.splitext(base_name)[0] + ".safetensors"
-    os.makedirs(output_dir, exist_ok=True)
-    output_path = os.path.join(output_dir, output_name)
-    # Save to safetensors format
-    print(f"Saving model to: {output_path}")
-    try:
-        save_file(state_dict, output_path)
-        print(f"Successfully saved safetensors file: {output_path}")
-        return output_path
-    except Exception as e:
-        raise RuntimeError(f"Failed to save safetensors file: {e}")
-def convert_directory(input_dir, output_dir=None, recursive=False, file_pattern="*.pth"):
-    """
-    Convert all PyTorch .pth models in a directory to safetensors format.
-    Args:
-        input_dir (str): Input directory containing PyTorch models
-        output_dir (str, optional): Output directory for safetensors files. Defaults to input_dir.
-        recursive (bool, optional): Whether to recursively search for models in subdirectories. Defaults to False.
-        file_pattern (str, optional): File pattern to match. Defaults to "*.pth".
-    Returns:
-        list: List of paths to converted safetensors files
-    """
-    if output_dir is None:
-        output_dir = input_dir
-    converted_files = []
-    # Find all PyTorch files
-    input_path = Path(input_dir)
-    if recursive:
-        pth_files = list(input_path.rglob(file_pattern))
-    else:
-        pth_files = list(input_path.glob(file_pattern))
-    if not pth_files:
-        print(f"No PyTorch files found in {input_dir} with pattern {file_pattern}")
-        return converted_files
-    print(f"Found {len(pth_files)} PyTorch files to convert")
-    # Convert each file
-    for pth_file in pth_files:
-        relative_path = pth_file.relative_to(input_path)
-        target_dir = Path(output_dir) / relative_path.parent
-        target_dir.mkdir(parents=True, exist_ok=True)
-        output_name = pth_file.stem + ".safetensors"
-        try:
-            converted_file = convert_pth_to_safetensors(
-                str(pth_file),
-                str(target_dir),
-                output_name
-            )
-            converted_files.append(converted_file)
-        except Exception as e:
-            print(f"Error converting {pth_file}: {e}")
-    return converted_files
-def main():
-    parser = argparse.ArgumentParser(description="Convert PyTorch .pth models to safetensors format")
-    parser.add_argument("input", help="Input PyTorch model file or directory")
-    parser.add_argument("--output", "-o", help="Output file or directory for safetensors files")
-    parser.add_argument("--recursive", "-r", action="store_true",
-                        help="Recursively search for PyTorch files in subdirectories")
-    parser.add_argument("--pattern", "-p", default="*.pth",
-                        help="File pattern to match when searching directories (default: *.pth)")
-    args = parser.parse_args()
-    input_path = Path(args.input)
-    if input_path.is_file():
-        # Convert single file
-        output_dir = os.path.dirname(args.output) if args.output else None
-        output_name = os.path.basename(args.output) if args.output else None
-        try:
-            converted_file = convert_pth_to_safetensors(str(input_path), output_dir, output_name)
-            print(f"Conversion completed: {converted_file}")
-        except Exception as e:
-            print(f"Error: {e}")
-            return 1
-    else:
-        # Convert directory
-        try:
-            converted_files = convert_directory(
-                str(input_path),
-                args.output,
-                args.recursive,
-                args.pattern
-            )
-            print(f"Converted {len(converted_files)} files")
-        except Exception as e:
-            print(f"Error: {e}")
-            return 1
-    return 0
-if __name__ == "__main__":
-    exit(main())