Google Colab¶

Google Colab works well for Rex when the weights stay remote and the notebook uses either an HTTPS host or a mounted Drive path served over HTTP.

Install in Colab¶

!pip install "rex-framework[pytorch]" -q
import rex, torch
print(rex.__version__, torch.__version__)

Load from an HTTPS host¶

from rex.api.config import RexConfig
import rex

MANIFEST_URL = "https://your-static-host.com/rex_output/manifest.json"
CHUNKS_BASE_URL = "https://your-static-host.com/rex_output/weights"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL
config.cache.max_memory_cache_bytes = 1024 * 1024 * 1024
config.scheduler.enable_prefetch = True
config.scheduler.prefetch_window = 4

runtime = rex.load_model(MANIFEST_URL, config=config)

Load from Google Drive¶

from google.colab import drive
drive.mount("/content/drive")

from rex.api.config import RexConfig
import rex

# Cell 2b — serve the mounted directory locally
!rex-serve --dir /content/drive/MyDrive/rex_output/weights --port 8080 &

MANIFEST_PATH = "/content/drive/MyDrive/rex_output/manifest.json"
CHUNKS_BASE_URL = "http://127.0.0.1:8080"

config = RexConfig()
config.storage.base_url = CHUNKS_BASE_URL

runtime = rex.load_model(MANIFEST_PATH, config=config)

Use Colab secrets¶

from google.colab import userdata

config.storage.auth_token = f"Bearer {userdata.get('REX_AUTH_TOKEN')}"

GPU check¶

import torch
print("CUDA available:", torch.cuda.is_available())
print("Device:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU")