Guía rápida de OmniGen¶

En este ejemplo, entrenaremos un Lycoris LoKr para OmniGen, enfocado en mejoras generales de rendimiento T2I (no entrenamiento de edición/instrucción por ahora).

Requisitos de hardware¶

OmniGen es un modelo de tamaño bastante moderado, alrededor de 3.8B parámetros; usa el VAE de SDXL pero no usa un codificador de texto. En su lugar, OmniGen usa IDs de tokens nativos como entradas y se comporta como un modelo multimodal.

El uso de memoria durante el entrenamiento aún no se conoce, pero se espera que quepa fácilmente en una tarjeta de 24G con tamaño de lote 2 o 3. El modelo puede cuantizarse, ahorrando más VRAM.

OmniGen tiene una arquitectura extraña respecto a otros modelos entrenables por SimpleTuner;

• Actualmente, solo se admite entrenamiento t2i (texto‑a‑imagen), donde la salida del modelo se alinea con los prompts de entrenamiento y las imágenes de entrada.
• Aún no se admite entrenamiento imagen‑a‑imagen, pero podría en el futuro.
• Este modo permite proporcionar una segunda imagen como entrada, y el modelo la usa como condicionamiento/referencia para la salida.
• El valor de pérdida al entrenar OmniGen es muy alto, y no se sabe por qué.

Requisitos previos¶

Asegúrate de tener python instalado; SimpleTuner funciona bien con 3.10 a 3.12.

Puedes comprobarlo ejecutando:

python --version

Si no tienes python 3.12 instalado en Ubuntu, puedes intentar lo siguiente:

apt -y install python3.13 python3.13-venv

Dependencias de la imagen de contenedor¶

Para Vast, RunPod y TensorDock (entre otros), lo siguiente funcionará en una imagen CUDA 12.2-12.8 para habilitar la compilación de extensiones CUDA:

apt -y install nvidia-cuda-toolkit

Instalación¶

Instala SimpleTuner vía pip:

pip install 'simpletuner[cuda]'

# CUDA 13 / Blackwell users (NVIDIA B-series GPUs)
pip install 'simpletuner[cuda13]' --extra-index-url https://download.pytorch.org/whl/cu130

Para instalación manual o entorno de desarrollo, consulta la documentación de instalación.

Configuración del entorno¶

Para ejecutar SimpleTuner, necesitas configurar un archivo de configuración, los directorios del dataset y del modelo, y un archivo de configuración del dataloader.

Archivo de configuración¶

Un script experimental, configure.py, puede permitirte omitir por completo esta sección mediante una configuración interactiva paso a paso. Incluye funciones de seguridad que ayudan a evitar errores comunes.

Nota: Esto no configura tu dataloader. Aún tendrás que hacerlo manualmente más adelante.

Para ejecutarlo:

simpletuner configure

⚠️ Para usuarios ubicados en países donde Hugging Face Hub no es fácilmente accesible, debes agregar HF_ENDPOINT=https://hf-mirror.com a tu ~/.bashrc o ~/.zshrc dependiendo de cuál $SHELL use tu sistema.

Si prefieres configurarlo manualmente:

Copia config/config.json.example a config/config.json:

cp config/config.json.example config/config.json

Allí, probablemente necesitarás modificar las siguientes variables:

• model_type - Configúralo en lora.
• lora_type - Configúralo en lycoris.
• model_family - Configúralo en omnigen.
• model_flavour - Configúralo en v1
• output_dir - Configúralo al directorio donde quieres guardar tus checkpoints y las imágenes de validación. Se recomienda usar una ruta completa aquí.
• train_batch_size - para una tarjeta de 24G con gradient checkpointing completo, puede ser tan alto como 6.
• validation_resolution - Este checkpoint de OmniGen es un modelo 1024px, debes configurarlo en 1024x1024 o una de las otras resoluciones compatibles de OmniGen.
• Se pueden especificar otras resoluciones separándolas con comas: 1024x1024,1280x768,2048x2048
• validation_guidance - Usa el valor que suelas seleccionar en inferencia para OmniGen; un valor más bajo alrededor de 2.5-3.0 da resultados más realistas
• validation_num_inference_steps - Usa alrededor de 30

• use_ema - Configurar esto en true ayudará mucho a obtener un resultado más suavizado junto con tu checkpoint principal.

• optimizer - Puedes usar cualquier optimizador con el que te sientas cómodo, pero usaremos adamw_bf16 para este ejemplo.

• mixed_precision - Se recomienda configurarlo en bf16 para la configuración de entrenamiento más eficiente, o no (pero consumirá más memoria y será más lento).
• gradient_checkpointing - Desactivarlo será lo más rápido, pero limita el tamaño de los lotes. Es necesario activarlo para obtener el menor uso de VRAM.

Usuarios multi-GPU pueden consultar este documento para información sobre cómo configurar la cantidad de GPUs a usar.

Tu config.json se verá algo así al final:

{
    "validation_torch_compile": "false",
    "validation_step_interval": 200,
    "validation_seed": 42,
    "validation_resolution": "1024x1024",
    "validation_prompt": "A photo-realistic image of a cat",
    "validation_num_inference_steps": "20",
    "validation_guidance": 2.0,
    "validation_guidance_rescale": "0.0",
    "vae_cache_ondemand": true,
    "vae_batch_size": 1,
    "train_batch_size": 1,
    "tracker_run_name": "eval_loss_test1",
    "seed": 42,
    "resume_from_checkpoint": "latest",
    "resolution": 1024,
    "resolution_type": "pixel_area",
    "report_to": "tensorboard",
    "output_dir": "output/models-omnigen",
    "optimizer": "adamw_bf16",
    "num_train_epochs": 0,
    "num_eval_images": 1,
    "model_type": "lora",
    "model_family": "omnigen",
    "mixed_precision": "bf16",
    "minimum_image_size": 0,
    "max_train_steps": 10000,
    "max_grad_norm": 0.01,
    "lycoris_config": "config/lycoris_config.json",
    "lr_warmup_steps": 100,
    "lr_scheduler": "constant",
    "lora_type": "lycoris",
    "learning_rate": "4e-5",
    "gradient_checkpointing": "true",
    "grad_clip_method": "value",
    "eval_steps_interval": 100,
    "disable_benchmark": false,
    "data_backend_config": "config/omnigen/multidatabackend.json",
    "checkpoints_total_limit": 5,
    "checkpoint_step_interval": 500,
    "caption_dropout_probability": 0.0,
    "base_model_precision": "no_change",
    "aspect_bucket_rounding": 2
}

Y un archivo simple config/lycoris_config.json - ten en cuenta que se puede eliminar FeedForward para mayor estabilidad de entrenamiento.

{
    "algo": "lokr",
    "multiplier": 1.0,
    "linear_dim": 10000,
    "linear_alpha": 1,
    "factor": 16,
    "apply_preset": {
        "target_module": [
            "Attention",
            "FeedForward"
        ],
        "module_algo_map": {
            "Attention": {
                "factor": 16
            },
            "FeedForward": {
                "factor": 8
            }
        }
    }
}

Funciones experimentales avanzadas¶

{
  "nickname": "the prompt goes here",
  "another_nickname": "another prompt goes here"
}

  "--user_prompt_library": "config/user_prompt_library.json",

{
    "anime_<token>": "a breathtaking anime-style portrait of <token>, capturing her essence with vibrant colors and expressive features",
    "chef_<token>": "a high-quality, detailed photograph of <token> as a sous-chef, immersed in the art of culinary creation",
    "just_<token>": "a lifelike and intimate portrait of <token>, showcasing her unique personality and charm",
    "cinematic_<token>": "a cinematic, visually stunning photo of <token>, emphasizing her dramatic and captivating presence",
    "elegant_<token>": "an elegant and timeless portrait of <token>, exuding grace and sophistication",
    "adventurous_<token>": "a dynamic and adventurous photo of <token>, captured in an exciting, action-filled moment",
    "mysterious_<token>": "a mysterious and enigmatic portrait of <token>, shrouded in shadows and intrigue",
    "vintage_<token>": "a vintage-style portrait of <token>, evoking the charm and nostalgia of a bygone era",
    "artistic_<token>": "an artistic and abstract representation of <token>, blending creativity with visual storytelling",
    "futuristic_<token>": "a futuristic and cutting-edge portrayal of <token>, set against a backdrop of advanced technology",
    "woman": "a beautifully crafted portrait of a woman, highlighting her natural beauty and unique features",
    "man": "a powerful and striking portrait of a man, capturing his strength and character",
    "boy": "a playful and spirited portrait of a boy, capturing youthful energy and innocence",
    "girl": "a charming and vibrant portrait of a girl, emphasizing her bright personality and joy",
    "family": "a heartwarming and cohesive family portrait, showcasing the bonds and connections between loved ones"
}

Pérdida de evaluación estable¶

Si deseas usar pérdida MSE estable para puntuar el rendimiento del modelo, consulta este documento para información sobre cómo configurar e interpretar la pérdida de evaluación.

Vistas previas de validación¶

SimpleTuner admite la transmisión de vistas previas de validación intermedias durante la generación usando modelos Tiny AutoEncoder. Esto te permite ver las imágenes de validación generándose paso a paso en tiempo real mediante callbacks de webhook.

Para habilitarlo:

{
  "validation_preview": true,
  "validation_preview_steps": 1
}

Requisitos: - Configuración de webhook - Validación habilitada

Configura validation_preview_steps en un valor más alto (p. ej., 3 o 5) para reducir la sobrecarga del Tiny AutoEncoder. Con validation_num_inference_steps=20 y validation_preview_steps=5, recibirás imágenes de vista previa en los pasos 5, 10, 15 y 20.

Desplazamiento del calendario de flujo¶

Actualmente, OmniGen está hardcodeado para usar su propia formulación de flow-matching, y el schedule shift no se aplica.

Consideraciones del dataset¶

Es crucial contar con un dataset sustancial para entrenar tu modelo. Hay limitaciones en el tamaño del dataset, y debes asegurarte de que sea lo suficientemente grande para entrenar tu modelo de forma efectiva. Ten en cuenta que el tamaño mínimo del dataset es train_batch_size * gradient_accumulation_steps y también mayor que vae_batch_size. El dataset no será utilizable si es demasiado pequeño.

ℹ️ Con pocas imágenes, podrías ver el mensaje no images detected in dataset - aumentar el valor de repeats superará esta limitación.

Dependiendo del dataset que tengas, necesitarás configurar el directorio del dataset y el archivo de configuración del dataloader de manera diferente. En este ejemplo, usaremos pseudo-camera-10k como dataset.

Crea un documento --data_backend_config (config/multidatabackend.json) que contenga esto:

[
  {
    "id": "pseudo-camera-10k-omnigen",
    "type": "local",
    "crop": true,
    "crop_aspect": "square",
    "crop_style": "center",
    "resolution": 512,
    "minimum_image_size": 512,
    "maximum_image_size": 512,
    "target_downsample_size": 512,
    "resolution_type": "pixel_area",
    "cache_dir_vae": "cache/vae/omnigen/pseudo-camera-10k",
    "instance_data_dir": "datasets/pseudo-camera-10k",
    "disabled": false,
    "skip_file_discovery": "",
    "caption_strategy": "filename",
    "metadata_backend": "discovery",
    "repeats": 0,
    "is_regularisation_data": true
  },
  {
    "id": "dreambooth-subject",
    "type": "local",
    "crop": false,
    "resolution": 1024,
    "minimum_image_size": 1024,
    "maximum_image_size": 1024,
    "target_downsample_size": 1024,
    "resolution_type": "pixel_area",
    "cache_dir_vae": "cache/vae/omnigen/dreambooth-subject",
    "instance_data_dir": "datasets/dreambooth-subject",
    "caption_strategy": "instanceprompt",
    "instance_prompt": "the name of your subject goes here",
    "metadata_backend": "discovery",
    "repeats": 1000
  },
  {
    "id": "dreambooth-subject-512",
    "type": "local",
    "crop": false,
    "resolution": 512,
    "minimum_image_size": 512,
    "maximum_image_size": 512,
    "target_downsample_size": 512,
    "resolution_type": "pixel_area",
    "cache_dir_vae": "cache/vae/omnigen/dreambooth-subject-512",
    "instance_data_dir": "datasets/dreambooth-subject",
    "caption_strategy": "instanceprompt",
    "instance_prompt": "the name of your subject goes here",
    "metadata_backend": "discovery",
    "repeats": 1000
  },
  {
    "id": "text-embeds",
    "type": "local",
    "dataset_type": "text_embeds",
    "default": true,
    "cache_dir": "cache/text/omnigen",
    "disabled": false,
    "write_batch_size": 128
  }
]

Consulta las opciones y requisitos de caption_strategy en DATALOADER.md.

ℹ️ Se admite ejecutar datasets de 512px y 1024px de forma concurrente, y podría resultar en mejor convergencia.

ℹ️ OmniGen no genera embeddings de codificador de texto, pero aún se requiere definir uno (por ahora).

Mi configuración para OmniGen fue muy básica y se veía así, ya que usé entrenamiento con pérdida de evaluación estable:

[
    {
        "id": "something-special-to-remember-by",
        "type": "local",
        "instance_data_dir": "/datasets/pseudo-camera-10k/train",
        "minimum_image_size": 1024,
        "maximum_image_size": 1536,
        "target_downsample_size": 1024,
        "resolution": 1024,
        "resolution_type": "pixel_area",
        "caption_strategy": "filename",
        "cache_dir_vae": "cache/vae/omnigen",
        "vae_cache_clear_each_epoch": false,
        "crop": true,
        "crop_aspect": "square"
    },
    {
        "id": "omnigen-eval",
        "type": "local",
        "dataset_type": "eval",
        "crop": true,
        "crop_aspect": "square",
        "instance_data_dir": "/datasets/test_datasets/squares",
        "resolution": 1024,
        "minimum_image_size": 1024,
        "maximum_image_size": 1024,
        "target_downsample_size": 1024,
        "resolution_type": "pixel_area",
        "cache_dir_vae": "cache/vae/omnigen-eval",
        "caption_strategy": "filename"
    },
    {
        "id": "text-embed-cache",
        "dataset_type": "text_embeds",
        "default": true,
        "type": "local",
        "cache_dir": "cache/text/omnigen"
    }
]

Luego, crea un directorio datasets:

mkdir -p datasets
pushd datasets
    huggingface-cli download --repo-type=dataset bghira/pseudo-camera-10k --local-dir=pseudo-camera-10k
    mkdir dreambooth-subject
    # place your images into dreambooth-subject/ now
popd

Esto descargará alrededor de 10k muestras de fotografías a tu directorio datasets/pseudo-camera-10k, que se creará automáticamente.

Tus imágenes de Dreambooth deben ir en el directorio datasets/dreambooth-subject.

Iniciar sesión en WandB y Huggingface Hub¶

Querrás iniciar sesión en WandB y HF Hub antes de empezar el entrenamiento, especialmente si usas --push_to_hub y --report_to=wandb.

Si vas a subir elementos a un repositorio Git LFS manualmente, también deberías ejecutar git config --global credential.helper store

Ejecuta los siguientes comandos:

wandb login

y

huggingface-cli login

Sigue las instrucciones para iniciar sesión en ambos servicios.

Ejecutar el entrenamiento¶

Desde el directorio de SimpleTuner, tienes varias opciones para iniciar el entrenamiento:

Opción 1 (Recomendada - instalación con pip):

pip install 'simpletuner[cuda]'

# CUDA 13 / Blackwell users (NVIDIA B-series GPUs)
pip install 'simpletuner[cuda13]' --extra-index-url https://download.pytorch.org/whl/cu130
simpletuner train

Opción 2 (Método de git clone):

simpletuner train

Opción 3 (Método heredado - aún funciona):

./train.sh

Esto iniciará el caché a disco de text embeds y salidas del VAE.

Para más información, consulta los documentos de dataloader y tutorial.

Notas y consejos de solución de problemas¶

Configuración de VRAM más baja¶

La configuración de VRAM más baja de OmniGen aún no se conoce, pero se espera que sea similar a lo siguiente:

• OS: Ubuntu Linux 24
• GPU: Un solo dispositivo NVIDIA CUDA (10G, 12G)
• Memoria del sistema: alrededor de 50G de memoria del sistema
• Precisión del modelo base: int8-quanto (o fp8-torchao, int8-torchao siguen perfiles de uso de memoria similares)
• Optimizador: Lion 8Bit Paged, bnb-lion8bit-paged
• Resolución: 1024px
• Tamaño de lote: 1, cero pasos de acumulación de gradiente
• DeepSpeed: deshabilitado / sin configurar
• PyTorch: 2.7+
• Usar --quantize_via=cpu para evitar error outOfMemory durante el arranque en tarjetas <=16G.
• Habilitar --gradient_checkpointing
• Usar una configuración LoRA o Lycoris pequeña (p. ej., rango LoRA 1 o factor Lokr 25)

NOTA: El pre-caché de embeddings del VAE y salidas del codificador de texto puede usar más memoria y aun así producir OOM. Si sucede, se pueden habilitar tiling y slicing del VAE. Los codificadores de texto pueden descargarse a CPU durante el caché del VAE con offload_during_startup=true. Para evitar uso de disco del caché VAE en datasets grandes, usa --vae_cache_disable.

La velocidad fue de aproximadamente 3.4 iteraciones por segundo en una AMD 7900XTX usando Pytorch 2.7 y ROCm 6.3.

Pérdida con máscara¶

Si estás entrenando un sujeto o estilo y te gustaría enmascarar uno u otro, consulta la sección de entrenamiento con pérdida enmascarada de la guía de Dreambooth.

Cuantización¶

No probado a fondo (aún).

Tasas de aprendizaje¶

LoRA (--lora_type=standard)¶

No compatible.

LoKr (--lora_type=lycoris)¶

• Tasas de aprendizaje suaves son mejores para LoKr (1e-4 con AdamW, 2e-5 con Lion)
• Otros algoritmos necesitan más exploración.
• Configurar is_regularisation_data tiene impacto/efecto desconocido con OmniGen (no probado)

Artefactos de imagen¶

OmniGen tiene una respuesta desconocida a artefactos de imagen, aunque usa el VAE de SDXL y tiene limitaciones idénticas de detalle fino.

Si surge algún problema de calidad de imagen, por favor abre un issue en Github.

Aspect bucketing¶

Este modelo tiene una respuesta desconocida a datos bucketed por aspecto. La experimentación será útil.

Valores de pérdida altos¶

OmniGen tiene un valor de pérdida muy alto, y no se sabe por qué. Se recomienda ignorar el valor de pérdida y enfocarse en la calidad visual de las imágenes generadas.