Upload app.py

app.py

@@ -0,0 +1,217 @@
+import gradio as gr
+import os
+import requests
+import shutil
+import uuid
+from ftplib import FTP
+from spandrel import ImageModelDescriptor, ModelLoader
+import torch
+import subprocess
+
+
+# 新增日志开关
+log_to_terminal = True
+
+# 新增日志函数
+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"
+        else:
+            log += f"执行命令: {command} 成功\n"
+    except Exception as e:
+        log += f"执行命令: {command} 失败，错误信息: {str(e)}\n"
+
+    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
+
+# 修改为字典类型
+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://')
+        if isinstance(input1, str) and input1.startswith(supported_protocols):
+            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, "检查下载状态", "跳过下载")
+                yield [], log
+            else:
+                # 生成唯一文件名
+                file_name = str(uuid.uuid4()) + input_suffix
+                file_path = os.path.join(os.getcwd(), file_name)
+                if url.startswith('ftp://'):
+                    try:
+                        # 解析 ftp 地址
+                        parts = url.replace('ftp://', '').split('/')
+                        host = parts[0]
+                        remote_file_path = '/'.join(parts[1:])
+                        ftp = FTP(host)
+                        ftp.login()
+                        with open(file_path, 'wb') as f:
+                            ftp.retrbinary('RETR ' + remote_file_path, f.write)
+                        ftp.quit()
+                        downloaded_files[url] = file_path
+                        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)}")
+                        yield [], log
+                        return
+                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
+                            log += f"文件下载成功: {file_path}\n"
+                            yield [], log
+                        else:
+                            log += f"文件下载失败，状态码: {response.status_code}\n"
+                            yield [], log
+                            return
+
+        elif input1 is not None:
+            file_path = input1.name
+            log += f"使用上传的文件: {file_path}\n"
+            print_log(task_id, input2, "使用上传文件", "开始")
+            yield [], log
+        else:
+            log += "未提供有效文件或地址\n"
+            print_log(task_id, input2, "检查文件输入", "失败 (无有效输入)")
+            yield [], log
+            return
+
+        # 生成新文件夹用于暂存结果
+        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, "创建临时文件夹", "完成")
+        yield [], log
+
+        # 解析输入的字符串为数组
+        try:
+            # 尝试解析 shape0_str
+            shape0 = [int(x) for x in shape0_str.split(',')] if shape0_str else [0, 0, 0, 0]
+            # 检查 shape0 是否为 4 个元素，如果不是则设置为全 0
+            if len(shape0) != 4:
+                shape0 = [0, 0, 0, 0]
+            # 尝试解析 shape1_str
+            shape1 = [int(x) for x in shape1_str.split(',')] if shape1_str else [0, 0, 0, 0]
+            # 检查 shape1 是否为 4 个元素，如果不是则设置为全 0
+            if len(shape1) != 4:
+                shape1 = [0, 0, 0, 0]
+        except ValueError:
+            # 如果解析过程中出现 ValueError，将 shape0 和 shape1 设置为全 0
+            shape0 = [0, 0, 0, 0]
+            shape1 = [0, 0, 0, 0]
+            log += "输入的 shape 字符串格式不正确，请使用逗号分隔的整数。\n"
+            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, "调用处理函数", "完成")
+        yield output_files, log
+    except Exception as e:
+        log += f"发生错误: {str(e)}\n"
+        print_log(task_id, input2, "整体处理", f"失败: {str(e)}")
+        yield [], log
+
+
+# 创建 Gradio 界面
+with gr.Blocks() as demo:
+    gr.Markdown("文件处理界面")
+    with gr.Row():
+        # 左侧列，包含输入组件和按钮
+        with gr.Column():
+            with gr.Row():
+                input1 = gr.Textbox(label="粘贴地址")
+                input1_file = gr.File(label="上传文件")
+            input2 = gr.Textbox(label="自定义文件名")
+            # 修改为字符串输入控件
+            shape0_str = gr.Textbox(label="shape0 (逗号分隔的整数)", value="1,3,128,128")
+            shape1_str = gr.Textbox(label="shape1 (逗号分隔的整数)", value="0,0,0,0")
+            start_button = gr.Button("开始")
+        # 右侧列，包含输出组件和日志文本框
+        with gr.Column():
+            output = gr.File(label="输出文件", file_count="multiple")
+            log_textbox = gr.Textbox(label="日志", lines=10, interactive=False)
+
+    # 绑定事件，修改输入参数
+    start_button.click(
+        fn=start_process,
+        inputs=[input1_file if input1_file.value else input1, input2, shape0_str, shape1_str],
+        outputs=[output, log_textbox]
+    )
+
+demo.launch()