Update ad_gen.py

playground/demos/ad_gen/ad_gen.py

@@ -2,92 +2,57 @@ import random
 import os
 from dotenv import load_dotenv
 from swarms.models import OpenAIChat
-from playground.models.stable_diffusion import StableDiffusion
-from swarms.structs import Agent, SequentialWorkflow
+from swarms.structs import Agent
+from swarms.models.stable_diffusion import StableDiffusion
 
 load_dotenv()
 openai_api_key = os.getenv("OPENAI_API_KEY")
 stability_api_key = os.getenv("STABILITY_API_KEY")
 
 # Initialize the language model and image generation model
-llm = OpenAIChat(
-    openai_api_key=openai_api_key, temperature=0.5, max_tokens=3000
-)
+llm = OpenAIChat(openai_api_key=openai_api_key, temperature=0.5, max_tokens=3000)
 sd_api = StableDiffusion(api_key=stability_api_key)
 
-
-def run_task(description, product_name, agent, **kwargs):
-    full_description = (  # Incorporate product name into the task
-        f"{description} about {product_name}"
-    )
-    result = agent.run(task=full_description, **kwargs)
-    return result
-
-
-# Creative Concept Generator
-class ProductPromptGenerator:
+# Creative Concept Generator for Product Ads
+class ProductAdConceptGenerator:
     def __init__(self, product_name):
         self.product_name = product_name
         self.themes = [
-            "lightning",
-            "sunset",
-            "ice cave",
-            "space",
-            "forest",
-            "ocean",
-            "mountains",
-            "cityscape",
+            "futuristic", "rustic", "luxurious", "minimalistic", "vibrant", "elegant",
+            "retro", "urban", "ethereal", "surreal", "artistic", "tech-savvy",
+            "vintage", "natural", "sophisticated", "playful", "dynamic", "serene", "lasers," "lightning"
         ]
-        self.styles = [
-            "translucent",
-            "floating in mid-air",
-            "expanded into pieces",
-            "glowing",
-            "mirrored",
-            "futuristic",
+        self.contexts = [
+            "in an everyday setting", "in a rave setting", "in an abstract environment", 
+            "in an adventurous context", "surrounded by nature", "in a high-tech setting",
+            "in a historical context", "in a busy urban scene", "in a tranquil and peaceful setting",
+            "against a backdrop of city lights", "in a surreal dreamscape", "in a festive atmosphere",
+            "in a luxurious setting", "in a playful and colorful background", "in an ice cave setting",
+            "in a serene and calm landscape"
         ]
-        self.contexts = ["high realism product ad (extremely creative)"]
 
-    def generate_prompt(self):
+    def generate_concept(self):
         theme = random.choice(self.themes)
-        style = random.choice(self.styles)
         context = random.choice(self.contexts)
-        return f"{theme} inside a {style} {self.product_name}, {context}"
-
+        return f"An ad for {self.product_name} that embodies a {theme} theme {context}"
 
 # User input
-product_name = input(
-    "Enter a product name for ad creation (e.g., 'PS5', 'AirPods', 'Kirkland"
-    " Vodka'): "
-)
+product_name = input("Enter a product name for ad creation (e.g., 'PS5', 'AirPods', 'Kirkland Vodka'): ")
+social_media_platform = input("Enter a social media platform (e.g., 'Facebook', 'Twitter', 'Instagram'): ")
 
 # Generate creative concept
-prompt_generator = ProductPromptGenerator(product_name)
-creative_prompt = prompt_generator.generate_prompt()
-
-# Run tasks using Agent
-concept_flow = Agent(llm=llm, max_loops=1, dashboard=False)
-design_flow = Agent(llm=llm, max_loops=1, dashboard=False)
-copywriting_flow = Agent(llm=llm, max_loops=1, dashboard=False)
+concept_generator = ProductAdConceptGenerator(product_name)
+creative_concept = concept_generator.generate_concept()
 
-# Execute tasks
-concept_result = run_task(
-    "Generate a creative concept", product_name, concept_flow
-)
-design_result = run_task(
-    "Suggest visual design ideas", product_name, design_flow
-)
-copywriting_result = run_task(
-    "Create compelling ad copy for the product photo",
-    product_name,
-    copywriting_flow,
-)
+# Generate product image based on the creative concept
+image_paths = sd_api.run(creative_concept)
 
-# Generate product image
-image_paths = sd_api.run(creative_prompt)
+# Generate ad copy
+ad_copy_agent = Agent(llm=llm, max_loops=1)
+ad_copy_prompt = f"Write a compelling {social_media_platform} ad copy for a product photo showing {product_name} {creative_concept}."
+ad_copy = ad_copy_agent.run(task=ad_copy_prompt)
 
 # Output the results
-print("Creative Concept:", concept_result)
-print("Design Ideas:", design_result)
-print("Ad Copy:", copywriting_result)
+print("Creative Concept:", creative_concept)
 print("Image Path:", image_paths[0] if image_paths else "No image generated")
+print("Ad Copy:", ad_copy)