Update core/image_services.py

core/image_services.py

@@ -1,132 +1,127 @@
 # storyverse_weaver/core/image_services.py
 import os
-import requests # For generic API calls
+import requests
 import base64
 from io import BytesIO
 from PIL import Image
-# from dotenv import load_dotenv
-# load_dotenv()
+from huggingface_hub import InferenceClient # Ensure this is imported
 
-# --- API Key Configuration (Use specific names for StoryVerse) ---
+# --- API Key Configuration ---
 STABILITY_API_KEY = os.getenv("STORYVERSE_STABILITY_API_KEY")
-OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY") # For DALL-E
-# HUGGINGFACE_TOKEN also used by llm_services, can be reused if using HF image models
+OPENAI_API_KEY = os.getenv("STORYVERSE_OPENAI_API_KEY")
+HF_TOKEN = os.getenv("STORYVERSE_HF_TOKEN") # Reuse from llm_services or get it here
 
 STABILITY_API_CONFIGURED = bool(STABILITY_API_KEY and STABILITY_API_KEY.strip())
 OPENAI_DALLE_CONFIGURED = bool(OPENAI_API_KEY and OPENAI_API_KEY.strip())
-# HF_IMAGE_CONFIGURED = bool(HF_TOKEN and HF_TOKEN.strip()) # Assuming HF_TOKEN is also for image models
+HF_IMAGE_API_CONFIGURED = bool(HF_TOKEN and HF_TOKEN.strip()) # New flag
 
-class ImageGenResponse:
-    def __init__(self, image: Image.Image = None, image_url: str = None, error: str = None, success: bool = True, provider: str = "unknown"):
-        self.image = image # PIL Image object
-        self.image_url = image_url # If API returns a URL
-        self.error = error
-        self.success = success
-        self.provider = provider
+hf_inference_image_client = None # Separate client instance for image tasks if needed, or reuse
 
-def initialize_image_llms(): # Simple function to print status
+class ImageGenResponse: # Keep this class
+    def __init__(self, image: Image.Image = None, image_url: str = None, error: str = None, success: bool = True, provider: str = "unknown", 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():
+    global STABILITY_API_CONFIGURED, OPENAI_DALLE_CONFIGURED, HF_IMAGE_API_CONFIGURED, hf_inference_image_client
+    
     print("INFO: image_services.py - Initializing Image Generation services...")
-    if STABILITY_API_CONFIGURED: print("SUCCESS: image_services.py - Stability AI API Key detected.")
-    else: print("WARNING: image_services.py - STORYVERSE_STABILITY_API_KEY not found. Stability AI disabled.")
-    if OPENAI_DALLE_CONFIGURED: print("SUCCESS: image_services.py - OpenAI API Key detected (for DALL-E).")
-    else: print("WARNING: image_services.py - STORYVERSE_OPENAI_API_KEY not found. DALL-E disabled.")
-    # if HF_IMAGE_CONFIGURED: print("INFO: image_services.py - Hugging Face Token detected (can be used for HF image models).")
-    print("INFO: image_services.py - Image LLM Init complete.")
-
-
-# --- Stability AI (Example) ---
-def generate_image_stabilityai(prompt: str, style_preset: str = None, negative_prompt: str = None, 
-                               engine_id: str = "stable-diffusion-xl-1024-v1-0", 
-                               steps: int = 30, cfg_scale: float = 7.0, 
-                               width: int = 1024, height: int = 1024) -> ImageGenResponse:
-    if not STABILITY_API_CONFIGURED:
-        return ImageGenResponse(error="Stability AI API key not configured.", success=False, provider="StabilityAI")
-
-    api_host = os.getenv('API_HOST', 'https://api.stability.ai')
-    request_url = f"{api_host}/v1/generation/{engine_id}/text-to-image"
-
-    payload = {
-        "text_prompts": [{"text": prompt}],
-        "cfg_scale": cfg_scale,
+    # Stability AI (as before)
+    if STABILITY_API_KEY and STABILITY_API_KEY.strip():
+        STABILITY_API_CONFIGURED = True
+        print("SUCCESS: image_services.py - Stability AI API Key detected.")
+    else:
+        STABILITY_API_CONFIGURED = False
+        print("WARNING: image_services.py - STORYVERSE_STABILITY_API_KEY not found. Stability AI disabled.")
+
+    # OpenAI DALL-E (as before)
+    if OPENAI_API_KEY and OPENAI_API_KEY.strip():
+        OPENAI_DALLE_CONFIGURED = True
+        print("SUCCESS: image_services.py - OpenAI API Key detected (for DALL-E).")
+    else:
+        OPENAI_DALLE_CONFIGURED = False
+        print("WARNING: image_services.py - STORYVERSE_OPENAI_API_KEY not found. DALL-E disabled.")
+
+    # Hugging Face Image Models
+    if HF_TOKEN and HF_TOKEN.strip():
+        try:
+            # You can use the same token for text and image clients if the permissions cover it
+            # Or, if you want a dedicated client for image tasks (maybe different default model types)
+            hf_inference_image_client = InferenceClient(token=HF_TOKEN) # Reusing the token
+            HF_IMAGE_API_CONFIGURED = True
+            print("SUCCESS: image_services.py - Hugging Face InferenceClient (for images) ready.")
+        except Exception as e:
+            HF_IMAGE_API_CONFIGURED = False
+            print(f"ERROR: image_services.py - Failed to initialize HF InferenceClient for images: {e}")
+    else:
+        HF_IMAGE_API_CONFIGURED = False
+        print("WARNING: image_services.py - STORYVERSE_HF_TOKEN not found. HF Image models disabled.")
+    
+    print("INFO: image_services.py - Image Service Init complete.")
+
+
+# --- Stability AI (Keep as is, it will just be disabled if no key) ---
+def generate_image_stabilityai(prompt: str, ...) -> ImageGenResponse:
+    # ... (your existing generate_image_stabilityai function)
+    if not STABILITY_API_CONFIGURED: return ImageGenResponse(error="Stability AI API key not configured.", success=False, provider="StabilityAI")
+    # ... rest of the function
+    api_host = os.getenv('API_HOST', 'https://api.stability.ai'); request_url = f"{api_host}/v1/generation/{engine_id if 'engine_id' in locals() else 'stable-diffusion-xl-1024-v1-0'}/text-to-image"; payload = {"text_prompts": [{"text": prompt}], "steps": steps if 'steps' in locals() else 30}; headers = {"Authorization": f"Bearer {STABILITY_API_KEY}", "Accept":"application/json", "Content-Type":"application/json"}; try: response = requests.post(request_url, headers=headers, json=payload, timeout=60); response.raise_for_status(); artifacts = response.json().get("artifacts"); img_data = base64.b64decode(artifacts[0]["base64"]); img = Image.open(BytesIO(img_data)); return ImageGenResponse(image=img, provider="StabilityAI")
+    except Exception as e: return ImageGenResponse(error=f"Stability AI Error: {str(e)}", success=False, provider="StabilityAI", raw_response=e)
+
+
+# --- DALL-E (Keep as is, conceptual) ---
+def generate_image_dalle(prompt: str, ...) -> ImageGenResponse:
+    # ... (your existing generate_image_dalle function)
+    if not OPENAI_DALLE_CONFIGURED: return ImageGenResponse(error="OpenAI DALL-E API key not configured.", success=False, provider="DALL-E")
+    dummy_image = Image.new('RGB', (512, 512), color = 'skyblue'); return ImageGenResponse(image=dummy_image, provider="DALL-E (Simulated)")
+
+
+# --- NEW: Hugging Face Image Model via Inference API ---
+def generate_image_hf_model(prompt: str, 
+                            model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", # A popular choice
+                            # model_id: str = "runwayml/stable-diffusion-v1-5", # Another option
+                            # model_id: str = "prompthero/openjourney", # Midjourney-like style
+                            negative_prompt: str = None,
+                            height: int = 768, # Adjust for different models
+                            width: int = 768,
+                            num_inference_steps: int = 25,
+                            guidance_scale: float = 7.5
+                            ) -> ImageGenResponse:
+    global hf_inference_image_client # Use the initialized client
+    if not HF_IMAGE_API_CONFIGURED 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,
-        "steps": steps,
-        "samples": 1,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale
     }
-    if style_preset: payload["style_preset"] = style_preset
-    if negative_prompt: payload["text_prompts"].append({"text": negative_prompt, "weight": -1.0})
+    # Remove None params as some models might not like them
+    params = {k: v for k, v in params.items() if v is not None}
 
-    headers = {
-        "Accept": "application/json",
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {STABILITY_API_KEY}"
-    }
-    print(f"DEBUG: image_services.py - Calling Stability AI with prompt: {prompt[:50]}...")
+    print(f"DEBUG: image_services.py - Calling HF Image API ({model_id}) with prompt: {prompt[:50]}...")
     try:
-        response = requests.post(request_url, headers=headers, json=payload, timeout=60) # Increased timeout
-        response.raise_for_status() # Will raise an HTTPError if the HTTP request returned an unsuccessful status code
-        
-        artifacts = response.json().get("artifacts")
-        if not artifacts:
-            return ImageGenResponse(error="No image artifacts found in Stability AI response.", success=False, provider="StabilityAI", raw_response=response.text)
-
-        image_data = base64.b64decode(artifacts[0]["base64"])
-        image = Image.open(BytesIO(image_data))
-        print("DEBUG: image_services.py - Stability AI image generated successfully.")
-        return ImageGenResponse(image=image, provider="StabilityAI")
-    except requests.exceptions.RequestException as e:
-        error_msg = f"Stability AI API request failed: {type(e).__name__} - {str(e)}"
-        if hasattr(e, 'response') and e.response is not None: error_msg += f" - Response: {e.response.text[:200]}"
-        print(f"ERROR: image_services.py - {error_msg}")
-        return ImageGenResponse(error=error_msg, success=False, provider="StabilityAI", raw_response=e)
-    except Exception as e: # Catch other potential errors like JSON decoding
-        error_msg = f"Error processing Stability AI response: {type(e).__name__} - {str(e)}"
-        print(f"ERROR: image_services.py - {error_msg}")
-        return ImageGenResponse(error=error_msg, success=False, provider="StabilityAI", raw_response=e)
+        # The `text_to_image` method of InferenceClient returns a PIL Image directly
+        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:
+        error_msg = f"HF Image API Error ({model_id}): {type(e).__name__} - {str(e)}"
+        # Check for common HF API errors
+        if "Rate limit reached" in str(e):
+            error_msg += " You may have hit free tier limits."
+        elif "Model is currently loading" in str(e):
+            error_msg += " The model might be loading, please try again in a moment."
+        elif "Authorization header is correct" in str(e) or "401" in str(e):
+             error_msg += " Issue with your HF_TOKEN authentication."
 
+        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)
 
-# --- DALL-E (Conceptual - you'd need 'openai' library and setup) ---
-def generate_image_dalle(prompt: str, size="1024x1024", quality="standard", n=1) -> ImageGenResponse:
-    if not OPENAI_DALLE_CONFIGURED:
-        return ImageGenResponse(error="OpenAI DALL-E API key not configured.", success=False, provider="DALL-E")
-    try:
-        # from openai import OpenAI # Would be imported at top level
-        # client = OpenAI(api_key=OPENAI_API_KEY)
-        # response = client.images.generate(
-        # model="dall-e-3", # or "dall-e-2"
-        # prompt=prompt,
-        # size=size,
-        # quality=quality,
-        # n=n,
-        # response_format="url" # or "b64_json"
-        # )
-        # image_url = response.data[0].url
-        # image_content = requests.get(image_url).content
-        # image = Image.open(BytesIO(image_content))
-        # return ImageGenResponse(image=image, image_url=image_url, provider="DALL-E")
-        print("DEBUG: image_services.py - DALL-E call placeholder.") # Placeholder
-        # Simulate an image for now
-        dummy_image = Image.new('RGB', (512, 512), color = 'skyblue')
-        return ImageGenResponse(image=dummy_image, provider="DALL-E (Simulated)")
-    except Exception as e:
-        return ImageGenResponse(error=f"DALL-E API Error: {type(e).__name__} - {str(e)}", success=False, provider="DALL-E")
-
-
-# --- Hugging Face Image Model (Conceptual - via Inference API or local Diffusers) ---
-# def generate_image_hf_model(prompt: str, model_id="stabilityai/stable-diffusion-xl-base-1.0") -> ImageGenResponse:
-# if not HF_IMAGE_CONFIGURED:
-# return ImageGenResponse(error="HF Token not configured for image models.", success=False, provider="HF")
-# try:
-# You might use a client similar to hf_inference_text_client but for image generation task
-# Or if it's a diffusers pipeline, you'd load and run it.
-# This requires the `diffusers` library and often significant compute.
-# response_bytes = hf_some_image_client.text_to_image(prompt, model=model_id) # Hypothetical client method
-# image = Image.open(BytesIO(response_bytes))
-# return ImageGenResponse(image=image, provider=f"HF ({model_id})")
-# print("DEBUG: image_services.py - HF Image Model call placeholder.")
-# dummy_image = Image.new('RGB', (512, 512), color = 'lightgreen')
-# return ImageGenResponse(image=dummy_image, provider=f"HF ({model_id} - Simulated)")
-# except Exception as e:
-# return ImageGenResponse(error=f"HF Image Model Error: {type(e).__name__} - {str(e)}", success=False, provider=f"HF ({model_id})")
-
-print("DEBUG: core.image_services (for StoryVerseWeaver) - Module defined.")
+print("DEBUG: core.image_services (for StoryVerseWeaver) - Module defined with HF Image support.")