Working with Replicate in Pixeltable

Pixeltable's Replicate integration enables you to access Replicate's models via the Replicate API.

Prerequisites

• A Replicate account with an API token.

Important Notes

• Replicate usage may incur costs based on your Replicate plan.
• Be mindful of sensitive data and consider security measures when integrating with external services.

First you'll need to install required libraries and enter a Replicate API token.

%pip install -qU pixeltable replicate

import os
import getpass

if 'REPLICATE_API_TOKEN' not in os.environ:
    os.environ['REPLICATE_API_TOKEN'] = getpass.getpass('Replicate API Token:')

Now let's create a Pixeltable directory to hold the tables for our demo.

import pixeltable as pxt

# Remove the `replicate_demo` directory and its contents, if it exists
pxt.drop_dir('replicate_demo', force=True)
pxt.create_dir('replicate_demo')

Connected to Pixeltable database at:
postgresql+psycopg://postgres:@/pixeltable?host=/Users/asiegel/.pixeltable/pgdata
Created directory `replicate_demo`.

<pixeltable.catalog.dir.Dir at 0x334624190>

Chat Completions

Create a Table: In Pixeltable, create a table with columns to represent your input data and the columns where you want to store the results from Replicate.

from pixeltable.functions.replicate import run

# Create a table in Pixeltable and pick a model hosted on Replicate with some parameters

t = pxt.create_table('replicate_demo.chat', {'prompt': pxt.String})

input = {
    'system_prompt': 'You are a helpful assistant.',
    'prompt': t.prompt,
    # These parameters are optional and can be used to tune model behavior:
    'max_tokens': 300,
    'top_p': 0.9,
    'temperature': 0.8
}
t['output'] = run(input, ref='meta/meta-llama-3-8b-instruct')

Created table `chat`.
Added 0 column values with 0 errors.

# Parse the response into a new column
t['response'] = pxt.functions.string.join('', t.output)

Added 0 column values with 0 errors.

# Start a conversation
t.insert(prompt='What foods are rich in selenium?')
t.select(t.prompt, t.response).show()

Computing cells: 100%|████████████████████████████████████████████| 4/4 [00:07<00:00,  1.89s/ cells]
Inserting rows into `chat`: 1 rows [00:00, 171.89 rows/s]
Computing cells: 100%|████████████████████████████████████████████| 4/4 [00:07<00:00,  1.89s/ cells]
Inserted 1 row with 0 errors.

Image Generation

Here's an example that shows how to use Replicate's image generation models with Pixeltable. We'll use the FLUX Schnell model.

t = pxt.create_table('replicate_demo.images', {'prompt': pxt.String})

input = {
    'prompt': t.prompt,
    'go_fast': True,
    'megapixels': '1'
}
t['output'] = run(input, ref='black-forest-labs/flux-schnell')

Created table `images`.
Added 0 column values with 0 errors.

t.insert(prompt='Draw a pencil sketch of a friendly dinosaur playing tennis in a cornfield.')

Computing cells: 100%|████████████████████████████████████████████| 2/2 [00:00<00:00,  2.17 cells/s]
Inserting rows into `images`: 1 rows [00:00, 198.61 rows/s]
Computing cells: 100%|████████████████████████████████████████████| 2/2 [00:00<00:00,  2.14 cells/s]
Inserted 1 row with 0 errors.

UpdateStatus(num_rows=1, num_computed_values=2, num_excs=0, updated_cols=[], cols_with_excs=[])

t.select(t.prompt, t.output).collect()

We see that Replicate returns our image as an array containing a single URL. To turn it into an actual image, we cast the string to type pxt.Image in a new computed column:

t['image'] = t.output[0].astype(pxt.Image)
t.select(t.image).collect()

Computing cells: 100%|████████████████████████████████████████████| 1/1 [00:00<00:00, 93.29 cells/s]
Added 1 column value with 0 errors.

Learn More

To learn more about advanced techniques like RAG operations in Pixeltable, check out the RAG Operations in Pixeltable tutorial.

If you have any questions, don't hesitate to reach out.