We raised a Series A! Read a post from our CEO, Zhen Lu: 1M devs and the cloud we're building next.

Running JAX Diffusion Models on Runpod

Curious about JAX-based diffusion models? This post walks through setting up and running them on Runpod using our GPU pods. It covers environment setup.

Running JAX Diffusion Models on Runpod

Here's a quick and dirty guide to getting Huemin's JAX diffusion to run on runpod!

Today I figured I'd switch things up a bit and leave you a gif quick start instead of the usual article. It only takes a few minutes to get started with this and all the deps have been sorted out for you in the nb. If you need the notebook file that you should upload to your instance, you can download it from one of these two places: Link 1, Link 2

Jupyter notebook cell with Python code starting a JAX diffusion image generation run

Please note that it will sit here for quite a few minutes while it downloads models and prepared for the first render cycle:

__wf_reserved_inherit

Be patient and it will start rendering when it's ready. I was getting around 1.5-2s/it on the demo example.

Enjoy!

Credits to: Huemin (@huemin_art), for the notebook, Alexander Redde (@alexanderredde3) for working out the deps, and nsheppherd (@nshepperd1) and Rivers Have Wings (@rivershavewings) for Jax Diffusion

Author profile: Zhen Lu

Related articles

View All
What's new in Runpod Serverless: Faster cold starts, batch inference, and no-Docker deploys

What's new in Runpod Serverless: Faster cold starts, batch inference, and no-Docker deploys

Whether you're already running production endpoints on Runpod or you're sizing us up for the first time, here's a plain-language tour of what Runpod Serverless does today, why it's faster and cheaper than it was six months ago, and how to deploy your first endpoint in minutes.

All

Build what’s next.

Build, train, and scale AI workloads on Runpod with cloud GPUs, Serverless, and Clusters.