webscout 8.2.9__py3-none-any.whl → 8.3.1__py3-none-any.whl

webscout/Provider/TTI/magicstudio.py

@@ -4,7 +4,10 @@ import uuid
 import time
 import tempfile
 from typing import Optional, List
-from webscout.Provider.TTI.utils import ImageData, ImageResponse
+from webscout.Provider.TTI.utils import (
+    ImageData,
+    ImageResponse
+)
 from webscout.Provider.TTI.base import TTICompatibleProvider, BaseImages
 from io import BytesIO
 from webscout.litagent import LitAgent
@@ -14,11 +17,13 @@ try:
 except ImportError:
     Image = None
 
+
 class Images(BaseImages):
     def __init__(self, client):
         self._client = client
 
-    def create(self,
+    def create(
+        self,
         model: str = "magicstudio",
         prompt: str = None,
         n: int = 1,
@@ -29,7 +34,7 @@ class Images(BaseImages):
         aspect_ratio: str = None,
         timeout: int = 60,
         image_format: str = "jpg",
-        **kwargs
+        **kwargs,
     ) -> ImageResponse:
         if not prompt:
             raise ValueError("Prompt is required!")
@@ -43,7 +48,7 @@ class Images(BaseImages):
             "Origin": "https://magicstudio.com",
             "Referer": "https://magicstudio.com/ai-art-generator/",
             "DNT": "1",
-            "Sec-GPC": "1"
+            "Sec-GPC": "1",
         }
         session = requests.Session()
         session.headers.update(headers)
@@ -57,7 +62,11 @@ class Images(BaseImages):
                 "user_is_subscribed": "false",
                 "client_id": uuid.uuid4().hex,
             }
-            resp = session.post(api_url, data=form_data, timeout=timeout)
+            resp = session.post(
+                api_url,
+                data=form_data,
+                timeout=timeout,
+            )
             resp.raise_for_status()
             img_bytes = resp.content
             # Convert to png or jpeg in memory if needed
@@ -65,7 +74,10 @@ class Images(BaseImages):
                 with BytesIO(img_bytes) as input_io:
                     with Image.open(input_io) as im:
                         out_io = BytesIO()
-                        if image_format.lower() == "jpeg" or image_format.lower() == "jpg":
+                        if (
+                            image_format.lower() == "jpeg"
+                            or image_format.lower() == "jpg"
+                        ):
                             im = im.convert("RGB")
                             im.save(out_io, format="JPEG")
                         else:
@@ -73,37 +85,43 @@ class Images(BaseImages):
                         img_bytes = out_io.getvalue()
             images.append(img_bytes)
             if response_format == "url":
+
                 def upload_file_with_retry(img_bytes, image_format, max_retries=3):
                     ext = "jpg" if image_format.lower() in ("jpeg", "jpg") else "png"
                     for attempt in range(max_retries):
                         tmp_path = None
                         try:
-                            with tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False) as tmp:
+                            with tempfile.NamedTemporaryFile(
+                                suffix=f".{ext}", delete=False
+                            ) as tmp:
                                 tmp.write(img_bytes)
                                 tmp.flush()
                                 tmp_path = tmp.name
-                            with open(tmp_path, 'rb') as f:
+                            with open(tmp_path, "rb") as f:
                                 files = {
-                                    'fileToUpload': (f'image.{ext}', f, f'image/{ext}')
+                                    "fileToUpload": (f"image.{ext}", f, f"image/{ext}")
                                 }
-                                data = {
-                                    'reqtype': 'fileupload',
-                                    'json': 'true'
-                                }
-                                headers = {'User-Agent': agent.random()}
+                                data = {"reqtype": "fileupload", "json": "true"}
+                                headers = {"User-Agent": agent.random()}
                                 if attempt > 0:
-                                    headers['Connection'] = 'close'
-                                resp = requests.post("https://catbox.moe/user/api.php", files=files, data=data, headers=headers, timeout=timeout)
+                                    headers["Connection"] = "close"
+                                resp = requests.post(
+                                    "https://catbox.moe/user/api.php",
+                                    files=files,
+                                    data=data,
+                                    headers=headers,
+                                    timeout=timeout,
+                                )
                                 if resp.status_code == 200 and resp.text.strip():
                                     text = resp.text.strip()
-                                    if text.startswith('http'):
+                                    if text.startswith("http"):
                                         return text
                                     try:
                                         result = resp.json()
                                         if "url" in result:
                                             return result["url"]
                                     except Exception:
-                                        if 'http' in text:
+                                        if "http" in text:
                                             return text
                         except Exception:
                             if attempt < max_retries - 1:
@@ -115,22 +133,27 @@ class Images(BaseImages):
                                 except Exception:
                                     pass
                     return None
+
                 def upload_file_alternative(img_bytes, image_format):
                     try:
-                        ext = "jpg" if image_format.lower() in ("jpeg", "jpg") else "png"
-                        with tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False) as tmp:
+                        ext = (
+                            "jpg" if image_format.lower() in ("jpeg", "jpg") else "png"
+                        )
+                        with tempfile.NamedTemporaryFile(
+                            suffix=f".{ext}", delete=False
+                        ) as tmp:
                             tmp.write(img_bytes)
                             tmp.flush()
                             tmp_path = tmp.name
                         try:
                             if not os.path.isfile(tmp_path):
                                 return None
-                            with open(tmp_path, 'rb') as img_file:
-                                files = {'file': img_file}
-                                alt_resp = requests.post('https://0x0.st', files=files)
+                            with open(tmp_path, "rb") as img_file:
+                                files = {"file": img_file}
+                                alt_resp = requests.post("https://0x0.st", files=files)
                                 alt_resp.raise_for_status()
                                 image_url = alt_resp.text.strip()
-                                if not image_url.startswith('http'):
+                                if not image_url.startswith("http"):
                                     return None
                                 return image_url
                         except Exception:
@@ -142,32 +165,36 @@ class Images(BaseImages):
                                 pass
                     except Exception:
                         return None
+
                 uploaded_url = upload_file_with_retry(img_bytes, image_format)
                 if not uploaded_url:
                     uploaded_url = upload_file_alternative(img_bytes, image_format)
                 if uploaded_url:
                     urls.append(uploaded_url)
                 else:
-                    raise RuntimeError("Failed to upload image to catbox.moe using all available methods")
+                    raise RuntimeError(
+                        "Failed to upload image to catbox.moe using all available methods"
+                    )
         result_data = []
         if response_format == "url":
             for url in urls:
                 result_data.append(ImageData(url=url))
         elif response_format == "b64_json":
             import base64
+
             for img in images:
                 b64 = base64.b64encode(img).decode("utf-8")
                 result_data.append(ImageData(b64_json=b64))
         else:
             raise ValueError("response_format must be 'url' or 'b64_json'")
         from time import time as _time
-        return ImageResponse(
-            created=int(_time()),
-            data=result_data
-        )
+
+        return ImageResponse(created=int(_time()), data=result_data)
+
 
 class MagicStudioAI(TTICompatibleProvider):
     AVAILABLE_MODELS = ["magicstudio"]
+
     def __init__(self):
         self.api_endpoint = "https://ai-api.magicstudio.com/api/ai-art-generator"
         self.session = requests.Session()
@@ -182,15 +209,19 @@ class MagicStudioAI(TTICompatibleProvider):
         }
         self.session.headers.update(self.headers)
         self.images = Images(self)
+
     @property
     def models(self):
         class _ModelList:
             def list(inner_self):
                 return type(self).AVAILABLE_MODELS
+
         return _ModelList()
 
+
 if __name__ == "__main__":
     from rich import print
+
     client = MagicStudioAI()
     response = client.images.create(
         prompt="A cool cyberpunk city at night",

webscout/Provider/TTI/piclumen.py

@@ -1,6 +1,9 @@
 import requests
 from typing import Optional, List, Dict, Any
-from webscout.Provider.TTI.utils import ImageData, ImageResponse
+from webscout.Provider.TTI.utils import (
+    ImageData,
+    ImageResponse
+)
 from webscout.Provider.TTI.base import TTICompatibleProvider, BaseImages
 from io import BytesIO
 import os
@@ -14,11 +17,13 @@ try:
 except ImportError:
     Image = None
 
+
 class Images(BaseImages):
     def __init__(self, client):
         self._client = client
 
-    def create(self,
+    def create(
+        self,
         model: str,
         prompt: str,
         n: int = 1,
@@ -29,7 +34,7 @@ class Images(BaseImages):
         aspect_ratio: str = "1:1",
         timeout: int = 60,
         image_format: str = "jpeg",
-        **kwargs
+        **kwargs,
     ) -> ImageResponse:
         """
         image_format: "png" or "jpeg"
@@ -46,32 +51,35 @@ class Images(BaseImages):
             for attempt in range(max_retries):
                 tmp_path = None
                 try:
-                    with tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False) as tmp:
+                    with tempfile.NamedTemporaryFile(
+                        suffix=f".{ext}", delete=False
+                    ) as tmp:
                         tmp.write(img_bytes)
                         tmp.flush()
                         tmp_path = tmp.name
-                    with open(tmp_path, 'rb') as f:
-                        files = {
-                            'fileToUpload': (f'image.{ext}', f, f'image/{ext}')
-                        }
-                        data = {
-                            'reqtype': 'fileupload',
-                            'json': 'true'
-                        }
-                        headers = {'User-Agent': agent.random()}
+                    with open(tmp_path, "rb") as f:
+                        files = {"fileToUpload": (f"image.{ext}", f, f"image/{ext}")}
+                        data = {"reqtype": "fileupload", "json": "true"}
+                        headers = {"User-Agent": agent.random()}
                         if attempt > 0:
-                            headers['Connection'] = 'close'
-                        resp = requests.post("https://catbox.moe/user/api.php", files=files, data=data, headers=headers, timeout=timeout)
+                            headers["Connection"] = "close"
+                        resp = requests.post(
+                            "https://catbox.moe/user/api.php",
+                            files=files,
+                            data=data,
+                            headers=headers,
+                            timeout=timeout,
+                        )
                         if resp.status_code == 200 and resp.text.strip():
                             text = resp.text.strip()
-                            if text.startswith('http'):
+                            if text.startswith("http"):
                                 return text
                             try:
                                 result = resp.json()
                                 if "url" in result:
                                     return result["url"]
                             except json.JSONDecodeError:
-                                if 'http' in text:
+                                if "http" in text:
                                     return text
                 except Exception:
                     if attempt < max_retries - 1:
@@ -94,12 +102,12 @@ class Images(BaseImages):
                 try:
                     if not os.path.isfile(tmp_path):
                         return None
-                    with open(tmp_path, 'rb') as img_file:
-                        files = {'file': img_file}
-                        response = requests.post('https://0x0.st', files=files)
+                    with open(tmp_path, "rb") as img_file:
+                        files = {"file": img_file}
+                        response = requests.post("https://0x0.st", files=files)
                         response.raise_for_status()
                         image_url = response.text.strip()
-                        if not image_url.startswith('http'):
+                        if not image_url.startswith("http"):
                             return None
                         return image_url
                 except Exception:
@@ -117,11 +125,11 @@ class Images(BaseImages):
             resp = self._client.session.post(
                 self._client.api_endpoint,
                 json=payload,
-                timeout=timeout
+                timeout=timeout,
             )
             resp.raise_for_status()
             # Piclumen returns image/jpeg directly
-            if resp.headers.get('content-type') == 'image/jpeg':
+            if resp.headers.get("content-type") == "image/jpeg":
                 img_bytes = resp.content
                 # Convert to png or jpeg in memory
                 with BytesIO(img_bytes) as input_io:
@@ -141,7 +149,9 @@ class Images(BaseImages):
                     if uploaded_url:
                         urls.append(uploaded_url)
                     else:
-                        raise RuntimeError("Failed to upload image to catbox.moe using all available methods")
+                        raise RuntimeError(
+                            "Failed to upload image to catbox.moe using all available methods"
+                        )
             else:
                 raise RuntimeError("No image data received from Piclumen")
 
@@ -151,6 +161,7 @@ class Images(BaseImages):
                 result_data.append(ImageData(url=url))
         elif response_format == "b64_json":
             import base64
+
             for img in images:
                 b64 = base64.b64encode(img).decode("utf-8")
                 result_data.append(ImageData(b64_json=b64))
@@ -158,15 +169,12 @@ class Images(BaseImages):
             raise ValueError("response_format must be 'url' or 'b64_json'")
 
         from time import time as _time
-        return ImageResponse(
-            created=int(_time()),
-            data=result_data
-        )
+
+        return ImageResponse(created=int(_time()), data=result_data)
+
 
 class PiclumenAI(TTICompatibleProvider):
-    AVAILABLE_MODELS = [
-        "piclumen-v1"
-    ]
+    AVAILABLE_MODELS = ["piclumen-v1"]
 
     def __init__(self):
         self.api_endpoint = "https://s9.piclumen.art/comfy/api/generate-image"
@@ -188,10 +196,13 @@ class PiclumenAI(TTICompatibleProvider):
         class _ModelList:
             def list(inner_self):
                 return type(self).AVAILABLE_MODELS
+
         return _ModelList()
 
+
 if __name__ == "__main__":
     from rich import print
+
     client = PiclumenAI()
     response = client.images.create(
         model="piclumen-v1",
@@ -200,4 +211,4 @@ if __name__ == "__main__":
         n=2,
         timeout=30,
     )
-    print(response)
+    print(response)