大活

art-agent/backend/venv/Lib/site-packages/replicate/helpers.py

@@ -0,0 +1,192 @@
+import base64
+import io
+import mimetypes
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from types import GeneratorType
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Optional
+
+import httpx
+
+if TYPE_CHECKING:
+    from replicate.client import Client
+    from replicate.file import FileEncodingStrategy
+
+
+try:
+    import numpy as np  # type: ignore
+
+    HAS_NUMPY = True
+except ImportError:
+    HAS_NUMPY = False
+
+
+# pylint: disable=too-many-return-statements
+def encode_json(
+    obj: Any,  # noqa: ANN401
+    client: "Client",
+    file_encoding_strategy: Optional["FileEncodingStrategy"] = None,
+) -> Any:  # noqa: ANN401
+    """
+    Return a JSON-compatible version of the object.
+    """
+
+    if isinstance(obj, dict):
+        return {
+            key: encode_json(value, client, file_encoding_strategy)
+            for key, value in obj.items()
+        }
+    if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
+        return [encode_json(value, client, file_encoding_strategy) for value in obj]
+    if isinstance(obj, Path):
+        with obj.open("rb") as file:
+            return encode_json(file, client, file_encoding_strategy)
+    if isinstance(obj, io.IOBase):
+        if file_encoding_strategy == "base64":
+            return base64_encode_file(obj)
+        return client.files.create(obj).urls["get"]
+    if HAS_NUMPY:
+        if isinstance(obj, np.integer):  # type: ignore
+            return int(obj)
+        if isinstance(obj, np.floating):  # type: ignore
+            return float(obj)
+        if isinstance(obj, np.ndarray):  # type: ignore
+            return obj.tolist()
+    return obj
+
+
+async def async_encode_json(
+    obj: Any,  # noqa: ANN401
+    client: "Client",
+    file_encoding_strategy: Optional["FileEncodingStrategy"] = None,
+) -> Any:  # noqa: ANN401
+    """
+    Asynchronously return a JSON-compatible version of the object.
+    """
+
+    if isinstance(obj, dict):
+        return {
+            key: (await async_encode_json(value, client, file_encoding_strategy))
+            for key, value in obj.items()
+        }
+    if isinstance(obj, (list, set, frozenset, GeneratorType, tuple)):
+        return [
+            (await async_encode_json(value, client, file_encoding_strategy))
+            for value in obj
+        ]
+    if isinstance(obj, Path):
+        with obj.open("rb") as file:
+            return await async_encode_json(file, client, file_encoding_strategy)
+    if isinstance(obj, io.IOBase):
+        if file_encoding_strategy == "base64":
+            # TODO: This should ideally use an async based file reader path.
+            return base64_encode_file(obj)
+        return (await client.files.async_create(obj)).urls["get"]
+    if HAS_NUMPY:
+        if isinstance(obj, np.integer):  # type: ignore
+            return int(obj)
+        if isinstance(obj, np.floating):  # type: ignore
+            return float(obj)
+        if isinstance(obj, np.ndarray):  # type: ignore
+            return obj.tolist()
+    return obj
+
+
+def base64_encode_file(file: io.IOBase) -> str:
+    """
+    Base64 encode a file.
+
+    Args:
+        file: A file handle to upload.
+    Returns:
+        str: A base64-encoded data URI.
+    """
+
+    file.seek(0)
+    body = file.read()
+
+    # Ensure the file handle is in bytes
+    body = body.encode("utf-8") if isinstance(body, str) else body
+    encoded_body = base64.b64encode(body).decode("utf-8")
+
+    mime_type = (
+        mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream"
+    )
+    return f"data:{mime_type};base64,{encoded_body}"
+
+
+class FileOutput(httpx.SyncByteStream, httpx.AsyncByteStream):
+    """
+    An object that can be used to read the contents of an output file
+    created by running a Replicate model.
+    """
+
+    url: str
+    """
+    The file URL.
+    """
+
+    _client: "Client"
+
+    def __init__(self, url: str, client: "Client") -> None:
+        self.url = url
+        self._client = client
+
+    def read(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(",", 1)
+            return base64.b64decode(encoded)
+
+        with self._client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return response.read()
+
+    def __iter__(self) -> Iterator[bytes]:
+        if self.url.startswith("data:"):
+            yield self.read()
+            return
+
+        with self._client._client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            yield from response.iter_bytes()
+
+    async def aread(self) -> bytes:
+        if self.url.startswith("data:"):
+            _, encoded = self.url.split(",", 1)
+            return base64.b64decode(encoded)
+
+        async with self._client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            return await response.aread()
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        if self.url.startswith("data:"):
+            yield await self.aread()
+            return
+
+        async with self._client._async_client.stream("GET", self.url) as response:
+            response.raise_for_status()
+            async for chunk in response.aiter_bytes():
+                yield chunk
+
+    def __str__(self) -> str:
+        return self.url
+
+
+def transform_output(value: Any, client: "Client") -> Any:
+    """
+    Transform the output of a prediction to a `FileOutput` object if it's a URL.
+    """
+
+    def transform(obj: Any) -> Any:
+        if isinstance(obj, Mapping):
+            return {k: transform(v) for k, v in obj.items()}
+        if isinstance(obj, Sequence) and not isinstance(obj, str):
+            return [transform(item) for item in obj]
+        if isinstance(obj, str) and (
+            obj.startswith("https:") or obj.startswith("data:")
+        ):
+            return FileOutput(obj, client)
+        return obj
+
+    return transform(value)