app.py

import gradio as gr
import torch, torchvision
import torch.nn.functional as F
import numpy as np
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'

# Load the pretrained pipeline
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)

# And the core function to generate an image given the relevant inputs
def generate():
    x = torch.randn(8, 4, 64, 64).to(device)
    # Minimal sampling loop
    for i, t in tqdm(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
    grid = torchvision.utils.make_grid(x, nrow=4)
    im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
    im.convert("RGBA").save("test.png")
    return im

# See the gradio docs for the types of inputs and outputs available
outputs = gr.Image(label="result")

# Setting up a minimal interface to our function:
demo = gr.Interface(
    fn=generate,
    inputs=None,
    outputs=outputs,
)

# And launching
if __name__ == "__main__":
    demo.launch(enable_queue=True)