Update app.py

app.py

@@ -4,33 +4,39 @@ from diffusers import DiffusionPipeline
 from PIL import Image, ImageDraw, ImageFont
 
 # Load the model (make sure to use a model that exists on Hugging Face)
-model = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32).to(device)
 
 def generate_image(caption):
     # Generate the image from the caption
-    image = model(caption).images[0]
-    
-    # Convert to PIL Image for drawing
-    image = image.convert("RGBA")
-    
-    # Create a draw object
-    draw = ImageDraw.Draw(image)
-    
-    # Define font size and color
-    font_size = 40
-    font_color = "white"
-    
-    # Load a font
-    font = ImageFont.load_default()  # You can specify a TTF font file if needed
-    
-    # Calculate text size and position
-    text_width, text_height = draw.textsize(caption, font=font)
-    text_position = ((image.width - text_width) // 2, 10)  # Centered at the top
-    
-    # Draw the text on the image
-    draw.text(text_position, caption, font=font, fill=font_color)
-    
-    return image
+    try:
+        with torch.no_grad():
+            image = model(caption).images[0]
+        
+        # Convert to PIL Image for drawing
+        image = image.convert("RGBA")
+        
+        # Create a draw object
+        draw = ImageDraw.Draw(image)
+        
+        # Define font size and color
+        font_size = 40
+        font_color = "white"
+        
+        # Load a font
+        font = ImageFont.load_default()  # You can specify a TTF font file if needed
+        
+        # Calculate text size and position
+        text_width, text_height = draw.textsize(caption, font=font)
+        text_position = ((image.width - text_width) // 2, 10)  # Centered at the top
+        
+        # Draw the text on the image
+        draw.text(text_position, caption, font=font, fill=font_color)
+        
+        return image
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return None
 
 # Create the Gradio interface
 with gr.Blocks() as demo: