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:
@@ -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>
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user