Diffusers' ControlNet Implementation Subjective Evaluation

plot_controlnet_sample.py

# Diffusers' ControlNet Implementation Subjective Evaluation
# https://github.com/takuma104/diffusers/tree/controlnet

import einops
import numpy as np
import pytest
import torch

from diffusers import StableDiffusionControlNetPipeline
from diffusers.utils import load_image

import matplotlib.pyplot as plt
from PIL import Image

plt.rcParams["figure.figsize"] = (10,5)
plt.rcParams['figure.facecolor'] = 'white'

model_id_sd15_canny = "takuma104/control_sd15_canny"
test_prompt = "best quality, extremely detailed, illustration, looking at viewer"
test_negative_prompt = (
    "longbody, lowres, bad anatomy, bad hands, missing fingers, "
    + "pubic hair,extra digit, fewer digits, cropped, worst quality, low quality"
)

def pil_image_to_control(image):
    # code from https://github.com/lllyasviel/ControlNet/blob/main/gradio_canny2image.py
    control = torch.from_numpy(np.array(image).copy()).float().cuda() / 255.0
    control = torch.stack([control for _ in range(1)], dim=0)
    control = einops.rearrange(control, "b h w c -> b c h w").clone()
    return control

def generate_image(seed, control):
    image = pipe(
        prompt=test_prompt,
        negative_prompt=test_negative_prompt,
        guidance_scale=9.0,
        num_inference_steps=20,
        generator=torch.Generator(device="cuda").manual_seed(seed),
        controlnet_hint=control,
    ).images[0]
    return image

def save_reference_image(seed, fn):
    # reference image generated by https://gist.github.com/takuma104/6cdb6d9aa27f67462f11554cccdf4b34
    output_ref_image = load_image(
        f"https://huggingface.co/takuma104/controlnet_dev/resolve/main/vermeer_canny_edged_seed_{seed}.png"
    )
    output_ref_image.save(fn)

def render_figure(fn):
    def plot_row(axs, fn_prefix, name):
        for i, ax in enumerate(axs):
            if i == 0:
                ax.set_title(f'Control ({name})')
                ax.imshow(Image.open('control.png'))
            else:
                ax.set_title(f'Seed={i-1} ({name})')
                ax.imshow(Image.open(f'{fn_prefix}_{i-1}.png'))

    fig, axs = plt.subplots(2, 5)
    for ax in axs.flatten():
        ax.set_aspect('equal', 'box')
        ax.axis('off')
    plot_row(axs[0], 'gen', 'ours')
    plot_row(axs[1], 'ref', 'ref impl.')
    fig.tight_layout()
    fig.savefig(fn)


if __name__ == '__main__':
    pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id_sd15_canny).to("cuda")
    pipe.enable_attention_slicing(1)

    canny_edged_image = load_image(
        "https://huggingface.co/takuma104/controlnet_dev/resolve/main/vermeer_canny_edged.png"
    )
    canny_edged_image.save('control.png')
    control = pil_image_to_control(canny_edged_image)

    for seed in range(4):
        image = generate_image(seed=seed, control=control)
        image.save(f"gen_{seed}.png")

    for seed in range(4):
        save_reference_image(seed=seed, fn=f'ref_{seed}.png')

    render_figure('fig_diffusers_controlnet.png')