ChartPipeline / scripts /image_generation.py
Ray1ee01's picture
Upload folder using huggingface_hub
58e6885 verified
Raw
History Blame Contribute Delete
2.75 kB
import os
from openai import OpenAI
from PIL import Image
from io import BytesIO
import base64
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import api_key, base_url
# OpenAI API configuration
API_KEY = api_key
API_PROVIDER = base_url
client = OpenAI(
api_key=API_KEY,
base_url=API_PROVIDER,
)
'''
response = client.chat.completions.create(
model="V_2A",
messages=[
{
"role": "user",
"content": 'A cartt in a infographic chart for "Theresa May's Attributes Decline Over Time". No gradients, shadows, or complex details. Avoid photorealistic rendering.'
}
]
)
'''
# Using text-only input
response = client.chat.completions.create(
model="gemini-2.0-flash-exp",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": 'A pictogram for Italian Cuisine. No gradients, shadows, or complex details. Avoid photorealistic rendering. No text, no title.',
}
],
},
],
modalities=["text", "image"],
temperature=0.7,
)
try:
# Print basic response information
print(f"Creation time: {response.created}")
print(f"Token usage: {response.usage.total_tokens}")
# Check if multi_mod_content field exists
if (
hasattr(response.choices[0].message, "multi_mod_content")
and response.choices[0].message.multi_mod_content is not None
):
print("\nResponse content:")
for part in response.choices[0].message.multi_mod_content:
if "text" in part and part["text"] is not None:
print(part["text"])
# Process image content
elif "inline_data" in part and part["inline_data"] is not None:
print("\n🖼️ [Image content received]")
image_data = base64.b64decode(part["inline_data"]["data"])
mime_type = part["inline_data"].get("mime_type", "image/png")
print(f"Image type: {mime_type}")
image = Image.open(BytesIO(image_data))
image.show()
# Save image
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "generated_image.jpg")
image.save(output_path)
print(f"✅ Image saved to: {output_path}")
else:
print("No valid multimodal response received, check response structure")
except Exception as e:
print(f"Error processing response: {str(e)}")