Mencipta Perkhidmatan API Penjanaan Imej dikuasakan AI dengan FLUX, Python dan Peresap

FLUX (oleh Black Forest Labs) telah mengambil alih dunia penjanaan imej AI oleh ribut sejak beberapa bulan lalu. Ia bukan sahaja telah mengalahkan Stable Diffusion (raja sumber terbuka terdahulu) pada banyak penanda aras, ia juga telah mengatasi model proprietari seperti Dall-E atau Midjourney dalam beberapa metrik.

Tetapi bagaimanakah anda akan menggunakan FLUX pada salah satu apl anda? Seseorang mungkin terfikir untuk menggunakan hos tanpa pelayan seperti Replicate dan lain-lain, tetapi ini boleh menjadi sangat mahal dengan cepat, dan mungkin tidak memberikan fleksibiliti yang anda perlukan. Di situlah mencipta pelayan FLUX tersuai anda sendiri berguna.

Dalam artikel ini, kami akan membimbing anda membuat pelayan FLUX anda sendiri menggunakan Python. Pelayan ini akan membolehkan anda menjana imej berdasarkan gesaan teks melalui API mudah. Sama ada anda menjalankan pelayan ini untuk kegunaan peribadi atau menggunakannya sebagai sebahagian daripada aplikasi pengeluaran, panduan ini akan membantu anda bermula.

Prasyarat

Sebelum menyelami kod, mari pastikan anda menyediakan alatan dan perpustakaan yang diperlukan:

• Python: Anda memerlukan Python 3 dipasang pada mesin anda, sebaik-baiknya versi 3.10.
• obor: Rangka kerja pembelajaran mendalam yang akan kami gunakan untuk menjalankan FLUX.
• peresap: Menyediakan akses kepada model FLUX.
• transformer: Kebergantungan yang diperlukan penyebar.
• potongan ayat: Diperlukan untuk menjalankan tokenizer FLUX
• protobuf: Diperlukan untuk menjalankan FLUX
• mempercepatkan: Membantu memuatkan model FLUX dengan lebih cekap dalam sesetengah kes.
• fastapi: Rangka kerja untuk mencipta pelayan web yang boleh menerima permintaan penjanaan imej.
• uvicorn: Diperlukan untuk menjalankan pelayan FastAPI.
• psutil: Membolehkan kami menyemak jumlah RAM yang terdapat pada mesin kami.

Anda boleh memasang semua perpustakaan dengan menjalankan arahan berikut: pip install torch diffusers transformers sentencepiece protobuf accelerate fastapi uvicorn.

Jika anda menggunakan Mac dengan cip M1 atau M2, anda harus menyediakan PyTorch dengan Metal untuk prestasi optimum. Ikuti panduan rasmi PyTorch with Metal sebelum meneruskan.

Anda juga perlu memastikan anda mempunyai sekurang-kurangnya 12 GB VRAM jika anda merancang untuk menjalankan FLUX pada peranti GPU. Atau sekurang-kurangnya 12 GB RAM untuk berjalan pada CPU/MPS (yang akan menjadi lebih perlahan).

Langkah 1: Sediakan Persekitaran

Mari mulakan skrip dengan memilih peranti yang betul untuk menjalankan inferens berdasarkan perkakasan yang kami gunakan.

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

Anda boleh menentukan cpu, cuda (untuk GPU NVIDIA) atau mps (untuk Shader Prestasi Logam Apple). Skrip kemudian menyemak sama ada peranti yang dipilih tersedia dan menimbulkan pengecualian jika tidak.

Langkah 2: Memuatkan Model FLUX

Seterusnya, kami memuatkan model FLUX. Kami akan memuatkan model dalam ketepatan fp16 yang akan menjimatkan memori kami tanpa banyak kehilangan kualiti.

Pada ketika ini, anda mungkin diminta untuk mengesahkan dengan HuggingFace, kerana model FLUX berpagar. Untuk mengesahkan dengan jayanya, anda perlu membuat akaun HuggingFace, pergi ke halaman model, terima syarat dan kemudian buat token HuggingFace daripada tetapan akaun anda dan tambahkannya pada mesin anda sebagai pembolehubah persekitaran 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")

Di sini, kami memuatkan model FLUX menggunakan perpustakaan penyebar. Model yang kami gunakan ialah black-forest-labs/FLUX.1-dev, dimuatkan dalam ketepatan fp16.

Terdapat juga model suling langkah waktu bernama FLUX Schnell yang mempunyai inferens lebih pantas, tetapi menghasilkan imej yang kurang terperinci, serta model FLUX Pro yang sumber tertutup.
Kami akan menggunakan penjadual Euler di sini, tetapi anda boleh mencubanya. Anda boleh membaca lebih lanjut mengenai penjadual di sini.
Memandangkan penjanaan imej boleh menjadi intensif sumber, adalah penting untuk mengoptimumkan penggunaan memori, terutamanya apabila dijalankan pada CPU atau peranti dengan memori terhad.

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)

Kod ini menyemak jumlah memori yang tersedia dan membolehkan pemotongan perhatian jika sistem mempunyai kurang daripada 64 GB RAM. Penghirisan perhatian mengurangkan penggunaan memori semasa penjanaan imej, yang penting untuk peranti dengan sumber terhad.

Langkah 3: Mencipta API dengan FastAPI

Seterusnya, kami akan menyediakan pelayan FastAPI, yang akan menyediakan API untuk menjana imej.

# 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 ialah rangka kerja yang popular untuk membina API web dengan Python. Dalam kes ini, kami menggunakannya untuk mencipta pelayan yang boleh menerima permintaan untuk penjanaan imej. Kami juga menggunakan perisian tengah GZip untuk memampatkan respons, yang amat berguna apabila menghantar semula imej dalam format base64.

Dalam persekitaran pengeluaran, anda mungkin mahu menyimpan imej yang dijana dalam baldi S3 atau storan awan lain dan mengembalikan URL dan bukannya rentetan berkod base64, untuk memanfaatkan CDN dan pengoptimuman lain.

Langkah 4: Menentukan Model Permintaan

Kami kini perlu menentukan model untuk permintaan yang akan diterima oleh API kami.

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)

Model GenerateRequest ini mentakrifkan parameter yang diperlukan untuk menjana imej. Medan gesaan ialah perihalan teks imej yang ingin anda buat. Medan lain termasuk dimensi imej, bilangan langkah inferens dan saiz kelompok.

Langkah 5: Mencipta Titik Akhir Penjanaan Imej

Sekarang, mari buat titik akhir yang akan mengendalikan permintaan penjanaan imej.

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

Titik akhir ini mengendalikan proses penjanaan imej. Ia mula-mula mengesahkan bahawa ketinggian dan lebar adalah gandaan 8, seperti yang diperlukan oleh FLUX. Ia kemudian menjana imej berdasarkan gesaan yang disediakan dan mengembalikannya sebagai rentetan berkod base64.

Langkah 6: Memulakan Pelayan

Akhir sekali, mari tambahkan beberapa kod untuk memulakan pelayan apabila skrip dijalankan.

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)

Kod ini memulakan pelayan FastAPI pada port 8000, menjadikannya boleh diakses bukan sahaja dari http://localhost:8000 tetapi juga dari peranti lain pada rangkaian yang sama menggunakan alamat IP mesin hos, terima kasih kepada pengikatan 0.0.0.0.

Langkah 7: Menguji Pelayan Anda Secara Setempat

Sekarang pelayan FLUX anda sedang aktif dan berjalan, tiba masanya untuk mengujinya. Anda boleh menggunakan curl, alat baris arahan untuk membuat permintaan HTTP, untuk berinteraksi dengan pelayan anda:

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

Arahan ini hanya akan berfungsi pada sistem berasaskan UNIX dengan utiliti curl, jq dan base64 dipasang. Ia juga mungkin mengambil masa sehingga beberapa minit untuk disiapkan bergantung pada perkakasan yang mengehos pelayan FLUX.

Kesimpulan

Tahniah! Anda telah berjaya mencipta pelayan FLUX anda sendiri menggunakan Python. Persediaan ini membolehkan anda menjana imej berdasarkan gesaan teks melalui API mudah. Jika anda tidak berpuas hati dengan hasil model FLUX asas, anda mungkin mempertimbangkan untuk memperhalusi model untuk prestasi yang lebih baik pada kes penggunaan tertentu.

Kod penuh

Anda mungkin menemui kod penuh yang digunakan dalam panduan ini di bawah:

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)

Atas ialah kandungan terperinci Mencipta Perkhidmatan API Penjanaan Imej dikuasakan AI dengan FLUX, Python dan Peresap. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!