| """
|
| Image Generation
|
| Part of SOVEREIGN PYTHON LLM ENGINE
|
|
|
| Generate images using various providers (DALL-E, Stable Diffusion, etc.).
|
| """
|
|
|
| from typing import Any
|
| from pathlib import Path
|
| from dataclasses import dataclass
|
| import base64
|
| from io import BytesIO
|
|
|
| from ...core.evidence import WORMLedger
|
|
|
|
|
| @dataclass
|
| class GeneratedImage:
|
| """Generated image result"""
|
| image_bytes: bytes
|
| format: str
|
| size: tuple[int, int]
|
| prompt: str
|
| model: str
|
| metadata: dict[str, Any]
|
|
|
|
|
| class ImageGenerator:
|
| """
|
| Image generation using multiple providers.
|
|
|
| Supports:
|
| - OpenAI DALL-E 3
|
| - Stable Diffusion (via API or local)
|
| - Amazon Bedrock (Stable Diffusion)
|
| """
|
|
|
| def __init__(
|
| self,
|
| provider: str = "openai",
|
| model: str | None = None,
|
| worm_ledger: WORMLedger | None = None
|
| ):
|
| """
|
| Initialize image generator.
|
|
|
| Args:
|
| provider: Provider name ("openai", "bedrock", "stability")
|
| model: Optional model override
|
| worm_ledger: Optional WORM ledger
|
| """
|
| self.provider = provider
|
| self.model = model
|
| self.worm_ledger = worm_ledger
|
|
|
| async def generate(
|
| self,
|
| prompt: str,
|
| size: str = "1024x1024",
|
| quality: str = "standard",
|
| n: int = 1
|
| ) -> list[GeneratedImage]:
|
| """
|
| Generate images from prompt.
|
|
|
| Args:
|
| prompt: Text prompt
|
| size: Image size ("1024x1024", "1792x1024", etc.)
|
| quality: Quality level ("standard", "hd")
|
| n: Number of images
|
|
|
| Returns:
|
| List of GeneratedImage
|
| """
|
| if self.provider == "openai":
|
| return await self._generate_openai(prompt, size, quality, n)
|
| elif self.provider == "bedrock":
|
| return await self._generate_bedrock(prompt, size)
|
| elif self.provider == "stability":
|
| return await self._generate_stability(prompt, size)
|
| else:
|
| raise ValueError(f"Unknown provider: {self.provider}")
|
|
|
| async def _generate_openai(
|
| self,
|
| prompt: str,
|
| size: str,
|
| quality: str,
|
| n: int
|
| ) -> list[GeneratedImage]:
|
| """Generate using OpenAI DALL-E"""
|
| try:
|
| from openai import AsyncOpenAI
|
| except ImportError:
|
| raise ImportError("openai is required. Install: pip install openai")
|
|
|
| client = AsyncOpenAI()
|
|
|
| response = await client.images.generate(
|
| model=self.model or "dall-e-3",
|
| prompt=prompt,
|
| size=size,
|
| quality=quality,
|
| n=n,
|
| response_format="b64_json"
|
| )
|
|
|
| results = []
|
|
|
| for image_data in response.data:
|
|
|
| image_bytes = base64.b64decode(image_data.b64_json)
|
|
|
|
|
| width, height = map(int, size.split("x"))
|
|
|
| results.append(GeneratedImage(
|
| image_bytes=image_bytes,
|
| format="png",
|
| size=(width, height),
|
| prompt=prompt,
|
| model=self.model or "dall-e-3",
|
| metadata={
|
| "revised_prompt": image_data.revised_prompt
|
| }
|
| ))
|
|
|
|
|
| if self.worm_ledger:
|
| await self.worm_ledger.append({
|
| "event": "image_generate",
|
| "provider": "openai",
|
| "prompt": prompt,
|
| "images_count": len(results)
|
| })
|
|
|
| return results
|
|
|
| async def _generate_bedrock(
|
| self,
|
| prompt: str,
|
| size: str
|
| ) -> list[GeneratedImage]:
|
| """Generate using AWS Bedrock (Stable Diffusion)"""
|
| try:
|
| import boto3
|
| import json
|
| except ImportError:
|
| raise ImportError("boto3 is required")
|
|
|
| client = boto3.client("bedrock-runtime", region_name="us-east-1")
|
|
|
|
|
| width, height = map(int, size.split("x"))
|
|
|
|
|
| request_body = {
|
| "text_prompts": [{"text": prompt}],
|
| "cfg_scale": 7,
|
| "steps": 30,
|
| "width": width,
|
| "height": height
|
| }
|
|
|
|
|
| response = client.invoke_model(
|
| modelId=self.model or "stability.stable-diffusion-xl-v1",
|
| body=json.dumps(request_body)
|
| )
|
|
|
|
|
| response_body = json.loads(response["body"].read())
|
|
|
| results = []
|
|
|
| for artifact in response_body.get("artifacts", []):
|
| image_bytes = base64.b64decode(artifact["base64"])
|
|
|
| results.append(GeneratedImage(
|
| image_bytes=image_bytes,
|
| format="png",
|
| size=(width, height),
|
| prompt=prompt,
|
| model=self.model or "stability.stable-diffusion-xl-v1",
|
| metadata={}
|
| ))
|
|
|
|
|
| if self.worm_ledger:
|
| await self.worm_ledger.append({
|
| "event": "image_generate",
|
| "provider": "bedrock",
|
| "prompt": prompt,
|
| "images_count": len(results)
|
| })
|
|
|
| return results
|
|
|
| async def _generate_stability(
|
| self,
|
| prompt: str,
|
| size: str
|
| ) -> list[GeneratedImage]:
|
| """Generate using Stability AI API"""
|
| try:
|
| import httpx
|
| except ImportError:
|
| raise ImportError("httpx is required")
|
|
|
|
|
| raise NotImplementedError("Stability AI provider not fully implemented")
|
|
|
|
|
| async def save_image(image: GeneratedImage, output_path: Path) -> None:
|
| """
|
| Save generated image to file.
|
|
|
| Args:
|
| image: Generated image
|
| output_path: Where to save
|
| """
|
| output_path.write_bytes(image.image_bytes)
|
|
|
|
|
|
|
| async def generate_image_tool(
|
| prompt: str,
|
| provider: str = "bedrock",
|
| size: str = "1024x1024",
|
| output_path: str | None = None
|
| ) -> dict:
|
| """
|
| Tool wrapper for image generation.
|
|
|
| Args:
|
| prompt: Image prompt
|
| provider: Provider ("openai", "bedrock")
|
| size: Image size
|
| output_path: Optional path to save
|
|
|
| Returns:
|
| Dictionary with result
|
| """
|
| generator = ImageGenerator(provider=provider)
|
| images = await generator.generate(prompt, size=size)
|
|
|
| result = {
|
| "images_count": len(images),
|
| "prompt": prompt,
|
| "size": size,
|
| "model": images[0].model if images else None
|
| }
|
|
|
|
|
| if output_path and images:
|
| await save_image(images[0], Path(output_path))
|
| result["saved_to"] = output_path
|
|
|
| return result
|
|
|
