init

gradio_dmd.py

@@ -0,0 +1,71 @@
+import argparse
+import time
+
+import gradio as gr
+import torch
+from diffusers import DiffusionPipeline, UNet2DConditionModel
+
+from scheduling_dmd import DMDScheduler
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--unet-path", type='Lykon/dreamshaper-8')
+parser.add_argument("--model-path", type='aaronb/dreamshaper-8-dmd-1kstep')
+args = parser.parse_args()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+unet = UNet2DConditionModel.from_pretrained(args.unet_path)
+pipe = DiffusionPipeline.from_pretrained(args.model_path, unet=unet)
+pipe.scheduler = DMDScheduler.from_config(pipe.scheduler.config)
+pipe.to(device=device, dtype=torch.float16)
+
+
+def predict(prompt, seed=1231231):
+    generator = torch.manual_seed(seed)
+    last_time = time.time()
+
+    image = pipe(
+        prompt,
+        num_inference_steps=1,
+        guidance_scale=0.0,
+        generator=generator,
+    ).images[0]
+
+    print(f"Pipe took {time.time() - last_time} seconds")
+    return image
+
+
+css = """
+#container{
+    margin: 0 auto;
+    max-width: 40rem;
+}
+#intro{
+    max-width: 100%;
+    text-align: center;
+    margin: 0 auto;
+}
+"""
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="container"):
+        gr.Markdown(
+            """# Distribution Matching Distillation
+            """,
+            elem_id="intro",
+        )
+        with gr.Row():
+            with gr.Row():
+                prompt = gr.Textbox(placeholder="Insert your prompt here:", scale=5, container=False)
+                generate_bt = gr.Button("Generate", scale=1)
+
+        image = gr.Image(type="filepath")
+        with gr.Accordion("Advanced options", open=False):
+            seed = gr.Slider(randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1)
+
+        inputs = [prompt, seed]
+        generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+
+demo.queue(api_open=False)
+demo.launch(show_api=False)

scheduling_dmd.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import List, Tuple, Union, Optional
+
+import torch
+from diffusers import DDPMScheduler
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class DMDSchedulerOutput(BaseOutput):
+    pred_original_sample: Optional[torch.FloatTensor] = None
+
+
+class DMDScheduler(DDPMScheduler):
+    def set_timesteps(
+        self,
+        num_inference_steps: Optional[int] = None,
+        device: Union[str, torch.device] = None,
+        timesteps: Optional[List[int]] = None,
+    ):
+        self.timesteps = torch.tensor([self.config.num_train_timesteps-1]).long().to(device)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        generator=None,
+        return_dict: bool = True,
+    ) -> Union[DMDSchedulerOutput, Tuple]:
+        t = self.config.num_train_timesteps - 1
+
+        # 1. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[t]
+        beta_prod_t = 1 - alpha_prod_t
+
+        if self.config.prediction_type == "epsilon":
+            pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
+                " `v_prediction`  for the DDPMScheduler."
+            )
+
+        if not return_dict:
+            return (pred_original_sample,)
+
+        return DMDSchedulerOutput(pred_original_sample=pred_original_sample)