Flux.1-dev: Run a top image generation model on Modal

All posts

Back

Model Library

December 16, 2024•5 minute read

Yiren Lu@YirenLu

Solutions Engineer

What is Flux.1-dev?

Flux.1-dev is a powerful text-to-image model from Black Forest Labs. Together with Stable Diffusion 3.5, it is one of the best text-to-image models on the market.

To run the following code, you will need to:

Create an account at modal.com
Run pip install modal to install the modal Python package
Run modal setup to authenticate (if this doesn’t work, try python -m modal setup)
Copy the code below into a file called app.py
Run modal run app.py

Please note that this code is not optimized. For a more detailed example of how to run Flux fast with torch.compile, refer here.

import time
from io import BytesIO
from pathlib import Path

import modal


diffusers_commit_sha = "81cf3b2f155f1de322079af28f625349ee21ec6b"

cuda_dev_image = modal.Image.from_registry(
    "nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.11"
).entrypoint([])

flux_image = (
    cuda_dev_image.apt_install(
        "git",
        "libglib2.0-0",
        "libsm6",
        "libxrender1",
        "libxext6",
        "ffmpeg",
        "libgl1",
    )
    .pip_install(
        "invisible_watermark==0.2.0",
        "transformers==4.44.0",
        "huggingface_hub[hf_transfer]==0.26.2",
        "accelerate==0.33.0",
        "safetensors==0.4.4",
        "sentencepiece==0.2.0",
        "torch==2.5.0",
        f"git+https://github.com/huggingface/diffusers.git@{diffusers_commit_sha}",
        "numpy<2",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

app = modal.App("flux", image=flux_image)

with flux_image.imports():
    import diffusers
    import torch


@app.cls(gpu="H100", timeout=3600, secrets=modal.Secret("huggingface"))
class Model:
    @modal.enter()
    def enter(self):
        self.pipe = diffusers.FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
        ).to("cuda")

    @modal.method()
    def inference(self, prompt: str) -> bytes:
        print("Generating image...")
        image = self.pipe(
            prompt,
            output_type="pil",
            num_inference_steps=4,
        ).images[0]

        byte_stream = BytesIO()
        image.save(byte_stream, format="JPEG")
        return byte_stream.getvalue()


@app.local_entrypoint()
def main(prompt: str = "A majestic dragon soaring over snow-capped mountains"):
    output_dir = Path("/tmp/flux")
    output_dir.mkdir(exist_ok=True)

    t0 = time.time()
    image_bytes = Model().inference.remote(prompt)
    print(f"Generation time: {time.time() - t0:.2f} seconds")

    output_path = output_dir / "output.jpg"
    output_path.write_bytes(image_bytes)
    print(f"Saved to {output_path}")

What is Flux.1-dev?

Example code for running the Flux.1-dev embedding model on Modal

Ship your first app in minutes.