update omnigen

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

CHANGELOG.md

@@ -1,8 +1,8 @@
 # Change Log for SD.Next
 
-## Update for 2024-10-22
+## Update for 2024-10-23
 
-### Highlights for 2024-10-22
+### Highlights for 2024-10-23
 
 A month later and with nearly 300 commits, here is the latest [SD.Next](https://github.com/vladmandic/automatic) update!  
 
@@ -32,13 +32,13 @@ A month later and with nearly 300 commits, here is the latest [SD.Next](https://
 
 #### Otherwise notable
 
-- Several of [Flux.1](https://huggingface.co/black-forest-labs/FLUX.1-dev) optimizations and new quantization types  
+- Tons of work on **dynamic quantization** that can be applied *on-the-fly* during model load to any model type (*you do not need to use pre-quantized models*)  
+  Supported quantization engines include `BitsAndBytes`, `TorchAO`, `Optimum.quanto`, `NNCF` compression, and more...  
 - Auto-detection of best available **device/dtype** settings for your platform and GPU reduces neeed for manual configuration  
   *Note*: This is a breaking change to default settings and its recommended to check your preferred settings after upgrade  
 - Full rewrite of **sampler options**, not far more streamlined with tons of new options to tweak scheduler behavior  
 - Improved **LoRA** detection and handling for all supported models  
-- Tons of work on **dynamic quantization** that can be applied *on-the-fly* during model load to any model type (*you do not need to use pre-quantized models*)  
-  Supported quantization engines include `BitsAndBytes`, `TorchAO`, `Optimum.quanto`, `NNCF` compression, and more...  
+- Several of [Flux.1](https://huggingface.co/black-forest-labs/FLUX.1-dev) optimizations and new quantization types  
 
 Oh, and we've compiled a full table with list of top-30 (*how many have you tried?*) popular text-to-image generative models,  
 their respective parameters and architecture overview: [Models Overview](https://github.com/vladmandic/automatic/wiki/Models)  
@@ -47,7 +47,7 @@ And there are also other goodies like multiple *XYZ grid* improvements, addition
 [README](https://github.com/vladmandic/automatic/blob/master/README.md) | [CHANGELOG](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867)
 
 
-### Details for 2024-10-22
+### Details for 2024-10-23
 
 - **reprocess**
   - new top-level button: reprocess latent from your history of generated image(s)  

modules/omnigen/model.py

@@ -1,21 +1,18 @@
 # The code is revised from DiT
 import os
+import math
 import torch
 import torch.nn as nn
 import numpy as np
-import math
-from typing import Dict
-
+from safetensors.torch import load_file
 from diffusers.loaders import PeftAdapterMixin
-from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
 from huggingface_hub import snapshot_download
-
 from .transformer import Phi3Config, Phi3Transformer
 
 
 def modulate(x, shift, scale):
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
- 
+
 
 class TimestepEmbedder(nn.Module):
     """
@@ -165,37 +162,35 @@ class OmniGen(nn.Module, PeftAdapterMixin):
         self.out_channels = in_channels
         self.patch_size = patch_size
         self.pos_embed_max_size = pos_embed_max_size
-
         hidden_size = transformer_config.hidden_size
-
         self.x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
         self.input_x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
-
         self.time_token = TimestepEmbedder(hidden_size)
         self.t_embedder = TimestepEmbedder(hidden_size)
-        
         self.pe_interpolation = pe_interpolation
         pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, interpolation_scale=self.pe_interpolation, base_size=64)
         self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=True)
-
         self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
-
         self.initialize_weights()
-
         self.llm = Phi3Transformer(config=transformer_config)
         self.llm.config.use_cache = False
-    
+
     @classmethod
-    def from_pretrained(cls, model_name):
-        if not os.path.exists(os.path.join(model_name, 'model.pt')):
-            cache_folder = os.getenv('HF_HUB_CACHE')
+    def from_pretrained(cls, model_name: str, cache_dir: str=None):
+        if not os.path.exists(os.path.join(model_name, 'model.pt')) and not os.path.exists(os.path.join(model_name, 'model.safetensors')):
+            cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
             model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_folder,
+                                           cache_dir=cache_dir,
                                            ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
         config = Phi3Config.from_pretrained(model_name)
         model = cls(config)
-        ckpt = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
-        model.load_state_dict(ckpt)
+        if os.path.exists(os.path.join(model_name, 'model.pt')):
+            state_dict = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
+        elif os.path.exists(os.path.join(model_name, 'model.safetensors')):
+            state_dict = load_file(os.path.join(model_name, 'model.safetensors'))
+        else:
+            raise ValueError(f"OmniGen: Could not find model file in {model_name}")
+        model.load_state_dict(state_dict)
         return model
 
     def initialize_weights(self):
@@ -208,7 +203,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
                 if module.bias is not None:
                     nn.init.constant_(module.bias, 0)
         self.apply(_basic_init)
-        
+
         # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
         w = self.x_embedder.proj.weight.data
         nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
@@ -282,7 +277,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
                     latent = self.input_x_embedder(latent)
                 else:
                     latent = self.x_embedder(latent)
-                pos_embed = self.cropped_pos_embed(height, width)    
+                pos_embed = self.cropped_pos_embed(height, width)
                 latent = latent + pos_embed
                 if padding is not None:
                     latent = torch.cat([latent, padding], dim=-2)
@@ -300,21 +295,16 @@ class OmniGen(nn.Module, PeftAdapterMixin):
                 latents = self.input_x_embedder(latents)
             else:
                 latents = self.x_embedder(latents)
-            pos_embed = self.cropped_pos_embed(height, width)  
+            pos_embed = self.cropped_pos_embed(height, width)
             latents = latents + pos_embed
             num_tokens = latents.size(1)
             shapes = [height, width]
         return latents, num_tokens, shapes
 
-    
     def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True):
-        """
-        
-        """
         input_is_list = isinstance(x, list)
         x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent)
-        time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1)   
-        
+        time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1)
         if input_img_latents is not None:
             input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True)
         if input_ids is not None:
@@ -325,7 +315,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
                     condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx]
                     input_img_inx += 1
             if input_img_latents is not None:
-                assert input_img_inx == len(input_latents) 
+                assert input_img_inx == len(input_latents)
 
             input_emb = torch.cat([condition_embeds, time_token, x], dim=1)
         else:
@@ -355,7 +345,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
     def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache):
         """
         Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """        
+        """
         self.llm.config.use_cache = use_kv_cache
         model_out, past_key_values = self.forward(x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, past_key_values=past_key_values, return_past_key_values=True)
         if use_img_cfg:
@@ -366,7 +356,6 @@ class OmniGen(nn.Module, PeftAdapterMixin):
             cond, uncond = torch.split(model_out, len(model_out) // 2, dim=0)
             cond = uncond + cfg_scale * (cond - uncond)
             model_out = [cond, cond]
-        
         return torch.cat(model_out, dim=0), past_key_values
 
 
@@ -374,7 +363,7 @@ class OmniGen(nn.Module, PeftAdapterMixin):
     def forward_with_separate_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, return_past_key_values=True):
         """
         Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """        
+        """
         self.llm.config.use_cache = use_kv_cache
         if past_key_values is None:
             past_key_values = [None] * len(attention_mask)
@@ -399,9 +388,4 @@ class OmniGen(nn.Module, PeftAdapterMixin):
             model_out = [cond, cond]
         else:
             return model_out[0]
-        
         return torch.cat(model_out, dim=0), pask_key_values
-
-
-
-

modules/omnigen/pipeline.py

@@ -40,9 +40,9 @@ class OmniGenPipeline():
         self.vae = vae
         self.model = model
         self.processor = processor
-        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-        self.dtype: torch.dtype = torch.bfloat16,
-        self.separate_cfg_infer: bool = True,
+        self.device = None
+        self.dtype: None
+        self.separate_cfg_infer: bool = True
         self.use_kv_cache: bool = False
         # omnigen does not inherit from diffusionpipeline so we hack it
         self._internal_dict = { # pylint: disable=protected-access
@@ -62,10 +62,8 @@ class OmniGenPipeline():
         processor = OmniGenProcessor.from_pretrained(model_name)
         if os.path.exists(os.path.join(model_name, "vae")):
             vae = AutoencoderKL.from_pretrained(os.path.join(model_name, "vae"))
-        elif vae_path is not None:
-            vae = AutoencoderKL.from_pretrained(vae_path)
         else:
-            vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
+            vae = AutoencoderKL.from_pretrained(vae_path or "stabilityai/sdxl-vae")
         return cls(vae, model, processor)
 
     def merge_lora(self, lora_path: str):

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)