Master OmniGen Setup on Mac M1, M2, M3: Your Easy Step-by-Step Guide

Setting up OmniGen on Apple Silicon Macs can be a bit tricky due to compatibility issues with the MPS (Metal Performance Shaders) backend. Here’s a comprehensive guide to help you get OmniGen running smoothly on your Mac M1, M2, or M3.

1. Create a Conda Environment

Start by creating and activating a new Conda environment to isolate your OmniGen installation:

conda create -n omnigen python=3.10
conda activate omnigen

2. Install PyTorch with MPS Support

On Apple Silicon Macs, PyTorch uses MPS for GPU acceleration. Install PyTorch and the necessary libraries using pip:

pip install torch torchvision torchaudio

To confirm that PyTorch is installed correctly and MPS is available, run:

python -c "import torch; print(torch.__version__, torch.backends.mps.is_available())"

Make sure that the output confirms MPS is available (True).

3. Clone and Install OmniGen

Clone the OmniGen repository and install it in editable mode:

git clone https://github.com/staoxiao/OmniGen.git
cd OmniGen
pip install -e .

Check the Installation:

To ensure OmniGen is installed successfully, run:

python -c "import OmniGen; print('OmniGen installed successfully')"

4. Create and Modify the omnigen_example.py File

Create omnigen_example.py with the following content:

from OmniGen import OmniGenPipeline
import torch

# Check if MPS (Metal Performance Shaders) is available
device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")

# Load the OmniGen model
pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")

# Move the model to the specified device
pipe.to(device)

# Generate an image from text
images = pipe(
    prompt="A curly-haired man in a red shirt is drinking tea.",
    height=512,  # Reduced size to manage memory usage
    width=512,
    guidance_scale=2.5,
    seed=0,
    use_kv_cache=False,  # Disable key-value cache
    offload_kv_cache=False  # Disable offloading
)

# Save the output image
images[0].save("example_t2i.png")

This script adjusts the device settings and reduces the image dimensions to prevent memory issues on Apple Silicon Macs.

5. Run the omnigen_example.py Script

To generate an image using OmniGen, execute the script:

python omnigen_example.py

This should create and save the image example_t2i.png without running into memory problems.

Additional Tips

• Memory Management: If you encounter memory-related errors, try reducing the image size further or setting environment variables to optimize memory allocation:

export PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.5

• Resource Monitoring: Monitor your system’s resource usage to ensure optimal performance.

By following these steps, you’ll have OmniGen up and running on your Mac M1, M2, or M3, leveraging the MPS backend for efficient processing. If you have any issues, refer back to the guide to troubleshoot and optimize your setup. Enjoy creating images with OmniGen on your Apple Silicon Mac!