add scripts

download_from_hf.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+# HF_ENDPOINT=https://hf-mirror.com python download_from_hf.py
+
+repo_id = "stabilityai/stable-diffusion-3-medium-diffusers"
+# repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
+
+local_path = str(Path("/Users/jeqin/work/code/sd/models") / repo_id)
+# snapshot_download(repo_id, local_dir=local_path, local_dir_use_symlinks=False,
+#                   # ignore_patterns=[".gitattributes", "*.md", "*.bin", "*.safetensors", "*.ckpt", "*.onnx"],
+#                   allow_patterns=pattern_list)
+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"]
+                  )
+
+

print_model.py

@@ -0,0 +1,21 @@
+from safetensors import safe_open
+
+
+def print_keys(model_path):
+    tensors = {}
+    with safe_open(model_path, framework="pt") as f:
+        print("key length: ", len(f.keys()))
+        for k in f.keys():
+            v = f.get_tensor(k)
+            if v.shape:
+                print(f"{k}: {type(v)}, {v.shape}")
+            else:
+                print(f"{k}: {v}")
+            tensors[k] = v
+
+
+if __name__ == '__main__':
+    model = "../models/lora_1.5/ColoringBook-sd15.safetensors"
+    # model = "../models/lora_2.1/pytorch_lora_weights-sd21-comfyui.safetensors"
+    print(model)
+    print_keys(model)

run_sd_with_lora.py

@@ -0,0 +1,98 @@
+from diffusers import AutoPipelineForText2Image
+from diffusers import TCDScheduler, LCMScheduler
+from diffusers.utils import make_image_grid
+import torch
+from PIL import Image
+import time
+from pathlib import Path
+import argparse
+
+device = torch.device("mps")
+
+folder = Path("../models/")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Generate images from a textual prompt using stable diffusion"
+    )
+    parser.add_argument("prompt")
+    parser.add_argument("--model", choices=["sd1.5", "sd2", "realistic", "sdxl", "sdxl-turbo"], default="sdxl")
+    parser.add_argument("--n_images", type=int, default=1)
+    parser.add_argument("--steps", type=int)
+    parser.add_argument("--guidance_scale", type=float)
+    parser.add_argument("--negative_prompt", default="")
+    parser.add_argument("--output", default="out.png")
+    parser.add_argument("--img-size", type=int, default=512)
+    parser.add_argument("--lora", type=str, default=None)
+    parser.add_argument("--lora-scale", type=float, default=None)
+    args = parser.parse_args()
+
+    if args.model == "sdxl-turbo":
+        model_path = folder / "stabilityai/sdxl-turbo"
+        args.guidance_scale = args.guidance_scale or 0.0
+        args.steps = args.steps or 2
+    elif args.model == "sdxl":
+        model_path = folder / "stabilityai/stable-diffusion-xl-base-1.0"
+        args.guidance_scale = args.guidance_scale or 7.5
+        args.steps = args.steps or 20
+    elif args.model == "sd2":
+        model_path = folder / "stabilityai/stable-diffusion-2-1-base"
+        args.guidance_scale = args.guidance_scale or 7.5
+        args.steps = args.steps or 20
+    elif args.model == "sd1.5":
+        model_path = folder / "runwayml/stable-diffusion-v1-5"
+        args.guidance_scale = args.guidance_scale or 7.5
+        args.steps = args.steps or 20
+    elif args.model == "realistic":
+        model_path = folder / "SG161222/Realistic_Vision_V3.0_VAE"
+        args.guidance_scale = args.guidance_scale or 5  # 3,5-7
+        args.steps = args.steps or 10
+    else:
+        raise ValueError(f"Unknown model: {args.model}")
+
+    print("*" * 10, "configurations")
+    print(f"model: {args.model}\nimage number: {args.n_images}\nsteps: {args.steps}\n"
+          f"guidance_scale: {args.guidance_scale}\nnegative_prompt: {args.negative_prompt}\noutput: {args.output}\n"
+          f"img-size: {args.img_size}\nlora: {args.lora}\nlora-scale: {args.lora_scale}\nprompt:{args.prompt}\n")
+    print("*" * 10)
+
+    t0 = time.time()
+    pipe = AutoPipelineForText2Image.from_pretrained(model_path,
+                                                     torch_dtype=torch.float16).to(device)
+    t1 = time.time()
+    print(f"load model time: {(t1 - t0):.3f}")
+
+    if args.lora:
+        t_load_lora = time.time()
+        # pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        pipe.load_lora_weights(args.lora)
+        # pipe.fuse_lora()
+    t2 = time.time()
+    print(f"load lora time: {(t2 - t1):.3f}")
+
+    if args.lora_scale:
+        output = pipe(prompt=args.prompt,
+                      height=args.img_size,
+                      width=args.img_size,
+                      num_inference_steps=args.steps,
+                      num_images_per_prompt=args.n_images,
+                      guidance_scale=args.guidance_scale,
+                      cross_attention_kwargs={'scale': args.lora_scale},
+                      # generator=torch.Generator(device="mps").manual_seed(0)
+                      )
+    else:
+        output = pipe(prompt=args.prompt,
+                      height=args.img_size,
+                      width=args.img_size,
+                      num_inference_steps=args.steps,
+                      num_images_per_prompt=args.n_images,
+                      guidance_scale=args.guidance_scale,
+                      # generator=torch.Generator(device="mps").manual_seed(0)
+                      )
+    t3 = time.time()
+    print(f"generate image time: {(t3 - t2):.3f}")
+    img = make_image_grid(output.images, rows=1, cols=args.n_images)
+    img.save(args.output)
+    print(f"save image to: {args.output}")
+    print(f"output image size: {img.size}")
+    print(f"total time: {(time.time() - t0):.3f}")

test_configs.py

@@ -0,0 +1,189 @@
+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 model with default config
+# base_models = ["sdxl"]
+# sd1.5 512, sd2 768, realistic 512, sdxl 1024, sdxl-turbo 1024
+# sd1.5 512, sd2 512, realistic 512, sdxl 1024, sdxl-turbo 512
+# model(guidance_scale, steps), sd1.5(7.5, 20-50), sd2(7.5, 20-50),realistic(3-5, 10) sdxl(7.5, 20-50), sdxl-turbo(0.0, 2-5 )
+
+base_models = {
+    "sd1.5": {
+        # "steps": 40,
+        "img-size": 512,
+        # "guidance_scale": 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,
+        # "n_images": "8"
+        # "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": {
+        # "steps": 20,
+        "img-size": 512,
+        # "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": {
+        # "n_images": "6",
+        "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_diffusers.py

@@ -0,0 +1,121 @@
+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"guidance_scale: (.+)", 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)
+    update_lora_time = re.search(r"load lora 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,
+            update_lora_time, generate_time, total_time, out_image, out_image_size)
+
+
+def _get_cmd(prompt, **kwargs):
+    base_cmd = f'python run_sd_with_lora.py "{prompt}"'
+    for k, v in kwargs.items():
+        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}
+            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', 'guidance scale', 'img size', 'img number', 'load model', 'update lora',
+         'generate image', 'total time', 'output image', 'output image size']
+    ]
+    result = test_base_model(result)
+    # result = test_lora(result)
+
+    with open("result_diffuser.csv", 'w', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerows(result)
+
+
+if __name__ == '__main__':
+    main()