Update app.py

app.py

@@ -1,26 +1,9 @@
-import gradio as gr
-from diffusers import StableDiffusionPipeline
 import torch
+from diffusers import StableDiffusionPipeline
 
-# Load the pre-trained model from Hugging Face (CompVis is the original model repository)
-pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
-
-# Ensure that the model runs on CPU
-pipe = pipe.to("cpu")  # Use CPU instead of CUDA (GPU)
-
-def generate_image(prompt):
-    # Generate image based on the user's prompt
-    image = pipe(prompt).images[0]
-
-    return image
-
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=generate_image
-,        # Function to call when the user submits a prompt
-    inputs="text",            # Input is a text field for the user to type their prompt
-    outputs="image",          # Output is an image, generated by the model
-    live=True                  # Make the interface responsive in real-time
-)
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
-iface.launch()
+prompt = "a photo of an astronaut riding a horse on mars"
+image = pipe(prompt).images[0]
+image