Création d'un service API de génération d'images basé sur l'IA avec FLUX, Python et Diffusers

FLUX (de Black Forest Labs) a pris d'assaut le monde de la génération d'images IA au cours des derniers mois. Non seulement il a battu Stable Diffusion (l'ancien roi de l'open source) sur de nombreux benchmarks, mais il a également surpassé les modèles propriétaires comme Dall-E ou Midjourney dans certains paramètres.

Mais comment feriez-vous pour utiliser FLUX sur l’une de vos applications ? On pourrait penser à utiliser des hôtes sans serveur comme Replicate et autres, mais ceux-ci peuvent devenir très coûteux très rapidement et peuvent ne pas offrir la flexibilité dont vous avez besoin. C'est là que la création de votre propre serveur FLUX personnalisé s'avère utile.

Dans cet article, nous vous guiderons dans la création de votre propre serveur FLUX à l'aide de Python. Ce serveur vous permettra de générer des images basées sur des invites textuelles via une simple API. Que vous utilisiez ce serveur pour un usage personnel ou que vous le déployiez dans le cadre d'une application de production, ce guide vous aidera à démarrer.

Conditions préalables

Avant de plonger dans le code, assurons-nous que vous disposez des outils et des bibliothèques nécessaires :

• Python : vous aurez besoin de Python 3 installé sur votre machine, de préférence la version 3.10.
• torch : Le framework d'apprentissage profond que nous utiliserons pour exécuter FLUX.
• diffuseurs : Donne accès au modèle FLUX.
• transformateurs : Dépendance requise des diffuseurs.
• sentencepiece : requis pour exécuter le tokenizer FLUX
• protobuf : requis pour exécuter FLUX
• accélérer : permet de charger le modèle FLUX plus efficacement dans certains cas.
• fastapi : Framework pour créer un serveur Web pouvant accepter les requêtes de génération d'images.
• uvicorn : requis pour exécuter le serveur FastAPI.
• psutil : Nous permet de vérifier la quantité de RAM qu'il y a sur notre machine.

Vous pouvez installer toutes les bibliothèques en exécutant la commande suivante : pip install torch diffusers transformers sentencepiece protobuf accélérer fastapi uvicorn.

Si vous utilisez un Mac avec une puce M1 ou M2, vous devez configurer PyTorch avec Metal pour des performances optimales. Suivez le guide officiel PyTorch with Metal avant de continuer.

Vous devrez également vous assurer que vous disposez d'au moins 12 Go de VRAM si vous envisagez d'exécuter FLUX sur un appareil GPU. Ou au moins 12 Go de RAM pour fonctionner sur CPU/MPS (ce qui sera plus lent).

Étape 1 : Configuration de l'environnement

Commençons le script en choisissant le bon appareil pour exécuter l'inférence en fonction du matériel que nous utilisons.

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

Vous pouvez spécifier cpu, cuda (pour les GPU NVIDIA) ou mps (pour les Metal Performance Shaders d'Apple). Le script vérifie ensuite si le périphérique sélectionné est disponible et déclenche une exception si ce n'est pas le cas.

Étape 2 : Chargement du modèle FLUX

Ensuite, nous chargeons le modèle FLUX. Nous chargerons le modèle avec une précision fp16, ce qui nous fera gagner de la mémoire sans trop de perte de qualité.

À ce stade, il vous sera peut-être demandé de vous authentifier auprès de HuggingFace, car le modèle FLUX est sécurisé. Pour vous authentifier avec succès, vous devrez créer un compte HuggingFace, accéder à la page du modèle, accepter les conditions, puis créer un jeton HuggingFace à partir des paramètres de votre compte et l'ajouter sur votre machine en tant que variable d'environnement HF_TOKEN.

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

Ici, nous chargeons le modèle FLUX en utilisant la bibliothèque des diffuseurs. Le modèle que nous utilisons est black-forest-labs/FLUX.1-dev, chargé avec une précision fp16.

Il existe également un modèle distillé par pas de temps nommé FLUX Schnell qui a une inférence plus rapide, mais produit des images moins détaillées, ainsi qu'un modèle FLUX Pro qui est à source fermée.
Nous utiliserons ici le planificateur Euler, mais vous pouvez l'expérimenter. Vous pouvez en savoir plus sur les planificateurs ici.
Étant donné que la génération d'images peut être gourmande en ressources, il est crucial d'optimiser l'utilisation de la mémoire, en particulier lorsqu'elle est exécutée sur un processeur ou un appareil avec une mémoire limitée.

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

Ce code vérifie la mémoire totale disponible et active le découpage d'attention si le système dispose de moins de 64 Go de RAM. Le découpage de l'attention réduit l'utilisation de la mémoire lors de la génération d'images, ce qui est essentiel pour les appareils aux ressources limitées.

Étape 3 : Création de l'API avec FastAPI

Ensuite, nous allons configurer le serveur FastAPI, qui fournira une API pour générer des images.

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

FastAPI est un framework populaire pour créer des API Web avec Python. Dans ce cas, nous l'utilisons pour créer un serveur capable d'accepter les demandes de génération d'images. Nous utilisons également le middleware GZip pour compresser la réponse, ce qui est particulièrement utile lors du renvoi d'images au format base64.

Dans un environnement de production, vous souhaiterez peut-être stocker les images générées dans un compartiment S3 ou un autre stockage cloud et renvoyer les URL au lieu des chaînes codées en base64, pour profiter d'un CDN et d'autres optimisations.

Étape 4 : Définition du modèle de demande

Nous devons maintenant définir un modèle pour les requêtes que notre API acceptera.

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

Ce modèle GenerateRequest définit les paramètres requis pour générer une image. Le champ d'invite est la description textuelle de l'image que vous souhaitez créer. D'autres champs incluent les dimensions de l'image, le nombre d'étapes d'inférence et la taille du lot.

Étape 5 : Création du point de terminaison de génération d'images

Maintenant, créons le point de terminaison qui gérera les demandes de génération d'images.

device = 'cuda' # can also be 'cpu' or 'mps'

import os

# MPS support in PyTorch is not yet fully implemented
if device == 'mps':
  os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch

if device == 'mps' and not torch.backends.mps.is_available():
      raise Exception("Device set to MPS, but MPS is not available")
elif device == 'cuda' and not torch.cuda.is_available():
      raise Exception("Device set to CUDA, but CUDA is not available")

Ce point de terminaison gère le processus de génération d’images. Il valide d'abord que la hauteur et la largeur sont des multiples de 8, comme l'exige FLUX. Il génère ensuite des images en fonction de l'invite fournie et les renvoie sous forme de chaînes codées en base64.

Étape 6 : Démarrage du serveur

Enfin, ajoutons du code pour démarrer le serveur lorsque le script est exécuté.

from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline
import psutil

model_name = "black-forest-labs/FLUX.1-dev"

print(f"Loading {model_name} on {device}")

pipeline = FluxPipeline.from_pretrained(
      model_name,

      # Diffusion models are generally trained on fp32, but fp16
      # gets us 99% there in terms of quality, with just half the (V)RAM
      torch_dtype=torch.float16,

      # Ensure we don't load any dangerous binary code
      use_safetensors=True

      # We are using Euler here, but you can also use other samplers
      scheduler=FlowMatchEulerDiscreteScheduler()
).to(device)

Ce code démarre le serveur FastAPI sur le port 8000, le rendant accessible non seulement depuis http://localhost:8000 mais également depuis d'autres appareils sur le même réseau en utilisant l'adresse IP de la machine hôte, grâce à la liaison 0.0.0.0.

Étape 7 : tester votre serveur localement

Maintenant que votre serveur FLUX est opérationnel, il est temps de le tester. Vous pouvez utiliser curl, un outil de ligne de commande permettant d'effectuer des requêtes HTTP, pour interagir avec votre serveur :

# Recommended if running on MPS or CPU with < 64 GB of RAM
total_memory = psutil.virtual_memory().total
total_memory_gb = total_memory / (1024 ** 3)
if (device == 'cpu' or device == 'mps') and total_memory_gb < 64:
      print("Enabling attention slicing")
      pipeline.enable_attention_slicing()

Cette commande ne fonctionnera que sur les systèmes UNIX avec les utilitaires curl, jq et base64 installés. Cela peut également prendre jusqu'à quelques minutes en fonction du matériel hébergeant le serveur FLUX.

Conclusion

Félicitations ! Vous avez créé avec succès votre propre serveur FLUX en utilisant Python. Cette configuration vous permet de générer des images basées sur des invites textuelles via une simple API. Si vous n'êtes pas satisfait des résultats du modèle FLUX de base, vous pouvez envisager d'affiner le modèle pour obtenir des performances encore meilleures dans des cas d'utilisation spécifiques.

Code complet

Vous pouvez trouver le code complet utilisé dans ce guide ci-dessous :

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field, conint, confloat
from fastapi.middleware.gzip import GZipMiddleware
from io import BytesIO
import base64

app = FastAPI()

# We will be returning the image as a base64 encoded string
# which we will want compressed
app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

Ce qui précède est le contenu détaillé de. pour plus d'informations, suivez d'autres articles connexes sur le site Web de PHP en chinois!