add: initial files.

README.md

@@ -1,6 +1,6 @@
 ---
-title: Evaluate Sd Schedulers
-emoji: 👁
+title: Evaluate StableDiffusionPipeline with Different Schedulers
+emoji: ⏰
 colorFrom: gray
 colorTo: indigo
 sdk: gradio

app.py

@@ -0,0 +1,205 @@
+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
+
+
+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 = f"![img_grid_{scheduler_name}](./{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,
+        ),
+        gr.Dropdown(
+            [
+                "EulerDiscreteScheduler",
+                "PNDMScheduler",
+                "LMSDiscreteScheduler",
+                "DPMSolverMultistepScheduler",
+                "DDIMScheduler",
+            ],
+            value=["LMSDiscreteScheduler"],
+            multiselect=True,
+        ),
+    ],
+    outputs=[gr.Markdown().style()],
+    allow_flagging=False,
+)
+demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+torchmetrics[image]
+transformers 
+diffusers
+accelerate
+numpy