1
1
import os
2
2
import base64
3
- import aiohttp
4
3
import imghdr
5
4
import requests
6
- import mimetypes
7
5
from dotenv import load_dotenv
8
6
import google .generativeai as genai
9
7
from io import BytesIO
10
8
from PIL import Image
11
9
from helpers .prompt import SYSTEM_PROMPT
10
+
12
11
load_dotenv ()
13
12
13
+ # API Keys
14
14
STABILITY_API_KEY = os .getenv ("STABILITY_API_KEY" )
15
15
GEMINI_API_KEY = os .getenv ("GEMINI_API_KEY" )
16
16
17
+ # API Host and Engine Configuration
17
18
API_HOST = "https://api.stability.ai"
18
- ENGINE_ID = "stable-diffusion-v1-6"
19
+ ENGINE_ID = "stable-diffusion-v1-6"
19
20
21
+ # Gemini AI Configuration
20
22
genai .configure (api_key = GEMINI_API_KEY )
21
23
generation_config = {
22
24
"temperature" : 1 ,
31
33
system_instruction = SYSTEM_PROMPT ,
32
34
)
33
35
36
+ # Get Text Response from Gemini Model
34
37
def get_response (conversation ):
35
38
try :
36
39
chat_session = model .start_chat (history = conversation )
@@ -40,30 +43,32 @@ def get_response(conversation):
40
43
print (f"Error in get_response: { e } " )
41
44
return "Uhg my brain hurts, can you say that again?"
42
45
46
+ # Recognize Image and Generate Content Based on Prompt
43
47
async def recognize_image (image_data , prompt ):
44
48
try :
45
49
image_format = imghdr .what (None , image_data )
46
- if image_format is None :
50
+ if not image_format :
47
51
return "Failed to determine image format"
48
52
49
53
image = Image .open (BytesIO (image_data ))
50
54
response = model .generate_content ([prompt , image ])
51
-
52
55
return response .text if response and hasattr (response , 'text' ) else "Failed to generate a response from the Gemini model"
53
56
except Exception as e :
54
57
print (f"Error in recognize_image: { e } " )
55
58
return None
56
59
60
+ # Generate Image from Text Prompt
57
61
def get_image (text ):
58
62
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
+
61
66
response = requests .post (
62
67
f"{ API_HOST } /v1/generation/{ ENGINE_ID } /text-to-image" ,
63
68
headers = {
64
69
"Content-Type" : "application/json" ,
65
70
"Accept" : "application/json" ,
66
- "Authorization" : f"Bearer { STABILITY_API_KEY } "
71
+ "Authorization" : f"Bearer { STABILITY_API_KEY } " ,
67
72
},
68
73
json = {
69
74
"text_prompts" : [{"text" : text }],
@@ -75,23 +80,23 @@ def get_image(text):
75
80
},
76
81
)
77
82
response .raise_for_status ()
78
- data = response .json ()
79
- image_data = data ["artifacts" ][0 ]["base64" ]
83
+ image_data = response .json ()["artifacts" ][0 ]["base64" ]
80
84
return base64 .b64decode (image_data )
81
85
except Exception as e :
82
86
print (f"Error in get_image: { e } " )
83
87
return None
84
-
88
+
89
+ # Edit Image Based on Prompt
85
90
def edit_image (image_bytes , prompt ):
86
91
try :
87
92
response = requests .post (
88
93
f"{ API_HOST } /v1/generation/{ ENGINE_ID } /image-to-image" ,
89
94
headers = {
90
95
"Accept" : "application/json" ,
91
- "Authorization" : f"Bearer { STABILITY_API_KEY } "
96
+ "Authorization" : f"Bearer { STABILITY_API_KEY } " ,
92
97
},
93
98
files = {
94
- "init_image" : ('init_image.png' , image_bytes , 'image/png' )
99
+ "init_image" : ('init_image.png' , image_bytes , 'image/png' ),
95
100
},
96
101
data = {
97
102
"image_strength" : 0.35 ,
@@ -100,12 +105,11 @@ def edit_image(image_bytes, prompt):
100
105
"cfg_scale" : 7 ,
101
106
"samples" : 1 ,
102
107
"steps" : 30 ,
103
- }
108
+ },
104
109
)
105
110
106
111
response .raise_for_status ()
107
- data = response .json ()
108
- image_data = data ["artifacts" ][0 ]["base64" ]
112
+ image_data = response .json ()["artifacts" ][0 ]["base64" ]
109
113
return base64 .b64decode (image_data )
110
114
except Exception as e :
111
115
print (f"Error in edit_image: { e } " )
0 commit comments