Files
ai-images/Source/AIImages/Window_AIImage.cs

788 lines
28 KiB
C#
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
using System;
using System.Linq;
using System.Threading;
using AIImages.Helpers;
using AIImages.Models;
using AIImages.Services;
using RimWorld;
using UnityEngine;
using Verse;
#pragma warning disable IDE1006 // Naming Styles
namespace AIImages
{
/// <summary>
/// Окно для просмотра персонажа и генерации AI изображений
/// </summary>
[System.Diagnostics.CodeAnalysis.SuppressMessage(
"Style",
"IDE1006:Naming Styles",
Justification = "RimWorld Window naming convention"
)]
[System.Diagnostics.CodeAnalysis.SuppressMessage(
"Minor Code Smell",
"S101:Types should be named in PascalCase",
Justification = "RimWorld Window naming convention"
)]
public class Window_AIImage : Window
{
private Pawn pawn;
private PawnAppearanceData appearanceData;
private StableDiffusionSettings generationSettings;
private Texture2D generatedImage;
private bool isGenerating = false;
private string generationStatus = "";
private CancellationTokenSource cancellationTokenSource;
// Прогресс генерации
private double generationProgress = 0.0;
private int currentStep = 0;
private int totalSteps = 0;
private double etaSeconds = 0.0;
// Сервисы (получаем через DI)
private readonly IPawnDescriptionService pawnDescriptionService;
private readonly IPromptGeneratorService promptGeneratorService;
private readonly IStableDiffusionApiService apiService;
public Window_AIImage(Pawn pawn)
{
this.pawn = pawn;
this.doCloseX = true;
this.doCloseButton = true;
this.forcePause = false;
this.absorbInputAroundWindow = false;
this.draggable = true;
this.preventCameraMotion = false;
// Получаем сервисы через DI контейнер
var services = AIImagesMod.Services;
pawnDescriptionService = services.PawnDescriptionService;
promptGeneratorService = services.PromptGeneratorService;
apiService = services.ApiService;
// Извлекаем данные персонажа
RefreshPawnData();
}
public override Vector2 InitialSize => new Vector2(900f, 800f);
private Vector2 scrollPosition = Vector2.zero;
private Vector2 rightColumnScrollPosition = Vector2.zero;
private Vector2 promptScrollPosition = Vector2.zero;
private Vector2 negativePromptScrollPosition = Vector2.zero;
private float copiedMessageTime = 0f;
/// <summary>
/// Обновляет данные персонажа
/// </summary>
private void RefreshPawnData()
{
appearanceData = pawnDescriptionService.ExtractAppearanceData(pawn);
generationSettings = AIImagesMod.Settings.ToStableDiffusionSettings();
// Загружаем сохраненный портрет, если есть
LoadSavedPortrait();
}
/// <summary>
/// Загружает сохраненный портрет персонажа
/// </summary>
private void LoadSavedPortrait()
{
if (PawnPortraitHelper.HasPortrait(pawn))
{
generatedImage = PawnPortraitHelper.LoadPortrait(pawn);
if (generatedImage != null)
{
generationStatus = "AIImages.Generation.LoadedFromSave".Translate();
}
}
}
/// <summary>
/// Освобождает ресурсы при закрытии окна
/// </summary>
public override void PreClose()
{
base.PreClose();
// Отменяем генерацию, если она идет
if (isGenerating && cancellationTokenSource != null)
{
cancellationTokenSource.Cancel();
}
// Освобождаем CancellationTokenSource
cancellationTokenSource?.Dispose();
}
/// <summary>
/// Обновляет текущую пешку в окне
/// </summary>
public void UpdatePawn(Pawn newPawn)
{
this.pawn = newPawn;
RefreshPawnData();
// Очищаем старое изображение при смене персонажа
generatedImage = null;
generationStatus = "";
generationProgress = 0.0;
}
/// <summary>
/// Получить текущую пешку
/// </summary>
public Pawn CurrentPawn => pawn;
/// <summary>
/// Вызывается каждый кадр для обновления окна
/// </summary>
public override void WindowUpdate()
{
base.WindowUpdate();
// Проверяем, изменилась ли выбранная пешка
Pawn selectedPawn = Find.Selector.SelectedPawns.FirstOrDefault(p =>
p.IsColonist && p.Spawned && p.Faction == Faction.OfPlayer
);
// Если выбрана новая колонистская пешка, обновляем окно
if (selectedPawn != null && selectedPawn != pawn)
{
UpdatePawn(selectedPawn);
}
// Уменьшаем таймер сообщения о копировании
if (copiedMessageTime > 0f)
{
copiedMessageTime -= Time.deltaTime;
}
}
/// <summary>
/// Асинхронная генерация изображения с поддержкой отмены
/// </summary>
private async System.Threading.Tasks.Task GenerateImageAsync()
{
if (isGenerating)
return;
// Создаем новый CancellationTokenSource
cancellationTokenSource?.Dispose();
cancellationTokenSource = new CancellationTokenSource();
isGenerating = true;
generationStatus = "AIImages.Generation.InProgress".Translate();
generationProgress = 0.0;
currentStep = 0;
totalSteps = generationSettings.Steps;
try
{
// Генерируем промпты
string positivePrompt = promptGeneratorService.GeneratePositivePrompt(
appearanceData,
generationSettings
);
string negativePrompt = promptGeneratorService.GenerateNegativePrompt(
generationSettings
);
// Создаем запрос
var request = new GenerationRequest
{
Prompt = positivePrompt,
NegativePrompt = negativePrompt,
Steps = generationSettings.Steps,
CfgScale = generationSettings.CfgScale,
Width = generationSettings.Width,
Height = generationSettings.Height,
Sampler = generationSettings.Sampler,
Scheduler = generationSettings.Scheduler,
Seed = generationSettings.Seed,
Model = AIImagesMod.Settings.apiEndpoint,
SaveImagesToServer = AIImagesMod.Settings.saveImagesToServer,
};
// Создаем отдельный CancellationTokenSource для мониторинга прогресса
var progressCts = new CancellationTokenSource();
var progressTask = MonitorProgressAsync(progressCts.Token);
// Генерируем изображение с поддержкой отмены
var result = await apiService.GenerateImageAsync(
request,
cancellationTokenSource.Token
);
// Останавливаем мониторинг прогресса
progressCts.Cancel();
try
{
await progressTask;
}
catch (OperationCanceledException)
{
// Ожидаемое исключение при остановке мониторинга
}
finally
{
progressCts.Dispose();
}
if (result.Success)
{
// Загружаем текстуру
generatedImage = new Texture2D(2, 2);
generatedImage.LoadImage(result.ImageData);
generationStatus = "AIImages.Generation.Success".Translate();
generationProgress = 1.0;
// Сохраняем путь к портрету на персонаже
PawnPortraitHelper.SavePortraitPath(pawn, result.SavedPath);
Messages.Message(
"AIImages.Generation.SavedTo".Translate(result.SavedPath),
MessageTypeDefOf.PositiveEvent
);
}
else
{
generationStatus =
$"AIImages.Generation.Failed".Translate() + ": {result.ErrorMessage}";
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
{
isGenerating = false;
}
}
/// <summary>
/// Мониторит прогресс генерации и обновляет UI
/// </summary>
private async System.Threading.Tasks.Task MonitorProgressAsync(
CancellationToken cancellationToken
)
{
try
{
while (!cancellationToken.IsCancellationRequested)
{
var progress = await apiService.GetProgressAsync(cancellationToken);
if (progress != null && progress.IsActive)
{
generationProgress = progress.Progress;
currentStep = progress.CurrentStep;
totalSteps = progress.TotalSteps;
etaSeconds = progress.EtaRelative;
Log.Message(
$"[AI Images] Progress: {progress.Progress:P} - Step {progress.CurrentStep}/{progress.TotalSteps} - ETA: {progress.EtaRelative:F1}s"
);
}
// Обновляем каждые 500ms
await System.Threading.Tasks.Task.Delay(500, cancellationToken);
}
}
catch (OperationCanceledException)
{
// Ожидаемое исключение при остановке
}
catch (Exception ex)
{
Log.Warning($"[AI Images] Progress monitoring error: {ex.Message}");
}
}
/// <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;
// Заголовок
Text.Font = GameFont.Medium;
Widgets.Label(
new Rect(0f, curY, inRect.width, 40f),
"AIImages.Window.Title".Translate()
);
curY += 45f;
// Имя пешки
Text.Font = GameFont.Small;
Widgets.Label(
new Rect(0f, curY, inRect.width, 30f),
"AIImages.Window.PawnLabel".Translate(pawn.NameShortColored.Resolve())
);
curY += 40f;
// Разделитель
Widgets.DrawLineHorizontal(0f, curY, inRect.width);
curY += 10f;
// Разделяем на две колонки: левая - информация, правая - изображение
float leftColumnWidth = inRect.width * 0.55f;
float rightColumnWidth = inRect.width * 0.42f;
float columnGap = inRect.width * 0.03f;
// Левая колонка - прокручиваемая информация
Rect leftColumnRect = new Rect(0f, curY, leftColumnWidth, inRect.height - curY);
DrawLeftColumn(leftColumnRect);
// Правая колонка - превью и управление
Rect rightColumnRect = new Rect(
leftColumnWidth + columnGap,
curY,
rightColumnWidth,
inRect.height - curY
);
DrawRightColumn(rightColumnRect);
}
private void DrawLeftColumn(Rect rect)
{
Rect scrollViewRect = new Rect(0f, 0f, rect.width - 20f, CalculateContentHeight());
Widgets.BeginScrollView(rect, ref scrollPosition, scrollViewRect);
float contentY = 0f;
// Портрет персонажа (если есть)
if (generatedImage != null)
{
float portraitSize = Mathf.Min(scrollViewRect.width - 20f, 200f);
Rect portraitRect = new Rect(
(scrollViewRect.width - portraitSize) / 2f,
contentY,
portraitSize,
portraitSize
);
GUI.DrawTexture(portraitRect, generatedImage);
contentY += portraitSize + 15f;
}
// Секция "Внешность"
Text.Font = GameFont.Medium;
Widgets.Label(
new Rect(10f, contentY, scrollViewRect.width - 20f, 30f),
"AIImages.Appearance.SectionTitle".Translate()
);
contentY += 35f;
Text.Font = GameFont.Small;
string appearanceText = pawnDescriptionService.GetAppearanceDescription(pawn);
float appearanceHeight = Text.CalcHeight(appearanceText, scrollViewRect.width - 30f);
Widgets.Label(
new Rect(20f, contentY, scrollViewRect.width - 30f, appearanceHeight),
appearanceText
);
contentY += appearanceHeight + 20f;
// Разделитель
Widgets.DrawLineHorizontal(10f, contentY, scrollViewRect.width - 20f);
contentY += 15f;
// Секция "Одежда"
Text.Font = GameFont.Medium;
Widgets.Label(
new Rect(10f, contentY, scrollViewRect.width - 20f, 30f),
"AIImages.Apparel.SectionTitle".Translate()
);
contentY += 35f;
Text.Font = GameFont.Small;
string apparelText = pawnDescriptionService.GetApparelDescription(pawn);
float apparelHeight = Text.CalcHeight(apparelText, scrollViewRect.width - 30f);
Widgets.Label(
new Rect(20f, contentY, scrollViewRect.width - 30f, apparelHeight),
apparelText
);
Widgets.EndScrollView();
}
private void DrawRightColumn(Rect rect)
{
// Рассчитываем высоту контента для скролла
float contentHeight = CalculateRightColumnHeight(rect);
Rect scrollViewRect = new Rect(0f, 0f, rect.width - 20f, contentHeight);
Widgets.BeginScrollView(rect, ref rightColumnScrollPosition, scrollViewRect);
float curY = 0f;
curY = DrawImagePreview(scrollViewRect, curY);
curY = DrawGenerationStatus(scrollViewRect, curY);
curY = DrawProgressBar(scrollViewRect, curY);
curY = DrawGenerationButton(scrollViewRect, curY);
// Позитивный промпт секция
Text.Font = GameFont.Medium;
Widgets.Label(
new Rect(0f, curY, scrollViewRect.width, 30f),
"AIImages.Prompt.PositiveTitle".Translate()
);
curY += 35f;
// Получаем позитивный промпт
Text.Font = GameFont.Tiny;
string positivePrompt = promptGeneratorService.GeneratePositivePrompt(
appearanceData,
generationSettings
);
// Фиксированная высота для области промпта
float promptBoxHeight = 100f;
float actualPositiveHeight = Text.CalcHeight(
positivePrompt,
scrollViewRect.width - 20f
);
Rect positiveOuterRect = new Rect(0f, curY, scrollViewRect.width, promptBoxHeight);
Rect positiveViewRect = new Rect(
0f,
0f,
scrollViewRect.width - 20f,
actualPositiveHeight
);
// Рисуем фон
Widgets.DrawBoxSolid(positiveOuterRect, new Color(0.1f, 0.3f, 0.1f, 0.5f));
// Рисуем промпт с прокруткой
Widgets.BeginScrollView(
positiveOuterRect.ContractedBy(5f),
ref promptScrollPosition,
positiveViewRect
);
Widgets.Label(
new Rect(0f, 0f, positiveViewRect.width, positiveViewRect.height),
positivePrompt
);
Widgets.EndScrollView();
curY += promptBoxHeight + 10f;
// Негативный промпт секция
Text.Font = GameFont.Medium;
Widgets.Label(
new Rect(0f, curY, scrollViewRect.width, 30f),
"AIImages.Prompt.NegativeTitle".Translate()
);
curY += 35f;
// Получаем негативный промпт
Text.Font = GameFont.Tiny;
string negativePrompt = promptGeneratorService.GenerateNegativePrompt(
generationSettings
);
float actualNegativeHeight = Text.CalcHeight(
negativePrompt,
scrollViewRect.width - 20f
);
Rect negativeOuterRect = new Rect(0f, curY, scrollViewRect.width, promptBoxHeight);
Rect negativeViewRect = new Rect(
0f,
0f,
scrollViewRect.width - 20f,
actualNegativeHeight
);
// Рисуем фон (красноватый для негативного)
Widgets.DrawBoxSolid(negativeOuterRect, new Color(0.3f, 0.1f, 0.1f, 0.5f));
// Рисуем промпт с прокруткой
Widgets.BeginScrollView(
negativeOuterRect.ContractedBy(5f),
ref negativePromptScrollPosition,
negativeViewRect
);
Widgets.Label(
new Rect(0f, 0f, negativeViewRect.width, negativeViewRect.height),
negativePrompt
);
Widgets.EndScrollView();
curY += promptBoxHeight + 10f;
// Кнопки копирования промптов
if (
Widgets.ButtonText(
new Rect(0f, curY, scrollViewRect.width / 2f - 5f, 30f),
"AIImages.Prompt.CopyPositive".Translate()
)
)
{
string positiveForCopy = promptGeneratorService.GeneratePositivePrompt(
appearanceData,
generationSettings
);
GUIUtility.systemCopyBuffer = positiveForCopy;
copiedMessageTime = 2f;
}
if (
Widgets.ButtonText(
new Rect(
scrollViewRect.width / 2f + 5f,
curY,
scrollViewRect.width / 2f - 5f,
30f
),
"AIImages.Prompt.CopyNegative".Translate()
)
)
{
string negativeForCopy = promptGeneratorService.GenerateNegativePrompt(
generationSettings
);
GUIUtility.systemCopyBuffer = negativeForCopy;
copiedMessageTime = 2f;
}
curY += 35f;
// Кнопка обновления данных
if (
Widgets.ButtonText(
new Rect(0f, curY, scrollViewRect.width, 30f),
"AIImages.Window.Refresh".Translate()
)
)
{
RefreshPawnData();
}
// Сообщение о копировании
if (copiedMessageTime > 0f)
{
curY += 35f;
GUI.color = new Color(0f, 1f, 0f, copiedMessageTime / 2f);
Widgets.Label(
new Rect(0f, curY, scrollViewRect.width, 25f),
"AIImages.Prompt.Copied".Translate()
);
GUI.color = Color.white;
}
Widgets.EndScrollView();
}
private float CalculateContentHeight()
{
float height = 0f;
// Портрет персонажа (если есть)
if (generatedImage != null)
{
float portraitSize = Mathf.Min(400f, 200f);
height += portraitSize + 15f;
}
// Заголовок "Внешность"
height += 35f;
// Текст внешности
string appearanceText = pawnDescriptionService.GetAppearanceDescription(pawn);
height += Text.CalcHeight(appearanceText, 400f) + 20f;
// Разделитель
height += 15f;
// Заголовок "Одежда"
height += 35f;
// Текст одежды
string apparelText = pawnDescriptionService.GetApparelDescription(pawn);
height += Text.CalcHeight(apparelText, 400f) + 20f;
// Дополнительный отступ
height += 50f;
return height;
}
private float CalculateRightColumnHeight(Rect rect)
{
float height = 0f;
float contentWidth = rect.width - 20f;
// Превью изображения
if (generatedImage != null)
{
height += 200f + 10f;
}
else if (!isGenerating)
{
height += 100f + 10f;
}
// Статус генерации
if (!string.IsNullOrEmpty(generationStatus))
{
height += 30f;
}
// Прогресс бар
if (isGenerating && generationProgress > 0.0)
{
height += 30f;
}
// Кнопка генерации
height += 40f;
// Позитивный промпт
height += 35f; // Заголовок
height += 100f + 10f; // Бокс
// Негативный промпт
height += 35f; // Заголовок
height += 100f + 10f; // Бокс
// Кнопки копирования
height += 35f;
// Кнопка обновления
height += 35f;
// Сообщение о копировании
if (copiedMessageTime > 0f)
{
height += 30f;
}
return height + 50f; // Дополнительный отступ
}
private float DrawImagePreview(Rect rect, float curY)
{
if (generatedImage != null)
{
float imageSize = Mathf.Min(rect.width, 400f);
Rect imageRect = new Rect(
rect.x + (rect.width - imageSize) / 2f,
rect.y + curY,
imageSize,
imageSize
);
GUI.DrawTexture(imageRect, generatedImage);
return curY + imageSize + 10f;
}
float placeholderSize = Mathf.Min(rect.width, 300f);
Rect placeholderRect = new Rect(
rect.x + (rect.width - placeholderSize) / 2f,
rect.y + curY,
placeholderSize,
placeholderSize
);
Widgets.DrawBoxSolid(placeholderRect, new Color(0.2f, 0.2f, 0.2f));
Text.Anchor = TextAnchor.MiddleCenter;
Widgets.Label(placeholderRect, "AIImages.Generation.NoImage".Translate());
Text.Anchor = TextAnchor.UpperLeft;
return curY + placeholderSize + 10f;
}
private float DrawGenerationStatus(Rect rect, float curY)
{
if (string.IsNullOrEmpty(generationStatus))
return curY;
Text.Font = GameFont.Small;
float statusHeight = Text.CalcHeight(generationStatus, rect.width);
Widgets.Label(
new Rect(rect.x, rect.y + curY, rect.width, statusHeight),
generationStatus
);
return curY + statusHeight + 10f;
}
private float DrawProgressBar(Rect rect, float curY)
{
if (!isGenerating || generationProgress <= 0.0)
return curY;
Rect progressBarRect = new Rect(rect.x, rect.y + curY, rect.width, 24f);
string progressText;
if (totalSteps > 0)
{
progressText =
$"{(generationProgress * 100):F1}% - Step {currentStep}/{totalSteps}";
if (etaSeconds > 0)
{
progressText += $" - ETA: {etaSeconds:F0}s";
}
}
else
{
progressText = $"{(generationProgress * 100):F1}%";
}
Widgets.FillableBar(progressBarRect, (float)generationProgress);
Text.Font = GameFont.Tiny;
Text.Anchor = TextAnchor.MiddleCenter;
Widgets.Label(progressBarRect, progressText);
Text.Anchor = TextAnchor.UpperLeft;
return curY + 30f;
}
private float DrawGenerationButton(Rect rect, float curY)
{
Text.Font = GameFont.Small;
string buttonLabel = isGenerating
? "AIImages.Generation.Cancel".Translate()
: "AIImages.Generation.Generate".Translate();
if (Widgets.ButtonText(new Rect(rect.x, rect.y + curY, rect.width, 35f), buttonLabel))
{
if (isGenerating)
CancelGeneration();
else
StartGeneration();
}
return curY + 40f;
}
}
}