ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #205 from elder-plinius/master

Browse Source 
stable-diffusion and ad-gen demo
pull/209/head
Eternal Reclaimer 1 year ago committed by GitHub
parent 7705973569 e113a77dc8
commit 3d3dddaf0c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 185 additions and 112 deletions
Show Stats Download Patch File Download Diff File

60
playground/demos/ad_gen/ad_gen.py
Unescape View File

@ -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")

112
playground/models/stable_diffusion.py
Unescape View File

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

125
swarms/models/stable_diffusion.py
Unescape View File

@ -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
Write Preview
Loading…
Cancel
Save