app.py

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

# 定义 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"任务{task_id}: {filename}, [{status}] {stage}")

# 修改 start_process 函数，处理新增输入
def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
    global task_counter
    task_counter += 1
    task_id = task_counter
    print_log(task_id, input2,  input1, "input1")
    log = "转换过程非常慢，请耐心等待。显示文件列表不代表转换完成。如果未发生错误，转换结束会显示”任务完成“\n"
    yield [], log
    if input2 == None or input2.strip() == "":
        split_input = os.path.splitext(os.path.basename(input1))
        if len(split_input) > 1:
            suffix = split_input[1].split('?')[0].lower()
            if  suffix not in [".pth" , ".safetensors" , ".ckpt"]:
                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)
        log += f"未提供文件名，使用{input2}\n"
        print_log(task_id, input2,  f"未提供文件名，使用{input2}", "修正")
        yield [], 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]
                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(task_id) + 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
                        print_log(task_id, input2, "下载文件", "成功")
                        log += f"文件下载成功: {file_path}\n"
                        yield [], log
                    except Exception as e:
                        print_log(task_id, input2, "下载文件", f"失败 (FTP): {str(e)}")
                        log += f"FTP 文件下载失败: {str(e)}\n"
                        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
                            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:
            print("check file" , input1, os.path.exists(input1))
            file_path = input1
            log += f"使用上传的文件: {file_path}\n"
            print_log(task_id, input2, "使用上传文件", "开始")
            yield [], log
        else:
            log += "未提供有效文件或地址\n"
            print_log(task_id, input2, "检查文件输入", "失败 (无有效输入)")
            yield [], log
            return

        # 检查文件大小
        try:
            file_size = os.path.getsize(file_path) / 1024 /1024  # 转换为 KB
            if file_size > 100 :
                log += f"文件太大，建议 100MB 以内，当前文件大小为 {file_size } MB。\n"
                print_log(task_id, input2, "文件太大（"+ file_size +"MB）", "失败")
                yield [], log
                return
        except Exception as e:
            log += f"获取文件大小失败: {str(e)}\n"
            print_log(task_id, input2, "检查文件大小", f"失败: {str(e)}")
            yield [], log
            return

        # 生成新文件夹用于暂存结果
        output_folder = os.path.join(os.getcwd(), str(uuid.uuid4()))
        os.makedirs(output_folder, exist_ok=True)
        print_log(task_id, input2, "创建临时文件夹", "完成")
        log += f"创建临时文件夹: {output_folder}\n生成张量\n"
        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

        # 以下是 process_file 函数的代码
        # 使用 torch.rand 生成 input_shape
        print_log(task_id, input2, "生成输入张量", "开始")
        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:
            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


        width_ratio = 4
        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模型", "完成")

            # 获取输出
            example_output = traced_torch_model(example_input)
            width_ratio = example_output.shape[2] / example_input.shape[2]
            print_log(task_id, input2, "获得缩放倍率="+ str(width_ratio)+", 输出shape="+str(list(example_output.shape)), "完成")
            log+= ("获得缩放倍率="+str(width_ratio)+", 输出shape="+str(list(example_output.shape))+"\n")
            yield [], log

        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:
                    # 
                    if log_to_terminal:
                        print(output.strip())
                        log += output.strip() + '\n'
                        yield [], log
            returncode = process.poll()
            if returncode != 0:
                log += f"执行命令失败，返回码: {returncode}，命令: {command} \n"
            else:
                log += f"执行命令成功: {command} \n"
        except Exception as e:
            log += f"执行命令: {command} 失败，错误信息: {str(e)}\n"
            
        # 查找 output_folder 目录下以 .ncnn.bin 和 .ncnn.param 结尾的文件
        bin_files = [f for f in os.listdir(output_folder) if f.endswith('.ncnn.bin')]
        param_files = [f for f in os.listdir(output_folder) if f.endswith('.ncnn.param')]
        

        if bin_files and param_files:
            param_file = os.path.join(output_folder, param_files[0])
            bin_file = os.path.join(output_folder, bin_files[0])
            import zipfile
            # 压缩包名称
            zip_file_name = os.path.join(output_folder, f"models-{input2}.zip")
            # 压缩包内文件夹名称
            zip_folder_name = f"models-{input2}"
            # 重命名后的文件名
            scale = int(width_ratio)
            new_bin_name = f"x{scale}.bin"
            new_param_name = f"x{scale}.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
        else:
            log += f"未找到 ncnn 文件\n"
            print_log(task_id, input2, "查找 ncnn 文件", "失败")
            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"
        print_log(task_id, input2, "执行命令", "完成")
        yield output_files, log
    except Exception as e:
        log += f"发生错误: {e}\n"
        print_log(task_id, input2, e , f"失败")
        yield [], log

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("文件处理界面")
    with gr.Row():
        # 左侧列，包含输入组件和按钮
        with gr.Column():
            # 添加文本提示
            gr.Markdown("请输入的url，或者上传一个文件。限制文件为小于100M的*.pth模型")
            with gr.Row():
                input1 = gr.Textbox(label="粘贴地址")
                # 新增文件上传组件
                input1_file = gr.File(label="上传文件", file_types=[".pth", ".safetensors", ".ckpt"])
            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")
            with gr.Row():
                start_button = gr.Button("开始")
                # 添加取消按钮
                cancel_button = gr.Button("取消")
        # 右侧列，包含输出组件和日志文本框
        with gr.Column():
            output = gr.File(label="输出文件", file_count="multiple")
            log_textbox = gr.Textbox(label="日志", lines=10, interactive=False)

    # 绑定事件，修改输入参数
    process = start_button.click(
        fn=start_process,
        inputs=[input1_file if input1_file.value else input1, input2, shape0_str, shape1_str],
        outputs=[output, log_textbox]
    )
    # 为取消按钮添加点击事件绑定，使用 cancels 属性取消 start_process 任务
    cancel_button.click(
        fn=None,
        inputs=None,
        outputs=None,
        cancels=[process]
    )

    # 添加范例
    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://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-WTP-ColorDS.pth", "", "1,3,128,128", "0,0,0,0"],
    ]
    gr.Examples(
        examples=examples,
        inputs=[input1, input2, shape0_str, shape1_str],
        outputs=[output, log_textbox],
        fn=start_process
    )

demo.launch()