Spaces:
Sleeping
Sleeping
Update core/image_services.py
Browse files- core/image_services.py +90 -38
core/image_services.py
CHANGED
@@ -9,36 +9,50 @@ from openai import OpenAI # For DALL-E
|
|
9 |
|
10 |
# --- API Key Configuration ---
|
11 |
OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Primary for DALL-E
|
12 |
-
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback
|
13 |
|
14 |
OPENAI_DALLE_CONFIGURED = False
|
15 |
-
HF_IMAGE_API_CONFIGURED = False
|
16 |
hf_inference_image_client = None
|
17 |
openai_client = None
|
18 |
|
19 |
-
class ImageGenResponse:
|
20 |
def __init__(self, image: Image.Image = None, image_url: str = None,
|
21 |
error: str = None, success: bool = True,
|
22 |
provider: str = "Unknown Image Gen", model_id_used: str = None):
|
23 |
-
self.image
|
24 |
-
|
|
|
|
|
|
|
|
|
25 |
|
26 |
-
def
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
27 |
global OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client, openai_client, OPENAI_API_KEY, HF_TOKEN
|
28 |
|
29 |
-
|
|
|
30 |
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
|
31 |
|
32 |
-
print("INFO: image_services.py - Initializing Image Generation services (DALL-E primary)...")
|
33 |
|
34 |
# OpenAI DALL-E (Primary)
|
35 |
if OPENAI_API_KEY and OPENAI_API_KEY.strip():
|
36 |
print("INFO: image_services.py - STORYVERSE_OPENAI_API_KEY found.")
|
37 |
try:
|
38 |
openai_client = OpenAI(api_key=OPENAI_API_KEY)
|
39 |
-
#
|
40 |
-
#
|
41 |
-
#
|
|
|
|
|
42 |
OPENAI_DALLE_CONFIGURED = True
|
43 |
print("SUCCESS: image_services.py - OpenAI DALL-E client configured.")
|
44 |
except Exception as e:
|
@@ -58,7 +72,7 @@ def initialize_image_llms(): # Renamed to reflect image services
|
|
58 |
print("SUCCESS: image_services.py - Hugging Face InferenceClient (for fallback images) ready.")
|
59 |
except Exception as e:
|
60 |
HF_IMAGE_API_CONFIGURED = False
|
61 |
-
print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for fallback images: {e}")
|
62 |
hf_inference_image_client = None
|
63 |
else:
|
64 |
HF_IMAGE_API_CONFIGURED = False
|
@@ -66,8 +80,13 @@ def initialize_image_llms(): # Renamed to reflect image services
|
|
66 |
|
67 |
print(f"INFO: image_services.py - Image Service Init complete. DALL-E Ready: {OPENAI_DALLE_CONFIGURED}, HF Image (Fallback) Ready: {HF_IMAGE_API_CONFIGURED}")
|
68 |
|
69 |
-
def is_dalle_ready():
|
70 |
-
|
|
|
|
|
|
|
|
|
|
|
71 |
|
72 |
# --- OpenAI DALL-E ---
|
73 |
def generate_image_dalle(prompt: str,
|
@@ -77,7 +96,7 @@ def generate_image_dalle(prompt: str,
|
|
77 |
n: int = 1,
|
78 |
response_format: str = "b64_json" # Get image data directly
|
79 |
) -> ImageGenResponse:
|
80 |
-
global openai_client
|
81 |
if not is_dalle_ready() or not openai_client:
|
82 |
return ImageGenResponse(error="OpenAI DALL-E API not configured.", success=False, provider="DALL-E", model_id_used=model)
|
83 |
|
@@ -89,53 +108,86 @@ def generate_image_dalle(prompt: str,
|
|
89 |
size=size,
|
90 |
quality=quality,
|
91 |
n=n,
|
92 |
-
response_format=response_format
|
93 |
)
|
94 |
|
95 |
if response_format == "b64_json":
|
96 |
if not response.data or not response.data[0].b64_json:
|
97 |
-
return ImageGenResponse(error="No image data in DALL-E b64_json response.", success=False, provider="DALL-E", model_id_used=model)
|
98 |
image_data = base64.b64decode(response.data[0].b64_json)
|
99 |
image = Image.open(BytesIO(image_data))
|
100 |
print(f"DEBUG: image_services.py - DALL-E image generated successfully ({model}).")
|
101 |
return ImageGenResponse(image=image, provider="DALL-E", model_id_used=model)
|
102 |
-
elif response_format == "url":
|
103 |
if not response.data or not response.data[0].url:
|
104 |
-
return ImageGenResponse(error="No image URL in DALL-E response.", success=False, provider="DALL-E", model_id_used=model)
|
105 |
image_url = response.data[0].url
|
106 |
-
|
107 |
img_content_response = requests.get(image_url, timeout=30)
|
108 |
img_content_response.raise_for_status()
|
109 |
image = Image.open(BytesIO(img_content_response.content))
|
110 |
print(f"DEBUG: image_services.py - DALL-E image downloaded successfully ({model}).")
|
111 |
return ImageGenResponse(image=image, image_url=image_url, provider="DALL-E", model_id_used=model)
|
|
|
|
|
112 |
|
113 |
except Exception as e:
|
114 |
error_msg = f"DALL-E API Error ({model}): {type(e).__name__} - {str(e)}"
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
119 |
|
120 |
print(f"ERROR: image_services.py - {error_msg}")
|
121 |
return ImageGenResponse(error=error_msg, success=False, provider="DALL-E", model_id_used=model, raw_response=e)
|
122 |
|
123 |
# --- Hugging Face Image Model (Fallback) ---
|
124 |
-
def generate_image_hf_model(prompt: str,
|
125 |
-
|
126 |
-
|
|
|
|
|
|
|
|
|
|
|
127 |
global hf_inference_image_client
|
128 |
-
if not is_hf_image_api_ready() or not hf_inference_image_client:
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
"
|
133 |
-
"
|
134 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
135 |
try:
|
136 |
-
image_result: Image.Image = hf_inference_image_client.text_to_image(
|
|
|
|
|
|
|
|
|
|
|
137 |
return ImageGenResponse(image=image_result, provider="HF Image API", model_id_used=model_id)
|
138 |
-
except Exception as e:
|
139 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
140 |
|
141 |
-
print("DEBUG: core.image_services (DALL-E Primary for StoryVerseWeaver) - Module defined.")
|
|
|
9 |
|
10 |
# --- API Key Configuration ---
|
11 |
OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Primary for DALL-E
|
12 |
+
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback & text
|
13 |
|
14 |
OPENAI_DALLE_CONFIGURED = False
|
15 |
+
HF_IMAGE_API_CONFIGURED = False
|
16 |
hf_inference_image_client = None
|
17 |
openai_client = None
|
18 |
|
19 |
+
class ImageGenResponse:
|
20 |
def __init__(self, image: Image.Image = None, image_url: str = None,
|
21 |
error: str = None, success: bool = True,
|
22 |
provider: str = "Unknown Image Gen", model_id_used: str = None):
|
23 |
+
self.image = image
|
24 |
+
self.image_url = image_url
|
25 |
+
self.error = error
|
26 |
+
self.success = success
|
27 |
+
self.provider = provider
|
28 |
+
self.model_id_used = model_id_used
|
29 |
|
30 |
+
def __str__(self):
|
31 |
+
status = "Success" if self.success else "Failed"
|
32 |
+
details = f"Image URL: {self.image_url}" if self.image_url else ("Image data present" if self.image else "No image data")
|
33 |
+
if self.error:
|
34 |
+
details = f"Error: {self.error}"
|
35 |
+
return f"ImageGenResponse(Provider: {self.provider}, Model: {self.model_id_used or 'N/A'}, Status: {status}, Details: {details})"
|
36 |
+
|
37 |
+
def initialize_image_llms(): # "LLMs" here is a bit of a misnomer for image services, but kept for consistency
|
38 |
global OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client, openai_client, OPENAI_API_KEY, HF_TOKEN
|
39 |
|
40 |
+
# Ensure keys are fetched within this function's scope if not already module-level and populated
|
41 |
+
OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY")
|
42 |
HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
|
43 |
|
44 |
+
print("INFO: image_services.py - Initializing Image Generation services (DALL-E primary, HF fallback)...")
|
45 |
|
46 |
# OpenAI DALL-E (Primary)
|
47 |
if OPENAI_API_KEY and OPENAI_API_KEY.strip():
|
48 |
print("INFO: image_services.py - STORYVERSE_OPENAI_API_KEY found.")
|
49 |
try:
|
50 |
openai_client = OpenAI(api_key=OPENAI_API_KEY)
|
51 |
+
# A lightweight way to test if the client is configured and key is somewhat valid:
|
52 |
+
# try:
|
53 |
+
# openai_client.models.list() # This makes a quick API call
|
54 |
+
# except Exception as test_e:
|
55 |
+
# raise Exception(f"OpenAI client initialized but test call failed: {test_e}") from test_e
|
56 |
OPENAI_DALLE_CONFIGURED = True
|
57 |
print("SUCCESS: image_services.py - OpenAI DALL-E client configured.")
|
58 |
except Exception as e:
|
|
|
72 |
print("SUCCESS: image_services.py - Hugging Face InferenceClient (for fallback images) ready.")
|
73 |
except Exception as e:
|
74 |
HF_IMAGE_API_CONFIGURED = False
|
75 |
+
print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for fallback images: {type(e).__name__} - {e}")
|
76 |
hf_inference_image_client = None
|
77 |
else:
|
78 |
HF_IMAGE_API_CONFIGURED = False
|
|
|
80 |
|
81 |
print(f"INFO: image_services.py - Image Service Init complete. DALL-E Ready: {OPENAI_DALLE_CONFIGURED}, HF Image (Fallback) Ready: {HF_IMAGE_API_CONFIGURED}")
|
82 |
|
83 |
+
def is_dalle_ready():
|
84 |
+
global OPENAI_DALLE_CONFIGURED
|
85 |
+
return OPENAI_DALLE_CONFIGURED
|
86 |
+
|
87 |
+
def is_hf_image_api_ready():
|
88 |
+
global HF_IMAGE_API_CONFIGURED
|
89 |
+
return HF_IMAGE_API_CONFIGURED
|
90 |
|
91 |
# --- OpenAI DALL-E ---
|
92 |
def generate_image_dalle(prompt: str,
|
|
|
96 |
n: int = 1,
|
97 |
response_format: str = "b64_json" # Get image data directly
|
98 |
) -> ImageGenResponse:
|
99 |
+
global openai_client # Use the initialized client
|
100 |
if not is_dalle_ready() or not openai_client:
|
101 |
return ImageGenResponse(error="OpenAI DALL-E API not configured.", success=False, provider="DALL-E", model_id_used=model)
|
102 |
|
|
|
108 |
size=size,
|
109 |
quality=quality,
|
110 |
n=n,
|
111 |
+
response_format=response_format
|
112 |
)
|
113 |
|
114 |
if response_format == "b64_json":
|
115 |
if not response.data or not response.data[0].b64_json:
|
116 |
+
return ImageGenResponse(error="No image data in DALL-E b64_json response.", success=False, provider="DALL-E", model_id_used=model, raw_response=response)
|
117 |
image_data = base64.b64decode(response.data[0].b64_json)
|
118 |
image = Image.open(BytesIO(image_data))
|
119 |
print(f"DEBUG: image_services.py - DALL-E image generated successfully ({model}).")
|
120 |
return ImageGenResponse(image=image, provider="DALL-E", model_id_used=model)
|
121 |
+
elif response_format == "url": # If you choose to get URL
|
122 |
if not response.data or not response.data[0].url:
|
123 |
+
return ImageGenResponse(error="No image URL in DALL-E response.", success=False, provider="DALL-E", model_id_used=model, raw_response=response)
|
124 |
image_url = response.data[0].url
|
125 |
+
print(f"DEBUG: image_services.py - DALL-E image URL received ({model}): {image_url}. Attempting download...")
|
126 |
img_content_response = requests.get(image_url, timeout=30)
|
127 |
img_content_response.raise_for_status()
|
128 |
image = Image.open(BytesIO(img_content_response.content))
|
129 |
print(f"DEBUG: image_services.py - DALL-E image downloaded successfully ({model}).")
|
130 |
return ImageGenResponse(image=image, image_url=image_url, provider="DALL-E", model_id_used=model)
|
131 |
+
else:
|
132 |
+
return ImageGenResponse(error=f"Unsupported DALL-E response_format: {response_format}", success=False, provider="DALL-E", model_id_used=model)
|
133 |
|
134 |
except Exception as e:
|
135 |
error_msg = f"DALL-E API Error ({model}): {type(e).__name__} - {str(e)}"
|
136 |
+
# Attempt to get more details from OpenAI error structure
|
137 |
+
if hasattr(e, 'response') and e.response is not None:
|
138 |
+
try:
|
139 |
+
err_data = e.response.json()
|
140 |
+
if 'error' in err_data and 'message' in err_data['error']:
|
141 |
+
error_msg += f" - OpenAI Message: {err_data['error']['message']}"
|
142 |
+
elif hasattr(e.response, 'text'):
|
143 |
+
error_msg += f" - API Response: {e.response.text[:200]}"
|
144 |
+
except: # Fallback if parsing response fails
|
145 |
+
if hasattr(e.response, 'text'): error_msg += f" - API Response: {e.response.text[:200]}"
|
146 |
+
elif hasattr(e, 'message'):
|
147 |
+
error_msg += f" - Detail: {e.message}"
|
148 |
|
149 |
print(f"ERROR: image_services.py - {error_msg}")
|
150 |
return ImageGenResponse(error=error_msg, success=False, provider="DALL-E", model_id_used=model, raw_response=e)
|
151 |
|
152 |
# --- Hugging Face Image Model (Fallback) ---
|
153 |
+
def generate_image_hf_model(prompt: str,
|
154 |
+
model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", # Default HF model
|
155 |
+
negative_prompt: str = None,
|
156 |
+
height: int = 768,
|
157 |
+
width: int = 768,
|
158 |
+
num_inference_steps: int = 25,
|
159 |
+
guidance_scale: float = 7.0
|
160 |
+
) -> ImageGenResponse:
|
161 |
global hf_inference_image_client
|
162 |
+
if not is_hf_image_api_ready() or not hf_inference_image_client:
|
163 |
+
return ImageGenResponse(error="Hugging Face API (for images) not configured.", success=False, provider="HF Image API", model_id_used=model_id)
|
164 |
+
|
165 |
+
params = {
|
166 |
+
"negative_prompt": negative_prompt,
|
167 |
+
"height": height,
|
168 |
+
"width": width,
|
169 |
+
"num_inference_steps": num_inference_steps,
|
170 |
+
"guidance_scale": guidance_scale
|
171 |
+
}
|
172 |
+
params = {k: v for k, v in params.items() if v is not None}
|
173 |
+
|
174 |
+
print(f"DEBUG: image_services.py - Calling HF Image API ({model_id}) with prompt: {prompt[:70]}...")
|
175 |
try:
|
176 |
+
image_result: Image.Image = hf_inference_image_client.text_to_image(
|
177 |
+
prompt,
|
178 |
+
model=model_id,
|
179 |
+
**params
|
180 |
+
)
|
181 |
+
print(f"DEBUG: image_services.py - HF Image API ({model_id}) image generated successfully.")
|
182 |
return ImageGenResponse(image=image_result, provider="HF Image API", model_id_used=model_id)
|
183 |
+
except Exception as e:
|
184 |
+
error_msg = f"HF Image API Error ({model_id}): {type(e).__name__} - {str(e)}"
|
185 |
+
if "Rate limit reached" in str(e): error_msg += " You may have hit free tier limits for HF Inference API."
|
186 |
+
elif "Model is currently loading" in str(e) or "estimated_time" in str(e).lower(): error_msg += " The HF model may be loading, please try again in a moment."
|
187 |
+
elif "Authorization" in str(e) or "401" in str(e): error_msg += " Authentication issue with your STORYVERSE_HF_TOKEN."
|
188 |
+
elif "does not seem to support task text-to-image" in str(e): error_msg = f"Model {model_id} on HF may not support text-to-image or is misconfigured for Inference API."
|
189 |
+
|
190 |
+
print(f"ERROR: image_services.py - {error_msg}")
|
191 |
+
return ImageGenResponse(error=error_msg, success=False, provider="HF Image API", model_id_used=model_id, raw_response=e)
|
192 |
|
193 |
+
print("DEBUG: core.image_services (DALL-E Primary, HF Fallback for StoryVerseWeaver) - Module defined.")
|