From 2c10018f58dc91b9e2c05bcf54b25bf032b12336 Mon Sep 17 00:00:00 2001 From: pliny <133052465+elder-plinius@users.noreply.github.com> Date: Mon, 27 Nov 2023 08:23:23 -0800 Subject: [PATCH 1/3] Update stable_diffusion.py --- playground/models/stable_diffusion.py | 91 +++++++++++++++------------ 1 file changed, 52 insertions(+), 39 deletions(-) diff --git a/playground/models/stable_diffusion.py b/playground/models/stable_diffusion.py index 3bb77c39..f45d5892 100644 --- a/playground/models/stable_diffusion.py +++ b/playground/models/stable_diffusion.py @@ -1,6 +1,7 @@ -import os import base64 +import os import requests +import uuid from dotenv import load_dotenv from typing import List @@ -8,41 +9,67 @@ load_dotenv() class StableDiffusion: """ - A class to interact with the Stable Diffusion API for image generation. + 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 of the Stable Diffusion API. + The host URL for the Stable Diffusion API. engine_id : str - The ID of the Stable Diffusion engine. - headers : dict - The headers for the API request. + 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 generated images will be saved. + Directory where the 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. + __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"): + 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): """ - Initializes the StableDiffusion class with the provided API key and host. + 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 of the Stable Diffusion API. Default is "https://api.stability.ai". + 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", @@ -51,45 +78,35 @@ class StableDiffusion: 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]: + def run(self, task: str) -> List[str]: """ - Generates images based on a text prompt. + Generates an image based on a given text prompt. Parameters: ----------- - prompt : str + task : 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. + A list of file paths where the generated images are saved. Raises: ------- Exception: - If the API response is not 200 (OK). + 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": prompt}], - "cfg_scale": cfg_scale, - "height": height, - "width": width, - "samples": samples, - "steps": steps, + "text_prompts": [{"text": task}], + "cfg_scale": self.cfg_scale, + "height": self.height, + "width": self.width, + "samples": self.samples, + "steps": self.steps, }, ) @@ -99,14 +116,10 @@ class StableDiffusion: 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") + 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 - -# Usage example: -# sd = StableDiffusion("your-api-key") -# images = sd.generate_image("A scenic landscape with mountains") -# print(images) From ae4d53e3e595eddd5d9bb6ea7388c4f85a3c5bef Mon Sep 17 00:00:00 2001 From: pliny <133052465+elder-plinius@users.noreply.github.com> Date: Mon, 27 Nov 2023 08:24:31 -0800 Subject: [PATCH 2/3] Create ad_gen.py --- playground/demos/ad_gen/ad_gen.py | 60 +++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 playground/demos/ad_gen/ad_gen.py 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") From 59e7573d8b067c04dbcbda95f2311a6a36ad6d45 Mon Sep 17 00:00:00 2001 From: pliny <133052465+elder-plinius@users.noreply.github.com> Date: Mon, 27 Nov 2023 11:14:53 -0800 Subject: [PATCH 3/3] Rename stable_diffusion.py to stable_diffusion.py --- {playground => swarms}/models/stable_diffusion.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {playground => swarms}/models/stable_diffusion.py (100%) diff --git a/playground/models/stable_diffusion.py b/swarms/models/stable_diffusion.py similarity index 100% rename from playground/models/stable_diffusion.py rename to swarms/models/stable_diffusion.py