Upload 5 files

app.py

@@ -0,0 +1,252 @@
+import gradio as gr
+import requests
+import os
+import subprocess
+from typing import Union
+from datetime import datetime
+from pth2onnx import convert_pth_to_onnx
+from urllib.parse import urlparse
+from typing import Optional
+
+# 新增日志开关
+log_to_terminal = True
+
+# 新增全局任务计数器
+task_counter = 0
+
+# 新增日志函数
+def print_log(task_id, stage, status):
+    if log_to_terminal:
+        print(f"任务{task_id}: [{status}] {stage}")
+
+# 使用 MNN 库自带的转换工具
+def convertmnn(onnx_path: str, mnn_path: str, fp16=False):
+    param = ['mnnconvert', '-f', 'ONNX', '--modelFile', onnx_path, '--MNNModel', mnn_path, '--bizCode', 'biz', '--info', '--detectSparseSpeedUp']
+    if fp16:
+        param.append('--fp16')
+    subprocess.run(param, check=True)
+
+def download_file(url: str, save_path: str):
+    response = requests.get(url)
+    with open(save_path, 'wb') as f:
+        f.write(response.content)
+
+def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min: int) -> Optional[str]:
+    """
+    从URL下载文件到指定文件夹，并进行文件大小检查
+    
+    Args:
+        url: 要下载的文件URL
+        folder: 保存文件的目标文件夹路径
+        filesize_max: 最大允许文件大小(字节)，超过此值将中断下载并删除文件
+        filesize_min: 最小允许文件大小(字节)，小于此值将删除文件并返回None
+    
+    Returns:
+        成功下载的文件名，如果下载失败或文件大小不符合要求则返回None
+    """
+    # 解析URL获取文件名
+    parsed_url = urlparse(url)
+    filename = os.path.basename(parsed_url.path)
+    if not filename:
+        return None  # 无法从URL获取文件名
+    
+    # 确保目标文件夹存在
+    os.makedirs(folder, exist_ok=True)
+    save_path = os.path.join(folder, filename)
+    
+    try:
+        # 发送HTTP请求，流式下载
+        with requests.get(url, stream=True, timeout=10) as response:
+            response.raise_for_status()  # 检查HTTP错误状态
+            
+            # 获取文件总大小(如果服务器提供)
+            total_size = int(response.headers.get('content-length', 0))
+            if total_size > filesize_max:
+                return None  # 文件大小超过最大值，不下载
+            
+            downloaded_size = 0
+            with open(save_path, 'wb') as file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:  # 过滤空块
+                        downloaded_size += len(chunk)
+                        # 检查是否超过最大允许大小
+                        if downloaded_size > filesize_max:
+                            file.close()
+                            os.remove(save_path)
+                            return None
+                        file.write(chunk)
+        
+        # 下载完成后检查最小文件大小
+        if os.path.getsize(save_path) < filesize_min:
+            os.remove(save_path)
+            return None
+        
+        return filename
+    
+    except Exception as e:
+        # 发生异常时清理文件
+        if os.path.exists(save_path):
+            os.remove(save_path)
+        return None
+
+# 原 process_model 函数重命名为 _process_model
+async def _process_model(model_input: Union[str, gr.File], tilesize: int, log_box: gr.Textbox, output_dir: str):
+    global task_counter
+    task_id = task_counter
+    log = ('初始化日志记录...')
+    print_log(task_id, '初始化日志记录', '开始')
+    yield [],[], log
+
+    # 处理输入模型
+    if isinstance(model_input, str):  # 处理链接
+        if model_input.startswith(('http://', 'https://')):
+            log += ( f'正在下载模型文件: {model_input}')
+            print_log(task_id, f'正在下载模型文件: {model_input}', '开始')
+            yield  [],[], log
+
+            # 下载文件到output文件夹
+            filename = download_file2folder(
+                url=model_input,
+                folder=output_dir,
+                filesize_max=200*1024*1024,  # 200MB
+                filesize_min=1024     # 1KB
+            )
+            input_path = os.path.join(output_dir, filename)
+            log += ( f'模型文件已下载到: {input_path}')
+            print_log(task_id, f'模型文件已下载到: {input_path}', '完成')
+            yield  [],[], log
+        else:
+            input_path = model_input
+            log += ( f'使用本地文件: {input_path}')
+            print_log(task_id, f'使用本地文件: {input_path}', '开始')
+            yield  [],[], log
+    else:
+        input_path = model_input.name
+        log += ( f'已上传模型文件: {input_path}')
+        print_log(task_id, f'已上传模型文件: {input_path}', '开始')
+        yield  [],[], log
+
+
+    if not input_path:
+        log += ( f'未获得正确的模型文件')
+        print_log(task_id, f'未获得正确的模型文件', '错误')
+        yield  [],[], log
+        return
+
+
+    if input_path.endswith('.onnx'):
+        onnx_path = input_path
+        log += ( '输入已是 ONNX 模型，直接使用...')
+        print_log(task_id, '输入已是 ONNX 模型，直接使用', '开始')
+        yield [],[], log
+    else:
+        print_log(task_id, f'转换 PTH 模型为 ONNX, folder={output_dir}', '开始')
+        onnx_path = convert_pth_to_onnx(input_path, tilesize=tilesize, output_folder=output_dir)
+        if onnx_path:
+            log += ( f'成功生成ONNX模型: {onnx_path}')
+            print_log(task_id, f'生成ONNX模型: {onnx_path}', '完成')
+        else:
+            log += ( '生成ONNX模型失败')
+            print_log(task_id, '生成ONNX模型', '错误')
+            yield [], [], log
+            return
+
+
+    # 转换为 MNN 模型
+    output_name= os.path.splitext(os.path.basename(onnx_path))[0]
+    mnn_path = os.path.join(output_dir,  f'{output_name}.mnn')
+    try:
+        log += ( '正在将 ONNX 模型转换为 MNN 格式...')
+        print_log(task_id, '正在将 ONNX 模型转换为 MNN 格式', '开始')
+        convertmnn(onnx_path, mnn_path)
+        yield onnx_path,[], log
+    except Exception as e:
+        log += ( f'转换 MNN 模型时出错: {str(e)}')
+        print_log(task_id, f'转换 MNN 模型时出错: {str(e)}', '错误')
+        yield onnx_path,[], log
+
+    print_log(task_id, '模型转换任务完成', '完成')
+    
+    # 转换为 MNN 模型后对文件检查
+    if os.path.exists(mnn_path) and os.path.getsize(mnn_path) > 1024:  # 1KB = 1024 bytes
+        log += ( f'MNN 模型已保存到: {mnn_path}')
+    else:
+        log += ( 'MNN 模型生成失败或文件大小不足1KB')
+        mnn_path = None
+    
+    yield onnx_path, mnn_path, log
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 模型转换工具")
+    with gr.Row():
+        with gr.Column():
+            input_type = gr.Radio(['模型链接', '上传模型文件'], label='输入类型')
+            url_input = gr.Textbox(label='模型链接')
+            file_input = gr.File(label='上传模型文件', visible=False)
+
+            
+            def show_input(input_type):
+                if input_type == '模型链接':
+                    return gr.update(visible=True), gr.update(visible=False)
+                else:
+                    return gr.update(visible=False), gr.update(visible=True)
+            
+            input_type.change(show_input, inputs=input_type, outputs=[url_input, file_input])
+            
+            tilesize = gr.Number(label="Tilesize", value=0, precision=0)
+            convert_btn = gr.Button("开始转换")
+        with gr.Column():
+            log_box = gr.Textbox(label="转换日志", lines=10, interactive=False)
+    with gr.Row():
+        onnx_output = gr.File(label="ONNX 模型输出")
+        mnn_output = gr.File(label="MNN 模型输出")
+
+    async def process_model(input_type, url_input, file_input, tilesize, log_box):
+        if input_type == '模型链接' and url_input:
+            model_input = url_input
+        elif input_type == '上传模型文件' and file_input:
+            model_input = file_input
+        else:
+            # 改为通过yield返回错误日志
+            log = '\n请选择输入类型并提供有效的输入！'
+            yield None, None, log
+            return
+
+        # 创建不重名的输出目录
+        global task_counter
+        task_counter += 1
+        output_dir = os.path.join(os.getcwd(), f"output_{task_counter}")
+        os.makedirs(output_dir, exist_ok=True)
+        
+        onnx_path = None
+        mnn_path = None
+        # 调用重命名后的函数
+        async for result in _process_model(model_input, int(tilesize), log_box, output_dir):
+            if isinstance(result, tuple) and len(result) == 3:
+                onnx_path, mnn_path, log_box = result
+            elif isinstance(result, tuple) and len(result) == 2:
+                # 处理纯日志yield
+                _, process_log = result
+                yield None, None, process_log
+            yield onnx_path, mnn_path, log_box
+
+    convert_btn.click(
+        process_model,
+        inputs=[input_type, url_input, file_input, tilesize, log_box],
+        outputs=[onnx_output, mnn_output, log_box],
+        api_name="convert_model"
+    )
+
+    # 将示例移至底部并包裹在列组件中
+    examples_column = gr.Column(visible=True)
+    with examples_column:
+        examples = [
+            ["模型链接", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"],
+            ["模型链接", "https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth"],
+            ["模型链接", "https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth"],
+            ["模型链接", "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-WTP-ColorDS.pth"],
+            ["模型链接", "https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV4/2x-AnimeSharpV4_RCAN_fp16_op17.onnx"]
+        ]
+        example_input = gr.Examples(examples=examples, inputs=[input_type, url_input], label='示例模型链接')
+    
+demo.launch()

onnx2mnn2.bat

@@ -0,0 +1,121 @@
+@echo off
+setlocal enabledelayedexpansion
+
+rem ���̶�shape��fp16����ͼƬ����
+rem ����Ƿ��в�������
+if "%~1"=="" (
+    rem ������������ʾ�û������ļ�·��
+    set /p onnx_path="������ .onnx �ļ���·��: "
+) else (
+    rem ��������ʹ�õ�һ��������Ϊ�ļ�·��
+    set "onnx_path=%~1"
+    cd /d %~dp0
+
+)
+
+:main
+rem ����ļ��Ƿ�������Ƿ��� .onnx ��β
+if not exist "!onnx_path!" (
+    echo ����: �ļ������ڣ�����������·����
+    goto loop
+)
+
+set "onnx_ext=!onnx_path:~-5!"
+if /i "!onnx_ext!" neq ".onnx" (
+    echo ����: �ļ����� .onnx ��ʽ������������·����
+    goto loop
+)
+
+rem ȥ��Ŀ¼·���ͺ�׺��
+for %%f in ("!onnx_path!") do set "onnx_name=%%~nf"
+echo ��ǰ����Ŀ¼��: %cd%
+rem ִ�� onnx2ncnn.exe
+MNNConvert -f ONNX  --modelFile "!onnx_path!" --MNNModel "!onnx_name!.mnn"  --bizCode biz --fp16  --info  --detectSparseSpeedUp
+
+rem ��� fp16 �ļ��Ƿ����
+if exist  "!onnx_name!.mnn"   (
+    rem ����
+) else (
+    echo δ���ģ���ļ�������������·����
+    goto loop
+)
+
+
+
+:next_step
+rem ��ʾ�û�ѡ����һ������
+echo.
+echo ��ѡ����һ������:
+echo 1. ��һ��ѭ����Ĭ�ϣ�
+echo 2. ����ģ�ͣ�ʹ�� test.png��
+echo 3. �����Զ���ͼƬ·��������ģ��
+set /p choice="������ѡ�� (1/2/3): "
+if "%choice%"=="" set "choice=1"
+
+if "%choice%"=="1" (
+    goto loop
+) else if "%choice%"=="2" (
+    set "test_image=test.png"
+    if exist "!test_image!" (
+        mnnsr-ncnn.exe -i "!test_image!" -o "!onnx_name!.png" -m "!onnx_name!.mnn" -s 0
+    ) else (
+        echo ����: �ļ� "!test_image!" �����ڣ�������ѡ��
+        goto next_step
+    )
+) else if "%choice%"=="3" (
+    set /p custom_image="�������Զ���ͼƬ·��: "
+    if exist "!custom_image!" (
+        mnnsr-ncnn.exe -i "!custom_image!" -o "!onnx_name!.png" -m "!onnx_name!.mnn" -s 0
+    ) else (
+        echo ����: �ļ� "!custom_image!" �����ڣ�������ѡ��
+        goto next_step
+    )
+) else (
+    echo ��Ч��ѡ�������ѡ��
+    goto next_step
+)
+
+rem �������� PNG �ļ��Ƿ����
+if exist "!onnx_name!.png" (
+    echo ����ļ����ɳɹ�: "!onnx_name!.png"
+    start "" "!onnx_name!.png"
+) else (
+    echo ��֤ʧ��: ����ļ������ڣ���ѡ���Ƿ�ɾ��ģ���ļ���
+    set /p delete_model="�Ƿ���ģ���ļ� "!onnx_name!.mnn" ? (y/n, Ĭ��n): "
+    if /i "!delete_model!"=="y" (
+        echo ģ���ļ�������
+    ) else (
+        del  "!onnx_name!.mnn"
+        echo ģ���ļ���ɾ����
+    )
+    goto loop
+)
+
+rem ѯ���û�ѡ����һ������
+echo.
+echo ��ѡ����һ������:
+echo 1. ������ļ��У�Ĭ�ϣ�
+echo 2. ɾ��ģ���ļ��Ͳ���ͼ
+echo 3. ��һ��ѭ��
+set /p next_choice="������ѡ�� (1/2/3): "
+if "%next_choice%"=="" set "next_choice=1"
+
+if "%next_choice%"=="1" (
+    start "" .
+) else if "%next_choice%"=="2" (
+    del "!onnx_name!.png"
+    del "!onnx_name!.mnn"
+    echo ģ���ļ��Ͳ���ͼ��ɾ����
+) else if "%next_choice%"=="3" (
+    goto loop
+) else (
+    echo ��Ч��ѡ�������ѡ��
+    goto next_choice
+)
+
+
+
+:loop
+echo ====================================
+set /p onnx_path="������ .onnx �ļ���·��: "
+goto main

pth2onnx.bat

@@ -0,0 +1,36 @@
+@echo off
+setlocal enabledelayedexpansion
+
+rem ����Ƿ��в�������
+if "%~1"=="" (
+    rem ������������ʾ�û������ļ�·��
+    set /p pth_path="������ .pth �ļ���·��: "
+) else (
+    rem ��������ʹ�õ�һ��������Ϊ�ļ�·��
+    set "pth_path=%~1"
+    cd /d %~dp0
+
+)
+
+:main
+rem ����ļ��Ƿ�������Ƿ��� .pth ��β
+if not exist "!pth_path!" (
+    echo ����: �ļ������ڣ�����������·����
+    goto loop
+)
+
+set "ext=!pth_path:~-4!"
+if /i "!ext!" neq ".pth" (
+    echo ����: �ļ����� .pth ��ʽ������������·����
+    goto loop
+)
+
+echo ��ǰ����Ŀ¼��: %cd%
+rem ִ�� pth2onnx.py
+rem python .\pth2onnx.py "!pth_path!" --fp16 --simplify --channel !channel!
+python .\pth2onnx.py "!pth_path!"  --fp16
+
+:loop
+echo ====================================
+set /p pth_path="������ .pth �ļ���·��: "
+goto main

pth2onnx.py

@@ -0,0 +1,149 @@
+import argparse
+import torch
+import os
+import re
+import onnx
+from spandrel import ImageModelDescriptor, ModelLoader
+from onnxsim import simplify
+
+def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tilesize: int = 64, use_fp16: bool=False, simplify_model: bool=False, min_size: int = 1024*1024, output_folder: str=None):
+    """
+    Loads a PyTorch model from a .pth file using Spandrel and converts it to ONNX format.
+
+    Args:
+        pth_path: Path to the input .pth model file.
+        onnx_path: Path to save the output .onnx file.
+        channel: Number of input channels for the model.
+        use_fp16: Boolean to determine if the model should be converted to half precision.
+        simplify_model: Boolean to determine if the ONNX model should be simplified.
+    """
+
+    print(f"Loading model from: {pth_path}")
+    try:
+        # Use Spandrel to load the model architecture and state dict
+        model_descriptor = ModelLoader().load_from_file(pth_path)
+
+        # Ensure it's the expected type from Spandrel
+        if not isinstance(model_descriptor, ImageModelDescriptor):
+            print(f"Error: Expected ImageModelDescriptor, but got {type(model_descriptor)}")
+            print("Please ensure the .pth file is compatible with Spandrel's loading mechanism.")
+            return False
+
+
+        # Get the underlying torch.nn.Module
+        torch_model = model_descriptor.model
+
+        # Set the model to evaluation mode (important for dropout, batchnorm layers)
+        torch_model.eval()
+
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return False
+
+    if channel == 0:
+        channel = model_descriptor.input_channels
+    if tilesize<1:
+        tilesize = 64
+    example_input = torch.randn(1, channel, tilesize, tilesize)
+    print("Model input channels:", channel, "tile size:", tilesize)
+
+    if use_fp16:
+        if torch.cuda.is_available():
+            torch_model.cuda()
+            example_input = example_input.cuda()
+        else:
+            print("Warning: no CUDA device")
+        torch_model.half()
+        example_input = example_input.half()  # 转换为半精度输入
+    print(f"Model loaded successfully: {type(torch_model).__name__}")
+
+    if output_folder:
+        os.makedirs(output_folder, exist_ok=True)
+    
+    if onnx_path is None:
+        base_path, _ = os.path.splitext(pth_path)
+        if output_folder:
+            base_path = os.path.join(output_folder, os.path.basename(base_path))
+
+        scale = model_descriptor.scale
+        # 判断 pth_path 的文件名是否包含 xs 或者 sx，x 为大小写字母 x，s 为 int scale
+        filename = os.path.basename(pth_path).upper()
+        pattern = f'(^|[_-])({scale}X|X{scale})([_-]|$)'
+        if re.search(pattern, filename):
+            print(f'文件名 {filename} 包含匹配模式。')
+        else:
+            base_path = f"{base_path}-x{scale}"
+        
+        onnx_path = base_path + ("-Grayscale" if channel==1 else "") + ("-fp16.onnx" if use_fp16 else ".onnx")
+        
+    # 处理相对路径情况
+    # elif output_folder and not os.path.isabs(onnx_path):
+    elif output_folder:
+        onnx_path = os.path.join(output_folder, onnx_path)
+        
+    print(f"output_folder: {output_folder}, onnx_path: {onnx_path}")
+
+    try:
+        # Export the model
+        torch.onnx.export(
+            torch_model,                 # The model instance
+            example_input,               # An example input tensor
+            onnx_path,                   # Where to save the model (file path)
+            export_params=True,          # Store the trained parameter weights inside the model file
+            opset_version=11,            # The ONNX version to export the model to (choose based on target runtime)
+            do_constant_folding=True,    # Whether to execute constant folding for optimization
+            input_names=['input'],       # The model's input names
+            output_names=['output'],     # The model's output names
+            dynamic_axes={               # Allow variable input/output dimensions
+                "input": {0: "batch_size", 2: "height", 3: "width"}, # Batch, H, W can vary
+                "output": {0: "batch_size", 2: "height", 3: "width"},# Batch, H, W can vary
+            }
+        )
+        print(f"ONNX export successful: {onnx_path}")
+
+        # Optional: Simplify the ONNX model
+        if simplify_model:
+            model = onnx.load(onnx_path)
+            model_simplified, _ = simplify(model)
+            onnx.save(model_simplified, onnx_path)
+            print(f"ONNX model simplified successfully: {onnx_path}")
+
+        # 添加文件验证逻辑
+        if os.path.exists(onnx_path):
+            file_size = os.path.getsize(onnx_path)
+            if file_size > min_size:
+                return onnx_path
+            
+            os.remove(onnx_path)
+            print(f"文件大小不足 {min_size} 字节，已删除无效文件")
+        return ""
+
+    except Exception as e:
+        print(f"导出失败: {e}")
+        return ""
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description='Convert PyTorch model to ONNX model.')
+    parser.add_argument('--pthpath', type=str, required=True, help='Path to the PyTorch model file.')
+    parser.add_argument('--onnxpath', type=str, default=None, help='Path to save the ONNX model file.')
+    parser.add_argument('--channel', type=int, default=0, help='Channel parameter.')
+    parser.add_argument('--tilesize', type=int, default=0, help='Tilesize parameter.')
+    parser.add_argument('--fp16', action='store_true', help='Use FP16 precision.')
+    parser.add_argument('--simplify', action='store_true', help='Simplify the ONNX model.')
+    args = parser.parse_args()
+
+    success =  convert_pth_to_onnx(
+        pth_path=args.pth_path,
+        onnx_path=args.onnx_path,
+        channel=args.channel,
+        tilesize=args.tilesize,
+        use_fp16=args.use_fp16,
+        simplify_model=args.simplify_model
+    )
+
+    if success:
+        print("Conversion process finished.")
+    else:
+        print("Conversion process failed.")
+        exit(1)  # Exit with error code

requirements.txt

@@ -0,0 +1,6 @@
+spandrel
+torch
+pnnx
+onnx
+mnn
+gradio