update inference code

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/model.py

@@ -312,7 +312,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
         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):
+    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, offload_model:bool=False):
         """
         
         """
@@ -335,7 +335,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
             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 = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, offload_model=offload_model)
         output, past_key_values = output.last_hidden_state, output.past_key_values
         if input_is_list:
             image_embedding = output[:, -max(num_tokens):]
@@ -357,12 +357,9 @@ class OmniGen(nn.Module, PeftAdapterMixin):
         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.
-        """        
+    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, offload_model):      
         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)
+        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, offload_model=offload_model)
         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)
@@ -376,10 +373,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
 
 
     @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.
-        """        
+    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, offload_model):
         self.llm.config.use_cache = use_kv_cache
         if past_key_values is None:
             past_key_values = [None] * len(attention_mask)
@@ -390,7 +384,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
 
         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])
+            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=past_key_values[i], return_past_key_values=True, offload_model=offload_model)
             model_out.append(temp_out)
             pask_key_values.append(temp_pask_key_values)
 

OmniGen/pipeline.py

@@ -1,6 +1,7 @@
 import os
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Union
+import gc
 
 from PIL import Image
 import numpy as np
@@ -33,7 +34,7 @@ EXAMPLE_DOC_STRING = """
         >>> prompt = "A woman holds a bouquet of flowers and faces the camera"
         >>> image = pipe(
         ...     prompt,
-        ...     guidance_scale=3.0,
+        ...     guidance_scale=2.5,
         ...     num_inference_steps=50,
         ... ).images[0]
         >>> image.save("t2i.png")
@@ -41,7 +42,7 @@ EXAMPLE_DOC_STRING = """
 """
 
 
-
+90
 class OmniGenPipeline:
     def __init__(
         self,
@@ -53,10 +54,21 @@ class OmniGenPipeline:
         self.model = model
         self.processor = processor
 
-        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-        self.model.to(self.device)
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        elif is_torch_npu_available():
+            self.device = torch.device("npu")
+        else:
+            logger.info("Don't detect any available devices, using CPU instead")
+            self.device = torch.device("cpu")
+
+        self.model.to(torch.bfloat16)
         self.model.eval()
-        self.vae.to(self.device)
+        self.vae.eval()
+
+        self.model_cpu_offload = False
 
     @classmethod
     def from_pretrained(cls, model_name, vae_path: str=None):
@@ -84,7 +96,6 @@ class OmniGenPipeline:
         model = PeftModel.from_pretrained(self.model, lora_path)
         model.merge_and_unload()
 
-        
         self.model = model
     
     def to(self, device: Union[str, torch.device]):
@@ -92,6 +103,7 @@ class OmniGenPipeline:
             device = torch.device(device)
         self.model.to(device)
         self.vae.to(device)
+        self.device = device
 
     def vae_encode(self, x, dtype):
         if self.vae.config.shift_factor is not None:
@@ -107,6 +119,17 @@ class OmniGenPipeline:
             return [x.to(self.device) for x in data]
         return data.to(self.device)
 
+    def enable_model_cpu_offload(self):
+        self.model_cpu_offload = True
+        self.model.to("cpu")
+        self.vae.to("cpu")
+        torch.cuda.empty_cache()  # Clear VRAM
+        gc.collect()  # Run garbage collection to free system RAM
+    
+    def disable_model_cpu_offload(self):
+        self.model_cpu_offload = False
+        self.model.to(self.device)
+        self.vae.to(self.device)
 
     @torch.no_grad()
     @replace_example_docstring(EXAMPLE_DOC_STRING)
@@ -120,8 +143,12 @@ class OmniGenPipeline:
         guidance_scale: float = 3,
         use_img_guidance: bool = True,
         img_guidance_scale: float = 1.6,
-        separate_cfg_infer: bool = False,
+        max_input_image_size: int = 1024,
+        separate_cfg_infer: bool = True,
+        offload_model: bool = False,
         use_kv_cache: bool = True,
+        offload_kv_cache: bool = True,
+        use_input_image_size_as_output: bool = False,
         dtype: torch.dtype = torch.bfloat16,
         seed: int = None,
         ):
@@ -149,31 +176,50 @@ class OmniGenPipeline:
                 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). 
+            max_input_image_size (`int`, *optional*, defaults to 1024): the maximum size of input image, which will be used to crop the input image to the maximum size
             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.
+            offload_kv_cache (`bool`, *optional*, defaults to True): offload the cached key and value to cpu, which can save memory but slow down the generation silightly
+            offload_model (`bool`, *optional*, defaults to False): offload the model to cpu, which can save memory but slow down the generation
+            use_input_image_size_as_output (bool, defaults to False): whether to use the input image size as the output image size, which can be used for single-image input, e.g., image editing task
+            seed (`int`, *optional*):
+                A random seed for generating output. 
+            dtype (`torch.dtype`, *optional*, defaults to `torch.bfloat16`):
+                data type for the model
         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"
+        # check inputs:
+        if use_input_image_size_as_output:
+            assert isinstance(prompt, str) and len(input_images) == 1, "if you want to make sure the output image have the same size as the input image, please only input one image instead of multiple input images"
+        else:
+            assert height%16 == 0 and width%16 == 0, "The height and width must be a multiple of 16."
         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
+        
+        # set model and processor
+        if max_input_image_size != self.processor.max_image_size:
+            self.processor = OmniGenProcessor(self.processor.text_tokenizer, max_image_size=max_input_image_size)
+        if offload_model:
+            self.enable_model_cpu_offload()
+        else:
+            self.disable_model_cpu_offload()
 
-        input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=separate_cfg_infer)
+        input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=separate_cfg_infer, use_input_image_size_as_output=use_input_image_size_as_output)
 
         num_prompt = len(prompt)
         num_cfg = 2 if use_img_guidance else 1
+        if use_input_image_size_as_output:
+            if separate_cfg_infer:
+                height, width = input_data['input_pixel_values'][0][0].shape[-2:]
+            else:
+                height, width = input_data['input_pixel_values'][0].shape[-2:]
         latent_size_h, latent_size_w = height//8, width//8
 
         if seed is not None:
@@ -183,6 +229,7 @@ class OmniGenPipeline:
         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)
 
+        if input_images is not None and self.model_cpu_offload: self.vae.to(self.device)
         input_img_latents = []
         if separate_cfg_infer:
             for temp_pixel_values in input_data['input_pixel_values']:
@@ -195,6 +242,10 @@ class OmniGenPipeline:
             for img in input_data['input_pixel_values']:
                 img = self.vae_encode(img.to(self.device), dtype)
                 input_img_latents.append(img)
+        if input_images is not None and self.model_cpu_offload:
+            self.vae.to('cpu')
+            torch.cuda.empty_cache()  # Clear VRAM
+            gc.collect()  # Run garbage collection to free system RAM
 
         model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']), 
             input_img_latents=input_img_latents, 
@@ -204,7 +255,9 @@ class OmniGenPipeline:
             cfg_scale=guidance_scale,
             img_cfg_scale=img_guidance_scale,
             use_img_cfg=use_img_guidance,
-            use_kv_cache=use_kv_cache)
+            use_kv_cache=use_kv_cache,
+            offload_model=offload_model,
+            )
         
         if separate_cfg_infer:
             func = self.model.forward_with_separate_cfg
@@ -212,16 +265,38 @@ class OmniGenPipeline:
             func = self.model.forward_with_cfg
         self.model.to(dtype)
 
+        if self.model_cpu_offload:
+            for name, param in self.model.named_parameters():
+                if 'layers' in name and 'layers.0' not in name:
+                    param.data = param.data.cpu()
+                else:
+                    param.data = param.data.to(self.device)
+            for buffer_name, buffer in self.model.named_buffers():
+                setattr(self.model, buffer_name, buffer.to(self.device))
+        # else:
+        #     self.model.to(self.device)
+
         scheduler = OmniGenScheduler(num_steps=num_inference_steps)
-        samples = scheduler(latents, func, model_kwargs, use_kv_cache=use_kv_cache)
+        samples = scheduler(latents, func, model_kwargs, use_kv_cache=use_kv_cache, offload_kv_cache=offload_kv_cache)
         samples = samples.chunk((1+num_cfg), dim=0)[0]
 
+        if self.model_cpu_offload:
+            self.model.to('cpu')
+            torch.cuda.empty_cache()  
+            gc.collect()  
+
+        self.vae.to(self.device)
         samples = samples.to(torch.float32)
         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   
         samples = self.vae.decode(samples).sample
+
+        if self.model_cpu_offload:
+            self.vae.to('cpu')
+            torch.cuda.empty_cache()  
+            gc.collect()  
         
         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()
@@ -229,4 +304,6 @@ class OmniGenPipeline:
         for i, sample in enumerate(output_samples):  
             output_images.append(Image.fromarray(sample))
         
+        torch.cuda.empty_cache()  # Clear VRAM
+        gc.collect()              # Run garbage collection to free system RAM
         return output_images

OmniGen/processor.py

@@ -108,6 +108,7 @@ class OmniGenProcessor:
                 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,
+                use_input_image_size_as_output: bool=False,
                 ) -> Dict:
 
         if input_images is None:
@@ -138,7 +139,10 @@ class OmniGenProcessor:
                 else:
                     img_cfg_mllm_input = neg_mllm_input
 
-            input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
+            if use_input_image_size_as_output:
+                input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [mllm_input['pixel_values'][0].size(-2), mllm_input['pixel_values'][0].size(-1)]))
+            else:
+                input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
 
         if separate_cfg_input:
             return self.separate_collator(input_data)
@@ -295,7 +299,6 @@ class OmniGenSeparateCollator(OmniGenCollator):
         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 = [], [], [], [], [], []
 

OmniGen/scheduler.py

@@ -1,6 +1,116 @@
-import torch
 from tqdm import tqdm
-from transformers.cache_utils import Cache, DynamicCache
+from typing import Optional, Dict, Any, Tuple, List
+import gc
+
+import torch
+from transformers.cache_utils import Cache, DynamicCache, OffloadedCache
+
+
+
+class OmniGenCache(DynamicCache):
+    def __init__(self, 
+                    num_tokens_for_img: int, offload_kv_cache: bool=False) -> None:
+        if not torch.cuda.is_available():
+            raise RuntimeError("OffloadedCache can only be used with a GPU")
+        super().__init__()
+        self.original_device = []
+        self.prefetch_stream = torch.cuda.Stream()
+        self.num_tokens_for_img = num_tokens_for_img
+        self.offload_kv_cache = offload_kv_cache
+
+    def prefetch_layer(self, layer_idx: int):
+        "Starts prefetching the next layer cache"
+        if layer_idx < len(self):
+            with torch.cuda.stream(self.prefetch_stream):
+                # Prefetch next layer tensors to GPU
+                device = self.original_device[layer_idx]
+                self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device, non_blocking=True)
+                self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device, non_blocking=True)
+
+    
+    def evict_previous_layer(self, layer_idx: int):
+        "Moves the previous layer cache to the CPU"
+        if len(self) > 2:
+            # We do it on the default stream so it occurs after all earlier computations on these tensors are done
+            if layer_idx == 0: 
+                prev_layer_idx = -1
+            else:
+                prev_layer_idx = (layer_idx - 1) % len(self)
+            self.key_cache[prev_layer_idx] = self.key_cache[prev_layer_idx].to("cpu", non_blocking=True)
+            self.value_cache[prev_layer_idx] = self.value_cache[prev_layer_idx].to("cpu", non_blocking=True)
+
+
+    def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]:
+        "Gets the cache for this layer to the device. Prefetches the next and evicts the previous layer."
+        if layer_idx < len(self):
+            if self.offload_kv_cache:
+                # Evict the previous layer if necessary
+                torch.cuda.current_stream().synchronize()
+                self.evict_previous_layer(layer_idx)
+                # Load current layer cache to its original device if not already there
+                original_device = self.original_device[layer_idx]
+                # self.prefetch_stream.synchronize(original_device)
+                torch.cuda.synchronize(self.prefetch_stream)
+                key_tensor = self.key_cache[layer_idx]
+                value_tensor = self.value_cache[layer_idx]
+                
+                # Prefetch the next layer
+                self.prefetch_layer((layer_idx + 1) % len(self))
+            else:
+                key_tensor = self.key_cache[layer_idx]
+                value_tensor = self.value_cache[layer_idx]
+            return (key_tensor, value_tensor)
+        else:
+            raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}")
+        
+       
+    def update(
+        self,
+        key_states: torch.Tensor, 
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
+        Parameters:
+            key_states (`torch.Tensor`):
+                The new key states to cache.
+            value_states (`torch.Tensor`):
+                The new value states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+            cache_kwargs (`Dict[str, Any]`, `optional`):
+                Additional arguments for the cache subclass. No additional arguments are used in `OffloadedCache`.
+        Return:
+            A tuple containing the updated key and value states.
+        """
+        # Update the cache
+        if len(self.key_cache) < layer_idx:
+            raise ValueError("OffloadedCache does not support model usage where layers are skipped. Use DynamicCache.")
+        elif len(self.key_cache) == layer_idx:
+            # only cache the states for condition tokens
+            key_states = key_states[..., :-(self.num_tokens_for_img+1), :]
+            value_states = value_states[..., :-(self.num_tokens_for_img+1), :]
+
+             # Update the number of seen tokens
+            if layer_idx == 0:
+                self._seen_tokens += key_states.shape[-2]
+                
+            self.key_cache.append(key_states)
+            self.value_cache.append(value_states)
+            self.original_device.append(key_states.device)
+            if self.offload_kv_cache:
+                self.evict_previous_layer(layer_idx)
+            return self.key_cache[layer_idx], self.value_cache[layer_idx]
+        else:
+            # only cache the states for condition tokens
+            key_tensor, value_tensor = self[layer_idx]
+            k = torch.cat([key_tensor, key_states], dim=-2)
+            v = torch.cat([value_tensor, value_states], dim=-2)
+            return k, v
+
+
 
 class OmniGenScheduler:
     def __init__(self, num_steps: int=50, time_shifting_factor: int=1):
@@ -12,12 +122,13 @@ class OmniGenScheduler:
         self.sigma = t
     
     def crop_kv_cache(self, past_key_values, num_tokens_for_img):
+        # return 
         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)
+        # 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):
@@ -32,24 +143,39 @@ class OmniGenScheduler:
             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
+    def crop_cache(self, cache, num_tokens_for_img):
+        for i in range(len(cache.key_cache)):
+            cache.key_cache[i] = cache.key_cache[i][..., :-(num_tokens_for_img+1), :]
+            cache.value_cache[i] = cache.value_cache[i][..., :-(num_tokens_for_img+1), :]
+        
+        return cache
+
+    def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True, offload_kv_cache: bool=True):
+        num_tokens_for_img = z.size(-1)*z.size(-2) // 4
+        if isinstance(model_kwargs['input_ids'], list):
+            cache = [OmniGenCache(num_tokens_for_img, offload_kv_cache) for _ in range(len(model_kwargs['input_ids']))] if use_kv_cache else None
+        else:
+            cache = OmniGenCache(num_tokens_for_img, offload_kv_cache) if use_kv_cache else None
+        results = {}
         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)
+            pred, cache = func(z, timesteps, past_key_values=cache, **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)
+                if isinstance(cache, list):
+                    model_kwargs['input_ids'] = [None] * len(cache)
                 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)
+
+        del cache
+        torch.cuda.empty_cache()  
+        gc.collect()
         return z
 
+    

OmniGen/transformer.py

@@ -29,6 +29,34 @@ class Phi3Transformer(Phi3Model):
     Args:
         config: Phi3Config
     """
+    def prefetch_layer(self, layer_idx: int, device: torch.device):
+        "Starts prefetching the next layer cache"
+        with torch.cuda.stream(self.prefetch_stream):
+            # Prefetch next layer tensors to GPU
+            for name, param in self.layers[layer_idx].named_parameters():
+                param.data = param.data.to(device, non_blocking=True)
+
+    def evict_previous_layer(self, layer_idx: int):
+        "Moves the previous layer cache to the CPU"
+        prev_layer_idx = layer_idx - 1
+        for name, param in self.layers[prev_layer_idx].named_parameters():
+            param.data = param.data.to("cpu", non_blocking=True)
+            
+    def get_offlaod_layer(self, layer_idx: int, device: torch.device):
+        # init stream
+        if not hasattr(self, "prefetch_stream"):
+            self.prefetch_stream = torch.cuda.Stream()
+
+        # delete previous layer
+        torch.cuda.current_stream().synchronize()
+        self.evict_previous_layer(layer_idx)
+        
+        # make sure the current layer is ready
+        torch.cuda.synchronize(self.prefetch_stream)
+
+        # load next layer
+        self.prefetch_layer((layer_idx + 1) % len(self.layers), device)
+        
 
     def forward(
         self,
@@ -42,6 +70,7 @@ class Phi3Transformer(Phi3Model):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
+        offload_model: Optional[bool] = False,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -75,16 +104,16 @@ class Phi3Transformer(Phi3Model):
                     "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
                 )
 
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
+        # 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 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
@@ -104,7 +133,10 @@ class Phi3Transformer(Phi3Model):
         all_self_attns = () if output_attentions else None
         next_decoder_cache = None
 
+        layer_idx = -1
         for decoder_layer in self.layers:
+            layer_idx += 1
+
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
@@ -120,6 +152,8 @@ class Phi3Transformer(Phi3Model):
                     cache_position,
                 )
             else:
+                if offload_model and not self.training:
+                    self.get_offlaod_layer(layer_idx, device=inputs_embeds.device)
                 layer_outputs = decoder_layer(
                     hidden_states,
                     attention_mask=attention_mask,
@@ -142,6 +176,7 @@ class Phi3Transformer(Phi3Model):
 
         # add hidden states from the last decoder layer
         if output_hidden_states:
+            print('************')
             all_hidden_states += (hidden_states,)
 
         next_cache = next_decoder_cache if use_cache else None

README.md

@@ -1,13 +1,180 @@
----
-title: OmniGen
-emoji: 🖼
-colorFrom: purple
-colorTo: red
-sdk: gradio
-sdk_version: 5.0.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-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=#1-news>News</a> |
+        <a href=#3-methodology>Methodology</a> |
+        <a href=#4-what-can-omnigen-do>Capabilities</a> |
+        <a href=#5-quick-start>Quick Start</a> |
+        <a href="#6-finetune">Finetune</a> |
+        <a href="#license">License</a> |
+        <a href="#citation">Citation</a>
+    <p>
+</h4>
+
+
+
+## 1. News
+- 2024-10-28: We release new version of inference code, optimizing the memory usage and time cost. You can refer to [docs/inference.md](docs/inference.md#requiremented-resources) for detailed information.
+- 2024-10-22: :fire: We release the code for OmniGen. Inference: [docs/inference.md](docs/inference.md) Train: [docs/fine-tuning.md](docs/fine-tuning.md) 
+- 2024-10-22: :fire: We release the first version of OmniGen. Model Weight: [Shitao/OmniGen-v1](https://huggingface.co/Shitao/OmniGen-v1) HF Demo: [🤗](https://huggingface.co/spaces/Shitao/OmniGen)  
+
+
+## 2. 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: [email protected], [email protected], [email protected]. We welcome any feedback to help us improve the model.
+
+
+
+
+## 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 input multiple images and use simple, general language to refer to the objects within those images. OmniGen can automatically recognize the necessary objects in each image and generate new images based on them. No additional operations, such as image cropping or face detection, are required.
+![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
+
+
+### Using OmniGen
+Install via Github:
+```bash
+git clone https://github.com/staoxiao/OmniGen.git
+cd OmniGen
+pip install -e .
+```
+
+Here are some examples:
+```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,
+    guidance_scale=2.5, 
+    img_guidance_scale=1.6,
+    seed=0
+)
+images[0].save("example_ti2i.png")  # save output PIL image
+```
+- For thre required resources and the method to run OmniGen efficiently, please refer to [docs/inference.md#requiremented-resources](docs/inference.md#requiremented-resources).
+- For more examples for image generation, you can refer to [inference.ipynb](inference.ipynb) and [inference_demo.ipynb](inference_demo.ipynb)
+- For more details about the argument in inference, please refer to [docs/inference.md](docs/inference.md). 
+
+
+### Using Diffusers
+Coming soon.
+
+
+### Gradio Demo
+
+We construct an online demo in [Huggingface](https://huggingface.co/spaces/Shitao/OmniGen).
+
+For the local gradio demo, you need to install `pip install gradio spaces` , and then you can run:
+```python
+pip install gradio spaces
+python app.py
+```
+
+
+
+## 6. Finetune
+We provide a training script `train.py` to fine-tune OmniGen. 
+Here is a toy example about LoRA finetune:
+```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 10 \
+    --epochs 200 \
+    --log_every 1 \
+    --results_dir ./results/toy_finetune_lora
+```
+
+Please refer to [docs/fine-tuning.md](docs/fine-tuning.md) for more details (e.g. full finetune).
+
+
+
+## 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

@@ -1,6 +1,7 @@
 import gradio as gr
 from PIL import Image
 import os
+
 import spaces
 
 from OmniGen import OmniGenPipeline
@@ -9,11 +10,11 @@ pipe = OmniGenPipeline.from_pretrained(
     "Shitao/OmniGen-v1"
 )
 
-@spaces.GPU(duration=180)
-# 示例处理函数：生成图像
-def generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer):
+@spaces.GPU(duration=300)
+def generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model,
+            use_input_image_size_as_output):
     input_images = [img1, img2, img3]
-    # 去除 None
+    # Delete None
     input_images = [img for img in input_images if img is not None]
     if len(input_images) == 0:
         input_images = None
@@ -24,25 +25,18 @@ def generate_image(text, img1, img2, img3, height, width, guidance_scale, img_gu
         height=height,
         width=width,
         guidance_scale=guidance_scale,
-        img_guidance_scale=1.6,
+        img_guidance_scale=img_guidance_scale,
         num_inference_steps=inference_steps,
-        separate_cfg_infer=True, # set False can speed up the inference process
-        use_kv_cache=False,
+        separate_cfg_infer=separate_cfg_infer, 
+        use_kv_cache=True,
+        offload_kv_cache=True,
+        offload_model=offload_model,
+        use_input_image_size_as_output=use_input_image_size_as_output,
         seed=seed,
-        # separate_cfg_infer=separate_cfg_infer,
     )
     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():
@@ -59,6 +53,8 @@ def get_example():
             50,
             0,
             True,
+            False,
+            False,
         ],
         [
             "The woman in <img><|image_1|></img> waves her hand happily in the crowd",
@@ -72,6 +68,8 @@ def get_example():
             50,
             128,
             True,
+            False,
+            False,
         ],
         [
             "A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
@@ -85,6 +83,8 @@ def get_example():
             50,
             0,
             True,
+            False,
+            False,
         ],
         [
             "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>.",
@@ -98,6 +98,8 @@ def get_example():
             50,
             168,
             True,
+            False,
+            False,
         ],
         [
             "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>",
@@ -111,6 +113,8 @@ def get_example():
             50,
             60,
             True,
+            False,
+            False,
         ],
         [
             "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>",
@@ -124,6 +128,8 @@ def get_example():
             50,
             66,
             True,
+            False,
+            False,
         ],
         [
             "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",
@@ -137,147 +143,185 @@ def get_example():
             50,
             0,
             True,
+            False,
+            False,
         ],
         [
             "<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,
+            None,
+            None,
             2.5,
             1.6,
             50,
             222,
             True,
+            False,
+            True,
         ],
         [
             "Detect the skeleton of human in this image: <img><|image_1|></img>.",
             "./imgs/test_cases/control.jpg",
             None,
             None,
-            1024,
-            1024,
+            None,
+            None,
             2.0,
             1.6,
             50,
             0,
             True,
+            False,
+            True,
         ],
         [
             "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,
+            None,
+            None,
             2,
             1.6,
             50,
             42,
             True,
+            False,
+            True,
         ],
         [
             "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,
+            None,
+            None,
             2.0,
             1.6,
             50,
             123,
             True,
+            False,
+            True,
         ],
         [
             "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,
+            None,
+            None,
             2.0,
             1.6,
             50,
             1,
             True,
+            False,
+            True,
         ],
         [
             "<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,
+            None,
+            None,
             2.5,
             1.6,
             50,
             0,
             True,
+            False,
+            True,
         ],
         [
             "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,
+            224,
+            224,
             2.5,
             1.6,
             50,
             1,
             True,
+            False,
+            False,
         ],
     ]
     return case
 
-def run_for_examples(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer,):    
-    return generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer,)
+def run_for_examples(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model,
+            use_input_image_size_as_output):    
+    return generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed, separate_cfg_infer, offload_model,
+            use_input_image_size_as_output)
 
 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:
+- For out of memory or time cost, you can refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources) to select a appropriate setting.
 - 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. 
 - 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>`.
+- 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>`. 
+- For image editing task and controlnet task, we recommend to set the height and width of output image as the same as input image. For example, if you want to edit a 512x512 image, you should set the height and width of output image as 512x512. You also can set the `use_input_image_size_as_output` to automatically set the height and width of output image as the same as input image.
+
+
+"""
+
+article = """
+---
+**Citation** 
+<br> 
+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}
+}
+```
+**Contact**
+<br>
+If you have any questions, please feel free to open an issue or directly reach us out via email.
 """
 
-separate_cfg_infer_arg = False
 
-# Gradio 接口
+# 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():
-            # 文本输入框
+            # text prompt
             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):
-                # 图片上传框
+                # input images
                 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")
 
-            # 高度和宽度滑块
+            # slider
             height_input = gr.Slider(
-                label="Height", minimum=256, maximum=2048, value=1024, step=16
+                label="Height", minimum=128, maximum=2048, value=1024, step=16
             )
             width_input = gr.Slider(
-                label="Width", minimum=256, maximum=2048, value=1024, step=16
+                label="Width", minimum=128, 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
             )
@@ -295,17 +339,24 @@ with gr.Blocks() as demo:
             )
 
             separate_cfg_infer = gr.Checkbox(
-                label="separate_cfg_infer", info="enable separate cfg infer"
+                label="separate_cfg_infer", info="Whether to use separate inference process for different guidance. This will reduce the memory cost.", value=True,
+            )
+            offload_model = gr.Checkbox(
+                label="offload_model", info="Offload model to CPU, which will significantly reduce the memory cost but slow down the generation speed. You can cancle separate_cfg_infer and set offload_model=True. If both separate_cfg_infer and offload_model be True, further reduce the memory, but slowest generation", value=False,
+            )
+            use_input_image_size_as_output = gr.Checkbox(
+                label="use_input_image_size_as_output", info="Automatically adjust the output image size to be same as input image size. For editing and controlnet task, it can make sure the output image has the same size with input image leading to better performance", value=False,
             )
 
-            # 生成按钮
+            # generate
             generate_button = gr.Button("Generate Image")
+            
 
         with gr.Column():
-            # 输出图像框
+            # output image
             output_image = gr.Image(label="Output Image")
 
-    # 按钮点击事件
+    # click
     generate_button.click(
         generate_image,
         inputs=[
@@ -320,6 +371,8 @@ with gr.Blocks() as demo:
             num_inference_steps,
             seed_input,
             separate_cfg_infer,
+            offload_model,
+            use_input_image_size_as_output,
         ],
         outputs=output_image,
     )
@@ -339,9 +392,13 @@ with gr.Blocks() as demo:
             num_inference_steps,
             seed_input,
             separate_cfg_infer,
+            offload_model,
+            use_input_image_size_as_output,
         ],
         outputs=output_image,
     )
 
-# 启动应用
+    gr.Markdown(article)
+
+# launch
 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,167 @@
+# 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,
+    guidance_scale=2.5, 
+    img_guidance_scale=1.6,
+    max_input_image_size=1024,
+    separate_cfg_infer=True, 
+    use_kv_cache=True,
+    offload_kv_cache=True,
+    offload_model=False,
+    use_input_image_size_as_output=False,
+    seed=0,
+)
+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.
+- `max_input_image_size`: the maximum size of input image, which will be used to crop the input image to the maximum size. A smaller number will result in faster generation speed and lower memory cost.
+- `separate_cfg_infer`: Whether to use separate inference process for CFG guidance. If set to True, memory cost will be lower. Default is True.
+- `use_kv_cache`: Whether to use key-value cache. Default is True.
+- `offload_kv_cache`:  offload the cached key and value to cpu, which can save memory but slow down the generation silightly. Default is True.
+- `offload_model`: offload the model to cpu, which can save memory but slow down the generation. Default is False.
+- `use_input_image_size_as_output`: whether to use the input image size as the output image size, which can be used for single-image input, e.g., image editing task. Default is False.
+- `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
+- For out of memory or time cost, you can refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources) to select a appropriate setting.
+- 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. 
+- 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>`. 
+- For image editing task and controlnet task, we recommend to set the height and width of output image as the same
+as input image. For example, if you want to edit a 512x512 image, you should set the height and width of output image as 512x512. You also can set the `use_input_image_size_as_output` to automatically set the height and width of output image as the same as input image.
+
+
+## Requiremented Resources
+
+We are currently experimenting with some techniques to reduce memory usage and improve speed, including `use_kv_cache, offload_kv_cache, separate_cfg_infer, offload_model`, which you can enable in the pipeline. 
+The default setting is`use_kv_cache=True, offload_kv_cache=True, separate_cfg_infer=True, offload_model=False`. 
+To reduce memory consumption while maintaining inference speed, quantization is also a method worth exploring and is left for future work.
+
+We conducted experiments on the A800 and RTX 3090. The memory requirements and inference times are shown in the table below. You can choose the appropriate settings based on your available resources.
+
+**Overall, the text-to-image task requires minimal memory and time costs, comparable to other latest text-to-image models. However, when using input images, the computational cost increases. Memory usage can be reduced by extending the processing time.**
+
+
+- Different image size. 
+
+Different image size (`max_input_image_size` is the max size of input image, `height` and `width` are the size of output image) with the default inference settings (`use_kv_cache=True,offload_kv_cache=True,separate_cfg_infer=True`)
+
+For A800 GPU:  
+| Settings     |  Only Text | Text + Single Image |  Text + Two Images    |
+|:-------------|:----------:|:-------------------:|:---------------------:|
+| max_input_image_size=1024,height=1024,width=1024 | 9G, 31s   | 12G, 1m6s  | 13G, 1m20s  |
+| max_input_image_size=512,height=1024,width=1024 | 9G, 31s   | 10G, 50s  | 10G, 54s |
+| max_input_image_size=768,height=768,width=768 | 9G, 16s  | 10G, 32s  | 10G, 37s  |
+| max_input_image_size=512,height=512,width=512 | 9G, 7s   | 9G, 14s  | 9G, 15s  |
+
+For RTX 3090 GPU(24G):
+| Settings     |  Only Text | Text + Single Image |  Text + Two Images    |
+|:-------------|:----------:|:-------------------:|:---------------------:|
+| max_input_image_size=1024,height=1024,width=1024 | 9G, 1m17s   | 12G, 2m46s  | 13G, 3m23s  |
+| max_input_image_size=512,height=1024,width=1024 | 9G, 1m18s   | 10G, 2m8s  | 10G, 2m18s |
+| max_input_image_size=768,height=768,width=768 | 9G, 41s  | 10G, 1m22s  | 10G, 1m38s  |
+| max_input_image_size=512,height=512,width=512 | 9G, 19s   | 9G, 36s  | 9G, 43s  |
+
+
+You can set smaller `max_input_image_size` to reduce memory usage, but note that the generation quality may be lower.
+And please set the `height` and `width` the same as the size of input image for image editing task.
+
+
+- Different inference settings
+
+Default image size: height=1024, width=1024, max_input_image_size=1024
+
+For A800 GPU:  
+| Settings     |  Only Text | Text + Single Image |  Text + Two Images    |
+|:-------------|:----------:|:-------------------:|:---------------------:|
+| use_kv_cache | 18G, 30s   | 36G, 1m  | 48G, 1m13s |
+| use_kv_cache,offload_kv_cache | 10G, 30s   | 14G, 1m10s | 17G, 1m30s  |
+| use_kv_cache,offload_kv_cache,separate_cfg_infer | 9G, 31s   | 12G, 1m6s  | 13G, 1m20s  |
+| use_kv_cache,offload_kv_cache,offload_model | 4G, 55s   | 7G, 1m30s  | 11G, 1m48s  |
+| use_kv_cache,offload_kv_cache,separate_cfg_infer,offload_model | 3G, 1m23s   | 5G, 2m19s   | 6G, 2m30s |
+
+For RTX 3090 GPU(24G):
+| Settings     |  Only Text | Text + Single Image |  Text + Two Images    |
+|:-------------|:----------:|:-------------------:|:---------------------:|
+| use_kv_cache | 18G, 1m14s   | OOM  | OOM  |
+| use_kv_cache,offload_kv_cache | 10G, 1m17s   | 14G, 3m11s | 17G, 4m3s  |
+| use_kv_cache,offload_kv_cache,separate_cfg_infer | 9G, 1m18s   | 12G, 2m46s  | 13G, 3m21s  |
+| use_kv_cache,offload_kv_cache,offload_model | 4G,3m1s    | 7G, 4m14s  | 11G, 5m4s  |
+| use_kv_cache,offload_kv_cache,separate_cfg_infer,offload_model | 3G, 4m56s   | 5G, 7m49s   | 6G, 8m6s |

requirements.txt

@@ -1,9 +1,11 @@
-accelerate
-diffusers
-invisible_watermark
-torch
-transformers
-xformers
-timm
-peft
-safetensors
+torch==2.3.1
+transformers==4.45.2
+datasets==2.20.0
+accelerate==0.26.1
+jupyter==1.0.0
+numpy==1.26.3
+pillow==10.2.0
+torch==2.3.1
+peft==0.9.0
+diffusers==0.30.3
+timm==0.9.16

setup.py

@@ -5,7 +5,7 @@ with open("README.md", mode="r", encoding="utf-8") as readme_file:
 
 setup(
     name='OmniGen',
-    version='1.0.0',
+    version='1.0.3',
     description='OmniGen',
     long_description=readme,
     long_description_content_type="text/markdown",
@@ -14,10 +14,13 @@ setup(
     packages=find_packages(),
     include_package_data=True,
     install_requires=[
-        'torch>=1.6.0',
-        'transformers>=4.41.0',
+        'torch==2.3.1',
+        'transformers==4.45.2',
         'datasets',
-        'accelerate>=0.20.1',
-        'diffusers>=0.30.3'
+        'accelerate==0.26.1',
+        'diffusers==0.30.3',
+        "timm",
+        "peft==0.9.0",
+        "safetensors"
     ],
-)
+)

train.py

@@ -0,0 +1,373 @@
+import json
+from time import time
+import argparse
+import logging
+import os
+from pathlib import Path
+import math
+
+import numpy as np
+from PIL import Image
+from copy import deepcopy
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torchvision import transforms
+
+from accelerate import Accelerator
+from accelerate.utils import ProjectConfiguration, set_seed
+from diffusers.optimization import get_scheduler
+from accelerate.utils import DistributedType
+from peft import LoraConfig, set_peft_model_state_dict, PeftModel, get_peft_model
+from peft.utils import get_peft_model_state_dict
+from huggingface_hub import snapshot_download
+from safetensors.torch import save_file
+
+from diffusers.models import AutoencoderKL
+
+from OmniGen import OmniGen, OmniGenProcessor
+from OmniGen.train_helper import DatasetFromJson, TrainDataCollator
+from OmniGen.train_helper import training_losses
+from OmniGen.utils import (
+    create_logger,
+    update_ema,
+    requires_grad,
+    center_crop_arr,
+    crop_arr,
+    vae_encode,
+    vae_encode_list
+)
+
+def main(args):
+    # Setup accelerator:
+    from accelerate import DistributedDataParallelKwargs as DDPK
+    kwargs = DDPK(find_unused_parameters=False)
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_dir=args.results_dir,
+        kwargs_handlers=[kwargs],
+        )
+    device = accelerator.device
+    accelerator.init_trackers("tensorboard_log", config=args.__dict__)
+
+    # Setup an experiment folder:
+    checkpoint_dir = f"{args.results_dir}/checkpoints"  # Stores saved model checkpoints
+    logger = create_logger(args.results_dir)
+    if accelerator.is_main_process:
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger.info(f"Experiment directory created at {args.results_dir}")
+        json.dump(args.__dict__, open(os.path.join(args.results_dir, 'train_args.json'), 'w'))
+
+
+    # Create model:    
+    if not os.path.exists(args.model_name_or_path):
+        cache_folder = os.getenv('HF_HUB_CACHE')
+        args.model_name_or_path = snapshot_download(repo_id=args.model_name_or_path,
+                                        cache_dir=cache_folder,
+                                        ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
+        logger.info(f"Downloaded model to {args.model_name_or_path}")
+    model = OmniGen.from_pretrained(args.model_name_or_path)
+    model.llm.config.use_cache = False
+    model.llm.gradient_checkpointing_enable()
+    model = model.to(device)
+
+    if args.vae_path is None:
+        print(args.model_name_or_path)
+        vae_path = os.path.join(args.model_name_or_path, "vae")
+        if os.path.exists(vae_path):
+            vae = AutoencoderKL.from_pretrained(vae_path).to(device)
+        else:
+            logger.info("No VAE found in model, downloading stabilityai/sdxl-vae from HF")
+            logger.info("If you have VAE in local folder, please specify the path with --vae_path")
+            vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae").to(device)
+    else:
+        vae = AutoencoderKL.from_pretrained(args.vae_path).to(device)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+    vae.to(dtype=torch.float32)
+    model.to(weight_dtype)
+
+    processor = OmniGenProcessor.from_pretrained(args.model_name_or_path)
+
+    requires_grad(vae, False)
+    if args.use_lora:
+        if accelerator.distributed_type == DistributedType.FSDP:
+            raise NotImplementedError("FSDP does not support LoRA")
+        requires_grad(model, False)
+        transformer_lora_config = LoraConfig(
+            r=args.lora_rank,
+            lora_alpha=args.lora_rank,
+            init_lora_weights="gaussian",
+            target_modules=["qkv_proj", "o_proj"],
+        )
+        model.llm.enable_input_require_grads()
+        model = get_peft_model(model, transformer_lora_config)
+        model.to(weight_dtype)
+        transformer_lora_parameters = list(filter(lambda p: p.requires_grad, model.parameters()))
+        for n,p in model.named_parameters():
+            print(n, p.requires_grad)
+        opt = torch.optim.AdamW(transformer_lora_parameters, lr=args.lr, weight_decay=args.adam_weight_decay)
+    else:
+        opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.adam_weight_decay)
+
+    ema = None
+    if args.use_ema:
+        ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+        requires_grad(ema, False)
+    
+
+    # Setup data:
+    crop_func = crop_arr
+    if not args.keep_raw_resolution:
+        crop_func = center_crop_arr
+    image_transform = transforms.Compose([
+        transforms.Lambda(lambda pil_image: crop_func(pil_image, args.max_image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    ])
+    
+    dataset = DatasetFromJson(json_file=args.json_file, 
+    image_path=args.image_path,
+    processer=processor,
+    image_transform=image_transform,
+    max_input_length_limit=args.max_input_length_limit,
+    condition_dropout_prob=args.condition_dropout_prob,
+    keep_raw_resolution=args.keep_raw_resolution
+    )
+    collate_fn = TrainDataCollator(pad_token_id=processor.text_tokenizer.eos_token_id, hidden_size=model.llm.config.hidden_size, keep_raw_resolution=args.keep_raw_resolution)
+
+    loader = DataLoader(
+        dataset,
+        collate_fn=collate_fn,
+        batch_size=args.batch_size_per_device,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        prefetch_factor=2,
+    )
+    
+    if accelerator.is_main_process:
+        logger.info(f"Dataset contains {len(dataset):,}")
+
+    num_update_steps_per_epoch = math.ceil(len(loader) / args.gradient_accumulation_steps)
+    max_train_steps = args.epochs * num_update_steps_per_epoch
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=opt,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=max_train_steps * args.gradient_accumulation_steps,
+    )
+
+    # Prepare models for training:
+    model.train()  # important! This enables embedding dropout for classifier-free guidance
+    
+    if ema is not None:
+        update_ema(ema, model, decay=0)  # Ensure EMA is initialized with synced weights
+        ema.eval()  # EMA model should always be in eval mode
+    
+
+    if ema is not None:
+        model, ema = accelerator.prepare(model, ema)
+    else:
+        model = accelerator.prepare(model)
+
+    opt, loader, lr_scheduler = accelerator.prepare(opt, loader, lr_scheduler)
+    
+    
+    # Variables for monitoring/logging purposes:
+    train_steps, log_steps = 0, 0
+    running_loss = 0
+    start_time = time()
+    
+    if accelerator.is_main_process:
+        logger.info(f"Training for {args.epochs} epochs...")
+    for epoch in range(args.epochs):
+        if accelerator.is_main_process:
+            logger.info(f"Beginning epoch {epoch}...")
+        
+        for data in loader:
+            with accelerator.accumulate(model):
+                with torch.no_grad():
+                    output_images = data['output_images']
+                    input_pixel_values = data['input_pixel_values']
+                    if isinstance(output_images, list):
+                        output_images = vae_encode_list(vae, output_images, weight_dtype)
+                        if input_pixel_values is not None:
+                            input_pixel_values = vae_encode_list(vae, input_pixel_values, weight_dtype)
+                    else:
+                        output_images = vae_encode(vae, output_images, weight_dtype)
+                        if input_pixel_values is not None:
+                            input_pixel_values = vae_encode(vae, input_pixel_values, weight_dtype)
+                   
+
+                model_kwargs = dict(input_ids=data['input_ids'], input_img_latents=input_pixel_values, input_image_sizes=data['input_image_sizes'], attention_mask=data['attention_mask'], position_ids=data['position_ids'], padding_latent=data['padding_images'], past_key_values=None, return_past_key_values=False)
+                
+                loss_dict = training_losses(model, output_images, model_kwargs)
+                loss = loss_dict["loss"].mean()
+
+                running_loss += loss.item()
+                accelerator.backward(loss)
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+                opt.step()
+                lr_scheduler.step()
+                opt.zero_grad()
+
+                log_steps += 1
+                train_steps += 1
+
+                accelerator.log({"training_loss": loss.item()}, step=train_steps)
+                if train_steps % args.gradient_accumulation_steps == 0:
+                    if accelerator.sync_gradients and ema is not None: 
+                        update_ema(ema, model)
+                    
+                if train_steps % (args.log_every * args.gradient_accumulation_steps) == 0 and train_steps > 0:
+                    torch.cuda.synchronize()
+                    end_time = time()
+                    steps_per_sec = log_steps / args.gradient_accumulation_steps / (end_time - start_time)
+                    # Reduce loss history over all processes:
+                    avg_loss = torch.tensor(running_loss / log_steps, device=device)
+                    dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+                    avg_loss = avg_loss.item() / accelerator.num_processes 
+                        
+                    if accelerator.is_main_process:
+                        cur_lr = opt.param_groups[0]["lr"]
+                        logger.info(f"(step={int(train_steps/args.gradient_accumulation_steps):07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}, Epoch: {train_steps/len(loader)}, LR: {cur_lr}")
+
+                    # Reset monitoring variables:
+                    running_loss = 0
+                    log_steps = 0
+                    start_time = time()
+
+
+            if train_steps % (args.ckpt_every * args.gradient_accumulation_steps) == 0 and train_steps > 0:
+                if accelerator.distributed_type == DistributedType.FSDP:
+                    state_dict = accelerator.get_state_dict(model)
+                    ema_state_dict = accelerator.get_state_dict(ema) if ema is not None else None
+                else:
+                    if not args.use_lora:
+                        state_dict = model.module.state_dict()
+                        ema_state_dict = accelerator.get_state_dict(ema) if ema is not None else None
+
+                if accelerator.is_main_process:
+                    if args.use_lora:
+                        checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}/"
+                        os.makedirs(checkpoint_path, exist_ok=True)
+
+                        model.module.save_pretrained(checkpoint_path)
+                    else:
+                        checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}/"
+                        os.makedirs(checkpoint_path, exist_ok=True)
+                        torch.save(state_dict, os.path.join(checkpoint_path, "model.pt"))
+                        processor.text_tokenizer.save_pretrained(checkpoint_path)
+                        model.llm.config.save_pretrained(checkpoint_path)
+                        if ema_state_dict is not None:
+                            checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}_ema"
+                            os.makedirs(checkpoint_path, exist_ok=True)
+                            torch.save(state_dict, os.path.join(checkpoint_path, "model.pt"))
+                            processor.text_tokenizer.save_pretrained(checkpoint_path)
+                            model.llm.config.save_pretrained(checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+
+            dist.barrier()
+    accelerator.end_training()
+    model.eval()  
+    
+    if accelerator.is_main_process:
+        logger.info("Done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--results_dir", type=str, default="results")
+    parser.add_argument("--model_name_or_path", type=str, default="OmniGen")
+    parser.add_argument("--json_file", type=str)
+    parser.add_argument("--image_path", type=str, default=None)
+    parser.add_argument("--epochs", type=int, default=1400)
+    parser.add_argument("--batch_size_per_device", type=int, default=1)
+    parser.add_argument("--vae_path", type=str, default=None) 
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--log_every", type=int, default=100)
+    parser.add_argument("--ckpt_every", type=int, default=20000)
+    parser.add_argument("--max_grad_norm", type=float, default=1.0)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--max_input_length_limit", type=int, default=1024)
+    parser.add_argument("--condition_dropout_prob", type=float, default=0.1)
+    parser.add_argument("--adam_weight_decay", type=float, default=0.0)
+    parser.add_argument(
+        "--keep_raw_resolution",
+        action="store_true",
+        help="multiple_resolutions",
+    )
+    parser.add_argument("--max_image_size", type=int, default=1344)
+
+    parser.add_argument(
+            "--use_lora",
+            action="store_true",
+        )
+    parser.add_argument(
+            "--lora_rank",
+            type=int, 
+            default=8
+        )
+
+    parser.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Whether or not to use ema.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    ) 
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=1000, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="bf16",
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+
+
+    args = parser.parse_args()
+    assert args.max_image_size % 16 == 0, "Image size must be divisible by 16."
+
+    main(args)
+
+