File size: 4,687 Bytes
746bf27 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 0b1e3c3 9462d72 0b1e3c3 9462d72 746bf27 9462d72 746bf27 19a39c5 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 19a39c5 9462d72 746bf27 9462d72 746bf27 9462d72 746bf27 9462d72 |
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 114 115 116 117 118 119 120 121 122 123 124 |
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=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
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"
pipeline = StableDiffusionPipeline.from_pretrained(
model_id,
use_auth_token=HUGGINGFACE_API_KEY,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = pipeline.to(device)
# Generate the image
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()
|