Skip to content

Commit 9979875

Browse files
authored
Update ai.py
1 parent 92eef44 commit 9979875

File tree

1 file changed

+20
-16
lines changed

1 file changed

+20
-16
lines changed

helpers/ai.py

+20-16
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,24 @@
11
import os
22
import base64
3-
import aiohttp
43
import imghdr
54
import requests
6-
import mimetypes
75
from dotenv import load_dotenv
86
import google.generativeai as genai
97
from io import BytesIO
108
from PIL import Image
119
from helpers.prompt import SYSTEM_PROMPT
10+
1211
load_dotenv()
1312

13+
# API Keys
1414
STABILITY_API_KEY = os.getenv("STABILITY_API_KEY")
1515
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
1616

17+
# API Host and Engine Configuration
1718
API_HOST = "https://api.stability.ai"
18-
ENGINE_ID= "stable-diffusion-v1-6"
19+
ENGINE_ID = "stable-diffusion-v1-6"
1920

21+
# Gemini AI Configuration
2022
genai.configure(api_key=GEMINI_API_KEY)
2123
generation_config = {
2224
"temperature": 1,
@@ -31,6 +33,7 @@
3133
system_instruction=SYSTEM_PROMPT,
3234
)
3335

36+
# Get Text Response from Gemini Model
3437
def get_response(conversation):
3538
try:
3639
chat_session = model.start_chat(history=conversation)
@@ -40,30 +43,32 @@ def get_response(conversation):
4043
print(f"Error in get_response: {e}")
4144
return "Uhg my brain hurts, can you say that again?"
4245

46+
# Recognize Image and Generate Content Based on Prompt
4347
async def recognize_image(image_data, prompt):
4448
try:
4549
image_format = imghdr.what(None, image_data)
46-
if image_format is None:
50+
if not image_format:
4751
return "Failed to determine image format"
4852

4953
image = Image.open(BytesIO(image_data))
5054
response = model.generate_content([prompt, image])
51-
5255
return response.text if response and hasattr(response, 'text') else "Failed to generate a response from the Gemini model"
5356
except Exception as e:
5457
print(f"Error in recognize_image: {e}")
5558
return None
5659

60+
# Generate Image from Text Prompt
5761
def get_image(text):
5862
try:
59-
if STABILITY_API_KEY is None:
60-
raise Exception("Missing Stability API key.")
63+
if not STABILITY_API_KEY:
64+
raise ValueError("Missing Stability API key.")
65+
6166
response = requests.post(
6267
f"{API_HOST}/v1/generation/{ENGINE_ID}/text-to-image",
6368
headers={
6469
"Content-Type": "application/json",
6570
"Accept": "application/json",
66-
"Authorization": f"Bearer {STABILITY_API_KEY}"
71+
"Authorization": f"Bearer {STABILITY_API_KEY}",
6772
},
6873
json={
6974
"text_prompts": [{"text": text}],
@@ -75,23 +80,23 @@ def get_image(text):
7580
},
7681
)
7782
response.raise_for_status()
78-
data = response.json()
79-
image_data = data["artifacts"][0]["base64"]
83+
image_data = response.json()["artifacts"][0]["base64"]
8084
return base64.b64decode(image_data)
8185
except Exception as e:
8286
print(f"Error in get_image: {e}")
8387
return None
84-
88+
89+
# Edit Image Based on Prompt
8590
def edit_image(image_bytes, prompt):
8691
try:
8792
response = requests.post(
8893
f"{API_HOST}/v1/generation/{ENGINE_ID}/image-to-image",
8994
headers={
9095
"Accept": "application/json",
91-
"Authorization": f"Bearer {STABILITY_API_KEY}"
96+
"Authorization": f"Bearer {STABILITY_API_KEY}",
9297
},
9398
files={
94-
"init_image": ('init_image.png', image_bytes, 'image/png')
99+
"init_image": ('init_image.png', image_bytes, 'image/png'),
95100
},
96101
data={
97102
"image_strength": 0.35,
@@ -100,12 +105,11 @@ def edit_image(image_bytes, prompt):
100105
"cfg_scale": 7,
101106
"samples": 1,
102107
"steps": 30,
103-
}
108+
},
104109
)
105110

106111
response.raise_for_status()
107-
data = response.json()
108-
image_data = data["artifacts"][0]["base64"]
112+
image_data = response.json()["artifacts"][0]["base64"]
109113
return base64.b64decode(image_data)
110114
except Exception as e:
111115
print(f"Error in edit_image: {e}")

0 commit comments

Comments
 (0)