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.

This commit is contained in:
Leonid Pershin
2025-10-26 19:10:45 +03:00
parent 3434927342
commit 02b0143186
11 changed files with 974 additions and 174 deletions

View File

@@ -1,6 +1,9 @@
using System;
using System.Linq;
using System.Threading;
using AIImages.Helpers;
using AIImages.Models;
using AIImages.Services;
using RimWorld;
using UnityEngine;
using Verse;
@@ -31,6 +34,12 @@ namespace AIImages
private Texture2D generatedImage;
private bool isGenerating = false;
private string generationStatus = "";
private CancellationTokenSource cancellationTokenSource;
// Сервисы (получаем через DI)
private readonly IPawnDescriptionService pawnDescriptionService;
private readonly IPromptGeneratorService promptGeneratorService;
private readonly IStableDiffusionApiService apiService;
public Window_AIImage(Pawn pawn)
{
@@ -42,6 +51,12 @@ namespace AIImages
this.draggable = true;
this.preventCameraMotion = false;
// Получаем сервисы через DI контейнер
var services = AIImagesMod.Services;
pawnDescriptionService = services.PawnDescriptionService;
promptGeneratorService = services.PromptGeneratorService;
apiService = services.ApiService;
// Извлекаем данные персонажа
RefreshPawnData();
}
@@ -57,10 +72,27 @@ namespace AIImages
/// </summary>
private void RefreshPawnData()
{
appearanceData = AIImagesMod.PawnDescriptionService.ExtractAppearanceData(pawn);
appearanceData = pawnDescriptionService.ExtractAppearanceData(pawn);
generationSettings = AIImagesMod.Settings.ToStableDiffusionSettings();
}
/// <summary>
/// Освобождает ресурсы при закрытии окна
/// </summary>
public override void PreClose()
{
base.PreClose();
// Отменяем генерацию, если она идет
if (isGenerating && cancellationTokenSource != null)
{
cancellationTokenSource.Cancel();
}
// Освобождаем CancellationTokenSource
cancellationTokenSource?.Dispose();
}
/// <summary>
/// Обновляет текущую пешку в окне
/// </summary>
@@ -101,24 +133,28 @@ namespace AIImages
}
/// <summary>
/// Асинхронная генерация изображения
/// Асинхронная генерация изображения с поддержкой отмены
/// </summary>
private async System.Threading.Tasks.Task GenerateImage()
private async System.Threading.Tasks.Task GenerateImageAsync()
{
if (isGenerating)
return;
// Создаем новый CancellationTokenSource
cancellationTokenSource?.Dispose();
cancellationTokenSource = new CancellationTokenSource();
isGenerating = true;
generationStatus = "AIImages.Generation.InProgress".Translate();
try
{
// Генерируем промпты
string positivePrompt = AIImagesMod.PromptGeneratorService.GeneratePositivePrompt(
string positivePrompt = promptGeneratorService.GeneratePositivePrompt(
appearanceData,
generationSettings
);
string negativePrompt = AIImagesMod.PromptGeneratorService.GenerateNegativePrompt(
string negativePrompt = promptGeneratorService.GenerateNegativePrompt(
generationSettings
);
@@ -137,8 +173,11 @@ namespace AIImages
Model = AIImagesMod.Settings.apiEndpoint,
};
// Генерируем изображение
var result = await AIImagesMod.ApiService.GenerateImageAsync(request);
// Генерируем изображение с поддержкой отмены
var result = await apiService.GenerateImageAsync(
request,
cancellationTokenSource.Token
);
if (result.Success)
{
@@ -159,10 +198,19 @@ namespace AIImages
Messages.Message(generationStatus, MessageTypeDefOf.RejectInput);
}
}
catch (OperationCanceledException)
{
generationStatus = "AIImages.Generation.Cancelled".Translate();
Log.Message("[AI Images] Generation cancelled by user");
}
catch (Exception ex)
{
generationStatus = $"Error: {ex.Message}";
Log.Error($"[AI Images] Generation error: {ex}");
Messages.Message(
$"AIImages.Generation.Error".Translate() + ": {ex.Message}",
MessageTypeDefOf.RejectInput
);
}
finally
{
@@ -170,6 +218,26 @@ namespace AIImages
}
}
/// <summary>
/// Запускает генерацию изображения (обертка для безопасного fire-and-forget)
/// </summary>
private void StartGeneration()
{
AsyncHelper.FireAndForget(GenerateImageAsync(), "Image Generation");
}
/// <summary>
/// Отменяет генерацию изображения
/// </summary>
private void CancelGeneration()
{
if (isGenerating && cancellationTokenSource != null)
{
cancellationTokenSource.Cancel();
generationStatus = "AIImages.Generation.Cancelling".Translate();
}
}
public override void DoWindowContents(Rect inRect)
{
float curY = 0f;
@@ -229,9 +297,7 @@ namespace AIImages
contentY += 35f;
Text.Font = GameFont.Small;
string appearanceText = AIImagesMod.PawnDescriptionService.GetAppearanceDescription(
pawn
);
string appearanceText = pawnDescriptionService.GetAppearanceDescription(pawn);
float appearanceHeight = Text.CalcHeight(appearanceText, scrollViewRect.width - 30f);
Widgets.Label(
new Rect(20f, contentY, scrollViewRect.width - 30f, appearanceHeight),
@@ -252,7 +318,7 @@ namespace AIImages
contentY += 35f;
Text.Font = GameFont.Small;
string apparelText = AIImagesMod.PawnDescriptionService.GetApparelDescription(pawn);
string apparelText = pawnDescriptionService.GetApparelDescription(pawn);
float apparelHeight = Text.CalcHeight(apparelText, scrollViewRect.width - 30f);
Widgets.Label(
new Rect(20f, contentY, scrollViewRect.width - 30f, apparelHeight),
@@ -308,18 +374,33 @@ namespace AIImages
curY += statusHeight + 10f;
}
// Кнопка генерации
// Кнопка генерации/отмены
Text.Font = GameFont.Small;
if (
Widgets.ButtonText(
new Rect(rect.x, rect.y + curY, rect.width, 35f),
isGenerating
? "AIImages.Generation.Generating".Translate()
: "AIImages.Generation.Generate".Translate()
) && !isGenerating
)
if (isGenerating)
{
_ = GenerateImage();
// Показываем кнопку отмены во время генерации
if (
Widgets.ButtonText(
new Rect(rect.x, rect.y + curY, rect.width, 35f),
"AIImages.Generation.Cancel".Translate()
)
)
{
CancelGeneration();
}
}
else
{
// Показываем кнопку генерации
if (
Widgets.ButtonText(
new Rect(rect.x, rect.y + curY, rect.width, 35f),
"AIImages.Generation.Generate".Translate()
)
)
{
StartGeneration();
}
}
curY += 40f;
@@ -333,7 +414,7 @@ namespace AIImages
// Получаем промпт
Text.Font = GameFont.Tiny;
string promptText = AIImagesMod.PromptGeneratorService.GeneratePositivePrompt(
string promptText = promptGeneratorService.GeneratePositivePrompt(
appearanceData,
generationSettings
);
@@ -411,9 +492,7 @@ namespace AIImages
height += 35f;
// Текст внешности
string appearanceText = AIImagesMod.PawnDescriptionService.GetAppearanceDescription(
pawn
);
string appearanceText = pawnDescriptionService.GetAppearanceDescription(pawn);
height += Text.CalcHeight(appearanceText, 400f) + 20f;
// Разделитель
@@ -423,7 +502,7 @@ namespace AIImages
height += 35f;
// Текст одежды
string apparelText = AIImagesMod.PawnDescriptionService.GetApparelDescription(pawn);
string apparelText = pawnDescriptionService.GetApparelDescription(pawn);
height += Text.CalcHeight(apparelText, 400f) + 20f;
// Дополнительный отступ