more examples

OmniGen/pipeline.py

@@ -16,6 +16,7 @@ from diffusers.utils import (
     scale_lora_layers,
     unscale_lora_layers,
 )
+from safetensors.torch import load_file
 
 from OmniGen import OmniGen, OmniGenProcessor, OmniGenScheduler
 
@@ -59,12 +60,12 @@ class OmniGenPipeline:
 
     @classmethod
     def from_pretrained(cls, model_name, vae_path: str=None):
-        if not os.path.exists(model_name):
+        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'])
+                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5', 'model.pt'])
             logger.info(f"Downloaded model to {model_name}")
         model = OmniGen.from_pretrained(model_name)
         processor = OmniGenProcessor.from_pretrained(model_name)
@@ -82,6 +83,8 @@ class OmniGenPipeline:
     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]):

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()))))

app.py

@@ -11,7 +11,7 @@ pipe = OmniGenPipeline.from_pretrained(
 
 @spaces.GPU(duration=180)
 # 示例处理函数：生成图像
-def generate_image(text, img1, img2, img3, height, width, guidance_scale, inference_steps, seed):
+def generate_image(text, img1, img2, img3, height, width, guidance_scale, img_guidance_scale, inference_steps, seed):
     input_images = [img1, img2, img3]
     # 去除 None
     input_images = [img for img in input_images if img is not None]
@@ -26,7 +26,7 @@ def generate_image(text, img1, img2, img3, height, width, guidance_scale, infere
         guidance_scale=guidance_scale,
         img_guidance_scale=1.6,
         num_inference_steps=inference_steps,
-        separate_cfg_infer=True,
+        separate_cfg_infer=True, # set False can speed up the inference process
         use_kv_cache=False,
         seed=seed,
     )
@@ -47,26 +47,28 @@ def generate_image(text, img1, img2, img3, height, width, guidance_scale, infere
 def get_example():
     case = [
         [
-            "A vintage camera placed on the ground, ejecting a swirling cloud of Polaroid-style photographs into the air. The photos, showing landscapes, wildlife, and travel scenes, seem to defy gravity, floating upward in a vortex of motion. The camera emits a glowing, smoky light from within, enhancing the magical, surreal atmosphere. The dark background contrasts with the illuminated photos and camera, creating a dreamlike, nostalgic scene filled with vibrant colors and dynamic movement. Scattered photos are visible on the ground, further contributing to the idea of an explosion of captured memories.",
+            "A curly-haired man in a red shirt is drinking tea.",
             None,
             None,
             None,
             1024,
             1024,
             2.5,
+            1.6,
             50,
             0,
         ],
         [
-            "A woman <img><|image_1|></img> in a wedding dress. Next to her is a black-haired man.",
+            "The woman in <img><|image_1|></img> waves her hand happily in the crowd",
             "./imgs/test_cases/yifei2.png",
             None,
             None,
             1024,
             1024,
             2.5,
+            1.9,
             50,
-            0,
+            128,
         ],
         [
             "A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.",
@@ -76,17 +78,55 @@ def get_example():
             1024,
             1024,
             2.5,
+            1.6,
             50,
             0,
         ],
         [
-            "Two men are celebrating with raised glasses in a restaurant. A man is <img><|image_1|></img>. The other man is <img><|image_2|></img>.",
-            "./imgs/test_cases/young_musk.jpg",
-            "./imgs/test_cases/young_trump.jpeg",
+            "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,
         ],
@@ -98,6 +138,7 @@ def get_example():
             1024,
             1024,
             2.5,
+            1.6,
             50,
             222,
         ],
@@ -109,6 +150,7 @@ def get_example():
             1024,
             1024,
             2.0,
+            1.6,
             50,
             0,
         ],
@@ -120,6 +162,7 @@ def get_example():
             1024,
             1024,
             2,
+            1.6,
             50,
             42,
         ],
@@ -131,9 +174,22 @@ def get_example():
             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",
@@ -142,25 +198,27 @@ def get_example():
             1024,
             1024,
             2.5,
+            1.6,
             50,
             0,
         ],
         [
-            "Three guitars are displayed side by side on a rustic wooden stage, each showcasing its unique character and style. The left guitar is <img><|image_1|><\/img>. The middle guitar is <img><|image_2|><\/img>. The right guitars is <img><|image_3|><\/img>.",
-            "./imgs/test_cases/guitar1.png",
-            "./imgs/test_cases/guitar1.png",
-            "./imgs/test_cases/guitar1.png",
+            "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,
-            0,
+            1,
         ],
     ]
     return case
 
-def run_for_examples(text, img1, img2, img3, height, width, guidance_scale, inference_steps, seed):    
-    return generate_image(text, img1, img2, img3, height, width, guidance_scale, inference_steps, seed)
+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.
@@ -168,6 +226,13 @@ OmniGen is a unified image generation model that you can use to perform various
 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 接口
@@ -197,7 +262,11 @@ with gr.Blocks() as demo:
 
             # 引导尺度输入
             guidance_scale_input = gr.Slider(
-                label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1
+                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(
@@ -226,6 +295,7 @@ with gr.Blocks() as demo:
             height_input,
             width_input,
             guidance_scale_input,
+            img_guidance_scale_input,
             num_inference_steps,
             seed_input,
         ],
@@ -243,6 +313,7 @@ with gr.Blocks() as demo:
             height_input,
             width_input,
             guidance_scale_input,
+            img_guidance_scale_input,
             num_inference_steps,
             seed_input,
         ],