mgbam commited on
Commit
e6e5425
·
verified ·
1 Parent(s): 886d0b0

Update core/image_services.py

Browse files
Files changed (1) hide show
  1. 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 # For fallback image model
16
  hf_inference_image_client = None
17
  openai_client = None
18
 
19
- class ImageGenResponse: # Keep this class
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, self.image_url, self.error, self.success, self.provider, self.model_id_used = \
24
- image, image_url, error, success, provider, model_id_used
 
 
 
 
25
 
26
- def initialize_image_llms(): # Renamed to reflect image services
 
 
 
 
 
 
 
27
  global OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client, openai_client, OPENAI_API_KEY, HF_TOKEN
28
 
29
- OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Ensure it's loaded here
 
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
- # Simple test: list models (lightweight call, though DALL-E models aren't listed this way usually)
40
- # A better test would be a very cheap image call if possible, or assume ready if client inits.
41
- # For now, client initialization is the main check.
 
 
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(): return OPENAI_DALLE_CONFIGURED
70
- def is_hf_image_api_ready(): return HF_IMAGE_API_CONFIGURED # Still useful for fallback
 
 
 
 
 
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 # Get base64 encoded image
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
- # Download the image from URL
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
- if hasattr(e, 'response') and e.response is not None and hasattr(e.response, 'text'):
116
- error_msg += f" - API Response: {e.response.text[:200]}"
117
- elif hasattr(e, 'message'): # OpenAI specific error structure
118
- error_msg += f" - OpenAI Message: {e.message}"
 
 
 
 
 
 
 
 
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, model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", ...) -> ImageGenResponse:
125
- # ... (This function remains the same as the one from "I don't have Stability api what do I do?" response)
126
- # ... (It uses hf_inference_image_client)
 
 
 
 
 
127
  global hf_inference_image_client
128
- if not is_hf_image_api_ready() or not hf_inference_image_client: return ImageGenResponse(error="HF Image API not configured.", success=False, provider="HF Image API", model_id_used=model_id)
129
- params = { # Default params, ensure they are passed from app.py or orchestrator
130
- "negative_prompt": negative_prompt if 'negative_prompt' in locals() else None,
131
- "height": height if 'height' in locals() else 768, "width": width if 'width' in locals() else 768,
132
- "num_inference_steps": num_inference_steps if 'num_inference_steps' in locals() else 25,
133
- "guidance_scale": guidance_scale if 'guidance_scale' in locals() else 7.0
134
- }; params = {k: v for k,v in params.items() if v is not None}
 
 
 
 
 
 
135
  try:
136
- image_result: Image.Image = hf_inference_image_client.text_to_image(prompt, model=model_id, **params)
 
 
 
 
 
137
  return ImageGenResponse(image=image_result, provider="HF Image API", model_id_used=model_id)
138
- except Exception as e: return ImageGenResponse(error=f"HF Image API Error ({model_id}): {e}", success=False, provider="HF Image API", model_id_used=model_id, raw_response=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.")