update

.gitignore

@@ -3,3 +3,4 @@
 *.bin
 *.onnx
 *.param
+*.zip

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import os
+from gradio.themes import default
 import requests
 import shutil
 import uuid
@@ -23,12 +24,14 @@ def print_log(task_id, filename, stage, status):
         print(f"任务{task_id}: {filename}, [{status}] {stage}")
 
 # 修改 start_process 函数，处理新增输入
-def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
+def start_process(input_file, input_url, input2, shape0_str, shape1_str, output_type="+ Dynamic MNN & NCNN", input_suffix=".pth"):
     global task_counter
     task_counter += 1
     task_id = task_counter
+    input1 = input_file if input_file else input_url
     print_log(task_id, input2,  input1, "input1")
     log = "转换过程非常慢，请耐心等待。显示文件列表不代表转换完成。如果未发生错误，转换结束会显示”任务完成“\n"
+    output_files = []
     yield [], log
     if input2 == None or input2.strip() == "":
         split_input = os.path.splitext(os.path.basename(input1))
@@ -115,8 +118,8 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         # 检查文件大小
         try:
             file_size = os.path.getsize(file_path) / 1024 /1024  # 转换为 KB
-            if file_size > 100 :
-                log += f"文件太大，建议 100MB 以内，当前文件大小为 {file_size } MB。\n"
+            if file_size > 200 :
+                log += f"文件太大，建议 200MB 以内，当前文件大小为 {file_size } MB。\n"
                 print_log(task_id, input2, "文件太大（"+ file_size +"MB）", "失败")
                 yield [], log
                 return
@@ -148,7 +151,7 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             # 如果解析过程中出现 ValueError，将 shape0 和 shape1 设置为全 0
             shape0 = [0, 0, 0, 0]
             shape1 = [0, 0, 0, 0]
-            log += "输入的 shape 字符串格式不正确，请使用逗号分隔的整数。\n"
+            log += f"输入的 shape 字符串格式不正确，请使用逗号分隔的整数。shape0_str={shape0_str},shape1_str={shape1_str}\n"
             yield [], log
             return
 
@@ -157,15 +160,15 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
         print_log(task_id, input2, "生成输入张量", "开始")
         log += "生成张量…\n"
         yield [], log
-        pt_path = output_folder + "/" + input2 + ".pt"
-        # onnx_path = output_folder + "/" + input2 + ".onnx"
+        output_base =  output_folder + "/" + input2
+        pt_path = output_base + ".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:
+        if "Fixed" in output_type and 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:
+        elif "Fixed" in output_type and input_tensor0 is not None:
             example_input = input_tensor0
             command = f"pnnx {pt_path} inputshape={str(shape0).replace(' ', '')}"
         else:
@@ -219,61 +222,129 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
             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, "完成")
+        # 转换为 ONNX 模型
+        if str(width_ratio) in input2:
+            onnx_path = output_base + ".onnx"
+        else:
+            onnx_path = output_base + "-x" + str(width_ratio) + ".onnx"
+        if os.path.exists(onnx_path):
+            print_log(task_id, input2, "转换为ONNX模型", "跳过")
+            log += "跳过转换为ONNX模型\n"
             yield [], log
         else:
-            log += f"未找到 ncnn 文件\n"
-            print_log(task_id, input2, "查找 ncnn 文件", "失败")
+            print_log(task_id, input2, "转换为ONNX模型", "开始")
+            log += "转换为ONNX模型…\n"
             yield [], log
+            torch.onnx.export(torch_model, example_input, onnx_path, opset_version=17, input_names=["input"], output_names=["output"])
+
+
+        # 转换为 mnn 模型
+        if "MNN" in output_type:
+            
+            if str(width_ratio) in input2:
+                mnn_path = output_base + ".mnn"
+            else:
+                mnn_path = output_base + "-x" + str(width_ratio) + ".mnn"
+
+            mnn_config = ""
+            if "Fixed" in output_type and input_tensor0 is not None:
+                mnn_config = output_base + ".mnnconfig"
+                with open(mnn_config, 'w') as f:
+                    if input_tensor1 is not None:
+                        f.write(f"input_names = input0, input1\n")
+                        f.write(f"input_dims = {'x'.join(map(str, shape0))}, {'x'.join(map(str, shape0))},\n")
+                    else:
+                        f.write(f"input_names = input\n")
+                        f.write(f"input_dims = {'x'.join(map(str, shape0))}\n")
+
+            if os.path.exists(mnn_path):
+                print_log(task_id, input2, "转换为MNN模型", "跳过")
+                log += "跳过转换为MNN模型\n"
+                yield [], log
+            else:
+                print_log(task_id, input2, "转换为MNN模型", "开始")
+                log += "转换为MNN模型…\n"
+                mnn_command = f"MNNConvert -f ONNX  --modelFile \"{onnx_path}\" --MNNModel \"{mnn_path}\"  --bizCode biz --fp16 --info --detectSparseSpeedUp"
+                if mnn_config:
+                    mnn_command += f" --inputConfigFile \"{mnn_config}\""
+                try:
+                    # 使用 subprocess.Popen 执行命令
+                    process = subprocess.Popen(mnn_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"执行mnn命令失败，返回码: {returncode}，命令: {mnn_command} \n"
+                    else:
+                        log += f"执行命令成功: {mnn_command} \n"
+                except Exception as e:
+                    log += f"执行命令: {mnn_command} 失败，错误信息: {str(e)}\n"
+
+        if "NCNN" in output_type:
+            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"
@@ -296,7 +367,11 @@ with gr.Blocks() as demo:
                 input1 = gr.Textbox(label="粘贴地址")
                 # 新增文件上传组件
                 input1_file = gr.File(label="上传文件", file_types=[".pth", ".safetensors", ".ckpt"])
-            input2 = gr.Textbox(label="自定义文件名")
+            
+            
+            with gr.Row():
+                input2 = gr.Textbox(label="自定义文件名")
+                output_type = gr.Dropdown(choices=["TorchScript & ONNX", "+ Fixed MNN & NCNN", "+ Dynamic MNN & NCNN", "Fixed NCNN", "Dynamic NCNN", "Fixed MNN", "Dynamic MNN" ], value="+ Dynamic MNN & NCNN", label="模型类型")
             # 修改为字符串输入控件
             shape0_str = gr.Textbox(label="shape0 (逗号分隔的整数)", value="1,3,128,128")
             shape1_str = gr.Textbox(label="shape1 (逗号分隔的整数)", value="0,0,0,0")
@@ -312,7 +387,7 @@ with gr.Blocks() as demo:
     # 绑定事件，修改输入参数
     process = start_button.click(
         fn=start_process,
-        inputs=[input1_file if input1_file.value else input1, input2, shape0_str, shape1_str],
+        inputs=[input1_file, input1, input2, output_type, shape0_str, shape1_str,output_type],
         outputs=[output, log_textbox]
     )
     # 为取消按钮添加点击事件绑定，使用 cancels 属性取消 start_process 任务
@@ -325,14 +400,14 @@ with gr.Blocks() as demo:
 
     # 添加范例
     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"],
+        ["","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],
+        inputs=[input1_file, input1, input2, shape0_str, shape1_str],
         outputs=[output, log_textbox],
         fn=start_process
     )

requirements.txt

@@ -1,4 +1,5 @@
 spandrel
 torch
 pnnx
-onnx
+onnx
+mnn