update

app.py

@@ -197,7 +197,8 @@ def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min:
             os.remove(save_path)
         return None
 
-async def _process_model(model_input: Union[str, gr.File], tilesize: int, output_dir: str,task_id:int,fp16:bool,onnxsim:bool,opset:int):
+async def _process_model(model_input: Union[str, gr.File], tilesize: int, output_dir: str,task_id:int,fp16:bool,onnxsim:bool,opset:int,dynamic_axes:bool):
+
 
     log = ('初始化日志记录...\n')
     print_log(task_id, '初始化日志记录', '开始')
@@ -226,7 +227,7 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
         yield [],log
     else:
         print_log(task_id, f'转换 PTH 模型为 ONNX, folder={output_dir}', '开始')
-        onnx_path = convert_pth_to_onnx(input_path, tilesize=tilesize, output_folder=output_dir,use_fp16=fp16, simplify_model=onnxsim, opset=opset)
+        onnx_path = convert_pth_to_onnx(input_path, tilesize=tilesize, output_folder=output_dir,use_fp16=fp16, simplify_model=onnxsim, opset=opset, dynamic_axes=dynamic_axes)
         if onnx_path:
             log += ( f'成功生成ONNX模型: {onnx_path}\n')
             print_log(task_id, f'生成ONNX模型: {onnx_path}', '完成')
@@ -277,12 +278,11 @@ with gr.Blocks() as demo:
                     return gr.update(visible=False), gr.update(visible=True)
             
             input_type.change(show_input, inputs=input_type, outputs=[url_input, file_input])
-            
-            tilesize = gr.Number(label="Tilesize", value=0, precision=0)
-            # 添加fp16和try_run复选框
+            tilesize = gr.Number(label="Dummy input width/height, default 64", value=64, precision=0)
+            opset = gr.Number(label="ONNX export opset version, suggest 9/11/13/16/17/18", value=13, precision=0)
             fp16 = gr.Checkbox(label="FP16", value=False)
             onnxsim = gr.Checkbox(label="ONNX export simplify model", value=False)
-            opset = gr.Number(label="ONNX export opset version, suggest 9/11/13/16/17/18", value=13, precision=0)
+            dynamic_axes = gr.Checkbox(label="ONNX input apply dynamic axes", value=True)
             try_run = gr.Checkbox(label="MNNSR test", value=False)
             convert_btn = gr.Button("Run")
         with gr.Column():
@@ -298,7 +298,8 @@ with gr.Blocks() as demo:
                     return gr.update(visible=False)
             try_run.change(show_try_run, inputs=try_run, outputs=img_output)
 
-    async def process_model(input_type, url_input, file_input, tilesize, fp16, onnxsim, opset, try_run):
+    async def process_model(input_type, url_input, file_input, tilesize, fp16, onnxsim, opset, dynamic_axes, try_run):
+
 
         global task_counter
         task_counter += 1
@@ -347,7 +348,8 @@ with gr.Blocks() as demo:
         onnx_path = None
         mnn_path = None
         # 调用重命名后的函数
-        async for result in _process_model(model_input, tilesize if tilesize>0 else 64, output_dir, task_counter, fp16, onnxsim, opset):
+        async for result in _process_model(model_input, tilesize if tilesize>0 else 64, output_dir, task_counter, fp16, onnxsim, opset, dynamic_axes):
+
             if isinstance(result, tuple) and len(result) == 3:
                 onnx_path, mnn_path, process_log = result
                 yield onnx_path, mnn_path, log+process_log, None
@@ -360,7 +362,7 @@ with gr.Blocks() as demo:
             if mnn_path: 
                 if try_run:
                     print_log(task_counter, f'测试模型: {mnn_path}', '开始')
-                    processed_image_np, load_time, infer_time = modelTest_for_gradio(mnn_path, "./sample.jpg", int(tilesize), 0)
+                    processed_image_np, load_time, infer_time = modelTest_for_gradio(mnn_path, "./sample.jpg", tilesize if tilesize>0 and dynamic_axes else 0, 0)
                     processed_image_pil = Image.fromarray(cv2.cvtColor(processed_image_np, cv2.COLOR_BGR2RGB))
                     # processed_image_pil = Image.fromarray(processed_image_np)
                     yield onnx_path, mnn_path, log+process_log+f"MNNSR 加载模型用时 {load_time:.4f} 秒, 推理({tilesize} px)用时 {infer_time:.4f} 秒", processed_image_pil
@@ -370,7 +372,7 @@ with gr.Blocks() as demo:
 
     convert_btn.click(
         process_model,
-        inputs=[input_type, url_input, file_input, tilesize, fp16, onnxsim, opset, try_run],
+        inputs=[input_type, url_input, file_input, tilesize, fp16, onnxsim, opset, dynamic_axes, try_run],
         outputs=[onnx_output, mnn_output, log_box, img_output],
         api_name="convert_nmm_model"
     )

pth2onnx.py

@@ -6,7 +6,7 @@ import onnx
 from spandrel import ImageModelDescriptor, ModelLoader
 from onnxsim import simplify
 
-def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tilesize: int = 64, use_fp16: bool=False, simplify_model: bool=False, min_size: int = 1024*1024, output_folder: str=None, opset: int = 11):
+def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tilesize: int = 64, use_fp16: bool=False, simplify_model: bool=False, min_size: int = 1024*1024, output_folder: str=None, opset: int = 11, dynamic_axes: bool = True):
     """
     Loads a PyTorch model from a .pth file using Spandrel and converts it to ONNX format.
 
@@ -86,6 +86,14 @@ def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tiles
     print(f"ONNX model exporting...")
     try:
         # Export the model
+        if dynamic_axes:
+            axes = {
+                "input": {2: "height", 3: "width"}, 
+                "output": {2: "height", 3: "width"},
+            }
+        else:
+            axes = {}
+
         torch.onnx.export(
             torch_model,                 # The model instance
             example_input,               # An example input tensor
@@ -95,10 +103,7 @@ def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tiles
             do_constant_folding=True,    # Whether to execute constant folding for optimization
             input_names=['input'],       # The model's input names
             output_names=['output'],     # The model's output names
-            dynamic_axes={               # Allow variable input/output dimensions
-                "input": {0: "batch_size", 2: "height", 3: "width"}, # Batch, H, W can vary
-                "output": {0: "batch_size", 2: "height", 3: "width"},# Batch, H, W can vary
-            }
+            dynamic_axes=axes
         )
         print(f"ONNX model export successful: {onnx_path}")
 
@@ -133,16 +138,18 @@ if __name__ == "__main__":
     parser.add_argument('--fp16', action='store_true', help='Use FP16 precision.')
     parser.add_argument('--simplify', action='store_true', help='Simplify the ONNX model.')
     parser.add_argument('--opset', type=int, default=11, help='ONNX opset version.')
+    parser.add_argument('--fixed_axes', action='store_true', help='Use dynamic axes.')
     args = parser.parse_args()
 
     success = convert_pth_to_onnx(
-        pth_path=args.pth_path,
-        onnx_path=args.onnx_path,
+        pth_path=args.pthpath,
+        onnx_path=args.onnxpath,
         channel=args.channel,
         tilesize=args.tilesize,
-        use_fp16=args.use_fp16,
-        simplify_model=args.simplify_model,
+        use_fp16=args.fp16,
+        simplify_model=args.simplify,
         opset=args.opset,
+        dynamic_axes= not args.fixed_axes,
     )
 
     if success: