You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
94 lines
2.7 KiB
94 lines
2.7 KiB
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
|
|
|
|
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
|
|
)
|
|
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:
|
|
def __init__(self, product_name):
|
|
self.product_name = product_name
|
|
self.themes = [
|
|
"lightning",
|
|
"sunset",
|
|
"ice cave",
|
|
"space",
|
|
"forest",
|
|
"ocean",
|
|
"mountains",
|
|
"cityscape",
|
|
]
|
|
self.styles = [
|
|
"translucent",
|
|
"floating in mid-air",
|
|
"expanded into pieces",
|
|
"glowing",
|
|
"mirrored",
|
|
"futuristic",
|
|
]
|
|
self.contexts = ["high realism product ad (extremely creative)"]
|
|
|
|
def generate_prompt(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}"
|
|
|
|
|
|
# User input
|
|
product_name = input(
|
|
"Enter a product name for ad creation (e.g., 'PS5', 'AirPods', 'Kirkland"
|
|
" Vodka'): "
|
|
)
|
|
|
|
# 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)
|
|
|
|
# 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
|
|
image_paths = sd_api.run(creative_prompt)
|
|
|
|
# Output the results
|
|
print("Creative Concept:", concept_result)
|
|
print("Design Ideas:", design_result)
|
|
print("Ad Copy:", copywriting_result)
|
|
print("Image Path:", image_paths[0] if image_paths else "No image generated")
|