Update app.py

app.py

@@ -38,15 +38,12 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
                 print_log(task_id, input2,  "不支持此文件的格式 suffix="+suffix, "错误")
                 log += f"不支持此文件的格式\n"
                 return [] , log
+            input2 = split_input[0]
         print_log(task_id, input2, "检查文件名", "开始")
         log += f"检查文件名…\n"
         yield [], log
         if input2 == None or input2.strip() == "":
             input2 = str(task_id)
-
-        input1 = split_input[0]
-        if input2 == "":
-            input2 = str(task_id)
         log += f"未提供文件名，使用{input2}\n"
         print_log(task_id, input2,  f"未提供文件名，使用{input2}", "修正")
         yield [], log
@@ -161,7 +158,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"
+        # 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:
@@ -200,6 +197,8 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         print_log(task_id, input2, "获得模型对象", "完成")
         yield [], log
 
+
+        width_ratio = 4
         if os.path.exists(pt_path):
             print_log(task_id, input2, "转换为TorchScript模型", "跳过")
             log += "跳过转换为TorchScript模型\n"
@@ -213,6 +212,13 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             traced_torch_model.save(output_folder + "/" + input2 + ".pt")
             print_log(task_id, input2, "转换为TorchScript模型", "完成")
 
+            # 获取输出
+            example_output = traced_torch_model(example_input)
+            width_ratio = example_output.shape[2] / example_input.shape[2]
+            print_log(task_id, input2, "获得缩放倍率="+width_ratio+", 输出shape="+example_output.shape, "完成")
+            log+= "获得缩放倍率="+width_ratio+", 输出shape="+example_output.shape+"\n"
+            yield [], log
+
         print_log(task_id, input2, "执行命令" + command, "开始")
         log += "执行命令…\n"
         yield [], log
@@ -237,6 +243,28 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
                 log += f"执行命令成功: {command} \n"
         except Exception as e:
             log += f"执行命令: {command} 失败，错误信息: {str(e)}\n"
+            
+        # 检查文件是否存在
+        bin_file = os.path.join(output_folder, input2 + ".ncnn.bin")
+        param_file = os.path.join(output_folder, input2 + ".ncnn.param")
+        if os.path.exists(bin_file) and os.path.exists(param_file):
+            import zipfile
+            # 压缩包名称
+            zip_file_name = os.path.join(output_folder, f"models-{input2}.zip")
+            # 压缩包内文件夹名称
+            zip_folder_name = f"models-{input2}"
+            # 重命名后的文件名
+            new_bin_name = f"x{width_ratio}.bin"
+            new_param_name = f"x{width_ratio}.param"
+            # 创建压缩包
+            with zipfile.ZipFile(zip_file_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                # 写入重命名后的.bin文件
+                zipf.write(bin_file, os.path.join(zip_folder_name, new_bin_name))
+                # 写入重命名后的.param文件
+                zipf.write(param_file, os.path.join(zip_folder_name, new_param_name))
+            log += f"已创建压缩包: {zip_file_name}\n"
+            print_log(task_id, input2, "创建压缩包"+zip_file_name, "完成")
+            yield [], log
 
         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"任务完成\n"