Spaces:
Running
Running
File size: 2,422 Bytes
0213535 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from __future__ import annotations
from pathlib import Path
import gradio as gr
import numpy as np
import torch
from model import PocketDenoiser
from PIL import Image
from safetensors.torch import load_file
STEPS = 50
ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "pocket-diffusion"
MODEL = PocketDenoiser(diffusion_steps=STEPS)
MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
MODEL.eval()
def generate_digit(label: int, seed: int, guidance: float) -> Image.Image:
generator = torch.Generator().manual_seed(seed)
betas = torch.linspace(1e-4, 0.025, STEPS)
alphas = 1.0 - betas
cumulative = torch.cumprod(alphas, dim=0)
pixels = torch.randn(1, 64, generator=generator)
labels = torch.tensor([label])
null_labels = torch.tensor([10])
with torch.no_grad():
for step in reversed(range(STEPS)):
timesteps = torch.tensor([step])
conditional = MODEL(pixels, timesteps, labels)
unconditional = MODEL(pixels, timesteps, null_labels)
noise_prediction = unconditional + guidance * (conditional - unconditional)
alpha = alphas[step]
mean = (
pixels - (1 - alpha) / torch.sqrt(1 - cumulative[step]) * noise_prediction
) / torch.sqrt(alpha)
if step:
pixels = mean + torch.sqrt(betas[step]) * torch.randn(
pixels.shape,
generator=generator,
)
else:
pixels = mean
image = torch.clamp((pixels[0] + 1) / 2, 0, 1).reshape(8, 8).numpy()
array = np.clip(image * 255, 0, 255).astype(np.uint8)
return Image.fromarray(array, mode="L").resize(
(512, 512),
Image.Resampling.NEAREST,
)
with gr.Blocks(title="PocketDiffusion") as demo:
gr.Markdown("# PocketDiffusion\nGenerate a digit through 50 reverse-denoising steps.")
with gr.Row():
label = gr.Slider(0, 9, value=8, step=1, label="Digit")
seed = gr.Slider(0, 100_000, value=2032, step=1, label="Noise seed")
guidance = gr.Slider(1.0, 4.0, value=3.0, step=0.1, label="Guidance")
output = gr.Image(value=generate_digit(8, 2032, 3.0), label="Generated glyph")
button = gr.Button("Denoise", variant="primary")
button.click(generate_digit, inputs=[label, seed, guidance], outputs=output)
if __name__ == "__main__":
demo.launch()
|