Update core/image_services.py

core/image_services.py

@@ -9,36 +9,50 @@ from openai import OpenAI # For DALL-E
 
 # --- API Key Configuration ---
 OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Primary for DALL-E
-HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback
+HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback & text
 
 OPENAI_DALLE_CONFIGURED = False
-HF_IMAGE_API_CONFIGURED = False # For fallback image model
+HF_IMAGE_API_CONFIGURED = False 
 hf_inference_image_client = None 
 openai_client = None
 
-class ImageGenResponse: # Keep this class
+class ImageGenResponse:
     def __init__(self, image: Image.Image = None, image_url: str = None, 
                  error: str = None, success: bool = True, 
                  provider: str = "Unknown Image Gen", model_id_used: str = None):
-        self.image, self.image_url, self.error, self.success, self.provider, self.model_id_used = \
-            image, image_url, error, success, provider, model_id_used
+        self.image = image
+        self.image_url = image_url
+        self.error = error
+        self.success = success
+        self.provider = provider
+        self.model_id_used = model_id_used
 
-def initialize_image_llms(): # Renamed to reflect image services
+    def __str__(self):
+        status = "Success" if self.success else "Failed"
+        details = f"Image URL: {self.image_url}" if self.image_url else ("Image data present" if self.image else "No image data")
+        if self.error:
+            details = f"Error: {self.error}"
+        return f"ImageGenResponse(Provider: {self.provider}, Model: {self.model_id_used or 'N/A'}, Status: {status}, Details: {details})"
+
+def initialize_image_llms(): # "LLMs" here is a bit of a misnomer for image services, but kept for consistency
     global OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client, openai_client, OPENAI_API_KEY, HF_TOKEN
     
-    OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Ensure it's loaded here
+    # Ensure keys are fetched within this function's scope if not already module-level and populated
+    OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY")
     HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
 
-    print("INFO: image_services.py - Initializing Image Generation services (DALL-E primary)...")
+    print("INFO: image_services.py - Initializing Image Generation services (DALL-E primary, HF fallback)...")
     
     # OpenAI DALL-E (Primary)
     if OPENAI_API_KEY and OPENAI_API_KEY.strip():
         print("INFO: image_services.py - STORYVERSE_OPENAI_API_KEY found.")
         try:
             openai_client = OpenAI(api_key=OPENAI_API_KEY)
-            # Simple test: list models (lightweight call, though DALL-E models aren't listed this way usually)
-            # A better test would be a very cheap image call if possible, or assume ready if client inits.
-            # For now, client initialization is the main check.
+            # A lightweight way to test if the client is configured and key is somewhat valid:
+            # try:
+            #     openai_client.models.list() # This makes a quick API call
+            # except Exception as test_e:
+            #     raise Exception(f"OpenAI client initialized but test call failed: {test_e}") from test_e
             OPENAI_DALLE_CONFIGURED = True
             print("SUCCESS: image_services.py - OpenAI DALL-E client configured.")
         except Exception as e:
@@ -58,7 +72,7 @@ def initialize_image_llms(): # Renamed to reflect image services
             print("SUCCESS: image_services.py - Hugging Face InferenceClient (for fallback images) ready.")
         except Exception as e:
             HF_IMAGE_API_CONFIGURED = False
-            print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for fallback images: {e}")
+            print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for fallback images: {type(e).__name__} - {e}")
             hf_inference_image_client = None
     else:
         HF_IMAGE_API_CONFIGURED = False
@@ -66,8 +80,13 @@ def initialize_image_llms(): # Renamed to reflect image services
     
     print(f"INFO: image_services.py - Image Service Init complete. DALL-E Ready: {OPENAI_DALLE_CONFIGURED}, HF Image (Fallback) Ready: {HF_IMAGE_API_CONFIGURED}")
 
-def is_dalle_ready(): return OPENAI_DALLE_CONFIGURED
-def is_hf_image_api_ready(): return HF_IMAGE_API_CONFIGURED # Still useful for fallback
+def is_dalle_ready(): 
+    global OPENAI_DALLE_CONFIGURED
+    return OPENAI_DALLE_CONFIGURED
+
+def is_hf_image_api_ready(): 
+    global HF_IMAGE_API_CONFIGURED
+    return HF_IMAGE_API_CONFIGURED
 
 # --- OpenAI DALL-E ---
 def generate_image_dalle(prompt: str, 
@@ -77,7 +96,7 @@ def generate_image_dalle(prompt: str,
                          n: int = 1,
                          response_format: str = "b64_json" # Get image data directly
                         ) -> ImageGenResponse:
-    global openai_client
+    global openai_client # Use the initialized client
     if not is_dalle_ready() or not openai_client:
         return ImageGenResponse(error="OpenAI DALL-E API not configured.", success=False, provider="DALL-E", model_id_used=model)
     
@@ -89,53 +108,86 @@ def generate_image_dalle(prompt: str,
             size=size,
             quality=quality,
             n=n,
-            response_format=response_format # Get base64 encoded image
+            response_format=response_format 
         )
         
         if response_format == "b64_json":
             if not response.data or not response.data[0].b64_json:
-                return ImageGenResponse(error="No image data in DALL-E b64_json response.", success=False, provider="DALL-E", model_id_used=model)
+                return ImageGenResponse(error="No image data in DALL-E b64_json response.", success=False, provider="DALL-E", model_id_used=model, raw_response=response)
             image_data = base64.b64decode(response.data[0].b64_json)
             image = Image.open(BytesIO(image_data))
             print(f"DEBUG: image_services.py - DALL-E image generated successfully ({model}).")
             return ImageGenResponse(image=image, provider="DALL-E", model_id_used=model)
-        elif response_format == "url":
+        elif response_format == "url": # If you choose to get URL
             if not response.data or not response.data[0].url:
-                 return ImageGenResponse(error="No image URL in DALL-E response.", success=False, provider="DALL-E", model_id_used=model)
+                 return ImageGenResponse(error="No image URL in DALL-E response.", success=False, provider="DALL-E", model_id_used=model, raw_response=response)
             image_url = response.data[0].url
-            # Download the image from URL
+            print(f"DEBUG: image_services.py - DALL-E image URL received ({model}): {image_url}. Attempting download...")
             img_content_response = requests.get(image_url, timeout=30)
             img_content_response.raise_for_status()
             image = Image.open(BytesIO(img_content_response.content))
             print(f"DEBUG: image_services.py - DALL-E image downloaded successfully ({model}).")
             return ImageGenResponse(image=image, image_url=image_url, provider="DALL-E", model_id_used=model)
+        else:
+            return ImageGenResponse(error=f"Unsupported DALL-E response_format: {response_format}", success=False, provider="DALL-E", model_id_used=model)
             
     except Exception as e:
         error_msg = f"DALL-E API Error ({model}): {type(e).__name__} - {str(e)}"
-        if hasattr(e, 'response') and e.response is not None and hasattr(e.response, 'text'):
-             error_msg += f" - API Response: {e.response.text[:200]}"
-        elif hasattr(e, 'message'): # OpenAI specific error structure
-             error_msg += f" - OpenAI Message: {e.message}"
+        # Attempt to get more details from OpenAI error structure
+        if hasattr(e, 'response') and e.response is not None:
+            try:
+                err_data = e.response.json()
+                if 'error' in err_data and 'message' in err_data['error']:
+                    error_msg += f" - OpenAI Message: {err_data['error']['message']}"
+                elif hasattr(e.response, 'text'):
+                     error_msg += f" - API Response: {e.response.text[:200]}"
+            except: # Fallback if parsing response fails
+                 if hasattr(e.response, 'text'): error_msg += f" - API Response: {e.response.text[:200]}"
+        elif hasattr(e, 'message'): 
+             error_msg += f" - Detail: {e.message}"
 
         print(f"ERROR: image_services.py - {error_msg}")
         return ImageGenResponse(error=error_msg, success=False, provider="DALL-E", model_id_used=model, raw_response=e)
 
 # --- Hugging Face Image Model (Fallback) ---
-def generate_image_hf_model(prompt: str, model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", ...) -> ImageGenResponse:
-    # ... (This function remains the same as the one from "I don't have Stability api what do I do?" response)
-    # ... (It uses hf_inference_image_client)
+def generate_image_hf_model(prompt: str, 
+                            model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", # Default HF model
+                            negative_prompt: str = None,
+                            height: int = 768, 
+                            width: int = 768,
+                            num_inference_steps: int = 25,
+                            guidance_scale: float = 7.0
+                           ) -> ImageGenResponse:
     global hf_inference_image_client 
-    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)
-    params = { # Default params, ensure they are passed from app.py or orchestrator
-        "negative_prompt": negative_prompt if 'negative_prompt' in locals() else None,
-        "height": height if 'height' in locals() else 768, "width": width if 'width' in locals() else 768,
-        "num_inference_steps": num_inference_steps if 'num_inference_steps' in locals() else 25,
-        "guidance_scale": guidance_scale if 'guidance_scale' in locals() else 7.0
-    }; params = {k: v for k,v in params.items() if v is not None}
+    if not is_hf_image_api_ready() or not hf_inference_image_client: 
+        return ImageGenResponse(error="Hugging Face API (for images) not configured.", success=False, provider="HF Image API", model_id_used=model_id)
+    
+    params = {
+        "negative_prompt": negative_prompt,
+        "height": height,
+        "width": width,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale
+    }
+    params = {k: v for k, v in params.items() if v is not None}
+
+    print(f"DEBUG: image_services.py - Calling HF Image API ({model_id}) with prompt: {prompt[:70]}...")
     try:
-        image_result: Image.Image = hf_inference_image_client.text_to_image(prompt, model=model_id, **params)
+        image_result: Image.Image = hf_inference_image_client.text_to_image(
+            prompt,
+            model=model_id,
+            **params 
+        )
+        print(f"DEBUG: image_services.py - HF Image API ({model_id}) image generated successfully.")
         return ImageGenResponse(image=image_result, provider="HF Image API", model_id_used=model_id)
-    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)
-
+    except Exception as e:
+        error_msg = f"HF Image API Error ({model_id}): {type(e).__name__} - {str(e)}"
+        if "Rate limit reached" in str(e): error_msg += " You may have hit free tier limits for HF Inference API."
+        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."
+        elif "Authorization" in str(e) or "401" in str(e): error_msg += " Authentication issue with your STORYVERSE_HF_TOKEN."
+        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."
+        
+        print(f"ERROR: image_services.py - {error_msg}")
+        return ImageGenResponse(error=error_msg, success=False, provider="HF Image API", model_id_used=model_id, raw_response=e)
 
-print("DEBUG: core.image_services (DALL-E Primary for StoryVerseWeaver) - Module defined.")
+print("DEBUG: core.image_services (DALL-E Primary, HF Fallback for StoryVerseWeaver) - Module defined.")