Spaces:
Sleeping
Sleeping
| import importlib | |
| from typing import List | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from diffusers import StableDiffusionPipeline | |
| from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure | |
| from image_utils import make_grid, numpy_to_pil | |
| from metrics_utils import compute_main_metrics, compute_psnr_or_ssim | |
| from report_utils import add_psnr_ssim_to_report, prepare_report | |
| SEED = 0 | |
| WEIGHT_DTYPE = torch.float16 | |
| TITLE = "Evaluate Schedulers with StableDiffusionPipeline 🧨" | |
| ABSTRACT = """ | |
| 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. | |
| """ | |
| DESCRIPTION = """ | |
| #### Hoes does it work? | |
| * 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 | |
| * When selecting a model checkpoint, if you select "Other" you will have the option to provide a custom Stable Diffusion checkpoint. | |
| * 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. | |
| """ | |
| psnr_fn = PeakSignalNoiseRatio() | |
| ssim_fn = StructuralSimilarityIndexMeasure() | |
| 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 run( | |
| prompt: str, | |
| num_images_per_prompt: int, | |
| num_inference_steps: int, | |
| checkpoint: str, | |
| other_finedtuned_checkpoints: str = None, | |
| schedulers_to_test: List[str] = None, | |
| ssim: bool = False, | |
| psnr: bool = False, | |
| progress=gr.Progress(), | |
| ): | |
| progress(0, desc="Starting...") | |
| if checkpoint == "Other" and other_finedtuned_checkpoints == "": | |
| return "❌ No legit checkpoint provided ❌" | |
| elif checkpoint == "Other": | |
| checkpoint = other_finedtuned_checkpoints | |
| all_images = {} | |
| scheduler_images = {} | |
| # Set up the pipeline | |
| sd_pipeline, original_scheduler_config = initialize_pipeline(checkpoint) | |
| sd_pipeline.set_progress_bar_config(disable=True) | |
| # Prepare latents to start generation and the prompts. | |
| latents = get_latents(num_images_per_prompt) | |
| prompts = [prompt] * num_images_per_prompt | |
| original_scheduler_name = original_scheduler_config._class_name | |
| schedulers_to_test.append(original_scheduler_name) | |
| # Start generating the images and computing their scores. | |
| for scheduler_name in progress.tqdm(schedulers_to_test): | |
| if scheduler_name != original_scheduler_name: | |
| 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, | |
| latents=latents, | |
| 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_main_metrics(cur_scheduler_images, prompts), | |
| } | |
| } | |
| ) | |
| scheduler_images.update({scheduler_name: cur_scheduler_images}) | |
| torch.cuda.empty_cache() | |
| # Prepare output report. | |
| output_str = "" | |
| for scheduler_name in all_images: | |
| output_str += prepare_report(scheduler_name, all_images[scheduler_name]) | |
| # Append PSNR or SSIM if needed. | |
| if len(schedulers_to_test) > 1: | |
| ssim_scores = psnr_scores = None | |
| if ssim: | |
| ssim_scores = compute_psnr_or_ssim( | |
| ssim_fn, scheduler_images, original_scheduler_name | |
| ) | |
| if psnr: | |
| psnr_scores = compute_psnr_or_ssim( | |
| psnr_fn, scheduler_images, original_scheduler_name | |
| ) | |
| if len(schedulers_to_test) > 1: | |
| ssim_psnr_str = add_psnr_ssim_to_report( | |
| original_scheduler_name, ssim_scores, psnr_scores | |
| ) | |
| if ssim_psnr_str != "": | |
| output_str += ssim_psnr_str | |
| return output_str | |
| with gr.Blocks(title="Scheduler Evaluation") as demo: | |
| gr.Markdown(f"## {TITLE}\n\n\n\n{ABSTRACT}") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Text( | |
| max_lines=1, placeholder="a painting of a dog", label="prompt" | |
| ) | |
| num_images_per_prompt = gr.Slider( | |
| 3, 10, value=3, step=1, label="num_images_per_prompt" | |
| ) | |
| num_inference_steps = gr.Slider( | |
| 10, 100, value=50, step=1, label="num_inference_steps" | |
| ) | |
| 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, | |
| label="model_ckpt", | |
| ) | |
| other_finedtuned_checkpoints = gr.Textbox( | |
| visible=False, | |
| interactive=True, | |
| placeholder="valhalla/sd-pokemon-model", | |
| label="custom_checkpoint", | |
| ) | |
| 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, | |
| label="schedulers_to_test", | |
| ) | |
| ssim = gr.Checkbox(label="Compute SSIM") | |
| psnr = gr.Checkbox(label="Compute PSNR") | |
| evaluation_button = gr.Button(value="Submit") | |
| with gr.Column(): | |
| report = gr.Markdown(label="Evaluation Report").style() | |
| evaluation_button.click( | |
| run, | |
| inputs=[ | |
| prompt, | |
| num_images_per_prompt, | |
| num_inference_steps, | |
| model_ckpt, | |
| other_finedtuned_checkpoints, | |
| schedulers_to_test, | |
| ssim, | |
| psnr, | |
| ], | |
| outputs=report, | |
| ) | |
| gr.Markdown(f"{DESCRIPTION}") | |
| demo.queue().launch(debug=True) | |