Skip to content

Commit f624214

Browse files
celesteanglmcelesteanglm
authored
Merge pull request #54 from chinesepowered/plugin_image
imagegen plugin
2 parents 0d1c751 + bdc5f53 commit f624214

5 files changed

Lines changed: 291 additions & 0 deletions

File tree

‎plugins/imagegen/README.md‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# Image Generator Plugin for GAME SDK
2+
3+
TogetherAI has a free flux schnell image generator, which we'll use here because it's free.
4+
Can change endpoints, params, etc to suit your needs.
5+
6+
## Features
7+
- Generate images based on prompt
8+
- Receive images as temporary URL or B64 objects
9+
10+
## Available Functions
11+
12+
1. `generate_image(prompt: str)` - Generates image based on prompt
13+
14+
## Setup and configuration
15+
1. Set the following environment variables:
16+
- `TOGETHER_API_KEY`: Create an API key by [creating an account](https://together.ai).
17+
18+
2. Import and initialize the plugin to use in your worker:
19+
```python
20+
from plugins.imagegen.imagegen_plugin import ImageGenPlugin
21+
22+
imagegen_plugin = ImageGenPlugin(
23+
api_key=os.environ.get("TOGETHER_API_KEY"),
24+
)
25+
```
26+
27+
**Basic worker example:**
28+
```python
29+
def get_state_fn(function_result: FunctionResult, current_state: dict) -> dict:
30+
"""
31+
Update state based on the function results
32+
"""
33+
init_state = {}
34+
35+
if current_state is None:
36+
return init_state
37+
38+
# Update state with the function result info
39+
current_state.update(function_result.info)
40+
41+
return current_state
42+
43+
generate_image_worker = Worker(
44+
api_key=os.environ.get("TOGETHER_API_KEY"),
45+
description="Worker for generating AI images based on prompt",
46+
get_state_fn=get_state_fn,
47+
action_space=[
48+
imagegen_plugin.get_function("generate_image"),
49+
],
50+
)
51+
52+
generate_image_worker.run("Cute anime girl with twitter logo")
53+
```
54+
55+
## Running examples
56+
57+
To run the examples showcased in the plugin's directory, follow these steps:
58+
59+
1. Install dependencies:
60+
```
61+
poetry install
62+
```
63+
64+
2. Set up environment variables:
65+
```
66+
export GAME_API_KEY="your-game-api-key"
67+
export TOGETHER_API_KEY="your-together-api-key"
68+
```
69+
70+
3. Run example scripts:
71+
72+
Example worker:
73+
```
74+
poetry run python ./examples/example-worker.py
75+
```
76+
77+
Example agent:
78+
```
79+
poetry run python ./examples/example-agent.py
80+
```

‎plugins/imagegen/__init__.py‎

Whitespace-only changes.

‎plugins/imagegen/examples/example-worker.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import os
2+
from game_sdk.game.worker import Worker
3+
from game_sdk.game.custom_types import FunctionResult
4+
from imagegen_game_sdk.imagegen_plugin import ImageGenPlugin
5+
6+
def get_state_fn(function_result: FunctionResult, current_state: dict) -> dict:
7+
"""
8+
Update state based on the function results
9+
"""
10+
init_state = {}
11+
12+
if current_state is None:
13+
return init_state
14+
15+
# Update state with the function result info
16+
current_state.update(function_result.info)
17+
18+
return current_state
19+
20+
21+
imagegen_plugin = ImageGenPlugin(
22+
api_key=os.environ.get("TOGETHER_API_KEY"),
23+
)
24+
25+
# Create worker
26+
imagegen_worker = Worker(
27+
api_key=os.environ.get("GAME_API_KEY"),
28+
description="Worker specialized in using AI image generator to generate images",
29+
get_state_fn=get_state_fn,
30+
action_space=[
31+
imagegen_plugin.get_function("generate_image"),
32+
],
33+
)
34+
35+
# # Run example query
36+
queries = [
37+
"Cute anime character with Twitter logo on outfit",
38+
]
39+
for query in queries:
40+
print("-" * 100)
41+
print(f"Query: {query}")
42+
imagegen_worker.run(query)

‎plugins/imagegen/imagegen_game_sdk/imagegen_plugin.py‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
from typing import Dict, List, Optional, Tuple
2+
from game_sdk.game.custom_types import Argument, Function, FunctionResultStatus
3+
import requests
4+
import os
5+
6+
DEFAULT_BASE_API_URL = "https://api.together.xyz/v1/images/generations"
7+
8+
9+
class ImageGenPlugin:
10+
"""
11+
AI Image Generation plugin using Together.ai API.
12+
13+
Requires:
14+
- Together.ai API key
15+
16+
Example:
17+
client = ImageGenPlugin(
18+
api_key="your-together-api-key",
19+
api_url="https://api.together.xyz/v1/images/generations",
20+
)
21+
22+
generate_image_fn = client.get_function("generate_image")
23+
"""
24+
def __init__(
25+
self,
26+
api_key: Optional[str] = os.environ.get("TOGETHER_API_KEY"),
27+
api_url: Optional[str] = DEFAULT_BASE_API_URL,
28+
):
29+
self.api_key = api_key
30+
self.api_url = api_url
31+
32+
# Available client functions
33+
self._functions: Dict[str, Function] = {
34+
"generate_image": Function(
35+
fn_name="generate_image",
36+
fn_description="Generates AI generated image based on prompt.",
37+
args=[
38+
Argument(
39+
name="prompt",
40+
description="The prompt for image generation model. Example: A dog in the park",
41+
type="string",
42+
),
43+
Argument(
44+
name="width",
45+
description="Width of generated image, up to 1440 px. Default should be 1024 unless other sizes specifically needed.",
46+
type="int",
47+
),
48+
Argument(
49+
name="height",
50+
description="Height of generated image, up to 1440 px. Default should be 1024 unless other sizes specifically needed.",
51+
type="int",
52+
),
53+
],
54+
hint="This function is used to generate an AI image based on prompt",
55+
executable=self.generate_image,
56+
),
57+
}
58+
59+
@property
60+
def available_functions(self) -> List[str]:
61+
"""Get list of available function names."""
62+
return list(self._functions.keys())
63+
64+
def get_function(self, fn_name: str) -> Function:
65+
"""
66+
Get a specific function by name.
67+
68+
Args:
69+
fn_name: Name of the function to retrieve
70+
71+
Raises:
72+
ValueError: If function name is not found
73+
74+
Returns:
75+
Function object
76+
"""
77+
if fn_name not in self._functions:
78+
raise ValueError(
79+
f"Function '{fn_name}' not found. Available functions: {', '.join(self.available_functions)}"
80+
)
81+
return self._functions[fn_name]
82+
83+
def generate_image(self, prompt: str, width: int = 1024, height: int = 1024, **kwargs) -> str:
84+
"""Generate image based on prompt.
85+
86+
Returns:
87+
str URL of image (need to save since temporal)
88+
"""
89+
# API endpoint for image generation
90+
url = DEFAULT_BASE_API_URL
91+
92+
# Prepare headers for the request
93+
headers = {
94+
"Authorization": f"Bearer {self.api_key}",
95+
"Content-Type": "application/json",
96+
}
97+
98+
# Prepare request payload
99+
payload = {
100+
"model": "black-forest-labs/FLUX.1-schnell-Free",
101+
"prompt": prompt,
102+
"width": width,
103+
"height": height,
104+
"steps": 1,
105+
"n": 1,
106+
"response_format": "url",
107+
}
108+
109+
try:
110+
# Make the API request
111+
response = requests.post(self.api_url, headers=headers, json=payload)
112+
response.raise_for_status()
113+
114+
# Extract the image URL from the response
115+
response_data = response.json()
116+
image_url = response_data["data"][0]["url"]
117+
118+
return (
119+
FunctionResultStatus.DONE,
120+
f"The generated image is: {image_url}",
121+
{
122+
"prompt": prompt,
123+
"image_url": image_url,
124+
},
125+
)
126+
except Exception as e:
127+
print(f"An error occurred while generating image: {str(e)}")
128+
return (
129+
FunctionResultStatus.FAILED,
130+
f"An error occurred while while generating image: {str(e)}",
131+
{
132+
"prompt": prompt,
133+
},
134+
)

‎plugins/imagegen/pyproject.toml‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
[build-system]
2+
requires = ["hatchling"]
3+
build-backend = "hatchling.build"
4+
5+
[project]
6+
name = "imagegen_game_sdk"
7+
version = "0.0.1"
8+
authors = [
9+
{ name = "Chinese Powered Labs"},
10+
]
11+
description = "Image generator for GAME by Virtuals"
12+
readme = "README.md"
13+
requires-python = ">=3.9"
14+
classifiers = [
15+
"Programming Language :: Python :: 3",
16+
"Programming Language :: Python :: 3.9",
17+
"Programming Language :: Python :: 3.10",
18+
"Programming Language :: Python :: 3.11",
19+
"License :: OSI Approved :: MIT License",
20+
"Operating System :: OS Independent",
21+
"Development Status :: 3 - Alpha",
22+
"Intended Audience :: Developers",
23+
"Topic :: Software Development :: Libraries :: Python Modules",
24+
]
25+
dependencies = [
26+
"game-sdk>=0.1.1",
27+
"requests (>=2.32.3,<3.0.0)",
28+
]
29+
30+
[tool.hatch.build.targets.wheel]
31+
packages = ["imagegen_game_sdk"]
32+
33+
[project.urls]
34+
"Homepage" = "https://github.com/game-by-virtuals/game-python/plugins/imagegen"
35+
"Bug Tracker" = "https://github.com/game-by-virtuals/game-python"

0 commit comments

Comments
 (0)