update

app.py

@@ -196,11 +196,15 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             print_log(task_id, input2, "转换为ONNX模型", "开始")
             log += "转换为ONNX模型…\n"
             yield [], log
-            # 使用 torch.onnx.export 进行模型转换
-            # 将列表转换为元组
-            shape_tuple = tuple(shape0)
-            torch.onnx.export(torch_model, torch.rand(shape_tuple), onnx_path, verbose=True, input_names=["data"], output_names=["output"])
-
+            try:
+                # 使用 torch.onnx.export 进行模型转换
+                # 将列表转换为元组
+                shape_tuple = tuple(shape0)
+                torch.onnx.export(torch_model, torch.rand(shape_tuple), onnx_path, verbose=True, input_names=["data"], output_names=["output"])
+            except Exception as e:
+                print_log(task_id, input2, "转换为ONNX模型"+e, f"失败")
+                log += f"转换为ONNX模型失败: {e}\n"
+                yield [], log
 
         if os.path.exists(pt_path):
             print_log(task_id, input2, "转换为TorchScript模型", "跳过")

requirements.txt

@@ -1,3 +1,4 @@
 spandrel
 torch
 pnnx
+onnx