using System; using System.Collections.Generic; using System.IO; using System.Net.Http; using System.Threading; using System.Threading.Tasks; using AIImages.Models; using StableDiffusionNet; using StableDiffusionNet.Exceptions; using StableDiffusionNet.Interfaces; using StableDiffusionNet.Models.Requests; using Verse; namespace AIImages.Services { /// /// Адаптер для Stable Diffusion API (AUTOMATIC1111 WebUI) /// Использует библиотеку StableDiffusionNet для работы с API /// public class StableDiffusionNetAdapter : IStableDiffusionApiService, IDisposable { private readonly IStableDiffusionClient _client; 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) ); } // Определяем путь для сохранения _saveFolderPath = Path.Combine(GenFilePaths.SaveDataFolderPath, savePath); // Создаем папку, если не существует if (!Directory.Exists(_saveFolderPath)) { Directory.CreateDirectory(_saveFolderPath); } // Создаем клиент StableDiffusion используя Builder _client = new StableDiffusionClientBuilder() .WithBaseUrl(apiEndpoint) .WithTimeout(300) // 5 минут в секундах .WithRetry(retryCount: 3, retryDelayMilliseconds: 1000) .Build(); Log.Message( $"[AI Images] StableDiffusion adapter initialized with endpoint: {apiEndpoint}" ); } public async Task GenerateImageAsync( GenerationRequest request, CancellationToken cancellationToken = default ) { ThrowIfDisposed(); if (request == null) { return GenerationResult.Failure("Request cannot be null"); } try { Log.Message( $"[AI Images] Starting image generation with prompt: {request.Prompt.Substring(0, Math.Min(50, request.Prompt.Length))}..." ); // Маппируем наш запрос на запрос библиотеки StableDiffusionNet var sdRequest = new TextToImageRequest { Prompt = request.Prompt, NegativePrompt = request.NegativePrompt, Steps = request.Steps, CfgScale = request.CfgScale, Width = request.Width, Height = request.Height, SamplerName = request.Sampler, Scheduler = request.Scheduler, Seed = request.Seed, SaveImages = request.SaveImagesToServer, // Сохранять ли изображения на сервере }; // Выполняем запрос через библиотеку (с встроенной retry логикой) var response = await _client.TextToImage.GenerateAsync( sdRequest, cancellationToken ); if (response?.Images == null || response.Images.Count == 0) { return GenerationResult.Failure("No images returned from API"); } // Декодируем изображение из base64 byte[] imageData = Convert.FromBase64String(response.Images[0]); // Сохраняем изображение string fileName = $"pawn_{DateTime.Now:yyyyMMdd_HHmmss}.png"; string fullPath = Path.Combine(_saveFolderPath, fileName); await File.WriteAllBytesAsync(fullPath, imageData); Log.Message($"[AI Images] Image generated successfully and saved to: {fullPath}"); return GenerationResult.SuccessResult(imageData, fullPath, request); } 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}"); return GenerationResult.Failure($"Connection error: {ex.Message}"); } catch (StableDiffusionException ex) { Log.Error($"[AI Images] StableDiffusion API error: {ex.Message}"); return GenerationResult.Failure($"API Error: {ex.Message}"); } catch (Exception ex) { Log.Error($"[AI Images] Unexpected error: {ex.Message}\n{ex.StackTrace}"); return GenerationResult.Failure($"Error: {ex.Message}"); } } public async Task GetProgressAsync( CancellationToken cancellationToken = default ) { ThrowIfDisposed(); try { // Используем Progress сервис библиотеки var progress = await _client.Progress.GetProgressAsync(cancellationToken); // Маппируем на наш тип return new GenerationProgress { Progress = progress.Progress, CurrentStep = progress.State?.SamplingStep ?? 0, TotalSteps = progress.State?.SamplingSteps ?? 0, EtaRelative = progress.EtaRelative, IsActive = progress.Progress > 0 && progress.Progress < 1.0, }; } catch (Exception ex) { Log.Warning($"[AI Images] Failed to get progress: {ex.Message}"); // Возвращаем пустой прогресс при ошибке return new GenerationProgress { Progress = 0, CurrentStep = 0, TotalSteps = 0, EtaRelative = 0, IsActive = false, }; } } public async Task CheckApiAvailability( string apiEndpoint, CancellationToken cancellationToken = default ) { ThrowIfDisposed(); try { // Используем встроенный метод PingAsync библиотеки return await _client.PingAsync(cancellationToken); } catch (Exception ex) { Log.Warning($"[AI Images] API check failed: {ex.Message}"); return false; } } public async Task> GetAvailableModels( string apiEndpoint, CancellationToken cancellationToken = default ) { ThrowIfDisposed(); try { // Используем Models сервис библиотеки var models = await _client.Models.GetModelsAsync(cancellationToken); var modelNames = new List(); if (models != null) { foreach (var model in models) { // Используем Title или ModelName в зависимости от того, что доступно modelNames.Add(model.Title ?? model.ModelName); } } Log.Message($"[AI Images] Found {modelNames.Count} models"); return modelNames; } catch (Exception ex) { Log.Error($"[AI Images] Failed to load models: {ex.Message}"); return new List(); } } public async Task> GetAvailableSamplers( string apiEndpoint, CancellationToken cancellationToken = default ) { ThrowIfDisposed(); try { // Используем Samplers сервис библиотеки var samplers = await _client.Samplers.GetSamplersAsync(cancellationToken); var samplerNames = new List(); if (samplers != null) { samplerNames.AddRange(samplers); } Log.Message($"[AI Images] Found {samplerNames.Count} samplers"); return samplerNames; } catch (Exception ex) { Log.Warning($"[AI Images] Failed to load samplers: {ex.Message}"); return GetDefaultSamplers(); } } public async Task> GetAvailableSchedulers( string apiEndpoint, CancellationToken cancellationToken = default ) { ThrowIfDisposed(); try { // Используем Schedulers сервис библиотеки (доступен с версии 1.1.1) var schedulers = await _client.Schedulers.GetSchedulersAsync(cancellationToken); var schedulerNames = new List(); if (schedulers != null) { schedulerNames.AddRange(schedulers); } Log.Message($"[AI Images] Found {schedulerNames.Count} schedulers"); return schedulerNames; } catch (Exception ex) { Log.Warning($"[AI Images] Failed to load schedulers: {ex.Message}"); return GetDefaultSchedulers(); } } private List GetDefaultSamplers() { return new List { "Euler a", "Euler", "LMS", "Heun", "DPM2", "DPM2 a", "DPM++ 2S a", "DPM++ 2M", "DPM++ SDE", "DPM fast", "DPM adaptive", "LMS Karras", "DPM2 Karras", "DPM2 a Karras", "DPM++ 2S a Karras", "DPM++ 2M Karras", "DPM++ SDE Karras", "DDIM", "PLMS", }; } private List GetDefaultSchedulers() { return new List { "Automatic", "Uniform", "Karras", "Exponential", "Polyexponential", "SGM Uniform", }; } private void ThrowIfDisposed() { if (_disposed) { throw new ObjectDisposedException(nameof(StableDiffusionNetAdapter)); } } public void Dispose() { if (_disposed) return; // Dispose клиента StableDiffusion if (_client is IDisposable disposableClient) { disposableClient.Dispose(); } _disposed = true; } } }