File size: 1,462 Bytes
4fd500e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
from PIL import Image
import io
import os
import requests
from crewai_tools import tool
# API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
# hugginface_api_key = os.getenv("HUGGINGFACE_API_KEY")
# headers = {"Authorization": f"Bearer {hugginface_api_key}"}
# def query(payload):
# response = requests.post(API_URL, headers=headers, json=payload)
# return response.content
# def generate_image(input: str):
# """Create an image based on input"""
# image_bytes = query({
# "inputs": input,
# })
# # You can access the image with PIL.Image for example
# image = Image.open(io.BytesIO(image_bytes))
# name = input.split(" ")[0]
# image.save(f"images/generate_{name}.jpg")
# return image
@tool("Image Generate Tool")
def image_generate(context: str):
"""This is Image generate tool"""
API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
hugginface_api_key = os.getenv("HUGGINGFACE_API_KEY")
headers = {"Authorization": f"Bearer {hugginface_api_key}"}
def query(payload):
response = requests.post(API_URL, headers=headers, json=payload)
return response.content
image_bytes = query({
"inputs": context,
})
name = context.split(" ")[0]
image = Image.open(io.BytesIO(image_bytes)).resize((1024, 1024))
image.save(f"outputs/generate_{name}.jpg")
return image
|