Add support for image URL (#198)

* Add support for image URL

* bump version to 1.3.1

* fix cli

Author: orangetin

pyproject.toml

@@ -12,7 +12,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry]
 name = "together"
-version = "1.3.0"
+version = "1.3.1"
 authors = [
     "Together AI <[email protected]>"
 ]

src/together/cli/api/images.py

@@ -1,5 +1,6 @@
 import base64
 import pathlib
+import requests
 
 import click
 from PIL import Image
@@ -70,8 +71,18 @@ def generate(
     for i, choice in enumerate(response.data):
         assert isinstance(choice, ImageChoicesData)
 
+        data = None
+        if choice.b64_json:
+            data = base64.b64decode(choice.b64_json)
+        elif choice.url:
+            data = requests.get(choice.url).content
+
+        if not data:
+            click.echo(f"Image [{i + 1}/{len(response.data)}] is empty")
+            continue
+
         with open(f"{output}/{prefix}{choice.index}.png", "wb") as f:
-            f.write(base64.b64decode(choice.b64_json))
+            f.write(data)
 
         click.echo(
             f"Image [{i + 1}/{len(response.data)}] saved to {output}/{prefix}{choice.index}.png"

src/together/types/images.py

@@ -28,7 +28,9 @@ class ImageChoicesData(BaseModel):
     # response index
     index: int
     # base64 image response
-    b64_json: str
+    b64_json: str | None = None
+    # URL hosting image
+    url: str | None = None
 
 
 class ImageResponse(BaseModel):