tumuyan2 commited on
Commit
0dc4c15
·
1 Parent(s): fafe4e2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +32 -17
app.py CHANGED
@@ -8,6 +8,16 @@ from ftplib import FTP
8
  from spandrel import ImageModelDescriptor, ModelLoader
9
  import torch
10
  import subprocess
 
 
 
 
 
 
 
 
 
 
11
 
12
  # 定义 downloaded_files 变量
13
  downloaded_files = {}
@@ -233,18 +243,24 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
233
 
234
  # 获取输出
235
  example_output = traced_torch_model(example_input)
236
- width_ratio = example_output.shape[2] / example_input.shape[2]
237
- print_log(task_id, input2, "获得缩放倍率="+ str(width_ratio)+", 输出shape="+str(list(example_output.shape)), "完成")
238
- log+= ("获得缩放倍率="+str(width_ratio)+", 输出shape="+str(list(example_output.shape))+"\n")
239
- yield [], log
 
 
 
 
 
240
 
 
241
 
242
  # 转换为 ONNX 模型
243
  if "ONNX" in output_type or "NCNN" in output_type or "MNN" in output_type:
244
- if str(width_ratio) in input2 or width_ratio >0:
245
  onnx_path = output_base + ".onnx"
246
  else:
247
- onnx_path = output_base + "-x" + str(int(width_ratio)) + ".onnx"
248
  if os.path.exists(onnx_path):
249
  print_log(task_id, input2, "转换为ONNX模型", "跳过")
250
  log += "跳过转换为ONNX模型\n"
@@ -259,10 +275,10 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
259
  # 转换为 mnn 模型
260
  if "MNN" in output_type:
261
 
262
- if str(width_ratio) in input2:
263
  mnn_path = output_base + ".mnn"
264
  else:
265
- mnn_path = output_base + "-x" + str(width_ratio) + ".mnn"
266
 
267
  mnn_config = ""
268
  if "Fixed" in output_type and input_tensor0 is not None:
@@ -300,17 +316,17 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
300
  yield [], log
301
  returncode = process.poll()
302
  if returncode != 0:
303
- print_log(task_id, input2, f"转换为MNN模型,返回码: {returncode},命令: {mnn_command} ", "失败")
304
  log += f"执行mnn命令失败,返回码: {returncode},命令: {mnn_command} \n"
305
  else:
306
- log += f"执行命令成功: {mnn_command} \n"
307
  except Exception as e:
308
- log += f"执行命令: {mnn_command} 失败,错误信息: {str(e)}\n"
309
  print_log(task_id, input2, f"转换为MNN模型,错误信息: {str(e)}", "错误")
310
 
311
  if "NCNN" in output_type:
312
- print_log(task_id, input2, "执行命令" + command, "开始")
313
- log += "执行命令…\n"
314
  yield [], log
315
 
316
  try:
@@ -328,12 +344,12 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
328
  yield [], log
329
  returncode = process.poll()
330
  if returncode != 0:
331
- log += f"执行命令失败,返回码: {returncode},命令: {command} \n"
332
  print_log(task_id, input2, f"返回码: {returncode},命令: {command} ", "失败")
333
  else:
334
- log += f"执行命令成功: {command} \n"
335
  except Exception as e:
336
- log += f"执行命令: {command} 失败,错误信息: {str(e)}\n"
337
  print_log(task_id, input2, f"错误信息: {str(e)}", "错误")
338
 
339
  # 查找 output_folder 目录下以 .ncnn.bin 和 .ncnn.param 结尾的文件
@@ -350,7 +366,6 @@ def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_
350
  # 压缩包内文件夹名称
351
  zip_folder_name = f"models-{input2}"
352
  # 重命名后的文件名
353
- scale = int(width_ratio)
354
  new_bin_name = f"x{scale}.bin"
355
  new_param_name = f"x{scale}.param"
356
  # 创建压缩包
 
8
  from spandrel import ImageModelDescriptor, ModelLoader
9
  import torch
10
  import subprocess
11
+ import spandrel
12
+ import ncnn
13
+ import MNN
14
+
15
+ print("Torch",torch.__version__)
16
+ print("Gradio",gr.__version__)
17
+ print("Spandrel", spandrel.__version__)
18
+ print("NCNN", ncnn.__version__)
19
+ print("MNN", MNN.__version__)
20
+
21
 
22
  # 定义 downloaded_files 变量
23
  downloaded_files = {}
 
243
 
244
  # 获取输出
245
  example_output = traced_torch_model(example_input)
246
+ if isinstance(example_output, torch.Tensor):
247
+ width_ratio = example_output.shape[2] / example_input.shape[2]
248
+ print_log(task_id, input2, "获得缩放倍率="+ str(width_ratio)+", 输出shape="+str(list(example_output.shape)), "完成")
249
+ log+= ("获得缩放倍率="+str(width_ratio)+", 输出shape="+str(list(example_output.shape))+"\n")
250
+ yield [], log
251
+ else:
252
+ print_log(task_id, input2, "Traced torch model输出" + type(example_output), "错误")
253
+ log+="Traced torch model输出" + type(example_output)+ "错误\n"
254
+ yield [], log
255
 
256
+ scale = int(width_ratio)
257
 
258
  # 转换为 ONNX 模型
259
  if "ONNX" in output_type or "NCNN" in output_type or "MNN" in output_type:
260
+ if str(scale) in input2 or scale <1:
261
  onnx_path = output_base + ".onnx"
262
  else:
263
+ onnx_path = output_base + "-x" + str(scale) + ".onnx"
264
  if os.path.exists(onnx_path):
265
  print_log(task_id, input2, "转换为ONNX模型", "跳过")
266
  log += "跳过转换为ONNX模型\n"
 
275
  # 转换为 mnn 模型
276
  if "MNN" in output_type:
277
 
278
+ if str(scale) in input2 or scale < 1:
279
  mnn_path = output_base + ".mnn"
280
  else:
281
+ mnn_path = output_base + "-x" + str(scale) + ".mnn"
282
 
283
  mnn_config = ""
284
  if "Fixed" in output_type and input_tensor0 is not None:
 
316
  yield [], log
317
  returncode = process.poll()
318
  if returncode != 0:
319
+ print_log(task_id, input2, f"转换为MNN模型,返回码: {returncode},命令: {mnn_command} ", "错误")
320
  log += f"执行mnn命令失败,返回码: {returncode},命令: {mnn_command} \n"
321
  else:
322
+ log += f"执行mnn命令成功: {mnn_command} \n"
323
  except Exception as e:
324
+ log += f"执行mnn命令: {mnn_command} 失败,错误信息: {str(e)}\n"
325
  print_log(task_id, input2, f"转换为MNN模型,错误信息: {str(e)}", "错误")
326
 
327
  if "NCNN" in output_type:
328
+ print_log(task_id, input2, "执行ncnn命令" + command, "开始")
329
+ log += "执行ncnn命令…\n"
330
  yield [], log
331
 
332
  try:
 
344
  yield [], log
345
  returncode = process.poll()
346
  if returncode != 0:
347
+ log += f"执行ncnn命令失败,返回码: {returncode},命令: {command} \n"
348
  print_log(task_id, input2, f"返回码: {returncode},命令: {command} ", "失败")
349
  else:
350
+ log += f"执行ncnn命令成功: {command} \n"
351
  except Exception as e:
352
+ log += f"执行ncnn命令: {command} 失败,错误信息: {str(e)}\n"
353
  print_log(task_id, input2, f"错误信息: {str(e)}", "错误")
354
 
355
  # 查找 output_folder 目录下以 .ncnn.bin 和 .ncnn.param 结尾的文件
 
366
  # 压缩包内文件夹名称
367
  zip_folder_name = f"models-{input2}"
368
  # 重命名后的文件名
 
369
  new_bin_name = f"x{scale}.bin"
370
  new_param_name = f"x{scale}.param"
371
  # 创建压缩包