""" 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, format_type: str = "image", on_progress: Optional[Callable[[float, str], None]] = None ) -> List[GeneratedImage]: """ Generate image from a text prompt. Args: prompt: Description of desired image format_type: Format type - "image" (generate_image_5), "multi" (generate_image_multi_5), or "3d" (generate_image_3d) on_progress: Callback for progress updates Returns: List of generated images """ # Select tool based on format type if format_type == "multi": tool_name = "generate_image_multi_5" elif format_type == "3d": tool_name = "generate_image_3d" else: tool_name = "generate_image_5" result = await self.client.submit_tool_task( tool_name=tool_name, parameters={ "prompt": prompt }, 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 = [] # Handle various response formats raw_images = data.get("images", []) if not raw_images: # Check for single image URL 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): # Raw URL string 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 # Determine extension from URL 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()