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
import pixeltable.functions as pxtf

pxt.create_dir('gemini_demo', if_exists='replace_force')

from google.genai.types import GenerateContentConfigDict

# 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=500, temperature=1.0, top_p=0.95, top_k=40
)
t.add_computed_column(
    output=pxtf.gemini.generate_content(
        t.input, model='gemini-2.5-flash', config=config
    )
)

# Ask Gemini for short responses so results are easy to inspect
t.insert(
    [
        {'input': 'Write a two-sentence story about a magic backpack.'},
        {'input': 'Tell a one-line science joke.'},
    ]
)

# Parse the response text into its own column
t.add_computed_column(
    response=t.output['candidates'][0]['content']['parts'][0]['text'],
    if_exists='replace',
)

t.select(t.input, t.response).collect()

nano_t = pxt.create_table(
    'gemini_demo/nano_banana_images', {'prompt': pxt.String}
)

config = GenerateContentConfigDict(
    response_modalities=['IMAGE'], image_config={'aspect_ratio': '16:9'}
)
nano_t.add_computed_column(
    response=pxtf.gemini.generate_content(
        nano_t.prompt, model='gemini-2.5-flash-image', config=config
    )
)
nano_t.add_computed_column(
    generated_image=(
        nano_t.response.candidates[0]
        .content.parts[0]
        .inline_data.data.astype(pxt.Image)
    )
)

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

nano_t.select(nano_t.prompt, nano_t.generated_image).collect()

from google.genai.types import GenerateImagesConfigDict

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

config = GenerateImagesConfigDict(aspect_ratio='16:9')
images_t.add_computed_column(
    generated_image=pxtf.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.select(images_t.prompt, images_t.generated_image).collect()

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

videos_t.add_computed_column(
    generated_video=pxtf.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.select(videos_t.prompt, videos_t.generated_video).collect()

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

nano_t.select(nano_t.prompt, nano_t.generated_video).collect()