Enhance AIImages mod by adding cancellation support for image generation, improving user experience with localized strings for cancellation actions in English and Russian. Refactor service integration for better dependency management and update AIImages.dll to reflect these changes.

2025-10-26 19:10:45 +03:00
parent 3434927342
commit 02b0143186
11 changed files with 974 additions and 174 deletions
--- a/Source/AIImages/Services/IStableDiffusionApiService.cs
+++ b/Source/AIImages/Services/IStableDiffusionApiService.cs
@@ -1,4 +1,5 @@
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using AIImages.Models;

@@ -12,26 +13,41 @@ namespace AIImages.Services
        /// <summary>
        /// Генерирует изображение на основе запроса
        /// </summary>
-        Task<GenerationResult> GenerateImageAsync(GenerationRequest request);
+        Task<GenerationResult> GenerateImageAsync(
+            GenerationRequest request,
+            CancellationToken cancellationToken = default
+        );

        /// <summary>
        /// Проверяет доступность API
        /// </summary>
-        Task<bool> CheckApiAvailability(string apiEndpoint);
+        Task<bool> CheckApiAvailability(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        );

        /// <summary>
        /// Получает список доступных моделей с API
        /// </summary>
-        Task<List<string>> GetAvailableModels(string apiEndpoint);
+        Task<List<string>> GetAvailableModels(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        );

        /// <summary>
        /// Получает список доступных сэмплеров
        /// </summary>
-        Task<List<string>> GetAvailableSamplers(string apiEndpoint);
+        Task<List<string>> GetAvailableSamplers(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        );

        /// <summary>
        /// Получает список доступных schedulers
        /// </summary>
-        Task<List<string>> GetAvailableSchedulers(string apiEndpoint);
+        Task<List<string>> GetAvailableSchedulers(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        );
    }
 }
--- a/Source/AIImages/Services/ServiceContainer.cs
+++ b/Source/AIImages/Services/ServiceContainer.cs
@@ -0,0 +1,83 @@
+using System;
+using AIImages.Settings;
+using Verse;
+
+namespace AIImages.Services
+{
+    /// <summary>
+    /// Простой DI контейнер для управления зависимостями мода
+    /// Используется вместо статического Service Locator паттерна
+    /// </summary>
+    public class ServiceContainer : IDisposable
+    {
+        private bool _disposed;
+
+        // Сервисы
+        public IPawnDescriptionService PawnDescriptionService { get; }
+        public IPromptGeneratorService PromptGeneratorService { get; }
+        public IStableDiffusionApiService ApiService { get; private set; }
+
+        // Настройки
+        public AIImagesModSettings Settings { get; }
+
+        public ServiceContainer(AIImagesModSettings settings)
+        {
+            if (settings == null)
+            {
+                throw new ArgumentNullException(nameof(settings));
+            }
+
+            Settings = settings;
+
+            // Инициализируем сервисы с внедрением зависимостей
+            PawnDescriptionService = new PawnDescriptionService();
+            PromptGeneratorService = new AdvancedPromptGenerator();
+
+            // Создаем API сервис с текущими настройками
+            RecreateApiService();
+
+            Log.Message("[AI Images] ServiceContainer initialized successfully");
+        }
+
+        /// <summary>
+        /// Пересоздает API сервис с новыми настройками (например, когда изменился endpoint)
+        /// </summary>
+        public void RecreateApiService()
+        {
+            // Освобождаем старый сервис, если он был
+            if (ApiService is IDisposable disposable)
+            {
+                disposable.Dispose();
+            }
+
+            // Создаем новый с актуальными настройками
+            ApiService = new StableDiffusionNetAdapter(Settings.apiEndpoint, Settings.savePath);
+
+            Log.Message($"[AI Images] API service recreated with endpoint: {Settings.apiEndpoint}");
+        }
+
+        /// <summary>
+        /// Проверяет, нужно ли пересоздать API сервис (например, если изменился endpoint)
+        /// </summary>
+        public bool ShouldRecreateApiService(string newEndpoint)
+        {
+            // Можно добавить более сложную логику проверки
+            return Settings.apiEndpoint != newEndpoint;
+        }
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            // Освобождаем ресурсы API сервиса
+            if (ApiService is IDisposable disposable)
+            {
+                disposable.Dispose();
+            }
+
+            _disposed = true;
+            Log.Message("[AI Images] ServiceContainer disposed");
+        }
+    }
+}
--- a/Source/AIImages/Services/StableDiffusionNetAdapter.cs
+++ b/Source/AIImages/Services/StableDiffusionNetAdapter.cs
@@ -1,8 +1,10 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Net.Http;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using AIImages.Models;
 using Newtonsoft.Json;
@@ -11,29 +13,60 @@ using Verse;
 namespace AIImages.Services
 {
    /// <summary>
-    /// Сервис для работы с Stable Diffusion API (AUTOMATIC1111 WebUI)
+    /// Адаптер для Stable Diffusion API (AUTOMATIC1111 WebUI)
+    /// TODO: В будущем можно мигрировать на библиотеку StableDiffusionNet когда API будет полностью совместимо
    /// </summary>
-    public class StableDiffusionApiService : IStableDiffusionApiService
+    public class StableDiffusionNetAdapter : IStableDiffusionApiService, IDisposable
    {
-        private readonly HttpClient httpClient;
-        private readonly string saveFolderPath;
-
-        public StableDiffusionApiService(string savePath = "AIImages/Generated")
+        // Shared HttpClient для предотвращения socket exhaustion
+        // См: https://learn.microsoft.com/en-us/dotnet/fundamentals/networking/http/httpclient-guidelines
+        private static readonly HttpClient _sharedHttpClient = new HttpClient
        {
-            httpClient = new HttpClient { Timeout = TimeSpan.FromMinutes(5) };
+            Timeout = TimeSpan.FromMinutes(5),
+        };
+
+        private readonly string _apiEndpoint;
+        private readonly string _saveFolderPath;
+        private bool _disposed;
+
+        public StableDiffusionNetAdapter(string apiEndpoint, string savePath = "AIImages/Generated")
+        {
+            if (string.IsNullOrEmpty(apiEndpoint))
+            {
+                throw new ArgumentException(
+                    "API endpoint cannot be null or empty",
+                    nameof(apiEndpoint)
+                );
+            }
+
+            _apiEndpoint = apiEndpoint;

            // Определяем путь для сохранения
-            saveFolderPath = Path.Combine(GenFilePaths.SaveDataFolderPath, savePath);
+            _saveFolderPath = Path.Combine(GenFilePaths.SaveDataFolderPath, savePath);

            // Создаем папку, если не существует
-            if (!Directory.Exists(saveFolderPath))
+            if (!Directory.Exists(_saveFolderPath))
            {
-                Directory.CreateDirectory(saveFolderPath);
+                Directory.CreateDirectory(_saveFolderPath);
            }
+
+            Log.Message(
+                $"[AI Images] StableDiffusion adapter initialized with endpoint: {apiEndpoint}"
+            );
        }

-        public async Task<GenerationResult> GenerateImageAsync(GenerationRequest request)
+        public async Task<GenerationResult> GenerateImageAsync(
+            GenerationRequest request,
+            CancellationToken cancellationToken = default
+        )
        {
+            ThrowIfDisposed();
+
+            if (request == null)
+            {
+                return GenerationResult.Failure("Request cannot be null");
+            }
+
            try
            {
                Log.Message(
@@ -59,9 +92,13 @@ namespace AIImages.Services
                string jsonRequest = JsonConvert.SerializeObject(apiRequest);
                var content = new StringContent(jsonRequest, Encoding.UTF8, "application/json");

-                // Отправляем запрос
-                string endpoint = $"{request.Model}/sdapi/v1/txt2img";
-                HttpResponseMessage response = await httpClient.PostAsync(endpoint, content);
+                // Отправляем запрос с поддержкой cancellation
+                string endpoint = $"{_apiEndpoint}/sdapi/v1/txt2img";
+                HttpResponseMessage response = await _sharedHttpClient.PostAsync(
+                    endpoint,
+                    content,
+                    cancellationToken
+                );

                if (!response.IsSuccessStatusCode)
                {
@@ -85,7 +122,7 @@ namespace AIImages.Services

                // Сохраняем изображение
                string fileName = $"pawn_{DateTime.Now:yyyyMMdd_HHmmss}.png";
-                string fullPath = Path.Combine(saveFolderPath, fileName);
+                string fullPath = Path.Combine(_saveFolderPath, fileName);
                await File.WriteAllBytesAsync(fullPath, imageData);

                Log.Message($"[AI Images] Image generated successfully and saved to: {fullPath}");
@@ -94,8 +131,14 @@ namespace AIImages.Services
            }
            catch (TaskCanceledException)
            {
+                Log.Warning("[AI Images] Request timeout. Generation took too long.");
                return GenerationResult.Failure("Request timeout. Generation took too long.");
            }
+            catch (OperationCanceledException)
+            {
+                Log.Message("[AI Images] Image generation was cancelled");
+                return GenerationResult.Failure("Generation cancelled");
+            }
            catch (HttpRequestException ex)
            {
                Log.Error($"[AI Images] HTTP error: {ex.Message}");
@@ -108,12 +151,20 @@ namespace AIImages.Services
            }
        }

-        public async Task<bool> CheckApiAvailability(string apiEndpoint)
+        public async Task<bool> CheckApiAvailability(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        )
        {
+            ThrowIfDisposed();
+
            try
            {
                string endpoint = $"{apiEndpoint}/sdapi/v1/sd-models";
-                HttpResponseMessage response = await httpClient.GetAsync(endpoint);
+                HttpResponseMessage response = await _sharedHttpClient.GetAsync(
+                    endpoint,
+                    cancellationToken
+                );
                return response.IsSuccessStatusCode;
            }
            catch (Exception ex)
@@ -123,12 +174,20 @@ namespace AIImages.Services
            }
        }

-        public async Task<List<string>> GetAvailableModels(string apiEndpoint)
+        public async Task<List<string>> GetAvailableModels(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        )
        {
+            ThrowIfDisposed();
+
            try
            {
                string endpoint = $"{apiEndpoint}/sdapi/v1/sd-models";
-                HttpResponseMessage response = await httpClient.GetAsync(endpoint);
+                HttpResponseMessage response = await _sharedHttpClient.GetAsync(
+                    endpoint,
+                    cancellationToken
+                );

                if (!response.IsSuccessStatusCode)
                    return new List<string>();
@@ -155,12 +214,20 @@ namespace AIImages.Services
            }
        }

-        public async Task<List<string>> GetAvailableSamplers(string apiEndpoint)
+        public async Task<List<string>> GetAvailableSamplers(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        )
        {
+            ThrowIfDisposed();
+
            try
            {
                string endpoint = $"{apiEndpoint}/sdapi/v1/samplers";
-                HttpResponseMessage response = await httpClient.GetAsync(endpoint);
+                HttpResponseMessage response = await _sharedHttpClient.GetAsync(
+                    endpoint,
+                    cancellationToken
+                );

                if (!response.IsSuccessStatusCode)
                    return GetDefaultSamplers();
@@ -187,12 +254,20 @@ namespace AIImages.Services
            }
        }

-        public async Task<List<string>> GetAvailableSchedulers(string apiEndpoint)
+        public async Task<List<string>> GetAvailableSchedulers(
+            string apiEndpoint,
+            CancellationToken cancellationToken = default
+        )
        {
+            ThrowIfDisposed();
+
            try
            {
                string endpoint = $"{apiEndpoint}/sdapi/v1/schedulers";
-                HttpResponseMessage response = await httpClient.GetAsync(endpoint);
+                HttpResponseMessage response = await _sharedHttpClient.GetAsync(
+                    endpoint,
+                    cancellationToken
+                );

                if (!response.IsSuccessStatusCode)
                    return GetDefaultSchedulers();
@@ -258,8 +333,25 @@ namespace AIImages.Services
            };
        }

+        private void ThrowIfDisposed()
+        {
+            if (_disposed)
+            {
+                throw new ObjectDisposedException(nameof(StableDiffusionNetAdapter));
+            }
+        }
+
+        public void Dispose()
+        {
+            if (_disposed)
+                return;
+
+            // Не dispose shared HttpClient - он используется глобально
+            _disposed = true;
+        }
+
        // Вспомогательные классы для десериализации JSON ответов
-#pragma warning disable S3459, S1144 // Properties set by JSON deserializer
+#pragma warning disable S3459, S1144, IDE1006 // Properties set by JSON deserializer
        private sealed class Txt2ImgResponse
        {
            public string[] images { get; set; }
@@ -280,6 +372,6 @@ namespace AIImages.Services
        {
            public string name { get; set; }
        }
-#pragma warning restore S3459, S1144
+#pragma warning restore S3459, S1144, IDE1006
    }
 }