WiNE-iNEFF's picture
Update app.py
eb4468f
raw
history blame
3.46 kB
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
from tqdm import tqdm
device = (
"mps"
if torch.backends.mps.is_available()
else "cuda"
if torch.cuda.is_available()
else "cpu"
)
pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion'
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
# Set up the scheduler
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
scheduler.set_timesteps(num_inference_steps=40)
def show_images_save(x):
"""Given a batch of images x, make a grid and convert to PIL"""
x = x * 0.5 + 0.5 # Map from (-1, 1) back to (0, 1)
grid = torchvision.utils.make_grid(x, nrow=4)
grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
grid_im.save(f"test.png")
return grid_im
def generate():
x = torch.randn(1, 4, 64, 64).to(device)
# Minimal sampling loop
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
# View the results
return show_images_save(x)
def ex():
t = time()
print(ctime(t))
return generate(), generate(), generate(), generate()
demo = gr.Blocks(css="#img_size {max-height: 128px} .container {max-width: 730px; margin: auto;} .min-h-\[15rem\]{min-height: 5rem !important;}")
with demo:
gr.HTML(
"""
<div style="text-align: center; margin: 0 auto;">
<div style="display: inline-flex;align-items: center;gap: 0.8rem;font-size: 1.75rem;">
<h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
Minecraft Skin Diffusion
</h1>
</div>
<p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
Gradio demo for Minecraft Skin Diffusion. This is simple Unconditional Diffusion Model that will help you generate skins for game Minecraft.
</p>
</div>
"""
)
with gr.Column():
with gr.Row().style(equal_height=True):
out = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
out2 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
with gr.Row().style(equal_height=True):
out3 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
out4 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
greet_btn = gr.Button("Generate")
greet_btn.click(fn=ex, inputs=None, outputs=[out, out2, out3, out4])
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> |
<center>
<img src='https://visitor-badge.glitch.me/badge?page_id=WiNE-iNEFF_MinecraftSkin-Diffusion' alt='visitor badge'>
</center>
</div>
</div>
"""
)
demo.launch()