From 1d36e19996d7ceafc5283e54eda02c5ae9dd403b Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Fri, 7 Jul 2023 18:52:04 +0900
Subject: [PATCH] Fix PNDMScheduler for DirectML.

---
 modules/dml/hijack/__init__.py  |  1 +
 modules/dml/hijack/diffusers.py | 47 +++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 modules/dml/hijack/diffusers.py

diff --git a/modules/dml/hijack/__init__.py b/modules/dml/hijack/__init__.py
index d8cc0913a..7e3424775 100644
--- a/modules/dml/hijack/__init__.py
+++ b/modules/dml/hijack/__init__.py
@@ -3,3 +3,4 @@ import modules.dml.hijack.stablediffusion
 import modules.dml.hijack.torch
 import modules.dml.hijack.realesrgan_model
 import modules.dml.hijack.plms
+import modules.dml.hijack.diffusers
diff --git a/modules/dml/hijack/diffusers.py b/modules/dml/hijack/diffusers.py
new file mode 100644
index 000000000..42888fb8c
--- /dev/null
+++ b/modules/dml/hijack/diffusers.py
@@ -0,0 +1,47 @@
+import diffusers
+
+def _get_prev_sample(self, sample, timestep, prev_timestep, model_output):
+    # See formula (9) of PNDM paper https://arxiv.org/pdf/2202.09778.pdf
+    # this function computes x_(t−δ) using the formula of (9)
+    # Note that x_t needs to be added to both sides of the equation
+
+    # Notation (<variable name> -> <name in paper>
+    # alpha_prod_t -> α_t
+    # alpha_prod_t_prev -> α_(t−δ)
+    # beta_prod_t -> (1 - α_t)
+    # beta_prod_t_prev -> (1 - α_(t−δ))
+    # sample -> x_t
+    # model_output -> e_θ(x_t, t)
+    # prev_sample -> x_(t−δ)
+    print(sample) # DML Solution: PNDM Sampling does not work without this print. (because it depends on PLMS)
+    alpha_prod_t = self.alphas_cumprod[timestep]
+    alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+    beta_prod_t = 1 - alpha_prod_t
+    beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+    if self.config.prediction_type == "v_prediction":
+        model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
+    elif self.config.prediction_type != "epsilon":
+        raise ValueError(
+            f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `v_prediction`"
+        )
+
+    # corresponds to (α_(t−δ) - α_t) divided by
+    # denominator of x_t in formula (9) and plus 1
+    # Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) =
+    # sqrt(α_(t−δ)) / sqrt(α_t))
+    sample_coeff = (alpha_prod_t_prev / alpha_prod_t) ** (0.5)
+
+    # corresponds to denominator of e_θ(x_t, t) in formula (9)
+    model_output_denom_coeff = alpha_prod_t * beta_prod_t_prev ** (0.5) + (
+        alpha_prod_t * beta_prod_t * alpha_prod_t_prev
+    ) ** (0.5)
+
+    # full formula (9)
+    prev_sample = (
+        sample_coeff * sample - (alpha_prod_t_prev - alpha_prod_t) * model_output / model_output_denom_coeff
+    )
+
+    return prev_sample
+
+diffusers.PNDMScheduler._get_prev_sample = _get_prev_sample