How to Run Flux Schnell Locally on M3 Max MacBook Pro: A Complete Guide

Master Flux Schnell on Your M3 Max MacBook Pro

Flux Schnell is a powerful open-source text-to-image model designed for efficient image generation with high-quality outputs. In this guide, we’ll show you how to set up and run Flux Schnell locally on an M3 Max MacBook Pro, enabling you to leverage its capabilities to create stunning visuals in just 30 seconds per image.

Requirements to Run Flux Schnell Locally

Before getting started, 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 Guide to Run Flux Schnell Locally

Step 1: Set Up the Environment

Create a new Conda environment and activate it:

conda create -n flux python=3.11 conda activate flux

Step 2: Install Dependencies

Install the necessary 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 in your terminal:

python flux_schnell.py

The script 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 current directory.

Understanding the Code

Importing Dependencies

The script uses PyTorch and the Diffusers library:

import torch from diffusers import FluxPipeline import diffusers

These modules are crucial for tensor operations and leveraging the Flux Schnell model.

Modifying the Rope Function

To ensure compatibility with Apple’s MPS device, the script modifies the Rotary Position Embedding function:

_flux_rope = diffusers.models.transformers.transformer_flux.rope def new_flux_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: if pos.device.type == "mps": return _flux_rope(pos.to("cpu"), dim, theta).to(device=pos.device) else: return _flux_rope(pos, dim, theta)

Loading the Model

The model is loaded from the Hugging Face hub:

pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", revision='refs/pr/1', torch_dtype=torch.bfloat16 ).to("mps")

Generating the Image

The following section handles image generation:

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]

Saving the Image

Finally, the generated image is saved:

out.save("flux_image.png")

Optimizing Performance

To achieve optimal results, consider these tips:

• Close unnecessary applications: Free up RAM and CPU resources.
• Adjust image dimensions: Use smaller dimensions for faster generation or larger ones for better quality.
• Experiment with num_inference_steps: Higher values improve quality but increase generation time.
• Utilize bfloat16 precision: Already implemented for efficiency.
• Batch processing: Generate multiple images in a single run for efficiency.

Troubleshooting

If you face issues while running Flux Schnell, try the following solutions:

• Update dependencies: Ensure all packages are up to date.
• Check resource usage: Monitor system usage via Activity Monitor.
• Clear PyTorch cache: Use torch.cuda.empty_cache().
• Restart the kernel: For Jupyter users, restarting the kernel can resolve persistent errors.

Conclusion

Setting up and running Flux Schnell locally on your M3 Max MacBook Pro allows you to create high-quality images from textual descriptions efficiently. By following this guide, you can harness the power of AI-driven image generation, optimize performance, and troubleshoot any issues along the way. Explore the endless possibilities of Flux Schnell and elevate your creative projects.