From 8115224846a638bab9078f830ba538b54d3330f5 Mon Sep 17 00:00:00 2001
From: malvm <malcevvadim261@gmail.com>
Date: Mon, 20 Jul 2026 16:23:48 +0400
Subject: [PATCH] fix weights conversion

---
 src/yolo_webui/export_runner.py  | 10 ++++++++--
 src/yolo_webui/static/app.js     |  5 ++++-
 src/yolo_webui/static/index.html | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/yolo_webui/export_runner.py b/src/yolo_webui/export_runner.py
index 31693e3..66cb388 100644
--- a/src/yolo_webui/export_runner.py
+++ b/src/yolo_webui/export_runner.py
@@ -34,6 +34,9 @@ def main() -> int:
     half = config_dict.get("half", False)
     int8 = config_dict.get("int8", False)
     dynamic = config_dict.get("dynamic", False)
+    simplify = config_dict.get("simplify", False)
+    batch = config_dict.get("batch", 1)
+    workspace = config_dict.get("workspace", 4)
 
     if not model_path or not Path(model_path).exists():
         print(f"Model file not found: {model_path}")
@@ -46,14 +49,17 @@ def main() -> int:
         model = YOLO(model_path)
         
         print(f"Запуск экспорта в формат {export_format}...", flush=True)
-        print(f"Параметры: imgsz={imgsz}, half={half}, int8={int8}, dynamic={dynamic}", flush=True)
+        print(f"Параметры: imgsz={imgsz}, half={half}, int8={int8}, dynamic={dynamic}, simplify={simplify}, batch={batch}, workspace={workspace}", flush=True)
         
         exported_path = model.export(
             format=export_format,
             imgsz=imgsz,
             half=half,
             int8=int8,
-            dynamic=dynamic
+            dynamic=dynamic,
+            simplify=simplify,
+            batch=batch,
+            workspace=workspace
         )
         
         print(f"Экспорт завершен успешно.", flush=True)
diff --git a/src/yolo_webui/static/app.js b/src/yolo_webui/static/app.js
index 23f99b2..5bcc4f1 100644
--- a/src/yolo_webui/static/app.js
+++ b/src/yolo_webui/static/app.js
@@ -1124,7 +1124,10 @@ document.addEventListener('DOMContentLoaded', () => {
             imgsz: readNumber('export-imgsz', 640, true),
             half: document.getElementById('export-half').checked,
             int8: document.getElementById('export-int8').checked,
-            dynamic: document.getElementById('export-dynamic').checked
+            dynamic: document.getElementById('export-dynamic').checked,
+            simplify: document.getElementById('export-simplify').checked,
+            batch: readNumber('export-batch', 1, true),
+            workspace: readNumber('export-workspace', 4, true)
         };
 
         try {
diff --git a/src/yolo_webui/static/index.html b/src/yolo_webui/static/index.html
index 7b3578f..0a41920 100644
--- a/src/yolo_webui/static/index.html
+++ b/src/yolo_webui/static/index.html
@@ -432,6 +432,7 @@
                                 <option value="onnx">ONNX (.onnx)</option>
                                 <option value="engine">TensorRT (.engine)</option>
                                 <option value="openvino">OpenVINO (_openvino_model/)</option>
+                                <option value="triton">Triton Inference Server (triton/)</option>
                                 <option value="torchscript">TorchScript (.torchscript)</option>
                                 <option value="coreml">CoreML (.mlpackage)</option>
                                 <option value="tflite">TFLite (.tflite)</option>
@@ -445,6 +446,14 @@
                                 <label for="export-imgsz">Размер изображения (imgsz)</label>
                                 <input type="number" id="export-imgsz" name="export-imgsz" value="640">
                             </div>
+                            <div class="field">
+                                <label for="export-batch" title="Максимальный размер батча. Важен при Dynamic=True для TensorRT и Triton.">Макс. батч (batch)</label>
+                                <input type="number" id="export-batch" name="export-batch" value="1" min="1">
+                            </div>
+                            <div class="field">
+                                <label for="export-workspace" title="Размер памяти для сборки TensorRT (ГБ).">Workspace (ГБ)</label>
+                                <input type="number" id="export-workspace" name="export-workspace" value="4" min="1" max="64">
+                            </div>
                         </div>
                         <div class="row">
                             <div class="toggle-row" style="padding: 10px 0;">
@@ -477,6 +486,16 @@
                                     <span class="switch-slider"></span>
                                 </label>
                             </div>
+                            <div class="toggle-row" style="padding: 10px 0;">
+                                <div class="toggle-label" title="Только для ONNX. Оптимизирует и упрощает граф модели.">
+                                    <h3>Simplify</h3>
+                                    <p>Упростить граф (ONNX)</p>
+                                </div>
+                                <label class="switch-container">
+                                    <input type="checkbox" id="export-simplify">
+                                    <span class="switch-slider"></span>
+                                </label>
+                            </div>
                         </div>
                     </div>
                 </form>