Upload 81 files

.gitattributes

@@ -33,3 +33,27 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/edit.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/entity.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/reasoning.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/same_pose.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/skeletal2img.png filter=lfs diff=lfs merge=lfs -text
+imgs/demo_cases/t2i_woman_with_book.png filter=lfs diff=lfs merge=lfs -text
+imgs/overall.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/referring.png filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/1.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/2.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/3.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/4.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/Amanda.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/control.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/mckenna.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/two_man.jpg filter=lfs diff=lfs merge=lfs -text
+imgs/test_cases/woman.png filter=lfs diff=lfs merge=lfs -text
+toy_data/images/cat.png filter=lfs diff=lfs merge=lfs -text
+toy_data/images/dog2.jpeg filter=lfs diff=lfs merge=lfs -text
+toy_data/images/dog3.jpeg filter=lfs diff=lfs merge=lfs -text
+toy_data/images/dog4.jpeg filter=lfs diff=lfs merge=lfs -text
+toy_data/images/dog5.jpeg filter=lfs diff=lfs merge=lfs -text
+toy_data/images/walking.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,170 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# myfile
+.results/
+local.ipynb
+convert_to_safetensor.py
+ttt.ipynb
+imgs/ttt/
+*.bak

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 VectorSpaceLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

OmniGen/__init__.py

@@ -0,0 +1,4 @@
+from .model import OmniGen
+from .processor import OmniGenProcessor
+from .scheduler import OmniGenScheduler
+from .pipeline import OmniGenPipeline

OmniGen/model.py

@@ -0,0 +1,468 @@
+# The code is revised from DiT
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from typing import Dict
+import torch.nn.functional as F
+
+from diffusers.loaders import PeftAdapterMixin
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+from OmniGen.transformer import Phi3Config, Phi3Transformer
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t, dtype=torch.float32):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    if isinstance(grid_size, int):
+        grid_size = (grid_size, grid_size)
+
+    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
+    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class PatchEmbedMR(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+    def __init__(
+            self,
+            patch_size: int = 2,
+            in_chans: int = 4,
+            embed_dim: int = 768,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
+        return x
+
+
+class Int8Quantized(nn.Module):
+    def __init__(self, tensor, scale_factor=None):
+        super().__init__()
+        if scale_factor is None:
+            max_val = torch.max(torch.abs(tensor))
+            scale_factor = max_val / 127.0
+        # Store quantized weights and scale factor
+        self.register_buffer('quantized_weight', torch.round(tensor / scale_factor).to(torch.int8))
+        self.register_buffer('scale_factor', torch.tensor(scale_factor))
+
+    def forward(self, dtype=None):
+        # Dequantize and convert to specified dtype
+        weight = self.quantized_weight.float() * self.scale_factor
+        if dtype is not None:
+            weight = weight.to(dtype)
+        return weight
+
+
+
+class QuantizedLinear(nn.Module):
+    def __init__(self, weight, bias=None):
+        super().__init__()
+        self.weight_quantized = Int8Quantized(weight)
+        if bias is not None:
+            self.register_buffer('bias', bias)
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        # Dequantize weight to match input dtype
+        weight = self.weight_quantized(dtype=x.dtype)
+        return F.linear(x, weight, self.bias)
+
+
+class OmniGen(nn.Module, PeftAdapterMixin):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        transformer_config: Phi3Config,
+        patch_size=2,
+        in_channels=4,
+        pe_interpolation: float = 1.0,
+        pos_embed_max_size: int = 192,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.pos_embed_max_size = pos_embed_max_size
+
+        hidden_size = transformer_config.hidden_size
+
+        self.x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
+        self.input_x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
+
+        self.time_token = TimestepEmbedder(hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+
+        self.pe_interpolation = pe_interpolation
+        pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, interpolation_scale=self.pe_interpolation, base_size=64)
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=True)
+
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+
+        self.initialize_weights()
+
+        self.llm = Phi3Transformer(config=transformer_config)
+        self.llm.config.use_cache = False
+
+    def _quantize_module(self, module):
+        """
+        Quantize a module to 8-bit precision
+        """
+        for name, child in module.named_children():
+            if isinstance(child, nn.Linear):
+                setattr(module, name, QuantizedLinear(child.weight.data, child.bias.data if child.bias is not None else None))
+            elif isinstance(child, nn.LayerNorm):
+                # Skip quantization for LayerNorm
+                continue
+            else:
+                self._quantize_module(child)
+
+    @classmethod
+    def from_pretrained(cls, model_name, quantize=False):  # Add quantize parameter
+        if not os.path.exists(model_name):
+            cache_folder = os.getenv('HF_HUB_CACHE')
+            model_name = snapshot_download(repo_id=model_name,
+                                           cache_dir=cache_folder,
+                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
+        config = Phi3Config.from_pretrained(model_name)
+        model = cls(config)
+        if os.path.exists(os.path.join(model_name, 'model.safetensors')):
+            print("Loading safetensors")
+            ckpt = load_file(os.path.join(model_name, 'model.safetensors'))
+        else:
+            ckpt = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
+
+        # Load weights first
+        model.load_state_dict(ckpt)
+
+        # Only quantize if explicitly requested
+        if quantize:
+            print("Quantizing weights to 8-bit...")
+            model._quantize_module(model.llm)
+
+        return model
+    def initialize_weights(self):
+        assert not hasattr(self, "llama")
+
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+        
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        w = self.input_x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        nn.init.normal_(self.time_token.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.time_token.mlp[2].weight, std=0.02)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x, h, w):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+
+        x = x.reshape(shape=(x.shape[0], h//self.patch_size, w//self.patch_size, self.patch_size, self.patch_size, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h, w))
+        return imgs
+
+
+    def cropped_pos_embed(self, height, width):
+        """Crops positional embeddings for SD3 compatibility."""
+        if self.pos_embed_max_size is None:
+            raise ValueError("`pos_embed_max_size` must be set for cropping.")
+
+        height = height // self.patch_size
+        width = width // self.patch_size
+        if height > self.pos_embed_max_size:
+            raise ValueError(
+                f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
+            )
+        if width > self.pos_embed_max_size:
+            raise ValueError(
+                f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
+            )
+
+        top = (self.pos_embed_max_size - height) // 2
+        left = (self.pos_embed_max_size - width) // 2
+        spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
+        spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
+        # print(top, top + height, left, left + width, spatial_pos_embed.size())
+        spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
+        return spatial_pos_embed
+
+
+    def patch_multiple_resolutions(self, latents, padding_latent=None, is_input_images:bool=False):
+        if isinstance(latents, list):
+            return_list = False
+            if padding_latent is None:
+                padding_latent = [None] * len(latents)
+                return_list = True
+            patched_latents, num_tokens, shapes = [], [], []
+            for latent, padding in zip(latents, padding_latent):
+                height, width = latent.shape[-2:]
+                if is_input_images:
+                    latent = self.input_x_embedder(latent)
+                else:
+                    latent = self.x_embedder(latent)
+                pos_embed = self.cropped_pos_embed(height, width)    
+                latent = latent + pos_embed
+                if padding is not None:
+                    latent = torch.cat([latent, padding], dim=-2)
+                patched_latents.append(latent)
+
+                num_tokens.append(pos_embed.size(1))
+                shapes.append([height, width])
+            if not return_list:
+                latents = torch.cat(patched_latents, dim=0)
+            else:
+                latents = patched_latents
+        else:
+            height, width = latents.shape[-2:]
+            if is_input_images:
+                latents = self.input_x_embedder(latents)
+            else:
+                latents = self.x_embedder(latents)
+            pos_embed = self.cropped_pos_embed(height, width)  
+            latents = latents + pos_embed
+            num_tokens = latents.size(1)
+            shapes = [height, width]
+        return latents, num_tokens, shapes
+
+    
+    def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True):
+        """
+        
+        """
+        input_is_list = isinstance(x, list)
+        x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent)
+        time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1)   
+        
+        if input_img_latents is not None:
+            input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True)
+        if input_ids is not None:
+            condition_embeds = self.llm.embed_tokens(input_ids).clone()
+            input_img_inx = 0
+            for b_inx in input_image_sizes.keys():
+                for start_inx, end_inx in input_image_sizes[b_inx]:
+                    condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx]
+                    input_img_inx += 1
+            if input_img_latents is not None:
+                assert input_img_inx == len(input_latents) 
+
+            input_emb = torch.cat([condition_embeds, time_token, x], dim=1)
+        else:
+            input_emb = torch.cat([time_token, x], dim=1)
+        output = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values)
+        output, past_key_values = output.last_hidden_state, output.past_key_values
+        if input_is_list:
+            image_embedding = output[:, -max(num_tokens):]
+            time_emb = self.t_embedder(timestep, dtype=x.dtype)
+            x = self.final_layer(image_embedding, time_emb)
+            latents = []
+            for i in range(x.size(0)):
+                latent = x[i:i+1, :num_tokens[i]]
+                latent = self.unpatchify(latent, shapes[i][0], shapes[i][1])
+                latents.append(latent)
+        else:
+            image_embedding = output[:, -num_tokens:]
+            time_emb = self.t_embedder(timestep, dtype=x.dtype)
+            x = self.final_layer(image_embedding, time_emb)
+            latents = self.unpatchify(x, shapes[0], shapes[1])
+
+        if return_past_key_values:
+            return latents, past_key_values
+        return latents
+
+    @torch.no_grad()
+    def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache):
+        """
+        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """        
+        self.llm.config.use_cache = use_kv_cache
+        model_out, past_key_values = self.forward(x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, past_key_values=past_key_values, return_past_key_values=True)
+        if use_img_cfg:
+            cond, uncond, img_cond = torch.split(model_out, len(model_out) // 3, dim=0)
+            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
+            model_out = [cond, cond, cond]
+        else:
+            cond, uncond = torch.split(model_out, len(model_out) // 2, dim=0)
+            cond = uncond + cfg_scale * (cond - uncond)
+            model_out = [cond, cond]
+        
+        return torch.cat(model_out, dim=0), past_key_values
+
+
+    @torch.no_grad()
+    def forward_with_separate_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, return_past_key_values=True):
+        """
+        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """        
+        self.llm.config.use_cache = use_kv_cache
+        if past_key_values is None:
+            past_key_values = [None] * len(attention_mask)
+
+        x = torch.split(x, len(x) // len(attention_mask), dim=0)
+        timestep = timestep.to(x[0].dtype)
+        timestep = torch.split(timestep, len(timestep) // len(input_ids), dim=0)
+
+        model_out, pask_key_values = [], []
+        for i in range(len(input_ids)):
+            temp_out, temp_pask_key_values = self.forward(x[i], timestep[i], input_ids[i], input_img_latents[i], input_image_sizes[i], attention_mask[i], position_ids[i], past_key_values[i])
+            model_out.append(temp_out)
+            pask_key_values.append(temp_pask_key_values)
+
+        if len(model_out) == 3:
+            cond, uncond, img_cond = model_out
+            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
+            model_out = [cond, cond, cond]
+        elif len(model_out) == 2:
+            cond, uncond = model_out
+            cond = uncond + cfg_scale * (cond - uncond)
+            model_out = [cond, cond]
+        else:
+            return model_out[0]
+        
+        return torch.cat(model_out, dim=0), pask_key_values
+
+
+
+

OmniGen/pipeline.py

@@ -0,0 +1,289 @@
+import os
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from PIL import Image
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from peft import LoraConfig, PeftModel
+from diffusers.models import AutoencoderKL
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from safetensors.torch import load_file
+
+from OmniGen import OmniGen, OmniGenProcessor, OmniGenScheduler
+
+import gc  # For clearing unused objects
+
+logger = logging.get_logger(__name__)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from OmniGen import OmniGenPipeline
+        >>> pipe = FluxControlNetPipeline.from_pretrained(
+        ...     base_model
+        ... )
+        >>> prompt = "A woman holds a bouquet of flowers and faces the camera"
+        >>> image = pipe(
+        ...     prompt,
+        ...     guidance_scale=3.0,
+        ...     num_inference_steps=50,
+        ... ).images[0]
+        >>> image.save("t2i.png")
+        ```
+"""
+
+class OmniGenPipeline:
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        model: OmniGen,
+        processor: OmniGenProcessor,
+
+    ):
+        self.vae = vae
+        self.model = model
+        self.processor = processor
+
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.model.to(self.device)
+        self.model.eval()
+        self.vae.to(self.device)
+
+    @classmethod
+    def from_pretrained(cls, model_name, vae_path: str=None, Quantization: bool=False):
+        if not os.path.exists(model_name) or (not os.path.exists(os.path.join(model_name, 'model.safetensors')) and model_name == "Shitao/OmniGen-v1"):
+            logger.info("Model not found, downloading...")
+            cache_folder = os.getenv('HF_HUB_CACHE')
+            model_name = snapshot_download(repo_id=model_name,
+                                        cache_dir=cache_folder,
+                                        ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5', 'model.pt'])
+            logger.info(f"Downloaded model to {model_name}")
+
+        # Pass Quantization parameter to OmniGen's from_pretrained
+        model = OmniGen.from_pretrained(model_name, quantize=Quantization)
+
+        processor = OmniGenProcessor.from_pretrained(model_name)
+
+        if os.path.exists(os.path.join(model_name, "vae")):
+            vae = AutoencoderKL.from_pretrained(os.path.join(model_name, "vae"))
+        elif vae_path is not None:
+            vae = AutoencoderKL.from_pretrained(vae_path)
+        else:
+            logger.info(f"No VAE found in {model_name}, downloading stabilityai/sdxl-vae from HF")
+            vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
+
+        return cls(vae, model, processor)
+
+
+
+    def merge_lora(self, lora_path: str):
+        model = PeftModel.from_pretrained(self.model, lora_path)
+        model.merge_and_unload()
+
+        self.model = model
+
+    def to(self, device: Union[str, torch.device]):
+        if isinstance(device, str):
+            device = torch.device(device)
+        self.model.to(device)
+        self.vae.to(device)
+
+    def vae_encode(self, x, dtype):
+        if self.vae.config.shift_factor is not None:
+            x = self.vae.encode(x).latent_dist.sample()
+            x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        else:
+            x = self.vae.encode(x).latent_dist.sample().mul_(self.vae.config.scaling_factor)
+        x = x.to(dtype)
+        return x
+
+    def move_to_device(self, data):
+        if isinstance(data, list):
+            return [x.to(self.device) for x in data]
+        return data.to(self.device)
+
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        input_images: Union[List[str], List[List[str]]] = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 3,
+        use_img_guidance: bool = True,
+        img_guidance_scale: float = 1.6,
+        separate_cfg_infer: bool = False,
+        use_kv_cache: bool = True,
+        dtype: torch.dtype = torch.bfloat16,
+        seed: int = None,
+        Quantization: bool = False,
+        ):
+
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide the image generation.
+            input_images (`List[str]` or `List[List[str]]`, *optional*):
+                The list of input images. We will replace the "<|image_i|>" in prompt with the 1-th image in list.
+            height (`int`, *optional*, defaults to 1024):
+                The height in pixels of the generated image. The number must be a multiple of 16.
+            width (`int`, *optional*, defaults to 1024):
+                The width in pixels of the generated image. The number must be a multiple of 16.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 4.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            use_img_guidance (`bool`, *optional*, defaults to True):
+                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800).
+            img_guidance_scale (`float`, *optional*, defaults to 1.6):
+                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800).
+            separate_cfg_infer (`bool`, *optional*, defaults to False):
+                Perform inference on images with different guidance separately; this can save memory when generating images of large size at the expense of slower inference.
+            use_kv_cache (`bool`, *optional*, defaults to True): enable kv cache to speed up the inference
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+        Examples:
+
+        Returns:
+            A list with the generated images.
+        """
+
+
+        assert height%16 == 0 and width%16 == 0
+        if separate_cfg_infer:
+            use_kv_cache = False
+            # raise "Currently, don't support both use_kv_cache and separate_cfg_infer"
+        if input_images is None:
+            use_img_guidance = False
+        if isinstance(prompt, str):
+            prompt = [prompt]
+            input_images = [input_images] if input_images is not None else None
+
+
+        input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=separate_cfg_infer)
+
+        num_prompt = len(prompt)
+        num_cfg = 2 if use_img_guidance else 1
+        latent_size_h, latent_size_w = height // 8, width // 8
+
+        if seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+        else:
+            generator = None
+        latents = torch.randn(num_prompt, 4, latent_size_h, latent_size_w, device=self.device, generator=generator)
+        latents = torch.cat([latents] * (1 + num_cfg), 0).to(dtype)
+
+
+        # Load VAE into VRAM (GPU) in bfloat16
+        self.vae.to(self.device, dtype=torch.bfloat16)
+
+
+
+
+        input_img_latents = []
+        if separate_cfg_infer:
+            for temp_pixel_values in input_data['input_pixel_values']:
+                temp_input_latents = []
+                for img in temp_pixel_values:
+                    img = self.vae_encode(img.to(self.device, dtype=torch.bfloat16), dtype)
+
+                    temp_input_latents.append(img)
+                input_img_latents.append(temp_input_latents)
+        else:
+            for img in input_data['input_pixel_values']:
+                img = self.vae_encode(img.to(self.device, dtype=torch.bfloat16), dtype)
+
+                input_img_latents.append(img)
+
+
+
+        model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']),
+            input_img_latents=input_img_latents,
+            input_image_sizes=input_data['input_image_sizes'],
+            attention_mask=self.move_to_device(input_data["attention_mask"]),
+            position_ids=self.move_to_device(input_data["position_ids"]),
+            cfg_scale=guidance_scale,
+            img_cfg_scale=img_guidance_scale,
+            use_img_cfg=use_img_guidance,
+            use_kv_cache=use_kv_cache)
+
+
+        #unlode vae to cpu
+        self.vae.to('cpu')
+        torch.cuda.empty_cache()  # Clear VRAM
+        gc.collect()  # Run garbage collection to free system RAM
+
+
+
+        if separate_cfg_infer:
+            func = self.model.forward_with_separate_cfg
+        else:
+            func = self.model.forward_with_cfg
+
+
+        #move main model to gpu
+        self.model.to(self.device, dtype=dtype)
+
+
+        scheduler = OmniGenScheduler(num_steps=num_inference_steps)
+        samples = scheduler(latents, func, model_kwargs, use_kv_cache=use_kv_cache)
+        samples = samples.chunk((1 + num_cfg), dim=0)[0]
+
+        if self.vae.config.shift_factor is not None:
+            samples = samples / self.vae.config.scaling_factor + self.vae.config.shift_factor
+        else:
+            samples = samples / self.vae.config.scaling_factor
+
+        #unlode main model to cpu
+        self.model.to('cpu')
+        torch.cuda.empty_cache()  # Clear VRAM
+        gc.collect()  # Run garbage collection to free system RAM
+
+        # Move samples to GPU and ensure they are in bfloat16 (for the VAE)
+        samples = samples.to(self.device, dtype=torch.bfloat16)
+
+        # Load VAE into VRAM (GPU) in bfloat16
+        self.vae.to(self.device, dtype=torch.bfloat16)
+
+        # Decode the samples using the VAE
+        samples = self.vae.decode(samples).sample
+
+        #unlode vae to cpu
+        self.vae.to('cpu')
+        torch.cuda.empty_cache()  # Clear VRAM
+        gc.collect()  # Run garbage collection to free system RAM
+
+
+        # Convert samples back to float32 for further processing
+        samples = samples.to(torch.float32)
+
+
+        # Convert samples to uint8 for final image output
+        output_samples = (samples * 0.5 + 0.5).clamp(0, 1) * 255
+        output_samples = output_samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+        # Create output images
+        output_images = []
+        for i, sample in enumerate(output_samples):
+            output_images.append(Image.fromarray(sample))
+
+        # Return the generated images
+        return output_images

OmniGen/processor.py

@@ -0,0 +1,335 @@
+import os
+import re
+from typing import Dict, List
+import json
+
+import torch
+import numpy as np
+import random
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoTokenizer
+from huggingface_hub import snapshot_download
+
+from OmniGen.utils import (
+    create_logger,
+    update_ema,
+    requires_grad,
+    center_crop_arr,
+    crop_arr,
+)
+
+
+
+
+class OmniGenProcessor:
+    def __init__(self, 
+                text_tokenizer, 
+                max_image_size: int=1024):
+        self.text_tokenizer = text_tokenizer
+        self.max_image_size = max_image_size
+
+        self.image_transform = transforms.Compose([
+            transforms.Lambda(lambda pil_image: crop_arr(pil_image, max_image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+        ])
+
+        self.collator = OmniGenCollator()
+        self.separate_collator = OmniGenSeparateCollator()
+
+    @classmethod
+    def from_pretrained(cls, model_name):
+        if not os.path.exists(model_name):
+            cache_folder = os.getenv('HF_HUB_CACHE')
+            model_name = snapshot_download(repo_id=model_name,
+                                           cache_dir=cache_folder,
+                                           allow_patterns="*.json")
+        text_tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+        return cls(text_tokenizer)
+
+
+    def process_image(self, image):
+        image = Image.open(image).convert('RGB')
+        return self.image_transform(image)
+    
+    def process_multi_modal_prompt(self, text, input_images):
+        text = self.add_prefix_instruction(text)
+        if input_images is None or len(input_images) == 0:
+            model_inputs = self.text_tokenizer(text)
+            return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None}
+
+        pattern = r"<\|image_\d+\|>"
+        prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)] 
+
+        for i in range(1, len(prompt_chunks)):
+            if prompt_chunks[i][0] == 1:
+                prompt_chunks[i] = prompt_chunks[i][1:]
+
+        image_tags = re.findall(pattern, text) 
+        image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
+
+        unique_image_ids = sorted(list(set(image_ids)))
+        assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}"
+        # total images must be the same as the number of image tags
+        assert len(unique_image_ids) == len(input_images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images"
+        
+        input_images = [input_images[x-1] for x in image_ids]
+
+        all_input_ids = []
+        img_inx = []
+        idx = 0
+        for i in range(len(prompt_chunks)):
+            all_input_ids.extend(prompt_chunks[i])
+            if i != len(prompt_chunks) -1:
+                start_inx = len(all_input_ids)
+                size = input_images[i].size(-2) *  input_images[i].size(-1) // 16 // 16
+                img_inx.append([start_inx, start_inx+size])
+                all_input_ids.extend([0]*size)
+
+        return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx}
+
+
+    def add_prefix_instruction(self, prompt):
+        user_prompt = '<|user|>\n'
+        generation_prompt = 'Generate an image according to the following instructions\n'
+        assistant_prompt = '<|assistant|>\n<|diffusion|>'
+        prompt_suffix = "<|end|>\n"
+        prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}"
+        return prompt
+
+
+    def __call__(self, 
+                instructions: List[str], 
+                input_images: List[List[str]] = None,
+                height: int = 1024,
+                width: int = 1024,
+                negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.",
+                use_img_cfg: bool = True,
+                separate_cfg_input: bool = False,
+                ) -> Dict:
+
+        if input_images is None:
+            use_img_cfg = False
+        if isinstance(instructions, str):
+            instructions = [instructions]
+            input_images = [input_images]
+        
+        input_data = []
+        for i in range(len(instructions)):
+            cur_instruction = instructions[i]
+            cur_input_images = None if input_images is None else input_images[i]
+            if cur_input_images is not None and len(cur_input_images) > 0:
+                cur_input_images = [self.process_image(x) for x in cur_input_images]
+            else:
+                cur_input_images = None
+                assert "<img><|image_1|></img>" not in cur_instruction
+            
+            mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images)
+
+        
+            neg_mllm_input, img_cfg_mllm_input = None, None
+            neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None)
+            if use_img_cfg:
+                if cur_input_images is not None and len(cur_input_images) >= 1:
+                    img_cfg_prompt = [f"<img><|image_{i+1}|></img>" for i in range(len(cur_input_images))]
+                    img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images)
+                else:
+                    img_cfg_mllm_input = neg_mllm_input
+
+            input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
+
+        if separate_cfg_input:
+            return self.separate_collator(input_data)
+        return self.collator(input_data)
+
+
+
+
+class OmniGenCollator:
+    def __init__(self, pad_token_id=2, hidden_size=3072):
+        self.pad_token_id = pad_token_id
+        self.hidden_size = hidden_size
+    
+    def create_position(self, attention_mask, num_tokens_for_output_images):
+        position_ids = []
+        text_length = attention_mask.size(-1)
+        img_length = max(num_tokens_for_output_images)  
+        for mask in attention_mask:
+            temp_l = torch.sum(mask)
+            temp_position = [0]*(text_length-temp_l) + [i for i in range(temp_l+img_length+1)] # we add a time embedding into the sequence, so add one more token
+            position_ids.append(temp_position)
+        return torch.LongTensor(position_ids)
+    
+    def create_mask(self, attention_mask, num_tokens_for_output_images):
+        extended_mask = []
+        padding_images = []
+        text_length = attention_mask.size(-1)
+        img_length = max(num_tokens_for_output_images)
+        seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token
+        inx = 0
+        for mask in attention_mask:
+            temp_l = torch.sum(mask)
+            pad_l = text_length - temp_l
+
+            temp_mask = torch.tril(torch.ones(size=(temp_l+1, temp_l+1)))
+
+            image_mask = torch.zeros(size=(temp_l+1, img_length))
+            temp_mask = torch.cat([temp_mask, image_mask], dim=-1)
+
+            image_mask = torch.ones(size=(img_length, temp_l+img_length+1))
+            temp_mask = torch.cat([temp_mask, image_mask], dim=0)
+
+            if pad_l > 0:
+                pad_mask = torch.zeros(size=(temp_l+1+img_length, pad_l))
+                temp_mask = torch.cat([pad_mask, temp_mask], dim=-1)
+
+                pad_mask = torch.ones(size=(pad_l, seq_len))
+                temp_mask = torch.cat([pad_mask, temp_mask], dim=0)
+
+            true_img_length = num_tokens_for_output_images[inx]
+            pad_img_length = img_length - true_img_length
+            if pad_img_length > 0:
+                temp_mask[:, -pad_img_length:] = 0
+                temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size))
+            else:
+                temp_padding_imgs = None
+            
+            extended_mask.append(temp_mask.unsqueeze(0))
+            padding_images.append(temp_padding_imgs)
+            inx += 1
+        return torch.cat(extended_mask, dim=0), padding_images
+    
+    def adjust_attention_for_input_images(self, attention_mask, image_sizes):
+        for b_inx in image_sizes.keys():
+            for start_inx, end_inx in image_sizes[b_inx]:
+                attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1
+
+        return attention_mask
+    
+    def pad_input_ids(self, input_ids, image_sizes):
+        max_l = max([len(x) for x in input_ids])
+        padded_ids = []
+        attention_mask = []
+        new_image_sizes = []
+
+        for i in range(len(input_ids)):
+            temp_ids = input_ids[i]
+            temp_l = len(temp_ids)
+            pad_l = max_l - temp_l
+            if pad_l == 0:
+                attention_mask.append([1]*max_l)
+                padded_ids.append(temp_ids)
+            else:
+                attention_mask.append([0]*pad_l+[1]*temp_l)
+                padded_ids.append([self.pad_token_id]*pad_l+temp_ids)
+            
+            if i in image_sizes:
+                new_inx = []
+                for old_inx in image_sizes[i]:
+                    new_inx.append([x+pad_l for x in old_inx])
+                image_sizes[i] = new_inx
+
+        return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes
+
+
+    def process_mllm_input(self, mllm_inputs, target_img_size):
+        num_tokens_for_output_images = []
+        for img_size in target_img_size:
+            num_tokens_for_output_images.append(img_size[0]*img_size[1]//16//16)
+
+        pixel_values, image_sizes = [], {}
+        b_inx = 0
+        for x in mllm_inputs:
+            if x['pixel_values'] is not None:
+                pixel_values.extend(x['pixel_values'])
+                for size in x['image_sizes']:
+                    if b_inx not in image_sizes:
+                        image_sizes[b_inx] = [size]
+                    else:
+                        image_sizes[b_inx].append(size)
+            b_inx += 1     
+        pixel_values = [x.unsqueeze(0) for x in pixel_values]
+
+        
+        input_ids = [x['input_ids'] for x in mllm_inputs]
+        padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes)
+        position_ids = self.create_position(attention_mask, num_tokens_for_output_images)
+        attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images)
+        attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes)
+
+        return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes
+    
+    
+    def __call__(self, features):
+        mllm_inputs = [f[0] for f in features]
+        cfg_mllm_inputs = [f[1] for f in features]
+        img_cfg_mllm_input = [f[2] for f in features]
+        target_img_size = [f[3] for f in features]
+
+        
+        if img_cfg_mllm_input[0] is not None:
+            mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input
+            target_img_size = target_img_size + target_img_size + target_img_size
+        else:
+            mllm_inputs = mllm_inputs + cfg_mllm_inputs
+            target_img_size = target_img_size + target_img_size
+
+
+        all_padded_input_ids, all_position_ids, all_attention_mask, all_padding_images, all_pixel_values, all_image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
+
+        data = {"input_ids": all_padded_input_ids,
+        "attention_mask": all_attention_mask,
+        "position_ids": all_position_ids,
+        "input_pixel_values": all_pixel_values,
+        "input_image_sizes": all_image_sizes,
+        "padding_images": all_padding_images,
+        }
+        return data
+
+
+class OmniGenSeparateCollator(OmniGenCollator):
+    def __call__(self, features):
+        mllm_inputs = [f[0] for f in features]
+        cfg_mllm_inputs = [f[1] for f in features]
+        img_cfg_mllm_input = [f[2] for f in features]
+        target_img_size = [f[3] for f in features]
+
+        
+        all_padded_input_ids, all_attention_mask, all_position_ids, all_pixel_values, all_image_sizes, all_padding_images = [], [], [], [], [], []
+
+
+        padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
+        all_padded_input_ids.append(padded_input_ids)
+        all_attention_mask.append(attention_mask)
+        all_position_ids.append(position_ids)
+        all_pixel_values.append(pixel_values)
+        all_image_sizes.append(image_sizes)
+        all_padding_images.append(padding_images)
+
+        if cfg_mllm_inputs[0] is not None:
+            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(cfg_mllm_inputs, target_img_size)
+            all_padded_input_ids.append(padded_input_ids)
+            all_attention_mask.append(attention_mask)
+            all_position_ids.append(position_ids)
+            all_pixel_values.append(pixel_values)
+            all_image_sizes.append(image_sizes)
+            all_padding_images.append(padding_images)
+        if img_cfg_mllm_input[0] is not None:
+            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(img_cfg_mllm_input, target_img_size)
+            all_padded_input_ids.append(padded_input_ids)
+            all_attention_mask.append(attention_mask)
+            all_position_ids.append(position_ids)
+            all_pixel_values.append(pixel_values)
+            all_image_sizes.append(image_sizes)
+            all_padding_images.append(padding_images)
+
+        data = {"input_ids": all_padded_input_ids,
+        "attention_mask": all_attention_mask,
+        "position_ids": all_position_ids,
+        "input_pixel_values": all_pixel_values,
+        "input_image_sizes": all_image_sizes,
+        "padding_images": all_padding_images,
+        }
+        return data

OmniGen/scheduler.py

@@ -0,0 +1,55 @@
+import torch
+from tqdm import tqdm
+from transformers.cache_utils import Cache, DynamicCache
+
+class OmniGenScheduler:
+    def __init__(self, num_steps: int=50, time_shifting_factor: int=1):
+        self.num_steps = num_steps
+        self.time_shift = time_shifting_factor
+
+        t = torch.linspace(0, 1, num_steps+1)
+        t = t / (t + time_shifting_factor - time_shifting_factor * t)
+        self.sigma = t
+    
+    def crop_kv_cache(self, past_key_values, num_tokens_for_img):
+        crop_past_key_values = ()
+        for layer_idx in range(len(past_key_values)):
+            key_states, value_states = past_key_values[layer_idx][:2]
+            crop_past_key_values += ((key_states[..., :-(num_tokens_for_img+1), :], value_states[..., :-(num_tokens_for_img+1), :], ),)
+        return crop_past_key_values
+        # return DynamicCache.from_legacy_cache(crop_past_key_values)
+
+    def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img):
+        if isinstance(position_ids, list):
+            for i in range(len(position_ids)):
+                position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):]
+        else:
+            position_ids = position_ids[:, -(num_tokens_for_img+1):]
+        return position_ids
+
+    def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img):
+        if isinstance(attention_mask, list):
+            return [x[..., -(num_tokens_for_img+1):, :] for x in attention_mask]
+        return attention_mask[..., -(num_tokens_for_img+1):, :]
+
+    def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True):
+        past_key_values = None
+        for i in tqdm(range(self.num_steps)):
+            timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i]
+            pred, temp_past_key_values = func(z, timesteps, past_key_values=past_key_values, **model_kwargs)
+            sigma_next = self.sigma[i+1]
+            sigma = self.sigma[i]
+            z = z + (sigma_next - sigma) * pred
+            if i == 0 and use_kv_cache:
+                num_tokens_for_img = z.size(-1)*z.size(-2) // 4
+                if isinstance(temp_past_key_values, list):
+                    past_key_values = [self.crop_kv_cache(x, num_tokens_for_img) for x in temp_past_key_values]
+                    model_kwargs['input_ids'] = [None] * len(temp_past_key_values)
+                else:
+                    past_key_values = self.crop_kv_cache(temp_past_key_values, num_tokens_for_img)
+                    model_kwargs['input_ids'] = None
+
+                model_kwargs['position_ids'] = self.crop_position_ids_for_cache(model_kwargs['position_ids'], num_tokens_for_img)
+                model_kwargs['attention_mask'] = self.crop_attention_mask_for_cache(model_kwargs['attention_mask'], num_tokens_for_img)
+        return z
+

OmniGen/train_helper/__init__.py

@@ -0,0 +1,2 @@
+from .data import DatasetFromJson, TrainDataCollator
+from .loss import training_losses

OmniGen/train_helper/data.py

@@ -0,0 +1,116 @@
+import os
+import datasets
+from datasets import load_dataset, ClassLabel, concatenate_datasets
+import torch
+import numpy as np
+import random
+from PIL import Image
+import json
+import copy
+# import torchvision.transforms as T
+from torchvision import transforms
+import pickle 
+import re
+
+from OmniGen import OmniGenProcessor
+from OmniGen.processor import OmniGenCollator
+
+
+class DatasetFromJson(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        json_file: str, 
+        image_path: str,
+        processer: OmniGenProcessor,
+        image_transform,
+        max_input_length_limit: int = 18000,
+        condition_dropout_prob: float = 0.1,
+        keep_raw_resolution: bool = True, 
+    ):
+        
+        self.image_transform = image_transform
+        self.processer = processer
+        self.condition_dropout_prob = condition_dropout_prob
+        self.max_input_length_limit = max_input_length_limit
+        self.keep_raw_resolution = keep_raw_resolution
+
+        self.data = load_dataset('json', data_files=json_file)['train']
+        self.image_path = image_path
+
+    def process_image(self, image_file):
+        if self.image_path is not None:
+            image_file = os.path.join(self.image_path, image_file)
+        image = Image.open(image_file).convert('RGB')
+        return self.image_transform(image)
+
+    def get_example(self, index):
+        example = self.data[index]
+        
+        instruction, input_images, output_image = example['instruction'], example['input_images'], example['output_image']
+        if random.random() < self.condition_dropout_prob:
+            instruction = '<cfg>'
+            input_images = None
+        if input_images is not None:
+            input_images = [self.process_image(x) for x in input_images]
+        mllm_input = self.processer.process_multi_modal_prompt(instruction, input_images)
+
+        output_image = self.process_image(output_image)
+            
+        return (mllm_input, output_image)
+
+
+    def __getitem__(self, index):
+        return self.get_example(index)
+        for _ in range(8):
+            try:
+                mllm_input, output_image = self.get_example(index)
+                if len(mllm_input['input_ids']) > self.max_input_length_limit:
+                    raise RuntimeError(f"cur number of tokens={len(mllm_input['input_ids'])}, larger than max_input_length_limit={self.max_input_length_limit}")
+                return mllm_input, output_image
+            except Exception as e:
+                print("error when loading data: ", e)
+                print(self.data[index])
+                index = random.randint(0, len(self.data)-1)
+        raise RuntimeError("Too many bad data.")
+    
+
+    def __len__(self):
+        return len(self.data)
+
+
+
+class TrainDataCollator(OmniGenCollator):
+    def __init__(self, pad_token_id: int, hidden_size: int, keep_raw_resolution: bool):
+        self.pad_token_id = pad_token_id
+        self.hidden_size = hidden_size
+        self.keep_raw_resolution = keep_raw_resolution
+
+    def __call__(self, features):
+        mllm_inputs = [f[0] for f in features]
+
+        output_images = [f[1].unsqueeze(0) for f in features]
+        target_img_size = [[x.size(-2), x.size(-1)] for x in output_images]
+
+        all_padded_input_ids, all_position_ids, all_attention_mask, all_padding_images, all_pixel_values, all_image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
+
+        if not self.keep_raw_resolution:
+            output_image = torch.cat(output_image, dim=0)
+            if len(pixel_values) > 0:
+                all_pixel_values = torch.cat(all_pixel_values, dim=0)
+            else:
+                all_pixel_values = None
+
+        data = {"input_ids": all_padded_input_ids,
+        "attention_mask": all_attention_mask,
+        "position_ids": all_position_ids,
+        "input_pixel_values": all_pixel_values,
+        "input_image_sizes": all_image_sizes,
+        "padding_images": all_padding_images,
+        "output_images": output_images,
+        }
+        return data
+
+
+
+
+

OmniGen/train_helper/loss.py

@@ -0,0 +1,68 @@
+import torch
+
+
+def sample_x0(x1):
+    """Sampling x0 & t based on shape of x1 (if needed)
+    Args:
+      x1 - data point; [batch, *dim]
+    """
+    if isinstance(x1, (list, tuple)):
+        x0 = [torch.randn_like(img_start) for img_start in x1]
+    else:
+        x0 = torch.randn_like(x1)
+
+    return x0
+
+def sample_timestep(x1):
+    u = torch.normal(mean=0.0, std=1.0, size=(len(x1),))
+    t = 1 / (1 + torch.exp(-u))
+    t = t.to(x1[0])
+    return t
+
+
+def training_losses(model, x1, model_kwargs=None, snr_type='uniform'):
+    """Loss for training torche score model
+    Args:
+    - model: backbone model; could be score, noise, or velocity
+    - x1: datapoint
+    - model_kwargs: additional arguments for torche model
+    """
+    if model_kwargs == None:
+        model_kwargs = {}
+
+    B = len(x1)
+
+    x0 = sample_x0(x1)
+    t = sample_timestep(x1)
+
+    if isinstance(x1, (list, tuple)):
+        xt = [t[i] * x1[i] + (1 - t[i]) * x0[i] for i in range(B)]
+        ut = [x1[i] - x0[i] for i in range(B)]
+    else:
+        dims = [1] * (len(x1.size()) - 1)
+        t_ = t.view(t.size(0), *dims)
+        xt = t_ * x1 + (1 - t_) * x0
+        ut = x1 - x0
+
+    model_output = model(xt, t, **model_kwargs)
+
+    terms = {}
+
+    if isinstance(x1, (list, tuple)):
+        assert len(model_output) == len(ut) == len(x1)
+        for i in range(B):
+            terms["loss"] = torch.stack(
+            [((ut[i] - model_output[i]) ** 2).mean() for i in range(B)],
+            dim=0,
+            )
+    else:
+        terms["loss"] = mean_flat(((model_output - ut) ** 2))
+
+    return terms
+
+
+def mean_flat(x):
+    """
+    Take torche mean over all non-batch dimensions.
+    """
+    return torch.mean(x, dim=list(range(1, len(x.size()))))

OmniGen/transformer.py

@@ -0,0 +1,159 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from huggingface_hub import snapshot_download
+
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers import Phi3Config, Phi3Model
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class Phi3Transformer(Phi3Model):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
+    We only modified the attention mask
+    Args:
+        config: Phi3Config
+    """
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # kept for BC (non `Cache` `past_key_values` inputs)
+        return_legacy_cache = False
+        if use_cache and not isinstance(past_key_values, Cache):
+            return_legacy_cache = True
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            else:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+                logger.warning_once(
+                    "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
+                    "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
+                    "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
+                )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if attention_mask is not None and attention_mask.dim() == 3:
+            dtype = inputs_embeds.dtype
+            min_dtype = torch.finfo(dtype).min
+            attention_mask = (1 - attention_mask) * min_dtype
+            attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype)
+        else:
+            raise 
+            # causal_mask = self._update_causal_mask(
+            #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+            # )
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if return_legacy_cache:
+            next_cache = next_cache.to_legacy_cache()
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+

OmniGen/utils.py

@@ -0,0 +1,110 @@
+import logging
+
+from PIL import Image
+import torch
+import numpy as np
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    logging.basicConfig(
+        level=logging.INFO,
+        format='[\033[34m%(asctime)s\033[0m] %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+    )
+    logger = logging.getLogger(__name__)
+    return logger
+
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = dict(ema_model.named_parameters())
+    for name, param in model.named_parameters():
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+
+
+def crop_arr(pil_image, max_image_size):
+    while min(*pil_image.size) >= 2 * max_image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    if max(*pil_image.size) > max_image_size:
+        scale = max_image_size / max(*pil_image.size)
+        pil_image = pil_image.resize(
+            tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+        )
+    
+    if min(*pil_image.size) < 16:
+        scale = 16 / min(*pil_image.size)
+        pil_image = pil_image.resize(
+            tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+        )
+    
+    arr = np.array(pil_image)
+    crop_y1 = (arr.shape[0] % 16) // 2
+    crop_y2 = arr.shape[0] % 16 - crop_y1
+
+    crop_x1 = (arr.shape[1] % 16) // 2
+    crop_x2 = arr.shape[1] % 16 - crop_x1
+
+    arr = arr[crop_y1:arr.shape[0]-crop_y2, crop_x1:arr.shape[1]-crop_x2]    
+    return Image.fromarray(arr)
+
+
+
+def vae_encode(vae, x, weight_dtype):
+    if x is not None:
+        if vae.config.shift_factor is not None:
+            x = vae.encode(x).latent_dist.sample()
+            x = (x - vae.config.shift_factor) * vae.config.scaling_factor
+        else:
+            x = vae.encode(x).latent_dist.sample().mul_(vae.config.scaling_factor)
+        x = x.to(weight_dtype)
+    return x
+
+def vae_encode_list(vae, x, weight_dtype):
+    latents = []
+    for img in x:
+        img = vae_encode(vae, img, weight_dtype)
+        latents.append(img)
+    return latents
+

README.md

@@ -1,14 +1,93 @@
----
-title: OmniGen
-emoji: 👁
-colorFrom: purple
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.4.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: OmniGen
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<h1 align="center">OmniGen: Unified Image Generation</h1>
+
+
+<p align="center">
+    <a href="">
+        <img alt="Build" src="https://img.shields.io/badge/Project%20Page-OmniGen-yellow">
+    </a>
+    <a href="https://arxiv.org/abs/2409.11340">
+            <img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-2409.11340-b31b1b.svg">
+    </a>
+    <a href="https://huggingface.co/spaces/Shitao/OmniGen">
+        <img alt="License" src="https://img.shields.io/badge/HF%20Demo-🤗-lightblue">
+    </a>
+    <a href="https://huggingface.co/Shitao/OmniGen-v1">
+        <img alt="Build" src="https://img.shields.io/badge/HF%20Model-🤗-yellow">
+    </a>
+</p>
+
+<h4 align="center">
+    <p>
+        <a href=#2-news>Credits for Quantized version</a> |
+        <a href=#3-methodology>Methodology</a> |
+        <a href=#4-what-can-omnigen-do>Capabilities</a> |
+        <a href="#license">License</a> |
+        <a href="#citation">Citation</a>
+    <p>
+</h4>
+
+
+## 1. Overview
+
+OmniGen is a unified image generation model that can generate a wide range of images from multi-modal prompts. It is designed to be simple, flexible and easy to use. We provide [inference code](#5-quick-start) so that everyone can explore more functionalities of OmniGen.
+
+Existing image generation models often require loading several additional network modules (such as ControlNet, IP-Adapter, Reference-Net, etc.) and performing extra preprocessing steps (e.g., face detection, pose estimation, cropping, etc.) to generate a satisfactory image. However, **we believe that the future image generation paradigm should be more simple and flexible, that is, generating various images directly through arbitrarily multi-modal instructions without the need for additional plugins and operations, similar to how GPT works in language generation.** 
+
+Due to the limited resources, OmniGen still has room for improvement. We will continue to optimize it, and hope it inspire more universal image generation models. You can also easily fine-tune OmniGen without worrying about designing networks for specific tasks; you just need to prepare the corresponding data, and then run the [script](#6-finetune). Imagination is no longer limited; everyone can construct any image generation task, and perhaps we can achieve very interesting, wonderful and creative things.
+
+If you have any questions, ideas or interesting tasks you want OmniGen to accomplish, feel free to discuss with us: 2906698981@qq.com, wangyueze@tju.edu.cn, zhengliu1026@gmail.com. We welcome any feedback to help us improve the model.
+
+
+
+## 2. Credits for Quantized version
+- https://github.com/Manni1000
+
+
+
+## 3. Methodology
+
+You can see details in our [paper](https://arxiv.org/abs/2409.11340). 
+
+
+## 4. What Can OmniGen do?
+
+
+OmniGen is a unified image generation model that you can use to perform various tasks, including but not limited to text-to-image generation, subject-driven generation, Identity-Preserving Generation, image editing, and image-conditioned generation. **OmniGen don't need additional plugins or operations, it can automatically identify the features (e.g., required object, human pose, depth mapping) in input images according the text prompt.**
+We showcase some examples in [inference.ipynb](inference.ipynb). And in [inference_demo.ipynb](inference_demo.ipynb), we show an interesting pipeline to generate and modify a image.
+
+Here is the illustration of OmniGen's capabilities: 
+- You can control the image generation flexibly via OmniGen
+![demo](./imgs/demo_cases.png)
+- Referring Expression Generation: You can generate images by simply referring to objects, and OmniGen will automatically recognize the required objects in the image.
+![demo](./imgs/referring.png)
+
+If you are not entirely satisfied with certain functionalities or wish to add new capabilities, you can try [fine-tuning OmniGen](#6-finetune).
+
+
+
+## 5. Quick Start
+
+### Please refer youtube video for installation
+
+https://www.youtube.com/watch?v=9ZXmXA2AJZ4
+
+
+## License
+This repo is licensed under the [MIT License](LICENSE). 
+
+
+## Citation
+If you find this repository useful, please consider giving a star ⭐ and citation
+```
+@article{xiao2024omnigen,
+  title={Omnigen: Unified image generation},
+  author={Xiao, Shitao and Wang, Yueze and Zhou, Junjie and Yuan, Huaying and Xing, Xingrun and Yan, Ruiran and Wang, Shuting and Huang, Tiejun and Liu, Zheng},
+  journal={arXiv preprint arXiv:2409.11340},
+  year={2024}
+}
+```
+
+
+
+
+

app.py

@@ -0,0 +1,359 @@
+import gradio as gr
+from PIL import Image
+import os
+from threading import Lock
+
+from OmniGen import OmniGenPipeline
+
+class OmniGenManager:
+    def __init__(self):
+        self.pipe = None
+        self.lock = Lock()
+        self.current_quantization = None
+
+    def get_pipeline(self, quantization: bool) -> OmniGenPipeline:
+        """
+        Get or initialize the pipeline with the specified quantization setting.
+        Uses a lock to ensure thread safety.
+        """
+        with self.lock:
+            # Only reinitialize if quantization setting changed or pipeline doesn't exist
+            if self.pipe is None or self.current_quantization != quantization:
+                # Clear any existing pipeline
+                if self.pipe is not None:
+                    del self.pipe
+                    self.pipe = None
+
+                # Initialize new pipeline
+                self.pipe = OmniGenPipeline.from_pretrained(
+                    "Shitao/OmniGen-v1",
+                    Quantization=quantization
+                )
+                self.current_quantization = quantization
+
+            return self.pipe
+
+# Create a single instance of the manager
+pipeline_manager = OmniGenManager()
+
+def generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, quantization):
+    input_images = [img1, img2, img3]
+    # 去除 None
+    input_images = [img for img in input_images if img is not None]
+    if len(input_images) == 0:
+        input_images = None
+
+    # Get or initialize pipeline with current settings
+    pipe = pipeline_manager.get_pipeline(quantization)
+
+    # Generate image
+    output = pipe(
+        prompt=text,
+        input_images=input_images,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        img_guidance_scale=1.6,
+        num_inference_steps=inference_steps,
+        separate_cfg_infer=True, # set False can speed up the inference process
+        use_kv_cache=False,
+        seed=seed,
+    )
+    img = output[0]
+    return img
+# def generate_image(text, img1, img2, img3, height, width, guidance_scale, inference_steps):
+#     input_images = []
+#     if img1:
+#         input_images.append(Image.open(img1))
+#     if img2:
+#         input_images.append(Image.open(img2))
+#     if img3:
+#         input_images.append(Image.open(img3))
+        
+#     return input_images[0] if input_images else None
+
+
+def get_example():
+    case = [
+        [
+            "A curly-haired man in a red shirt is drinking tea.",
+            None,
+            None,
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            0,
+        ],
+        [
+            "The woman in <img><|image_1|></img> waves her hand happily in the crowd",
+            "./imgs/test_cases/zhang.png",
+            None,
+            None,
+            1024,
+            1024,
+            2.5,
+            1.9,
+            50,
+            128,
+        ],
+        [
+            "A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
+            "./imgs/test_cases/two_man.jpg",
+            None,
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            0,
+        ],
+        [
+            "Two woman are raising fried chicken legs in a bar. A woman is <img><|image_1|></img>. The other woman is <img><|image_2|></img>.",
+            "./imgs/test_cases/mckenna.jpg",
+            "./imgs/test_cases/Amanda.jpg",
+            None,
+            1024,
+            1024,
+            2.5,
+            1.8,
+            50,
+            168,
+        ],
+        [
+            "A man and a short-haired woman with a wrinkled face are standing in front of a bookshelf in a library. The man is the man in the middle of <img><|image_1|></img>, and the woman is oldest woman in <img><|image_2|></img>",
+            "./imgs/test_cases/1.jpg",
+            "./imgs/test_cases/2.jpg",
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            60,
+        ],
+        [
+            "A man and a woman are sitting at a classroom desk. The man is the man with yellow hair in <img><|image_1|></img>. The woman is the woman on the left of <img><|image_2|></img>",
+            "./imgs/test_cases/3.jpg",
+            "./imgs/test_cases/4.jpg",
+            None,
+            1024,
+            1024,
+            2.5,
+            1.8,
+            50,
+            66,
+        ],
+        [
+            "The flower <img><|image_1|><\/img> is placed in the vase which is in the middle of <img><|image_2|><\/img> on a wooden table of a living room",
+            "./imgs/test_cases/rose.jpg",
+            "./imgs/test_cases/vase.jpg",
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            0,
+        ],
+        [
+            "<img><|image_1|><img>\n Remove the woman's earrings. Replace the mug with a clear glass filled with sparkling iced cola.",
+            "./imgs/demo_cases/t2i_woman_with_book.png",
+            None,
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            222,
+        ],
+        [
+            "Detect the skeleton of human in this image: <img><|image_1|></img>.",
+            "./imgs/test_cases/control.jpg",
+            None,
+            None,
+            1024,
+            1024,
+            2.0,
+            1.6,
+            50,
+            0,
+        ],
+        [
+            "Generate a new photo using the following picture and text as conditions: <img><|image_1|><img>\n A young boy is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
+            "./imgs/demo_cases/skeletal.png",
+            None,
+            None,
+            1024,
+            1024,
+            2,
+            1.6,
+            50,
+            42,
+        ],
+        [
+            "Following the pose of this image <img><|image_1|><img>, generate a new photo: A young boy is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
+            "./imgs/demo_cases/edit.png",
+            None,
+            None,
+            1024,
+            1024,
+            2.0,
+            1.6,
+            50,
+            123,
+        ],
+        [
+            "Following the depth mapping of this image <img><|image_1|><img>, generate a new photo: A young girl is sitting on a sofa in the library, holding a book. His hair is neatly combed, and a faint smile plays on his lips, with a few freckles scattered across his cheeks. The library is quiet, with rows of shelves filled with books stretching out behind him.",
+            "./imgs/demo_cases/edit.png",
+            None,
+            None,
+            1024,
+            1024,
+            2.0,
+            1.6,
+            50,
+            1,
+        ],
+        [
+            "<img><|image_1|><\/img> What item can be used to see the current time? Please remove it.",
+            "./imgs/test_cases/watch.jpg",
+            None,
+            None,
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            0,
+        ],
+        [
+            "According to the following examples, generate an output for the input.\nInput: <img><|image_1|></img>\nOutput: <img><|image_2|></img>\n\nInput: <img><|image_3|></img>\nOutput: ",
+            "./imgs/test_cases/icl1.jpg",
+            "./imgs/test_cases/icl2.jpg",
+            "./imgs/test_cases/icl3.jpg",
+            1024,
+            1024,
+            2.5,
+            1.6,
+            50,
+            1,
+        ],
+    ]
+    return case
+
+def run_for_examples(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed):    
+    return generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed)
+
+description = """
+OmniGen is a unified image generation model that you can use to perform various tasks, including but not limited to text-to-image generation, subject-driven generation, Identity-Preserving Generation, and image-conditioned generation.
+
+For multi-modal to image generation, you should pass a string as `prompt`, and a list of image paths as `input_images`. The placeholder in the prompt should be in the format of `<img><|image_*|></img>` (for the first image, the placeholder is <img><|image_1|></img>. for the second image, the the placeholder is <img><|image_2|></img>).
+For example, use an image of a woman to generate a new image:
+prompt = "A woman holds a bouquet of flowers and faces the camera. Thw woman is \<img\>\<|image_1|\>\</img\>."
+
+Tips:
+- Oversaturated: If the image appears oversaturated, please reduce the `guidance_scale`.
+- Low-quality: More detailed prompt will lead to better results. 
+- Animate Style: If the genereate images is in animate style, you can try to add `photo` to the prompt`.
+- Edit generated image. If you generate a image by omnigen and then want to edit it, you cannot use the same seed to edit this image. For example, use seed=0 to generate image, and should use seed=1 to edit this image.
+- For image editing tasks, we recommend placing the image before the editing instruction. For example, use `<img><|image_1|></img> remove suit`, rather than `remove suit <img><|image_1|></img>`.
+"""
+
+# Gradio 接口
+with gr.Blocks() as demo:
+    gr.Markdown("# OmniGen: Unified Image Generation [paper](https://arxiv.org/abs/2409.11340) [code](https://github.com/VectorSpaceLab/OmniGen)")
+    gr.Markdown(description)
+    with gr.Row():
+        with gr.Column():
+            # 文本输入框
+            prompt_input = gr.Textbox(
+                label="Enter your prompt, use <img><|image_i|></img> to represent i-th input image", placeholder="Type your prompt here..."
+            )
+
+            with gr.Row(equal_height=True):
+                # 图片上传框
+                image_input_1 = gr.Image(label="<img><|image_1|></img>", type="filepath")
+                image_input_2 = gr.Image(label="<img><|image_2|></img>", type="filepath")
+                image_input_3 = gr.Image(label="<img><|image_3|></img>", type="filepath")
+
+            # 高度和宽度滑块
+            height_input = gr.Slider(
+                label="Height", minimum=256, maximum=2048, value=1024, step=16
+            )
+            width_input = gr.Slider(
+                label="Width", minimum=256, maximum=2048, value=1024, step=16
+            )
+
+            # 引导尺度输入
+            guidance_scale_input = gr.Slider(
+                label="Guidance Scale", minimum=1.0, maximum=5.0, value=2.5, step=0.1
+            )
+
+            img_guidance_scale_input = gr.Slider(
+                label="img_guidance_scale", minimum=1.0, maximum=2.0, value=1.6, step=0.1
+            )
+
+            num_inference_steps = gr.Slider(
+                label="Inference Steps", minimum=1, maximum=100, value=50, step=1
+            )
+
+            Quantization = gr.Checkbox(
+                label="Low VRAM (8-bit Quantization)", value=True
+            )
+
+            seed_input = gr.Slider(
+                label="Seed", minimum=0, maximum=2147483647, value=42, step=1
+            )
+
+            # 生成按钮
+            generate_button = gr.Button("Generate Image")
+
+        with gr.Column():
+            # 输出图像框
+            output_image = gr.Image(label="Output Image")
+
+    # 按钮点击事件
+    generate_button.click(
+        generate_image,
+        inputs=[
+            prompt_input,
+            image_input_1,
+            image_input_2,
+            image_input_3,
+            height_input,
+            width_input,
+            guidance_scale_input,
+            img_guidance_scale_input,
+            num_inference_steps,
+            seed_input,
+            Quantization,
+        ],
+        outputs=output_image,
+    )
+
+    gr.Examples(
+        examples=get_example(),
+        fn=run_for_examples,
+        inputs=[
+            prompt_input,
+            image_input_1,
+            image_input_2,
+            image_input_3,
+            height_input,
+            width_input,
+            guidance_scale_input,
+            img_guidance_scale_input,
+            num_inference_steps,
+            seed_input,
+            Quantization,
+        ],
+        outputs=output_image,
+    )
+
+# 启动应用
+demo.launch()

docs/fine-tuning.md

@@ -0,0 +1,172 @@
+# Fine-tuning OmniGen
+
+Fine-tuning Omnigen can better help you handle specific image generation tasks. For example, by fine-tuning on a person's images, you can generate multiple pictures of that person while maintaining task consistency.
+
+A lot of previous work focused on designing new networks to facilitate specific tasks. For instance, ControlNet was proposed to handle image conditions, and IP-Adapter was constructed to maintain ID features. If you want to perform new tasks, you need to build new architectures and repeatedly debug them. Adding and adjusting extra network parameters is usually time-consuming and labor-intensive, which is not user-friendly and cost-efficient enough. However, with Omnigen, all of this becomes very simple.
+
+By comparison, Omnigen can accept multi-modal conditional inputs and has been pre-trained on various tasks. You can fine-tune it on any task without designing specialized networks like ControlNet or IP-Adapter for a specific task.
+
+**All you need to do is prepare the data and start training. You can break the limitations of previous models, allowing Omnigen to accomplish a variety of interesting tasks, even those that have never been done before.**
+
+
+## Installation
+
+```bash
+git clone https://github.com/VectorSpaceLab/OmniGen.git
+cd OmniGen
+pip install -e .
+```
+
+
+## Full fine-tuning
+
+### Fine-tuning command
+
+```bash
+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 100 \
+    --epochs 100 \
+    --log_every 1 \
+    --results_dir ./results/toy_finetune
+```
+
+Some important arguments:
+- `num_processes`: number of GPU to use for training
+- `model_name_or_path`: path to the pretrained model
+- `json_file`: path to the json file containing the training data, e.g., ./toy_data/toy_data.jsonl
+- `image_path`: path to the image folder, e.g., ./toy_data/images
+- `batch_size_per_device`: batch size per device
+- `lr`: learning rate
+- `keep_raw_resolution`: whether to keep the original resolution of the image, if not, all images will be resized to (max_image_size, max_image_size)
+- `max_image_size`: max image size
+- `gradient_accumulation_steps`: number of steps to accumulate gradients
+- `ckpt_every`: number of steps to save checkpoint
+- `epochs`: number of epochs
+- `log_every`: number of steps to log
+- `results_dir`: path to the results folder
+
+The data format of json_file is as follows:
+```
+{
+    "instruction": str, 
+    "input_images": [str, str, ...], 
+    "output_images": str
+}
+```
+You can see a toy example in `./toy_data/toy_data.jsonl`.
+
+If an OOM(Out of Memory) issue occurs, you can try to decrease the `batch_size_per_device` or `max_image_size`. You can also try to use LoRA instead of full fine-tuning.
+
+
+### Inference
+
+The checkpoint can be found at `{results_dir}/checkpoints/*`. You can use the following command to load saved checkpoint:
+```python
+from OmniGen import OmniGenPipeline
+
+pipe = OmniGenPipeline.from_pretrained("checkpoint_path")  # e.g., ./results/toy_finetune/checkpoints/0000200
+```
+
+
+
+
+
+## LoRA fine-tuning
+LoRA fine-tuning is a simple way to fine-tune OmniGen with less GPU memory. To use lora, you should add `--use_lora` and `--lora_rank` to the command.
+
+```bash
+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 100 \
+    --epochs 100 \
+    --log_every 1 \
+    --results_dir ./results/toy_finetune_lora
+```
+
+### Inference
+
+The checkpoint can be found at `{results_dir}/checkpoints/*`. You can use the following command to load checkpoint:
+```python
+from OmniGen import OmniGenPipeline
+
+pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")
+pipe.merge_lora("checkpoint_path")  # e.g., ./results/toy_finetune_lora/checkpoints/0000100
+```
+
+
+## A simple example
+
+Here is an example for learning new concepts: "sks dog". We use five images of one dog from [dog-example](https://huggingface.co/datasets/diffusers/dog-example). 
+
+The json file is `./toy_data/toy_subject_data.jsonl`, and the images have been saved in `./toy_data/images`.
+
+```bash
+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 100 \
+    --epochs 200 \
+    --log_every 1 \
+    --results_dir ./results/toy_finetune_lora
+```
+
+After training, you can use the following command to generate images:
+```python
+from OmniGen import OmniGenPipeline
+
+pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")
+pipe.merge_lora("checkpoint_path") # e.g., ./results/toy_finetune_lora/checkpoints/0000200
+
+images = pipe(
+    prompt="a photo of sks dog running in the snow", 
+    height=1024, 
+    width=1024, 
+    guidance_scale=3
+)
+images[0].save("example_sks_dog_snow.png")
+```

docs/inference.md

@@ -0,0 +1,96 @@
+# Inference with OmniGen
+
+To handle some complex tasks, image generation models are becoming increasingly sophisticated, leading to more and more cumbersome workflows. Existing image generation models like SD and Flux require loading many additional network modules (such as ControlNet, IP-Adapter, Reference-Net) and extra preprocessing steps (e.g., face detection, pose detection, image cropping) to generate a satisfactory image. This complex workflow is not user-friendly. We believe that future image generation models should be simpler, generating various images directly through instructions, similar to how GPT works in language generation.
+
+Therefore, we propose OmniGen, a model capable of handling various image generation tasks within a single framework. The goal of OmniGen is to complete various image generation tasks without relying on any additional components or image preprocessing steps. OmniGen supports tasks including text-to-image generation, image editing, subject-driven image generation, and classical vision tasks, among others. More capabilities can be found in our examples. We provide inference code so you can explore more unknown functionalities yourself.
+
+
+
+## Install
+```bash
+git clone https://github.com/staoxiao/OmniGen.git
+cd OmniGen
+pip install -e .
+```
+
+
+
+## Generate Images
+You can use the following code to generate images:
+```python
+from OmniGen import OmniGenPipeline
+
+pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1")
+
+# Text to Image
+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")  # save output PIL Image
+
+# Multi-modal to Image
+# In prompt, we use the placeholder to represent the image. The image placeholder should be in the format of <img><|image_*|></img>
+# You can add multiple images in the input_images. Please ensure that each image has its placeholder. For example, for the list input_images [img1_path, img2_path], the prompt needs to have two placeholders: <img><|image_1|></img>, <img><|image_2|></img>.
+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,
+    separate_cfg_infer=False,  # if OOM, you can set separate_cfg_infer=True 
+    guidance_scale=2.5, 
+    img_guidance_scale=1.6
+)
+images[0].save("example_ti2i.png")  # save output PIL image
+```
+
+Some important arguments:
+- `guidance_scale`: The strength of the guidance. Based on our experience, it is usually best to set it between 2 and 3. The higher the value, the more similar the generated image will be to the prompt. If the image appears oversaturated, please reduce the scale. 
+- `height` and `width`: The height and width of the generated image. The default value is 1024x1024. OmniGen support any size, but these number must be divisible by 16.
+- `num_inference_steps`: The number of steps to take in the diffusion process. The higher the value, the more detailed the generated image will be.
+- `separate_cfg_infer`: Whether to use separate inference process for CFG guidance. If set to True, memory cost will be lower but the generation speed will be slower. Default is False.
+- `use_kv_cache`: Whether to use key-value cache. Default is True.
+- `seed`: The seed for random number generator.
+
+**More examples please refer to [inference.ipynb](../inference.ipynb)**
+
+
+#### Input data
+OmniGen can accept multi-modal input data. Specifically, you should pass two arguments: `prompt` and `input_images`.
+For text to image generation, you can pass a string as `prompt`, or pass a list of strings as `prompt` to generate multiple images.
+
+For multi-modal to image generation, you should pass a string as `prompt`, and a list of image paths as `input_images`. The placeholder in the prompt should be in the format of `<img><|image_*|></img>`.
+For example, if you want to generate an image with a person holding a bouquet of flowers, you can pass the following prompt:
+```
+prompt = "A woman holds a bouquet of flowers and faces the camera. Thw woman is <img><|image_1|></img>."
+input_images = ["./imgs/test_cases/liuyifei.png"]
+```
+The placeholder `<|image_1|>` will be replaced by the image at `input_images[0]`, i.e., `./imgs/test_cases/liuyifei.png`.
+
+If you want to generate multiple images, you can pass a list of prompts and a list of image paths. For example:
+```
+prompt = ["A woman holds a bouquet of flowers and faces the camera.", "A woman holds a bouquet of flowers and faces the camera. Thw woman is <img><|image_1|></img>."]
+input_images = [[], ["./imgs/test_cases/liuyifei.png"]]
+```
+
+
+#### Gradio Demo
+We have constructed a online demo in [Huggingface](https://huggingface.co/spaces/Shitao/OmniGen).
+
+For the local gradio demo, you can run with the following command:
+```python
+python app.py
+```
+
+
+## Tips
+- OOM issue: If you encounter OOM issue, you can try to set `separate_cfg_infer=True`. This will reduce the memory usage but increase the generation latecy. You also can reduce the size of the image, e.g., `height=768, width=512`.
+- Oversaturated: If the image appears oversaturated, please reduce the `guidance_scale`.
+- Not match the prompt: If the image does not match the prompt, please try to increase the `guidance_scale`.
+- Low-quality: More detailed prompt will lead to better results. Besides, larger size of the image (`height` and `width`) will also help.
+- Animate Style: If the genereate images is in animate style, you can try to add `photo` to the prompt`.
+- Edit generated image. If you generate a image by omnigen and then want to edit it, you cannot use the same seed to edit this image. For example, use seed=0 to generate image, and should use seed=1 to edit this image.
+- For image editing tasks, we recommend placing the image before the editing instruction. For example, use `<img><|image_1|></img> remove suit`, rather than `remove suit <img><|image_1|></img>`.