update App

.gitignore

@@ -0,0 +1,5 @@
+*.pt
+*.pth
+*.bin
+*.onnx
+*.param

app.py

@@ -8,89 +8,40 @@ from spandrel import ImageModelDescriptor, ModelLoader
 import torch
 import subprocess
 
+# 定义 downloaded_files 变量
+downloaded_files = {}
 
 # 新增日志开关
 log_to_terminal = True
 
+# 新增全局任务计数器
+task_counter = 0
+
 # 新增日志函数
 def print_log(task_id, filename, stage, status):
     if log_to_terminal:
         print(f"任务 ID: {task_id}, 文件名: {filename}, 状态: [{status}], 阶段: {stage}")
 
-# 修改 process_file 函数，接收 shape 数组
-def process_file(task_id, file_path, model_name, output_folder, shape0, shape1):
-    # 使用 torch.rand 生成 input_shape
-    log = ""
-    print_log(task_id, model_name, "生成输入张量", "开始")
-    pt_path = output_folder + "/" + model_name + ".pt"
-    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:
-        example_input = (input_tensor0, input_tensor1)
-        # 修改此处，去除 shape 字符串中的空格
-        command = f"pnnx {pt_path} inputshape={str(shape0).replace(' ', '')} inputshape2={str(shape1).replace(' ', '')}"
-    elif input_tensor0 is not None:
-        example_input = input_tensor0
-        command = f"pnnx {pt_path} inputshape={str(shape0).replace(' ', '')}"
-    else:
-        example_input = input_tensor1
-        command = f"pnnx {pt_path}"
-    print_log(task_id, model_name, "生成输入张量", "完成")
-
-    # 确保 output_folder 存在
-    if not os.path.exists(output_folder):
-        os.makedirs(output_folder)
-
-    # load a model from disk
-    model = ModelLoader().load_from_file(file_path)
-
-    # make sure it's an image to image model
-    assert isinstance(model, ImageModelDescriptor)
-
-    print_log(task_id, model_name, "加载模型", "完成")
-    # send it to the GPU and put it in inference mode
-    # model.cuda().eval()
-    model.eval()
-    torch_model = model.model
-    print_log(task_id, model_name, "获得模型对象", "完成")
-    if os.path.exists(pt_path):
-        print_log(task_id, model_name, "转换为TorchScript模型", "跳过")
-    else:
-        print_log(task_id, model_name, "转换为TorchScript模型", "开始")
-        # 使用 torch.jit.trace 进行模型转换
-        traced_torch_model = torch.jit.trace(torch_model, example_input)
-        traced_torch_model.save(output_folder + "/" + model_name + ".pt")
-        print_log(task_id, model_name, "转换为TorchScript模型", "完成")
-
-    print_log(task_id, model_name, "运行命令"+command, "开始")
-
-    try:
-        # 使用 subprocess.Popen 执行命令
-        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
-        while True:
-            output = process.stdout.readline()
-            if output == '' and process.poll() is not None:
-                break
-            if output:
-                log += output.strip() + '\n'
-                if log_to_terminal:
-                    print(output.strip())
-        returncode = process.poll()
-        if returncode != 0:
-            log += f"执行命令: {command} 失败，返回码: {returncode}\n"
+# 修改 start_process 函数，处理新增输入
+def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
+    global task_counter
+    task_counter += 1
+    task_id = task_counter
+    log = "转换过程非常慢，请耐心等待。显示文件列表不代表转换完成。如果未发生错误，转换结束会显示”任务完成“\n"
+    yield [], log
+    if input2 == None or input2.strip() == "":
+        if isinstance(input1, str):
+            input2 = os.path.splitext(os.path.basename(input1))[0]
         else:
-            log += f"执行命令: {command} 成功\n"
-    except Exception as e:
-        log += f"执行命令: {command} 失败，错误信息: {str(e)}\n"
+            input2 = os.path.splitext(os.path.basename(input1.name))[0]
 
-    return [os.path.join(output_folder, f) for f in os.listdir(output_folder) if os.path.isfile(os.path.join(output_folder, f))], log
+        if input2 == "":
+            input2 = str(task_id)
+        log += f"未提供文件名，使用{input2}\n"
+        print_log(task_id, input2,  f"未提供文件名，使用{input2}", "修正")
+        yield [], log
+        input2 = "output"
 
-# 修改为字典类型
-downloaded_files = {}
-# 修改 start_process 函数，处理新增输入
-def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
-    task_id = str(uuid.uuid4())
-    log = ""
     try:
         # 判断 input1 是地址还是文件，增加对 ftp 和 webdav 协议的支持
         supported_protocols = ('http://', 'https://', 'ftp://', 'webdav://')
@@ -98,10 +49,13 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             url = input1
             if url in downloaded_files and os.path.exists(downloaded_files[url]):
                 file_path = downloaded_files[url]
-                log += f"跳过下载，文件已存在: {file_path}\n"
                 print_log(task_id, input2, "检查下载状态", "跳过下载")
+                log += f"跳过下载，文件已存在: {file_path}\n"
                 yield [], log
             else:
+                print_log(task_id, input2, "下载文件", "开始")
+                log += f"开始下载文件…\n"
+                yield [], log
                 # 生成唯一文件名
                 file_name = str(uuid.uuid4()) + input_suffix
                 file_path = os.path.join(os.getcwd(), file_name)
@@ -117,27 +71,29 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
                             ftp.retrbinary('RETR ' + remote_file_path, f.write)
                         ftp.quit()
                         downloaded_files[url] = file_path
+                        print_log(task_id, input2, "下载文件", "成功")
                         log += f"文件下载成功: {file_path}\n"
                         yield [], log
                     except Exception as e:
-                        log += f"FTP 文件下载失败: {str(e)}\n"
                         print_log(task_id, input2, "下载文件", f"失败 (FTP): {str(e)}")
+                        log += f"FTP 文件下载失败: {str(e)}\n"
                         yield [], log
                         return
-                else :
+                else:
                     if url.startswith(('http://', 'https://')):
                         response = requests.get(url)
                         if response.status_code == 200:
                             with open(file_path, 'wb') as f:
                                 f.write(response.content)
                             downloaded_files[url] = file_path
+                            print_log(task_id, input2, "下载文件", "成功")
                             log += f"文件下载成功: {file_path}\n"
                             yield [], log
                         else:
+                            print_log(task_id, input2, f"下载文件(HTTP): {response.status_code}", "失败")
                             log += f"文件下载失败，状态码: {response.status_code}\n"
                             yield [], log
                             return
-
         elif input1 is not None:
             file_path = input1.name
             log += f"使用上传的文件: {file_path}\n"
@@ -152,10 +108,9 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         # 生成新文件夹用于暂存结果
         output_folder = os.path.join(os.getcwd(), str(uuid.uuid4()))
         os.makedirs(output_folder, exist_ok=True)
-        log += f"创建临时文件夹: {output_folder}\n"
         print_log(task_id, input2, "创建临时文件夹", "完成")
+        log += f"创建临时文件夹: {output_folder}\n生成张量\n"
         yield [], log
-
         # 解析输入的字符串为数组
         try:
             # 尝试解析 shape0_str
@@ -176,18 +131,94 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             yield [], log
             return
 
-        # 调用处理函数，传递 shape 数组
-        output_files, process_log = process_file(task_id, file_path, input2, output_folder, shape0, shape1)
-        log += process_log
-        log += f"处理完成，输出文件: {output_files}\n"
-        print_log(task_id, input2, "调用处理函数", "完成")
+        # 以下是 process_file 函数的代码
+        # 使用 torch.rand 生成 input_shape
+        print_log(task_id, input2, "生成输入张量", "开始")
+        log += "生成张量…\n"
+        yield [], log
+        pt_path = output_folder + "/" + input2 + ".pt"
+        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:
+            example_input = (input_tensor0, input_tensor1)
+            # 修改此处，去除 shape 字符串中的空格
+            command = f"pnnx {pt_path} inputshape={str(shape0).replace(' ', '')} inputshape2={str(shape1).replace(' ', '')}"
+        elif input_tensor0 is not None:
+            example_input = input_tensor0
+            command = f"pnnx {pt_path} inputshape={str(shape0).replace(' ', '')}"
+        else:
+            example_input = input_tensor1
+            command = f"pnnx {pt_path}"
+        print_log(task_id, input2, "生成输入张量", "完成")
+
+
+        # 确保 output_folder 存在
+        if not os.path.exists(output_folder):
+            os.makedirs(output_folder)
+
+        print_log(task_id, input2, "加载模型", "开始")
+        log += "加载模型…\n"
+        yield [], log
+        # load a model from disk
+        model = ModelLoader().load_from_file(file_path)
+
+        # make sure it's an image to image model
+        assert isinstance(model, ImageModelDescriptor)
+
+        print_log(task_id, input2, "获得���型对象", "开始")
+        log += "获得模型对象…\n"
+        yield [], log
+        # send it to the GPU and put it in inference mode
+        # model.cuda().eval()
+        model.eval()
+        torch_model = model.model
+        print_log(task_id, input2, "获得模型对象", "完成")
+        yield [], log
+        if os.path.exists(pt_path):
+            print_log(task_id, input2, "转换为TorchScript模型", "跳过")
+            log += "跳过转换为TorchScript模型\n"
+            yield [], log
+        else:
+            print_log(task_id, input2, "转换为TorchScript模型", "开始")
+            log+= "转换为TorchScript模型…\n"
+            yield [], log
+            # 使用 torch.jit.trace 进行模型转换
+            traced_torch_model = torch.jit.trace(torch_model, example_input)
+            traced_torch_model.save(output_folder + "/" + input2 + ".pt")
+            print_log(task_id, input2, "转换为TorchScript模型", "完成")
+
+        print_log(task_id, input2, "执行命令" + command, "开始")
+        log += "执行命令…\n"
+        yield [], log
+
+        try:
+            # 使用 subprocess.Popen 执行命令
+            process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+            while True:
+                output = process.stdout.readline()
+                if output == '' and process.poll() is not None:
+                    break
+                if output:
+                    log += output.strip() + '\n'
+                    if log_to_terminal:
+                        print(output.strip())
+            returncode = process.poll()
+            if returncode != 0:
+                log += f"执行命令: {command} 失败，返回码: {returncode}\n"
+            else:
+                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"
+        print_log(task_id, input2, "执行命令", "完成")
         yield output_files, log
     except Exception as e:
         log += f"发生错误: {str(e)}\n"
-        print_log(task_id, input2, "整体处理", f"失败: {str(e)}")
+        print_log(task_id, input2,str(e) , f"失败")
         yield [], log
 
-
 # 创建 Gradio 界面
 with gr.Blocks() as demo:
     gr.Markdown("文件处理界面")