From b4941ea474a0817680e4cf7f96b1807dccfe44ff Mon Sep 17 00:00:00 2001 From: fapoverflow <50244958+fapoverflow@users.noreply.github.com> Date: Fri, 18 Jul 2025 17:06:53 +0200 Subject: [PATCH] Upgrade comfyui versions & add convert_2_safetensors service --- data/.gitignore | 1 + docker-compose.yml | 8 + services/comfy/Dockerfile | 4 +- services/convert2safetensors/Dockerfile | 26 +++ .../convert_2_safetensors.py | 160 ++++++++++++++++++ services/convert2safetensors/entrypoint.sh | 6 + services/convert2safetensors/requirements.txt | 5 + 7 files changed, 208 insertions(+), 2 deletions(-) create mode 100644 services/convert2safetensors/Dockerfile create mode 100644 services/convert2safetensors/convert_2_safetensors.py create mode 100644 services/convert2safetensors/entrypoint.sh create mode 100644 services/convert2safetensors/requirements.txt diff --git a/data/.gitignore b/data/.gitignore index f20063c..1241499 100644 --- a/data/.gitignore +++ b/data/.gitignore @@ -3,3 +3,4 @@ /embeddings /models /states +/convert diff --git a/docker-compose.yml b/docker-compose.yml index 5d94433..2fc68e3 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -81,3 +81,11 @@ services: - CLI_ARGS=--allow-code --xformers --enable-insecure-extension-access --api ports: - "${WEBUI_PORT:-7860}:7860" + + convertor: + <<: *base_service + profiles: [ "convert" ] + build: ./services/convert2safetensors + image: sd-convert + volumes: + - ./data/convert:/opt/convert2safetensors/models diff --git a/services/comfy/Dockerfile b/services/comfy/Dockerfile index 4684b5a..89ce958 100644 --- a/services/comfy/Dockerfile +++ b/services/comfy/Dockerfile @@ -1,6 +1,6 @@ # Defines the versions of ComfyUI, ComfyUI Manager, and PyTorch to use -ARG COMFYUI_VERSION=v0.3.43 -ARG COMFYUI_MANAGER_VERSION=3.33.3 +ARG COMFYUI_VERSION=v0.3.44 +ARG COMFYUI_MANAGER_VERSION=3.33.8 ARG PYTORCH_VERSION=2.7.1-cuda12.8-cudnn9-runtime # This image is based on the latest official PyTorch image, because it already contains CUDA, CuDNN, and PyTorch diff --git a/services/convert2safetensors/Dockerfile b/services/convert2safetensors/Dockerfile new file mode 100644 index 0000000..d6f144f --- /dev/null +++ b/services/convert2safetensors/Dockerfile @@ -0,0 +1,26 @@ +ARG PYTORCH_VERSION=2.7.1-cuda12.8-cudnn9-runtime +FROM pytorch/pytorch:${PYTORCH_VERSION} + +ENV DEBIAN_FRONTEND=noninteractive PIP_PREFER_BINARY=1 + +RUN apt update --assume-yes && \ +# apt install --assume-yes \ +# git \ +# sudo \ +# build-essential \ +# libgl1-mesa-glx \ +# libglib2.0-0 \ +# libsm6 \ +# libxext6 \ +# ffmpeg && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /opt/convert2safetensors +COPY . /opt/convert2safetensors/ +RUN chmod u+x /opt/convert2safetensors/entrypoint.sh + +RUN pip install --requirement /opt/convert2safetensors/requirements.txt + +ENV NVIDIA_VISIBLE_DEVICES=all PYTHONPATH="${PYTHONPATH}:${PWD}" CLI_ARGS="" +ENTRYPOINT ["/bin/bash", "/opt/convert2safetensors/entrypoint.sh"] diff --git a/services/convert2safetensors/convert_2_safetensors.py b/services/convert2safetensors/convert_2_safetensors.py new file mode 100644 index 0000000..086c9e0 --- /dev/null +++ b/services/convert2safetensors/convert_2_safetensors.py @@ -0,0 +1,160 @@ +import argparse +import json +import os +import shutil +import torch +from collections import defaultdict +from safetensors.torch import load_file, save_file +from tqdm import tqdm + + +def shared_pointers(tensors): + ptrs = defaultdict(list) + for k, v in tensors.items(): + ptrs[v.data_ptr()].append(k) + return [names for names in ptrs.values() if len(names) > 1] + + +def check_file_size(sf_filename, pt_filename): + sf_size = os.stat(sf_filename).st_size + pt_size = os.stat(pt_filename).st_size + if (sf_size - pt_size) / pt_size > 0.01: + raise RuntimeError(f"File size difference exceeds 1% between {sf_filename} and {pt_filename}") + + +def convert_file(pt_filename, sf_filename, copy_add_data=True, allow_pickle=False): + try: + loaded = torch.load(pt_filename, map_location="cpu", weights_only=not allow_pickle) + except Exception as e: + print(f"[WARNING] Failed to load '{pt_filename}': {e}") + return # Skip this file + loaded = loaded.get("state_dict", loaded) + shared = shared_pointers(loaded) + + for shared_weights in shared: + for name in shared_weights[1:]: + loaded.pop(name) + + loaded = {k: v.contiguous().half() for k, v in loaded.items()} + + source_folder = os.path.dirname(pt_filename) + dest_folder = os.path.dirname(sf_filename) + os.makedirs(dest_folder, exist_ok=True) + save_file(loaded, sf_filename, metadata={"format": "pt"}) + check_file_size(sf_filename, pt_filename) + if copy_add_data: + copy_additional_files(source_folder, dest_folder) + + reloaded = load_file(sf_filename) + for k, v in loaded.items(): + if not torch.equal(v, reloaded[k]): + raise RuntimeError(f"Mismatch in tensors for key {k}.") + + +def rename(pt_filename): + filename, ext = os.path.splitext(pt_filename) + local = f"{filename}.safetensors" + local = local.replace("pytorch_model", "model") + return local + + +def copy_additional_files(source_folder, dest_folder): + for file in os.listdir(source_folder): + file_path = os.path.join(source_folder, file) + if os.path.isfile(file_path) and not file.endswith(('.bin', '.py', '.pt', '.pth')): + shutil.copy(file_path, dest_folder) + + +def find_index_file(source_folder): + for file in os.listdir(source_folder): + if file.endswith('.bin.index.json'): + return file + return None + + +def convert_files(source_folder, dest_folder, delete_old, use_dest_folder, allow_pickle): + index_file = find_index_file(source_folder) + if index_file: + index_file = os.path.join(source_folder, index_file) + with open(index_file) as f: + index_data = json.load(f) + for pt_filename in tqdm(set(index_data["weight_map"].values()), desc="Converting files"): + full_pt_filename = os.path.join(source_folder, pt_filename) + sf_filename = os.path.join(dest_folder, rename(pt_filename)) if use_dest_folder else os.path.join( + os.path.dirname(full_pt_filename), rename(pt_filename)) + convert_file(full_pt_filename, sf_filename, copy_add_data=False, allow_pickle=allow_pickle) + if delete_old: + os.remove(full_pt_filename) + index_path = os.path.join(dest_folder, "model.safetensors.index.json") if use_dest_folder else os.path.join( + os.path.dirname(index_file), "model.safetensors.index.json") + with open(index_path, "w") as f: + new_map = {k: rename(v) for k, v in index_data["weight_map"].items()} + json.dump({**index_data, "weight_map": new_map}, f, indent=4) + else: + for pt_filename in tqdm(os.listdir(source_folder), desc="Converting files"): + if pt_filename.endswith(('.bin', '.pt', '.pth')): + full_pt_filename = os.path.join(source_folder, pt_filename) + sf_filename = os.path.join(dest_folder, rename(pt_filename)) if use_dest_folder else os.path.join( + os.path.dirname(full_pt_filename), rename(pt_filename)) + convert_file(full_pt_filename, sf_filename, copy_add_data=False, allow_pickle=allow_pickle) + if delete_old: + os.remove(full_pt_filename) + + copy_additional_files(source_folder, dest_folder) + + +def convert_batch(source_dir, dest_dir, delete_old, use_dest_folder, allow_pickle): + files_to_convert = [] + for root, _, files in os.walk(source_dir): + for file in files: + if file.endswith(('.bin', '.pt', '.pth')): + files_to_convert.append(os.path.join(root, file)) + for full_model_path in tqdm(files_to_convert, desc="Converting batch of models"): + sf_filename = os.path.join(dest_dir, os.path.relpath(rename(full_model_path), + source_dir)) if use_dest_folder else os.path.join( + os.path.dirname(full_model_path), rename(full_model_path)) + os.makedirs(os.path.dirname(sf_filename), exist_ok=True) + convert_file(full_model_path, sf_filename, copy_add_data=False, allow_pickle=allow_pickle) + if delete_old: + os.remove(full_model_path) + + +def main(): + parser = argparse.ArgumentParser(description="Convert PyTorch model files to SafeTensors format.") + parser.add_argument("--source", required=False, default=os.getcwd(), + help="Source folder containing .pth/.pt/.bin files") + parser.add_argument("--dest", required=False, + help="Destination folder for .safetensors files. If not specified, will use default.") + parser.add_argument("--delete-old", action="store_true", + help="Delete original PyTorch files after conversion") + parser.add_argument("--batch", action="store_true", + help="Recursively convert a batch of models from subdirectories") + parser.add_argument("--allow-pickle", action="store_true", + help="Allow loading legacy models with full object deserialization (trusted sources only!)") + + args = parser.parse_args() + + source_folder = os.path.abspath(args.source) + dest_folder = os.path.abspath(args.dest) if args.dest else os.path.join( + source_folder, os.path.basename(source_folder) + "_safetensors") + use_dest_folder = bool(args.dest) + + if args.batch: + convert_batch(source_folder, dest_folder, args.delete_old, use_dest_folder, args.allow_pickle) + else: + if any(fname.endswith(('.bin', '.pt', '.pth')) for fname in os.listdir(source_folder)): + if "pytorch_model.bin" in os.listdir(source_folder): + sf_filename = os.path.join(dest_folder, "model.safetensors") if use_dest_folder else os.path.join( + source_folder, "model.safetensors") + convert_file(os.path.join(source_folder, "pytorch_model.bin"), sf_filename, + copy_add_data=True, allow_pickle=args.allow_pickle) + if args.delete_old: + os.remove(os.path.join(source_folder, "pytorch_model.bin")) + else: + convert_files(source_folder, dest_folder, args.delete_old, use_dest_folder, args.allow_pickle) + else: + raise RuntimeError("No valid model files found in the specified directory.") + + +if __name__ == "__main__": + main() diff --git a/services/convert2safetensors/entrypoint.sh b/services/convert2safetensors/entrypoint.sh new file mode 100644 index 0000000..3500eb5 --- /dev/null +++ b/services/convert2safetensors/entrypoint.sh @@ -0,0 +1,6 @@ +#!/bin/bash +set -ex + +echo "Running container as $USER" +echo "Running container as $USER..." +python convert_2_safetensors.py --source ./models/input --dest ./models/output --allow-pickle diff --git a/services/convert2safetensors/requirements.txt b/services/convert2safetensors/requirements.txt new file mode 100644 index 0000000..0ed85e7 --- /dev/null +++ b/services/convert2safetensors/requirements.txt @@ -0,0 +1,5 @@ +safetensors +argparse +torch +tqdm +numpy