Update app.py

app.py

@@ -4,41 +4,28 @@ from PIL import Image, ImageDraw, ImageFont
 import os
 import gradio as gr  # Import Gradio for the interface
 
-# Function to Generate Image
-def generate_image(prompt, height=1024, width=1024):
-    """
-    Generate an image using Stable Diffusion from Hugging Face.
-
-    Args:
-        prompt (str): The text prompt to guide image generation.
-        height (int): Height of the generated image.
-        width (int): Width of the generated image.
 
-    Returns:
-        PIL Image: Generated image.
-    """
-    # Load Hugging Face access token from environment variables
+#Function to generate image
+def generate_image(prompt, height=512, width=512):
+    model_id = "stabilityai/stable-diffusion-2-1"
     HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
+
     if not HUGGINGFACE_API_KEY:
         raise ValueError("Hugging Face API key is not set. Export it as HUGGINGFACE_API_KEY.")
-    
-    # Load the Stable Diffusion model
-    model_id = "stabilityai/stable-diffusion-2-1"
+
+    # Use half-precision and reduce model load time
     pipeline = StableDiffusionPipeline.from_pretrained(
         model_id, 
         use_auth_token=HUGGINGFACE_API_KEY,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        torch_dtype=torch.float16,
+        revision="fp16"
     )
     
-    # Move to GPU if available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline = pipeline.to(device)
-    
-    # Generate the image
+    pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
     image = pipeline(prompt, height=height, width=width).images[0]
-    
     return image
 
+
 # Function to Add Text to Image
 def add_text_to_image(image, product_name, tagline, cta_text, font_size=50):
     """