Visko runs live AI video model on Google Cloud TPU

Visko has deployed its Orbis live world model on Google Cloud TPUs, with the full live inference pipeline reaching 1.07x playback speed on a single TPU slice.

The work combined video generation, super-resolution upscaling and synchronised audio in one pipeline, moving from initial work to production handover in 30 days. Orbis generates a continuous live video stream from a text or image prompt and lets viewers add new prompts during a session without restarting the scene.

The project addresses a different problem from conventional video generation tools, which can process requests offline while users wait. In Visko's system, each chunk of video and audio must be completed before the previous chunk finishes playing, or the stream stops.

That creates two constraints for engineers. Each session effectively runs at a batch size of one because every viewer interacts with an individual stream. Each new chunk also depends on the previous one, because the model carries forward a rolling memory of recent video latents.

The Orbis pipeline was designed to run on a single Cloud TPU slice using BF16 precision. The three models in the stack - the core video diffusion model, the upscaler and the audio generator - share the same hardware sequentially for each chunk.

The software stack combined JAX, torchax and Pallas. Google Cloud describes JAX as its Python array library for accelerators, while torchax lets PyTorch models run on Cloud TPUs through JAX. Pallas is used for custom TPU kernels.

According to the companies, the main performance burden was video denoising, which took nearly two-thirds of the total compute budget. That led the teams to focus on the model's attention and parallelism layers.

Key changes

One of the main changes was the use of Ulysses context parallelism across TPU chips. Rather than splitting model weights, the method splits the token sequence across chips, then uses all-to-all interconnect exchanges around the attention layer so each chip can process the full sequence for a portion of the attention heads.

The teams also integrated a custom Pallas splash attention kernel from MaxDiffusion, Google Cloud's open-source reference library for diffusion models on Cloud TPU. The method avoids storing the full attention matrix in high-bandwidth memory by computing attention in tiles held in on-chip memory.

A third step involved rewriting the audio model. All three models were first moved onto Cloud TPU with torchax, but the audio stage remained too slow, so the team rewrote it in native JAX.

That change produced a 5.8x speed-up for the audio model, according to the companies, and brought the full pipeline under the playback deadline. The video and upscaling models remained in torchax in the final setup.

Visko positioned the exercise as a guide for startup infrastructure teams moving AI workloads from prototypes into production. One practical lesson, it said, was to use torchax first to get PyTorch models running quickly, then rewrite only the stages that dominate latency.

It also said BF16 was more suitable than FP8 quantisation for continuous streams because small quantisation errors can build up across chunks over time. Another operational issue was startup delay from XLA compilation, which the teams reduced by caching compiled artefacts so worker restarts could be brought under a minute.

Orbis has been presented as a live model rather than a static clip generator. Visko said the platform can respond to new prompts during a stream, changing the scene without breaking continuity, with applications ranging from interactive entertainment and live commerce to AI companions, immersive education and robotics simulation.

Google Cloud framed the work as an example of how smaller AI companies can move from a PyTorch prototype to production on its TPU infrastructure. The collaboration was led by Google Cloud TPU engineers including Yixin Shi, James Huang and Elisa Tsai.

"Visko is bringing an exciting new class of real-time, steerable workloads to the Cloud TPU ecosystem. By co-engineering its live multimodal pipeline on a single TPU slice using JAX, torchax and custom Pallas kernels, we proved how quickly a lean startup team can go from PyTorch prototype to rendering at 1.07x playback speed on Google Cloud's AI Hypercomputer," said Dave Elliott, head of developer advocacy for AI at Google Cloud.