Duplicate from sincostanx/momentum-diffusion

Co-authored-by: Worameth Chinchuthakun <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,15 @@
+---
+title: Momentum Diffusion
+emoji: 🐠
+colorFrom: blue
+colorTo: yellow
+sdk: gradio
+sdk_version: 3.36.1
+python_version: 3.9.17
+app_file: app.py
+pinned: false
+license: cc-by-4.0
+duplicated_from: sincostanx/momentum-diffusion
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,227 @@
+import gradio as gr
+import torch
+from pipeline import CustomPipeline, setup_scheduler
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+# from easydict import EasyDict as edict
+
+original_pipe = None
+original_config = None
+device = None
+
+
+# def run_dpm_demo(id, prompt, beta, num_inference_steps, guidance_scale, seed, enable_token_merging):
+def run_dpm_demo(prompt, beta, num_inference_steps, guidance_scale, seed):
+    global original_pipe, original_config
+    pipe = CustomPipeline(**original_pipe.components)
+
+    seed = int(seed)
+    num_inference_steps = int(num_inference_steps)
+
+    scheduler = "DPM-Solver++"
+    params = {
+        "prompt": prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "method": "dpm"
+    }
+
+    # without momentum (equivalent to DPM-Solver++)
+    pipe = setup_scheduler(pipe, scheduler, beta=1.0, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    ori_image = pipe(**params).images[0]
+
+    # with momentum
+    pipe = setup_scheduler(pipe, scheduler, beta=beta, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    image = pipe(**params).images[0]
+
+    ori_image.save("temp1.png")
+    image.save("temp2.png")
+
+    return [ori_image, image]
+
+# def run_plms_demo(id, prompt, order, beta, momentum_type, num_inference_steps, guidance_scale, seed, enable_token_merging):
+def run_plms_demo(prompt, order, beta, momentum_type, num_inference_steps, guidance_scale, seed):
+    global original_pipe, original_config
+    pipe = CustomPipeline(**original_pipe.components)
+
+    seed = int(seed)
+    num_inference_steps = int(num_inference_steps)
+
+    scheduler = "PLMS"
+    method = "hb" if momentum_type == "Polyak's heavy ball" else "nt"
+    params = {
+        "prompt": prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "method": method
+    }
+
+    # without momentum (equivalent to PLMS)
+    pipe = setup_scheduler(pipe, scheduler, momentum_type=momentum_type, order=order, beta=1.0, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    ori_image = pipe(**params).images[0]
+
+    # with momentum
+    pipe = setup_scheduler(pipe, scheduler, momentum_type=momentum_type, order=order, beta=beta, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    image = pipe(**params).images[0]
+
+    return [ori_image, image]
+
+# def run_ghvb_demo(id, prompt, order, beta, num_inference_steps, guidance_scale, seed, enable_token_merging):
+def run_ghvb_demo(prompt, order, beta, num_inference_steps, guidance_scale, seed):
+    global original_pipe, original_config
+    pipe = CustomPipeline(**original_pipe.components)
+
+    seed = int(seed)
+    num_inference_steps = int(num_inference_steps)
+
+    scheduler = "GHVB"
+    params = {
+        "prompt": prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "method": "ghvb"
+    }
+
+    # without momentum (equivalent to PLMS)
+    pipe = setup_scheduler(pipe, scheduler, order=order, beta=1.0, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    ori_image = pipe(**params).images[0]
+
+    # with momentum
+    pipe = setup_scheduler(pipe, scheduler, order=order, beta=beta, original_config=original_config)
+    params["generator"] = torch.Generator(device=device).manual_seed(seed)
+    image = pipe(**params).images[0]
+
+    return [ori_image, image]
+
+if __name__ == "__main__":
+
+    demo = gr.Blocks()
+
+    inputs = {}
+    outputs = {}
+    buttons = {}
+
+    list_models = [
+        "Linaqruf/anything-v3.0",
+        "runwayml/stable-diffusion-v1-5",
+        "dreamlike-art/dreamlike-photoreal-2.0",
+    ]
+    for model_id in list_models:
+        pipeline = StableDiffusionPipeline.from_pretrained(model_id)
+        del pipeline
+        print(f"Downloaded {model_id}")
+
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # Momentum-Diffusion Demo
+
+            A novel sampling method for diffusion models based on momentum to reduce artifacts
+            
+            """
+        )
+        id = gr.Dropdown(list_models, label="Model ID", value="Linaqruf/anything-v3.0", allow_custom_value=True)
+        enable_token_merging = gr.Checkbox(label="Enable Token Merging", value=False)
+        # output = gr.Textbox()
+        buttons["select_model"] = gr.Button("Select")
+
+        with gr.Tab("GHVB", visible=False) as tab3:
+            prompt3 = gr.Textbox(label="Prompt", value="a cozy cafe", visible=False)
+
+            with gr.Row(visible=False) as row31:
+                order = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="order")
+                beta = gr.Slider(minimum=0, maximum=1, value=0.4, step=0.05, label="beta")
+                num_inference_steps = gr.Number(label="Number of steps", value=12)
+                guidance_scale = gr.Number(label="Guidance scale (cfg)", value=10)
+                seed = gr.Number(label="Seed", value=42)
+                
+            with gr.Row(visible=False) as row32:
+                out1 = gr.Image(label="PLMS", interactive=False)
+                out2 = gr.Image(label="GHVB", interactive=False)
+                
+            inputs["GHVB"] = [prompt3, order, beta, num_inference_steps, guidance_scale, seed]
+            outputs["GHVB"] = [out1, out2]
+            buttons["GHVB"] = gr.Button("Sample", visible=False)
+
+        with gr.Tab("PLMS", visible=False) as tab2:
+            prompt2 = gr.Textbox(label="Prompt", value="1girl", visible=False)
+
+            with gr.Row(visible=False) as row21:
+                order = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="order")
+                beta = gr.Slider(minimum=0, maximum=1, value=0.7, step=0.05, label="beta")
+                momentum_type = gr.Dropdown(["Polyak's heavy ball", "Nesterov"], label="Momentum Type", value="Polyak's heavy ball")
+                num_inference_steps = gr.Number(label="Number of steps", value=10)
+                guidance_scale = gr.Number(label="Guidance scale (cfg)", value=10)
+                seed = gr.Number(label="Seed", value=42)
+                
+            with gr.Row(visible=False) as row22:
+                out1 = gr.Image(label="Without momentum", interactive=False)
+                out2 = gr.Image(label="With momentum", interactive=False)
+
+            inputs["PLMS"] = [prompt2, order, beta, momentum_type, num_inference_steps, guidance_scale, seed]
+            outputs["PLMS"] = [out1, out2]
+            buttons["PLMS"] = gr.Button("Sample", visible=False)
+
+        with gr.Tab("DPM-Solver++", visible=False) as tab1:
+            prompt1 = gr.Textbox(label="Prompt", value="1girl", visible=False)
+
+            with gr.Row(visible=False) as row11:
+                beta = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.05, label="beta")
+                num_inference_steps = gr.Number(label="Number of steps", value=15)
+                guidance_scale = gr.Number(label="Guidance scale (cfg)", value=20)
+                seed = gr.Number(label="Seed", value=0)
+                
+            with gr.Row(visible=False) as row12:
+                out1 = gr.Image(label="Without momentum", interactive=False)
+                out2 = gr.Image(label="With momentum", interactive=False)
+
+            inputs["DPM-Solver++"] = [prompt1, beta, num_inference_steps, guidance_scale, seed]
+            outputs["DPM-Solver++"] = [out1, out2]
+            buttons["DPM-Solver++"] = gr.Button("Sample", visible=False)
+
+        def prepare_model(id, enable_token_merging):
+            global original_pipe, original_config, device
+            
+            if original_pipe is not None:
+                del original_pipe
+
+            original_pipe = CustomPipeline.from_pretrained(id)
+            device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+            original_pipe = original_pipe.to(device)
+
+            if enable_token_merging:
+                import tomesd
+                tomesd.apply_patch(original_pipe, ratio=0.5)
+                print("Enabled Token merging.")
+
+            original_config = original_pipe.scheduler.config
+            print(type(original_pipe))
+            print(original_config)
+
+            return {
+                row11: gr.update(visible=True),
+                row12: gr.update(visible=True),
+                row21: gr.update(visible=True),
+                row22: gr.update(visible=True),
+                row31: gr.update(visible=True),
+                row32: gr.update(visible=True),
+                prompt1: gr.update(visible=True),
+                prompt2: gr.update(visible=True),
+                prompt3: gr.update(visible=True),
+                buttons["DPM-Solver++"]: gr.update(visible=True),
+                buttons["PLMS"]: gr.update(visible=True),
+                buttons["GHVB"]: gr.update(visible=True),
+            }
+
+        all_outputs = [row11, row12, row21, row22, row31, row32, prompt1, prompt2, prompt3, buttons["DPM-Solver++"], buttons["PLMS"], buttons["GHVB"]]
+        buttons["select_model"].click(prepare_model, inputs=[id, enable_token_merging], outputs=all_outputs)
+        buttons["DPM-Solver++"].click(run_dpm_demo, inputs=inputs["DPM-Solver++"], outputs=outputs["DPM-Solver++"])
+        buttons["PLMS"].click(run_plms_demo, inputs=inputs["PLMS"], outputs=outputs["PLMS"])
+        buttons["GHVB"].click(run_ghvb_demo, inputs=inputs["GHVB"], outputs=outputs["GHVB"])
+
+    demo.launch()

momentum_scheduler.py

@@ -0,0 +1,385 @@
+import torch
+from diffusers import DPMSolverMultistepScheduler, UniPCMultistepScheduler
+from typing import List
+
+def AdamBmixer(order, ets, b=1):
+ 
+    cur_order = min(order, len(ets))
+    if cur_order == 1:
+        prime = b * ets[-1]  
+    elif cur_order == 2:
+        prime = ((2+b) * ets[-1] - (2-b)*ets[-2]) / 2
+    elif cur_order == 3:
+        prime = ((18+5*b) * ets[-1] - (24-8*b) * ets[-2] + (6-1*b) * ets[-3]) / 12
+    elif cur_order == 4:
+        prime = ((46+9*b) * ets[-1] - (78-19*b) * ets[-2] + (42-5*b) * ets[-3] - (10-b) * ets[-4]) / 24
+    elif cur_order == 5:
+        prime = ((1650+251*b) * ets[-1] - (3420-646*b) * ets[-2]
+                     + (2880-264*b) * ets[-3] - (1380-106*b) * ets[-4]
+                     + (270-19*b)* ets[-5]) / 720
+    else:
+        raise NotImplementedError
+    
+    prime = prime/b
+    return prime
+
+class PLMSWithHBScheduler():
+    """
+    PLMS with Polyak's Heavy Ball Momentum (HB) for diffusion ODEs.
+    We implement it as a wrapper for schedulers in diffusers (https://github.com/huggingface/diffusers)
+    
+    When order is an integer, this method is equivalent to PLMS without momentum.
+    """
+    def __init__(self, scheduler, order):
+        self.scheduler = scheduler
+        self.ets = []
+        self.update_order(order)
+        self.mixer = AdamBmixer
+        
+    def update_order(self, order):
+        self.order = order // 1  + 1 if order%1 > 0 else order // 1 
+        self.beta = order % 1 if order%1 > 0 else 1
+        self.vel = None
+ 
+    def clear(self):
+        self.ets = []
+        self.vel = None
+
+    def update_ets(self, val):
+        self.ets.append(val)
+        if len(self.ets) > self.order:
+            self.ets.pop(0)
+
+    def _step_with_momentum(self, grads):
+        self.update_ets(grads)
+        prime = self.mixer(self.order, self.ets, 1.0)
+        self.vel = (1 - self.beta) * self.vel + self.beta * prime
+        return self.vel
+
+    def step(
+        self,
+        grads: torch.FloatTensor,
+        timestep: int,
+        latents: torch.FloatTensor,
+        output_mode: str = "scale",
+    ):
+        if self.vel is None: self.vel = grads
+ 
+        if hasattr(self.scheduler, 'sigmas'):
+            step_index = (self.scheduler.timesteps == timestep).nonzero().item()
+            sigma = self.scheduler.sigmas[step_index]
+            sigma_next = self.scheduler.sigmas[step_index + 1]
+            del_g = sigma_next - sigma
+ 
+            update_val = self._step_with_momentum(grads)
+            return latents + del_g * update_val
+
+        elif isinstance(self.scheduler, DPMSolverMultistepScheduler):
+            step_index = (self.scheduler.timesteps == timestep).nonzero().item()
+            current_timestep = self.scheduler.timesteps[step_index]
+            prev_timestep = 0 if step_index == len(self.scheduler.timesteps) - 1 else self.scheduler.timesteps[step_index + 1]
+
+            alpha_prod_t = self.scheduler.alphas_cumprod[current_timestep]
+            alpha_bar_prev = self.scheduler.alphas_cumprod[prev_timestep]
+
+            s0 = torch.sqrt(alpha_prod_t)
+            s_1 = torch.sqrt(alpha_bar_prev)
+            g0 = torch.sqrt(1-alpha_prod_t)/s0
+            g_1 = torch.sqrt(1-alpha_bar_prev)/s_1
+            del_g = g_1 - g0
+ 
+            update_val = self._step_with_momentum(grads)
+            if output_mode in ["scale"]:
+                return (latents/s0  + del_g * update_val) * s_1
+            elif output_mode in ["back"]:
+                return latents + del_g * update_val * s_1
+            elif output_mode in ["front"]:
+                return latents + del_g * update_val * s0
+            else:
+                return latents + del_g * update_val
+        else:
+            raise NotImplementedError
+
+class GHVBScheduler(PLMSWithHBScheduler):
+    """
+    Generalizing Polyak's Heavy Bal (GHVB) for diffusion ODEs.
+    We implement it as a wrapper for schedulers in diffusers (https://github.com/huggingface/diffusers)
+    
+    When order is an integer, this method is equivalent to PLMS without momentum.
+    """
+    def _step_with_momentum(self, grads):
+        self.vel = (1 - self.beta) * self.vel + self.beta * grads
+        self.update_ets(self.vel)
+        prime = self.mixer(self.order, self.ets, self.beta)
+        return prime
+
+class PLMSWithNTScheduler(PLMSWithHBScheduler):
+    """
+    PLMS with Nesterov Momentum (NT) for diffusion ODEs.
+    We implement it as a wrapper for schedulers in diffusers (https://github.com/huggingface/diffusers)
+    
+    When order is an integer, this method is equivalent to PLMS without momentum.
+    """
+    def _step_with_momentum(self, grads):
+        self.update_ets(grads)
+        prime = self.mixer(self.order, self.ets, 1.0) # update v^{(2)}
+        self.vel = (1 - self.beta) * self.vel + self.beta * prime # update v^{(1)}
+        update_val = (1 - self.beta) * self.vel + self.beta * prime # update x
+        return update_val
+
+class MomentumDPMSolverMultistepScheduler(DPMSolverMultistepScheduler):
+    """
+    DPM-Solver++2M with HB momentum.
+    Currently support only algorithm_type = "dpmsolver++" and solver_type = "midpoint"
+
+    When beta = 1.0, this method is equivalent to DPM-Solver++2M without momentum.
+    """
+    def initialize_momentum(self, beta):
+        self.vel = None
+        self.beta = beta
+
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.FloatTensor],
+        timestep_list: List[int],
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        
+        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
+        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
+        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                diff = (D0 + 0.5 * D1)
+
+                if self.vel is None:
+                    self.vel = diff
+                else:
+                    self.vel = (1-self.beta)*self.vel + self.beta * diff
+                
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * self.vel
+                )
+            elif self.config.solver_type == "heun":
+                raise NotImplementedError(
+                    "{self.config.algorithm_type} with {self.config.solver_type} is currently not supported."
+                )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                raise NotImplementedError(
+                    "{self.config.algorithm_type} with {self.config.solver_type} is currently not supported."
+                )
+            elif self.config.solver_type == "heun":
+                raise NotImplementedError(
+                    "{self.config.algorithm_type} with {self.config.solver_type} is currently not supported."
+                )
+        return x_t
+
+class MomentumUniPCMultistepScheduler(UniPCMultistepScheduler):
+    """
+    UniPC with HB momentum.
+    Currently support only self.predict_x0 = True
+
+    When beta = 1.0, this method is equivalent to UniPC without momentum.
+    """
+    def initialize_momentum(self, beta):
+        self.vel_p = None
+        self.vel_c = None
+        self.beta = beta
+ 
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.FloatTensor,
+        prev_timestep: int,
+        sample: torch.FloatTensor,
+        order: int,
+    ) -> torch.FloatTensor:
+ 
+        timestep_list = self.timestep_list
+        model_output_list = self.model_outputs
+ 
+        s0, t = self.timestep_list[-1], prev_timestep
+        m0 = model_output_list[-1]
+        x = sample
+ 
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+ 
+        lambda_t, lambda_s0 = self.lambda_t[t], self.lambda_t[s0]
+        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
+        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+ 
+        h = lambda_t - lambda_s0
+        device = sample.device
+ 
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = timestep_list[-(i + 1)]
+            mi = model_output_list[-(i + 1)]
+            lambda_si = self.lambda_t[si]
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+ 
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+ 
+        R = []
+        b = []
+ 
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+ 
+        factorial_i = 1
+ 
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+ 
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+ 
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+ 
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
+        else:
+            D1s = None
+ 
+        if self.predict_x0:
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkchw->bchw", rhos_p, D1s)
+            else:
+                pred_res = 0
+ 
+            val = ( h_phi_1 * m0 + B_h * pred_res ) /sigma_t /h_phi_1
+            if self.vel_p is None:
+                self.vel_p = val
+            else:
+                self.vel_p = (1-self.beta)*self.vel_p + self.beta * val
+            self.vel_p = val
+ 
+            x_t = sigma_t  * (x/ sigma_s0 - alpha_t * self.vel_p * h_phi_1) 
+        else:
+            raise NotImplementedError
+ 
+        x_t = x_t.to(x.dtype)
+        return x_t
+ 
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.FloatTensor,
+        this_timestep: int,
+        last_sample: torch.FloatTensor,
+        this_sample: torch.FloatTensor,
+        order: int,
+    ) -> torch.FloatTensor:
+ 
+        timestep_list = self.timestep_list
+        model_output_list = self.model_outputs
+ 
+        s0, t = timestep_list[-1], this_timestep
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+ 
+        lambda_t, lambda_s0 = self.lambda_t[t], self.lambda_t[s0]
+        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
+        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
+ 
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+ 
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = timestep_list[-(i + 1)]
+            mi = model_output_list[-(i + 1)]
+            lambda_si = self.lambda_t[si]
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+ 
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+ 
+        R = []
+        b = []
+ 
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+ 
+        factorial_i = 1
+ 
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+ 
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+ 
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+ 
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+ 
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b)
+ 
+        if self.predict_x0:
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkchw->bchw", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+ 
+            val = (h_phi_1 * m0 + B_h * (corr_res + rhos_c[-1] * D1_t))/sigma_t/h_phi_1
+            if self.vel_c is None:
+                self.vel_c = val
+            else:
+                self.vel_c = (1-self.beta)*self.vel_c + self.beta * val
+
+            x_t = sigma_t  * (x/ sigma_s0 - alpha_t * self.vel_c * h_phi_1) 
+        else:
+            raise NotImplementedError
+        
+        x_t = x_t.to(x.dtype)
+        return x_t

pipeline.py

@@ -0,0 +1,236 @@
+import torch
+import math
+import numpy as np
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, UniPCMultistepScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from typing import Union, Optional, List, Callable, Dict, Any, Tuple
+from momentum_scheduler import (
+    GHVBScheduler,
+    PLMSWithHBScheduler,
+    PLMSWithNTScheduler,
+    MomentumDPMSolverMultistepScheduler,
+    MomentumUniPCMultistepScheduler,
+)
+
+available_solvers = {
+    "GHVB": GHVBScheduler,
+    "PLMS_HB": PLMSWithHBScheduler,
+    "PLMS_NT": PLMSWithNTScheduler,
+    "DPM-Solver++": MomentumDPMSolverMultistepScheduler,
+    "UniPC": MomentumUniPCMultistepScheduler,
+}
+
+def get_momentum_number(order, beta):
+    out = order if beta == 1.0 else order - (1 - beta)
+    return out
+
+def setup_scheduler(pipe, scheduler, momentum_type="Polyak's heavy ball", order=4.0, beta=1.0, original_config=None):
+    assert original_config is not None
+    
+    if scheduler in ["DPM-Solver++", "UniPC"]:
+        if momentum_type in ["Nesterov"]:
+            raise NotImplementedError(f"{scheduler} w/ Nesterov is not implemented.")
+
+        pipe.scheduler = available_solvers[scheduler].from_config(original_config)
+        pipe.scheduler.initialize_momentum(beta=beta)
+
+    elif scheduler in ["PLMS"]:
+        momentum_number = get_momentum_number(order, beta)
+        method = "PLMS_HB" if momentum_type == "Polyak's heavy ball" else "PLMS_NT"
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
+        pipe.init_scheduler(method=method, order=momentum_number)
+        pipe.clear_scheduler()
+
+    elif scheduler in ["GHVB"]:
+        momentum_number = get_momentum_number(order, beta)
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(original_config)
+        pipe.init_scheduler(method="GHVB", order=momentum_number)
+        pipe.clear_scheduler()
+
+    return pipe
+
+class CustomPipeline(StableDiffusionPipeline):
+    def clear_scheduler(self):
+        self.scheduler_uncond.clear()
+        self.scheduler_text.clear()
+    
+    def init_scheduler(self, method, order):
+        # equivalent to not applied numerical operator splitting since orders are the same
+        self.scheduler_uncond = available_solvers[method](self.scheduler, order)
+        self.scheduler_text = available_solvers[method](self.scheduler, order)
+
+    def get_noise(self, latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance):
+        # expand the latents if we are doing classifier free guidance
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+        latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+        # predict the noise residual
+        noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=prompt_embeds).sample
+
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            grads_a = guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        return noise_pred_uncond, grads_a
+
+    def denoising_step(
+            self,
+            latents,
+            prompt_embeds,
+            guidance_scale,
+            t,
+            do_classifier_free_guidance,
+            method,
+            extra_step_kwargs,
+    ):
+        noise_pred_uncond, grads_a = self.get_noise(
+            latents, prompt_embeds, guidance_scale, t, do_classifier_free_guidance
+        )
+        if method in ["dpm", "unipc"]:
+            latents = self.scheduler.step(noise_pred_uncond + grads_a, t, latents, **extra_step_kwargs).prev_sample
+        
+        elif method in ["hb", "ghvb", "nt"]:
+            latents = self.scheduler_uncond.step(noise_pred_uncond, t, latents, output_mode="scale")
+            latents = self.scheduler_text.step(grads_a, t, latents, output_mode='back')
+        else:
+            raise NotImplementedError
+
+        return latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        method="ghvb",
+    ):
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+        # print(timesteps)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latents = self.denoising_step(
+                    latents,
+                    prompt_embeds,
+                    guidance_scale,
+                    t,
+                    do_classifier_free_guidance,
+                    method,
+                    extra_step_kwargs,
+                )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        elif output_type == "pil":
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+            has_nsfw_concept = False
+
+            # 10. Convert to PIL
+            image = self.numpy_to_pil(image)
+        else:
+            # 8. Post-processing
+            image = self.decode_latents(latents)
+
+            # 9. Run safety checker
+            # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+            has_nsfw_concept = False
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+    def generate(self, params):
+        params["output_type"] = "latent"
+        ori_latents = self.__call__(**params)["images"]
+
+        with torch.no_grad():
+            latents = torch.clone(ori_latents)
+            image = self.decode_latents(latents)
+            image = self.numpy_to_pil(image)[0]
+
+        return image, ori_latents

requirements.txt

@@ -0,0 +1,96 @@
+accelerate==0.21.0
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==5.0.1
+annotated-types==0.5.0
+anyio==3.7.1
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2023.5.7
+charset-normalizer==3.2.0
+click==8.1.5
+cmake==3.26.4
+contourpy==1.1.0
+cycler==0.11.0
+diffusers==0.15.0
+exceptiongroup==1.1.2
+fastapi==0.100.0
+ffmpy==0.3.0
+filelock==3.12.2
+fonttools==4.41.0
+frozenlist==1.4.0
+fsspec==2023.6.0
+gradio==3.36.1
+gradio_client==0.2.9
+h11==0.14.0
+httpcore==0.17.3
+httpx==0.24.1
+huggingface-hub==0.16.4
+idna==3.4
+importlib-metadata==6.8.0
+importlib-resources==6.0.0
+Jinja2==3.1.2
+jsonschema==4.18.3
+jsonschema-specifications==2023.6.1
+kiwisolver==1.4.4
+linkify-it-py==2.0.2
+lit==16.0.6
+markdown-it-py==2.2.0
+MarkupSafe==2.1.3
+matplotlib==3.7.2
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.1
+numpy==1.25.1
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+orjson==3.9.2
+packaging==23.1
+pandas==2.0.3
+Pillow==10.0.0
+psutil==5.9.5
+pydantic==2.0.2
+pydantic_core==2.1.2
+pydub==0.25.1
+Pygments==2.15.1
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3
+PyYAML==6.0
+referencing==0.29.1
+regex==2023.6.3
+requests==2.31.0
+rpds-py==0.8.10
+semantic-version==2.10.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+tokenizers==0.13.3
+tomesd==0.1.3
+toolz==0.12.0
+torch==2.0.1
+tqdm==4.65.0
+transformers==4.28.1
+triton==2.0.0
+typing_extensions==4.7.1
+tzdata==2023.3
+uc-micro-py==1.0.2
+urllib3==2.0.3
+uvicorn==0.22.0
+websockets==11.0.3
+yarl==1.9.2
+zipp==3.16.1