Create image_services.py

core/image_services.py

@@ -0,0 +1,132 @@
+# storyverse_weaver/core/image_services.py
+import os
+import requests # For generic API calls
+import base64
+from io import BytesIO
+from PIL import Image
+# from dotenv import load_dotenv
+# load_dotenv()
+
+# --- API Key Configuration (Use specific names for StoryVerse) ---
+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
+
+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
+
+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
+
+def initialize_image_llms(): # Simple function to print status
+    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,
+        "height": height,
+        "width": width,
+        "steps": steps,
+        "samples": 1,
+    }
+    if style_preset: payload["style_preset"] = style_preset
+    if negative_prompt: payload["text_prompts"].append({"text": negative_prompt, "weight": -1.0})
+
+    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]}...")
+    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)
+
+
+# --- 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.")