September 19, 2024

Install and Run FLUX.1 [schnell] Text-to-Image model in Python on a Local Linux Computer

The YouTube tutorial given below explains how to install and run FLUX.1 -schnell text-to-image model in Python on a local Linux computer. Besides providing the tutorial, we also give codes for downloading the Flux. 1 schnell model files from the Hugging face website and for running the Flux. 1- schnell in Python. For the complete installation procedure see the YouTube video tutorial.

Background Information about FLUX.1-schnell and Hardware Requirements

Flux.1 [schnell] is a rectified flow transformer that is used to generate images from text descriptions. The rectified flow transformer has around 12 billion parameters. It is released under Apache 2.0 license, which means that this model can be used for personal, scientific, and commercial purposes. In this tutorial, we explain how to correctly install and how to run Flux 1 [schnell] in Python on a local Linux machine.

Software/Hardware prerequisites:

  • Official webpage states that the FLUX.1[schnell] can be run locally. However, they do not specify the required GPU memory requirements and other specs, or if they did, we could not find the specs. If you just blindly run the script on the official Hugging Face website, you will see that the script will not run unless GPU has 60 GB VRAM is available. However, you can actually run FLUX1.[schnell] on GPUs with less memory, by just adding a single line of code. We have tested FLUX1.[schnell] on NVIDIA 3090 with 24 GB VRAM and it takes around 30-60 seconds to generate an image. You can decrease this time by optimizing the code. More about this in our future video tutorials. To summarize, you will need a relatively modest GPU to run FLUX.1[schnell]. In the comment section in the YouTube video you can share with us your experiences in running FLUX.1[schnell] on different GPUs.
  • Most likely, you need the NVIDIA CUDA toolkit installed before you install the library. We are not 100% sure about this since we have NVIDIA CUDA toolkit on our system, that is, we did not test the installation without  NVIDIA CUDA toolkit. Install it just in case. Anyway, you will need it for some other machine-learning libraries. We created a tutorial on how to properly install NVIDIA CUDA toolkit in Linux Ubuntu.
  • You will need 50-100 GB of free space on your local disk in order to run FLUX.1[schnell] on the local computer. You will need this space to download all the files.

Python Codes for Downloading the Model and Running FLUX. 1 Schnell Locally

Here is the Python code for downloading the model weights from the Huggingface website (see the YouTube tutorial):

from huggingface_hub import snapshot_download

snapshot_download(repo_id="black-forest-labs/FLUX.1-schnell",
                  local_dir="/home/aleksandar/codes/flux1schnell")

Here is the Python code for running FLUX. 1. schnell locally (see the YouTube tutorial):

import matplotlib.pyplot as plt
import torch
from diffusers import FluxPipeline

# here, adjust the path
pipe = FluxPipeline.from_pretrained("/home/aleksandar/codes/flux1schnell", torch_dtype=torch.bfloat16)
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)

prompt = "Ancient greek soldier with a sword and a shield. Behind there are horses. In the background there is a mountain with snow."

image = pipe(
    prompt,
    guidance_scale=0.0,
    output_type="pil",
    num_inference_steps=4,
    max_sequence_length=256,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]


plt.imshow(image)
plt.show()
image.save("generated_imag.png")