File size: 5,402 Bytes
746bf27 0b1e3c3 746bf27 19a39c5 746bf27 19a39c5 746bf27 19a39c5 746bf27 19a39c5 746bf27 |
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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image, ImageDraw, ImageFont
import os
# 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
print("Loading Stable Diffusion model...")
model_id = "CompVis/stable-diffusion-v1-4" # You can change to v2 models if needed
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
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)
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.
"""
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
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()
# 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
return image
# Main function to generate advertisement
def generate_ad(brand_title, tagline, cta, brand_logo=None, product_image=None, custom_prompt=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.
"""
# 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
generated_image = generate_image(prompt)
# Add product and CTA text to the image
final_image = add_text_to_image(generated_image, brand_title, tagline, cta)
# Optionally, add brand logo and product image if provided
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
iface = gr.Interface(
fn=generate_ad,
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)"),
],
outputs=gr.Image(label="Generated Advertisement"),
)
# Launch the Gradio Interface
iface.launch()
|