Want to try out FLUX online without additional hassle?
You can run FLUX.1 Pro, Schnell and Dev at Anakin AI Now!
Flux Schnell is a robust text-to-image generation model capable of producing high-quality images from textual prompts. In this tutorial, we'll guide you through the steps to set up and run Flux Schnell locally on an M3 Max MacBook Pro. This configuration enables efficient image generation by leveraging up to 40GB of RAM, with each image taking around 30 seconds to produce.
Requirements to Run Flux.1 Pro on Mac Locally
Before we begin, ensure your system meets the following requirements:
- MacBook Pro with M3 Max chip
- macOS (latest version recommended)
- At least 40GB of available RAM
- Anaconda or Miniconda installed
- Internet connection for downloading dependencies
Step-by-Step Tutorial to Run Flux.1 Pro on Mac Locally
Step-by-Step Tutorial to Run Flux.1 Pro on Mac Locally
Step 1: Set Up the Environment
First, we’ll create a new Conda environment and activate it:
conda create -n flux python=3.11
conda activate flux
Step 2: Install Dependencies
Now, let’s install the required packages:
pip install torch==2.3.1
pip install git+https://github.com/huggingface/diffusers.git
pip install transformers==4.43.3 sentencepiece==0.2.0 accelerate==0.33.0 protobuf==5.27.3
Step 3: Prepare the Python Script
Create a new Python file named flux_schnell.py and add the following code:
import torch
from diffusers import FluxPipeline
import diffusers# Modify the rope function to handle MPS device
_flux_rope = diffusers.models.transformers.transformer_flux.rope
def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
assert dim % 2 == 0, "The dimension must be even."
if pos.device.type == "mps":
return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device)
else:
return _flux_rope(pos, dim, theta)
diffusers.models.transformers.transformer_flux.rope = new_flux_rope
# Load the Flux Schnell model
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
revision='refs/pr/1',
torch_dtype=torch.bfloat16
).to("mps")
# Define the prompt
prompt = "A cat holding a sign that says hello world"
# Generate the image
out = pipe(
prompt=prompt,
guidance_scale=0.,
height=1024,
width=1024,
num_inference_steps=4,
max_sequence_length=256,
).images[0]
# Save the generated image
out.save("flux_image.png")
Step 4: Run the Script
Execute the Python script:
python flux_schnell.py
This will generate an image based on the prompt “A cat holding a sign that says hello world” and save it as “flux_image.png” in the same directory.
Understanding the Code
Let’s break down the key components of the script:
Importing Dependencies
import torch
from diffusers import FluxPipeline
import diffusers
These lines import the necessary modules: PyTorch for tensor operations, FluxPipeline from the diffusers library to use the Flux Schnell model, and the diffusers module for additional utilities.
Modifying the Rope Function
_flux_rope = diffusers.models.transformers.transformer_flux.rope
def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
assert dim % 2 == 0, "The dimension must be even."
if pos.device.type == "mps":
return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device)
else:
return _flux_rope(pos, dim, theta)diffusers.models.transformers.transformer_flux.rope = new_flux_rope
This section modifies the rope (Rotary Position Embedding) function to handle the MPS (Metal Performance Shaders) device used in Apple Silicon chips. It ensures compatibility with the M3 Max chip by moving tensors to CPU when necessary.
Loading the Model
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
revision='refs/pr/1',
torch_dtype=torch.bfloat16
).to("mps")
This code loads the Flux Schnell model from the Hugging Face model hub. It specifies the model version, sets the data type to bfloat16 for efficiency, and moves the model to the MPS device.
Generating the Image
prompt = "A cat holding a sign that says hello world"out = pipe(
prompt=prompt,
guidance_scale=0.,
height=1024,
width=1024,
num_inference_steps=4,
max_sequence_length=256,
).images[0]
Here, we define the prompt and use the loaded model to generate an image. The parameters control various aspects of the generation process:
guidance_scale: Set to 0 for unconditional generation.heightandwidth: Set the dimensions of the output image to 1024x1024 pixels.num_inference_steps: Controls the number of denoising steps (4 in this case for faster generation).max_sequence_length: Limits the maximum length of the input sequence.
Saving the Image
out.save("flux_image.png")
Finally, this line saves the generated image as a PNG file named “flux_image.png” in the current directory.
Optimizing Performance
To make the most of your M3 Max MacBook Pro and achieve the ~30-second generation time with up to 40GB RAM usage, consider the following tips:
- Close unnecessary applications: Ensure other resource-intensive applications are closed to free up RAM and CPU resources.
- Adjust image dimensions: If you need faster generation times, you can reduce the
heightandwidthparameters. Conversely, you can increase them for higher quality at the cost of longer generation times. - Experiment with
num_inference_steps: Increasing this value can lead to higher quality images but will increase generation time. Decreasing it will speed up the process but may reduce image quality. - Use bfloat16 precision: As already implemented in the script, using
torch.bfloat16data type helps in reducing memory usage and potentially increasing speed without significant loss in quality. - Batch processing: If you need to generate multiple images, consider implementing batch processing to make more efficient use of the GPU.
Troubleshooting
If you encounter any issues while running Flux Schnell on your M3 Max MacBook Pro, try the following:
- Update dependencies: Ensure all packages are up to date, especially PyTorch and diffusers.
- Check CUDA compatibility: Although the M3 Max uses MPS, some CUDA-related errors might occur. Make sure you’re using the correct versions of PyTorch and related libraries for your system.
- Monitor resource usage: Use Activity Monitor to keep an eye on RAM and CPU usage. If you’re consistently hitting the limits, you may need to adjust the model parameters or free up more system resources.
- Clear cache: If you encounter memory-related errors, try clearing the PyTorch cache:
torch.cuda.empty_cache()
5. Restart kernel: If you’re using a Jupyter notebook or similar environment, try restarting the kernel if you encounter persistent errors.
Conclusion
Running Flux Schnell locally on your M3 Max MacBook Pro offers powerful text-to-image generation capabilities. With the setup outlined in this tutorial, you can create high-quality images from textual descriptions in about 30 seconds, utilizing up to 40GB of RAM. Keep in mind that performance may vary depending on your hardware and the complexity of the prompts.
Feel free to experiment with different parameters to achieve the best balance between speed and image quality for your needs. As AI models and hardware continue to advance, even more impressive results are on the horizon. Happy generating!