-
Notifications
You must be signed in to change notification settings - Fork 44.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(store): Generate AI images for store submissions (#9090)
Allow generating ai images for store submissions
- Loading branch information
Showing
6 changed files
with
275 additions
and
25 deletions.
There are no files selected for viewing
94 changes: 94 additions & 0 deletions
94
autogpt_platform/backend/backend/server/v2/store/image_gen.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,94 @@ | ||
import io | ||
import logging | ||
from enum import Enum | ||
|
||
import replicate | ||
import replicate.exceptions | ||
import requests | ||
from replicate.helpers import FileOutput | ||
|
||
from backend.data.graph import Graph | ||
from backend.util.settings import Settings | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
class ImageSize(str, Enum): | ||
LANDSCAPE = "1024x768" | ||
|
||
|
||
class ImageStyle(str, Enum): | ||
DIGITAL_ART = "digital art" | ||
|
||
|
||
async def generate_agent_image(agent: Graph) -> io.BytesIO: | ||
""" | ||
Generate an image for an agent using Flux model via Replicate API. | ||
Args: | ||
agent (Graph): The agent to generate an image for | ||
Returns: | ||
io.BytesIO: The generated image as bytes | ||
""" | ||
try: | ||
settings = Settings() | ||
|
||
if not settings.secrets.replicate_api_key: | ||
raise ValueError("Missing Replicate API key in settings") | ||
|
||
# Construct prompt from agent details | ||
prompt = f"App store image for AI agent that gives a cool visual representation of what the agent does: - {agent.name} - {agent.description}" | ||
|
||
# Set up Replicate client | ||
client = replicate.Client(api_token=settings.secrets.replicate_api_key) | ||
|
||
# Model parameters | ||
input_data = { | ||
"prompt": prompt, | ||
"width": 1024, | ||
"height": 768, | ||
"aspect_ratio": "4:3", | ||
"output_format": "jpg", | ||
"output_quality": 90, | ||
"num_inference_steps": 30, | ||
"guidance": 3.5, | ||
"negative_prompt": "blurry, low quality, distorted, deformed", | ||
"disable_safety_checker": True, | ||
} | ||
|
||
try: | ||
# Run model | ||
output = client.run("black-forest-labs/flux-pro", input=input_data) | ||
|
||
# Depending on the model output, extract the image URL or bytes | ||
# If the output is a list of FileOutput or URLs | ||
if isinstance(output, list) and output: | ||
if isinstance(output[0], FileOutput): | ||
image_bytes = output[0].read() | ||
else: | ||
# If it's a URL string, fetch the image bytes | ||
result_url = output[0] | ||
response = requests.get(result_url) | ||
response.raise_for_status() | ||
image_bytes = response.content | ||
elif isinstance(output, FileOutput): | ||
image_bytes = output.read() | ||
elif isinstance(output, str): | ||
# Output is a URL | ||
response = requests.get(output) | ||
response.raise_for_status() | ||
image_bytes = response.content | ||
else: | ||
raise RuntimeError("Unexpected output format from the model.") | ||
|
||
return io.BytesIO(image_bytes) | ||
|
||
except replicate.exceptions.ReplicateError as e: | ||
if e.status == 401: | ||
raise RuntimeError("Invalid Replicate API token") from e | ||
raise RuntimeError(f"Replicate API error: {str(e)}") from e | ||
|
||
except Exception as e: | ||
logger.exception("Failed to generate agent image") | ||
raise RuntimeError(f"Image generation failed: {str(e)}") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.