From b3665630edf2de8dd51ff58148ec45082fc030bd Mon Sep 17 00:00:00 2001
From: Oleksandr Bezdieniezhnykh <zxsanny@gmail.com>
Date: Wed, 3 Sep 2025 16:36:26 +0300
Subject: [PATCH] add warning ai status, fix ai availability handling 1

---
 Azaion.Annotator/Annotator.xaml               |  1 +
 Azaion.Annotator/AnnotatorEventHandler.cs     |  2 ++
 Azaion.Common/DTO/AIAvailabilityStatus.cs     |  2 ++
 .../Services/Inference/InferenceService.cs    | 32 +++++++++++++------
 Azaion.Inference/ai_availability_status.pyx   |  1 +
 Azaion.Inference/inference.pyx                |  5 +--
 6 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/Azaion.Annotator/Annotator.xaml b/Azaion.Annotator/Annotator.xaml
index ab57de6..ac7c451 100644
--- a/Azaion.Annotator/Annotator.xaml
+++ b/Azaion.Annotator/Annotator.xaml
@@ -496,6 +496,7 @@
 
             <Button
                 x:Name="AIDetectBtn"
+                IsEnabled="False"
                 Grid.Column="10"
                 Padding="2" Width="25"
                 Height="25"
diff --git a/Azaion.Annotator/AnnotatorEventHandler.cs b/Azaion.Annotator/AnnotatorEventHandler.cs
index 2dd6e57..3973894 100644
--- a/Azaion.Annotator/AnnotatorEventHandler.cs
+++ b/Azaion.Annotator/AnnotatorEventHandler.cs
@@ -483,5 +483,7 @@ public class AnnotatorEventHandler(
             mainWindow.AIDetectBtn.IsEnabled = e.Status == AIAvailabilityEnum.Enabled;
             mainWindow.StatusHelp.Text = e.ToString();
         });
+1        if (e.Status is AIAvailabilityEnum.Enabled or AIAvailabilityEnum.Error)
+            await inferenceService.CheckAIAvailabilityTokenSource.CancelAsync();
     }
 }
\ No newline at end of file
diff --git a/Azaion.Common/DTO/AIAvailabilityStatus.cs b/Azaion.Common/DTO/AIAvailabilityStatus.cs
index b45f865..717165c 100644
--- a/Azaion.Common/DTO/AIAvailabilityStatus.cs
+++ b/Azaion.Common/DTO/AIAvailabilityStatus.cs
@@ -10,6 +10,7 @@ public enum AIAvailabilityEnum
     Converting = 20,
     Uploading = 30,
     Enabled = 200,
+    Warning = 300,
     Error = 500
 }
 
@@ -27,6 +28,7 @@ public class AIAvailabilityStatusEvent : INotification
         { AIAvailabilityEnum.Converting, "Йде налаштування AI під Ваше залізо. (5-12 хвилин в залежності від моделі відеокарти, до 50 хв на старих GTX1650)" },
         { AIAvailabilityEnum.Uploading, "Йде зберігання AI" },
         { AIAvailabilityEnum.Enabled, "AI готовий для розпізнавання" },
+        { AIAvailabilityEnum.Warning, "Неможливо запустити AI наразі, йде налаштування під Ваше залізо" },
         { AIAvailabilityEnum.Error, "Помилка під час налаштування AI" }
     };
 
diff --git a/Azaion.Common/Services/Inference/InferenceService.cs b/Azaion.Common/Services/Inference/InferenceService.cs
index 7a816b1..57357ac 100644
--- a/Azaion.Common/Services/Inference/InferenceService.cs
+++ b/Azaion.Common/Services/Inference/InferenceService.cs
@@ -9,25 +9,37 @@ public interface IInferenceService
 {
     Task RunInference(List<string> mediaPaths, int tileSize, CancellationToken ct = default);
     CancellationTokenSource InferenceCancelTokenSource { get; set; }
+    CancellationTokenSource CheckAIAvailabilityTokenSource { get; set; }
     void StopInference();
 }
 
 // SHOULD BE ONLY ONE INSTANCE OF InferenceService. Do not add ANY NotificationHandler to it!
 // _inferenceCancelTokenSource should be created only once.
-public class InferenceService(
-    IInferenceClient client,
-    IAzaionApi azaionApi,
-    IOptions<AIRecognitionConfig> aiConfigOptions) : IInferenceService
+public class InferenceService : IInferenceService
 {
+    private readonly IInferenceClient _client;
+    private readonly IAzaionApi _azaionApi;
+    private readonly IOptions<AIRecognitionConfig> _aiConfigOptions;
+
+    public InferenceService(IInferenceClient client,
+        IAzaionApi azaionApi,
+        IOptions<AIRecognitionConfig> aiConfigOptions)
+    {
+        _client = client;
+        _azaionApi = azaionApi;
+        _aiConfigOptions = aiConfigOptions;
+        _ = Task.Run(async () => await CheckAIAvailabilityStatus());
+    }
+
     public CancellationTokenSource InferenceCancelTokenSource { get; set; } = new();
     public CancellationTokenSource CheckAIAvailabilityTokenSource { get; set; } = new();
 
-    public async Task CheckAIAvailabilityStatus()
+    private async Task CheckAIAvailabilityStatus()
     {
         CheckAIAvailabilityTokenSource = new CancellationTokenSource();
         while (!CheckAIAvailabilityTokenSource.IsCancellationRequested)
         {
-            client.Send(RemoteCommand.Create(CommandType.AIAvailabilityCheck));
+            _client.Send(RemoteCommand.Create(CommandType.AIAvailabilityCheck));
             await Task.Delay(10000, CheckAIAvailabilityTokenSource.Token);
         }
     }
@@ -35,16 +47,16 @@ public class InferenceService(
     public async Task RunInference(List<string> mediaPaths, int tileSize, CancellationToken ct = default)
     {
         InferenceCancelTokenSource = new CancellationTokenSource();
-        client.Send(RemoteCommand.Create(CommandType.Login, azaionApi.Credentials));
+        _client.Send(RemoteCommand.Create(CommandType.Login, _azaionApi.Credentials));
 
-        var aiConfig = aiConfigOptions.Value;
+        var aiConfig = _aiConfigOptions.Value;
         aiConfig.Paths = mediaPaths;
         aiConfig.TileSize = tileSize;
-        client.Send(RemoteCommand.Create(CommandType.Inference, aiConfig));
+        _client.Send(RemoteCommand.Create(CommandType.Inference, aiConfig));
 
         using var combinedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(ct, InferenceCancelTokenSource.Token);
         await combinedTokenSource.Token.AsTask();
     }
 
-    public void StopInference() => client.Stop();
+    public void StopInference() => _client.Stop();
 }
\ No newline at end of file
diff --git a/Azaion.Inference/ai_availability_status.pyx b/Azaion.Inference/ai_availability_status.pyx
index d467682..9f2e4d0 100644
--- a/Azaion.Inference/ai_availability_status.pyx
+++ b/Azaion.Inference/ai_availability_status.pyx
@@ -7,6 +7,7 @@ AIStatus2Text = {
     AIAvailabilityEnum.CONVERTING: "Converting",
     AIAvailabilityEnum.UPLOADING: "Uploading",
     AIAvailabilityEnum.ENABLED: "Enabled",
+    AIAvailabilityEnum.WARNING: "Warning",
     AIAvailabilityEnum.ERROR: "Error",
 }
 
diff --git a/Azaion.Inference/inference.pyx b/Azaion.Inference/inference.pyx
index 668dd55..34c8f22 100644
--- a/Azaion.Inference/inference.pyx
+++ b/Azaion.Inference/inference.pyx
@@ -91,7 +91,7 @@ cdef class Inference:
                     self.engine = TensorRTEngine(res.data)
                     self.ai_availability_status.set_status(AIAvailabilityEnum.ENABLED)
                 except Exception as e:
-                    self.ai_availability_status.set_status(AIAvailabilityEnum.ERROR, <str>str(e))
+                    self.ai_availability_status.set_status(AIAvailabilityEnum.WARNING, <str>str(e))
                     onnx_engine_bytes = self.get_onnx_engine_bytes()
                     self.ai_availability_status.set_status(AIAvailabilityEnum.CONVERTING)
                     model_bytes = TensorRTEngine.convert_from_onnx(onnx_engine_bytes)
@@ -100,7 +100,8 @@ cdef class Inference:
                     res = self.loader_client.upload_big_small_resource(model_bytes, <str> engine_filename, models_dir)
                     if res.err is not None:
                         self.ai_availability_status.set_status(AIAvailabilityEnum.ERROR, res.err)
-                    self.ai_availability_status.set_status(AIAvailabilityEnum.ENABLED)
+                    else:
+                        self.ai_availability_status.set_status(AIAvailabilityEnum.ENABLED)
             else:
                 self.engine = OnnxEngine(<bytes>self.get_onnx_engine_bytes())
             self.is_building_engine = False