Add custom prompt

app.py

@@ -4,32 +4,14 @@ 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
+import torch
+from diffusers import StableDiffusionPipeline
 
-    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
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
-# Streamlit app
-st.title("DDPM Image Generation")
-st.write("Generating and displaying an image using DDPM.")
+prompt = st.text_input('Insert here your prompt')
 
-# Generate and display the image
-generated_image = generate_image()
-st.image(generated_image, caption="Generated Image", use_column_width=True)
+image = pipe(prompt).images[0]