Working with Replicate in Pixeltable

%pip install -qU pixeltable replicate

import getpass
import os

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

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')

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.add_computed_column(
    output=run(input, ref='meta/meta-llama-3-8b-instruct')
)

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

# Start a conversation
t.insert([{'prompt': 'What foods are rich in selenium?'}])
t.select(t.prompt, t.response).show()

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

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

t.insert(
    [
        {
            'prompt': 'Draw a pencil sketch of a friendly dinosaur playing tennis in a cornfield.'
        }
    ]
)

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

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