update omnigen

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/omnigen/processor.py

@@ -1,30 +1,16 @@
 import os
 import re
 from typing import Dict, List
-import json
-
 import torch
-import numpy as np
-import random
-from PIL import Image
 from torchvision import transforms
 from transformers import AutoTokenizer
 from huggingface_hub import snapshot_download
-
-from .utils import (
-    create_logger,
-    update_ema,
-    requires_grad,
-    center_crop_arr,
-    crop_arr,
-)
-
-
+from .utils import crop_arr
 
 
 class OmniGenProcessor:
-    def __init__(self, 
-                text_tokenizer, 
+    def __init__(self,
+                text_tokenizer,
                 max_image_size: int=1024):
         self.text_tokenizer = text_tokenizer
         self.max_image_size = max_image_size
@@ -52,7 +38,7 @@ class OmniGenProcessor:
 
     def process_image(self, image):
         return self.image_transform(image)
-    
+
     def process_multi_modal_prompt(self, text, input_images):
         text = self.add_prefix_instruction(text)
         if input_images is None or len(input_images) == 0:
@@ -60,25 +46,25 @@ class OmniGenProcessor:
             return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None}
 
         pattern = r"<\|image_\d+\|>"
-        prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)] 
+        prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)]
 
         for i in range(1, len(prompt_chunks)):
             if prompt_chunks[i][0] == 1:
                 prompt_chunks[i] = prompt_chunks[i][1:]
 
-        image_tags = re.findall(pattern, text) 
+        image_tags = re.findall(pattern, text)
         image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
 
         unique_image_ids = sorted(list(set(image_ids)))
         assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}"
         # total images must be the same as the number of image tags
         assert len(unique_image_ids) == len(input_images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images"
-        
+
         input_images = [input_images[x-1] for x in image_ids]
 
         all_input_ids = []
         img_inx = []
-        idx = 0
+        _idx = 0
         for i in range(len(prompt_chunks)):
             all_input_ids.extend(prompt_chunks[i])
             if i != len(prompt_chunks) -1:
@@ -99,8 +85,8 @@ class OmniGenProcessor:
         return prompt
 
 
-    def __call__(self, 
-                instructions: List[str], 
+    def __call__(self,
+                instructions: List[str],
                 input_images: List[List[str]] = None,
                 height: int = 1024,
                 width: int = 1024,
@@ -114,7 +100,7 @@ class OmniGenProcessor:
         if isinstance(instructions, str):
             instructions = [instructions]
             input_images = [input_images]
-        
+
         input_data = []
         for i in range(len(instructions)):
             cur_instruction = instructions[i]
@@ -124,10 +110,9 @@ class OmniGenProcessor:
             else:
                 cur_input_images = None
                 assert "<img><|image_1|></img>" not in cur_instruction
-            
+
             mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images)
 
-        
             neg_mllm_input, img_cfg_mllm_input = None, None
             neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None)
             if use_img_cfg:
@@ -144,23 +129,21 @@ class OmniGenProcessor:
         return self.collator(input_data)
 
 
-
-
 class OmniGenCollator:
     def __init__(self, pad_token_id=2, hidden_size=3072):
         self.pad_token_id = pad_token_id
         self.hidden_size = hidden_size
-    
+
     def create_position(self, attention_mask, num_tokens_for_output_images):
         position_ids = []
         text_length = attention_mask.size(-1)
-        img_length = max(num_tokens_for_output_images)  
+        img_length = max(num_tokens_for_output_images)
         for mask in attention_mask:
             temp_l = torch.sum(mask)
             temp_position = [0]*(text_length-temp_l) + [i for i in range(temp_l+img_length+1)] # we add a time embedding into the sequence, so add one more token
             position_ids.append(temp_position)
         return torch.LongTensor(position_ids)
-    
+
     def create_mask(self, attention_mask, num_tokens_for_output_images):
         extended_mask = []
         padding_images = []
@@ -194,24 +177,24 @@ class OmniGenCollator:
                 temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size))
             else:
                 temp_padding_imgs = None
-            
+
             extended_mask.append(temp_mask.unsqueeze(0))
             padding_images.append(temp_padding_imgs)
             inx += 1
         return torch.cat(extended_mask, dim=0), padding_images
-    
+
     def adjust_attention_for_input_images(self, attention_mask, image_sizes):
         for b_inx in image_sizes.keys():
             for start_inx, end_inx in image_sizes[b_inx]:
                 attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1
 
         return attention_mask
-    
+
     def pad_input_ids(self, input_ids, image_sizes):
         max_l = max([len(x) for x in input_ids])
         padded_ids = []
         attention_mask = []
-        new_image_sizes = []
+        _new_image_sizes = []
 
         for i in range(len(input_ids)):
             temp_ids = input_ids[i]
@@ -223,7 +206,7 @@ class OmniGenCollator:
             else:
                 attention_mask.append([0]*pad_l+[1]*temp_l)
                 padded_ids.append([self.pad_token_id]*pad_l+temp_ids)
-            
+
             if i in image_sizes:
                 new_inx = []
                 for old_inx in image_sizes[i]:
@@ -248,10 +231,9 @@ class OmniGenCollator:
                         image_sizes[b_inx] = [size]
                     else:
                         image_sizes[b_inx].append(size)
-            b_inx += 1     
+            b_inx += 1
         pixel_values = [x.unsqueeze(0) for x in pixel_values]
 
-        
         input_ids = [x['input_ids'] for x in mllm_inputs]
         padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes)
         position_ids = self.create_position(attention_mask, num_tokens_for_output_images)
@@ -259,15 +241,13 @@ class OmniGenCollator:
         attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes)
 
         return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes
-    
-    
+
     def __call__(self, features):
         mllm_inputs = [f[0] for f in features]
         cfg_mllm_inputs = [f[1] for f in features]
         img_cfg_mllm_input = [f[2] for f in features]
         target_img_size = [f[3] for f in features]
 
-        
         if img_cfg_mllm_input[0] is not None:
             mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input
             target_img_size = target_img_size + target_img_size + target_img_size
@@ -295,10 +275,8 @@ class OmniGenSeparateCollator(OmniGenCollator):
         img_cfg_mllm_input = [f[2] for f in features]
         target_img_size = [f[3] for f in features]
 
-        
         all_padded_input_ids, all_attention_mask, all_position_ids, all_pixel_values, all_image_sizes, all_padding_images = [], [], [], [], [], []
 
-
         padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
         all_padded_input_ids.append(padded_input_ids)
         all_attention_mask.append(attention_mask)