Update app.py

app.py

@@ -1,11 +1,20 @@
-import subprocess
-subprocess.run(["pip", "install", "diffusers"])
+from diffusers import DiffusionPipeline
+import torch
 
-subprocess.check_call([sys.executable, "-m", "pip", "install", "torch", "torchvision", "torchaudio"])
+# Load the pipeline
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)
 
-from diffusers import DiffusionPipeline
-pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
-prompt = input("What do you want to see?")
+# If running on CUDA, ensure the model is moved to GPU for faster inference
+pipe.to("cuda")
+
+# Get the prompt from the user
+prompt = input("What do you want to see? ")
+
+# Generate the image
 image = pipe(prompt).images[0]
-image
-print("How cool is this!")
+
+# Display the image
+image.show()
+
+# Print a message
+print("How cool is this!")