From bf6f130f38405dc265ac0804dbb85b000a51feb7 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 23 Oct 2024 07:04:51 -0400 Subject: [PATCH] update omnigen Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 12 +++---- modules/omnigen/model.py | 60 ++++++++++++-------------------- modules/omnigen/pipeline.py | 10 +++--- modules/omnigen/processor.py | 66 ++++++++++++------------------------ 4 files changed, 54 insertions(+), 94 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a85098504..586abb9d6 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/modules/omnigen/model.py b/modules/omnigen/model.py index 7446ef080..3a42263d2 100644 --- a/modules/omnigen/model.py +++ b/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 - - - - diff --git a/modules/omnigen/pipeline.py b/modules/omnigen/pipeline.py index c15fd123d..e17afea5c 100644 --- a/modules/omnigen/pipeline.py +++ b/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): diff --git a/modules/omnigen/processor.py b/modules/omnigen/processor.py index 071f28423..ada813a8b 100644 --- a/modules/omnigen/processor.py +++ b/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 "<|image_1|>" 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)