File size: 1,405 Bytes
220a7fc |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import io
import os
from openai import OpenAI
from langchain.tools import StructuredTool, Tool
from io import BytesIO
import requests
import json
from io import BytesIO
import chainlit as cl
def get_image_name():
image_count = cl.user_session.get("image_count")
if image_count is None:
image_count = 0
else:
image_count += 1
cl.user_session.set("image_count", image_count)
return f"image-{image_count}"
def _generate_image(prompt: str):
client = OpenAI()
response = client.images.generate(
model="dall-e-3",
prompt=prompt,
size="1024x1024",
quality="standard",
n=1,
)
image_payload = requests.get(response.data[0].url, stream=True)
image_bytes = BytesIO(image_payload.content)
print(type(image_bytes))
name = get_image_name()
cl.user_session.set(name, image_bytes.getvalue())
cl.user_session.set("generated_image", name)
return name
def generate_image(prompt: str):
image_name = _generate_image(prompt)
return f"Here is {image_name}."
generate_image_format = '{{"prompt": "prompt"}}'
generate_image_tool = Tool.from_function(
func=generate_image,
name="GenerateImage",
description=f"Useful to create an image from a text prompt. Input should be a single string strictly in the following JSON format: {generate_image_format}",
return_direct=True,
)
|