LCM-LoRA vs Origin Model

lcm-vs-origin.py

# Copyright 2023 https://novita.ai

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import os
import time
import torch
from diffusers import LCMScheduler, AutoPipelineForText2Image, DPMSolverMultistepScheduler
import math

from PIL import Image, ImageDraw, ImageFont


def check_model_is_sdxl(model_path):
    return os.path.exists(os.path.join(model_path, "text_encoder_2"))


models_list = [
    # "models/checkpoint/darkSushiMixMix_225D",
    # "models/checkpoint/dreamshaper_8",
    # "models/checkpoint/cyberrealistic_classicV14_73029",
    # "models/checkpoint/arthemyComics_v50_126515",
    # "models/checkpoint/based65_v20_29967",
    # "models/checkpoint/CounterfeitV30_v30",
    # "models/checkpoint/endlessmix_v9Realmix_56039",
    # "models/checkpoint/cetusMix_cetusVersion0_7926",
    # "models/checkpoint/toonyou_beta5Unstable_72242",
    # "models/checkpoint/AnythingV5_v5PrtRE",
    # "models/checkpoint/CheckpointYesmix_v15_10395",
    # "models/checkpoint/protovisionXLHighFidelity3D_release0620Bakedvae"
]
# sdxl
models_list = [
    "models/checkpoint/sd_xl_base_1.0",
    "models/checkpoint/protovisionXLHighFidelity3D_release0620Bakedvae",
    "models/checkpoint/sdxlNijiV51_sdxlNijiV51_112807",
    # "models/checkpoint/zavychromaxl_v21_129006",
    "models/checkpoint/crystalClearXL_ccxl_97637",
    "models/checkpoint/sdxlYamersAnimeUltra_yamersAnimeV3_121537",
]

prompts_list = [
    "a photo of a cat",
    "a photo of a dog",
    "a photo of a bird",
    "a photo of a cute girl, 8 years",
    # "a photo of a cute boy, 8 years",
    # "a photo of a handsome man, 20 years",
]


def timeit(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        print(f"{func.__name__} - {args, kwargs} took {time.time() - start} seconds")
        return result
    return wrapper


@timeit
def generate_by_lcm(model, prompt, negative_prompt, steps, width, height, seed, sdxl):
    if sdxl:
        adapter_id = "latent-consistency/lcm-lora-sdxl"
    else:
        adapter_id = "latent-consistency/lcm-lora-sdv1-5"

    pipe = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
    pipe.safety_checker = None
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.enable_xformers_memory_efficient_attention()

    pipe.to("cuda")
    # load and fuse lcm lora
    pipe.load_lora_weights(adapter_id)
    pipe.fuse_lora(lora_scale=1.0)

    with torch.no_grad():
        images = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=0, width=width, height=height, generator=torch.manual_seed(seed)).images
    pipe.to("meta")

    del pipe
    return images


@timeit
def generate_by_origin(model, prompt, negative_prompt, steps, width, height, seed):
    pipe = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
    pipe.safety_checker = None
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe.enable_xformers_memory_efficient_attention()

    pipe.to("cuda")

    with torch.no_grad():
        images = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=7.5, width=width, height=height, generator=torch.manual_seed(seed)).images
    pipe.to("meta")
    del pipe
    return images


def make_image_grid1(images, rows: int, cols: int, resize: int = None, padding: int = 10):
    """
    Prepares a single grid of images with padding between images.
    """
    assert len(images) == rows * cols

    if resize is not None:
        images = [img.resize((resize, resize)) for img in images]

    w, h = images[0].size
    # Adjust grid size to include padding
    grid_width = cols * w + (cols - 1) * padding
    grid_height = rows * h + (rows - 1) * padding
    grid = Image.new("RGB", size=(grid_width, grid_height))

    for i, img in enumerate(images):
        # Calculate position with padding
        x = (i % cols) * (w + padding)
        y = (i // cols) * (h + padding)
        grid.paste(img, box=(x, y))

    return grid

# Function to merge images in a grid with labels on top and left


def make_image_grid2(images, rows: int, cols: int, row_labels, col_labels, resize: int = None, padding: int = 10, top_text_space: int = 50, left_text_space: int = 500):
    """
    Prepares a single grid of images with padding between images and adds text descriptions
    to the left of each row and above each column.
    """
    assert len(images) == rows * cols
    assert len(row_labels) == rows
    assert len(col_labels) == cols

    if resize is not None:
        images = [img.resize((resize, resize)) for img in images]

    w, h = images[0].size
    # Adjust grid size to include padding and text space
    grid_width = cols * w + (cols - 1) * padding + left_text_space
    grid_height = rows * h + (rows - 1) * padding + top_text_space
    grid = Image.new("RGB", size=(grid_width, grid_height), color='white')

    # Font settings for text
    font = ImageFont.truetype("Arial.ttf", 30)

    for i, img in enumerate(images):
        # Calculate position with padding
        x = (i % cols) * (w + padding) + left_text_space
        y = (i // cols) * (h + padding) + top_text_space
        grid.paste(img, box=(x, y))

    draw = ImageDraw.Draw(grid)

    # Add row labels
    for i, label in enumerate(row_labels):
        draw.text((padding, i * (h + padding) + left_text_space), label, fill="black", font=font)

    # Add column labels
    for i, label in enumerate(col_labels):
        draw.text((i * (w + padding) + top_text_space + w / 2, 0), label, fill="black", font=font)

    return grid


if __name__ == '__main__':
    torch.set_num_threads(1)
    top_labels = ['LCM (8 steps)', 'Origin (20 steps, DPM++ 2M)']
    left_labels = []
    final_images = []
    for model in models_list:
        left_labels.append(os.path.basename(model))
        if check_model_is_sdxl(model):
            lcm_images = [img.resize((512, 512)) for img in generate_by_lcm(model, prompt=prompts_list, negative_prompt=None, steps=8, width=1024, height=1024, seed=1234, sdxl=True)]
            origin_images = [img.resize((512, 512)) for img in generate_by_origin(model, prompt=prompts_list, negative_prompt=None, steps=20, width=1024, height=1024, seed=1234)]
        else:
            lcm_images = generate_by_lcm(model, prompt=prompts_list, negative_prompt=None, steps=8, width=512, height=512, seed=1234, sdxl=False)
            origin_images = generate_by_origin(model, prompt=prompts_list, negative_prompt=None, steps=20, width=512, height=512, seed=1234)
        lcm_images_merged = make_image_grid1(lcm_images, rows=len(lcm_images) // 2, cols=2, padding=0)
        origin_images_merged = make_image_grid1(origin_images, rows=len(origin_images) // 2, cols=2, padding=0)
        final_images.append(lcm_images_merged)
        final_images.append(origin_images_merged)

    # merge_images_with_labels(final_images, top_labels, left_labels, label_width=50).save("result.png")
    make_image_grid2(final_images, rows=len(final_images) // 2, cols=2, padding=50, row_labels=left_labels, col_labels=top_labels).save("result.png")

    # create image grid group by LCM and origin