Initial D

.gitignore

@@ -0,0 +1,4 @@
+!**/.gitkeep
+
+cache/*
+output/*

Pipfile

@@ -0,0 +1,17 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+diffusers = { version = "==0.7.*", extras = ["torch"] }
+gradio = "==3.9.*"
+scipy = "==1.9.*"
+torch = "==1.13.*"
+torchvision = "==0.14.*"
+transformers = "==4.24.*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.10"

README.md

@@ -1,6 +1,6 @@
 ---
-title: Sd Spritesheets
-emoji: 🌖
+title: Stable Diffusion Sprite Sheets
+emoji: 🚶‍♀️
 colorFrom: purple
 colorTo: gray
 sdk: gradio
@@ -8,5 +8,3 @@ sdk_version: 3.9
 app_file: app.py
 pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,48 @@
+from typing import List
+
+import time
+import gradio as gr
+from modules.sprites import generate_sides, build_gifs, build_spritesheet
+
+
+def generate(prompt, thresh):
+		timestamp = int(time.time())
+
+		sides = generate_sides(prompt, 3)[0]	
+		spritesheet = build_spritesheet(sides, prompt, timestamp=timestamp, thresh=thresh)[0]
+
+		filepaths = build_gifs(sides, prompt, save=True, timestamp=timestamp, thresh=thresh)[1]
+
+		return spritesheet, filepaths[0], filepaths[1], filepaths[2], filepaths[3]
+
+
+demo = gr.Blocks()
+
+with demo:
+	gr.Markdown("""
+		# Stable Diffusion Sprite Sheets
+
+		Generate a sprite sheet of pixel art character sides and their corresponding walk animations! Checkpoint by [Onodofthenorth](https://huggingface.co/Onodofthenorth/SD_PixelArt_SpriteSheet_Generator). Sprites are 32x32 pixels scaled up to 96x96. NSFW content replaced with blank sprites.
+	""")
+
+	with gr.Row():
+		with gr.Column():
+			prompt = gr.Textbox(label="Prompt", placeholder="Enter text prompt")
+			threshold = gr.Slider(label="Background removal threshold", placeholder="Tweak how strong the background removal is", minimum=0, maximum=255, value=128)
+
+		button = gr.Button("Generate")
+
+	with gr.Box():
+		with gr.Row():
+			spritesheet = gr.Image(label="Sprite Sheet")
+
+			with gr.Row():
+				front = gr.Image(label="Front")
+				back = gr.Image(label="Back")
+				left = gr.Image(label="Left")
+				right = gr.Image(label="Right")
+
+		button.click(fn=generate, inputs=[prompt, threshold], outputs=[spritesheet, front, back, left, right])
+
+demo.queue()
+demo.launch(show_api=False) 

modules/sprites.py

@@ -0,0 +1,266 @@
+import os
+import io
+import re
+import time
+import random
+import torch
+from typing import 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

requirements.txt

@@ -0,0 +1,80 @@
+-i https://pypi.org/simple
+accelerate==0.14.0
+aiohttp==3.8.3
+aiosignal==1.3.1
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==22.1.0
+bcrypt==4.0.1
+certifi==2022.9.24
+cffi==1.15.1
+charset-normalizer==2.1.1
+click==8.1.3
+contourpy==1.0.6
+cryptography==38.0.3
+cycler==0.11.0
+diffusers==0.7.2
+fastapi==0.86.0
+ffmpy==0.3.0
+filelock==3.8.0
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2022.11.0
+gradio==3.9.1
+h11==0.12.0
+httpcore==0.15.0
+httpx==0.23.0
+huggingface-hub==0.10.1
+idna==3.4
+importlib-metadata==5.0.0
+jinja2==3.1.2
+kiwisolver==1.4.4
+linkify-it-py==1.0.3
+markdown-it-py==2.1.0
+markupsafe==2.1.1
+matplotlib==3.6.2
+mdit-py-plugins==0.3.1
+mdurl==0.1.2
+multidict==6.0.2
+numpy==1.23.4
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+orjson==3.8.1
+packaging==21.3
+pandas==1.5.1
+paramiko==2.12.0
+pillow==9.3.0
+psutil==5.9.4
+pycparser==2.21
+pycryptodome==3.15.0
+pydantic==1.10.2
+pydub==0.25.1
+pynacl==1.5.0
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytz==2022.6
+pyyaml==6.0
+regex==2022.10.31
+requests==2.28.1
+rfc3986==1.5.0
+scipy==1.9.3
+setuptools==65.5.1
+six==1.16.0
+sniffio==1.3.0
+starlette==0.20.4
+tokenizers==0.13.2
+torch==1.13.0
+torchvision==0.14.0
+tqdm==4.64.1
+transformers==4.24.0
+typing-extensions==4.4.0
+uc-micro-py==1.0.1
+urllib3==1.26.12
+uvicorn==0.19.0
+websockets==10.4
+wheel==0.38.4
+yarl==1.8.1
+zipp==3.10.0