update scritps

download_models.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+from huggingface_hub import snapshot_download
+from mlx_app.stable_diffusion.models import _MODELS
+
+# HF_ENDPOINT=https://hf-mirror.com python download_models.py
+
+
+for repo_id, items in _MODELS.items():
+    snapshot_download(repo_id, local_dir=f'../models/{repo_id}', local_dir_use_symlinks=False,
+                      allow_patterns=list(items.values()))
+
+# snapshot_download(repo_id, local_dir=local_path, local_dir_use_symlinks=False,
+#                   ignore_patterns=[
+#                       ".gitattributes", "*.bin", "*.onnx", "*.ckpt", "*.onnx_data", "*.png", "*.jpg", "*.md"
+#                   ],
+#                   allow_patterns=["text_encoder/model.safetensors"]
+#                   )
+
+

test_configs.py

@@ -0,0 +1,178 @@
+from pathlib import Path
+
+output = Path("./base_model_images")
+if not output.exists():
+    output.mkdir(exist_ok=True, parents=True)
+
+prompts = {
+    0: "astronaut riding a horse",
+    1: "a cute corgi",
+    2: "A cinematic shot of a baby racoon wearing an intricate italian priest robe",
+    3: "portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hair",
+    4: "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
+}
+
+# base_models = ["sd1.5", "sd2", "realistic", "sdxl", "sdxl-turbo"]
+# base_models = ["sdxl-turbo"]
+# sd1.5 512, sd2 512, realistic 512, sdxl 1024, sdxl-turbo 512
+base_models = {
+    "sd1.5": {
+        "img-size": 512,
+        #     "steps": 40,
+        #     "cfg": 1.2,
+        # "loras": [
+            # {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lcm-sdv15.safetensors",
+            #     "trigger_words": ""
+            # },
+            # {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/ColoringBook-sd15.safetensors",
+            #     "trigger_words": "Coloring Book, ColoringBookAF"
+            # },
+            # {
+            #     "lora-scale": 0.8,
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lego-sd15.safetensors",
+            #     "trigger_words": "LEGO Creator"
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/toyglasses-sd15.safetensors",
+            #     "trigger_words": "<lora:toyglasses:1>toyglasses"
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/song_flower-sd15.safetensors",
+            #     "trigger_words": "Song Dynasty flower and bird painting"
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/3dillu-sd15.safetensors",
+            #     "trigger_words": ""
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/Icons-sd15.safetensors",
+            #     "trigger_words": " icons, ios icon app",
+            #     "lora-scale": 0.8,
+            # },{
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/CuteCartoon-sd15.safetensors",
+            #     "trigger_words": "Cartoon,CuteCartoonAF",
+            # },
+        # ]
+    },
+    "sd2": {
+        # "steps": 40,
+        "img-size": 512,
+        # "loras": [
+        #     {
+        #         "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/ColoringBook-sd21.safetensors",
+        #         "trigger_words": "ColoringBookAF",
+        #     },
+            #         {
+            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/TShirtDesign-sd21.safetensors",
+            #             "trigger_words": "T Shirt Design, TShirtDesignAF",
+            #         }, {
+            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/StudioGhibli-sd21.safetensors",
+            #             "trigger_words": "Studio Ghibli, StdGBRedmAF",
+            #         }, {
+            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/3D-sd21.safetensors",
+            #             "trigger_words": "3D Render Style, 3DRenderAF",
+            #         }, {
+            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/Stickers-sd21.safetensors",
+            #             "trigger_words": " Sticker",
+            #         },
+        # ]
+    },
+    "realistic": {
+        "img-size": 512,
+        # "steps": 20,
+        # "loras": [
+        #     {
+        #         "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lego-sd15.safetensors",
+        #         "trigger_words": "LEGO Creator",
+        #         "lora-scale": 0.8,  # Between 0.6-1.0, recommended to use 0.8.
+        #     },
+            # {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/toyglasses-sd15.safetensors",
+            #     "trigger_words": "<lora:toyglasses:1>toyglasses"
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/song_flower-sd15.safetensors",
+            #     "trigger_words": "Song Dynasty flower and bird painting"
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/3dillu-sd15.safetensors",
+            #     "trigger_words": ""
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/Icons-sd15.safetensors",
+            #     "trigger_words": " icons, ios icon app",
+            #     "lora-scale": 0.8,
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/CuteCartoon-sd15.safetensors",
+            #     "trigger_words": "Cartoon,CuteCartoonAF",
+            # },
+        # ],
+    },
+    "sdxl": {
+        # "steps": 40,
+        "img-size": 1024,
+        # "loras": [
+        #     {
+        #         "lora": "/Users/jeqin/work/code/sd/models/lora_xl/papercut_sdxl.safetensors",
+        #         "trigger_words": "papercut style",
+        #     },
+            # {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/toy_face-sdxl.safetensors",
+            #     "trigger_words": "toy_face",
+            #     "lora-scale": 0.9,
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Products10k-sdxl.safetensors",
+            #     "trigger_words": "",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/ikea_instructions-sdxl.safetensors",
+            #     "trigger_words": "",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Lego_sdxl.safetensors",
+            #     "trigger_words": "LEGO MiniFig",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/lego_cinematic_sdxl.safetensors",
+            #     "trigger_words": "Lego",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Emojis-sdxl.safetensors",
+            #     "trigger_words": "Emoji",
+            # },  {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/MoviePoster-sdxl.safetensors",
+            #     "trigger_words": "Movie Poster, MoviePosterAF",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/XboxAvatar-sdxl.safetensors",
+            #     "trigger_words": "XBOX AVATAR",
+            # },
+        # ],
+    },
+    "sdxl-turbo": {
+        "img-size": 512,
+        # "steps": 4,
+        # "loras": [
+        #     {
+        #         "lora": "/Users/jeqin/work/code/sd/models/lora_xl/papercut_sdxl.safetensors",
+        #         "trigger_words": "papercut style"
+        #     },
+            # {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/toy_face-sdxl.safetensors",
+            #     "trigger_words": "toy_face",
+            #     "lora-scale": 0.9,
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Products10k-sdxl.safetensors",
+            #     "trigger_words": "",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/ikea_instructions-sdxl.safetensors",
+            #     "trigger_words": "",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Lego_sdxl.safetensors",
+            #     "trigger_words": "LEGO MiniFig",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/lego_cinematic_sdxl.safetensors",
+            #     "trigger_words": "Lego",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Emojis-sdxl.safetensors",
+            #     "trigger_words": "Emoji",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/MoviePoster-sdxl.safetensors",
+            #     "trigger_words": "Movie Poster, MoviePosterAF",
+            # }, {
+            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/XboxAvatar-sdxl.safetensors",
+            #     "trigger_words": "XBOX AVATAR",
+            # },
+        # ]
+    }
+}

test_mlx.py

@@ -0,0 +1,126 @@
+import json
+import re
+import subprocess
+import csv
+from subprocess import CompletedProcess
+
+from test_configs import *
+
+
+def cmd(command: str, check=True, capture_output=False) -> CompletedProcess:
+    print(command)
+    if capture_output:
+        ret = subprocess.run(command, shell=True, check=check, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                             universal_newlines=True)
+    else:
+        ret = subprocess.run(command, shell=True, check=check)
+    print(ret.stdout)
+    return ret
+
+
+def parse_log(output):
+    """output example：
+    """
+    model_name = re.search(r"model: (.+)", output).group(1)
+    steps = re.search(r"steps: (.+)", output).group(1)
+    cfg_weight = re.search(r"cfg_weight: (.+)", output).group(1)
+    img_size = re.search(r"img-size: (.+)", output).group(1)
+    img_number = re.search(r"image number: (.+)", output).group(1)
+    load_model_time = re.search(r"load model time: (.+)", output).group(1)
+    preload_model = re.search(r"preload model time: (.+)", output).group(1)
+    update_lora_time = re.search(r"update lora time: (.+)", output).group(1)
+    quantize_time = re.search(r"quantize time: (.+)", output).group(1)
+    generate_time = re.search(r"generate image time: (.+)", output).group(1)
+    total_time = re.search(r"total time: (.+)", output).group(1)
+    out_image = re.search(r"save image to: (.+)", output).group(1)
+    out_image = '/'.join(out_image.split("/")[-2:])
+    out_image_size = re.search(r"output image size: \((.+)\)", output).group(1)
+    out_image_size = out_image_size.replace(', ', '*')
+    return (model_name, steps, cfg_weight, img_size, img_number, load_model_time,
+            preload_model, update_lora_time, quantize_time, generate_time, total_time, out_image, out_image_size)
+
+
+def _get_cmd(prompt, **kwargs):
+    base_cmd = f'python mlx_app/txt2image_lora.py "{prompt}"'
+    for k, v in kwargs.items():
+        if v is True:
+            base_cmd += f" --{k}"
+        else:
+            base_cmd += f" --{k} {v}"
+    return base_cmd
+
+
+def test_lora(result):
+    commands = {
+        "no_lora": [],
+        "no_trigger": [],
+        "with_trigger": []
+    }
+    for model, config in base_models.items():
+        loras = config.pop("loras")
+        for l in loras:
+            trigger_words = l.get("trigger_words")
+            lora_name = l.get("lora").split('/')[-1].split('.')[0]
+            for i, p in prompts.items():
+                # 1. run with no lora
+                # paras = {"model": model, "output": str(output / model / f"{lora_name}-{i}-a_no_lora.png")}
+                # paras.update(config)
+                # commands["no_lora"].append(_get_cmd(p, **paras))
+                #
+                # # 2. run with lora, but no trigger words
+                # paras = {"model": model, "output": str(output / model / f"{lora_name}-{i}-b_no_trigger.png")}
+                # paras.update(config)
+                # paras["lora"] = l.get("lora")
+                # if l.get("lora-scale"):
+                #     paras["lora-scale"] = l.get("lora-scale")
+                # commands["no_trigger"].append(_get_cmd(p, **paras))
+
+                # 3. run with lora, with trigger words
+                paras = {"model": model, "output": str(output / f"{model}-{lora_name}-{i}-c_with_trigger.png"),
+                         "n_images": 4}
+                paras.update(config)
+                paras["lora"] = l.get("lora")
+                if l.get("lora-scale"):
+                    paras["lora-scale"] = l.get("lora-scale")
+                p = f"{p}, {trigger_words}"
+                commands["with_trigger"].append(_get_cmd(p, **paras))
+
+    for _, cmds in commands.items():
+        for c in cmds:
+            try:
+                ret = cmd(c, capture_output=True)
+                result.append(parse_log(ret.stdout))
+            except Exception as e:
+                print("Exception: ", e)
+    return result
+
+
+def test_base_model(result: list):
+    for model, config in base_models.items():
+        for i, p in prompts.items():
+            paras = {"model": model, "output": str(output / f"{model}_{i}.png"), "n_images": 4, "decoding_batch_size": 4}
+            paras.update(config)
+            command = _get_cmd(p, **paras)
+            try:
+                ret = cmd(command, capture_output=True)
+                result.append(parse_log(ret.stdout))
+            except Exception as e:
+                print("Exception: ", e)
+    return result
+
+
+def main():
+    result = [
+        ['model name', 'steps', 'cfg_weight', 'img size', 'img number', 'load model', 'preload model', 'update lora',
+         'quantize', 'generate image', 'total time', 'output image', 'output image size']
+    ]
+    result = test_base_model(result)
+    # result = test_lora(result)
+
+    with open("result_mlx.csv", 'w', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerows(result)
+
+
+if __name__ == '__main__':
+    main()