File size: 5,014 Bytes
957256e
 
 
 
 
 
 
 
 
 
26cd782
957256e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
982cefe
957256e
 
 
982cefe
957256e
 
 
982cefe
 
957256e
 
 
 
 
982cefe
 
 
 
 
 
 
26cd782
 
982cefe
26cd782
982cefe
26cd782
957256e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26cd782
957256e
 
 
 
 
 
 
 
 
 
26cd782
 
 
 
ed5dd8e
 
26cd782
957256e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
"""
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()