Update app.py

app.py

@@ -1,8 +1,8 @@
 import torch
-import gradio as gr
 from diffusers import StableDiffusionPipeline
 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):
@@ -23,8 +23,7 @@ def generate_image(prompt, height=1024, width=1024):
         raise ValueError("Hugging Face API key is not set. Export it as HUGGINGFACE_API_KEY.")
     
     # Load the Stable Diffusion model
-    print("Loading Stable Diffusion model...")
-    model_id = "CompVis/stable-diffusion-v1-4"  # You can change to v2 models if needed
+    model_id = "stabilityai/stable-diffusion-2-1"
     pipeline = StableDiffusionPipeline.from_pretrained(
         model_id, 
         use_auth_token=HUGGINGFACE_API_KEY,
@@ -36,84 +35,47 @@ def generate_image(prompt, height=1024, width=1024):
     pipeline = pipeline.to(device)
     
     # Generate the image
-    print(f"Generating image for prompt: '{prompt}'...")
     image = pipeline(prompt, height=height, width=width).images[0]
     
     return image
 
-# Function to Add Emphasized Text to the Image (Product Name, Tagline, CTA)
+# Function to Add Text to Image
 def add_text_to_image(image, product_name, tagline, cta_text, font_size=50):
     """
-    Add clean and sharp text (product name, tagline, and CTA) to the generated image with emphasis.
-    
-    Args:
-        image (PIL Image): Generated image to add text to.
-        product_name (str): Product name to be emphasized.
-        tagline (str): Tagline to be emphasized.
-        cta_text (str): Call to action text to be emphasized.
-        font_size (int): The font size.
-    
-    Returns:
-        PIL Image: Image with added text.
+    Add clean and sharp text to the generated image.
     """
     draw = ImageDraw.Draw(image)
-    
     try:
-        # Load font for product name, tagline, and CTA
-        product_font = ImageFont.truetype("arial.ttf", font_size + 20)  # Larger for product name
-        tagline_font = ImageFont.truetype("arial.ttf", font_size)       # Slightly smaller for tagline
-        cta_font = ImageFont.truetype("arial.ttf", font_size - 10)      # Smaller for CTA
+        product_font = ImageFont.truetype("arial.ttf", font_size + 20)
+        tagline_font = ImageFont.truetype("arial.ttf", font_size)
+        cta_font = ImageFont.truetype("arial.ttf", font_size - 10)
     except IOError:
-        # Fallback font in case "arial.ttf" is not available
-        product_font = ImageFont.load_default()
-        tagline_font = ImageFont.load_default()
-        cta_font = ImageFont.load_default()
+        product_font = tagline_font = cta_font = ImageFont.load_default()
 
-    # Define positions for text
-    product_name_position = (50, 50)
-    tagline_position = (50, 150)
-    cta_position = (50, 250)
-
-    # Add product name in large font
-    draw.text(product_name_position, product_name, font=product_font, fill="white")
-    
-    # Add tagline below the product name
-    draw.text(tagline_position, tagline, font=tagline_font, fill="white")
-    
-    # Add CTA below tagline
-    draw.text(cta_position, cta_text, font=cta_font, fill="gold")  # CTA in gold to stand out
+    # Add product name, tagline, and CTA to the image
+    draw.text((50, 50), product_name, font=product_font, fill="white")
+    draw.text((50, 150), tagline, font=tagline_font, fill="white")
+    draw.text((50, 250), cta_text, font=cta_font, fill="gold")
     
     return image
 
 # Main function to generate advertisement
-def generate_ad(brand_title, tagline, cta, brand_logo=None, product_image=None, custom_prompt=None):
+def generate_advertisement(brand_title, tagline, cta, custom_prompt=None, brand_logo=None, product_image=None):
     """
-    Generate an advertisement image with the provided details.
-    
-    Args:
-        brand_title (str): Brand title for the advertisement.
-        tagline (str): Tagline for the advertisement.
-        cta (str): Call to action text.
-        brand_logo (file, optional): Brand logo image (optional).
-        product_image (file, optional): Product image (optional).
-        custom_prompt (str, optional): Custom prompt for image generation.
-
-    Returns:
-        PIL Image: Final advertisement image.
+    Generate advertisement image with text overlay and optional logo/product image.
     """
-    # Prepare the final prompt
     prompt = custom_prompt if custom_prompt else (
         f"An elegant advertisement for {brand_title}, featuring gold and white tones, "
         f"with a radiant and premium look. Product focus and beautiful typography for '{tagline}'."
     )
 
-    # Generate the image
+    # Generate the base image using Stable Diffusion
     generated_image = generate_image(prompt)
 
-    # Add product and CTA text to the image
+    # Overlay text (brand title, tagline, and CTA)
     final_image = add_text_to_image(generated_image, brand_title, tagline, cta)
 
-    # Optionally, add brand logo and product image if provided
+    # Optionally add logo and product images
     if brand_logo:
         logo = Image.open(brand_logo).resize((150, 150))
         final_image.paste(logo, (50, 350), logo.convert('RGBA'))
@@ -125,18 +87,37 @@ def generate_ad(brand_title, tagline, cta, brand_logo=None, product_image=None,
     return final_image
 
 # Gradio Interface
-iface = gr.Interface(
-    fn=generate_ad,
+def gradio_interface(brand_title, tagline, cta, custom_prompt, brand_logo, product_image):
+    """
+    Gradio interface wrapper to call the advertisement generation function.
+    """
+    # Generate the ad
+    ad_image = generate_advertisement(
+        brand_title=brand_title,
+        tagline=tagline,
+        cta=cta,
+        custom_prompt=custom_prompt,
+        brand_logo=brand_logo.name if brand_logo else None,
+        product_image=product_image.name if product_image else None
+    )
+    return ad_image
+
+# Gradio UI Layout
+interface = gr.Interface(
+    fn=gradio_interface,
     inputs=[
-        gr.Textbox(label="Brand Title", placeholder="Enter brand title"),
-        gr.Textbox(label="Tagline", placeholder="Enter tagline"),
-        gr.Textbox(label="CTA", placeholder="Enter Call to Action"),
-        gr.Image(type="file", label="Brand Logo (optional)", optional=True),
-        gr.Image(type="file", label="Product Image (optional)", optional=True),
-        gr.Textbox(label="Custom Prompt (optional)", placeholder="Enter a custom prompt (or leave empty)"),
+        gr.Textbox(label="Brand Title", placeholder="e.g., GlowWell Skin Serum"),
+        gr.Textbox(label="Tagline", placeholder="e.g., Radiance Redefined"),
+        gr.Textbox(label="Call to Action (CTA)", placeholder="e.g., Shop Now"),
+        gr.Textbox(label="Custom Prompt (Optional)", placeholder="Describe your ad style..."),
+        gr.File(label="Brand Logo (Optional)"),
+        gr.File(label="Product Image (Optional)")
     ],
-    outputs=gr.Image(label="Generated Advertisement"),
+    outputs=gr.Image(type="pil", label="Generated Advertisement"),
+    title="AI-Powered Advertisement Generator",
+    description="Generate stunning advertisements using Stable Diffusion. Provide brand details, and optionally upload images or add custom descriptions to create your perfect ad."
 )
 
-# Launch the Gradio Interface
-iface.launch()
+# Launch the interface
+if __name__ == "__main__":
+    interface.launch()