Generate images code

app.py

@@ -1,4 +1,35 @@
 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)
+
 
-x = st.slider("Select a value")
-st.write('Test number:', x)

requirements.txt

@@ -2,4 +2,5 @@ transformers
 streamlit
 diffusers
 pandas
-plotly
+plotly
+torch