updates to support async example

Author: dtmeadows

examples/demo.py renamed to examples/run_a_model.py

@@ -0,0 +1,20 @@
+import asyncio
+
+from replicate import AsyncReplicate
+
+client = AsyncReplicate()
+
+# https://replicate.com/stability-ai/sdxl
+model_version = "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b"
+prompts = [f"A chariot pulled by a team of {count} rainbow unicorns" for count in ["two", "four", "six", "eight"]]
+
+
+async def main() -> None:
+    # Create tasks with asyncio.gather directly
+    tasks = [client.run(model_version, input={"prompt": prompt}) for prompt in prompts]
+
+    results = await asyncio.gather(*tasks)
+    print(results)
+
+
+asyncio.run(main())

examples/run_async.py

@@ -14,6 +14,7 @@ dependencies = [
     "anyio>=3.5.0, <5",
     "distro>=1.7.0, <2",
     "sniffio",
+    "asyncio>=3.4.3",
 ]
 requires-python = ">= 3.8"
 classifiers = [

pyproject.toml

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Tuple, Union, Optional
-from typing_extensions import TypedDict
+import re
+from typing import Any, Dict, Tuple, Union, Optional, NamedTuple
 
 
 class Model:
@@ -35,12 +35,26 @@ class Version(BaseModel):
     """An OpenAPI description of the model inputs and outputs."""
 
 
-class ModelVersionIdentifier(TypedDict, total=False):
-    """A structure to identify a model version."""
+class ModelVersionIdentifier(NamedTuple):
+    """
+    A reference to a model version in the format owner/name or owner/name:version.
+    """
 
     owner: str
     name: str
-    version: str
+    version: Optional[str] = None
+
+    @classmethod
+    def parse(cls, ref: str) -> "ModelVersionIdentifier":
+        """
+        Split a reference in the format owner/name:version into its components.
+        """
+
+        match = re.match(r"^(?P<owner>[^/]+)/(?P<name>[^/:]+)(:(?P<version>.+))?$", ref)
+        if not match:
+            raise ValueError(f"Invalid reference to model version: {ref}. Expected format: owner/name:version")
+
+        return cls(match.group("owner"), match.group("name"), match.group("version"))
 
 
 def resolve_reference(
@@ -57,25 +71,13 @@ def resolve_reference(
     version_id = None
 
     if isinstance(ref, Model):
-        owner = ref.owner
-        name = ref.name
+        owner, name = ref.owner, ref.name
     elif isinstance(ref, Version):
+        version = ref
         version_id = ref.id
-    elif isinstance(ref, dict):
-        owner = ref.get("owner")
-        name = ref.get("name")
-        version_id = ref.get("version")
+    elif isinstance(ref, ModelVersionIdentifier):
+        owner, name, version_id = ref
     else:
-        # Check if the string is a version ID (assumed to be a hash-like string)
-        if "/" not in ref and len(ref) >= 32:
-            version_id = ref
-        else:
-            # Handle owner/name or owner/name/version format
-            parts = ref.split("/")
-            if len(parts) >= 2:
-                owner = parts[0]
-                name = parts[1]
-                if len(parts) >= 3:
-                    version_id = parts[2]
+        owner, name, version_id = ModelVersionIdentifier.parse(ref)
 
     return version, owner, name, version_id

src/replicate/lib/_models.py

@@ -242,7 +242,7 @@ async def async_run(
 def _has_output_iterator_array_type(version: Version) -> bool:
     schema = make_schema_backwards_compatible(version.openapi_schema, version.cog_version)
     output = schema.get("components", {}).get("schemas", {}).get("Output", {})
-    return output.get("type") == "array" and output.get("x-cog-array-type") == "iterator" # type: ignore[no-any-return]
+    return output.get("type") == "array" and output.get("x-cog-array-type") == "iterator"  # type: ignore[no-any-return]
 
 
 async def _make_async_iterator(list: List[Any]) -> AsyncIterator[Any]:

src/replicate/lib/_predictions.py

@@ -0,0 +1,27 @@
+from typing import Any, Dict, Optional
+
+from packaging import version
+
+# TODO: this code is shared with replicate's backend. Maybe we should put it in the Cog Python package as the source of truth?
+
+
+def version_has_no_array_type(cog_version: str) -> Optional[bool]:
+    """Iterators have x-cog-array-type=iterator in the schema from 0.3.9 onward"""
+    try:
+        return version.parse(cog_version) < version.parse("0.3.9")
+    except version.InvalidVersion:
+        return None
+
+
+def make_schema_backwards_compatible(
+    schema: Dict[str, Any],
+    cog_version: str,
+) -> Dict[str, Any]:
+    """A place to add backwards compatibility logic for our openapi schema"""
+
+    # If the top-level output is an array, assume it is an iterator in old versions which didn't have an array type
+    if version_has_no_array_type(cog_version):
+        output = schema["components"]["schemas"]["Output"]
+        if output.get("type") == "array":
+            output["x-cog-array-type"] = "iterator"
+    return schema

src/replicate/lib/_schema.py

@@ -274,7 +274,9 @@ def test_run_with_model_version_identifier(self, respx_mock: MockRouter) -> None
         # Case where version ID is provided
         respx_mock.post("/predictions").mock(return_value=httpx.Response(201, json=create_mock_prediction()))
 
-        identifier: ModelVersionIdentifier = {"owner": "test-owner", "name": "test-model", "version": "test-version-id"}
+        identifier = ModelVersionIdentifier(
+            owner="test-owner", name="test-model", version="test-version-id"
+        )
         output = self.client.run(identifier, input={"prompt": "test prompt"})
 
         assert output == "test output"
@@ -506,7 +508,9 @@ async def test_async_run_with_model_version_identifier(self, respx_mock: MockRou
         # Case where version ID is provided
         respx_mock.post("/predictions").mock(return_value=httpx.Response(201, json=create_mock_prediction()))
 
-        identifier: ModelVersionIdentifier = {"owner": "test-owner", "name": "test-model", "version": "test-version-id"}
+        identifier = ModelVersionIdentifier(
+            owner="test-owner", name="test-model", version="test-version-id"
+        )
         output = await self.client.run(identifier, input={"prompt": "test prompt"})
 
         assert output == "test output"