Spaces:
Runtime error
Runtime error
File size: 1,136 Bytes
4d585b4 757b45e 4d585b4 |
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 |
import streamlit as st
from diffusers import DDPMScheduler, UNet2DModel
from PIL import Image
import torch
import numpy as np
def generate_image():
scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
scheduler.set_timesteps(50)
sample_size = model.config.sample_size
noise = torch.randn((1, 3, sample_size, sample_size)).to("cuda")
input = noise
for t in scheduler.timesteps:
with torch.no_grad():
noisy_residual = model(input, t).sample()
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
input = prev_noisy_sample
image = (input / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
image = Image.fromarray((image * 255).round().astype("uint8"))
return image
# Streamlit app
st.title("DDPM Image Generation")
st.write("Generating and displaying an image using DDPM.")
# Generate and display the image
generated_image = generate_image()
st.image(generated_image, caption="Generated Image", use_column_width=True)
|