Update app.py

app.py

@@ -24,7 +24,7 @@ def print_log(task_id, filename, stage, status):
         print(f"任务{task_id}: {filename}, [{status}] {stage}")
 
 # 修改 start_process 函数，处理新增输入
-def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_type="+ Dynamic MNN & NCNN", input_suffix=".pth"):
+def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_type=["TorchScript", "ONNX", "MNN", "ONNX"], input_suffix=".pth"):
     global task_counter
     task_counter += 1
     task_id = task_counter
@@ -221,7 +221,7 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
             print_log(task_id, input2, "转换为TorchScript模型", "跳过")
             log += "跳过转换为TorchScript模型\n"
             yield [], log
-        else:
+        elif "TorchScript" in output_type:
             print_log(task_id, input2, "转换为TorchScript模型", "开始")
             log+= "转换为TorchScript模型…\n"
             yield [], log
@@ -239,19 +239,20 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
 
 
         # 转换为 ONNX 模型
-        if str(width_ratio) in input2:
-            onnx_path = output_base + ".onnx"
-        else:
-            onnx_path = output_base + "-x" + str(width_ratio) + ".onnx"
-        if os.path.exists(onnx_path):
-            print_log(task_id, input2, "转换为ONNX模型", "跳过")
-            log += "跳过转换为ONNX模型\n"
-            yield [], log
-        else:
-            print_log(task_id, input2, "转换为ONNX模型", "开始")
-            log += "转换为ONNX模型…\n"
-            yield [], log
-            torch.onnx.export(torch_model, example_input, onnx_path, opset_version=17, input_names=["input"], output_names=["output"])
+        if "ONNX" in output_type or "NCNN" in output_type or "MNN" in output_type:
+            if str(width_ratio) in input2:
+                onnx_path = output_base + ".onnx"
+            else:
+                onnx_path = output_base + "-x" + str(width_ratio) + ".onnx"
+            if os.path.exists(onnx_path):
+                print_log(task_id, input2, "转换为ONNX模型", "跳过")
+                log += "跳过转换为ONNX模型\n"
+                yield [], log
+            else:
+                print_log(task_id, input2, "转换为ONNX模型", "开始")
+                log += "转换为ONNX模型…\n"
+                yield [], log
+                torch.onnx.export(torch_model, example_input, onnx_path, opset_version=17, input_names=["input"], output_names=["output"])
 
 
         # 转换为 mnn 模型
@@ -390,8 +391,8 @@ with gr.Blocks() as demo:
             
             with gr.Row():
                 input2 = gr.Textbox(label="自定义文件名")
-                output_type = gr.Dropdown(choices=["TorchScript & ONNX", "+ Fixed MNN & NCNN", "+ Dynamic MNN & NCNN", "Fixed NCNN", "Dynamic NCNN", "Fixed MNN", "Dynamic MNN" ], value="+ Dynamic MNN & NCNN", label="模型类型")
-            # 修改为字符串输入控件
+                output_type = gr.Dropdown(  ["TorchScript", "ONNX", "Fixed", "MNN", "NCNN"], value=["TorchScript", "ONNX", "MNN", "ONNX"], multiselect=True, label="模型类型", info="1. 生成mnn和ncnn模型必须先生成onnx模型；2.如果选项中包含了Fixed，那么输出的onnx和mnn模型都使用固定shape的input。"
+        ),
             shape0_str = gr.Textbox(label="shape0 (逗号分隔的整数)", value="1,3,128,128")
             shape1_str = gr.Textbox(label="shape1 (逗号分隔的整数)", value="0,0,0,0")
             with gr.Row():
@@ -424,11 +425,11 @@ with gr.Blocks() as demo:
         ["","https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth", "", "1,3,128,128", "0,0,0,0"],
         ["","https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-WTP-ColorDS.pth", "", "1,3,128,128", "0,0,0,0"],
     ]
-    # gr.Examples(
-    #     examples=examples,
-    #     inputs=[input1_file, input1, input2, shape0_str, shape1_str],
-    #     outputs=[output, log_textbox],
-    #     fn=start_process
-    # )
+    gr.Examples(
+        examples=examples,
+        inputs=[input1_file, input1, input2, shape0_str, shape1_str],
+        outputs=[output, log_textbox],
+        fn=start_process
+    )
 
 demo.launch()