Add support for returning an iterator from use()

Author: aron

replicate/use.py

@@ -1,20 +1,20 @@
 # TODO
-# - [ ] Support downloading files and conversion into Path when schema is URL
-# - [ ] Support asyncio variant
-# - [ ] Support list outputs
+# - [x] Support downloading files and conversion into Path when schema is URL
+# - [x] Support list outputs
 # - [ ] Support iterator outputs
-# - [ ] Support text streaming
-# - [ ] Support file streaming
+# - [ ] Support helpers for working with ContatenateIterator
 # - [ ] Support reusing output URL when passing to new method
 # - [ ] Support lazy downloading of files into Path
-# - [ ] Support helpers for working with ContatenateIterator
+# - [ ] Support text streaming
+# - [ ] Support file streaming
+# - [ ] Support asyncio variant
 import inspect
 import os
 import tempfile
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Iterator, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 import httpx
@@ -66,6 +66,16 @@ def _has_concatenate_iterator_output_type(openapi_schema: dict) -> bool:
     return True
 
 
+def _has_iterator_output_type(openapi_schema: dict) -> bool:
+    """
+    Returns true if the model output type is an iterator (non-concatenate).
+    """
+    output = openapi_schema.get("components", {}).get("schemas", {}).get("Output", {})
+    return (
+        output.get("type") == "array" and output.get("x-cog-array-type") == "iterator"
+    )
+
+
 def _download_file(url: str) -> Path:
     """
     Download a file from URL to a temporary location and return the Path.
@@ -86,6 +96,23 @@ def _download_file(url: str) -> Path:
     return Path(temp_file.name)
 
 
+def _process_iterator_item(item: Any, openapi_schema: dict) -> Any:
+    """
+    Process a single item from an iterator output based on schema.
+    """
+    output_schema = openapi_schema.get("components", {}).get("schemas", {}).get("Output", {})
+
+    # For array/iterator types, check the items schema
+    if output_schema.get("type") == "array" and output_schema.get("x-cog-array-type") == "iterator":
+        items_schema = output_schema.get("items", {})
+        # If items are file URLs, download them
+        if items_schema.get("type") == "string" and items_schema.get("format") == "uri":
+            if isinstance(item, str) and item.startswith(("http://", "https://")):
+                return _download_file(item)
+
+    return item
+
+
 def _process_output_with_schema(output: Any, openapi_schema: dict) -> Any:
     """
     Process output data, downloading files based on OpenAPI schema.
@@ -159,7 +186,7 @@ class Run:
     prediction: Prediction
     schema: dict
 
-    def wait(self) -> Any:
+    def wait(self) -> Union[Any, Iterator[Any]]:
         """
         Wait for the prediction to complete and return its output.
         """
@@ -171,6 +198,13 @@ def wait(self) -> Any:
         if _has_concatenate_iterator_output_type(self.schema):
             return "".join(self.prediction.output)
 
+        # Return an iterator for iterator output types
+        if _has_iterator_output_type(self.schema) and self.prediction.output is not None:
+            return (
+                _process_iterator_item(chunk, self.schema)
+                for chunk in self.prediction.output
+            )
+
         # Process output for file downloads based on schema
         return _process_output_with_schema(self.prediction.output, self.schema)
 
@@ -286,8 +320,6 @@ def use(function_ref: str) -> Function:
 
     """
     if not _in_module_scope():
-        raise RuntimeError(
-            "You may only call cog.ext.pipelines.include at the top level."
-        )
+        raise RuntimeError("You may only call replicate.use() at the top level.")
 
     return Function(function_ref)

tests/test_use.py

@@ -1,4 +1,6 @@
 import os
+import types
+from pathlib import Path
 
 import httpx
 import pytest
@@ -417,8 +419,11 @@ async def test_use_iterator_of_strings_output(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as a list (iterators are returned as lists)
-    assert output == ["hello", "world", "test"]
+    # Assert that output is returned as an iterator
+    assert isinstance(output, types.GeneratorType)
+    # Convert to list to check contents
+    output_list = list(output)
+    assert output_list == ["hello", "world", "test"]
 
 
 @pytest.mark.asyncio
@@ -449,9 +454,6 @@ async def test_use_path_output(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as a Path object
-    from pathlib import Path
-
     assert isinstance(output, Path)
     assert output.exists()
     assert output.read_bytes() == b"fake image data"
@@ -497,9 +499,6 @@ async def test_use_list_of_paths_output(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as a list of Path objects
-    from pathlib import Path
-
     assert isinstance(output, list)
     assert len(output) == 2
     assert all(isinstance(path, Path) for path in output)
@@ -549,15 +548,15 @@ async def test_use_iterator_of_paths_output(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as a list of Path objects
-    from pathlib import Path
-
-    assert isinstance(output, list)
-    assert len(output) == 2
-    assert all(isinstance(path, Path) for path in output)
-    assert all(path.exists() for path in output)
-    assert output[0].read_bytes() == b"fake image 1 data"
-    assert output[1].read_bytes() == b"fake image 2 data"
+    # Assert that output is returned as an iterator of Path objects
+    assert isinstance(output, types.GeneratorType)
+    # Convert to list to check contents
+    output_list = list(output)
+    assert len(output_list) == 2
+    assert all(isinstance(path, Path) for path in output_list)
+    assert all(path.exists() for path in output_list)
+    assert output_list[0].read_bytes() == b"fake image 1 data"
+    assert output_list[1].read_bytes() == b"fake image 2 data"
 
 
 @pytest.mark.asyncio
@@ -681,9 +680,6 @@ async def test_use_object_output_with_file_properties(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as an object with file downloaded
-    from pathlib import Path
-
     assert isinstance(output, dict)
     assert output["text"] == "Generated text"
     assert output["count"] == 42
@@ -742,9 +738,6 @@ async def test_use_object_output_with_file_list_property(use_async_client):
     # Call function with prompt="hello world"
     output = hotdog_detector(prompt="hello world")
 
-    # Assert that output is returned as an object with files downloaded
-    from pathlib import Path
-
     assert isinstance(output, dict)
     assert output["text"] == "Generated text"
     assert isinstance(output["images"], list)