code_convert/convert_image_encoder.py

# import os
# from pathlib import Path
# from PIL import Image
# import onnx
# import onnx_graphsurgeon as gs
# import torch
# from onnx import shape_inference
# from packaging import version
# from polygraphy.backend.onnx.loader import fold_constants
# from torch.onnx import export

# from transformers import CLIPVisionModelWithProjection
# from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor

# is_torch_less_than_1_11 = version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")
# is_torch_2_0_1 = version.parse(version.parse(torch.__version__).base_version) == version.parse("2.0.1")

# class Optimizer:
#     def __init__(self, onnx_graph, verbose=False):
#         self.graph = gs.import_onnx(onnx_graph)
#         self.verbose = verbose

#     def info(self, prefix):
#         if self.verbose:
#             print(
#                 f"{prefix} .. {len(self.graph.nodes)} nodes, {len(self.graph.tensors().keys())} tensors, {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs"
#             )

#     def cleanup(self, return_onnx=False):
#         self.graph.cleanup().toposort()
#         if return_onnx:
#             return gs.export_onnx(self.graph)

#     def select_outputs(self, keep, names=None):
#         self.graph.outputs = [self.graph.outputs[o] for o in keep]
#         if names:
#             for i, name in enumerate(names):
#                 self.graph.outputs[i].name = name

#     def fold_constants(self, return_onnx=False):
#         onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)
#         self.graph = gs.import_onnx(onnx_graph)
#         if return_onnx:
#             return onnx_graph

#     def infer_shapes(self, return_onnx=False):
#         onnx_graph = gs.export_onnx(self.graph)
#         if onnx_graph.ByteSize() > 4147483648:
#             raise TypeError("ERROR: model size exceeds supported 2GB limit")
#         else:
#             onnx_graph = shape_inference.infer_shapes(onnx_graph)

#         self.graph = gs.import_onnx(onnx_graph)
#         if return_onnx:
#             return onnx_graph


# def optimize(onnx_graph, name, verbose):
#     opt = Optimizer(onnx_graph, verbose=verbose)
#     opt.info(name + ": original")
#     opt.cleanup()
#     opt.info(name + ": cleanup")
#     opt.fold_constants()
#     opt.info(name + ": fold constants")
#     # opt.infer_shapes()
#     # opt.info(name + ': shape inference')
#     onnx_opt_graph = opt.cleanup(return_onnx=True)
#     opt.info(name + ": finished")
#     return onnx_opt_graph


# class CLIPVisionProj(torch.nn.Module):
#     def __init__(self, clip_model) -> None:
#         super().__init__()
#         self.clip_model = clip_model

#     def forward(self, image_embedding):
#         result = self.clip_model(image_embedding,return_dict = False)
#         return result[0]
    
# def onnx_export(
#     model,
#     model_args: tuple,
#     output_path: Path,
#     ordered_input_names,
#     output_names,
#     dynamic_axes,
#     opset: int,
#     use_external_data_format=False,
#     verbose=False,  # Thêm tham số verbose
# ):
#     output_path.parent.mkdir(parents=True, exist_ok=True)
#     with torch.inference_mode(), torch.autocast("cuda"):
#         if is_torch_less_than_1_11:
#             export(
#                 model,
#                 model_args,
#                 f=output_path.as_posix(),
#                 input_names=ordered_input_names,
#                 output_names=output_names,
#                 dynamic_axes=dynamic_axes,
#                 do_constant_folding=True,
#                 use_external_data_format=use_external_data_format,
#                 enable_onnx_checker=True,
#                 opset_version=opset,
#                 verbose=verbose,  # Thêm verbose ở đây
#             )
#         else:
#             export(
#                 model,
#                 model_args,
#                 f=output_path.as_posix(),
#                 input_names=ordered_input_names,
#                 output_names=output_names,
#                 dynamic_axes=dynamic_axes,
#                 do_constant_folding=True,
#                 opset_version=opset,
#                 verbose=verbose,  # Thêm verbose ở đây
#             )

# def convert_models(
#     image_path:str, 
#     output_path:str,
#     opset:int=16,
# ):
#         dtype =  torch.float32
#         device = 'cpu'
#         image = Image.open(image_path)
#         image_encoder_processor = CLIPImageProcessor()
#         image_embedding = image_encoder_processor(image, return_tensors="pt").pixel_values
#         clip_model= CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter", subfolder = 'sdxl_models/image_encoder')
#         image_encoder = CLIPVisionProj(clip_model).to(device=device)
#         output_path = Path(output_path)

#         clip_path = output_path / "clip_vision_proj" / "model.onnx"
#         clip_optimize = output_path / 'clip_vision_proj' / 'optimize' / 'model.onnx'
#         #create folder for optimize clip
#         os.makedirs(output_path / 'optimize', exist_ok= True)
#         onnx_export(image_encoder,
#                     model_args= (image_embedding).to(dtype = torch.float32, device = device),
#                     output_path =clip_path,
#                     ordered_input_names= ['image_embedding'],
#                     output_names=["image_encoder"],
#                     dynamic_axes={'image_embedding': {0: 'Batch_size', 1: 'channel', 2: 'height', 3:'width'},
#                                  'image_encoder': {0:'Batch_size', 1: 'sequence_length'} },
#                     opset=opset,
#                     verbose=True,
#                     use_external_data_format=True, 
#                 )
#         clip_opt_graph =  onnx.load(clip_path)
#         onnx.save_model(
#             clip_opt_graph,
#             clip_optimize,  
#             save_as_external_data=True, 
#             all_tensors_to_one_file=True,  
#             convert_attribute=False,            
#             location="weights.pb",
#         )

# if __name__ == "__main__":
#     convert_models(image_path= '/home/SDXL_CNextAnimeCanny_IPAdapter_ONNX/code_inference/image_condition/control_canny_edge/condition_0.jpg',
#                       output_path='/home/new_onnx/image_encoder',
#                        opset=18,
#                      )


import onnx

def optimize_onnx_model(clip_path, output_model_path, weight_file="weights.pb"):
    # Load the ONNX model từ đường dẫn clip_path
    clip_opt_graph = onnx.load(clip_path)
    
    # Save optimized model, với weights được lưu riêng vào weight_file
    onnx.save_model(
        clip_opt_graph,
        output_model_path,  
        save_as_external_data=True,         # Lưu dữ liệu tensor lớn ra ngoài
        all_tensors_to_one_file=True,       # Gom tất cả tensor vào một file duy nhất
        location=weight_file,               # Đường dẫn đến file lưu weights
    )

# Sử dụng hàm
clip_path = "/home/new_onnx/image_encoder/clip_vision_proj/model.onnx"     # Đường dẫn model đã convert
output_model_path = "/home/new_onnx/image_encoder/optimize/model.onnx"            # Đường dẫn file model bạn muốn lưu
weight_file = "weights.pb"                  # Tên file chứa weights
optimize_onnx_model(clip_path, output_model_path, weight_file)