Shitao/OmniGen-v1

Instalación y uso con Diffusers:

pip install -U diffusers transformers accelerate

import torch
from diffusers import DiffusionPipeline

# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("Shitao/OmniGen-v1", dtype=torch.bfloat16, device_map="cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

Instalación desde GitHub:

git clone https://github.com/staoxiao/OmniGen.git
cd OmniGen
pip install -e .

Ejemplo de texto a imagen con OmniGen:

from OmniGen import OmniGenPipeline

pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")

images = pipe(
    prompt="A curly-haired man in a red shirt is drinking tea.",
    height=1024,
    width=1024,
    guidance_scale=2.5,
    seed=0,
)
images[0].save("example_t2i.png")

Ejemplo multimodal a imagen:

images = pipe(
    prompt="A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
    input_images=["./imgs/test_cases/two_man.jpg"],
    height=1024,
    width=1024,
    guidance_scale=2.5,
    img_guidance_scale=1.6,
    seed=0
)
images[0].save("example_ti2i.png")

Demo local con Gradio:

pip install gradio spaces
python app.py

Uso en Google Colab:

!git clone https://github.com/staoxiao/OmniGen.git
%cd OmniGen
!pip install -e .
!pip install gradio spaces
!python app.py --share

Ejemplo de fine-tuning LoRA:

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 1e-3 \
  --use_lora \
  --lora_rank 8 \
  --json_file ./toy_data/toy_subject_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 10 \
  --epochs 200 \
  --log_every 1 \
  --results_dir ./results/toy_finetune_lora