Update app.py

app.py

@@ -31,17 +31,12 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
     log = "转换过程非常慢，请耐心等待。显示文件列表不代表转换完成。如果未发生错误，转换结束会显示”任务完成“\n"
     yield [], log
     if input2 == None or input2.strip() == "":
-        if isinstance(input1, str):
-            input2 = os.path.splitext(os.path.basename(input1))[0]
-        else:
-            input2 = os.path.splitext(os.path.basename(input1.name))[0]
-
+        input2 = os.path.splitext(os.path.basename(input1))[0]
         if input2 == "":
             input2 = str(task_id)
         log += f"未提供文件名，使用{input2}\n"
         print_log(task_id, input2,  f"未提供文件名，使用{input2}", "修正")
         yield [], log
-        input2 = "output"
 
     try:
         # 判断 input1 是地址还是文件，增加对 ftp 和 webdav 协议的支持
@@ -153,6 +148,7 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         log += "生成张量…\n"
         yield [], log
         pt_path = output_folder + "/" + input2 + ".pt"
+        onnx_path = output_folder + "/" + input2 + ".onnx"
         input_tensor0 = torch.rand(shape0) if any(shape0) else None
         input_tensor1 = torch.rand(shape1) if any(shape1) else None
         if input_tensor0 is not None and input_tensor1 is not None:
@@ -190,6 +186,22 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         torch_model = model.model
         print_log(task_id, input2, "获得模型对象", "完成")
         yield [], log
+
+
+        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 进行模型转换
+            # 将列表转换为元组
+            shape_tuple = tuple(shape0)
+            torch.onnx.export(torch_model, torch.rand(shape_tuple), onnx_path, verbose=True, input_names=["data"], output_names=["output"])
+
+
         if os.path.exists(pt_path):
             print_log(task_id, input2, "转换为TorchScript模型", "跳过")
             log += "跳过转换为TorchScript模型\n"
@@ -222,14 +234,14 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
                         yield [], log
             returncode = process.poll()
             if returncode != 0:
-                log += f"执行命令: {command} 失败，返回码: {returncode}\n"
+                log += f"执行命令失败，返回码: {returncode}，命令: {command} \n"
             else:
-                log += f"执行命令: {command} 成功\n"
+                log += f"执行命令成功: {command} \n"
         except Exception as e:
             log += f"执行命令: {command} 失败，错误信息: {str(e)}\n"
 
         output_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if os.path.isfile(os.path.join(output_folder, f))]
-        log += f"任务完成，输出文件: {output_files}\n"
+        log += f"任务完成\n"
         print_log(task_id, input2, "执行命令", "完成")
         yield output_files, log
     except Exception as e:
@@ -243,8 +255,11 @@ with gr.Blocks() as demo:
     with gr.Row():
         # 左侧列，包含输入组件和按钮
         with gr.Column():
+            # 添加文本提示
+            gr.Markdown("请输入的url，或者上传一个文件。限制文件为小于100M的*.pth模型")
             with gr.Row():
                 input1 = gr.Textbox(label="粘贴地址")
+                # 新增文件上传组件
                 input1_file = gr.File(label="上传文件")
             input2 = gr.Textbox(label="自定义文件名")
             # 修改为字符串输入控件
@@ -275,13 +290,13 @@ with gr.Blocks() as demo:
 
     # 添加范例
     examples = [
-        ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", "", "1,3,128,128", "0,0,0,0"],
-        ["https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth", "", "1,3,128,128", "0,0,0,0"],
-        ["https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth", "", "1,3,128,128", "0,0,0,0"]
+        ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", None, "", "1,3,128,128", "0,0,0,0"],
+        ["https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth", None, "", "1,3,128,128", "0,0,0,0"],
+        ["https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth", None, "", "1,3,128,128", "0,0,0,0"]
     ]
     gr.Examples(
         examples=examples,
-        inputs=[input1, input2, shape0_str, shape1_str],
+        inputs=[input1, input1_file, input2, shape0_str, shape1_str],
         outputs=[output, log_textbox],
         fn=start_process
     )