app.py

import gradio as gr
import torch, torchvision
from torchvision import transforms
import torch.nn.functional as F
import numpy as np
from time import time, ctime
from PIL import Image, ImageColor
from diffusers import DDPMPipeline
from diffusers import DDIMScheduler

device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion-V2'
image_pipe = DDPMPipeline.from_pretrained(pipeline_name, use_safetensors=True).to(device)
try:
    scheduler = DDIMScheduler.from_pretrained(pipeline_name)
except:
    scheduler = DDIMScheduler.from_pretrained(pipeline_name, subfolder="scheduler")
scheduler.set_timesteps(num_inference_steps=25)

def show_images_save(x):
    grid = torchvision.utils.make_grid(x, nrow=4, padding=0)
    grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
    grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
    #CROP
    width, height = grid_im.size
    sav = []
    for w in range(width//64):
      for h in range(height//64):
        box = (w*64, h*64, (w+1)*64, (h+1)*64)
        sav.append(grid_im.crop(box))
    return sav

def color_loss(images, target_color=(0, 0, 0, 1)):
    target = (torch.tensor(target_color).to(images.device) * 2 - 1)
    target = target[None, :, None, None]
    error = torch.abs(images - target).mean()
    return error

def generate():
    print(ctime(time()))
    x = torch.randn(8, 4, 64, 64).to(device)
    for i, t in enumerate(scheduler.timesteps):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = scheduler.step(noise_pred, t, x).prev_sample
    return show_images_save(x)

def generate_g(color, guidance_loss_scale):
    print(f'--V2-- {ctime(time())}')
    target_color = ImageColor.getcolor(color, "RGBA")  # Target color as RGB
    target_color = [a / 255 for a in target_color]  # Rescale from (0, 255) to (0, 1)
    x = torch.randn(8, 4, 64, 64).to(device)
    for i, t in enumerate(scheduler.timesteps):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = x.detach().requires_grad_()
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample
        loss = color_loss(x0, target_color) * guidance_loss_scale
        cond_grad = -torch.autograd.grad(loss, x)[0]
        x = x.detach() + cond_grad
        x = scheduler.step(noise_pred, t, x).prev_sample
    return show_images_save(x)

demo = gr.Blocks(css=".tabs {max-width: 60%; margin-left: 20%; margin-right: 20%;} #column {max-width: 50%} .fixed-height {min-height: 0px !important} @media(min-width: 1280px){.fixed-height {min-height: 0px !important}}")
with demo:
    gr.HTML(
        """
        <div style="text-align: center; margin: 0 auto;">
            <div style="display: inline-flex; align-items: center; justify-content: center;">
                <img src='https://huggingface.co/spaces/WiNE-iNEFF/MinecraftSkin-Diffusion/resolve/main/MSD_7.png' width='45%'>
            </div>
            <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                Gradio demo for Minecraft Skin Diffusion. This is simple Unconditional Diffusion Model<br>that can generate skins for game Minecraft.
            </p>
        </div>
        """)
    with gr.Tabs():
        with gr.TabItem("V1"):
            with gr.Column():
                gall = gr.Gallery().style(grid=[4])
                greet_btn = gr.Button("Generate")
                greet_btn.click(fn=generate, outputs=gall)
        with gr.TabItem("V2"):
            with gr.Row():
                with gr.Column(min_width = 0, elem_id='column'):
                    picker = gr.ColorPicker(label="color", value="#55FFAA", min_width='0')
                    slide = gr.Slider(label="guidance_scale", minimum=0, maximum=100, value=50)
                    greet_btn = gr.Button("Generate", min_width = '0')
                with gr.Column(min_width = 0):
                  gall = gr.Gallery(min_width = '250').style(grid=[3])
                greet_btn.click(fn=generate_g, inputs=[picker, slide], outputs=gall)
    gr.HTML(
        """
        <div class="footer">
            <div style='text-align: center;'>Minecraft Skin Diffusion by <a href='https://twitter.com/wine_ineff' target='_blank'>Artsem Holub (WiNE-iNEFF)</a></div>
        </div>
        """)

demo.launch()