OmniGen Documentation

LoRA (Low-Rank Adaptation) is a memory-efficient way to fine-tune OmniGen. To use it, add the --use_lora and --lora_rank flags to your training command.

Additional Arguments for LoRA:

• --use_lora: Enables LoRA training.
• --lora_rank: The rank for LoRA adaptation.
• --condition_dropout_prob: Dropout probability for conditions.
• --max_input_length_limit: Limit for input length.

Example Command:

accelerate launch \
    --num_processes=1 \
    train.py \
    --model_name_or_path Shitao/OmniGen-v1 \
    --batch_size_per_device 2 \
    --condition_dropout_prob 0.01 \
    --lr 3e-4 \
    --use_lora \
    --lora_rank 8 \
    --json_file ./toy_data/toy_data.jsonl \
    --image_path ./toy_data/images \
    --max_input_length_limit 18000 \
    --keep_raw_resolution \
    --max_image_size 1024 \
    --gradient_accumulation_steps 1 \
    --ckpt_every 50 \
    --epochs 100 \
    --log_every 1 \
    --results_dir ./results/toy_finetune_lora

Run full fine-tuning using accelerate launch and train.py. This method updates all model parameters.

Key Arguments for train.py:

• --model_name_or_path: Path to the pretrained model (e.g., Shitao/OmniGen-v1).
• --json_file: Path to the training data .jsonl file.
• --image_path: Path to the folder containing training images.
• --batch_size_per_device: Batch size per GPU.
• --lr: Learning rate.
• --keep_raw_resolution: Boolean flag to keep original image resolution; otherwise, images are resized to max_image_size.
• --max_image_size: Maximum image dimension.
• --gradient_accumulation_steps: Number of steps to accumulate gradients.
• --ckpt_every: Number of steps between saving checkpoints.
• --epochs: Total number of training epochs.
• --results_dir: Directory where results and checkpoints will be saved.

Troubleshooting OOM (Out of Memory): If you encounter OOM errors, decrease --batch_size_per_device or --max_image_size, or switch to LoRA fine-tuning.

accelerate launch \
    --num_processes=1 \
    --use_fsdp \
    --fsdp_offload_params false \
    --fsdp_sharding_strategy SHARD_GRAD_OP \
    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
    --fsdp_transformer_layer_cls_to_wrap Phi3DecoderLayer \
    --fsdp_state_dict_type FULL_STATE_DICT \
    --fsdp_forward_prefetch false \
    --fsdp_use_orig_params True \
    --fsdp_cpu_ram_efficient_loading false \
    --fsdp_sync_module_states True \
    train.py \
    --model_name_or_path Shitao/OmniGen-v1 \
    --json_file ./toy_data/toy_data.jsonl \
    --image_path ./toy_data/images \
    --batch_size_per_device 1 \
    --lr 2e-5 \
    --keep_raw_resolution \
    --max_image_size 1024 \
    --gradient_accumulation_steps 1 \
    --ckpt_every 50 \
    --epochs 200 \
    --log_every 1 \
    --results_dir ./results/toy_finetune

OmniGen can generate new images based on specific objects or identities from input images without requiring external detection/cropping models.

Key Features:

• Multi-object identification: Use descriptive instructions in the prompt to identify specific objects in multiple input images using the placeholder format <img><|image_N|></img> (e.g., the man in <img><|image_1|></img>).
• Object extraction: Extract multiple objects from different images to combine them into a single new scene.

Usage: Pass a list of image paths to input_images and use the <img><|image_N|></img> syntax in your prompt to reference them.

prompt = "Two woman are raising fried chicken legs in a bar. A woman is <img><|image_1|></img>. Another woman is <img><|image_2|></img>."
input_images = ["./imgs/test_cases/mckenna.jpg", "./imgs/test_cases/Amanda.jpg"]
images = pipe(
    prompt=prompt, 
    input_images=input_images, 
    height=1024, 
    width=1024,
    guidance_scale=2.5, 
    img_guidance_scale=1.8, 
    max_input_image_size=1024,
    seed=168
)
images[0].show()

OmniGen can perform classical Computer Vision (CV) tasks (like depth map or skeleton detection) and image-conditional generation within a single model.

Tasks include:

• Condition Extraction: Prompt the model to detect features like "Detect the skeleton of human in this image: <img><|image_1|></img>."
• One-step Generation: Generate new images following specific conditions (pose, depth, etc.) directly by using phrases like "Following the human pose of this image <img><|image_1|></img>, generate a new photo: ...".

prompt = "Following the human pose of this image <img><|image_1|></img>, generate a new photo: An elderly man wearing a gold-framed glasses..."
input_images = ["./imgs/test_cases/control.jpg"]
images = pipe(
    prompt=prompt, 
    input_images=input_images, 
    height=1024, 
    width=1024,
    guidance_scale=2.5, 
    img_guidance_scale=1.6,
    seed=0
)
images[0].show()