Rotate pixel art using PixelLabs AI models

PixelLabRotate.cs

using UnityEngine;
using UnityEditor;
using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using UnityEngine.Networking;
using System.IO;

// Get an api key at https://www.pixellab.ai/pixellab-api
// Documentation for curl and example code: https://api.pixellab.ai/v1/docs

public class PixelLabRotate : EditorWindow
{
    private Texture2D referenceImage;
    private Texture2D processedReferenceImage;
    private int selectedDirections = 4;
    private Dictionary<string, Texture2D> generatedImages = new Dictionary<string, Texture2D>();
    private Dictionary<string, bool> isGeneratingDirection = new Dictionary<string, bool>();
    private string apiToken = "API KEY";
    private readonly string[] directionOptions = { "4 Directions", "8 Directions" };
    private readonly string[] viewOptions = { "high top-down", "low top-down", "side" };
    private readonly string[] cardinalDirections = { "north", "east", "south", "west" };
    private readonly string[] eightDirections = { "north", "north-east", "east", "south-east", "south", "south-west", "west", "north-west" };
    private int selectedFromViewIndex = 0;
    private int selectedToViewIndex = 0;
    private int selectedFromDirectionIndex = 2;
    private Vector2 scrollPosition;
    private float currentBalance = 0;
    private bool isGenerating = false;
    private float generationProgress = 0f;
    private bool shouldCancel = false;

    [MenuItem("Window/PixelLab/PixelLab Rotate")]
    public static void ShowWindow()
    {
        var window = GetWindow<PixelLabRotate>("PixelLab Rotate");
        window.minSize = new Vector2(600, 600);
        window.maxSize = new Vector2(600, 600);
    }

    async void OnEnable()
    {
        await UpdateBalance();
    }

    void OnDestroy()
    {
        if (generatedImages.Count > 0)
        {
            bool shouldClose = EditorUtility.DisplayDialog(
                "Unsaved Images",
                "You have unsaved generated images. Are you sure you want to close the window?",
                "Close Window",
                "Cancel"
            );

            if (!shouldClose)
            {
                ShowWindow();
            }
        }
    }

    private async Task UpdateBalance()
    {
        using (UnityWebRequest request = new UnityWebRequest("https://api.pixellab.ai/v1/balance", "GET"))
        {
            request.downloadHandler = new DownloadHandlerBuffer();
            request.SetRequestHeader("Authorization", $"Bearer {apiToken}");

            await request.SendWebRequest();

            if (request.result == UnityWebRequest.Result.Success)
            {
                string response = request.downloadHandler.text;
                var balanceData = JsonUtility.FromJson<BalanceResponse>(response);
                currentBalance = balanceData.usd;
                Repaint();
            }
        }
    }

    [Serializable]
    private class BalanceResponse
    {
        public float usd;
    }

    void OnGUI()
    {
        EditorGUILayout.BeginVertical(new GUIStyle { padding = new RectOffset(15, 15, 15, 15) });
        scrollPosition = EditorGUILayout.BeginScrollView(scrollPosition);

        EditorGUILayout.BeginVertical(EditorStyles.helpBox);
        var headerStyle = new GUIStyle(EditorStyles.boldLabel) { fontSize = 14, alignment = TextAnchor.MiddleLeft };
        EditorGUILayout.LabelField("PixelLab Rotate", headerStyle);
        EditorGUILayout.LabelField($"Balance: ${currentBalance:F2} USD", EditorStyles.miniLabel);
        EditorGUILayout.EndVertical();
        EditorGUILayout.Space(10);

        EditorGUILayout.BeginVertical(EditorStyles.helpBox);
        EditorGUILayout.LabelField("Reference Image", headerStyle);
        EditorGUILayout.Space(5);

        EditorGUILayout.BeginHorizontal();
        GUILayout.FlexibleSpace();
        Texture2D previousImage = referenceImage;
        referenceImage = (Texture2D)EditorGUILayout.ObjectField(referenceImage, typeof(Texture2D), false, GUILayout.Width(250));
        GUILayout.FlexibleSpace();
        EditorGUILayout.EndHorizontal();

        if (referenceImage != null && referenceImage != previousImage)
        {
            string path = AssetDatabase.GetAssetPath(referenceImage);
            TextureImporter importer = AssetImporter.GetAtPath(path) as TextureImporter;
            if (importer != null)
            {
                bool needsReimport = false;
                if (!importer.isReadable)
                {
                    importer.isReadable = true;
                    needsReimport = true;
                }
                if (importer.textureCompression != TextureImporterCompression.Uncompressed)
                {
                    importer.textureCompression = TextureImporterCompression.Uncompressed;
                    needsReimport = true;
                }
                if (importer.filterMode != FilterMode.Point)
                {
                    importer.filterMode = FilterMode.Point;
                    needsReimport = true;
                }
                if (needsReimport)
                {
                    AssetDatabase.ImportAsset(path);
                    Debug.Log($"Made texture readable, uncompressed and point filtered: {path}");
                }
            }

            processedReferenceImage = ProcessImage(referenceImage);
            if (processedReferenceImage != null)
            {
                processedReferenceImage.filterMode = FilterMode.Point;
            }
        }

        if (processedReferenceImage != null)
        {
            EditorGUILayout.Space(5);
            float imageSize = 128f;
            EditorGUILayout.BeginHorizontal();
            GUILayout.FlexibleSpace();

            var imageRect = GUILayoutUtility.GetRect(imageSize, imageSize);
            GUI.Box(imageRect, "", EditorStyles.helpBox);
            GUI.DrawTexture(new Rect(imageRect.x + 2, imageRect.y + 2, imageRect.width - 4, imageRect.height - 4), processedReferenceImage);

            GUILayout.FlexibleSpace();
            EditorGUILayout.EndHorizontal();
            EditorGUILayout.Space(5);
        }
        EditorGUILayout.EndVertical();
        EditorGUILayout.Space(10);

        EditorGUILayout.BeginVertical(EditorStyles.helpBox);
        EditorGUILayout.LabelField("View Settings", headerStyle);
        EditorGUILayout.Space(5);

        var popupStyle = new GUIStyle(EditorStyles.popup) { fixedHeight = 20 };
        selectedFromViewIndex = EditorGUILayout.Popup("From View", selectedFromViewIndex, viewOptions, popupStyle);
        selectedToViewIndex = EditorGUILayout.Popup("To View", selectedToViewIndex, viewOptions, popupStyle);
        EditorGUILayout.EndVertical();
        EditorGUILayout.Space(10);

        EditorGUILayout.BeginVertical(EditorStyles.helpBox);
        EditorGUILayout.LabelField("Direction Settings", headerStyle);
        EditorGUILayout.Space(5);

        int dirIndex = selectedDirections == 4 ? 0 : 1;
        dirIndex = EditorGUILayout.Popup("Number of Directions", dirIndex, directionOptions, popupStyle);
        selectedDirections = dirIndex == 0 ? 4 : 8;

        string[] currentDirections = selectedDirections == 4 ? cardinalDirections : eightDirections;
        selectedFromDirectionIndex = EditorGUILayout.Popup("Reference Direction", selectedFromDirectionIndex, currentDirections, popupStyle);
        EditorGUILayout.EndVertical();
        EditorGUILayout.Space(10);

        if (generatedImages.Count > 0)
        {
            EditorGUILayout.BeginVertical(EditorStyles.helpBox);
            EditorGUILayout.LabelField("Generated Images", headerStyle);
            EditorGUILayout.Space(5);
            DisplayGeneratedImages();
            EditorGUILayout.EndVertical();
            EditorGUILayout.Space(10);
        }

        EditorGUILayout.BeginVertical(EditorStyles.helpBox);

        var buttonStyle = new GUIStyle(GUI.skin.button)
        {
            fixedHeight = 30,
            fontStyle = FontStyle.Bold
        };

        GUI.enabled = !isGenerating && referenceImage != null;
        if (GUILayout.Button("Generate Rotations", buttonStyle))
        {
            if (referenceImage != null)
            {
                if (referenceImage.width > 128 || referenceImage.height > 128)
                {
                    EditorUtility.DisplayDialog("Error", "Image must be smaller than 128x128 pixels", "OK");
                    return;
                }
                shouldCancel = false;
                GenerateRotations();
            }
            else
            {
                EditorUtility.DisplayDialog("Error", "Please select a reference image first!", "OK");
            }
        }

        GUI.enabled = !isGenerating && generatedImages.Count > 0;
        if (GUILayout.Button("Save All Images", buttonStyle))
        {
            string outputPath = EditorUtility.SaveFolderPanel("Select Output Folder", "", "");
            if (!string.IsNullOrEmpty(outputPath))
            {
                string refImageName = Path.GetFileNameWithoutExtension(AssetDatabase.GetAssetPath(referenceImage));
                foreach (var kvp in generatedImages)
                {
                    string filePath = Path.Combine(outputPath, $"{refImageName}_{kvp.Key}.png");
                    File.WriteAllBytes(filePath, kvp.Value.EncodeToPNG());
                }
                AssetDatabase.Refresh();
                EditorUtility.DisplayDialog("Success", "Images saved successfully!", "OK");
            }
        }

        EditorGUILayout.EndVertical();

        GUI.enabled = true;
        EditorGUILayout.EndScrollView();
        EditorGUILayout.EndVertical();
    }

    private Texture2D ProcessImage(Texture2D source)
    {
        Texture2D processed = new Texture2D(source.width, source.height, TextureFormat.RGBA32, false);
        Color[] pixels = source.GetPixels();

        Color[] newPixels = new Color[pixels.Length];
        for (int i = 0; i < pixels.Length; i++)
        {
            if (pixels[i].a > 0)
            {
                newPixels[i] = pixels[i];
            }
            else
            {
                newPixels[i] = Color.clear;
            }
        }

        processed.SetPixels(newPixels);
        processed.Apply();
        processed.filterMode = FilterMode.Point;
        return processed;
    }

    private async void GenerateRotations()
    {
        if (referenceImage == null)
        {
            EditorUtility.DisplayDialog("Error", "Reference image is null", "OK");
            return;
        }

        if (referenceImage.width != referenceImage.height)
        {
            EditorUtility.DisplayDialog("Error", "Image must be square (16x16, 32x32, 64x64, or 128x128)", "OK");
            return;
        }

        int[] validSizes = { 16, 32, 64, 128 };
        if (Array.IndexOf(validSizes, referenceImage.width) == -1)
        {
            EditorUtility.DisplayDialog("Error", "Invalid image size. Must be 16x16, 32x32, 64x64, or 128x128", "OK");
            return;
        }

        isGenerating = true;
        generationProgress = 0f;

        string[] currentDirections = selectedDirections == 4 ? cardinalDirections : eightDirections;
        string fromDirection = currentDirections[selectedFromDirectionIndex];

        generatedImages.Clear();
        isGeneratingDirection.Clear();

        float progressPerDirection = 1f / currentDirections.Length;
        int completedDirections = 0;

        foreach (string direction in currentDirections)
        {
            if (shouldCancel)
            {
                break;
            }

            isGeneratingDirection[direction] = true;
            await GenerateRotatedImage(direction);
            isGeneratingDirection[direction] = false;
            completedDirections++;
            generationProgress = completedDirections * progressPerDirection;
            Repaint();

            await Task.Delay(1000);
        }

        await UpdateBalance();
        isGenerating = false;
        generationProgress = 0f;
        shouldCancel = false;
        Repaint();
    }

    private async Task GenerateRotatedImage(string toDirection)
    {
        if (processedReferenceImage == null)
        {
            Debug.LogError("Processed reference image is null");
            return;
        }

        byte[] imageBytes = processedReferenceImage.EncodeToPNG();

        if (imageBytes == null)
        {
            Debug.LogError("Failed to encode image to PNG");
            return;
        }

        string base64Image = Convert.ToBase64String(imageBytes);
        string fromView = viewOptions[selectedFromViewIndex];
        string toView = viewOptions[selectedToViewIndex];
        string[] currentDirections = selectedDirections == 4 ? cardinalDirections : eightDirections;
        string fromDirection = currentDirections[selectedFromDirectionIndex];

        string jsonBody = $@"{{
            ""image_size"": {{
                ""width"": {referenceImage.width},
                ""height"": {referenceImage.height}
            }},
            ""from_view"": ""{fromView}"",
            ""to_view"": ""{toView}"",
            ""from_direction"": ""{fromDirection}"",
            ""to_direction"": ""{toDirection}"",
            ""from_image"": {{
                ""type"": ""base64"",
                ""base64"": ""{base64Image}""
            }}
        }}";

        using (UnityWebRequest request = new UnityWebRequest("https://api.pixellab.ai/v1/rotate", "POST"))
        {
            byte[] bodyRaw = System.Text.Encoding.UTF8.GetBytes(jsonBody);
            request.uploadHandler = new UploadHandlerRaw(bodyRaw);
            request.downloadHandler = new DownloadHandlerBuffer();
            request.SetRequestHeader("Content-Type", "application/json");
            request.SetRequestHeader("Authorization", $"Bearer {apiToken}");

            await request.SendWebRequest();

            if (request.result == UnityWebRequest.Result.Success)
            {
                var response = JsonUtility.FromJson<RotateResponse>(request.downloadHandler.text);

                byte[] imageData = Convert.FromBase64String(response.image.base64);
                Texture2D newTexture = new Texture2D(2, 2);
                newTexture.LoadImage(imageData);
                newTexture.filterMode = FilterMode.Point;

                generatedImages[toDirection] = newTexture;
            }
        }
    }

    [Serializable]
    private class RotateResponse
    {
        public ImageData image;
    }

    [Serializable]
    private class ImageData
    {
        public string type;
        public string base64;
    }

    private void DisplayGeneratedImages()
    {
        float imageSize = 128f;
        float spacing = 15f;
        int imagesPerRow = 4;
        float totalWidth = (imageSize + spacing) * imagesPerRow - spacing;

        float totalHeight = selectedDirections == 4 ?
            imageSize + spacing :
            (imageSize + spacing) * 2 - spacing;

        EditorGUILayout.BeginVertical();

        if (isGenerating)
        {
            EditorGUILayout.Space(5);
            EditorGUILayout.BeginHorizontal();
            var progressRect = EditorGUILayout.GetControlRect(GUILayout.Height(22));
            EditorGUI.ProgressBar(progressRect, generationProgress, $"Generating... {(generationProgress * 100):F0}%");
            GUI.enabled = true;
            if (GUILayout.Button("Stop", GUILayout.Width(60), GUILayout.Height(22)))
            {
                shouldCancel = true;
            }
            EditorGUILayout.EndHorizontal();
            EditorGUILayout.Space(5);
        }

        EditorGUILayout.BeginHorizontal();
        Rect gridArea = EditorGUILayout.GetControlRect(
            GUILayout.Width(totalWidth),
            GUILayout.Height(((Mathf.Ceil(generatedImages.Count / (float)imagesPerRow)) * (imageSize + spacing + 40)) - spacing)
        );

        int currentIndex = 0;
        foreach (var kvp in generatedImages)
        {
            float xPos = gridArea.x + (currentIndex % imagesPerRow) * (imageSize + spacing);
            float yPos = gridArea.y + (currentIndex / imagesPerRow) * (imageSize + spacing + 40);

            // Make the box slightly larger to contain everything
            Rect groupRect = new Rect(xPos, yPos, imageSize, imageSize + 40);
            GUI.Box(groupRect, "", EditorStyles.helpBox);

            var labelStyle = new GUIStyle(EditorStyles.boldLabel) { alignment = TextAnchor.MiddleCenter };
            Rect labelRect = new Rect(xPos, yPos + 5, imageSize, 15);
            GUI.Label(labelRect, kvp.Key, labelStyle);

            Rect imageRect = new Rect(xPos + 5, yPos + 25, imageSize - 10, imageSize - 10);
            if (kvp.Value != null)
            {
                GUI.Box(imageRect, "", EditorStyles.helpBox);
                GUI.DrawTexture(new Rect(imageRect.x + 2, imageRect.y + 2, imageRect.width - 4, imageRect.height - 4), kvp.Value);
            }

            var buttonStyle = new GUIStyle(GUI.skin.button) { fixedHeight = 20 };
            Rect buttonRect = new Rect(xPos + (imageSize - 60) / 2, yPos + imageSize + 5, 60, 20);
            bool isLoading = isGeneratingDirection.ContainsKey(kvp.Key) && isGeneratingDirection[kvp.Key];

            if (isLoading)
            {
                GUI.Label(buttonRect, "Loading...", new GUIStyle(EditorStyles.miniLabel) { alignment = TextAnchor.MiddleCenter });
            }
            else if (GUI.Button(buttonRect, "Retry", buttonStyle))
            {
                isGeneratingDirection[kvp.Key] = true;
                _ = GenerateRotatedImage(kvp.Key).ContinueWith(_ =>
                {
                    isGeneratingDirection[kvp.Key] = false;
                    Repaint();
                });
            }

            currentIndex++;
        }

        EditorGUILayout.EndHorizontal();
        EditorGUILayout.EndVertical();
    }
}