Blog
ML Modeling with FLAIM and Hugging Face Diffusers
Steve Swoyer | 10 July 2024
![ML Modeling with FLAIM and Hugging Face Diffusers](/images/blog/canvas-on-easel.webp)
Diffuser models like Stable Diffusion 3.0 are lots of fun to play around with. But these models are more than just toys: you can build with them, too!
The fine folks at đ€ (Hugging Face) have made it surprisingly simple to do this, producing a virtual toolbox that lets you download and run Stable Diffusion 3 and other âDiffusers.â (Diffusers are a type of model that generates data by iteratively refining it across multiple steps.) Thanks to đ€âs Diffusers toolbox, running models like Stable Diffusion 3 is pretty close to painlessâonce you have it installed.
Flox makes it easy to create environments that contain đ€ Diffusers.
In fact, we created one you can download and use right now: FLAIM, the Flox AI Modeling environment. This short video shows what it can do:
In this blog, youâll learn how Flox equips you to build environments like FLAIM that deliver reproducible results across all phases of your SDLCâfrom local dev to CI and beyond. All without using containers, which canât give us the acceleration we need for projects like FLAIM.
How not to build a dev environment for running ML and AI models
Letâs deal with the elephant in the room: The solution is not to containerize Stable Diffusion.
In the first place, youâll need to build and maintain separate container images for popular machine architectures, like ARM and 64-bit x86. Second, getting Nvidiaâs CUDA toolkit running in containers isnât exactly a heavy lift, but it isnât super easy, either. And third, getting containerized ARM workloads to take advantage of Appleâs Metal GPU API is ⊠well, it really isnât easy.
The third one is the real showstopper. A solid one-third of professional developers use MacOS, more than six times the overall market share of that platform. (The data is from StackOverflowâs 2023 Developer Survey, but MacOs performed about the same in StackOverflowâs 2022 tally, too.)
So if you canât containerize Stable Diffusion, but you still need it as part of a reproducible development environment, what can you do?
The team at Flox would like to make a modest proposal.
A better way to build and run software
Why not run Stable Diffusion directly on your local system using Flox, a virtual environment manager and package manager rolled into one?
Flox works with *any* language or toolchain, and *any* language- or application-specific package manager. And by layering Flox environments, you can create the equivalent of multi-container runtimes on your local system, with transparent access to all of your tools, files, environment variables, secrets, etc. This isnât sleight of hand, much less magic, because Flox borrows from the same ideas and principles used to power Nix, the proven open-source package manager.
This walk-through showcases another unique Flox feature, running FLAIM from a remote environment. Flox remote environments are useful because theyâre ephemeral: you activate them when you need them, and they disappear when youâre done. This makes test-driving software painless.
First, weâll activate FLAIM from FloxHub. When you grab a Flox remote environment, it âactivatesâ in whatever directory youâre currently in. For this walk-through, weâll assume weâre in ~/dev/
.
~/dev % flox activate -r flox/flaim
You use the -r
switch plus the /
syntax to activate a Flox remote environment. This environment will cease to exist as soon as you type exit
and press the enter key.
Until then, it will provide a discrete context in which to run software.
Once activated, FLAIM attempts to enable GPU acceleration, if available. This is actually trickier than you might think! NVIDIAâs CUDA libraries donât live in the same locations across all Linux distributions.
In creating FLAIM, Floxâs Ross Turk built in some arresting visual cues, so youâll know if youâre running with GPU acceleration or not. As the screengrab below shows, these cues are hard to miss.
If you believe your FLAIM environment should be accelerated, but isnât, run the built in testaccel
alias to retry:
I donât use OSX, but Ross does. Hereâs what the output of testaccel
looks like on his MacBook Pro:
flox [flox/flaim (remote)] ~ % testaccel
Metal is available đ
When you try to run FLAIM on a non-GPU-accelerated platform, your output looks like this:
flox [flox/flaim (remote)] daedalus@hybris:~/dev$ testaccel
I only see a CPU đ
Things will still run in FLAIM, but make some popcorn. Itâll take awhile.
Putting FLAIM through its paces
Once FLAIM is active, you will have access to a few key packages:
~/dev % flox list
accelerate: python311Packages.accelerate (python3.11-accelerate-0.30.0)
diffusers: python311Packages.diffusers (python3.11-diffusers-0.29.0)
sentencepiece: python311Packages.sentencepiece (python3.11-sentencepiece-0.2.0)
transformers: python311Packages.transformers (python3.11-transformers-4.42.3)
pytorch: python311Packages.pytorch-bin (python3.11-torch-2.2.2)
You can now pass the following commands to the built-in Python interpreter, generating an image:
% python3
>>> import torch
>>> from diffusers import StableDiffusionPipeline
>>> pipe = StableDiffusionPipeline.from\_pretrained(
"runwayml/stable-diffusion-v1-5", torch\_dtype=torch.float16)
[this takes a while]
>>> pipe = pipe.to("cuda") # or "mps" on mac, or leave out for cpu
>>> image = pipe("a smiling brittany spaniel puppy").images[0]
>>> image.save("smiling-brittany-spaniel-puppy.png")
But no one should have to interact with a Python interpreter to generate images, right?! Happily, Ross built a script into FLAIM that does this for us. You can call it with the alias genimg
, passing your promptâenclosed in quotesâas part of the command.
flox [flox/flaim (remote)] daedalus@ethos:~/dev$ genimg \
"public domain steamboat willie letting his freak flag fly"
FLAIM is by no means a one-trick pony.
If you need a Python package that isnât already built into itâand thereâs a high likelihood of this!âyou can add it to requirements.txt:
daedalus@ethos:~/dev/myproject$ mkdir myproject && cd myproject
daedalus@ethos:~/dev/myproject$ echo "imgcat" > requirements.txt
daedalus@ethos:~/dev/myproject$ echo "click" >> requirements.txt
daedalus@ethos:~/dev/myproject$ flox activate -r flox/flaim
flox [flox/flaim (remote)] daedalus@ethos:~/dev/myproject$
Note: If you do create a requirements.txt file, you end up with a flaim-venv
virtual environment in the directory where you activated FLAIM. This means you can activate FLAIM in different projects, even if they have different dependency requirements.
DIY AIM
Thatâs basically it. But wait! Thereâs one more thing: Even though we built FLAIM for you, you could almost as easily Build Your Own AI Modeling (BYOAIM) environment. After all, one of the neatest things about Flox is that itâs simple to use and has just a few commands.
Another neat thing is how Flox borrows from Nix to give you a declarative way of managing and versioning software. A manifest.toml
file defines which versions of which packages are installed in your environment. Creating a Flox environment is as easy as adding the names of packages to the [install]
section of this file, which is nested in the ./.flox/env
folder inside your project directory.
For example, to BYOAIM, all youâve got to do is type
flox init
in your projectâs directory and copy the following (from FLAIM) package and version information into your manifest.toml
file.
[install]
pytorch.pkg-path = "python311Packages.pytorch-bin"
pytorch.version = "python3.11-torch-2.2.2"
pytorch.pkg-group = "torch"
accelerate.pkg-path = "python311Packages.accelerate"
accelerate.version = "python3.11-accelerate-0.30.0"
accelerate.pkg-group = "pkgs"
transformers.pkg-path = "python311Packages.transformers"
transformers.version = "python3.11-transformers-4.42.3"
transformers.pkg-group = "pkgs"
sentencepiece.pkg-path = "python311Packages.sentencepiece"
sentencepiece.version = "python3.11-sentencepiece-0.2.0"
sentencepiece.pkg-group = "pkgs"
diffusers.pkg-path = "python311Packages.diffusers"
diffusers.version = "python3.11-diffusers-0.29.0"
diffusers.pkg-group = "pkgs"
To customize your BYOAIM environment, just add environment variables, hooks, aliases, and other stuff (including functions and complete code blocks!) to the [vars]
, [hook]
, and [profile]
sections of your manifest.toml
.
Mic drop: Reproducible software development made simple
Youâve probably been frustrated with how difficult it is to build and deploy GenAI solutions, mainly because you canât guarantee that theyâll work on your teammatesâ systems, in CI, or in prod.
So maybe they still seem more like toys than useful tools to you.
But itâs time to rethink this. Flox gives you a way to ensure reproducibility and platform-optimized performance across the SDLC, whether youâre building collaboratively, pushing to CI, or deploying to prod. You can easily put Flox environments into containers, or you can (just as easily) build containers using Flox environments.
FLAIM is the latest example of this. Weâve also created a portable, reproducible, GPU-accelerated Flox environment for Ollama, which lets you run dozens of open-source LLMs. Weâve shipped environments that showcase how easy it is to access applications and servicesâlike Jupyter Notebookâin portable, reproducible ways.
All of this without using containers, while also taking full advantage of available hardware and operating system affordances.
Sound too good to be true? Download Flox and take FLAIM for a test-drive! And think about whether you have any projects where Flox can help you build and ship faster - with a catalog containing more than 1 million package and version combinations, the possibilities are endless!