[add] merge two onnx and related

Author: xiaominghan121

README_zh-CN.md

@@ -182,6 +182,19 @@
 
 `onnx-modifer`正在活跃地更新中:hammer_and_wrench:。 欢迎使用，提issue，如果有帮助的话，感谢给个:star:~
 
+## 合并两个onnx模型
+
+将两个模型文件拖动到页面，或在模型选择页面选择两个模型。等待加载...完成。
+修改完成后点击`download`下载。
+
+> **已知问题:** 
+> 1. 只能复制已知节点，未知节点需要先修改`metadata.json`,即先变为已知节点。
+> 2. 调试模式下按下Alt，进入断点并松开Alt，退出调试后，Alt会被认为一直是按下状态。需要重新按一次解除。
+
+## 复制或删除多个节点
+
+按住`Alt`键，点击第一个节点，`Alt`不松手点击第二个节点，完成选择。按下`j`键并松开，实现复制。按下`l`键并松开，实现删除。松开`Alt`取消选中状态。
+
 # 示例模型文件
 
 为方便测试，以下提供一些典型的样例模型文件，主要来自于[onnx model zoo](https://github.com/onnx/models)

app.py

@@ -1,6 +1,7 @@
 import argparse
 import logging
-from flask import Flask, render_template, request
+import time
+from flask import Flask, render_template, request, send_file
 from onnx_modifier import onnxModifier
 logging.basicConfig(level=logging.INFO)
 
@@ -24,6 +25,23 @@ def open_model():
 
     return 'OK', 200
 
+@app.route('/merge_model', methods=['POST'])
+def merge_model():
+
+    onnx_file1 = request.files['file0']
+    onnx_file2 = request.files['file1']
+    timestamp = time.time()
+    global onnx_modifier
+    onnx_modifier, stream ,merged_name = onnxModifier.merge(
+        onnx_file1.filename, onnx_file1.stream, 
+        onnx_file2.filename, onnx_file2.stream, 
+        "", str(int(timestamp)) + "_")
+    
+    return send_file(stream,
+                     mimetype='application/octet-stream',
+                     as_attachment=True,
+                     download_name=merged_name)
+
 
 @app.route('/download', methods=['POST'])
 def modify_and_download_model():

onnx_modifier.py

@@ -3,6 +3,7 @@
 # https://github.com/saurabh-shandilya/onnx-utils
 # https://stackoverflow.com/questions/52402448/how-to-read-individual-layers-weight-bias-values-from-onnx-model
 
+from io import BytesIO
 import os
 import copy
 import struct
@@ -39,13 +40,36 @@ def from_name_json_stream(cls, name, stream):
         return cls(name, model_proto)
 
     @classmethod
-    def from_name_protobuf_stream(cls, name, stream):
+    def from_name_protobuf_stream(cls, name, stream, prefix = ""):
         # https://leimao.github.io/blog/ONNX-IO-Stream/
         logging.info("loading model...")
         stream.seek(0)
         model_proto = onnx.load_model(stream, "protobuf", load_external_data=False)
+        model_proto = onnx.compose.add_prefix(model_proto, prefix=prefix)
         logging.info("load done!")
         return cls(name, model_proto)
+    
+    @classmethod
+    def merge(cls, name1, stream1, name2, stream2, prefix1, prefix2):
+        stream1.seek(0)
+        model_proto1 = onnx.load_model(stream1, "protobuf", load_external_data=False)
+        model_proto1 = onnx.compose.add_prefix(model_proto1, prefix=prefix1)
+        
+        stream2.seek(0)
+        model_proto2 = onnx.load_model(stream2, "protobuf", load_external_data=False)
+        model_proto2 = onnx.compose.add_prefix(model_proto2, prefix=prefix2)
+        
+        model_proto1.graph.input.extend(model_proto2.graph.input)
+        model_proto1.graph.node.extend(model_proto2.graph.node)
+        model_proto1.graph.initializer.extend(model_proto2.graph.initializer)
+        model_proto1.graph.output.extend(model_proto2.graph.output)
+        
+        merged_name = name1.split('.')[0] + "_" + name2.split('.')[0] + ".onnx"
+        byte_stream = BytesIO()
+        onnx.save_model(model_proto1, byte_stream)
+        byte_stream.seek(0)
+        return cls(merged_name, model_proto1), byte_stream, merged_name
+
 
     def reload(self):
         self.model_proto = copy.deepcopy(self.model_proto_backup)
@@ -460,6 +484,7 @@ def modify(self, modify_info):
         logging.debug("=== modify_info ===\n", modify_info)
 
         self.add_nodes(modify_info['added_node_info'], modify_info['node_states'])
+        self.change_initializer(modify_info['added_tensor'])
         self.change_initializer(modify_info['changed_initializer'])
         self.change_node_io_name(modify_info['node_renamed_io'])
         self.edit_inputs(modify_info['added_inputs'], modify_info['rebatch_info'])

static/index.js

@@ -116,8 +116,80 @@ host.BrowserHost = class {
         });
     }
 
+    // open the the model at flask server and reload
+    remote_open_model(obj_files) {
+        var files = Array.from(obj_files);
+        const acceptd_file = files.filter(file => this._view.accept(file.name)).slice(0, 2);
+        // console.log(file)
+        if(acceptd_file.length == 1)
+        {   
+            var file = acceptd_file[0];
+            this.upload_filename = file.name;
+            var form = new FormData();
+            form.append('file', file);
+            
+            // https://stackoverflow.com/questions/66039996/javascript-fetch-upload-files-to-python-flask-restful
+            fetch('/open_model', {
+                method: 'POST',
+                body: form
+            }).then(function (response) {
+                return response.text();
+            }).then(function (text) {
+                console.log('POST response: ');
+                // Should be 'OK' if everything was successful
+                console.log(text);
+            });
 
 
+            if (file) {
+                this._open(file, files);
+                this._view.modifier.clearGraph();
+            }
+        } else if (acceptd_file.length == 2)
+        {
+            var form = new FormData();
+            for(var i = 0; i < acceptd_file.length; i++)
+            {
+                form.append('file'+i, acceptd_file[i]);
+            }
+            // console.log(file)
+            // this.upload_filename = file.name;
+            let filename = 'unknown';
+            // https://stackoverflow.com/questions/66039996/javascript-fetch-upload-files-to-python-flask-restful
+            this._view.showLoading();
+            fetch('/merge_model', {
+                method: 'POST',
+                body: form
+            }).then(function (response) {
+                
+                const contentDisposition = response.headers.get('content-disposition');
+                if (contentDisposition) {
+                    const matches = contentDisposition.match(/filename[^;=\n]*=((['"]).*?\2|[^;\n]*)/i);
+                    if (matches && matches[1]) {
+                        filename = decodeURIComponent(matches[1].replace(/['"]/g, ''));
+                    }
+                }
+                if (!response.ok) {
+                    throw new Error(`HTTP error! status: ${response.status}`);
+                }
+                var blob = response.blob();
+                return blob;
+            }).then(blob => {
+                var file = new File([blob], filename);
+                // console.log('POST response: ');
+                // // Should be 'OK' if everything was successful
+                // console.log(text);
+                if (file) {
+                    this.upload_filename = file.name;
+                    files = [];
+                    files.push(file);
+                    this._open(file, files);
+                    this._view.modifier.clearGraph();
+                }
+            });
+        } 
+    }
+
     start() {
         this.window.addEventListener('error', (e) => {
             this.exception(e.error, true);
@@ -328,33 +400,14 @@ host.BrowserHost = class {
             });
             openFileDialog.addEventListener('change', (e) => {
                 if (e.target && e.target.files && e.target.files.length > 0) {
-                    const files = Array.from(e.target.files);
-                    const file = files.find((file) => this._view.accept(file.name));
-                    // console.log(file)
-                    this.upload_filename = file.name;
-                    var form = new FormData();
-                    form.append('file', file);
-
-                    // https://stackoverflow.com/questions/66039996/javascript-fetch-upload-files-to-python-flask-restful
-                    fetch('/open_model', {
-                        method: 'POST',
-                        body: form
-                    }).then(function (response) {
-                        return response.text();
-                    }).then(function (text) {
-                        console.log('POST response: ');
-                        // Should be 'OK' if everything was successful
-                        console.log(text);
-                    });
-
-
-                    if (file) {
-                        this._open(file, files);
-                        this._view.modifier.clearGraph();
-                    }
+                    this.remote_open_model(e.target.files);
                 }
             });
+        
         }
+
+
+
         const openModelButton = this.document.getElementById('load-model');
         if (openModelButton && openFileDialog) {
             openModelButton.addEventListener('click', () => {
@@ -379,27 +432,8 @@ host.BrowserHost = class {
         this.document.body.addEventListener('drop', (e) => {
             e.preventDefault();
             if (e.dataTransfer && e.dataTransfer.files && e.dataTransfer.files.length > 0) {
-                const files = Array.from(e.dataTransfer.files);
-                const file = files.find((file) => this._view.accept(file.name));
-                this.upload_filename = file.name;
-                var form = new FormData();
-                form.append('file', file);
-
-                // https://stackoverflow.com/questions/66039996/javascript-fetch-upload-files-to-python-flask-restful
-                fetch('/open_model', {
-                    method: 'POST',
-                    body: form
-                }).then(function (response) {
-                    return response.text();
-                }).then(function (text) {
-                    console.log('POST response: ');
-                    // Should be 'OK' if everything was successful
-                    console.log(text);
-                });
-                if (file) {
-                    this._open(file, files);
-                    this._view.modifier.clearGraph();
-                }
+                this.remote_open_model(e.dataTransfer.files);
+
             }
         });
 
@@ -588,6 +622,7 @@ host.BrowserHost = class {
             // 'modified_inputs_info' : this.arrayToObject(this.process_modified_inputs(this._view.modifier.inputModificationForSave,
             //     this._view.modifier.renameMap, this._view.modifier.name2NodeStates)),
             'rebatch_info' : this.mapToObjectRec(this._view.modifier.reBatchInfo),
+            'added_tensor' : this.mapToObjectRec(this._view.modifier.addedTensor),
             'changed_initializer' : this.mapToObjectRec(this._view.modifier.initializerEditInfo),
             'postprocess_args' : {'shapeInf' : this._view.modifier.downloadWithShapeInf, 'cleanUp' : this._view.modifier.downloadWithCleanUp}
         })

static/modifier.js

@@ -23,6 +23,8 @@ modifier.Modifier = class {
         this.downloadWithShapeInf = false;
         this.downloadWithCleanUp = false;
 
+        this.addedTensor = new Map();
+
     }
 
     loadModelGraph(model, graphs) {
@@ -79,13 +81,13 @@ modifier.Modifier = class {
     }
 
 
-    try_get_node_name(op_type)
+    try_get_node_name(op_type, input_node_id)
     {
-        var node_id = (this.addNodeKey++).toString();  // in case input (onnx) node has no name
+        var node_id = (input_node_id || this.addNodeKey++).toString();  // in case input (onnx) node has no name
         var modelNodeName = 'custom_added_' + op_type + node_id;
 
         if (this.addedNode.has(modelNodeName) || this.name2NodeStates.get(modelNodeName) ){
-            modelNodeName = this.randomString(16, 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ');
+            modelNodeName = try_get_node_name(op_type, Date.parse(new Date()));//this.randomString(16, 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ');
         }
         return modelNodeName;
     }
@@ -103,6 +105,69 @@ modifier.Modifier = class {
         this.applyAndUpdateView();
     }
 
+    // duplicate a node with ( _cp + unique_id ) as param suffix
+    duplicateNode(node_name, unique_id = "") {
+        //avoid to add a existed name node
+        var srcModelNode = this.name2ModelNode.get(node_name);
+        if (!srcModelNode.type){
+            return;
+        }
+        var dstModelNodeName = this.try_get_node_name(srcModelNode.type.name);
+
+        var properties = new Map();
+        properties.set('domain', "ai.onnx");
+        properties.set('op_type', srcModelNode.type.name);
+        properties.set('name', dstModelNodeName);
+
+
+        var attributes = new Map();
+        for (const attribute of srcModelNode.attributes) {
+            attributes.set(attribute.name, [
+                attribute.value?attribute.value.toString():"undefined", 
+                attribute.type||"undefined"])
+            // attributes.set(key, modelNode.attributes.get(key));
+        }
+
+        var outputs = new Map();
+        for (const output of srcModelNode.outputs) {
+            var dstNameList = [];
+            for (const srcArg of output.arguments) {
+                var dstName = srcArg.name + "_cp" + unique_id;
+                dstNameList.push([dstName, false]);
+                this.graph.copy_tensor(dstName, srcArg.name);
+
+            }
+            outputs.set(output.name, dstNameList);
+           
+        }
+
+        var inputs = new Map();
+        for (const input of srcModelNode.inputs) {
+            var dstNameList = [];
+            for (const srcArg of input.arguments) {
+                var dstName = srcArg.name + "_cp" + unique_id;
+                
+                if (this.graph._context._tensors && 
+                    this.graph._context._tensors.has(srcArg.name)) {
+                    this.graph.copy_tensor(dstName, srcArg.name);
+                    var initializer_info = this.graph.get_initializer_info(dstName);
+                    if(initializer_info) {
+                        this.addedTensor.set(dstName, initializer_info);
+                    }
+                    
+                }
+                dstNameList.push([dstName, false])
+            }  
+            inputs.set(input.name, dstNameList)
+         
+        }
+
+        this.addedNode.set(dstModelNodeName, 
+            new view.LightNodeInfo(properties, attributes, inputs, outputs));
+        this.applyAndUpdateView();
+    }
+
+
     addModelOutput(node_name) {
         var modelNode = this.name2ModelNode.get(node_name);
         // use a output argument as a proxy
@@ -402,6 +467,10 @@ modifier.Modifier = class {
         for (const [modelNodeName, node_info] of this.addedNode) {
             // console.log(node_info)
             var node = this.graph.make_custom_added_node(node_info);
+            if (!node) {
+                console.log("node not supported yet");
+                continue;
+            }
             // console.log(node)
 
             for (const input of node.inputs) {
@@ -489,6 +558,7 @@ modifier.Modifier = class {
             this.graph.reset_custom_added_node();
             this.graph.reset_custom_modified_outputs();
             this.graph.reset_custom_modified_inputs();
+            this.graph.reset_custom_added_tensors();
         }
         // reset load location
         var container = this.view._getElementById('graph');