From e1c9f14489c1413dc8ad922017de28986d900749 Mon Sep 17 00:00:00 2001
From: Pavel Yakovlev <zhmactep@gmail.com>
Date: Fri, 27 Jan 2023 20:37:37 +0300
Subject: [PATCH] feature: add openjourney models

---
 README.md            |  2 ++
 imaginairy/config.py | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 68bf47c..1f34679 100644
--- a/README.md
+++ b/README.md
@@ -283,6 +283,8 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -
 
 ## ChangeLog
 
+- feature: `openjourney-v1` and `openjourney-v2` models added. available via `--model openjourney-v2`
+
 **8.2.0**
 - feature: added `aimg system-info` command to help debug issues
 
diff --git a/imaginairy/config.py b/imaginairy/config.py
index a59efc4..f2296c5 100644
--- a/imaginairy/config.py
+++ b/imaginairy/config.py
@@ -24,6 +24,7 @@ class ModelConfig:
     weights_url_full: str = None
     forced_attn_precision: str = "default"
     default_negative_prompt: str = DEFAULT_NEGATIVE_PROMPT
+    alias: str = None
 
 
 midas_url = "https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt"
@@ -34,6 +35,7 @@ MODEL_CONFIGS = [
         config_path="configs/stable-diffusion-v1.yaml",
         weights_url="https://huggingface.co/bstddev/sd-v1-4/resolve/77221977fa8de8ab8f36fac0374c120bd5b53287/sd-v1-4.ckpt",
         default_image_size=512,
+        alias="sd14",
     ),
     ModelConfig(
         short_name="SD-1.5",
@@ -41,36 +43,42 @@ MODEL_CONFIGS = [
         weights_url="https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/889b629140e71758e1e0006e355c331a5744b4bf/v1-5-pruned-emaonly.ckpt",
         weights_url_full="https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/889b629140e71758e1e0006e355c331a5744b4bf/v1-5-pruned.ckpt",
         default_image_size=512,
+        alias="sd15",
     ),
     ModelConfig(
         short_name="SD-1.5-inpaint",
         config_path="configs/stable-diffusion-v1-inpaint.yaml",
         weights_url="https://huggingface.co/julienacquaviva/inpainting/resolve/2155ff7fe38b55f4c0d99c2f1ab9b561f8311ca7/sd-v1-5-inpainting.ckpt",
         default_image_size=512,
+        alias="sd15in",
     ),
     ModelConfig(
         short_name="SD-2.0",
         config_path="configs/stable-diffusion-v2-inference.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt",
         default_image_size=512,
+        alias="sd20",
     ),
     ModelConfig(
         short_name="SD-2.0-inpaint",
         config_path="configs/stable-diffusion-v2-inpainting-inference.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
         default_image_size=512,
+        alias="sd20in",
     ),
     ModelConfig(
         short_name="SD-2.1",
         config_path="configs/stable-diffusion-v2-inference.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt",
         default_image_size=512,
+        alias="sd21",
     ),
     ModelConfig(
         short_name="SD-2.1-inpaint",
         config_path="configs/stable-diffusion-v2-inpainting-inference.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
         default_image_size=512,
+        alias="sd21in",
     ),
     ModelConfig(
         short_name="SD-2.1-v",
@@ -84,19 +92,38 @@ MODEL_CONFIGS = [
         config_path="configs/stable-diffusion-v2-inference-v.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt",
         default_image_size=768,
+        alias="sd20v",
     ),
     ModelConfig(
         short_name="SD-2.0-depth",
         config_path="configs/stable-diffusion-v2-midas-inference.yaml",
         weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-depth/resolve/main/512-depth-ema.ckpt",
         default_image_size=512,
+        alias="sd20dep",
     ),
     ModelConfig(
-        short_name="edit",
+        short_name="instruct-pix2pix",
         config_path="configs/instruct-pix2pix.yaml",
         weights_url="https://huggingface.co/imaginairy/instruct-pix2pix/resolve/ea0009b3d0d4888f410a40bd06d69516d0b5a577/instruct-pix2pix-00-22000-pruned.ckpt",
         default_image_size=512,
         default_negative_prompt="",
+        alias="edit",
+    ),
+    ModelConfig(
+        short_name="openjourney-v1",
+        config_path="configs/stable-diffusion-v1.yaml",
+        weights_url="https://huggingface.co/prompthero/openjourney/resolve/7428477dad893424c92f6ea1cc29d45f6d1448c1/mdjrny-v4.safetensors",
+        default_image_size=512,
+        default_negative_prompt="",
+        alias="oj1",
+    ),
+    ModelConfig(
+        short_name="openjourney-v2",
+        config_path="configs/stable-diffusion-v1.yaml",
+        weights_url="https://huggingface.co/prompthero/openjourney-v2/resolve/47257274a40e93dab7fbc0cd2cfd5f5704cfeb60/openjourney-v2.ckpt",
+        default_image_size=512,
+        default_negative_prompt="",
+        alias="oj2",
     )
     # ModelConfig(
     #     short_name="SD-2.0-upscale",
@@ -107,5 +134,11 @@ MODEL_CONFIGS = [
 ]
 
 MODEL_CONFIG_SHORTCUTS = {m.short_name: m for m in MODEL_CONFIGS}
+for m in MODEL_CONFIGS:
+    if m.alias:
+        MODEL_CONFIG_SHORTCUTS[m.alias] = m
+
+MODEL_CONFIG_SHORTCUTS["openjourney"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"]
+MODEL_CONFIG_SHORTCUTS["oj"] = MODEL_CONFIG_SHORTCUTS["openjourney-v2"]
 
 MODEL_SHORT_NAMES = sorted(MODEL_CONFIG_SHORTCUTS.keys())