ImageGenerator / model.py
Jaspertw177's picture
Initial commit
ed69556
Raw
History Blame Contribute Delete
1.96 kB
from openai import AzureOpenAI
import json
from io import BytesIO
import requests
import re
import streamlit as st
def get_client():
client = AzureOpenAI(
api_version="2024-05-01-preview",
azure_endpoint=st.secrets['endpoint'],
api_key=st.secrets['key'],
)
return client
def generate_image(prompt):
client = get_client()
result = client.images.generate(
model="Dalle3",
prompt=prompt,
n=1
)
image_url = json.loads(result.model_dump_json())['data'][0]['url']
result = requests.get(image_url)
return BytesIO(result.content)
def generate_image_prompt(prompt):
payload = {
"messages": [
{
"role": "system",
"content": [
{
"type": "text",
"text": "You give a few examples of english prompts that help generate image base on user's input. Return prompts in bullet point"
}
]
},
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt
}
]
}
],
"temperature": 0.9,
"top_p": 0.95,
"max_tokens": 800
}
response = requests.post(st.secrets['completionendpoint'], headers={"Content-Type": "application/json", "api-key": st.secrets['key']}, json=payload)
response.raise_for_status() # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
return response.json()['choices'][0]['message']['content']
def process_image_prompt(response):
response = response.split('\n')
response = [re.sub(r"(?<!\\)['\"](.*?)(?<!\\)['\"]", r"\1", response[i]) for i in range(len(response))]
return response
if __name__ == "__main__":
response = generate_image_prompt('halong bay, vietnam')
response = process_image_prompt(response)