add initial stuff

Dockerfile

@@ -0,0 +1,25 @@
+FROM diffusers/diffusers-pytorch-compile-cuda
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,182 @@
+import gradio as gr
+import torch
+from diffusers import (
+    StableDiffusionXLControlNetPipeline,
+    DiffusionPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionInpaintPipeline,
+    StableDiffusionAdapterPipeline,
+    StableDiffusionControlNetPipeline,
+    StableDiffusionXLAdapterPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    StableDiffusionXLInpaintPipeline,
+    ControlNetModel,
+    T2IAdapter,
+)
+import time
+
+
+dtype = torch.float16
+device = torch.device("cuda")
+
+pipeline_mapping = {
+    "SD T2I": (DiffusionPipeline, "runwayml/stable-diffusion-v1-5"),
+    "SD I2I": (StableDiffusionImg2ImgPipeline, "runwayml/stable-diffusion-v1-5"),
+    "SD Inpainting": (
+        StableDiffusionInpaintPipeline,
+        "runwayml/stable-diffusion-inpainting",
+    ),
+    "SD ControlNet": (
+        StableDiffusionControlNetPipeline,
+        "runwayml/stable-diffusion-v1-5",
+        "lllyasviel/sd-controlnet-canny",
+    ),
+    "SD T2I Adapters": (
+        StableDiffusionAdapterPipeline,
+        "CompVis/stable-diffusion-v1-4" "TencentARC/t2iadapter_canny_sd14v1",
+    ),
+    "SDXL T2I": (DiffusionPipeline, "stabilityai/stable-diffusion-xl-base-1.0"),
+    "SDXL I2I": (
+        StableDiffusionXLImg2ImgPipeline,
+        "stabilityai/stable-diffusion-xl-base-1.0",
+    ),
+    "SDXL Inpainting": (
+        StableDiffusionXLInpaintPipeline,
+        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+    ),
+    "SDXL ControlNet": (
+        StableDiffusionXLControlNetPipeline,
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        "diffusers/controlnet-canny-sdxl-1.0",
+    ),
+    "SDXL T2I Adapters": (
+        StableDiffusionXLAdapterPipeline,
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        "TencentARC/t2i-adapter-canny-sdxl-1.0",
+    ),
+}
+
+
+def load_pipeline(
+    pipeline_to_benchmark: str,
+    use_channels_last: bool = False,
+    do_torch_compile: bool = False,
+):
+    # Get pipeline details.
+    pipeline_details = pipeline_mapping[pipeline_to_benchmark]
+    pipeline_cls = pipeline_details[0]
+    pipeline_ckpt = pipeline_details[1]
+
+    # Load adapter if needed.
+    if "ControlNet" in pipeline_to_benchmark:
+        controlnet_ckpt = pipeline_details[2]
+        controlnet = ControlNetModel.from_pretrained(
+            controlnet_ckpt, variant="fp16", torch_dtype=torch.float16
+        ).to(device)
+    elif "Adapters" in pipeline_to_benchmark:
+        adapter_clpt = pipeline_details[2]
+        adapter = T2IAdapter.from_pretrained(
+            adapter_clpt, variant="fp16", torch_dtype=torch.float16
+        ).to(device)
+
+    # Load pipeline.
+    if (
+        "ControlNet" not in pipeline_to_benchmark
+        or "Adapters" not in pipeline_to_benchmark
+    ):
+        pipeline = pipeline_cls.from_pretrained(
+            pipeline_ckpt, variant="fp16", torch_dtype=dtype
+        )
+
+    elif "ControlNet" in pipeline_to_benchmark:
+        pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, controlnet=controlnet)
+    elif "Adapters" in pipeline_to_benchmark:
+        pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, adapter=adapter)
+    pipeline.to(device)
+
+    # Optionally set memory layout.
+    if use_channels_last:
+        pipeline.unet.to(memory_format=torch.channels_last)
+
+    if hasattr(pipeline, "controlnet"):
+        pipeline.controlnet.to(memory_format=torch.channels_last)
+    elif hasattr(pipeline, "adapter"):
+        pipeline.adapter.to(memory_format=torch.channels_last)
+
+    # Optional torch compilation.
+    if do_torch_compile:
+        pipeline.unet = torch.compile(
+            pipeline.unet, mode="reduce-overhead", fullgraph=True
+        )
+        if hasattr(pipeline, "controlnet"):
+            pipeline.controlnet = torch.compile(
+                pipeline.controlnet, mode="reduce-overhead", fullgraph=True
+            )
+        elif hasattr(pipeline, "adapter"):
+            pipeline.adapter = torch.compile(
+                pipeline.adapter, mode="reduce-overhead", fullgraph=True
+            )
+
+    return pipeline
+
+
+def generate(
+    pipeline_to_benchmark: str,
+    num_images_per_prompt: int = 1,
+    use_channels_last: bool = False,
+    do_torch_compile: bool = False,
+):
+    print("Start...")
+    print("Torch version", torch.__version__)
+    print("Torch CUDA version", torch.version.cuda)
+
+    pipeline = load_pipeline(
+        pipeline_to_benchmark=pipeline_to_benchmark,
+        use_channels_last=use_channels_last,
+        do_torch_compile=do_torch_compile,
+    )
+    for _ in range(3):
+        prompt = 77 * "a"
+        num_inference_steps = 20
+        start_time = time.time()
+        _ = pipeline(
+            prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            num_inference_steps=num_inference_steps,
+        ).images
+        end_time = time.time()
+
+        print(f"For {num_inference_steps} steps", end_time - start_time)
+        print("Avg per step", (end_time - start_time) / num_inference_steps)
+
+
+with gr.Blocks() as demo:
+    do_torch_compile = gr.Checkbox(label="Enable torch.compile()?")
+    use_channels_last = gr.Checkbox(label="Use `channels_last` memory layout?")
+    pipeline_to_benchmark = (
+        gr.Dropdown(
+            list(pipeline_mapping.keys()),
+            value=["Stable Diffusion V1.5"],
+            multiselect=False,
+            label="Pipeline to benchmark",
+        ),
+    )
+    batch_size = gr.Slider(
+        label="Number of images per prompt",
+        minimum=1,
+        maximum=16,
+        step=1,
+        value=1,
+    )
+    btn = gr.Button("Benchmark!").style(
+        margin=False,
+        rounded=(False, True, True, False),
+        full_width=False,
+    )
+
+    btn.click(
+        fn=generate,
+        inputs=[pipeline_to_benchmark, batch_size, use_channels_last, do_torch_compile],
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+gradio
+diffusers @ git+https://github.com/huggingface/diffusers