Update app.py

app.py

@@ -6,13 +6,15 @@ from PIL import Image
 
 import spaces
 
+import torch
 from transformers import pipeline
 
 # Initialize Caption Generation Model
-get_caption = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base")
+get_caption = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)
 
 # Initialize Image Generation Model
-generate_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+generate_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", )
+pipe = generate_pipeline.to("cuda")
 
 def captioner(input: Image.Image) -> str:
     """
@@ -37,7 +39,7 @@ def generate(prompt: str) -> Image.Image:
     Returns:
         Image.Image: The generated image corresponding to the given prompt.
     """
-    return generate_pipeline(prompt).images[0]
+    return pipe(prompt).images[0]
 
 @spaces.GPU(duration=300)
 def caption_and_generate(image: Image.Image) -> list: