modules/sprites.py

import os
import io
import re
import time
import random
import torch
from typing import Dict, Final, List, Optional, Tuple, cast

from PIL import Image, ImageDraw, ImageEnhance
from PIL.Image import Image as PILImage
from diffusers import StableDiffusionPipeline

model_id: Final = "Onodofthenorth/SD_PixelArt_SpriteSheet_Generator"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, torch_dtype=torch.float16, cache_dir="cache"
)
pipe = pipe.to("cuda")

sprite_sides: Final = {
    "front": "PixelArtFSS",
    "right": "PixelArtRSS",
    "back": "PixelArtBSS",
    "left": "PixelArtLSS",
}


def torchGenerator(seed: Optional[int], max: int = 1024) -> Tuple[torch.Generator, int]:
    seed = seed or random.randrange(0, max)

    return torch.Generator("cuda").manual_seed(seed), seed


def generate(
    prompt: str,
    sfw_retries: int = 1,
    seed: Optional[int] = None,
) -> PILImage:
    """
    Generate a sprite image from a text description.

    Return a blank image if the model fails to generate a safe image.
    """

    generator = torchGenerator(seed)[0]
    image: PILImage | None = None

    for _ in range(sfw_retries):
        pipe_output = pipe(prompt, generator=generator, width=512, height=512)
        image = pipe_output.images[0]

        if not pipe_output.nsfw_content_detected[0]:
            break

        rand_seed = seed

        while rand_seed == seed:
            print(f"Regenerating `{prompt}` with different seed.")

            rand_seed = random.randrange(0, 1024)
            generator = torchGenerator(rand_seed)[0]

    return cast(PILImage, image)


def generate_sides(
    prompt: str, sfw_retries: int = 1, sides: Dict[str, str] = sprite_sides
) -> Tuple[Dict[str, PILImage], str]:
    """
    Generate sprite images from a text description of different sides.

    If both left and right side specified, duplicate and flip left side as the right side
    """

    print(f"Generating sprites for `{prompt}`")

    seed = random.randrange(0, 1024)
    sprites = {}

    # If both left and right side specified, duplicate and flip left side as the right side
    for side, label in sides.items():
        if side == "right" and "left" in sides and "right" in sides:
            continue

        sprites[side] = generate(f"({prompt}) [nsfw] [photograph] {label}", sfw_retries, seed)

    if "left" in sides and "right" in sides:
        sprites["right"] = sprites["left"].transpose(Image.Transpose.FLIP_LEFT_RIGHT)

    return sprites, prompt


def clean_sprite(
    image: PILImage,
    size: Tuple[int, int] = (192, 192),
    sharpness: float = 1.5,
    thresh: int = 128,
    rescaling: Optional[int] = None,
) -> PILImage:
    """
    Process image to be more sprite-like.

    `rescale` will first scale down by value, then up to specified size.
    """

    width, height = image.size
    sharpener = ImageEnhance.Sharpness(image)

    image = sharpener.enhance(sharpness)
    image = image.convert("RGBA")
    ImageDraw.floodfill(image, (0, 0), (255, 255, 255, 0), thresh=thresh)

    if type(rescaling) is int:
        image = image.resize(
            (int(width / rescaling), int(height / rescaling)),
            resample=Image.Resampling.NEAREST,
        )

    image = image.resize(size, resample=Image.Resampling.NEAREST)

    return image


def split_sprites(image: PILImage, size: Tuple[int, int] = (96, 96)) -> List[PILImage]:
    """Split sprite image into individual sides."""

    width, height = image.size
    w, h = size

    # fmt: off
    frames = [
        image.crop((
            0,
            int(h / 2),
            int(width / 4),
            int(height * 0.75),
        )),
        image.crop((
            int(width / 4),
            int(h / 2),
            int(width / 4) * 2,
            int(height * 0.75),
        )),
        image.crop((
            int(width / 4) * 2,
            int(h / 2),
            int(width / 4) * 3,
            int(height * 0.75),
        )),
        image.crop((
            int(width / 4) * 3,
            int(h / 2),
            width,
            int(height * 0.75),
        )),
    ]
    # fmt: on

    new_canvas = Image.new("RGBA", size, (255, 255, 255, 0))

    for i in range(len(frames)):
        canvas = new_canvas.copy()
        canvas.paste(frames[i], (int(w / 4), 0, int(w * 0.75), h))
        frames[i] = canvas

    return frames


def build_spritesheet(
    images: Dict[str, PILImage],
    text: str = "sd_pixelart",
    sprite_size: Tuple[int, int] = (96, 96),
    dir: str = "output",
    save: bool = False,
    timestamp: Optional[int] = None,
    thresh: int = 128,
) -> Tuple[PILImage, str | None]:
    """
    Build sprite sheet from sides.

    1. Clean and scale each image
    2. Split each image into individual frames
    3. Create a new spritesheet canvas for all sides[frames]
    4. Paste each individial frame onto canvas
    """

    frames = {}
    width, height = sprite_size
    text = re.sub(r"[^\w()[\]_-]", "", text)
    filepath = None

    for side, image in images.items():
        image = clean_sprite(image, (width * 2, height * 2), thresh=thresh)
        frames[side] = split_sprites(image, sprite_size)

    canvas = Image.new(
        "RGBA",
        (width * len(frames["front"]), height * len(frames)),
        (255, 255, 255, 0),
    )

    for j in range(len(frames["front"])):
        for k, side in enumerate(frames):
            canvas.paste(
                frames[side][j],
                (
                    j * width,
                    k * height,
                    j * width + width,
                    k * height + height,
                ),
            )

    spritesheet = io.BytesIO()
    canvas.save(spritesheet, "PNG")

    if save:
        timestamp = timestamp or int(time.time())
        filepath = os.path.join(dir, f"{timestamp}_{text}.png")
        canvas.save(filepath)

    return Image.open(spritesheet), filepath


def build_gifs(
    images: Dict[str, PILImage],
    text: str = "sd_spritesheet",
    dir: str = "output",
    duration: int | List[int] | Tuple[int, ...] = (300, 450, 300, 450),
    save: bool = False,
    timestamp: Optional[int] = None,
    thresh: int = 128,
) -> Tuple[Dict[str, List[PILImage]], List[str] | None]:
    """Build animated GIFs from side frames."""

    gifs = {}
    text = re.sub(r"[^\w()[\]_-]", "", text)
    filepaths = [] if save else None

    for side, image in images.items():
        image = clean_sprite(image, thresh=thresh)
        frames = split_sprites(image)

        gif = io.BytesIO()

        options = {
            "fp": gif,
            "format": "GIF",
            "save_all": True,
            "append_images": frames[1:],
            "disposal": 3,
            "duration": duration,
            "loop": 0,
        }

        frames[0].save(**options)
        gifs[side] = Image.open(gif)

        if save:
            timestamp = timestamp or int(time.time())
            filepath = os.path.join(dir, f"{timestamp}_{text}_{side}.gif")
            filepaths.append(filepath)

            options.update({"fp": filepath})
            frames[0].save(**options)

    return gifs, filepaths