Upload app.py

app.py

@@ -20,6 +20,8 @@ def convert_and_save_onnx(model, name: str, torch_input, out_dir: str, opset: in
         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"},
@@ -76,14 +78,10 @@ def convert_pipeline(model_path: str, opset: int = 17, verify: bool = True, opti
 
     # Generate dummy input
     if static:
-        sr = model_desc.size_requirements
-        min_size = getattr(sr, "size", 32)
-        multiple = getattr(sr, "multiple", 1)
-        h = ((min_size + multiple - 1) // multiple) * multiple
-        w = ((min_size + multiple - 1) // multiple) * multiple
-        torch_input = torch.randn(1, model_desc.input_channels, h, w, device="cpu")
+        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 or 3, 32, 32, device="cpu")
+        torch_input = torch.randn(1, model_desc.input_channels, 32, 32, device="cpu")
 
     out_dir = "./onnx"
     os.makedirs(out_dir, exist_ok=True)
@@ -137,14 +135,19 @@ def load_model(model_path: str):
             'input_channels': model.input_channels,
             'output_channels': model.output_channels,
             'scale': model.scale,
-            #'size_requirements': model.size_requirements,
-            'tags': model.tags
+            'tags': model.tags,
+            'supports_fp16': model.supports_half
+            #'supports_bf16': model.supports_bfloat16,
+            #'size_requirements': model.size_requirements
         }
         
-        return str(architecture_info)
+        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}"
+        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:
@@ -176,15 +179,18 @@ with gr.Blocks(title="PTH to ONNX Converter") as demo:
     install()
     spandrel_extra_arches.install()
 
-    gr.Textbox(value="This space is still under construction. Model outputs may not be correct or valid.", 
-               label="Notice", interactive=False)
+    gr.Textbox(value="This tool is still in development. Outputs may be incorrect.", label="Notice")
+    # with gr.Accordion("Model Cheat Sheet", open=False):
+    #     gr.Markdown("""
+    #     After converting a model, be sure to click the log button above to check the output for any errors or warnings.
+    #     """)
 
     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 - Usually fine to leave false. An FP32 ONNX model can be ran as FP16, but not the other way around.")
-    radio_static = gr.Radio(choices=["True", "False"], value="False", label="Static Shapes - Usually best left false unless the conversion fails or the resulting model doesn't work correctly.")
+    radio_fp16 = gr.Radio(choices=["True", "False"], value="False", label="FP16 - Export at half precision. Not supported by all models.")
+    radio_static = gr.Radio(choices=["True", "False"], value="False", label="Static Shapes - Some models may work better with static shapes. You may need to try both options.")
     radio_slim = gr.Radio(choices=["True", "False"], value="False", label="OnnxSlim - Can potentially optimize the model further, but usually has little effect and can cause the model to not work correctly.")
     
     process_button = gr.Button("Convert", interactive=True)
@@ -200,7 +206,7 @@ with gr.Blocks(title="PTH to ONNX Converter") as demo:
     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)
+    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)