diff --git a/.pylintrc b/.pylintrc
index c7d03d0c3..45869a8c3 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -30,6 +30,7 @@ ignore-paths=/usr/lib/.*$,
modules/vdm,
modules/xadapter,
modules/meissonic,
+ modules/omnigen,
repositories,
extensions-builtin/sd-webui-agent-scheduler,
extensions-builtin/sd-extension-chainner/nodes,
diff --git a/.ruff.toml b/.ruff.toml
index 3bc6de045..fe8ac4f87 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -25,6 +25,7 @@ exclude = [
"modules/vdm",
"modules/xadapter",
"modules/meissonic",
+ "modules/omnigen",
"repositories",
"extensions-builtin/sd-extension-chainner/nodes",
"extensions-builtin/sd-webui-agent-scheduler",
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6b5cd1d98..6dc018e31 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,8 @@
# Change Log for SD.Next
-## Update for 2024-10-21
+## Update for 2024-10-22
-### Highlights for 2024-10-21
+### Highlights for 2024-10-22
#### Workflow highlights
@@ -19,6 +19,7 @@
- New fine-tuned [CLiP-ViT-L]((https://huggingface.co/zer0int/CLIP-GmP-ViT-L-14)) 1st stage **text-encoders** used by SD15, SDXL, Flux.1, etc. brings additional details to your images
- New models:
+ [OmniGen](https://arxiv.org/pdf/2409.11340)
[CogView 3 Plus](https://huggingface.co/THUDM/CogView3-Plus-3B)
[Meissonic](https://github.com/viiika/Meissonic)
- Additional integration:
@@ -42,7 +43,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-21
+### Details for 2024-10-22
- **reprocess**
- new top-level button: reprocess latent from your history of generated image(s)
@@ -167,6 +168,16 @@ And there are also other goodies like multiple *XYZ grid* improvements, addition
- added native load mode for qint8/qint4 models
- avoid unet load if unchanged
+- [OmniGen](https://arxiv.org/pdf/2409.11340)
+ - Radical new model with pure LLM architecture based on Phi-3
+ - Select from *networks -> models -> reference*
+ - Can be used for text-to-image and image-to-image
+ - Image-to-image is *very* different, you need to specify in prompt what do you want to do
+ and add `|image|` placeholder where input image is used!
+ examples: `in |image| remove glasses from face`, `using depth map from |image|, create new image of a cute robot`
+ - Params used: prompt, steps, guidance scale for prompt guidance, refine guidance scale for image guidance
+ Recommended: guidance=3.0, refine-guidance=1.6
+
- [CogView 3 Plus](https://huggingface.co/THUDM/CogView3-Plus-3B)
- Select from *networks -> models -> reference*
- resolution width and height can be from 512px to 2048px and must be divisible by 32
diff --git a/html/reference.json b/html/reference.json
index 2cf490c00..0a44f130d 100644
--- a/html/reference.json
+++ b/html/reference.json
@@ -10,19 +10,19 @@
"path": "juggernautXL_juggXIByRundiffusion.safetensors@https://civitai.com/api/download/models/782002",
"preview": "juggernautXL_v9Rundiffusionphoto2.jpg",
"desc": "Showcase finetuned model based on Stable diffusion XL",
- "extras": "width: 1024, height: 1024, sampler: DEIS, steps: 20, cfg_scale: 6.0"
+ "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0"
},
"Juggernaut SD-XL X Hyper": {
"path": "Juggernaut_X_RunDiffusion_Hyper.safetensors@https://civitai.com/api/download/models/471120",
"preview": "juggernautXL_v9Rundiffusionphoto2.jpg",
"desc": "Showcase finetuned model based on Stable diffusion XL",
- "extras": "width: 1024, height: 1024, sampler: DEIS, steps: 20, cfg_scale: 6.0"
+ "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0"
},
"Juggernaut SD-XL IX Lightning": {
"path": "juggernautXL_v9Rdphoto2Lightning.safetensors@https://civitai.com/api/download/models/357609",
"preview": "juggernautXL_v9Rdphoto2Lightning.jpg",
"desc": "Showcase finetuned model based on Stable diffusion XL",
- "extras": "width: 1024, height: 1024, sampler: DPM SDE, steps: 6, cfg_scale: 2.0"
+ "extras": "sampler: DPM SDE, steps: 6, cfg_scale: 2.0"
},
"Juggernaut SD Reborn": {
"original": true,
@@ -49,7 +49,7 @@
"path": "dreamshaperXL_v21TurboDPMSDE.safetensors@https://civitai.com/api/download/models/351306",
"preview": "dreamshaperXL_v21TurboDPMSDE.jpg",
"desc": "Showcase finetuned model based on Stable diffusion XL",
- "extras": "width: 1024, height: 1024, sampler: DPM SDE, steps: 8, cfg_scale: 2.0"
+ "extras": "sampler: DPM SDE, steps: 8, cfg_scale: 2.0"
},
"SDXS DreamShaper 512": {
@@ -93,7 +93,7 @@
"path": "sd_xl_base_1.0.safetensors@https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors?download=true",
"preview": "sd_xl_base_1.0.jpg",
"desc": "Stable Diffusion XL (SDXL) is the latest AI image generation model that is tailored towards more photorealistic outputs with more detailed imagery and composition compared to previous SD models, including SD 2.1. It can make realistic faces, legible text within the images, and better image composition, all while using shorter and simpler prompts at a greatly increased base resolution of 1024x1024. Just like its predecessors, SDXL has the ability to generate image variations using image-to-image prompting, inpainting (reimagining of the selected parts of an image), and outpainting (creating new parts that lie outside the image borders).",
- "extras": "width: 1024, height: 1024, sampler: DEIS, steps: 20, cfg_scale: 6.0"
+ "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0"
},
"StabilityAI Stable Cascade": {
"path": "huggingface/stabilityai/stable-cascade",
@@ -101,7 +101,7 @@
"variant": "bf16",
"desc": "Stable Cascade is a diffusion model built upon the Würstchen architecture and its main difference to other models like Stable Diffusion is that it is working at a much smaller latent space. Why is this important? The smaller the latent space, the faster you can run inference and the cheaper the training becomes. How small is the latent space? Stable Diffusion uses a compression factor of 8, resulting in a 1024x1024 image being encoded to 128x128. Stable Cascade achieves a compression factor of 42, meaning that it is possible to encode a 1024x1024 image to 24x24, while maintaining crisp reconstructions. The text-conditional model is then trained in the highly compressed latent space. Previous versions of this architecture, achieved a 16x cost reduction over Stable Diffusion 1.5",
"preview": "stabilityai--stable-cascade.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 4.0, image_cfg_scale: 1.0"
+ "extras": "sampler: Default, cfg_scale: 4.0, image_cfg_scale: 1.0"
},
"StabilityAI Stable Cascade Lite": {
"path": "huggingface/stabilityai/stable-cascade-lite",
@@ -109,7 +109,7 @@
"variant": "bf16",
"desc": "Stable Cascade is a diffusion model built upon the Würstchen architecture and its main difference to other models like Stable Diffusion is that it is working at a much smaller latent space. Why is this important? The smaller the latent space, the faster you can run inference and the cheaper the training becomes. How small is the latent space? Stable Diffusion uses a compression factor of 8, resulting in a 1024x1024 image being encoded to 128x128. Stable Cascade achieves a compression factor of 42, meaning that it is possible to encode a 1024x1024 image to 24x24, while maintaining crisp reconstructions. The text-conditional model is then trained in the highly compressed latent space. Previous versions of this architecture, achieved a 16x cost reduction over Stable Diffusion 1.5",
"preview": "stabilityai--stable-cascade.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 4.0, image_cfg_scale: 1.0"
+ "extras": "sampler: Default, cfg_scale: 4.0, image_cfg_scale: 1.0"
},
"StabilityAI Stable Diffusion 3 Medium": {
"path": "huggingface/stabilityai/stable-diffusion-3-medium-diffusers",
@@ -118,7 +118,7 @@
"te3": null,
"desc": "Stable Diffusion 3 Medium is a Multimodal Diffusion Transformer (MMDiT) text-to-image model that features greatly improved performance in image quality, typography, complex prompt understanding, and resource-efficiency",
"preview": "stabilityai--stable-diffusion-3.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 7.0"
+ "extras": "sampler: Default, cfg_scale: 7.0"
},
"Black Forest Labs FLUX.1 Dev": {
@@ -126,43 +126,49 @@
"preview": "black-forest-labs--FLUX.1-dev.jpg",
"desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 3.5"
+ "extras": "sampler: Default, cfg_scale: 3.5"
},
"Black Forest Labs FLUX.1 Schnell": {
"path": "black-forest-labs/FLUX.1-schnell",
"preview": "black-forest-labs--FLUX.1-schnell.jpg",
"desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching. Trained using latent adversarial diffusion distillation, FLUX.1 [schnell] can generate high-quality images in only 1 to 4 steps",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 3.5"
+ "extras": "sampler: Default, cfg_scale: 3.5"
},
"Black Forest Labs FLUX.1 Dev qint8": {
"path": "Disty0/FLUX.1-dev-qint8",
"preview": "black-forest-labs--FLUX.1-dev.jpg",
"desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 3.5"
+ "extras": "sampler: Default, cfg_scale: 3.5"
},
"Black Forest Labs FLUX.1 Dev qint4": {
"path": "Disty0/FLUX.1-dev-qint4",
"preview": "black-forest-labs--FLUX.1-dev.jpg",
"desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 3.5"
+ "extras": "sampler: Default, cfg_scale: 3.5"
},
"Black Forest Labs FLUX.1 Dev nf4": {
"path": "sayakpaul/flux.1-dev-nf4",
"preview": "black-forest-labs--FLUX.1-dev.jpg",
"desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 3.5"
+ "extras": "sampler: Default, cfg_scale: 3.5"
},
+ "VectorSpaceLab OmniGen v1": {
+ "path": "Shitao/OmniGen-v1",
+ "desc": "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.",
+ "preview": "Shitao--OmniGen-v1.jpg",
+ "skip": true
+ },
+
"AuraFlow 0.3": {
"path": "fal/AuraFlow-v0.3",
"desc": "AuraFlow v0.3 is the fully open-sourced flow-based text-to-image generation model. The model was trained with more compute compared to the previous version, AuraFlow-v0.2. Compared to AuraFlow-v0.2, the model is fine-tuned on more aesthetic datasets and now supports various aspect ratio, (now width and height up to 1536 pixels).",
"preview": "fal--AuraFlow-v0.3.jpg",
- "skip": true,
- "extras": "width: 1024, height: 1024"
+ "skip": true
},
"Segmind Vega": {
@@ -171,7 +177,7 @@
"desc": "The Segmind-Vega Model is a distilled version of the Stable Diffusion XL (SDXL), offering a remarkable 70% reduction in size and an impressive 100% speedup while retaining high-quality text-to-image generation capabilities. Trained on diverse datasets, including Grit and Midjourney scrape data, it excels at creating a wide range of visual content based on textual prompts. Employing a knowledge distillation strategy, Segmind-Vega leverages the teachings of several expert models, including SDXL, ZavyChromaXL, and JuggernautXL, to combine their strengths and produce compelling visual outputs.",
"variant": "fp16",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 9.0"
+ "extras": "sampler: Default, cfg_scale: 9.0"
},
"Segmind SSD-1B": {
"path": "huggingface/segmind/SSD-1B",
@@ -179,7 +185,7 @@
"desc": "The Segmind Stable Diffusion Model (SSD-1B) offers a compact, efficient, and distilled version of the SDXL model. At 50% smaller and 60% faster than Stable Diffusion XL (SDXL), it provides quick and seamless performance without sacrificing image quality.",
"variant": "fp16",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 9.0"
+ "extras": "sampler: Default, cfg_scale: 9.0"
},
"Segmind Tiny": {
"path": "segmind/tiny-sd",
@@ -197,7 +203,7 @@
"path": "segmind/SegMoE-4x2-v0",
"preview": "segmind--SegMoE-4x2-v0.jpg",
"desc": "SegMoE-4x2-v0 is an untrained Segmind Mixture of Diffusion Experts Model generated using segmoe from 4 Expert SDXL models. SegMoE is a powerful framework for dynamically combining Stable Diffusion Models into a Mixture of Experts within minutes without training",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"Pixart-α XL 2 Medium": {
@@ -210,7 +216,7 @@
"path": "PixArt-alpha/PixArt-XL-2-1024-MS",
"desc": "PixArt-α is a Transformer-based T2I diffusion model whose image generation quality is competitive with state-of-the-art image generators (e.g., Imagen, SDXL, and even Midjourney), and the training speed markedly surpasses existing large-scale T2I models. Extensive experiments demonstrate that PIXART-α excels in image quality, artistry, and semantic control. It can directly generate 1024px images from text prompts within a single sampling process.",
"preview": "PixArt-alpha--PixArt-XL-2-1024-MS.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
+ "extras": "sampler: Default, cfg_scale: 2.0"
},
"Pixart-Σ Small": {
"path": "huggingface/PixArt-alpha/PixArt-Sigma-XL-2-512-MS",
@@ -224,21 +230,21 @@
"desc": "PixArt-Σ, a Diffusion Transformer model (DiT) capable of directly generating images at 4K resolution. PixArt-Σ represents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts.",
"preview": "PixArt-alpha--pixart_sigma_sdxlvae_T5_diffusers.jpg",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
+ "extras": "sampler: Default, cfg_scale: 2.0"
},
"Pixart-Σ Large": {
"path": "huggingface/PixArt-alpha/PixArt-Sigma-XL-2-2K-MS",
"desc": "PixArt-Σ, a Diffusion Transformer model (DiT) capable of directly generating images at 4K resolution. PixArt-Σ represents a significant advancement over its predecessor, PixArt-α, offering images of markedly higher fidelity and improved alignment with text prompts.",
"preview": "PixArt-alpha--pixart_sigma_sdxlvae_T5_diffusers.jpg",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
+ "extras": "sampler: Default, cfg_scale: 2.0"
},
"Tencent HunyuanDiT 1.2": {
"path": "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers",
"desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
"preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 2.0"
+ "extras": "sampler: Default, cfg_scale: 2.0"
},
"AlphaVLLM Lumina Next SFT": {
@@ -246,7 +252,7 @@
"desc": "The Lumina-Next-SFT is a Next-DiT model containing 2B parameters and utilizes Gemma-2B as the text encoder, enhanced through high-quality supervised fine-tuning (SFT).",
"preview": "Alpha-VLLM--Lumina-Next-SFT-diffusers.jpg",
"skip": true,
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"Kwai Kolors": {
@@ -274,7 +280,7 @@
"desc": "Kandinsky 3.0 is an open-source text-to-image diffusion model built upon the Kandinsky2-x model family. In comparison to its predecessors, Kandinsky 3.0 incorporates more data and specifically related to Russian culture, which allows to generate pictures related to Russin culture. Furthermore, enhancements have been made to the text understanding and visual quality of the model, achieved by increasing the size of the text encoder and Diffusion U-Net models, respectively.",
"preview": "kandinsky-community--kandinsky-3.jpg",
"variant": "fp16",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"Playground v1": {
@@ -299,13 +305,13 @@
"path": "playgroundai/playground-v2-1024px-aesthetic",
"desc": "Playground v2 is a diffusion-based text-to-image generative model. The model was trained from scratch by the research team at Playground. Images generated by Playground v2 are favored 2.5 times more than those produced by Stable Diffusion XL, according to Playground’s user study.",
"preview": "playgroundai--playground-v2-1024px-aesthetic.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"Playground v2.5": {
"path": "playground-v2.5-1024px-aesthetic.fp16.safetensors@https://huggingface.co/playgroundai/playground-v2.5-1024px-aesthetic/resolve/main/playground-v2.5-1024px-aesthetic.fp16.safetensors?download=true",
"desc": "Playground v2.5 is a diffusion-based text-to-image generative model, and a successor to Playground v2. Playground v2.5 is the state-of-the-art open-source model in aesthetic quality. Our user studies demonstrate that our model outperforms SDXL, Playground v2, PixArt-α, DALL-E 3, and Midjourney 5.2.",
"preview": "playgroundai--playground-v2-1024px-aesthetic.jpg",
- "extras": "width: 1024, height: 1024, sampler: DPM++ 2M EDM"
+ "extras": "sampler: DPM++ 2M EDM"
},
"CogView 3 Plus": {
@@ -339,7 +345,7 @@
"path": "warp-ai/wuerstchen",
"desc": "Würstchen is a diffusion model whose text-conditional model works in a highly compressed latent space of images. Why is this important? Compressing data can reduce computational costs for both training and inference by magnitudes. Training on 1024x1024 images, is way more expensive than training at 32x32. Usually, other works make use of a relatively small compression, in the range of 4x - 8x spatial compression. Würstchen takes this to an extreme. Through its novel design, we achieve a 42x spatial compression. Würstchen employs a two-stage compression, what we call Stage A and Stage B. Stage A is a VQGAN, and Stage B is a Diffusion Autoencoder (more details can be found in the paper). A third model, Stage C, is learned in that highly compressed latent space. This training requires fractions of the compute used for current top-performing models, allowing also cheaper and faster inference.",
"preview": "warp-ai--wuerstchen.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default, cfg_scale: 4.0, image_cfg_scale: 0.0"
+ "extras": "sampler: Default, cfg_scale: 4.0, image_cfg_scale: 0.0"
},
"KOALA 700M": {
"path": "huggingface/etri-vilab/koala-700m-llava-cap",
@@ -347,7 +353,7 @@
"skip": true,
"desc": "Fast text-to-image model, called KOALA, by compressing SDXL's U-Net and distilling knowledge from SDXL into our model. KOALA-700M can generate a 1024x1024 image in less than 1.5 seconds on an NVIDIA 4090 GPU, which is more than 2x faster than SDXL.",
"preview": "etri-vilab--koala-700m-llava-cap.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"Tsinghua UniDiffuser": {
"path": "thu-ml/unidiffuser-v1",
@@ -369,13 +375,13 @@
"path": "DeepFloyd/IF-I-M-v1.0",
"desc": "DeepFloyd-IF is a pixel-based text-to-image triple-cascaded diffusion model, that can generate pictures with new state-of-the-art for photorealism and language understanding. The result is a highly efficient model that outperforms current state-of-the-art models, achieving a zero-shot FID-30K score of 6.66 on the COCO dataset. It is modular and composed of frozen text mode and three pixel cascaded diffusion modules, each designed to generate images of increasing resolution: 64x64, 256x256, and 1024x1024.",
"preview": "DeepFloyd--IF-I-M-v1.0.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
},
"DeepFloyd IF Large": {
"path": "DeepFloyd/IF-I-L-v1.0",
"desc": "DeepFloyd-IF is a pixel-based text-to-image triple-cascaded diffusion model, that can generate pictures with new state-of-the-art for photorealism and language understanding. The result is a highly efficient model that outperforms current state-of-the-art models, achieving a zero-shot FID-30K score of 6.66 on the COCO dataset. It is modular and composed of frozen text mode and three pixel cascaded diffusion modules, each designed to generate images of increasing resolution: 64x64, 256x256, and 1024x1024.",
"preview": "DeepFloyd--IF-I-M-v1.0.jpg",
- "extras": "width: 1024, height: 1024, sampler: Default"
+ "extras": "sampler: Default"
}
}
diff --git a/models/Reference/MeissonFlow--Meissonic.jpg b/models/Reference/MeissonFlow--Meissonic.jpg
old mode 100755
new mode 100644
diff --git a/models/Reference/Shitao--OmniGen-v1.jpg b/models/Reference/Shitao--OmniGen-v1.jpg
new file mode 100644
index 000000000..f7bdb994f
Binary files /dev/null and b/models/Reference/Shitao--OmniGen-v1.jpg differ
diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py
new file mode 100644
index 000000000..9212193dd
--- /dev/null
+++ b/modules/model_omnigen.py
@@ -0,0 +1,29 @@
+def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
+ from modules import shared, devices, sd_models
+ repo_id = sd_models.path_to_repo(checkpoint_info.name)
+
+ # load
+ from modules.omnigen import OmniGenPipeline
+ pipe = OmniGenPipeline.from_pretrained(
+ model_name=repo_id,
+ vae_path='madebyollin/sdxl-vae-fp16-fix',
+ cache_dir=shared.opts.diffusers_dir,
+ )
+
+ # init
+ pipe.device = devices.device
+ pipe.dtype = devices.dtype
+ pipe.separate_cfg_infer = True
+ pipe.use_kv_cache = False
+ pipe.model.to(device=devices.device, dtype=devices.dtype)
+ pipe.model.eval()
+ pipe.vae.to(devices.device, dtype=devices.dtype)
+ devices.torch_gc()
+
+ # register
+ # from diffusers import pipelines
+ # pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen"] = pipe.__class__
+ # pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen"] = pipe.__class__
+ # pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen"] = pipe.__class__
+
+ return pipe
diff --git a/modules/model_stablecascade.py b/modules/model_stablecascade.py
index 444d11adc..6c23ea00a 100644
--- a/modules/model_stablecascade.py
+++ b/modules/model_stablecascade.py
@@ -1,5 +1,4 @@
import os
-import copy
import torch
import diffusers
from modules import shared, devices, sd_models
diff --git a/modules/modeldata.py b/modules/modeldata.py
index 27444aeaf..604ff4623 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -106,6 +106,8 @@ class Shared(sys.modules[__name__].__class__):
model_type = 'auraflow'
elif "Flux" in self.sd_model.__class__.__name__:
model_type = 'f1'
+ elif "OmniGen" in self.sd_model.__class__.__name__:
+ model_type = 'omnigen'
elif "CogVideo" in self.sd_model.__class__.__name__:
model_type = 'cogvideox'
else:
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 5894dad63..50eb0974b 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -291,10 +291,11 @@ def load_diffusers_models(clear=True):
mtime = os.path.getmtime(commit)
info = os.path.join(commit, "model_info.json")
index = os.path.join(commit, "model_index.json")
- if (not os.path.exists(index)) and (not os.path.exists(info)):
+ config = os.path.join(commit, "config.json")
+ if (not os.path.exists(index)) and (not os.path.exists(info)) and (not os.path.exists(config)):
debug(f'Diffusers skip model no info: {name}')
continue
- repo = { 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': commit, 'hash': snapshot, 'mtime': mtime, 'model_info': info, 'model_index': index }
+ repo = { 'name': name, 'filename': name, 'friendly': friendly, 'folder': folder, 'path': commit, 'hash': snapshot, 'mtime': mtime, 'model_info': info, 'model_index': index, 'model_config': config }
diffuser_repos.append(repo)
if os.path.exists(os.path.join(folder, 'hidden')):
continue
diff --git a/modules/modelstats.py b/modules/modelstats.py
index ff23afe25..8ef02ca35 100644
--- a/modules/modelstats.py
+++ b/modules/modelstats.py
@@ -21,8 +21,8 @@ class Module():
if hasattr(module, 'config'):
self.config = module.config
if isinstance(module, torch.nn.Module):
- self.device = module.device
- self.dtype = module.dtype
+ self.device = getattr(module, 'device', None)
+ self.dtype = getattr(module, 'dtype', None)
self.params = sum(p.numel() for p in module.parameters(recurse=True))
self.modules = len(list(module.modules()))
@@ -72,12 +72,12 @@ def analyze():
model = Model(shared.opts.sd_model_checkpoint)
if model.cls == '':
return model
- if not hasattr(shared.sd_model, '_internal_dict'):
- return model
+ if hasattr(shared.sd_model, '_internal_dict'):
+ keys = shared.sd_model._internal_dict.keys() # pylint: disable=protected-access
+ else:
+ keys = sd_models.get_signature(shared.sd_model).keys()
model.modules.clear()
- if not hasattr(shared.sd_model, '_internal_dict'):
- return model
- for k in shared.sd_model._internal_dict.keys(): # pylint: disable=protected-access
+ for k in keys: # pylint: disable=protected-access
component = getattr(shared.sd_model, k, None)
module = Module(k, component)
model.modules.append(module)
diff --git a/modules/omnigen/__init__.py b/modules/omnigen/__init__.py
new file mode 100644
index 000000000..40315a6f3
--- /dev/null
+++ b/modules/omnigen/__init__.py
@@ -0,0 +1,4 @@
+from .model import OmniGen
+from .processor import OmniGenProcessor
+from .scheduler import OmniGenScheduler
+from .pipeline import OmniGenPipeline
diff --git a/modules/omnigen/model.py b/modules/omnigen/model.py
new file mode 100644
index 000000000..7446ef080
--- /dev/null
+++ b/modules/omnigen/model.py
@@ -0,0 +1,407 @@
+# The code is revised from DiT
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from typing import Dict
+
+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):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t, dtype=torch.float32):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class FinalLayer(nn.Module):
+ """
+ The final layer of DiT.
+ """
+ def __init__(self, hidden_size, patch_size, out_channels):
+ super().__init__()
+ self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+ )
+
+ def forward(self, x, c):
+ shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+ x = modulate(self.norm_final(x), shift, scale)
+ x = self.linear(x)
+ return x
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1):
+ """
+ grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+ [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+ """
+ if isinstance(grid_size, int):
+ grid_size = (grid_size, grid_size)
+
+ grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
+ grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
+ grid = np.meshgrid(grid_w, grid_h) # here w goes first
+ grid = np.stack(grid, axis=0)
+
+ grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+ if cls_token and extra_tokens > 0:
+ pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+ return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+ assert embed_dim % 2 == 0
+
+ # use half of dimensions to encode grid_h
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
+
+ emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+ return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+ """
+ embed_dim: output dimension for each position
+ pos: a list of positions to be encoded: size (M,)
+ out: (M, D)
+ """
+ assert embed_dim % 2 == 0
+ omega = np.arange(embed_dim // 2, dtype=np.float64)
+ omega /= embed_dim / 2.
+ omega = 1. / 10000**omega # (D/2,)
+
+ pos = pos.reshape(-1) # (M,)
+ out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
+
+ emb_sin = np.sin(out) # (M, D/2)
+ emb_cos = np.cos(out) # (M, D/2)
+
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
+ return emb
+
+
+class PatchEmbedMR(nn.Module):
+ """ 2D Image to Patch Embedding
+ """
+ def __init__(
+ self,
+ patch_size: int = 2,
+ in_chans: int = 4,
+ embed_dim: int = 768,
+ bias: bool = True,
+ ):
+ super().__init__()
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+
+ def forward(self, x):
+ x = self.proj(x)
+ x = x.flatten(2).transpose(1, 2) # NCHW -> NLC
+ return x
+
+
+class OmniGen(nn.Module, PeftAdapterMixin):
+ """
+ Diffusion model with a Transformer backbone.
+ """
+ def __init__(
+ self,
+ transformer_config: Phi3Config,
+ patch_size=2,
+ in_channels=4,
+ pe_interpolation: float = 1.0,
+ pos_embed_max_size: int = 192,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ 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')
+ model_name = snapshot_download(repo_id=model_name,
+ cache_dir=cache_folder,
+ 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)
+ return model
+
+ def initialize_weights(self):
+ assert not hasattr(self, "llama")
+
+ # Initialize transformer layers:
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ 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]))
+ nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+ w = self.input_x_embedder.proj.weight.data
+ nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+ nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+
+ # Initialize timestep embedding MLP:
+ nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+ nn.init.normal_(self.time_token.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.time_token.mlp[2].weight, std=0.02)
+
+ # Zero-out output layers:
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+ nn.init.constant_(self.final_layer.linear.weight, 0)
+ nn.init.constant_(self.final_layer.linear.bias, 0)
+
+ def unpatchify(self, x, h, w):
+ """
+ x: (N, T, patch_size**2 * C)
+ imgs: (N, H, W, C)
+ """
+ c = self.out_channels
+
+ x = x.reshape(shape=(x.shape[0], h//self.patch_size, w//self.patch_size, self.patch_size, self.patch_size, c))
+ x = torch.einsum('nhwpqc->nchpwq', x)
+ imgs = x.reshape(shape=(x.shape[0], c, h, w))
+ return imgs
+
+
+ def cropped_pos_embed(self, height, width):
+ """Crops positional embeddings for SD3 compatibility."""
+ if self.pos_embed_max_size is None:
+ raise ValueError("`pos_embed_max_size` must be set for cropping.")
+
+ height = height // self.patch_size
+ width = width // self.patch_size
+ if height > self.pos_embed_max_size:
+ raise ValueError(
+ f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
+ )
+ if width > self.pos_embed_max_size:
+ raise ValueError(
+ f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
+ )
+
+ top = (self.pos_embed_max_size - height) // 2
+ left = (self.pos_embed_max_size - width) // 2
+ spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
+ spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
+ # print(top, top + height, left, left + width, spatial_pos_embed.size())
+ spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
+ return spatial_pos_embed
+
+
+ def patch_multiple_resolutions(self, latents, padding_latent=None, is_input_images:bool=False):
+ if isinstance(latents, list):
+ return_list = False
+ if padding_latent is None:
+ padding_latent = [None] * len(latents)
+ return_list = True
+ patched_latents, num_tokens, shapes = [], [], []
+ for latent, padding in zip(latents, padding_latent):
+ height, width = latent.shape[-2:]
+ if is_input_images:
+ latent = self.input_x_embedder(latent)
+ else:
+ latent = self.x_embedder(latent)
+ pos_embed = self.cropped_pos_embed(height, width)
+ latent = latent + pos_embed
+ if padding is not None:
+ latent = torch.cat([latent, padding], dim=-2)
+ patched_latents.append(latent)
+
+ num_tokens.append(pos_embed.size(1))
+ shapes.append([height, width])
+ if not return_list:
+ latents = torch.cat(patched_latents, dim=0)
+ else:
+ latents = patched_latents
+ else:
+ height, width = latents.shape[-2:]
+ if is_input_images:
+ latents = self.input_x_embedder(latents)
+ else:
+ latents = self.x_embedder(latents)
+ 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)
+
+ 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:
+ condition_embeds = self.llm.embed_tokens(input_ids).clone()
+ input_img_inx = 0
+ for b_inx in input_image_sizes.keys():
+ for start_inx, end_inx in input_image_sizes[b_inx]:
+ 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)
+
+ 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, past_key_values = output.last_hidden_state, output.past_key_values
+ if input_is_list:
+ image_embedding = output[:, -max(num_tokens):]
+ time_emb = self.t_embedder(timestep, dtype=x.dtype)
+ x = self.final_layer(image_embedding, time_emb)
+ latents = []
+ for i in range(x.size(0)):
+ latent = x[i:i+1, :num_tokens[i]]
+ latent = self.unpatchify(latent, shapes[i][0], shapes[i][1])
+ latents.append(latent)
+ else:
+ image_embedding = output[:, -num_tokens:]
+ time_emb = self.t_embedder(timestep, dtype=x.dtype)
+ x = self.final_layer(image_embedding, time_emb)
+ latents = self.unpatchify(x, shapes[0], shapes[1])
+
+ if return_past_key_values:
+ return latents, past_key_values
+ 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.
+ """
+ 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:
+ 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)
+ model_out = [cond, cond, cond]
+ else:
+ 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
+
+
+ @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.
+ """
+ self.llm.config.use_cache = use_kv_cache
+ if past_key_values is None:
+ past_key_values = [None] * len(attention_mask)
+
+ x = torch.split(x, len(x) // len(attention_mask), dim=0)
+ timestep = timestep.to(x[0].dtype)
+ timestep = torch.split(timestep, len(timestep) // len(input_ids), dim=0)
+
+ 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])
+ model_out.append(temp_out)
+ pask_key_values.append(temp_pask_key_values)
+
+ if len(model_out) == 3:
+ cond, uncond, img_cond = model_out
+ cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
+ model_out = [cond, cond, cond]
+ elif len(model_out) == 2:
+ cond, uncond = model_out
+ cond = uncond + cfg_scale * (cond - uncond)
+ 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
new file mode 100644
index 000000000..c15fd123d
--- /dev/null
+++ b/modules/omnigen/pipeline.py
@@ -0,0 +1,221 @@
+import os
+from typing import List, Union
+from PIL import Image
+import torch
+from huggingface_hub import snapshot_download
+from peft import PeftModel
+from diffusers.models import AutoencoderKL
+from diffusers.utils import replace_example_docstring
+from .model import OmniGen
+from .processor import OmniGenProcessor
+from .scheduler import OmniGenScheduler
+
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```py
+ >>> from OmniGen import OmniGenPipeline
+ >>> pipe = FluxControlNetPipeline.from_pretrained(
+ ... base_model
+ ... )
+ >>> prompt = "A woman holds a bouquet of flowers and faces the camera"
+ >>> image = pipe(
+ ... prompt,
+ ... guidance_scale=3.0,
+ ... num_inference_steps=50,
+ ... ).images[0]
+ >>> image.save("t2i.png")
+ ```
+"""
+
+
+class OmniGenPipeline():
+ def __init__(
+ self,
+ vae: AutoencoderKL,
+ model: OmniGen,
+ processor: OmniGenProcessor,
+ ):
+ super().__init__()
+ 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.use_kv_cache: bool = False
+ # omnigen does not inherit from diffusionpipeline so we hack it
+ self._internal_dict = { # pylint: disable=protected-access
+ 'vae': self.vae,
+ 'model': self.model,
+ 'processor': self.processor,
+ }
+
+ @classmethod
+ def from_pretrained(cls, model_name, vae_path: str=None, cache_dir: str=None):
+ if not os.path.exists(model_name):
+ cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
+ model_name = snapshot_download(repo_id=model_name,
+ cache_dir=cache_dir,
+ ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
+ model = OmniGen.from_pretrained(model_name)
+ 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")
+ return cls(vae, model, processor)
+
+ 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]):
+ if isinstance(device, str):
+ device = torch.device(device)
+ self.model.to(device)
+ self.vae.to(device)
+
+ def vae_encode(self, x, dtype):
+ x = x.to(dtype)
+ if self.vae.config.shift_factor is not None:
+ x = self.vae.encode(x).latent_dist.sample()
+ x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+ else:
+ x = self.vae.encode(x).latent_dist.sample().mul_(self.vae.config.scaling_factor)
+ x = x.to(dtype)
+ return x
+
+ def move_to_device(self, data):
+ if isinstance(data, list):
+ return [x.to(self.device) for x in data]
+ return data.to(self.device)
+
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Union[str, List[str]],
+ input_images: Union[List[str], List[List[str]]] = None,
+ height: int = 1024,
+ width: int = 1024,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 3,
+ use_img_guidance: bool = True,
+ img_guidance_scale: float = 1.6,
+ output_type: str = 'latent',
+ seed: int = None,
+ ):
+ r"""
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ prompt (`str` or `List[str]`):
+ The prompt or prompts to guide the image generation.
+ input_images (`List[str]` or `List[List[str]]`, *optional*):
+ The list of input images. We will replace the "<|image_i|>" in prompt with the 1-th image in list.
+ height (`int`, *optional*, defaults to 1024):
+ The height in pixels of the generated image. The number must be a multiple of 16.
+ width (`int`, *optional*, defaults to 1024):
+ The width in pixels of the generated image. The number must be a multiple of 16.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
+ guidance_scale (`float`, *optional*, defaults to 4.0):
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+ usually at the expense of lower image quality.
+ use_img_guidance (`bool`, *optional*, defaults to True):
+ 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).
+ self.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.
+ self.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.
+ Examples:
+
+ Returns:
+ A list with the generated images.
+ """
+ assert height%16 == 0 and width%16 == 0
+ if self.separate_cfg_infer:
+ self.use_kv_cache = False
+ # raise "Currently, don't support both self.use_kv_cache and self.separate_cfg_infer"
+ 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
+
+ input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=self.separate_cfg_infer)
+
+ num_prompt = len(prompt)
+ num_cfg = 2 if use_img_guidance else 1
+ latent_size_h, latent_size_w = height//8, width//8
+
+ if seed is not None:
+ generator = torch.Generator(device=self.device).manual_seed(seed)
+ else:
+ generator = None
+ 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(self.dtype)
+
+ input_img_latents = []
+ if self.separate_cfg_infer:
+ for temp_pixel_values in input_data['input_pixel_values']:
+ temp_input_latents = []
+ for img in temp_pixel_values:
+ img = self.vae_encode(img.to(self.device), self.dtype)
+ temp_input_latents.append(img)
+ input_img_latents.append(temp_input_latents)
+ else:
+ for img in input_data['input_pixel_values']:
+ img = self.vae_encode(img.to(self.device), self.dtype)
+ input_img_latents.append(img)
+
+ model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']),
+ input_img_latents=input_img_latents,
+ input_image_sizes=input_data['input_image_sizes'],
+ attention_mask=self.move_to_device(input_data["attention_mask"]),
+ position_ids=self.move_to_device(input_data["position_ids"]),
+ cfg_scale=guidance_scale,
+ img_cfg_scale=img_guidance_scale,
+ use_img_cfg=use_img_guidance,
+ use_kv_cache=self.use_kv_cache)
+
+ if self.separate_cfg_infer:
+ func = self.model.forward_with_separate_cfg
+ else:
+ func = self.model.forward_with_cfg
+ self.model.to(self.dtype)
+
+ scheduler = OmniGenScheduler(num_steps=num_inference_steps)
+ samples = scheduler(latents, func, model_kwargs, use_kv_cache=self.use_kv_cache)
+ samples = samples.chunk((1+num_cfg), dim=0)[0]
+
+ if output_type == 'latent':
+ output_images = { 'images': samples }
+ return output_images
+
+ samples = samples.to(self.vae.dtype)
+ 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
+
+ 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()
+ output_images = []
+ for _i, sample in enumerate(output_samples):
+ output_images.append(Image.fromarray(sample))
+
+ return output_images
diff --git a/modules/omnigen/processor.py b/modules/omnigen/processor.py
new file mode 100644
index 000000000..071f28423
--- /dev/null
+++ b/modules/omnigen/processor.py
@@ -0,0 +1,334 @@
+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,
+)
+
+
+
+
+class OmniGenProcessor:
+ def __init__(self,
+ text_tokenizer,
+ max_image_size: int=1024):
+ self.text_tokenizer = text_tokenizer
+ self.max_image_size = max_image_size
+
+ self.image_transform = transforms.Compose([
+ transforms.Lambda(lambda pil_image: crop_arr(pil_image, max_image_size)),
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+ ])
+
+ self.collator = OmniGenCollator()
+ self.separate_collator = OmniGenSeparateCollator()
+
+ @classmethod
+ def from_pretrained(cls, model_name):
+ if not os.path.exists(model_name):
+ cache_folder = os.getenv('HF_HUB_CACHE')
+ model_name = snapshot_download(repo_id=model_name,
+ cache_dir=cache_folder,
+ allow_patterns="*.json")
+ text_tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+ return cls(text_tokenizer)
+
+
+ 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:
+ model_inputs = self.text_tokenizer(text)
+ 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)]
+
+ 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_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
+ for i in range(len(prompt_chunks)):
+ all_input_ids.extend(prompt_chunks[i])
+ if i != len(prompt_chunks) -1:
+ start_inx = len(all_input_ids)
+ size = input_images[i].size(-2) * input_images[i].size(-1) // 16 // 16
+ img_inx.append([start_inx, start_inx+size])
+ all_input_ids.extend([0]*size)
+
+ return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx}
+
+
+ def add_prefix_instruction(self, prompt):
+ user_prompt = '<|user|>\n'
+ generation_prompt = 'Generate an image according to the following instructions\n'
+ assistant_prompt = '<|assistant|>\n<|diffusion|>'
+ prompt_suffix = "<|end|>\n"
+ prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}"
+ return prompt
+
+
+ def __call__(self,
+ instructions: List[str],
+ input_images: List[List[str]] = None,
+ height: int = 1024,
+ width: int = 1024,
+ 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,
+ ) -> Dict:
+
+ if input_images is None:
+ use_img_cfg = False
+ if isinstance(instructions, str):
+ instructions = [instructions]
+ input_images = [input_images]
+
+ input_data = []
+ for i in range(len(instructions)):
+ cur_instruction = instructions[i]
+ cur_input_images = None if input_images is None else input_images[i]
+ if cur_input_images is not None and len(cur_input_images) > 0:
+ cur_input_images = [self.process_image(x) for x in cur_input_images]
+ 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:
+ if cur_input_images is not None and len(cur_input_images) >= 1:
+ img_cfg_prompt = [f"![]()
<|image_{i+1}|>" for i in range(len(cur_input_images))]
+ img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images)
+ else:
+ img_cfg_mllm_input = neg_mllm_input
+
+ input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
+
+ if separate_cfg_input:
+ return self.separate_collator(input_data)
+ 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)
+ 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 = []
+ text_length = attention_mask.size(-1)
+ img_length = max(num_tokens_for_output_images)
+ seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token
+ inx = 0
+ for mask in attention_mask:
+ temp_l = torch.sum(mask)
+ pad_l = text_length - temp_l
+
+ temp_mask = torch.tril(torch.ones(size=(temp_l+1, temp_l+1)))
+
+ image_mask = torch.zeros(size=(temp_l+1, img_length))
+ temp_mask = torch.cat([temp_mask, image_mask], dim=-1)
+
+ image_mask = torch.ones(size=(img_length, temp_l+img_length+1))
+ temp_mask = torch.cat([temp_mask, image_mask], dim=0)
+
+ if pad_l > 0:
+ pad_mask = torch.zeros(size=(temp_l+1+img_length, pad_l))
+ temp_mask = torch.cat([pad_mask, temp_mask], dim=-1)
+
+ pad_mask = torch.ones(size=(pad_l, seq_len))
+ temp_mask = torch.cat([pad_mask, temp_mask], dim=0)
+
+ true_img_length = num_tokens_for_output_images[inx]
+ pad_img_length = img_length - true_img_length
+ if pad_img_length > 0:
+ temp_mask[:, -pad_img_length:] = 0
+ 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 = []
+
+ for i in range(len(input_ids)):
+ temp_ids = input_ids[i]
+ temp_l = len(temp_ids)
+ pad_l = max_l - temp_l
+ if pad_l == 0:
+ attention_mask.append([1]*max_l)
+ padded_ids.append(temp_ids)
+ 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]:
+ new_inx.append([x+pad_l for x in old_inx])
+ image_sizes[i] = new_inx
+
+ return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes
+
+
+ def process_mllm_input(self, mllm_inputs, target_img_size):
+ num_tokens_for_output_images = []
+ for img_size in target_img_size:
+ num_tokens_for_output_images.append(img_size[0]*img_size[1]//16//16)
+
+ pixel_values, image_sizes = [], {}
+ b_inx = 0
+ for x in mllm_inputs:
+ if x['pixel_values'] is not None:
+ pixel_values.extend(x['pixel_values'])
+ for size in x['image_sizes']:
+ if b_inx not in image_sizes:
+ image_sizes[b_inx] = [size]
+ else:
+ image_sizes[b_inx].append(size)
+ 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)
+ attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images)
+ 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
+ else:
+ mllm_inputs = mllm_inputs + cfg_mllm_inputs
+ target_img_size = target_img_size + target_img_size
+
+
+ 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)
+
+ 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,
+ }
+ return data
+
+
+class OmniGenSeparateCollator(OmniGenCollator):
+ 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]
+
+
+ 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)
+ all_position_ids.append(position_ids)
+ all_pixel_values.append(pixel_values)
+ all_image_sizes.append(image_sizes)
+ all_padding_images.append(padding_images)
+
+ if cfg_mllm_inputs[0] is not None:
+ padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(cfg_mllm_inputs, target_img_size)
+ all_padded_input_ids.append(padded_input_ids)
+ all_attention_mask.append(attention_mask)
+ all_position_ids.append(position_ids)
+ all_pixel_values.append(pixel_values)
+ all_image_sizes.append(image_sizes)
+ all_padding_images.append(padding_images)
+ if img_cfg_mllm_input[0] is not None:
+ padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(img_cfg_mllm_input, target_img_size)
+ all_padded_input_ids.append(padded_input_ids)
+ all_attention_mask.append(attention_mask)
+ all_position_ids.append(position_ids)
+ all_pixel_values.append(pixel_values)
+ all_image_sizes.append(image_sizes)
+ all_padding_images.append(padding_images)
+
+ 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,
+ }
+ return data
diff --git a/modules/omnigen/scheduler.py b/modules/omnigen/scheduler.py
new file mode 100644
index 000000000..0764fd8f0
--- /dev/null
+++ b/modules/omnigen/scheduler.py
@@ -0,0 +1,55 @@
+import torch
+from tqdm import tqdm
+from transformers.cache_utils import Cache, DynamicCache
+
+class OmniGenScheduler:
+ def __init__(self, num_steps: int=50, time_shifting_factor: int=1):
+ self.num_steps = num_steps
+ self.time_shift = time_shifting_factor
+
+ t = torch.linspace(0, 1, num_steps+1)
+ t = t / (t + time_shifting_factor - time_shifting_factor * t)
+ self.sigma = t
+
+ def crop_kv_cache(self, past_key_values, num_tokens_for_img):
+ 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)
+
+ def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img):
+ if isinstance(position_ids, list):
+ for i in range(len(position_ids)):
+ position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):]
+ else:
+ position_ids = position_ids[:, -(num_tokens_for_img+1):]
+ return position_ids
+
+ def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img):
+ if isinstance(attention_mask, list):
+ 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
+ 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)
+ 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)
+ 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)
+ return z
+
diff --git a/modules/omnigen/transformer.py b/modules/omnigen/transformer.py
new file mode 100644
index 000000000..f3bcdb15a
--- /dev/null
+++ b/modules/omnigen/transformer.py
@@ -0,0 +1,159 @@
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from huggingface_hub import snapshot_download
+
+from transformers.modeling_outputs import (
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers import Phi3Config, Phi3Model
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class Phi3Transformer(Phi3Model):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
+ We only modified the attention mask
+ Args:
+ config: Phi3Config
+ """
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if self.gradient_checkpointing and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+ )
+ use_cache = False
+
+ # kept for BC (non `Cache` `past_key_values` inputs)
+ return_legacy_cache = False
+ if use_cache and not isinstance(past_key_values, Cache):
+ return_legacy_cache = True
+ if past_key_values is None:
+ past_key_values = DynamicCache()
+ else:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ logger.warning_once(
+ "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
+ "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
+ "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
+ )
+
+ 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 attention_mask is not None and attention_mask.dim() == 3:
+ dtype = inputs_embeds.dtype
+ min_dtype = torch.finfo(dtype).min
+ attention_mask = (1 - attention_mask) * min_dtype
+ attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype)
+ else:
+ raise
+ # causal_mask = self._update_causal_mask(
+ # attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+ # )
+
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ cache_position,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ hidden_states = self.norm(hidden_states)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = next_decoder_cache if use_cache else None
+ if return_legacy_cache:
+ next_cache = next_cache.to_legacy_cache()
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
diff --git a/modules/omnigen/utils.py b/modules/omnigen/utils.py
new file mode 100644
index 000000000..5483d6eab
--- /dev/null
+++ b/modules/omnigen/utils.py
@@ -0,0 +1,109 @@
+import logging
+
+from PIL import Image
+import torch
+import numpy as np
+
+def create_logger(logging_dir):
+ """
+ Create a logger that writes to a log file and stdout.
+ """
+ logging.basicConfig(
+ level=logging.INFO,
+ format='[\033[34m%(asctime)s\033[0m] %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S',
+ handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+ )
+ logger = logging.getLogger(__name__)
+ return logger
+
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+ """
+ Step the EMA model towards the current model.
+ """
+ ema_params = dict(ema_model.named_parameters())
+ for name, param in model.named_parameters():
+ # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+ ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+
+
+def requires_grad(model, flag=True):
+ """
+ Set requires_grad flag for all parameters in a model.
+ """
+ for p in model.parameters():
+ p.requires_grad = flag
+
+
+def center_crop_arr(pil_image, image_size):
+ """
+ Center cropping implementation from ADM.
+ https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+ """
+ while min(*pil_image.size) >= 2 * image_size:
+ pil_image = pil_image.resize(
+ tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+ )
+
+ scale = image_size / min(*pil_image.size)
+ pil_image = pil_image.resize(
+ tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+ )
+
+ arr = np.array(pil_image)
+ crop_y = (arr.shape[0] - image_size) // 2
+ crop_x = (arr.shape[1] - image_size) // 2
+ return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+
+
+def crop_arr(pil_image, max_image_size):
+ while min(*pil_image.size) >= 2 * max_image_size:
+ pil_image = pil_image.resize(
+ tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+ )
+
+ if max(*pil_image.size) > max_image_size:
+ scale = max_image_size / max(*pil_image.size)
+ pil_image = pil_image.resize(
+ tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+ )
+
+ if min(*pil_image.size) < 16:
+ scale = 16 / min(*pil_image.size)
+ pil_image = pil_image.resize(
+ tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+ )
+
+ arr = np.array(pil_image)
+ crop_y1 = (arr.shape[0] % 16) // 2
+ crop_y2 = arr.shape[0] % 16 - crop_y1
+
+ crop_x1 = (arr.shape[1] % 16) // 2
+ crop_x2 = arr.shape[1] % 16 - crop_x1
+
+ arr = arr[crop_y1:arr.shape[0]-crop_y2, crop_x1:arr.shape[1]-crop_x2]
+ return Image.fromarray(arr)
+
+
+
+def vae_encode(vae, x, weight_dtype):
+ if x is not None:
+ if vae.config.shift_factor is not None:
+ x = vae.encode(x).latent_dist.sample()
+ x = (x - vae.config.shift_factor) * vae.config.scaling_factor
+ else:
+ x = vae.encode(x).latent_dist.sample().mul_(vae.config.scaling_factor)
+ x = x.to(weight_dtype)
+ return x
+
+def vae_encode_list(vae, x, weight_dtype):
+ latents = []
+ for img in x:
+ img = vae_encode(vae, img, weight_dtype)
+ latents.append(img)
+ return latents
diff --git a/modules/processing_args.py b/modules/processing_args.py
index 0e2999128..7296fccb8 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -19,7 +19,7 @@ def task_specific_kwargs(p, model):
is_img2img_model = bool('Zero123' in shared.sd_model.__class__.__name__)
if len(getattr(p, 'init_images', [])) > 0:
p.init_images = [p.convert('RGB') for p in p.init_images]
- if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE and not is_img2img_model:
+ if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE or len(getattr(p, 'init_images', [])) == 0 and not is_img2img_model:
p.ops.append('txt2img')
if hasattr(p, 'width') and hasattr(p, 'height'):
task_args = {
@@ -38,6 +38,14 @@ def task_specific_kwargs(p, model):
p.width = 8 * math.ceil(p.init_images[0].width / 8)
p.height = 8 * math.ceil(p.init_images[0].height / 8)
task_args['width'], task_args['height'] = p.width, p.height
+ if model.__class__.__name__ == 'OmniGenPipeline':
+ p.width = 16 * math.ceil(p.init_images[0].width / 16)
+ p.height = 16 * math.ceil(p.init_images[0].height / 16)
+ task_args = {
+ 'width': p.width,
+ 'height': p.height,
+ 'input_images': [p.init_images], # omnigen expects list-of-lists
+ }
elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0:
p.ops.append('instruct')
task_args = {
@@ -128,6 +136,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
args['pooled_prompt_embeds'] = p.positive_pooleds[0]
elif 'Flux' in model.__class__.__name__ and len(getattr(p, 'positive_pooleds', [])) > 0:
args['pooled_prompt_embeds'] = p.positive_pooleds[0]
+ if 'OmniGen' in model.__class__.__name__:
+ args['prompt'] = [p.replace('|image|', '<|image_1|>') for p in prompts]
else:
args['prompt'] = prompts
if 'negative_prompt' in possible:
@@ -174,6 +184,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
args['noise_sampler_seed'] = p.seeds
if 'guidance_scale' in possible:
args['guidance_scale'] = p.cfg_scale
+ if 'img_guidance_scale' in possible and hasattr(p, 'image_cfg_scale'):
+ args['img_guidance_scale'] = p.image_cfg_scale
if 'generator' in possible:
args['generator'] = get_generator(p)
if 'latents' in possible and getattr(p, "init_latent", None) is not None:
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 9b485d361..b76376fd1 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -99,6 +99,8 @@ def process_base(p: processing.StableDiffusionProcessing):
output = shared.sd_model(**base_args)
if isinstance(output, dict):
output = SimpleNamespace(**output)
+ if isinstance(output, list):
+ output = SimpleNamespace(images=output)
shared.history.add(output.images, info=processing.create_infotext(p), ops=p.ops)
timer.process.record('pipeline')
hidiffusion.unapply()
diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py
index 30bf18fdb..ef83834e5 100644
--- a/modules/processing_helpers.py
+++ b/modules/processing_helpers.py
@@ -446,6 +446,8 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
if not is_txt2img():
if use_denoise_start and shared.sd_model_type == 'sdxl':
steps = p.steps // (1 - p.refiner_start)
+ elif shared.sd_model_type == 'omnigen':
+ steps = p.steps
elif p.denoising_strength > 0:
steps = (p.steps // p.denoising_strength) + 1
else:
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 47d0c5947..ab7eca5e6 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -628,6 +628,9 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
if 'meissonic' in f.lower():
guess = 'Meissonic'
pipeline = 'custom'
+ if 'omnigen' in f.lower():
+ guess = 'OmniGen'
+ pipeline = 'custom'
if 'flux' in f.lower():
guess = 'FLUX'
if size > 11000 and size < 20000:
@@ -1090,6 +1093,9 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
elif model_type in ['Meissonic']: # forced pipeline
from modules.model_meissonic import load_meissonic
sd_model = load_meissonic(checkpoint_info, diffusers_load_config)
+ elif model_type in ['OmniGen']: # forced pipeline
+ from modules.model_omnigen import load_omnigen
+ sd_model = load_omnigen(checkpoint_info, diffusers_load_config)
except Exception as e:
shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
if debug_load:
@@ -1390,7 +1396,7 @@ class DiffusersTaskType(Enum):
def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
- if pipe.__class__.__name__ in ["StableVideoDiffusionPipeline", "LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL"]:
+ if pipe.__class__.__name__ in ["StableVideoDiffusionPipeline", "LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL", "OmniGenPipeline"]:
return DiffusersTaskType.IMAGE_2_IMAGE
elif pipe.__class__.__name__ == "StableDiffusionXLInstructPix2PixPipeline":
return DiffusersTaskType.INSTRUCT
@@ -1505,7 +1511,7 @@ def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionP
def clean_diffuser_pipe(pipe):
- if pipe is not None and shared.sd_model_type == 'sdxl' and 'requires_aesthetics_score' in pipe.config and hasattr(pipe, '_internal_dict'):
+ if pipe is not None and shared.sd_model_type == 'sdxl' and hasattr(pipe, 'config') and 'requires_aesthetics_score' in pipe.config and hasattr(pipe, '_internal_dict'):
debug_process(f'Pipeline clean: {pipe.__class__.__name__}')
# diffusers adds requires_aesthetics_score with img2img and complains if requires_aesthetics_score exist in txt2img
internal_dict = dict(pipe._internal_dict) # pylint: disable=protected-access
@@ -1523,7 +1529,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
return pipe
# skip specific pipelines
- if n in ['StableDiffusionReferencePipeline', 'StableDiffusionAdapterPipeline', 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline']:
+ if n in ['StableDiffusionReferencePipeline', 'StableDiffusionAdapterPipeline', 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline', 'OmniGenPipeline']:
return pipe
if 'Onnx' in pipe.__class__.__name__:
return pipe
diff --git a/modules/segmoe/segmoe_model.py b/modules/segmoe/segmoe_model.py
index 3dd5a8cb1..4b96527be 100644
--- a/modules/segmoe/segmoe_model.py
+++ b/modules/segmoe/segmoe_model.py
@@ -24,7 +24,7 @@ def remove_all_forward_hooks(model: torch.nn.Module) -> None:
for _name, child in model._modules.items(): # pylint: disable=protected-access
if child is not None:
if hasattr(child, "_forward_hooks"):
- child._forward_hooks: Dict[int, Callable] = OrderedDict()
+ child._forward_hooks = OrderedDict()
remove_all_forward_hooks(child)
diff --git a/wiki b/wiki
index b445dda53..bd142e7c8 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit b445dda532e2c0a1ffa4ed01451bd33a11a06658
+Subproject commit bd142e7c88368d8693b1efec4a3ec4ce1f3a0321