webscout 6.3__py3-none-any.whl → 6.4__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of webscout might be problematic. Click here for more details.
- webscout/AIauto.py +191 -176
- webscout/AIbase.py +0 -197
- webscout/AIutel.py +488 -1130
- webscout/Bing_search.py +250 -153
- webscout/DWEBS.py +151 -19
- webscout/Extra/__init__.py +2 -1
- webscout/Extra/autocoder/__init__.py +9 -0
- webscout/Extra/autocoder/autocoder_utiles.py +121 -0
- webscout/Extra/autocoder/rawdog.py +681 -0
- webscout/Extra/autollama.py +246 -195
- webscout/Extra/gguf.py +441 -416
- webscout/LLM.py +206 -43
- webscout/Litlogger/__init__.py +681 -0
- webscout/Provider/DARKAI.py +1 -1
- webscout/Provider/EDITEE.py +1 -1
- webscout/Provider/NinjaChat.py +1 -1
- webscout/Provider/PI.py +221 -207
- webscout/Provider/Perplexity.py +598 -598
- webscout/Provider/RoboCoders.py +206 -0
- webscout/Provider/TTI/AiForce/__init__.py +22 -0
- webscout/Provider/TTI/AiForce/async_aiforce.py +257 -0
- webscout/Provider/TTI/AiForce/sync_aiforce.py +242 -0
- webscout/Provider/TTI/Nexra/__init__.py +22 -0
- webscout/Provider/TTI/Nexra/async_nexra.py +286 -0
- webscout/Provider/TTI/Nexra/sync_nexra.py +258 -0
- webscout/Provider/TTI/PollinationsAI/__init__.py +23 -0
- webscout/Provider/TTI/PollinationsAI/async_pollinations.py +330 -0
- webscout/Provider/TTI/PollinationsAI/sync_pollinations.py +285 -0
- webscout/Provider/TTI/__init__.py +2 -4
- webscout/Provider/TTI/artbit/__init__.py +22 -0
- webscout/Provider/TTI/artbit/async_artbit.py +184 -0
- webscout/Provider/TTI/artbit/sync_artbit.py +176 -0
- webscout/Provider/TTI/blackbox/__init__.py +4 -0
- webscout/Provider/TTI/blackbox/async_blackbox.py +212 -0
- webscout/Provider/TTI/{blackboximage.py → blackbox/sync_blackbox.py} +199 -153
- webscout/Provider/TTI/deepinfra/__init__.py +4 -0
- webscout/Provider/TTI/deepinfra/async_deepinfra.py +227 -0
- webscout/Provider/TTI/deepinfra/sync_deepinfra.py +199 -0
- webscout/Provider/TTI/huggingface/__init__.py +22 -0
- webscout/Provider/TTI/huggingface/async_huggingface.py +199 -0
- webscout/Provider/TTI/huggingface/sync_huggingface.py +195 -0
- webscout/Provider/TTI/imgninza/__init__.py +4 -0
- webscout/Provider/TTI/imgninza/async_ninza.py +214 -0
- webscout/Provider/TTI/{imgninza.py → imgninza/sync_ninza.py} +209 -136
- webscout/Provider/TTI/talkai/__init__.py +4 -0
- webscout/Provider/TTI/talkai/async_talkai.py +229 -0
- webscout/Provider/TTI/talkai/sync_talkai.py +207 -0
- webscout/Provider/__init__.py +146 -139
- webscout/Provider/askmyai.py +2 -2
- webscout/Provider/cerebras.py +227 -219
- webscout/Provider/llama3mitril.py +0 -1
- webscout/Provider/mhystical.py +176 -0
- webscout/Provider/perplexitylabs.py +265 -0
- webscout/Provider/twitterclone.py +251 -245
- webscout/Provider/typegpt.py +359 -0
- webscout/__init__.py +28 -23
- webscout/__main__.py +5 -5
- webscout/cli.py +252 -280
- webscout/conversation.py +227 -0
- webscout/exceptions.py +161 -29
- webscout/litagent/__init__.py +172 -0
- webscout/litprinter/__init__.py +831 -0
- webscout/optimizers.py +270 -0
- webscout/prompt_manager.py +279 -0
- webscout/swiftcli/__init__.py +810 -0
- webscout/transcriber.py +479 -551
- webscout/update_checker.py +125 -0
- webscout/version.py +1 -1
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/METADATA +26 -45
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/RECORD +75 -45
- webscout/Provider/TTI/AIuncensoredimage.py +0 -103
- webscout/Provider/TTI/Nexra.py +0 -120
- webscout/Provider/TTI/PollinationsAI.py +0 -138
- webscout/Provider/TTI/WebSimAI.py +0 -142
- webscout/Provider/TTI/aiforce.py +0 -160
- webscout/Provider/TTI/artbit.py +0 -141
- webscout/Provider/TTI/deepinfra.py +0 -148
- webscout/Provider/TTI/huggingface.py +0 -155
- webscout/Provider/TTI/talkai.py +0 -116
- webscout/models.py +0 -23
- /webscout/{g4f.py → gpt4free.py} +0 -0
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/LICENSE.md +0 -0
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/WHEEL +0 -0
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/entry_points.txt +0 -0
- {webscout-6.3.dist-info → webscout-6.4.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,206 @@
|
|
|
1
|
+
import requests
|
|
2
|
+
import json
|
|
3
|
+
from typing import Any, Dict, Optional, Generator, Union
|
|
4
|
+
|
|
5
|
+
from webscout.AIutel import Optimizers
|
|
6
|
+
from webscout.AIutel import Conversation
|
|
7
|
+
from webscout.AIutel import AwesomePrompts
|
|
8
|
+
from webscout.AIbase import Provider
|
|
9
|
+
from webscout import exceptions
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class RoboCoders(Provider):
|
|
13
|
+
"""
|
|
14
|
+
A class to interact with the RoboCoders API.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
api_endpoint = "https://api.robocoders.ai/chat"
|
|
18
|
+
working = True
|
|
19
|
+
supports_message_history = True
|
|
20
|
+
default_model = "GeneralCodingAgent"
|
|
21
|
+
agent = [default_model, "RepoAgent", "FrontEndAgent"]
|
|
22
|
+
models = [*agent]
|
|
23
|
+
|
|
24
|
+
def __init__(
|
|
25
|
+
self,
|
|
26
|
+
is_conversation: bool = True,
|
|
27
|
+
max_tokens: int = 600,
|
|
28
|
+
timeout: int = 30,
|
|
29
|
+
intro: str = None,
|
|
30
|
+
filepath: str = None,
|
|
31
|
+
update_file: bool = True,
|
|
32
|
+
proxies: dict = {},
|
|
33
|
+
history_offset: int = 10250,
|
|
34
|
+
act: str = None,
|
|
35
|
+
model: str = default_model,
|
|
36
|
+
system_prompt: str = "You are a helpful coding assistant.",
|
|
37
|
+
):
|
|
38
|
+
"""Initializes the RoboCoders API client."""
|
|
39
|
+
if model not in self.models:
|
|
40
|
+
raise ValueError(f"Invalid model: {model}. Choose from: {', '.join(self.models)}")
|
|
41
|
+
|
|
42
|
+
self.session = requests.Session()
|
|
43
|
+
self.is_conversation = is_conversation
|
|
44
|
+
self.max_tokens_to_sample = max_tokens
|
|
45
|
+
self.timeout = timeout
|
|
46
|
+
self.last_response = {}
|
|
47
|
+
self.model = model
|
|
48
|
+
self.system_prompt = system_prompt
|
|
49
|
+
|
|
50
|
+
self.headers = {"Content-Type": "application/json"}
|
|
51
|
+
self.access_token = self._get_access_token() # Get token on initialization
|
|
52
|
+
if not self.access_token:
|
|
53
|
+
raise exceptions.AuthenticationError("Failed to get access token")
|
|
54
|
+
|
|
55
|
+
self.session_id = self._create_session() # Create session on initialization
|
|
56
|
+
if not self.session_id:
|
|
57
|
+
raise exceptions.SessionCreationError("Failed to create session")
|
|
58
|
+
|
|
59
|
+
self.headers["Authorization"] = f"Bearer {self.access_token}"
|
|
60
|
+
|
|
61
|
+
self.__available_optimizers = (
|
|
62
|
+
method
|
|
63
|
+
for method in dir(Optimizers)
|
|
64
|
+
if callable(getattr(Optimizers, method)) and not method.startswith("__")
|
|
65
|
+
)
|
|
66
|
+
Conversation.intro = (
|
|
67
|
+
AwesomePrompts().get_act(
|
|
68
|
+
act, raise_not_found=True, default=None, case_insensitive=True
|
|
69
|
+
)
|
|
70
|
+
if act
|
|
71
|
+
else intro or Conversation.intro
|
|
72
|
+
)
|
|
73
|
+
self.conversation = Conversation(
|
|
74
|
+
is_conversation, self.max_tokens_to_sample, filepath, update_file
|
|
75
|
+
)
|
|
76
|
+
self.conversation.history_offset = history_offset
|
|
77
|
+
self.session.proxies = proxies
|
|
78
|
+
|
|
79
|
+
def _get_access_token(self) -> Optional[str]:
|
|
80
|
+
"""Get access token for authentication."""
|
|
81
|
+
url_auth = "https://api.robocoders.ai/auth"
|
|
82
|
+
headers_auth = {
|
|
83
|
+
"accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
|
|
84
|
+
"accept-language": "en-US,en;q=0.9",
|
|
85
|
+
"user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36",
|
|
86
|
+
}
|
|
87
|
+
try:
|
|
88
|
+
response = self.session.get(url_auth, headers=headers_auth, timeout=self.timeout)
|
|
89
|
+
response.raise_for_status() # Raise exception for HTTP errors
|
|
90
|
+
text = response.text
|
|
91
|
+
return text.split('id="token">')[1].split("</pre>")[0].strip()
|
|
92
|
+
except (requests.exceptions.RequestException, IndexError) as e:
|
|
93
|
+
raise exceptions.APIConnectionError(f"Failed to get access token: {e}")
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def _create_session(self) -> Optional[str]:
|
|
97
|
+
"""Create a new chat session."""
|
|
98
|
+
url_create_session = "https://api.robocoders.ai/create-session"
|
|
99
|
+
headers_create_session = {"Authorization": f"Bearer {self.access_token}"}
|
|
100
|
+
try:
|
|
101
|
+
response = self.session.get(url_create_session, headers=headers_create_session, timeout=self.timeout)
|
|
102
|
+
response.raise_for_status()
|
|
103
|
+
data = response.json()
|
|
104
|
+
return data.get("sid")
|
|
105
|
+
except requests.exceptions.RequestException as e:
|
|
106
|
+
raise exceptions.APIConnectionError(f"Failed to create session: {e}")
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def ask(
|
|
111
|
+
self,
|
|
112
|
+
prompt: str,
|
|
113
|
+
stream: bool = False,
|
|
114
|
+
raw: bool = False,
|
|
115
|
+
optimizer: str = None,
|
|
116
|
+
conversationally: bool = False,
|
|
117
|
+
) -> Union[Dict, Generator[str, None, None]]:
|
|
118
|
+
"""
|
|
119
|
+
Sends a prompt to the RoboCoders API and returns the response.
|
|
120
|
+
"""
|
|
121
|
+
conversation_prompt = self.conversation.gen_complete_prompt(prompt)
|
|
122
|
+
if optimizer:
|
|
123
|
+
if optimizer in self.__available_optimizers:
|
|
124
|
+
conversation_prompt = getattr(Optimizers, optimizer)(
|
|
125
|
+
conversation_prompt if conversationally else prompt
|
|
126
|
+
)
|
|
127
|
+
else:
|
|
128
|
+
raise exceptions.FailedToGenerateResponseError(
|
|
129
|
+
f"Optimizer is not one of {self.__available_optimizers}"
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
data = {
|
|
133
|
+
"sid": self.session_id,
|
|
134
|
+
"prompt": conversation_prompt,
|
|
135
|
+
"agent": self.model,
|
|
136
|
+
}
|
|
137
|
+
|
|
138
|
+
def generate_chunks():
|
|
139
|
+
response = self._make_request(
|
|
140
|
+
"POST", self.api_endpoint, headers=self.headers, json=data, stream=True
|
|
141
|
+
)
|
|
142
|
+
for line in response.iter_lines():
|
|
143
|
+
if line:
|
|
144
|
+
try:
|
|
145
|
+
response_data = json.loads(line)
|
|
146
|
+
message = response_data.get("message", "")
|
|
147
|
+
if message:
|
|
148
|
+
yield message
|
|
149
|
+
except json.JSONDecodeError:
|
|
150
|
+
pass # Handle or log the error as needed
|
|
151
|
+
|
|
152
|
+
if stream:
|
|
153
|
+
streaming_text = ""
|
|
154
|
+
for chunk in generate_chunks():
|
|
155
|
+
streaming_text += chunk
|
|
156
|
+
yield chunk if raw else {"text": chunk}
|
|
157
|
+
self.last_response.update({"text": streaming_text}) #Update last_response
|
|
158
|
+
self.conversation.update_chat_history(prompt, streaming_text)
|
|
159
|
+
else:
|
|
160
|
+
full_response = "".join(generate_chunks())
|
|
161
|
+
self.last_response.update({"text": full_response})
|
|
162
|
+
self.conversation.update_chat_history(prompt, full_response)
|
|
163
|
+
return self.last_response
|
|
164
|
+
|
|
165
|
+
def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
|
|
166
|
+
try:
|
|
167
|
+
response = self.session.request(
|
|
168
|
+
method, url, timeout=self.timeout, **kwargs
|
|
169
|
+
)
|
|
170
|
+
response.raise_for_status()
|
|
171
|
+
return response
|
|
172
|
+
except requests.exceptions.RequestException as e:
|
|
173
|
+
raise exceptions.APIConnectionError(f"Failed to make request: {e}") from e
|
|
174
|
+
|
|
175
|
+
def chat(
|
|
176
|
+
self,
|
|
177
|
+
prompt: str,
|
|
178
|
+
stream: bool = False,
|
|
179
|
+
optimizer: str = None,
|
|
180
|
+
conversationally: bool = False,
|
|
181
|
+
) -> Union[str, Generator[str, None, None]]:
|
|
182
|
+
"""Generate response string or stream."""
|
|
183
|
+
|
|
184
|
+
if stream:
|
|
185
|
+
gen = self.ask(
|
|
186
|
+
prompt, stream=True, optimizer=optimizer, conversationally=conversationally
|
|
187
|
+
)
|
|
188
|
+
for chunk in gen:
|
|
189
|
+
yield self.get_message(chunk)
|
|
190
|
+
else:
|
|
191
|
+
return self.get_message(
|
|
192
|
+
self.ask(prompt, stream=False, optimizer=optimizer, conversationally=conversationally)
|
|
193
|
+
)
|
|
194
|
+
|
|
195
|
+
def get_message(self, response: Dict[str, Any]) -> str:
|
|
196
|
+
"""Retrieves message from response."""
|
|
197
|
+
assert isinstance(response, dict), "Response should be of dict data-type only"
|
|
198
|
+
return response["text"]
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
if __name__ == "__main__":
|
|
202
|
+
from rich import print
|
|
203
|
+
ai = RoboCoders(model="GeneralCodingAgent")
|
|
204
|
+
response = ai.chat("Can you help me with Python programming?", stream=True)
|
|
205
|
+
for chunk in response:
|
|
206
|
+
print(chunk, end='', flush=True)
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AiForce - Your go-to provider for generating fire images! 🔥
|
|
3
|
+
|
|
4
|
+
Examples:
|
|
5
|
+
>>> # Sync Usage
|
|
6
|
+
>>> from webscout import AiForceimager
|
|
7
|
+
>>> provider = AiForceimager()
|
|
8
|
+
>>> images = provider.generate("Cool art")
|
|
9
|
+
>>> paths = provider.save(images)
|
|
10
|
+
>>>
|
|
11
|
+
>>> # Async Usage
|
|
12
|
+
>>> from webscout import AsyncAiForceimager
|
|
13
|
+
>>> async def example():
|
|
14
|
+
... provider = AsyncAiForceimager()
|
|
15
|
+
... images = await provider.generate("Epic dragon")
|
|
16
|
+
... paths = await provider.save(images)
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from .sync_aiforce import AiForceimager
|
|
20
|
+
from .async_aiforce import AsyncAiForceimager
|
|
21
|
+
|
|
22
|
+
__all__ = ["AiForceimager", "AsyncAiForceimager"]
|
|
@@ -0,0 +1,257 @@
|
|
|
1
|
+
import aiohttp
|
|
2
|
+
import asyncio
|
|
3
|
+
import os
|
|
4
|
+
import time
|
|
5
|
+
from typing import List, Optional, Union, AsyncGenerator
|
|
6
|
+
from string import punctuation
|
|
7
|
+
from random import choice
|
|
8
|
+
from aiohttp import ClientError
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
|
|
11
|
+
from webscout.AIbase import AsyncImageProvider
|
|
12
|
+
from webscout.Litlogger import LitLogger, LogFormat, ColorScheme
|
|
13
|
+
from webscout.litagent import LitAgent
|
|
14
|
+
|
|
15
|
+
# Initialize our fire logger and agent 🔥
|
|
16
|
+
logger = LitLogger(
|
|
17
|
+
"AsyncAiForce",
|
|
18
|
+
format=LogFormat.MODERN_EMOJI,
|
|
19
|
+
color_scheme=ColorScheme.CYBERPUNK
|
|
20
|
+
)
|
|
21
|
+
agent = LitAgent()
|
|
22
|
+
|
|
23
|
+
class AsyncAiForceimager(AsyncImageProvider):
|
|
24
|
+
"""Your go-to async provider for generating fire images with AiForce! ⚡
|
|
25
|
+
|
|
26
|
+
Examples:
|
|
27
|
+
>>> # Basic usage
|
|
28
|
+
>>> provider = AsyncAiForceimager()
|
|
29
|
+
>>> async def example():
|
|
30
|
+
... images = await provider.generate("Cool art")
|
|
31
|
+
... paths = await provider.save(images)
|
|
32
|
+
>>>
|
|
33
|
+
>>> # Advanced usage
|
|
34
|
+
>>> provider = AsyncAiForceimager(timeout=120)
|
|
35
|
+
>>> async def example():
|
|
36
|
+
... images = await provider.generate(
|
|
37
|
+
... prompt="Epic dragon",
|
|
38
|
+
... amount=3,
|
|
39
|
+
... model="Flux-1.1-Pro",
|
|
40
|
+
... width=1024,
|
|
41
|
+
... height=1024
|
|
42
|
+
... )
|
|
43
|
+
... paths = await provider.save(images, name="dragon", dir="my_art")
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
AVAILABLE_MODELS = [
|
|
47
|
+
"stable-diffusion-xl-lightning",
|
|
48
|
+
"stable-diffusion-xl-base",
|
|
49
|
+
"Flux-1.1-Pro",
|
|
50
|
+
"ideogram",
|
|
51
|
+
"flux",
|
|
52
|
+
"flux-realism",
|
|
53
|
+
"flux-anime",
|
|
54
|
+
"flux-3d",
|
|
55
|
+
"flux-disney",
|
|
56
|
+
"flux-pixel",
|
|
57
|
+
"flux-4o",
|
|
58
|
+
"any-dark"
|
|
59
|
+
]
|
|
60
|
+
|
|
61
|
+
def __init__(self, timeout: int = 60, proxies: dict = {}, logging: bool = True):
|
|
62
|
+
"""Initialize your async AiForce provider with custom settings! ⚙️
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
timeout (int): Request timeout in seconds (default: 60)
|
|
66
|
+
proxies (dict): Proxy settings for requests (default: {})
|
|
67
|
+
logging (bool): Enable fire logging (default: True)
|
|
68
|
+
"""
|
|
69
|
+
self.api_endpoint = "https://api.airforce/imagine2"
|
|
70
|
+
self.headers = {
|
|
71
|
+
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
|
|
72
|
+
"Accept-Language": "en-US,en;q=0.5",
|
|
73
|
+
"Accept-Encoding": "gzip, deflate",
|
|
74
|
+
"User-Agent": agent.random()
|
|
75
|
+
}
|
|
76
|
+
self.timeout = timeout
|
|
77
|
+
self.proxies = proxies
|
|
78
|
+
self.prompt: str = "AI-generated image - webscout"
|
|
79
|
+
self.image_extension: str = "png"
|
|
80
|
+
self.logging = logging
|
|
81
|
+
if self.logging:
|
|
82
|
+
logger.info("AsyncAiForce provider initialized! 🚀")
|
|
83
|
+
|
|
84
|
+
async def generate(
|
|
85
|
+
self,
|
|
86
|
+
prompt: str,
|
|
87
|
+
amount: int = 1,
|
|
88
|
+
additives: bool = True,
|
|
89
|
+
model: str = "Flux-1.1-Pro",
|
|
90
|
+
width: int = 768,
|
|
91
|
+
height: int = 768,
|
|
92
|
+
seed: Optional[int] = None,
|
|
93
|
+
max_retries: int = 3,
|
|
94
|
+
retry_delay: int = 5
|
|
95
|
+
) -> List[bytes]:
|
|
96
|
+
"""Generate some fire images from your prompt asynchronously! ⚡
|
|
97
|
+
|
|
98
|
+
Examples:
|
|
99
|
+
>>> provider = AsyncAiForceimager()
|
|
100
|
+
>>> async def example():
|
|
101
|
+
... # Basic usage
|
|
102
|
+
... images = await provider.generate("Cool art")
|
|
103
|
+
... # Advanced usage
|
|
104
|
+
... images = await provider.generate(
|
|
105
|
+
... prompt="Epic dragon",
|
|
106
|
+
... amount=3,
|
|
107
|
+
... model="Flux-1.1-Pro"
|
|
108
|
+
... )
|
|
109
|
+
|
|
110
|
+
Args:
|
|
111
|
+
prompt (str): Your image description
|
|
112
|
+
amount (int): How many images you want (default: 1)
|
|
113
|
+
additives (bool): Make each prompt unique (default: True)
|
|
114
|
+
model (str): Model to use - check AVAILABLE_MODELS (default: "Flux-1.1-Pro")
|
|
115
|
+
width (int): Image width (default: 768)
|
|
116
|
+
height (int): Image height (default: 768)
|
|
117
|
+
seed (int, optional): Seed for reproducible results
|
|
118
|
+
max_retries (int): Max retry attempts if something fails (default: 3)
|
|
119
|
+
retry_delay (int): Seconds to wait between retries (default: 5)
|
|
120
|
+
|
|
121
|
+
Returns:
|
|
122
|
+
List[bytes]: Your generated images
|
|
123
|
+
|
|
124
|
+
Raises:
|
|
125
|
+
ValueError: If the inputs ain't valid
|
|
126
|
+
ClientError: If the API calls fail after retries
|
|
127
|
+
"""
|
|
128
|
+
assert bool(prompt), "Prompt cannot be null"
|
|
129
|
+
assert isinstance(amount, int), f"Amount should be an integer only not {type(amount)}"
|
|
130
|
+
assert amount > 0, "Amount should be greater than 0"
|
|
131
|
+
assert model in self.AVAILABLE_MODELS, f"Model should be one of {self.AVAILABLE_MODELS}"
|
|
132
|
+
|
|
133
|
+
ads = lambda: (
|
|
134
|
+
""
|
|
135
|
+
if not additives
|
|
136
|
+
else choice(punctuation)
|
|
137
|
+
+ choice(punctuation)
|
|
138
|
+
+ choice(punctuation)
|
|
139
|
+
+ choice(punctuation)
|
|
140
|
+
+ choice(punctuation)
|
|
141
|
+
)
|
|
142
|
+
|
|
143
|
+
self.prompt = prompt
|
|
144
|
+
response = []
|
|
145
|
+
|
|
146
|
+
if self.logging:
|
|
147
|
+
logger.info(f"Generating {amount} images with {model}... 🎨")
|
|
148
|
+
|
|
149
|
+
async with aiohttp.ClientSession(headers=self.headers) as session:
|
|
150
|
+
for i in range(amount):
|
|
151
|
+
url = f"{self.api_endpoint}?model={model}&prompt={prompt}&size={width}:{height}"
|
|
152
|
+
if seed:
|
|
153
|
+
url += f"&seed={seed}"
|
|
154
|
+
|
|
155
|
+
for attempt in range(max_retries):
|
|
156
|
+
try:
|
|
157
|
+
async with session.get(url, timeout=self.timeout, proxy=self.proxies.get('http')) as resp:
|
|
158
|
+
resp.raise_for_status()
|
|
159
|
+
response.append(await resp.read())
|
|
160
|
+
if self.logging:
|
|
161
|
+
logger.success(f"Generated image {i + 1}/{amount}! 🎨")
|
|
162
|
+
break
|
|
163
|
+
except ClientError as e:
|
|
164
|
+
if attempt == max_retries - 1:
|
|
165
|
+
if self.logging:
|
|
166
|
+
logger.error(f"Failed to generate image after {max_retries} attempts: {e} 😢")
|
|
167
|
+
raise
|
|
168
|
+
else:
|
|
169
|
+
if self.logging:
|
|
170
|
+
logger.warning(f"Attempt {attempt + 1} failed. Retrying in {retry_delay} seconds... 🔄")
|
|
171
|
+
await asyncio.sleep(retry_delay)
|
|
172
|
+
|
|
173
|
+
if self.logging:
|
|
174
|
+
logger.success("Images generated successfully! 🎉")
|
|
175
|
+
return response
|
|
176
|
+
|
|
177
|
+
async def save(
|
|
178
|
+
self,
|
|
179
|
+
response: Union[AsyncGenerator[bytes, None], List[bytes]],
|
|
180
|
+
name: Optional[str] = None,
|
|
181
|
+
dir: Optional[Union[str, Path]] = None,
|
|
182
|
+
filenames_prefix: str = "",
|
|
183
|
+
) -> List[str]:
|
|
184
|
+
"""Save your fire generated images asynchronously! 💾
|
|
185
|
+
|
|
186
|
+
Examples:
|
|
187
|
+
>>> provider = AsyncAiForceimager()
|
|
188
|
+
>>> async def example():
|
|
189
|
+
... images = await provider.generate("Cool art")
|
|
190
|
+
... # Save with default settings
|
|
191
|
+
... paths = await provider.save(images)
|
|
192
|
+
... # Save with custom name and directory
|
|
193
|
+
... paths = await provider.save(
|
|
194
|
+
... images,
|
|
195
|
+
... name="my_art",
|
|
196
|
+
... dir="my_images",
|
|
197
|
+
... filenames_prefix="test_"
|
|
198
|
+
... )
|
|
199
|
+
|
|
200
|
+
Args:
|
|
201
|
+
response (Union[AsyncGenerator[bytes, None], List[bytes]]): Your generated images
|
|
202
|
+
name (Optional[str]): Custom name for your images
|
|
203
|
+
dir (Optional[Union[str, Path]]): Where to save the images (default: current directory)
|
|
204
|
+
filenames_prefix (str): Prefix for your image files
|
|
205
|
+
|
|
206
|
+
Returns:
|
|
207
|
+
List[str]: Paths to your saved images
|
|
208
|
+
"""
|
|
209
|
+
save_dir = dir if dir else os.getcwd()
|
|
210
|
+
if not os.path.exists(save_dir):
|
|
211
|
+
os.makedirs(save_dir)
|
|
212
|
+
if self.logging:
|
|
213
|
+
logger.info(f"Created directory: {save_dir} 📁")
|
|
214
|
+
|
|
215
|
+
name = self.prompt if name is None else name
|
|
216
|
+
saved_paths = []
|
|
217
|
+
timestamp = int(time.time())
|
|
218
|
+
|
|
219
|
+
if self.logging:
|
|
220
|
+
logger.info(f"Saving {len(response)} images... 💾")
|
|
221
|
+
|
|
222
|
+
async def save_single_image(image_bytes: bytes, index: int) -> str:
|
|
223
|
+
filename = f"{filenames_prefix}{name}_{index}.{self.image_extension}"
|
|
224
|
+
filepath = os.path.join(save_dir, filename)
|
|
225
|
+
|
|
226
|
+
# Write file using asyncio
|
|
227
|
+
async with asyncio.Lock():
|
|
228
|
+
with open(filepath, "wb") as f:
|
|
229
|
+
f.write(image_bytes)
|
|
230
|
+
|
|
231
|
+
if self.logging:
|
|
232
|
+
logger.success(f"Saved image to: {filepath} 💾")
|
|
233
|
+
return filepath
|
|
234
|
+
|
|
235
|
+
# Handle both List[bytes] and AsyncGenerator
|
|
236
|
+
if isinstance(response, list):
|
|
237
|
+
image_list = response
|
|
238
|
+
else:
|
|
239
|
+
image_list = [chunk async for chunk in response]
|
|
240
|
+
|
|
241
|
+
tasks = [save_single_image(img, i) for i, img in enumerate(image_list)]
|
|
242
|
+
saved_paths = await asyncio.gather(*tasks)
|
|
243
|
+
if self.logging:
|
|
244
|
+
logger.success(f"Images saved successfully! Check {dir} 🎉")
|
|
245
|
+
return saved_paths
|
|
246
|
+
|
|
247
|
+
if __name__ == "__main__":
|
|
248
|
+
async def main():
|
|
249
|
+
bot = AsyncAiForceimager()
|
|
250
|
+
try:
|
|
251
|
+
resp = await bot.generate("A shiny red sports car speeding down a scenic mountain road", 1)
|
|
252
|
+
paths = await bot.save(resp)
|
|
253
|
+
print(paths)
|
|
254
|
+
except Exception as e:
|
|
255
|
+
logger.error(f"An error occurred: {e} 😢")
|
|
256
|
+
|
|
257
|
+
asyncio.run(main())
|