Working with Gemini in Pixeltable

%pip install -qU pixeltable google-genai

import os

vertex_enabled = (
    os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '').lower() == 'true'
)

# Option 1: Google AI Studio (API key)
# Set GEMINI_API_KEY or GOOGLE_API_KEY in your environment,
# or add api_key to the [gemini] section of $PIXELTABLE_HOME/config.toml

if (
    not vertex_enabled
    and 'GEMINI_API_KEY' not in os.environ
    and 'GOOGLE_API_KEY' not in os.environ
):
    import getpass

    os.environ['GEMINI_API_KEY'] = getpass.getpass(
        'Google AI Studio API Key:'
    )

# Option 2: Vertex AI (Application Default Credentials)
# Uncomment and set the following environment variables to use Vertex AI instead:
# os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'
# os.environ['GOOGLE_CLOUD_PROJECT'] = 'your-project-id'
# os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'  # optional, defaults to us-central1
# Then authenticate via: gcloud auth application-default login

import pixeltable as pxt

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

from google.genai.types import GenerateContentConfigDict
from pixeltable.functions import gemini

# Create a table in Pixeltable and pick a model hosted on Google AI Studio with some parameters

t = pxt.create_table('gemini_demo/text', {'input': pxt.String})

config = GenerateContentConfigDict(
    max_output_tokens=300, temperature=1.0, top_p=0.95, top_k=40
)
t.add_computed_column(
    output=gemini.generate_content(
        t.input, model='gemini-2.5-flash', config=config
    )
)

# Ask Gemini to generate some content based on the input
t.insert(
    [
        {'input': 'Write a story about a magic backpack.'},
        {'input': 'Tell me a science joke.'},
    ]
)

# Parse the response into a new column
t.add_computed_column(
    response=t.output['candidates'][0]['content']['parts'][0]['text']
)
t.select(t.input, t.response).head()

from google.genai.types import GenerateImagesConfigDict
from pixeltable.functions import gemini

images_t = pxt.create_table('gemini_demo/images', {'prompt': pxt.String})

config = GenerateImagesConfigDict(aspect_ratio='16:9')
images_t.add_computed_column(
    generated_image=gemini.generate_images(
        images_t.prompt, model='imagen-4.0-generate-001', config=config
    )
)

images_t.insert(
    [{'prompt': 'A friendly dinosaur playing tennis in a cornfield'}]
)

images_t.head()

videos_t = pxt.create_table('gemini_demo/videos', {'prompt': pxt.String})

videos_t.add_computed_column(
    generated_video=gemini.generate_videos(
        videos_t.prompt, model='veo-2.0-generate-001'
    )
)

videos_t.insert(
    [
        {
            'prompt': 'A giant pixel floating over the open ocean in a sea of data'
        }
    ]
)

videos_t.head()

images_t.add_computed_column(
    generated_video=gemini.generate_videos(
        image=images_t.generated_image, model='veo-2.0-generate-001'
    )
)

images_t.head()