diff --git a/playground/demos/ad_gen/ad_gen.py b/playground/demos/ad_gen/ad_gen.py new file mode 100644 index 00000000..0edde900 --- /dev/null +++ b/playground/demos/ad_gen/ad_gen.py @@ -0,0 +1,60 @@ +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 Flow, 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, flow, **kwargs): + full_description = f"{description} about {product_name}" # Incorporate product name into the task + result = flow.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 Flow +concept_flow = Flow(llm=llm, max_loops=1, dashboard=False) +design_flow = Flow(llm=llm, max_loops=1, dashboard=False) +copywriting_flow = Flow(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") diff --git a/playground/models/stable_diffusion.py b/playground/models/stable_diffusion.py deleted file mode 100644 index 3bb77c39..00000000 --- a/playground/models/stable_diffusion.py +++ /dev/null @@ -1,112 +0,0 @@ -import os -import base64 -import requests -from dotenv import load_dotenv -from typing import List - -load_dotenv() - -class StableDiffusion: - """ - A class to interact with the Stable Diffusion API for image generation. - - Attributes: - ----------- - api_key : str - The API key for accessing the Stable Diffusion API. - api_host : str - The host URL of the Stable Diffusion API. - engine_id : str - The ID of the Stable Diffusion engine. - headers : dict - The headers for the API request. - output_dir : str - Directory where generated images will be saved. - - Methods: - -------- - generate_image(prompt: str, cfg_scale: int, height: int, width: int, samples: int, steps: int) -> List[str]: - Generates images based on a text prompt and returns a list of file paths to the generated images. - """ - - def __init__(self, api_key: str, api_host: str = "https://api.stability.ai"): - """ - Initializes the StableDiffusion class with the provided API key and host. - - Parameters: - ----------- - api_key : str - The API key for accessing the Stable Diffusion API. - api_host : str - The host URL of the Stable Diffusion API. Default is "https://api.stability.ai". - """ - self.api_key = api_key - self.api_host = api_host - self.engine_id = "stable-diffusion-v1-6" - self.headers = { - "Authorization": f"Bearer {self.api_key}", - "Content-Type": "application/json", - "Accept": "application/json" - } - self.output_dir = "images" - os.makedirs(self.output_dir, exist_ok=True) - - def generate_image(self, prompt: str, cfg_scale: int = 7, height: int = 1024, width: int = 1024, samples: int = 1, steps: int = 30) -> List[str]: - """ - Generates images based on a text prompt. - - Parameters: - ----------- - prompt : str - The text prompt based on which the image will be generated. - cfg_scale : int - CFG scale parameter for image generation. Default is 7. - height : int - Height of the generated image. Default is 1024. - width : int - Width of the generated image. Default is 1024. - samples : int - Number of images to generate. Default is 1. - steps : int - Number of steps for the generation process. Default is 30. - - Returns: - -------- - List[str]: - A list of paths to the generated images. - - Raises: - ------- - Exception: - If the API response is not 200 (OK). - """ - response = requests.post( - f"{self.api_host}/v1/generation/{self.engine_id}/text-to-image", - headers=self.headers, - json={ - "text_prompts": [{"text": prompt}], - "cfg_scale": cfg_scale, - "height": height, - "width": width, - "samples": samples, - "steps": steps, - }, - ) - - if response.status_code != 200: - raise Exception(f"Non-200 response: {response.text}") - - data = response.json() - image_paths = [] - for i, image in enumerate(data["artifacts"]): - image_path = os.path.join(self.output_dir, f"v1_txt2img_{i}.png") - with open(image_path, "wb") as f: - f.write(base64.b64decode(image["base64"])) - image_paths.append(image_path) - - return image_paths - -# Usage example: -# sd = StableDiffusion("your-api-key") -# images = sd.generate_image("A scenic landscape with mountains") -# print(images) diff --git a/swarms/models/stable_diffusion.py b/swarms/models/stable_diffusion.py new file mode 100644 index 00000000..f45d5892 --- /dev/null +++ b/swarms/models/stable_diffusion.py @@ -0,0 +1,125 @@ +import base64 +import os +import requests +import uuid +from dotenv import load_dotenv +from typing import List + +load_dotenv() + +class StableDiffusion: + """ + A class to interact with the Stable Diffusion API for generating images from text prompts. + + Attributes: + ----------- + api_key : str + The API key for accessing the Stable Diffusion API. + api_host : str + The host URL for the Stable Diffusion API. + engine_id : str + The engine ID for the Stable Diffusion API. + cfg_scale : int + Configuration scale for image generation. + height : int + The height of the generated image. + width : int + The width of the generated image. + samples : int + The number of samples to generate. + steps : int + The number of steps for the generation process. + output_dir : str + Directory where the generated images will be saved. + + Methods: + -------- + __init__(self, api_key: str, api_host: str, cfg_scale: int, height: int, width: int, samples: int, steps: int): + Initializes the StableDiffusion instance with provided parameters. + + generate_image(self, task: str) -> List[str]: + Generates an image based on the provided text prompt and returns the paths of the saved images. + """ + + def __init__(self, api_key: str, api_host: str = "https://api.stability.ai", cfg_scale: int = 7, height: int = 1024, width: int = 1024, samples: int = 1, steps: int = 30): + """ + Initialize the StableDiffusion class with API configurations. + + Parameters: + ----------- + api_key : str + The API key for accessing the Stable Diffusion API. + api_host : str + The host URL for the Stable Diffusion API. + cfg_scale : int + Configuration scale for image generation. + height : int + The height of the generated image. + width : int + The width of the generated image. + samples : int + The number of samples to generate. + steps : int + The number of steps for the generation process. + """ + self.api_key = api_key + self.api_host = api_host + self.engine_id = "stable-diffusion-v1-6" + self.cfg_scale = cfg_scale + self.height = height + self.width = width + self.samples = samples + self.steps = steps + self.headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + "Accept": "application/json" + } + self.output_dir = "images" + os.makedirs(self.output_dir, exist_ok=True) + + def run(self, task: str) -> List[str]: + """ + Generates an image based on a given text prompt. + + Parameters: + ----------- + task : str + The text prompt based on which the image will be generated. + + Returns: + -------- + List[str]: + A list of file paths where the generated images are saved. + + Raises: + ------- + Exception: + If the API request fails and returns a non-200 response. + """ + response = requests.post( + f"{self.api_host}/v1/generation/{self.engine_id}/text-to-image", + headers=self.headers, + json={ + "text_prompts": [{"text": task}], + "cfg_scale": self.cfg_scale, + "height": self.height, + "width": self.width, + "samples": self.samples, + "steps": self.steps, + }, + ) + + if response.status_code != 200: + raise Exception(f"Non-200 response: {response.text}") + + data = response.json() + image_paths = [] + for i, image in enumerate(data["artifacts"]): + unique_id = uuid.uuid4() # Generate a unique identifier + image_path = os.path.join(self.output_dir, f"{unique_id}_v1_txt2img_{i}.png") + with open(image_path, "wb") as f: + f.write(base64.b64decode(image["base64"])) + image_paths.append(image_path) + + return image_paths