cedpsam's picture
Upload tool
c2fbae4 verified
from smolagents import Tool
from typing import Any, Optional
class SimpleTool(Tool):
name = "pollinations_generate"
description = "Generate an image from a prompt"
inputs = {'prompt': {'type': 'string', 'description': 'The prompt to generate the image from'}, 'width': {'type': 'integer', 'nullable': True, 'description': 'image width'}, 'height': {'type': 'integer', 'nullable': True, 'description': 'image height'}, 'seed': {'type': 'integer', 'nullable': True, 'description': 'A decimal integer Each seed generates a new image variation Returns generated image as a PILL image'}}
output_type = "object"
def forward(self, prompt: str, width: int = 1024, height: int = 1024, seed: int = 0) -> 'PIL.Image.Image':
"""
Generate an image from a prompt
Args:
prompt : The prompt to generate the image from
width : image width
height : image height
seed : A decimal integer Each seed generates a new image variation
Returns
generated image as a PILL image
"""
import PIL
from PIL import Image
import requests
import random
from io import BytesIO
if seed ==0:
seed=random.randint(1,8*1024**2)
else:
random.seed(seed)
url = f"https://pollinations.ai/p/{prompt}"
params={"width":width,"height":height,"seed":seed}
response = requests.get(url)
return Image.open(requests.get(url, stream=True,params=params).raw)