app.py

import os
import gradio as gr
import time
import numpy as np
import onnx
import onnxruntime as ort
import torch
from onnxconverter_common.float16 import convert_float_to_float16
from onnxslim import slim
from spandrel import ModelLoader, ImageModelDescriptor
import spandrel_extra_arches
from rich.traceback import install

def get_out_path(out_dir: str, name: str, opset: int, fp16: bool = False, optimized: bool = False) -> str:
    filename = f"{name}_fp{'16' if fp16 else '32'}_op{opset}{'_onnxslim' if optimized else ''}.onnx"
    return os.path.normpath(os.path.join(out_dir, filename))

def convert_and_save_onnx(model, name: str, torch_input, out_dir: str, opset: int, use_static_shapes: bool) -> tuple[onnx.ModelProto, str]:
    if use_static_shapes:
        dynamic_axes = None
        input_names = None
        output_names = None
        #input_names = ["input"]
        #output_names = ["output"]
    else:
        dynamic_axes = {
            "input": {0: "batch_size", 2: "width", 3: "height"},
            "output": {0: "batch_size", 2: "width", 3: "height"},
        }
        input_names = ["input"]
        output_names = ["output"]

    out_path = get_out_path(out_dir, name, opset, False)

    #if isinstance(model, ImageModelDescriptor):

    #this class was taken from chainner. Running the model through this seems to fix some issues with various arches.
    class FakeModel(torch.nn.Module):
        def __init__(self, model: ImageModelDescriptor):
            super().__init__()
            self.model = model

        def forward(self, x: torch.Tensor):
            return self.model(x)

    model = FakeModel(model)

    torch.onnx.export(
        model,
        (torch_input,),
        out_path,
        dynamo=False,
        verbose=False,
        opset_version=opset,
        dynamic_axes=dynamic_axes,
        input_names=input_names,
        output_names=output_names,
    )
    model_proto = onnx.load(out_path)
    return model_proto, out_path

def verify_onnx(model, torch_input, onnx_path: str) -> None:
    with torch.inference_mode():
        torch_output_np = model(torch_input).cpu().numpy()

    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

    try:
        ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        ort_inputs = {ort_session.get_inputs()[0].name: torch_input.cpu().numpy()}
        onnx_output = ort_session.run(None, ort_inputs)
        np.testing.assert_allclose(
            torch_output_np,
            onnx_output[0],
            rtol=1e-02,
            atol=1e-03,
        )
        print("ONNX output verified against PyTorch output successfully.")
    except AssertionError as e:
        print(f"ONNX verification completed with warnings: {e}")
        gr.Warning("ONNX verification completed with warnings")
    except Exception as e:
        print(f"ONNX verification failed: {e}")
        gr.Warning("ONNX verification failed")


def convert_pipeline(model_path: str, opset: int = 17, verify: bool = True, optimize: bool = True, fp16: bool = False, static: bool = False) -> str:
    loader = ModelLoader()
    model_desc = loader.load_from_file(model_path)
    assert isinstance(model_desc, ImageModelDescriptor)

    model = model_desc.model.to("cpu").eval()
    model_name = os.path.splitext(os.path.basename(model_path))[0]

    # Generate dummy input
    if static:
        height, width = 256, 256
        torch_input = torch.randn(1, model_desc.input_channels, height, width, device="cpu")
    else:
        torch_input = torch.randn(1, model_desc.input_channels, 32, 32, device="cpu")

    out_dir = "./onnx"
    os.makedirs(out_dir, exist_ok=True)

    # Convert to ONNX
    start_time = time.time()
    model_proto, out_path_fp32 = convert_and_save_onnx(
        model, model_name, torch_input, out_dir, opset, static
    )
    out_path = out_path_fp32

    print(f"Saved to {out_path_fp32} in {time.time() - start_time:.2f} seconds.")

    # Verify
    if verify:
        verify_onnx(model, torch_input, out_path_fp32)

    # Optimize
    if optimize:
        model_proto = slim(model_proto)
        session_opt = ort.SessionOptions()
        session_opt.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
        session_opt.optimized_model_filepath = get_out_path(out_dir, model_name, opset, False, True)
        ort.InferenceSession(out_path_fp32, session_opt)
        if verify:
            verify_onnx(model, torch_input, session_opt.optimized_model_filepath)
        model_proto = onnx.load(session_opt.optimized_model_filepath)
        out_path = session_opt.optimized_model_filepath

    # Convert to FP16
    if fp16:
        start_time = time.time()
        out_path = get_out_path(out_dir, model_name, opset, True, optimize)
        model_proto_fp16 = convert_float_to_float16(model_proto)
        onnx.save(model_proto_fp16, out_path)
        print(f"Saved to {out_path_fp32} in {time.time() - start_time:.2f} seconds.")

    return out_path

def load_model(model_path: str):
    if not model_path:
        return "Ready"
    
    loader = ModelLoader()
    try:
        model = loader.load_from_file(model_path)
        assert isinstance(model, ImageModelDescriptor)

        architecture_info = {
            'architecture_name': getattr(model.architecture, 'name', str(model.architecture)),
            'input_channels': model.input_channels,
            'output_channels': model.output_channels,
            'scale': model.scale,
            'tags': model.tags,
            'supports_fp16': model.supports_half
            #'supports_bf16': model.supports_bfloat16,
            #'size_requirements': model.size_requirements
        }
        
        if model.supports_half:
            return [str(architecture_info), gr.Radio(choices=["True", "False"], interactive=True, label="FP16 - Export at half precision. Not supported by all models.")]
        else: 
            return [str(architecture_info), gr.Radio(choices=["True", "False"], value="False", interactive=False, label="FP16 - Export at half precision. Not supported by all models.")]
        
    except Exception as e:
        return [f"Error loading model: {e}", gr.Radio(choices=["True", "False"], interactive=True, label="FP16 - Export at half precision. Not supported by all models.")]

def process_choices(opset, fp16, static, slim, file):
    if not file:
        print("No file loaded.")
        gr.Warning("No file loaded.")
        yield [gr.Button("Convert", interactive=True), gr.DownloadButton(label="💾 Download Converted Model", visible=False)]
        return
    
    # Convert string choices to boolean
    fp16 = fp16 == "True"
    static = static == "True" 
    slim = slim == "True"
    
    yield [gr.Button("Processing", interactive=False), gr.DownloadButton(label="💾 Download Converted Model", visible=False)]
    
    try:
        result = convert_pipeline(file, opset, True, slim, fp16, static)
        short_name = os.path.basename(result)
        yield [gr.Button("Convert", interactive=True), gr.DownloadButton(label=f"💾 {short_name}", value=result, visible=True)]
        return
    except Exception as e:
        print(f"{e}")
        gr.Warning("Conversion error.")
        yield [gr.Button("Convert", interactive=True), gr.DownloadButton(label="💾 Download Converted Model", visible=False)]
        return

# Create Gradio interface
with gr.Blocks(title="PTH to ONNX Converter") as demo:
    install()
    spandrel_extra_arches.install()

    file_upload = gr.File(label="Upload a PyTorch model", file_types=['.pth', '.pt', '.safetensors'])
    metadata = gr.Textbox(value="Ready", label="File Information")
    
    dropdown_opset = gr.Dropdown(choices=[17, 18, 19, 20], value=20, label="Opset")
    radio_fp16 = gr.Radio(choices=["True", "False"], value="False", label="FP16 - Not supported by all models. Not very useful because FP16 TRT engines can still be built from FP32 ONNX models.")
    radio_static = gr.Radio(choices=["True", "False"], value="False", label="Static Shapes - Might be required by some models, but can cause slower performance.")
    radio_slim = gr.Radio(choices=["True", "False"], value="False", label="OnnxSlim - Can cause issues in some models. I have not yet found any cases where it helps. May remove in the future.")
    
    gr.Markdown("After converting, click the logs button at the top to check for any errors or warnings.")
    process_button = gr.Button("Convert", interactive=True)
    file_output = gr.DownloadButton(label="💾 Download Converted Model", visible=False)

    gr.Markdown("""
    # Resources
    - [OpenModelDB](https://openmodeldb.info): Find upscaling models here
    - [VideoJaNai](https://github.com/the-database/VideoJaNai): For upscaling videos using ONNX models
    - [REAL Video Enhancer](https://github.com/TNTwise/REAL-Video-Enhancer): For upscaling videos using ONNX models
    """)
    
    process_button.click(fn=process_choices, 
                        inputs=[dropdown_opset, radio_fp16, radio_static, radio_slim, file_upload], 
                        outputs=[process_button, file_output])
    file_upload.upload(fn=load_model, inputs=file_upload, outputs=[metadata, radio_fp16])

if __name__ == "__main__":
    demo.launch(show_error=True, inbrowser=True, show_api=False, debug=False)