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()