Update core/image_services.py

core/image_services.py

@@ -1,93 +1,141 @@
 # storyverse_weaver/core/image_services.py
 import os
-import base64 # Still useful if HF API ever returns b64
+import requests 
+import base64
 from io import BytesIO
 from PIL import Image
-from huggingface_hub import InferenceClient # Main client for HF models
+from huggingface_hub import InferenceClient # For HF fallback
+from openai import OpenAI # For DALL-E
 
 # --- API Key Configuration ---
-HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # This is the key we'll use
+OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Primary for DALL-E
+HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # For fallback
 
-HF_IMAGE_API_CONFIGURED = False
+OPENAI_DALLE_CONFIGURED = False
+HF_IMAGE_API_CONFIGURED = False # For fallback image model
 hf_inference_image_client = None 
+openai_client = None
 
-class ImageGenResponse:
+class ImageGenResponse: # Keep this class
     def __init__(self, image: Image.Image = None, image_url: str = None, 
                  error: str = None, success: bool = True, 
-                 provider: str = "HF Image API", model_id_used: str = None):
+                 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
 
-def initialize_image_llms(): # Renamed for consistency, though it's not an "LLM" for images
-    global HF_IMAGE_API_CONFIGURED, hf_inference_image_client, HF_TOKEN
+def initialize_image_llms(): # Renamed to reflect image services
+    global OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client, openai_client, OPENAI_API_KEY, HF_TOKEN
     
-    HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # Ensure it's loaded here too
-    print("INFO: image_services.py - Initializing Image Generation services (HF Focus)...")
+    OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # Ensure it's loaded here
+    HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN")
+
+    print("INFO: image_services.py - Initializing Image Generation services (DALL-E primary)...")
     
+    # 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.
+            OPENAI_DALLE_CONFIGURED = True
+            print("SUCCESS: image_services.py - OpenAI DALL-E client configured.")
+        except Exception as e:
+            OPENAI_DALLE_CONFIGURED = False
+            print(f"ERROR: image_services.py - Failed to configure OpenAI DALL-E client: {type(e).__name__} - {e}")
+            openai_client = None
+    else:
+        OPENAI_DALLE_CONFIGURED = False
+        print("WARNING: image_services.py - STORYVERSE_OPENAI_API_KEY not found or empty. DALL-E disabled.")
+
+    # Hugging Face Image Models (Fallback)
     if HF_TOKEN and HF_TOKEN.strip():
+        print("INFO: image_services.py - STORYVERSE_HF_TOKEN found (for fallback image model).")
         try:
             hf_inference_image_client = InferenceClient(token=HF_TOKEN)
-            # Optional: Test with a quick model ping if desired, but client init is usually enough
-            # For instance, try to get model info for a known image model if API allows
-            # Or assume it's ready if client initializes without error.
             HF_IMAGE_API_CONFIGURED = True
-            print("SUCCESS: image_services.py - Hugging Face InferenceClient (for images) ready.")
+            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 images: {type(e).__name__} - {e}")
+            print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for fallback images: {e}")
             hf_inference_image_client = None
     else:
         HF_IMAGE_API_CONFIGURED = False
-        print("WARNING: image_services.py - STORYVERSE_HF_TOKEN not found or empty. HF Image models disabled.")
+        print("WARNING: image_services.py - STORYVERSE_HF_TOKEN not found or empty (for fallback image model).")
     
-    print(f"INFO: image_services.py - Image Service Init complete. HF Image API Configured: {HF_IMAGE_API_CONFIGURED}")
+    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_hf_image_api_ready(): # Getter function for app.py
-    global HF_IMAGE_API_CONFIGURED
-    return 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
 
-# --- Hugging Face Image Model via Inference API ---
-def generate_image_hf_model(prompt: str, 
-                            model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", # Default popular model
-                            negative_prompt: str = None,
-                            height: int = 768, # Common for SDXL
-                            width: int = 768,  # Common for SDXL
-                            num_inference_steps: int = 25,
-                            guidance_scale: float = 7.0 # Lower can be more creative, higher more prompt-adherent
-                           ) -> ImageGenResponse:
-    global hf_inference_image_client 
-    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, 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} # Clean out None params
-
-    print(f"DEBUG: image_services.py - Calling HF Image API ({model_id}) with prompt: {prompt[:70]}...")
+# --- OpenAI DALL-E ---
+def generate_image_dalle(prompt: str, 
+                         model: str = "dall-e-3", # or "dall-e-2"
+                         size: str = "1024x1024", 
+                         quality: str = "standard", # "standard" or "hd" for dall-e-3
+                         n: int = 1,
+                         response_format: str = "b64_json" # Get image data directly
+                        ) -> ImageGenResponse:
+    global openai_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)
+    
+    print(f"DEBUG: image_services.py - Calling DALL-E ({model}) with prompt: {prompt[:70]}...")
     try:
-        # InferenceClient's text_to_image method returns a PIL Image directly
-        image_result: Image.Image = hf_inference_image_client.text_to_image(
-            prompt,
-            model=model_id,
-            **params 
+        response = openai_client.images.generate(
+            model=model,
+            prompt=prompt,
+            size=size,
+            quality=quality,
+            n=n,
+            response_format=response_format # Get base64 encoded image
         )
-        # Some models might be on serverless inference endpoints that take longer
-        # The default timeout for InferenceClient is usually reasonable.
-        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:
-        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."
-        elif "Model is currently loading" in str(e) or "estimated_time" in str(e).lower(): error_msg += " Model may be loading, try again in a moment."
-        elif "Authorization" in str(e) or "401" in str(e): error_msg += " Authentication issue with your HF_TOKEN."
-        elif "does not seem to support task text-to-image" in str(e): error_msg = f"Model {model_id} may not support text-to-image or is misconfigured."
         
+        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)
+            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":
+            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)
+            image_url = response.data[0].url
+            # Download the image from URL
+            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)
+            
+    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}"
+
         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)
+        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)
+    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}
+    try:
+        image_result: Image.Image = hf_inference_image_client.text_to_image(prompt, model=model_id, **params)
+        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)
+
 
-print("DEBUG: core.image_services (HF Focus for StoryVerseWeaver) - Module defined.")
+print("DEBUG: core.image_services (DALL-E Primary for StoryVerseWeaver) - Module defined.")