Erstellen eines KI-gestützten API-Dienstes zur Bildgenerierung mit FLUX, Python und Diffusoren

FLUX (von Black Forest Labs) hat in den letzten Monaten die Welt der KI-Bildgenerierung im Sturm erobert. Es hat nicht nur Stable Diffusion (den früheren Open-Source-König) in vielen Benchmarks geschlagen, sondern in einigen Kennzahlen auch proprietäre Modelle wie Dall-E oder Midjourney übertroffen.

Aber wie würden Sie FLUX in einer Ihrer Apps verwenden? Man könnte darüber nachdenken, serverlose Hosts wie Replicate und andere zu verwenden, aber diese können sehr schnell sehr teuer werden und bieten möglicherweise nicht die Flexibilität, die Sie benötigen. Hier bietet sich die Erstellung Ihres eigenen benutzerdefinierten FLUX-Servers an.

In diesem Artikel führen wir Sie durch die Erstellung Ihres eigenen FLUX-Servers mit Python. Mit diesem Server können Sie über eine einfache API Bilder basierend auf Textaufforderungen generieren. Ganz gleich, ob Sie diesen Server für den persönlichen Gebrauch betreiben oder ihn als Teil einer Produktionsanwendung bereitstellen, dieser Leitfaden hilft Ihnen beim Einstieg.

Voraussetzungen

Bevor wir uns mit dem Code befassen, stellen wir sicher, dass Sie die erforderlichen Tools und Bibliotheken eingerichtet haben:

• Python: Auf Ihrem Computer muss Python 3 installiert sein, vorzugsweise Version 3.10.
• Torch: Das Deep-Learning-Framework, das wir zum Ausführen von FLUX verwenden werden.
• Diffusoren: Bietet Zugriff auf das FLUX-Modell.
• Transformatoren: Erforderliche Abhängigkeit von Diffusoren.
• Satzstück: Erforderlich, um den FLUX-Tokenizer auszuführen
• protobuf: Erforderlich zum Ausführen von FLUX
• beschleunigen: Hilft in einigen Fällen, das FLUX-Modell effizienter zu laden.
• fastapi: Framework zum Erstellen eines Webservers, der Bildgenerierungsanfragen akzeptieren kann.
• uvicorn: Erforderlich, um den FastAPI-Server auszuführen.
• psutil: Ermöglicht uns zu überprüfen, wie viel RAM auf unserem Computer vorhanden ist.

Sie können alle Bibliotheken installieren, indem Sie den folgenden Befehl ausführen: pip install Torch Diffusers Transformers Satzstück Protobuf Beschleunigung Fastapi Uvicorn.

Wenn Sie einen Mac mit einem M1- oder M2-Chip verwenden, sollten Sie PyTorch mit Metal für optimale Leistung einrichten. Befolgen Sie die offizielle PyTorch with Metal-Anleitung, bevor Sie fortfahren.

Sie müssen außerdem sicherstellen, dass Sie über mindestens 12 GB VRAM verfügen, wenn Sie FLUX auf einem GPU-Gerät ausführen möchten. Oder mindestens 12 GB RAM für die Ausführung auf CPU/MPS (was langsamer ist).

Schritt 1: Einrichten der Umgebung

Lassen Sie uns das Skript beginnen, indem wir basierend auf der von uns verwendeten Hardware das richtige Gerät zum Ausführen der Inferenz auswählen.

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")

Sie können CPU, Cuda (für NVIDIA-GPUs) oder MPS (für Apples Metal Performance Shader) angeben. Das Skript prüft dann, ob das ausgewählte Gerät verfügbar ist und löst eine Ausnahme aus, wenn dies nicht der Fall ist.

Schritt 2: Laden des FLUX-Modells

Als nächstes laden wir das FLUX-Modell. Wir laden das Modell mit fp16-Präzision, was uns etwas Speicherplatz spart, ohne dass es zu großen Qualitätsverlusten kommt.

An dieser Stelle werden Sie möglicherweise aufgefordert, sich bei HuggingFace zu authentifizieren, da das FLUX-Modell geschlossen ist. Für eine erfolgreiche Authentifizierung müssen Sie ein HuggingFace-Konto erstellen, zur Modellseite gehen, die Bedingungen akzeptieren und dann aus Ihren Kontoeinstellungen ein HuggingFace-Token erstellen und es auf Ihrem Computer als HF_TOKEN-Umgebungsvariable hinzufügen.

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")

Hier laden wir das FLUX-Modell mithilfe der Diffusorbibliothek. Das von uns verwendete Modell ist black-forest-labs/FLUX.1-dev, geladen in fp16-Präzision.

Es gibt auch ein zeitschrittdestilliertes Modell namens FLUX Schnell, das eine schnellere Inferenz bietet, aber weniger detaillierte Bilder ausgibt, sowie ein FLUX Pro-Modell, das Closed-Source ist.
Wir verwenden hier den Euler-Scheduler, aber Sie können damit experimentieren. Weitere Informationen zu Planern finden Sie hier.
Da die Bilderzeugung ressourcenintensiv sein kann, ist es wichtig, die Speichernutzung zu optimieren, insbesondere wenn sie auf einer CPU oder einem Gerät mit begrenztem Speicher ausgeführt wird.

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)

Dieser Code überprüft den gesamten verfügbaren Speicher und ermöglicht die Aufmerksamkeitsverteilung, wenn das System über weniger als 64 GB RAM verfügt. Aufmerksamkeits-Slicing reduziert den Speicherverbrauch während der Bilderzeugung, was für Geräte mit begrenzten Ressourcen unerlässlich ist.

Schritt 3: Erstellen der API mit FastAPI

Als nächstes richten wir den FastAPI-Server ein, der eine API zum Generieren von Bildern bereitstellt.

# 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 ist ein beliebtes Framework zum Erstellen von Web-APIs mit Python. In diesem Fall verwenden wir es, um einen Server zu erstellen, der Anfragen zur Bildgenerierung akzeptieren kann. Wir verwenden auch die GZip-Middleware, um die Antwort zu komprimieren, was besonders nützlich ist, wenn Bilder im Base64-Format zurückgesendet werden.

In einer Produktionsumgebung möchten Sie möglicherweise die generierten Bilder in einem S3-Bucket oder einem anderen Cloud-Speicher speichern und die URLs anstelle der Base64-codierten Zeichenfolgen zurückgeben, um von einem CDN und anderen Optimierungen zu profitieren.

Schritt 4: Definieren des Anforderungsmodells

Wir müssen jetzt ein Modell für die Anfragen definieren, die unsere API akzeptiert.

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)

Dieses GenerateRequest-Modell definiert die Parameter, die zum Generieren eines Bildes erforderlich sind. Das Eingabeaufforderungsfeld ist die Textbeschreibung des Bildes, das Sie erstellen möchten. Zu den weiteren Feldern gehören die Bildabmessungen, die Anzahl der Inferenzschritte und die Stapelgröße.

Schritt 5: Erstellen des Bildgenerierungsendpunkts

Jetzt erstellen wir den Endpunkt, der Bildgenerierungsanfragen verarbeitet.

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")

Dieser Endpunkt übernimmt den Bildgenerierungsprozess. Zunächst wird überprüft, ob Höhe und Breite ein Vielfaches von 8 sind, wie von FLUX gefordert. Anschließend generiert es Bilder basierend auf der bereitgestellten Eingabeaufforderung und gibt sie als Base64-codierte Zeichenfolgen zurück.

Schritt 6: Starten des Servers

Zuletzt fügen wir noch etwas Code hinzu, um den Server zu starten, wenn das Skript ausgeführt wird.

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)

Dieser Code startet den FastAPI-Server auf Port 8000 und macht ihn dank der 0.0.0.0-Bindung nicht nur von http://localhost:8000 aus zugänglich, sondern auch von anderen Geräten im selben Netzwerk über die IP-Adresse des Host-Computers.

Schritt 7: Testen Sie Ihren Server lokal

Da Ihr FLUX-Server nun betriebsbereit ist, ist es an der Zeit, ihn zu testen. Sie können Curl, ein Befehlszeilentool zum Senden von HTTP-Anfragen, verwenden, um mit Ihrem Server zu interagieren:

# 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()

Dieser Befehl funktioniert nur auf UNIX-basierten Systemen, auf denen die Dienstprogramme curl, jq und base64 installiert sind. Abhängig von der Hardware, auf der der FLUX-Server gehostet wird, kann es auch einige Minuten dauern, bis der Vorgang abgeschlossen ist.

Abschluss

Herzlichen Glückwunsch! Sie haben erfolgreich Ihren eigenen FLUX-Server mit Python erstellt. Mit diesem Setup können Sie Bilder basierend auf Textaufforderungen über eine einfache API generieren. Wenn Sie mit den Ergebnissen des FLUX-Basismodells nicht zufrieden sind, können Sie eine Feinabstimmung des Modells in Betracht ziehen, um bei bestimmten Anwendungsfällen eine noch bessere Leistung zu erzielen.

Vollständiger Code

Den vollständigen Code, der in dieser Anleitung verwendet wird, finden Sie unten:

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)

Das obige ist der detaillierte Inhalt vonErstellen eines KI-gestützten API-Dienstes zur Bildgenerierung mit FLUX, Python und Diffusoren. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!