demo-custom-css / app.py
sayakpaul's picture
sayakpaul HF staff
add: placeholder
d9658d9
import importlib
from functools import partial
from typing import List
import gradio as gr
import numpy as np
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
from torchmetrics.functional.multimodal import clip_score
from torchmetrics.image.inception import InceptionScore
SEED = 0
WEIGHT_DTYPE = torch.float16
TITLE = "Evaluate Schedulers with StableDiffusionPipeline 🧨"
DESCRIPTION = """
This Space allows you to quantitatively compare [different noise schedulers](https://huggingface.co/docs/diffusers/using-diffusers/schedulers) with a [`StableDiffusionPipeline`](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/overview).
One of the applications of this Space could be to evaluate different schedulers for a certain Stable Diffusion checkpoint for a fixed number of inference steps.
Here's how it works:
* The evaluator first sets a seed and then generates the initial noise which is passed as the initial latent to start the image generation process. It is done to ensure fair comparison.
* This initial latent is used every time the pipeline is run (with different schedulers).
* To quantify the quality of the generated images we use:
* [Inception Score](https://en.wikipedia.org/wiki/Inception_score)
* [Clip Score](https://arxiv.org/abs/2104.08718)
**Notes**:
* The default scheduler associated with the provided checkpoint is always used for reporting the scores.
* Increasing both the number of images per prompt and the number of inference steps could quickly build up the inference queue and thus
resulting in slowdowns.
"""
inception_score_fn = InceptionScore(normalize=True)
torch.manual_seed(SEED)
clip_score_fn = partial(clip_score, model_name_or_path="openai/clip-vit-base-patch16")
def make_grid(images, rows, cols):
w, h = images[0].size
grid = Image.new("RGB", size=(cols * w, rows * h))
for i, image in enumerate(images):
grid.paste(image, box=(i % cols * w, i // cols * h))
return grid
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_utils.py#L814
def numpy_to_pil(images):
"""
Convert a numpy image or a batch of images to a PIL image.
"""
if images.ndim == 3:
images = images[None, ...]
images = (images * 255).round().astype("uint8")
if images.shape[-1] == 1:
# special case for grayscale (single channel) images
pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
else:
pil_images = [Image.fromarray(image) for image in images]
return pil_images
def prepare_report(scheduler_name: str, results: dict):
image_grid = results["images"]
scores = results["scores"]
img_str = ""
image_name = f"{scheduler_name}_images.png"
image_grid.save(image_name)
img_str = img_str = f"![img_grid_{scheduler_name}](/file=./{image_name})\n"
report_str = f"""
\n\n## {scheduler_name}
### Sample images
{img_str}
### Scores
{scores}
\n\n
"""
return report_str
def initialize_pipeline(checkpoint: str):
sd_pipe = StableDiffusionPipeline.from_pretrained(
checkpoint, torch_dtype=WEIGHT_DTYPE
)
sd_pipe = sd_pipe.to("cuda")
original_scheduler_config = sd_pipe.scheduler.config
return sd_pipe, original_scheduler_config
def get_scheduler(scheduler_name: str):
schedulers_lib = importlib.import_module("diffusers", package="schedulers")
scheduler_abs = getattr(schedulers_lib, scheduler_name)
return scheduler_abs
def get_latents(num_images_per_prompt: int, seed=SEED):
generator = torch.manual_seed(seed)
latents = np.random.RandomState(seed).standard_normal(
(num_images_per_prompt, 4, 64, 64)
)
latents = torch.from_numpy(latents).to(device="cuda", dtype=WEIGHT_DTYPE)
return latents
def compute_metrics(images: np.ndarray, prompts: List[str]):
inception_score_fn.update(torch.from_numpy(images).permute(0, 3, 1, 2))
inception_score = inception_score_fn.compute()
images_int = (images * 255).astype("uint8")
clip_score = clip_score_fn(
torch.from_numpy(images_int).permute(0, 3, 1, 2), prompts
).detach()
return {
"inception_score (⬆️)": {
"mean": round(float(inception_score[0]), 4),
"std": round(float(inception_score[1]), 4),
},
"clip_score (⬆️)": round(float(clip_score), 4),
}
def run(
prompt: str,
num_images_per_prompt: int,
num_inference_steps: int,
checkpoint: str,
schedulers_to_test: List[str],
):
all_images = {}
sd_pipeline, original_scheduler_config = initialize_pipeline(checkpoint)
latents = get_latents(num_images_per_prompt)
prompts = [prompt] * num_images_per_prompt
images = sd_pipeline(
prompts,
latents=latents,
num_inference_steps=num_inference_steps,
output_type="numpy",
).images
original_scheduler_name = original_scheduler_config._class_name
all_images.update(
{
original_scheduler_name: {
"images": make_grid(numpy_to_pil(images), 1, num_images_per_prompt),
"scores": compute_metrics(images, prompts),
}
}
)
# print("First scheduler complete.")
for scheduler_name in schedulers_to_test:
if scheduler_name == original_scheduler_name:
continue
scheduler_cls = get_scheduler(scheduler_name)
current_scheduler = scheduler_cls.from_config(original_scheduler_config)
sd_pipeline.scheduler = current_scheduler
cur_scheduler_images = sd_pipeline(
prompts, num_inference_steps=num_inference_steps, output_type="numpy"
).images
all_images.update(
{
scheduler_name: {
"images": make_grid(
numpy_to_pil(cur_scheduler_images), 1, num_images_per_prompt
),
"scores": compute_metrics(cur_scheduler_images, prompts),
}
}
)
# print(f"{scheduler_name} complete.")
output_str = ""
for scheduler_name in all_images:
# print(f"scheduler_name: {scheduler_name}")
output_str += prepare_report(scheduler_name, all_images[scheduler_name])
# print(output_str)
return output_str
with gr.Blocks() as demo:
gr.HTML(f"<div align='center'{TITLE}</div>")
with gr.Row():
with gr.Column():
prompt = gr.Text(max_lines=1, placeholder="a painting of a dog")
num_images_per_prompt = gr.Slider(3, 10, value=3, step=1)
num_inference_steps = gr.Slider(10, 100, value=50, step=1)
model_ckpt = gr.Dropdown(
[
"CompVis/stable-diffusion-v1-4",
"runwayml/stable-diffusion-v1-5",
"stabilityai/stable-diffusion-2-base",
"Other"
],
value="CompVis/stable-diffusion-v1-4",
multiselect=False,
interactive=True,
)
other_finedtuned_checkpoints = gr.Text(visible=False, placeholder="valhalla/sd-pokemon-model")
model_ckpt.change(lambda x: gr.Dropdown.update(visible=x=="Other"), model_ckpt, other_finedtuned_checkpoints)
schedulers_to_test = gr.Dropdown(
[
"EulerDiscreteScheduler",
"PNDMScheduler",
"LMSDiscreteScheduler",
"DPMSolverMultistepScheduler",
"DDIMScheduler",
],
value=["LMSDiscreteScheduler"],
multiselect=True,
)
evaluation_button = gr.Button(value="Submit")
with gr.Column():
report = gr.Markdown(label="Evaluation Report")
demo.launch()