pixeltable 0.1.0__py3-none-any.whl → 0.2.4__py3-none-any.whl

pixeltable/functions/huggingface.py

@@ -0,0 +1,120 @@
+from typing import Any, Callable
+
+import PIL.Image
+import numpy as np
+
+import pixeltable as pxt
+import pixeltable.env as env
+import pixeltable.type_system as ts
+from pixeltable.func import Batch
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+def sentence_transformer(sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False) -> Batch[np.ndarray]:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import SentenceTransformer
+
+    model = _lookup_model(model_id, SentenceTransformer)
+
+    array = model.encode(sentences, normalize_embeddings=normalize_embeddings)
+    return [array[i] for i in range(array.shape[0])]
+
+
+@pxt.udf
+def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embeddings: bool = False) -> list:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import SentenceTransformer
+
+    model = _lookup_model(model_id, SentenceTransformer)
+
+    array = model.encode(sentences, normalize_embeddings=normalize_embeddings)
+    return [array[i].tolist() for i in range(array.shape[0])]
+
+
+@pxt.udf(batch_size=32)
+def cross_encoder(sentences1: Batch[str], sentences2: Batch[str], *, model_id: str) -> Batch[float]:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import CrossEncoder
+
+    model = _lookup_model(model_id, CrossEncoder)
+
+    array = model.predict([[s1, s2] for s1, s2 in zip(sentences1, sentences2)], convert_to_numpy=True)
+    return array.tolist()
+
+
+@pxt.udf
+def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> list:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import CrossEncoder
+
+    model = _lookup_model(model_id, CrossEncoder)
+
+    array = model.predict([[sentence1, s2] for s2 in sentences2], convert_to_numpy=True)
+    return array.tolist()
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
+    env.Env.get().require_package('transformers')
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
+    embeddings = model.get_text_features(**inputs).detach().numpy()
+    return [embeddings[i] for i in range(embeddings.shape[0])]
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndarray]:
+    env.Env.get().require_package('transformers')
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    inputs = processor(images=image, return_tensors='pt', padding=True)
+    embeddings = model.get_image_features(**inputs).detach().numpy()
+    return [embeddings[i] for i in range(embeddings.shape[0])]
+
+
+@pxt.udf(batch_size=32)
+def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
+    env.Env.get().require_package('transformers')
+    from transformers import DetrImageProcessor, DetrForObjectDetection
+
+    model = _lookup_model(model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'))
+    processor = _lookup_processor(model_id, lambda x: DetrImageProcessor.from_pretrained(x, revision='no_timm'))
+
+    inputs = processor(images=image, return_tensors='pt')
+    outputs = model(**inputs)
+
+    results = processor.post_process_object_detection(outputs, threshold=threshold)
+    return [
+        {
+            'scores': [score.item() for score in result['scores']],
+            'labels': [label.item() for label in result['labels']],
+            'label_text': [model.config.id2label[label.item()] for label in result['labels']],
+            'boxes': [box.tolist() for box in result['boxes']]
+        }
+        for result in results
+    ]
+
+
+def _lookup_model(model_id: str, create: Callable) -> Any:
+    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+    if key not in _model_cache:
+        _model_cache[key] = create(model_id)
+    return _model_cache[key]
+
+
+def _lookup_processor(model_id: str, create: Callable) -> Any:
+    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+    if key not in _processor_cache:
+        _processor_cache[key] = create(model_id)
+    return _processor_cache[key]
+
+
+_model_cache = {}
+_processor_cache = {}

pixeltable/functions/image.py

@@ -0,0 +1,16 @@
+import base64
+
+import PIL.Image
+
+from pixeltable.type_system import ImageType, StringType
+import pixeltable.func as func
+
+
+@func.udf
+def b64_encode(img: PIL.Image.Image, image_format: str = 'png') -> str:
+    # Encode this image as a b64-encoded png.
+    import io
+    bytes_arr = io.BytesIO()
+    img.save(bytes_arr, format=image_format)
+    b64_bytes = base64.b64encode(bytes_arr.getvalue())
+    return b64_bytes.decode('utf-8')

pixeltable/functions/openai.py

@@ -0,0 +1,256 @@
+import base64
+import io
+import pathlib
+import uuid
+from typing import Optional, TypeVar, Union, Callable
+
+import PIL.Image
+import numpy as np
+import openai
+import tenacity
+from openai._types import NOT_GIVEN, NotGiven
+
+import pixeltable as pxt
+import pixeltable.type_system as ts
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+def openai_client() -> openai.OpenAI:
+    return env.Env.get().get_client('openai', lambda api_key: openai.OpenAI(api_key=api_key))
+
+
+# Exponential backoff decorator using tenacity.
+# TODO(aaron-siegel): Right now this hardwires random exponential backoff with defaults suggested
+# by OpenAI. Should we investigate making this more customizable in the future?
+def _retry(fn: Callable) -> Callable:
+    return tenacity.retry(
+        retry=tenacity.retry_if_exception_type(openai.RateLimitError),
+        wait=tenacity.wait_random_exponential(min=1, max=60),
+        stop=tenacity.stop_after_attempt(6)
+    )(fn)
+
+
+#####################################
+# Audio Endpoints
+
+@pxt.udf(return_type=ts.AudioType())
+@_retry
+def speech(
+        input: str,
+        *,
+        model: str,
+        voice: str,
+        response_format: Optional[str] = None,
+        speed: Optional[float] = None
+) -> str:
+    content = openai_client().audio.speech.create(
+        input=input,
+        model=model,
+        voice=voice,
+        response_format=_opt(response_format),
+        speed=_opt(speed)
+    )
+    ext = response_format or 'mp3'
+    output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+    content.stream_to_file(output_filename, chunk_size=1 << 20)
+    return output_filename
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True),
+                 ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def transcriptions(
+        audio: str,
+        *,
+        model: str,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    transcription = openai_client().audio.transcriptions.create(
+        file=file,
+        model=model,
+        language=_opt(language),
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return transcription.dict()
+
+
+@pxt.udf(
+    param_types=[ts.AudioType(), ts.StringType(), ts.StringType(nullable=True), ts.FloatType(nullable=True)]
+)
+@_retry
+def translations(
+        audio: str,
+        *,
+        model: str,
+        prompt: Optional[str] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    file = pathlib.Path(audio)
+    translation = openai_client().audio.translations.create(
+        file=file,
+        model=model,
+        prompt=_opt(prompt),
+        temperature=_opt(temperature)
+    )
+    return translation.dict()
+
+
+#####################################
+# Chat Endpoints
+
+@pxt.udf
+@_retry
+def chat_completions(
+        messages: list,
+        *,
+        model: str,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[dict[str, int]] = None,
+        logprobs: Optional[bool] = None,
+        top_logprobs: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[dict] = None,
+        seed: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        tools: Optional[list[dict]] = None,
+        tool_choice: Optional[dict] = None,
+        user: Optional[str] = None
+) -> dict:
+    result = openai_client().chat.completions.create(
+        messages=messages,
+        model=model,
+        frequency_penalty=_opt(frequency_penalty),
+        logit_bias=_opt(logit_bias),
+        logprobs=_opt(logprobs),
+        top_logprobs=_opt(top_logprobs),
+        max_tokens=_opt(max_tokens),
+        n=_opt(n),
+        presence_penalty=_opt(presence_penalty),
+        response_format=_opt(response_format),
+        seed=_opt(seed),
+        stop=_opt(stop),
+        temperature=_opt(temperature),
+        top_p=_opt(top_p),
+        tools=_opt(tools),
+        tool_choice=_opt(tool_choice),
+        user=_opt(user)
+    )
+    return result.dict()
+
+
+@pxt.udf
+@_retry
+def vision(
+        prompt: str,
+        image: PIL.Image.Image,
+        *,
+        model: str = 'gpt-4-vision-preview'
+) -> str:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    bytes_arr = io.BytesIO()
+    image.save(bytes_arr, format='png')
+    b64_bytes = base64.b64encode(bytes_arr.getvalue())
+    b64_encoded_image = b64_bytes.decode('utf-8')
+    messages = [
+        {'role': 'user',
+         'content': [
+             {'type': 'text', 'text': prompt},
+             {'type': 'image_url', 'image_url': {
+                 'url': f'data:image/png;base64,{b64_encoded_image}'
+             }}
+         ]}
+    ]
+    result = openai_client().chat.completions.create(
+        messages=messages,
+        model=model
+    )
+    return result.choices[0].message.content
+
+
+#####################################
+# Embeddings Endpoints
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+@_retry
+def embeddings(
+        input: Batch[str],
+        *,
+        model: str,
+        user: Optional[str] = None
+) -> Batch[np.ndarray]:
+    result = openai_client().embeddings.create(
+        input=input,
+        model=model,
+        user=_opt(user),
+        encoding_format='float'
+    )
+    return [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+
+
+#####################################
+# Images Endpoints
+
+@pxt.udf
+@_retry
+def image_generations(
+        prompt: str,
+        *,
+        model: Optional[str] = None,
+        quality: Optional[str] = None,
+        size: Optional[str] = None,
+        style: Optional[str] = None,
+        user: Optional[str] = None
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = openai_client().images.generate(
+        prompt=prompt,
+        model=_opt(model),
+        quality=_opt(quality),
+        size=_opt(size),
+        style=_opt(style),
+        user=_opt(user),
+        response_format="b64_json"
+    )
+    b64_str = result.data[0].b64_json
+    b64_bytes = base64.b64decode(b64_str)
+    img = PIL.Image.open(io.BytesIO(b64_bytes))
+    img.load()
+    return img
+
+
+#####################################
+# Moderations Endpoints
+
+@pxt.udf
+@_retry
+def moderations(
+        input: str,
+        *,
+        model: Optional[str] = None
+) -> dict:
+    result = openai_client().moderations.create(
+        input=input,
+        model=_opt(model)
+    )
+    return result.dict()
+
+
+_T = TypeVar('_T')
+
+
+def _opt(arg: _T) -> Union[_T, NotGiven]:
+    return arg if arg is not None else NOT_GIVEN

pixeltable/functions/pil/image.py

@@ -1,9 +1,150 @@
-from pixeltable.type_system import FloatType, ImageType
-from pixeltable.function import Function
+from typing import Dict, Any, Tuple, Optional
 
+import PIL.Image
 
-alpha_composite = Function(
-    ImageType(), [ImageType(), ImageType()], module_name='PIL.Image', eval_symbol='alpha_composite')
-blend = Function(ImageType(), [ImageType(), ImageType(), FloatType()], module_name='PIL.Image', eval_symbol='blend')
-composite = Function(
-    ImageType(), [ImageType(), ImageType(), ImageType()], module_name='PIL.Image', eval_symbol='composite')
+from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType, BoolType
+import pixeltable.func as func
+
+
+def _caller_return_type(bound_args: Optional[Dict[str, Any]]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    return bound_args['self'].col_type
+
+@func.udf(
+    py_fn=PIL.Image.alpha_composite, return_type=ImageType(), param_types=[ImageType(), ImageType()])
+def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.blend, return_type=ImageType(), param_types=[ImageType(), ImageType(), FloatType()])
+def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.composite, return_type=ImageType(), param_types=[ImageType(), ImageType(), ImageType()])
+def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+
+
+# PIL.Image.Image methods
+
+# Image.convert()
+def _convert_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    assert 'self' in bound_args
+    assert 'mode' in bound_args
+    img_type = bound_args['self'].col_type
+    return ImageType(size=img_type.size, mode=bound_args['mode'])
+@func.udf(return_type=_convert_return_type, param_types=[ImageType(), StringType()])
+def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
+    return self.convert(mode)
+
+# Image.crop()
+def _crop_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    img_type = bound_args['self'].col_type
+    box = bound_args['box']
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=img_type.mode)
+    return ImageType()  # we can't compute the size statically
+@func.udf(
+    py_fn=PIL.Image.Image.crop, return_type=_crop_return_type,
+    param_types=[ImageType(), ArrayType((4,), dtype=IntType())])
+def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
+    pass
+
+# Image.getchannel()
+def _getchannel_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    img_type = bound_args['self'].col_type
+    return ImageType(size=img_type.size, mode='L')
+@func.udf(
+    py_fn=PIL.Image.Image.getchannel, return_type=_getchannel_return_type, param_types=[ImageType(), IntType()])
+def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
+    pass
+
+# Image.resize()
+def resize_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    assert 'size' in bound_args
+    return ImageType(size=bound_args['size'])
+@func.udf(return_type=resize_return_type, param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
+def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
+    return self.resize(size)
+
+# Image.rotate()
+@func.udf(return_type=ImageType(), param_types=[ImageType(), IntType()])
+def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
+    return self.rotate(angle)
+
+# Image.transform()
+@func.udf(return_type= _caller_return_type, param_types=[ImageType(), ArrayType((2,), dtype=IntType()), IntType()])
+def transform(self: PIL.Image.Image, size: Tuple[int, int], method: int) -> PIL.Image.Image:
+    return self.transform(size, method)
+
+@func.udf(
+    py_fn=PIL.Image.Image.effect_spread, return_type=_caller_return_type, param_types=[ImageType(), FloatType()])
+def effect_spread(self: PIL.Image.Image, distance: float) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.entropy, return_type=FloatType(), param_types=[ImageType(), ImageType(), JsonType()])
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> float:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbands, return_type=JsonType(), param_types=[ImageType()])
+def getbands(self: PIL.Image.Image) -> Tuple[str]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbbox, return_type=JsonType(), param_types=[ImageType()])
+def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
+def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getextrema, return_type=JsonType(), param_types=[ImageType()])
+def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getpalette, return_type=JsonType(), param_types=[ImageType(), StringType()])
+def getpalette(self: PIL.Image.Image, mode: str) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getpixel, return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
+def getpixel(self: PIL.Image.Image, xy: Tuple[int, int]) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
+def getprojection(self: PIL.Image.Image) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.quantize, return_type=ImageType(),
+    param_types=[ImageType(), IntType(), IntType(nullable=True), IntType(), IntType(nullable=True), IntType()])
+def quantize(
+        self: PIL.Image.Image, colors: int, method: int, kmeans: int, palette: int, dither: int) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.reduce, return_type=ImageType(), param_types=[ImageType(), IntType(), JsonType()])
+def reduce(self: PIL.Image.Image, factor: int, filter: Tuple[int]) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.transpose, return_type=_caller_return_type, param_types=[ImageType(), IntType()])
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+    pass

pixeltable/functions/string.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from pixeltable.type_system import StringType
+import pixeltable.func as func
+
+
+@func.udf(return_type=StringType(), param_types=[StringType()])
+def str_format(format_str: str, *args: Any, **kwargs: Any) -> str:
+    """ Return a formatted version of format_str, using substitutions from args and kwargs:
+    - {<int>} will be replaced by the corresponding element in args
+    - {<key>} will be replaced by the corresponding value in kwargs
+    """
+    return format_str.format(*args, **kwargs)

pixeltable/functions/together.py

@@ -0,0 +1,122 @@
+import base64
+import io
+from typing import Optional
+
+import PIL.Image
+import numpy as np
+import together
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+def together_client() -> together.Together:
+    return env.Env.get().get_client('together', lambda api_key: together.Together(api_key=api_key))
+
+
+@pxt.udf
+def completions(
+        prompt: str,
+        *,
+        model: str,
+        max_tokens: Optional[int] = None,
+        stop: Optional[list] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None
+) -> dict:
+    return together_client().completions.create(
+        prompt=prompt,
+        model=model,
+        max_tokens=max_tokens,
+        stop=stop,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model
+    ).dict()
+
+
+@pxt.udf
+def chat_completions(
+        messages: list[dict[str, str]],
+        *,
+        model: str,
+        max_tokens: Optional[int] = None,
+        stop: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        logprobs: Optional[int] = None,
+        echo: Optional[bool] = None,
+        n: Optional[int] = None,
+        safety_model: Optional[str] = None,
+        response_format: Optional[dict] = None,
+        tools: Optional[dict] = None,
+        tool_choice: Optional[dict] = None
+) -> dict:
+    return together_client().chat.completions.create(
+        messages=messages,
+        model=model,
+        max_tokens=max_tokens,
+        stop=stop,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        logprobs=logprobs,
+        echo=echo,
+        n=n,
+        safety_model=safety_model,
+        response_format=response_format,
+        tools=tools,
+        tool_choice=tool_choice
+    ).dict()
+
+
+@pxt.udf(batch_size=32, return_type=pxt.ArrayType((None,), dtype=pxt.FloatType()))
+def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+    result = together_client().embeddings.create(input=input, model=model)
+    return [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+
+
+@pxt.udf
+def image_generations(
+        prompt: str,
+        *,
+        model: str,
+        steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        negative_prompt: Optional[str] = None,
+) -> PIL.Image.Image:
+    # TODO(aaron-siegel): Decompose CPU/GPU ops into separate functions
+    result = together_client().images.generate(
+        prompt=prompt,
+        model=model,
+        steps=steps,
+        seed=seed,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt
+    )
+    b64_str = result.data[0].b64_json
+    b64_bytes = base64.b64decode(b64_str)
+    img = PIL.Image.open(io.BytesIO(b64_bytes))
+    img.load()
+    return img