Update AIImages mod to utilize StableDiffusionNet.Core for API interactions, enhancing image generation capabilities. Refactor service methods for improved clarity and efficiency, and update AIImages.dll to reflect these changes.

This commit is contained in:
Leonid Pershin
2025-10-26 22:29:27 +03:00
parent 02b0143186
commit 6e6e92df53
4 changed files with 54 additions and 123 deletions

View File

@@ -18,6 +18,6 @@
<ItemGroup>
<PackageReference Include="Krafs.Rimworld.Ref" Version="1.6.4566" />
<PackageReference Include="Lib.Harmony" Version="2.4.1" />
<PackageReference Include="StableDiffusionNet" Version="1.0.1" />
<PackageReference Include="StableDiffusionNet.Core" Version="1.1.1" />
</ItemGroup>
</Project>

View File

@@ -1,31 +1,25 @@
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;
using StableDiffusionNet;
using StableDiffusionNet.Exceptions;
using StableDiffusionNet.Interfaces;
using StableDiffusionNet.Models.Requests;
using Verse;
namespace AIImages.Services
{
/// <summary>
/// Адаптер для Stable Diffusion API (AUTOMATIC1111 WebUI)
/// TODO: В будущем можно мигрировать на библиотеку StableDiffusionNet когда API будет полностью совместимо
/// Использует библиотеку StableDiffusionNet для работы с API
/// </summary>
public class StableDiffusionNetAdapter : IStableDiffusionApiService, IDisposable
{
// Shared HttpClient для предотвращения socket exhaustion
// См: https://learn.microsoft.com/en-us/dotnet/fundamentals/networking/http/httpclient-guidelines
private static readonly HttpClient _sharedHttpClient = new HttpClient
{
Timeout = TimeSpan.FromMinutes(5),
};
private readonly string _apiEndpoint;
private readonly IStableDiffusionClient _client;
private readonly string _saveFolderPath;
private bool _disposed;
@@ -39,8 +33,6 @@ namespace AIImages.Services
);
}
_apiEndpoint = apiEndpoint;
// Определяем путь для сохранения
_saveFolderPath = Path.Combine(GenFilePaths.SaveDataFolderPath, savePath);
@@ -50,6 +42,13 @@ namespace AIImages.Services
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}"
);
@@ -73,52 +72,34 @@ namespace AIImages.Services
$"[AI Images] Starting image generation with prompt: {request.Prompt.Substring(0, Math.Min(50, request.Prompt.Length))}..."
);
// Формируем JSON запрос для AUTOMATIC1111 API
var apiRequest = new
// Маппируем наш запрос на запрос библиотеки StableDiffusionNet
var sdRequest = new TextToImageRequest
{
prompt = request.Prompt,
negative_prompt = request.NegativePrompt,
steps = request.Steps,
cfg_scale = request.CfgScale,
width = request.Width,
height = request.Height,
sampler_name = request.Sampler,
scheduler = request.Scheduler,
seed = request.Seed,
save_images = false,
send_images = true,
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 и SendImages не нужны - библиотека всегда возвращает изображения
};
string jsonRequest = JsonConvert.SerializeObject(apiRequest);
var content = new StringContent(jsonRequest, Encoding.UTF8, "application/json");
// Отправляем запрос с поддержкой cancellation
string endpoint = $"{_apiEndpoint}/sdapi/v1/txt2img";
HttpResponseMessage response = await _sharedHttpClient.PostAsync(
endpoint,
content,
// Выполняем запрос через библиотеку (с встроенной retry логикой)
var response = await _client.TextToImage.GenerateAsync(
sdRequest,
cancellationToken
);
if (!response.IsSuccessStatusCode)
{
string errorContent = await response.Content.ReadAsStringAsync();
Log.Error(
$"[AI Images] API request failed: {response.StatusCode} - {errorContent}"
);
return GenerationResult.Failure($"API Error: {response.StatusCode}");
}
string jsonResponse = await response.Content.ReadAsStringAsync();
var apiResponse = JsonConvert.DeserializeObject<Txt2ImgResponse>(jsonResponse);
if (apiResponse?.images == null || apiResponse.images.Length == 0)
if (response?.Images == null || response.Images.Count == 0)
{
return GenerationResult.Failure("No images returned from API");
}
// Декодируем изображение из base64
byte[] imageData = Convert.FromBase64String(apiResponse.images[0]);
byte[] imageData = Convert.FromBase64String(response.Images[0]);
// Сохраняем изображение
string fileName = $"pawn_{DateTime.Now:yyyyMMdd_HHmmss}.png";
@@ -144,6 +125,11 @@ namespace AIImages.Services
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}");
@@ -160,12 +146,8 @@ namespace AIImages.Services
try
{
string endpoint = $"{apiEndpoint}/sdapi/v1/sd-models";
HttpResponseMessage response = await _sharedHttpClient.GetAsync(
endpoint,
cancellationToken
);
return response.IsSuccessStatusCode;
// Используем встроенный метод PingAsync библиотеки
return await _client.PingAsync(cancellationToken);
}
catch (Exception ex)
{
@@ -183,24 +165,16 @@ namespace AIImages.Services
try
{
string endpoint = $"{apiEndpoint}/sdapi/v1/sd-models";
HttpResponseMessage response = await _sharedHttpClient.GetAsync(
endpoint,
cancellationToken
);
if (!response.IsSuccessStatusCode)
return new List<string>();
string jsonResponse = await response.Content.ReadAsStringAsync();
var models = JsonConvert.DeserializeObject<List<SdModel>>(jsonResponse);
// Используем Models сервис библиотеки
var models = await _client.Models.GetModelsAsync(cancellationToken);
var modelNames = new List<string>();
if (models != null)
{
foreach (var model in models)
{
modelNames.Add(model.title ?? model.model_name);
// Используем Title или ModelName в зависимости от того, что доступно
modelNames.Add(model.Title ?? model.ModelName);
}
}
@@ -223,25 +197,13 @@ namespace AIImages.Services
try
{
string endpoint = $"{apiEndpoint}/sdapi/v1/samplers";
HttpResponseMessage response = await _sharedHttpClient.GetAsync(
endpoint,
cancellationToken
);
if (!response.IsSuccessStatusCode)
return GetDefaultSamplers();
string jsonResponse = await response.Content.ReadAsStringAsync();
var samplers = JsonConvert.DeserializeObject<List<SdSampler>>(jsonResponse);
// Используем Samplers сервис библиотеки
var samplers = await _client.Samplers.GetSamplersAsync(cancellationToken);
var samplerNames = new List<string>();
if (samplers != null)
{
foreach (var sampler in samplers)
{
samplerNames.Add(sampler.name);
}
samplerNames.AddRange(samplers);
}
Log.Message($"[AI Images] Found {samplerNames.Count} samplers");
@@ -263,25 +225,13 @@ namespace AIImages.Services
try
{
string endpoint = $"{apiEndpoint}/sdapi/v1/schedulers";
HttpResponseMessage response = await _sharedHttpClient.GetAsync(
endpoint,
cancellationToken
);
if (!response.IsSuccessStatusCode)
return GetDefaultSchedulers();
string jsonResponse = await response.Content.ReadAsStringAsync();
var schedulers = JsonConvert.DeserializeObject<List<SdScheduler>>(jsonResponse);
// Используем Schedulers сервис библиотеки (доступен с версии 1.1.1)
var schedulers = await _client.Schedulers.GetSchedulersAsync(cancellationToken);
var schedulerNames = new List<string>();
if (schedulers != null)
{
foreach (var scheduler in schedulers)
{
schedulerNames.Add(scheduler.name);
}
schedulerNames.AddRange(schedulers);
}
Log.Message($"[AI Images] Found {schedulerNames.Count} schedulers");
@@ -346,32 +296,13 @@ namespace AIImages.Services
if (_disposed)
return;
// Не dispose shared HttpClient - он используется глобально
// Dispose клиента StableDiffusion
if (_client is IDisposable disposableClient)
{
disposableClient.Dispose();
}
_disposed = true;
}
// Вспомогательные классы для десериализации JSON ответов
#pragma warning disable S3459, S1144, IDE1006 // Properties set by JSON deserializer
private sealed class Txt2ImgResponse
{
public string[] images { get; set; }
}
private sealed class SdModel
{
public string title { get; set; }
public string model_name { get; set; }
}
private sealed class SdSampler
{
public string name { get; set; }
}
private sealed class SdScheduler
{
public string name { get; set; }
}
#pragma warning restore S3459, S1144, IDE1006
}
}