| """ |
| Image Service |
| |
| High-level service for image generation and editing. |
| Abstracts all API complexity from the UI layer. |
| """ |
|
|
| from typing import Callable, Optional, List |
| from dataclasses import dataclass |
|
|
| from ..api.client import StackNetClient |
|
|
|
|
| @dataclass |
| class GeneratedImage: |
| """Generated image result.""" |
| image_url: str |
| image_path: Optional[str] = None |
| prompt: Optional[str] = None |
| width: Optional[int] = None |
| height: Optional[int] = None |
|
|
|
|
| class ImageService: |
| """ |
| Service for image generation and editing. |
| |
| Provides clean interfaces for: |
| - Text-to-image generation |
| - Image-to-image editing/transformation |
| """ |
|
|
| def __init__(self, client: Optional[StackNetClient] = None): |
| self.client = client or StackNetClient() |
|
|
| async def generate_image( |
| self, |
| prompt: str, |
| style: Optional[str] = None, |
| aspect_ratio: Optional[str] = None, |
| on_progress: Optional[Callable[[float, str], None]] = None |
| ) -> List[GeneratedImage]: |
| """ |
| Generate image from a text prompt using generate_image_5 tool. |
| |
| Args: |
| prompt: Description of desired image |
| style: Style preset (Photorealistic, Digital Art, etc.) |
| aspect_ratio: Aspect ratio (1:1, 16:9, 9:16, etc.) |
| on_progress: Callback for progress updates |
| |
| Returns: |
| List of generated images |
| """ |
| full_prompt = prompt |
| if style and style != "Photorealistic": |
| full_prompt = f"{prompt}, {style.lower()} style" |
|
|
| |
| width, height = 1024, 1024 |
| if aspect_ratio == "16:9": |
| width, height = 1280, 720 |
| elif aspect_ratio == "9:16": |
| width, height = 720, 1280 |
| elif aspect_ratio == "4:3": |
| width, height = 1024, 768 |
| elif aspect_ratio == "3:4": |
| width, height = 768, 1024 |
|
|
| result = await self.client.submit_tool_task( |
| tool_name="generate_image_5", |
| parameters={ |
| "prompt": full_prompt, |
| "width": width, |
| "height": height |
| }, |
| on_progress=on_progress |
| ) |
|
|
| if not result.success: |
| raise Exception(result.error or "Image generation failed") |
|
|
| return self._parse_image_result(result.data, prompt) |
|
|
| async def edit_image( |
| self, |
| image_url: str, |
| edit_prompt: str, |
| strength: float = 0.5, |
| on_progress: Optional[Callable[[float, str], None]] = None |
| ) -> List[GeneratedImage]: |
| """ |
| Edit/transform an existing image using generate_image_edit_5 tool. |
| |
| Args: |
| image_url: URL to source image |
| edit_prompt: Edit instructions |
| strength: Edit strength (0.1 to 1.0) |
| on_progress: Progress callback |
| |
| Returns: |
| List of edited images |
| """ |
| result = await self.client.submit_tool_task( |
| tool_name="generate_image_edit_5", |
| parameters={ |
| "prompt": edit_prompt, |
| "image_url": image_url, |
| "strength": strength |
| }, |
| on_progress=on_progress |
| ) |
|
|
| if not result.success: |
| raise Exception(result.error or "Image editing failed") |
|
|
| return self._parse_image_result(result.data, edit_prompt) |
|
|
| def _parse_image_result(self, data: dict, prompt: str) -> List[GeneratedImage]: |
| """Parse API response into GeneratedImage objects.""" |
| images = [] |
|
|
| |
| raw_images = data.get("images", []) |
|
|
| if not raw_images: |
| |
| image_url = ( |
| data.get("image_url") or |
| data.get("imageUrl") or |
| data.get("url") or |
| data.get("content") |
| ) |
| if image_url: |
| raw_images = [{"url": image_url}] |
|
|
| for img_data in raw_images: |
| if isinstance(img_data, str): |
| |
| image_url = img_data |
| else: |
| image_url = ( |
| img_data.get("url") or |
| img_data.get("image_url") or |
| img_data.get("imageUrl") |
| ) |
|
|
| if image_url: |
| images.append(GeneratedImage( |
| image_url=image_url, |
| prompt=prompt, |
| width=img_data.get("width") if isinstance(img_data, dict) else None, |
| height=img_data.get("height") if isinstance(img_data, dict) else None |
| )) |
|
|
| return images |
|
|
| async def download_image(self, image: GeneratedImage) -> str: |
| """Download an image to local file.""" |
| if image.image_path: |
| return image.image_path |
|
|
| |
| url = image.image_url |
| if ".png" in url: |
| ext = ".png" |
| elif ".jpg" in url or ".jpeg" in url: |
| ext = ".jpg" |
| else: |
| ext = ".png" |
|
|
| filename = f"image_{hash(url) % 10000}{ext}" |
| image.image_path = await self.client.download_file(url, filename) |
| return image.image_path |
|
|
| def cleanup(self): |
| """Clean up temporary files.""" |
| self.client.cleanup() |
|
|
