Spaces:
Running
Running
| 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() |