Spaces:
Sleeping
Sleeping
File size: 4,328 Bytes
e65bcdc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import torch
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=512, width=512):
model_id = "CompVis/stable-diffusion-v1-4"
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.")
# Use half-precision and reduce model load time
pipeline = StableDiffusionPipeline.from_pretrained(
model_id,
use_auth_token=HUGGINGFACE_API_KEY,
torch_dtype=torch.float16,
revision="fp16"
)
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):
"""
Add clean and sharp text to the generated image.
"""
draw = ImageDraw.Draw(image)
try:
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:
product_font = tagline_font = cta_font = ImageFont.load_default()
# 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_advertisement(brand_title, tagline, cta, custom_prompt=None, brand_logo=None, product_image=None):
"""
Generate advertisement image with text overlay and optional logo/product image.
"""
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 base image using Stable Diffusion
generated_image = generate_image(prompt)
# Overlay text (brand title, tagline, and CTA)
final_image = add_text_to_image(generated_image, brand_title, tagline, cta)
# 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'))
if product_image:
product = Image.open(product_image).resize((300, 300))
final_image.paste(product, (250, 350), product.convert('RGBA'))
return final_image
# Gradio Interface
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="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(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 interface
if __name__ == "__main__":
interface.launch()
|