Servicio API de generación de imágenes con tecnología de IA con FLUX, Python y difusores: una guía rápida

En este artículo, te guiaremos en la creación de tu propio servidor FLUX con Python. Este servidor te permitirá generar imágenes basadas en indicaciones de texto a través de una API simple. Ya sea que estés ejecutando este servidor para uso personal o lo estés implementando como parte de una aplicación de producción, esta guía te ayudará a comenzar.

FLUX (de Black Forest Labs ) ha conquistado el mundo de la generación de imágenes con IA en los últimos meses. No solo ha superado a Stable Diffusion (el anterior rey del código abierto) en muchos puntos de referencia, sino que también ha superado a modelos propietarios como Dall-E o Midjourney en algunas métricas.

Pero, ¿cómo utilizarías FLUX en una de tus aplicaciones? Se podría pensar en utilizar hosts sin servidor como Replicate y otros, pero estos pueden volverse muy costosos muy rápidamente y es posible que no brinden la flexibilidad que necesitas. Ahí es donde resulta útil crear tu propio servidor FLUX personalizado.

Prerrequisitos

Antes de sumergirnos en el código, asegurémonos de tener configuradas las herramientas y bibliotecas necesarias:

• Python: necesitarás tener Python 3 instalado en tu máquina, preferiblemente la versión 3.10.
• torch : el marco de aprendizaje profundo que usaremos para ejecutar FLUX.
• diffusers : Proporciona acceso al modelo FLUX.
• transformers : Dependencia necesaria de difusores.
• sentencepiece : Necesario para ejecutar el tokenizador FLUX
• protobuf : necesario para ejecutar FLUX
• accelerate : ayuda a cargar el modelo FLUX de manera más eficiente en algunos casos.
• fastapi : Marco para crear un servidor web que pueda aceptar solicitudes de generación de imágenes.
• uvicorn : necesario para ejecutar el servidor FastAPI.
• psutil : Nos permite comprobar cuánta RAM hay en nuestra máquina.

Puede instalar todas las bibliotecas ejecutando el siguiente comando: pip install torch diffusers transformers sentencepiece protobuf accelerate fastapi uvicorn .

Si utiliza una Mac con un chip M1 o M2, debe configurar PyTorch con Metal para obtener un rendimiento óptimo. Siga la guía oficial de PyTorch con Metal antes de continuar.

También deberás asegurarte de tener al menos 12 GB de VRAM si planeas ejecutar FLUX en un dispositivo GPU. O al menos 12 GB de RAM para ejecutarlo en CPU/MPS (que será más lento).

Paso 1: Configuración del entorno

Comencemos el script eligiendo el dispositivo correcto para ejecutar la inferencia en función del hardware que estemos usando.

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

Puede especificar cpu , cuda (para GPU NVIDIA) o mps (para Metal Performance Shaders de Apple). Luego, el script verifica si el dispositivo seleccionado está disponible y genera una excepción si no lo está.

Paso 2: Carga del modelo FLUX

A continuación, cargamos el modelo FLUX. Lo cargaremos con precisión fp16, lo que nos permitirá ahorrar algo de memoria sin perder mucha calidad.

En este punto, es posible que se le solicite que se autentique con HuggingFace, ya que el modelo FLUX está bloqueado. Para autenticarse correctamente, deberá crear una cuenta de HuggingFace, ir a la página del modelo, aceptar los términos y luego crear un token de HuggingFace desde la configuración de su cuenta y agregarlo en su máquina como la variable de entorno HF_TOKEN .

 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)

Aquí, cargamos el modelo FLUX usando la biblioteca de difusores. El modelo que estamos usando es black-forest-labs/FLUX.1-dev , cargado en precisión fp16.

También existe un modelo destilado por pasos de tiempo llamado FLUX Schnell que tiene una inferencia más rápida, pero genera imágenes menos detalladas, así como un modelo FLUX Pro que es de código cerrado. Aquí usaremos el programador de Euler, pero puedes experimentar con él. Puedes leer más sobre los programadores aquí . Dado que la generación de imágenes puede consumir muchos recursos, es fundamental optimizar el uso de la memoria, especialmente cuando se ejecuta en una CPU o un dispositivo con memoria limitada.

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

Este código verifica la memoria total disponible y habilita la segmentación de atención si el sistema tiene menos de 64 GB de RAM. La segmentación de atención reduce el uso de memoria durante la generación de imágenes, lo cual es esencial para dispositivos con recursos limitados.

Paso 3: Creación de la API con FastAPI

A continuación, configuraremos el servidor FastAPI, que proporcionará una API para generar imágenes.

 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)

FastAPI es un marco popular para crear API web con Python. En este caso, lo estamos usando para crear un servidor que pueda aceptar solicitudes de generación de imágenes. También estamos usando el middleware GZip para comprimir la respuesta, lo que resulta particularmente útil cuando se envían imágenes en formato base64.

En un entorno de producción, es posible que desee almacenar las imágenes generadas en un depósito S3 u otro almacenamiento en la nube y devolver las URL en lugar de las cadenas codificadas en base64, para aprovechar una CDN y otras optimizaciones.

Paso 4: Definición del modelo de solicitud

Ahora necesitamos definir un modelo para las solicitudes que aceptará nuestra API.

 class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch")

Este modelo GenerateRequest define los parámetros necesarios para generar una imagen. El campo prompt es la descripción de texto de la imagen que desea crear. Otros campos incluyen las dimensiones de la imagen, la cantidad de pasos de inferencia y el tamaño del lote.

Paso 5: Creación del punto final de generación de imágenes

Ahora, creemos el punto final que manejará las solicitudes de generación de imágenes.

 @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URLs instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, }

Este punto final se encarga del proceso de generación de imágenes. Primero, valida que la altura y el ancho sean múltiplos de 8, como lo requiere FLUX. Luego, genera imágenes según el mensaje proporcionado y las devuelve como cadenas codificadas en base64.

Paso 6: Iniciar el servidor

Por último, agreguemos algo de código para iniciar el servidor cuando se ejecuta el script.

 @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)

Este código inicia el servidor FastAPI en el puerto 8000, haciéndolo accesible no solo desde http://localhost:8000 sino también desde otros dispositivos en la misma red usando la dirección IP de la máquina host, gracias al enlace 0.0.0.0 .

Paso 7: Prueba de tu servidor localmente

Ahora que el servidor FLUX está en funcionamiento, es hora de probarlo. Puede utilizar curl , una herramienta de línea de comandos para realizar solicitudes HTTP, para interactuar con el servidor:

 curl -X POST "http://localhost:8000/" \ -H "Content-Type: application/json" \ -d '{ "prompt": "A futuristic cityscape at sunset", "seed": 42, "height": 1024, "width": 1024, "cfg": 3.5, "steps": 50, "batch_size": 1 }' | jq -r '.images[0]' | base64 -d > test.png

Este comando solo funcionará en sistemas basados en UNIX con las utilidades curl , jq y base64 instaladas. También puede tardar unos minutos en completarse según el hardware que aloje el servidor FLUX.

Conclusión

¡Felicitaciones! Has creado con éxito tu propio servidor FLUX con Python. Esta configuración te permite generar imágenes basadas en indicaciones de texto a través de una API simple. Si no estás satisfecho con los resultados del modelo FLUX base, puedes considerar ajustar el modelo para obtener un rendimiento aún mejor en casos de uso específicos .

Código completo

Puede encontrar el código completo utilizado en esta guía a continuación:

 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") 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) # 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() 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) class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch") @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URL's instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, } @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)