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 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 fixed number of inference steps. Here's how it works: * The users provides: * An input prompt. * Number of images to generate with the prompt. * A checkpoint path compatible with `StableDiffusionPipeline`. You can either select one from the drop-down list or provide a valid path ("valhalla/sd-pokemon-model" for example). * Names of the schedulers to evaluate. * 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) """ 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): 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 demo = gr.Interface( run, inputs=[ gr.Text(max_lines=1, placeholder="a painting of a dog"), gr.Slider(3, 10, value=3), gr.Slider(10, 100, value=50), gr.Dropdown( [ "CompVis/stable-diffusion-v1-4", "runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-base", ], value="CompVis/stable-diffusion-v1-4", multiselect=False, interactive=True, ), gr.Dropdown( [ "EulerDiscreteScheduler", "PNDMScheduler", "LMSDiscreteScheduler", "DPMSolverMultistepScheduler", "DDIMScheduler", ], value=["LMSDiscreteScheduler"], multiselect=True, ), ], outputs=[gr.Markdown().style()], title=TITLE, description=DESCRIPTION, allow_flagging=False, ) demo.launch()