From 807cdd808c7af4b2619d14d4ed070d90b023dc69 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 15:16:36 -0700
Subject: [PATCH 1/8] Make enable and disable idempotent

---
 .gitignore                              |   4 +-
 DeepCache.egg-info/PKG-INFO             | 281 ++++++++++++++++++++++++
 DeepCache.egg-info/SOURCES.txt          |  29 +++
 DeepCache.egg-info/dependency_links.txt |   1 +
 DeepCache.egg-info/requires.txt         |   3 +
 DeepCache.egg-info/top_level.txt        |   1 +
 DeepCache/extension/deepcache.py        |  24 +-
 7 files changed, 335 insertions(+), 8 deletions(-)
 create mode 100644 DeepCache.egg-info/PKG-INFO
 create mode 100644 DeepCache.egg-info/SOURCES.txt
 create mode 100644 DeepCache.egg-info/dependency_links.txt
 create mode 100644 DeepCache.egg-info/requires.txt
 create mode 100644 DeepCache.egg-info/top_level.txt

diff --git a/.gitignore b/.gitignore
index d0d6b8e..3da4190 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,4 @@
 __pycache__
-output.png
\ No newline at end of file
+output.png
+
+.idea
\ No newline at end of file
diff --git a/DeepCache.egg-info/PKG-INFO b/DeepCache.egg-info/PKG-INFO
new file mode 100644
index 0000000..33b66e9
--- /dev/null
+++ b/DeepCache.egg-info/PKG-INFO
@@ -0,0 +1,281 @@
+Metadata-Version: 2.1
+Name: DeepCache
+Version: 0.1.0
+Summary: DeepCache: Accelerating Diffusion Models for Free
+Home-page: https://github.com/horseee/DeepCache
+Author: Xinyin Ma
+Author-email: maxinyin@u.nus.edu
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: diffusers
+Requires-Dist: transformers
+
+# DeepCache: Accelerating Diffusion Models for Free
+<div align="center">
+  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/example_compress.gif" width="100%" ></img>
+  <br>
+  <em>
+      (Results on Stable Diffusion v1.5. Left: 50 PLMS steps. Right: 2.3x acceleration upon 50 PLMS steps) 
+  </em>
+</div>
+<br>
+
+> **DeepCache: Accelerating Diffusion Models for Free**   
+> [Xinyin Ma](https://horseee.github.io/), [Gongfan Fang](https://fangggf.github.io/), [Xinchao Wang](https://sites.google.com/site/sitexinchaowang/)   
+> [Learning and Vision Lab](http://lv-nus.org/), National University of Singapore  
+> 🥯[[Arxiv]](https://arxiv.org/abs/2312.00858)🎄[[Project Page]](https://horseee.github.io/Diffusion_DeepCache/)
+
+### Why DeepCache  
+* 🚀 Training-free and almost lossless
+* 🚀 Support Stable Diffusion, Stable Diffusion XL, Stable Video Diffusion, Stable Diffusion Pipeline / XL Pipeline for Inpainting, Stable Diffusion Img2Img Pipeline, DDPM
+* 🚀 Compatible with sampling algorithms like DDIM and PLMS
+
+### Updates
+* **June 27, 2024**: 🔥Our new work [AsyncDiff](https://github.com/czg1225/AsyncDiff) enables parallel inference of diffusion models on multiple GPUs. Check our [paper](https://arxiv.org/abs/2406.06911) and [code](https://github.com/czg1225/AsyncDiff)!
+* **June 5, 2024**: 🔥Our new work, [Learning-to-Cache](https://github.com/horseee/learning-to-cache/tree/main), an improved version of DeepCache on DiT. Code and checkpoints are released.
+* **January 5, 2024**: 💥A doc page for DeepCache has been added in [Diffusers](https://github.com/huggingface/diffusers)! Check [here](https://huggingface.co/docs/diffusers/main/en/optimization/deepcache) for more information. Many thanks to the Diffusers team!
+* **December 26, 2023**: 🔥Update a plug-and-play implementation of DeepCache, no longer requiring any modifications of the diffuser's code! Check [here](https://github.com/horseee/DeepCache#usage) for the detailed usage! Big thanks to @[yuanshi9815](https://github.com/Yuanshi9815) for contributing the code!
+* **December 25, 2023**: A demo is available via Colab [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1COM9tfGvHSJ8Zn4tScZCGgt8nxAEOlAz?usp=sharing)
+* **December 21, 2023**: Release the code for **Stable Video Diffusion** and **Text2Video-Zero**. In the figure below, the upper row shows the original videos generated by SVD-XT, and the lower row is accelerated by DeepCache. For Text2Video-Zero, the results can be found [here](https://github.com/horseee/DeepCache/blob/master/assets/Text2Video-Zero.gif)
+<div align="center">
+  <img src="assets/svd.gif" width="90%" ></img>
+  <br>
+  <em>
+      (1.7x acceleration of SVD-XT) 
+  </em>
+</div>
+
+
+* **December 20, 2023**: Release the code for **DDPM**. See [here](https://github.com/horseee/DeepCache/tree/master/DeepCache#experiment-code-for-ddpm) for the experimental code and instructions.
+
+* **December 6, 2023**: Release the code for **Stable Diffusion XL**. The results of the `stabilityai/stable-diffusion-xl-base-1.0` are shown in the below figure, with the same prompts from the first figure.
+<div align="center">
+  <img src="assets/sdxl.png" width="90%" ></img>
+  <br>
+  <em>
+      (2.6x acceleration of Stable Diffusion XL) 
+  </em>
+</div>
+
+
+### Introduction
+We introduce **DeepCache**, a novel **training-free and almost lossless** paradigm that accelerates diffusion models from the perspective of model architecture. Utilizing the property of the U-Net, we reuse the high-level features while updating the low-level features in a very cheap way. DeepCache accelerates Stable Diffusion v1.5 by 2.3x with only a 0.05 decline in CLIP Score, and LDM-4-G(ImageNet) by 4.1x with a 0.22 decrease in FID.
+
+<div align="center">
+<img width="50%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/9ce3930c-c84c-4af8-8c6a-b6803a5a7b1d">
+</div>
+
+## Quick Start
+
+### Install
+```bash
+pip install DeepCache
+```
+
+### Usage
+
+```python
+import torch
+
+# Loading the original pipeline
+from diffusers import StableDiffusionPipeline
+pipe = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16).to("cuda:0")
+
+# Import the DeepCacheSDHelper
+from DeepCache import DeepCacheSDHelper
+helper = DeepCacheSDHelper(pipe=pipe)
+helper.set_params(
+    cache_interval=3,
+    cache_branch_id=0,
+)
+helper.enable()
+
+# Generate Image
+deepcache_image = pipe(
+        prompt,
+        output_type='pt'
+).images[0]
+helper.disable()
+```
+We here take the Stable Diffusion pipeline as an example. You can replace pipe with any variants of the Stable Diffusion pipeline, including choices like SDXL, SVD, and more. You can find examples in the [script](https://github.com/horseee/DeepCache#a-general-script-for-sd). The argument `cache_branch_id` specifies the selected skip branch. For the skip branches that are deeper, the model will engage them only during the caching steps, and exclude them during the retrieval steps. The argument `cache_interval` represents the interval for updating the cache.
+
+### A general script for SD
+```bash
+python main.py --model_type sdxl #Support [sdxl, sd1.5, sd2.1, svd, sd-inpaint, sdxl-inpaint, sd-img2img]
+```
+
+## Experimental code for DeepCache
+The above implementation does not require changes to the `forward` or `__call__` functions in the Diffusers pipeline, and is, therefore, more general. The following section is the experimental code that can be used to reproduce the results in the paper. It is implemented one by one for different model structures and pipelines, and thus, may not work properly due to the update of diffusers.
+
+### Setup
+```
+pip install diffusers==0.24.0 transformers
+```
+
+### Stable Diffusion XL
+```bash
+python stable_diffusion_xl.py --model stabilityai/stable-diffusion-xl-base-1.0
+```
+
+<details>
+<summary>Output:</summary>
+
+```bash
+Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  6.62it/s]
+2023-12-06 01:44:28,578 - INFO - Running baseline...
+100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:17<00:00,  2.93it/s]
+2023-12-06 01:44:46,095 - INFO - Baseline: 17.52 seconds
+Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.06it/s]
+2023-12-06 01:45:02,865 - INFO - Running DeepCache...
+100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:06<00:00,  8.01it/s]
+2023-12-06 01:45:09,573 - INFO - DeepCache: 6.71 seconds
+2023-12-06 01:45:10,678 - INFO - Saved to output.png. Done!
+```
+
+</details>
+
+You can add `--refine` at the end of the command to activate the refiner model for SDXL.
+
+### Stable Diffusion v1.5
+```bash
+python stable_diffusion.py --model runwayml/stable-diffusion-v1-5
+```
+
+<details>
+<summary>Output:</summary>
+
+```bash
+2023-12-03 16:18:13,636 - INFO - Loaded safety_checker as StableDiffusionSafetyChecker from `safety_checker` subfolder of runwayml/stable-diffusion-v1-5.
+2023-12-03 16:18:13,699 - INFO - Loaded vae as AutoencoderKL from `vae` subfolder of runwayml/stable-diffusion-v1-5.
+Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  5.88it/s]
+2023-12-03 16:18:22,837 - INFO - Running baseline...
+100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:03<00:00, 15.33it/s]
+2023-12-03 16:18:26,174 - INFO - Baseline: 3.34 seconds
+2023-12-03 16:18:26,174 - INFO - Running DeepCache...
+100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:01<00:00, 34.06it/s]
+2023-12-03 16:18:27,718 - INFO - DeepCache: 1.54 seconds
+2023-12-03 16:18:27,935 - INFO - Saved to output.png. Done!
+```
+
+</details>
+
+### Stable Diffusion v2.1
+
+```bash
+python stable_diffusion.py --model stabilityai/stable-diffusion-2-1
+```
+
+<details>
+  <summary>Output:</summary>
+  
+```bash
+2023-12-03 16:21:17,858 - INFO - Loaded feature_extractor as CLIPImageProcessor from `feature_extractor` subfolder of stabilityai/stable-diffusion-2-1.
+2023-12-03 16:21:17,864 - INFO - Loaded scheduler as DDIMScheduler from `scheduler` subfolder of stabilityai/stable-diffusion-2-1.
+Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00,  5.35it/s]
+2023-12-03 16:21:49,770 - INFO - Running baseline...
+100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:14<00:00,  3.42it/s]
+2023-12-03 16:22:04,551 - INFO - Baseline: 14.78 seconds
+2023-12-03 16:22:04,551 - INFO - Running DeepCache...
+100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:08<00:00,  6.10it/s]
+2023-12-03 16:22:12,911 - INFO - DeepCache: 8.36 seconds
+2023-12-03 16:22:13,417 - INFO - Saved to output.png. Done!
+```
+
+</details>
+
+Currently, our code supports the models that can be loaded by [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/v0.24.0/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline). You can specify the model name by the argument `--model`, which by default, is `runwayml/stable-diffusion-v1-5`. 
+
+### Stable Video Diffusion 
+```bash
+python stable_video_diffusion.py
+```
+
+<details>
+  <summary>Output:</summary>
+  
+```bash
+Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00,  8.36it/s]
+2023-12-21 04:56:47,329 - INFO - Running baseline...
+100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [01:27<00:00,  3.49s/it]
+2023-12-21 04:58:26,121 - INFO - Origin: 98.66 seconds
+Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 10.59it/s]
+2023-12-21 04:58:27,202 - INFO - Running DeepCache...
+100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [00:49<00:00,  1.96s/it]
+2023-12-21 04:59:26,607 - INFO - DeepCache: 59.31 seconds
+```
+
+</details>
+
+### DDPM and LDM
+Please check [here](https://github.com/horseee/DeepCache/tree/master/experiments) for the experimental code of DDPM and LDM.
+
+
+## Visualization
+
+Images in the upper line are the baselines and the images in the lower line are accelerated by DeepCache. 
+
+### Stable Diffusion v1.5 (2.15x Acceleration)
+<div align="center">
+  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/sd_examples_high_res.png" width="100%">
+</div>
+
+### LDM-4-G for ImageNet (6.96x Acceleration)
+<div align="center">
+  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/imagenet.png" width="100%">
+</div>
+
+### DDPM for LSUN Church & Bedroom (1.48x Acceleration)
+<div align="center">
+  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/bedroom.png" width="100%">
+  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/church.png" width="100%">
+</div>
+
+## Quantitative Results
+
+### ImageNet
+<div align="center">
+<img width="80%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/151d7639-2501-45cf-8de5-2af3bb5a354b">
+</div>
+
+
+### Stable Diffusion v1.5
+
+<div align="center">
+<img width="80%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/e9bd7a8e-07c8-4296-95a2-12d008995807">
+</div>
+
+More results can be found in our [paper](https://arxiv.org/abs/2312.00858)
+
+## Other Implementations and Plugins
+We sincerely thank the authors listed below who implemented DeepCache in plugins or other contexts. 
+
+* Diffusers: https://huggingface.co/docs/diffusers/main/en/optimization/deepcache
+* OneDiff Integration: https://github.com/siliconflow/onediff?tab=readme-ov-file#easy-to-use by @[SiliconFlow](https://github.com/siliconflow). OneDiff also has implementations for DeepCache on SVD, check [this](https://github.com/siliconflow/onediff/pull/438) for details.
+* Comfyui: https://gist.github.com/laksjdjf/435c512bc19636e9c9af4ee7bea9eb86 by @[laksjdjf](https://github.com/laksjdjf)
+* Colab & Gradio: https://github.com/camenduru/DeepCache-colab by @[camenduru](https://github.com/camenduru/DeepCache-colab)
+* WebUI: https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14210 by @[aria1th](https://github.com/aria1th)
+* SD.Next: https://github.com/vladmandic/automatic by @[vladmandic](https://github.com/vladmandic)
+* MLX_deepcache for running SD on M1 Mac: https://github.com/Barry0121/mlx_deepcache by @[Barry Xue](https://github.com/Barry0121)
+
+We warmly welcome contributions from everyone. Please feel free to reach out to us.
+
+
+## Bibtex
+```
+@inproceedings{ma2023deepcache,
+  title={DeepCache: Accelerating Diffusion Models for Free},
+  author={Ma, Xinyin and Fang, Gongfan and Wang, Xinchao},
+  booktitle={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  year={2024}
+}
+```
+
+
+
+
+
diff --git a/DeepCache.egg-info/SOURCES.txt b/DeepCache.egg-info/SOURCES.txt
new file mode 100644
index 0000000..6fb784b
--- /dev/null
+++ b/DeepCache.egg-info/SOURCES.txt
@@ -0,0 +1,29 @@
+LICENSE
+README.md
+setup.py
+DeepCache/__init__.py
+DeepCache/flops.py
+DeepCache.egg-info/PKG-INFO
+DeepCache.egg-info/SOURCES.txt
+DeepCache.egg-info/dependency_links.txt
+DeepCache.egg-info/requires.txt
+DeepCache.egg-info/top_level.txt
+DeepCache/extension/__init__.py
+DeepCache/extension/deepcache.py
+DeepCache/sd/__init__.py
+DeepCache/sd/pipeline_stable_diffusion.py
+DeepCache/sd/pipeline_text_to_video_zero.py
+DeepCache/sd/pipeline_utils.py
+DeepCache/sd/unet_2d_blocks.py
+DeepCache/sd/unet_2d_condition.py
+DeepCache/sdxl/__init__.py
+DeepCache/sdxl/pipeline_stable_diffusion_xl.py
+DeepCache/sdxl/pipeline_stable_diffusion_xl_img2img.py
+DeepCache/sdxl/pipeline_utils.py
+DeepCache/sdxl/unet_2d_blocks.py
+DeepCache/sdxl/unet_2d_condition.py
+DeepCache/svd/__init__.py
+DeepCache/svd/pipeline_stable_video_diffusion.py
+DeepCache/svd/pipeline_utils.py
+DeepCache/svd/unet_3d_blocks.py
+DeepCache/svd/unet_spatio_temporal_condition.py
\ No newline at end of file
diff --git a/DeepCache.egg-info/dependency_links.txt b/DeepCache.egg-info/dependency_links.txt
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/DeepCache.egg-info/dependency_links.txt
@@ -0,0 +1 @@
+
diff --git a/DeepCache.egg-info/requires.txt b/DeepCache.egg-info/requires.txt
new file mode 100644
index 0000000..de80785
--- /dev/null
+++ b/DeepCache.egg-info/requires.txt
@@ -0,0 +1,3 @@
+torch
+diffusers
+transformers
diff --git a/DeepCache.egg-info/top_level.txt b/DeepCache.egg-info/top_level.txt
new file mode 100644
index 0000000..924077e
--- /dev/null
+++ b/DeepCache.egg-info/top_level.txt
@@ -0,0 +1 @@
+DeepCache
diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index e96ddce..bf6ad43 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -1,15 +1,25 @@
 class DeepCacheSDHelper(object):
     def __init__(self, pipe=None):
-        if pipe is not None: self.pipe = pipe
+        if pipe is not None:
+            self.pipe = pipe
+        self._enabled = False
 
-    def enable(self, pipe=None):
-        assert self.pipe is not None
-        self.reset_states()
-        self.wrap_modules()
+    def enable(self):
+        if not self._enabled:
+            self.reset_states()
+            self.wrap_modules()
+            self._enabled = True
+            print("Enabling Deepcache")
+        else:
+            print("DeepCache is already enabled.")
 
     def disable(self):
-        self.unwrap_modules()
-        self.reset_states()
+        if self._enabled:
+            self.unwrap_modules()
+            self.reset_states()
+            print("Disabling Deepcache")
+        else:
+            print("DeepCache is already disabled.")
     
     def set_params(self,cache_interval=1, cache_branch_id=0, skip_mode='uniform'):
         cache_layer_id = cache_branch_id % 3

From a4895bc28ac046046afce2930282620b68bc4a38 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 15:17:15 -0700
Subject: [PATCH 2/8] removal

---
 .gitignore                              |   1 +
 DeepCache.egg-info/PKG-INFO             | 281 ------------------------
 DeepCache.egg-info/SOURCES.txt          |  29 ---
 DeepCache.egg-info/dependency_links.txt |   1 -
 DeepCache.egg-info/requires.txt         |   3 -
 DeepCache.egg-info/top_level.txt        |   1 -
 6 files changed, 1 insertion(+), 315 deletions(-)
 delete mode 100644 DeepCache.egg-info/PKG-INFO
 delete mode 100644 DeepCache.egg-info/SOURCES.txt
 delete mode 100644 DeepCache.egg-info/dependency_links.txt
 delete mode 100644 DeepCache.egg-info/requires.txt
 delete mode 100644 DeepCache.egg-info/top_level.txt

diff --git a/.gitignore b/.gitignore
index 3da4190..f81c83c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 __pycache__
 output.png
 
+
 .idea
\ No newline at end of file
diff --git a/DeepCache.egg-info/PKG-INFO b/DeepCache.egg-info/PKG-INFO
deleted file mode 100644
index 33b66e9..0000000
--- a/DeepCache.egg-info/PKG-INFO
+++ /dev/null
@@ -1,281 +0,0 @@
-Metadata-Version: 2.1
-Name: DeepCache
-Version: 0.1.0
-Summary: DeepCache: Accelerating Diffusion Models for Free
-Home-page: https://github.com/horseee/DeepCache
-Author: Xinyin Ma
-Author-email: maxinyin@u.nus.edu
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.6
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch
-Requires-Dist: diffusers
-Requires-Dist: transformers
-
-# DeepCache: Accelerating Diffusion Models for Free
-<div align="center">
-  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/example_compress.gif" width="100%" ></img>
-  <br>
-  <em>
-      (Results on Stable Diffusion v1.5. Left: 50 PLMS steps. Right: 2.3x acceleration upon 50 PLMS steps) 
-  </em>
-</div>
-<br>
-
-> **DeepCache: Accelerating Diffusion Models for Free**   
-> [Xinyin Ma](https://horseee.github.io/), [Gongfan Fang](https://fangggf.github.io/), [Xinchao Wang](https://sites.google.com/site/sitexinchaowang/)   
-> [Learning and Vision Lab](http://lv-nus.org/), National University of Singapore  
-> 🥯[[Arxiv]](https://arxiv.org/abs/2312.00858)🎄[[Project Page]](https://horseee.github.io/Diffusion_DeepCache/)
-
-### Why DeepCache  
-* 🚀 Training-free and almost lossless
-* 🚀 Support Stable Diffusion, Stable Diffusion XL, Stable Video Diffusion, Stable Diffusion Pipeline / XL Pipeline for Inpainting, Stable Diffusion Img2Img Pipeline, DDPM
-* 🚀 Compatible with sampling algorithms like DDIM and PLMS
-
-### Updates
-* **June 27, 2024**: 🔥Our new work [AsyncDiff](https://github.com/czg1225/AsyncDiff) enables parallel inference of diffusion models on multiple GPUs. Check our [paper](https://arxiv.org/abs/2406.06911) and [code](https://github.com/czg1225/AsyncDiff)!
-* **June 5, 2024**: 🔥Our new work, [Learning-to-Cache](https://github.com/horseee/learning-to-cache/tree/main), an improved version of DeepCache on DiT. Code and checkpoints are released.
-* **January 5, 2024**: 💥A doc page for DeepCache has been added in [Diffusers](https://github.com/huggingface/diffusers)! Check [here](https://huggingface.co/docs/diffusers/main/en/optimization/deepcache) for more information. Many thanks to the Diffusers team!
-* **December 26, 2023**: 🔥Update a plug-and-play implementation of DeepCache, no longer requiring any modifications of the diffuser's code! Check [here](https://github.com/horseee/DeepCache#usage) for the detailed usage! Big thanks to @[yuanshi9815](https://github.com/Yuanshi9815) for contributing the code!
-* **December 25, 2023**: A demo is available via Colab [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1COM9tfGvHSJ8Zn4tScZCGgt8nxAEOlAz?usp=sharing)
-* **December 21, 2023**: Release the code for **Stable Video Diffusion** and **Text2Video-Zero**. In the figure below, the upper row shows the original videos generated by SVD-XT, and the lower row is accelerated by DeepCache. For Text2Video-Zero, the results can be found [here](https://github.com/horseee/DeepCache/blob/master/assets/Text2Video-Zero.gif)
-<div align="center">
-  <img src="assets/svd.gif" width="90%" ></img>
-  <br>
-  <em>
-      (1.7x acceleration of SVD-XT) 
-  </em>
-</div>
-
-
-* **December 20, 2023**: Release the code for **DDPM**. See [here](https://github.com/horseee/DeepCache/tree/master/DeepCache#experiment-code-for-ddpm) for the experimental code and instructions.
-
-* **December 6, 2023**: Release the code for **Stable Diffusion XL**. The results of the `stabilityai/stable-diffusion-xl-base-1.0` are shown in the below figure, with the same prompts from the first figure.
-<div align="center">
-  <img src="assets/sdxl.png" width="90%" ></img>
-  <br>
-  <em>
-      (2.6x acceleration of Stable Diffusion XL) 
-  </em>
-</div>
-
-
-### Introduction
-We introduce **DeepCache**, a novel **training-free and almost lossless** paradigm that accelerates diffusion models from the perspective of model architecture. Utilizing the property of the U-Net, we reuse the high-level features while updating the low-level features in a very cheap way. DeepCache accelerates Stable Diffusion v1.5 by 2.3x with only a 0.05 decline in CLIP Score, and LDM-4-G(ImageNet) by 4.1x with a 0.22 decrease in FID.
-
-<div align="center">
-<img width="50%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/9ce3930c-c84c-4af8-8c6a-b6803a5a7b1d">
-</div>
-
-## Quick Start
-
-### Install
-```bash
-pip install DeepCache
-```
-
-### Usage
-
-```python
-import torch
-
-# Loading the original pipeline
-from diffusers import StableDiffusionPipeline
-pipe = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16).to("cuda:0")
-
-# Import the DeepCacheSDHelper
-from DeepCache import DeepCacheSDHelper
-helper = DeepCacheSDHelper(pipe=pipe)
-helper.set_params(
-    cache_interval=3,
-    cache_branch_id=0,
-)
-helper.enable()
-
-# Generate Image
-deepcache_image = pipe(
-        prompt,
-        output_type='pt'
-).images[0]
-helper.disable()
-```
-We here take the Stable Diffusion pipeline as an example. You can replace pipe with any variants of the Stable Diffusion pipeline, including choices like SDXL, SVD, and more. You can find examples in the [script](https://github.com/horseee/DeepCache#a-general-script-for-sd). The argument `cache_branch_id` specifies the selected skip branch. For the skip branches that are deeper, the model will engage them only during the caching steps, and exclude them during the retrieval steps. The argument `cache_interval` represents the interval for updating the cache.
-
-### A general script for SD
-```bash
-python main.py --model_type sdxl #Support [sdxl, sd1.5, sd2.1, svd, sd-inpaint, sdxl-inpaint, sd-img2img]
-```
-
-## Experimental code for DeepCache
-The above implementation does not require changes to the `forward` or `__call__` functions in the Diffusers pipeline, and is, therefore, more general. The following section is the experimental code that can be used to reproduce the results in the paper. It is implemented one by one for different model structures and pipelines, and thus, may not work properly due to the update of diffusers.
-
-### Setup
-```
-pip install diffusers==0.24.0 transformers
-```
-
-### Stable Diffusion XL
-```bash
-python stable_diffusion_xl.py --model stabilityai/stable-diffusion-xl-base-1.0
-```
-
-<details>
-<summary>Output:</summary>
-
-```bash
-Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  6.62it/s]
-2023-12-06 01:44:28,578 - INFO - Running baseline...
-100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:17<00:00,  2.93it/s]
-2023-12-06 01:44:46,095 - INFO - Baseline: 17.52 seconds
-Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.06it/s]
-2023-12-06 01:45:02,865 - INFO - Running DeepCache...
-100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:06<00:00,  8.01it/s]
-2023-12-06 01:45:09,573 - INFO - DeepCache: 6.71 seconds
-2023-12-06 01:45:10,678 - INFO - Saved to output.png. Done!
-```
-
-</details>
-
-You can add `--refine` at the end of the command to activate the refiner model for SDXL.
-
-### Stable Diffusion v1.5
-```bash
-python stable_diffusion.py --model runwayml/stable-diffusion-v1-5
-```
-
-<details>
-<summary>Output:</summary>
-
-```bash
-2023-12-03 16:18:13,636 - INFO - Loaded safety_checker as StableDiffusionSafetyChecker from `safety_checker` subfolder of runwayml/stable-diffusion-v1-5.
-2023-12-03 16:18:13,699 - INFO - Loaded vae as AutoencoderKL from `vae` subfolder of runwayml/stable-diffusion-v1-5.
-Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  5.88it/s]
-2023-12-03 16:18:22,837 - INFO - Running baseline...
-100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:03<00:00, 15.33it/s]
-2023-12-03 16:18:26,174 - INFO - Baseline: 3.34 seconds
-2023-12-03 16:18:26,174 - INFO - Running DeepCache...
-100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:01<00:00, 34.06it/s]
-2023-12-03 16:18:27,718 - INFO - DeepCache: 1.54 seconds
-2023-12-03 16:18:27,935 - INFO - Saved to output.png. Done!
-```
-
-</details>
-
-### Stable Diffusion v2.1
-
-```bash
-python stable_diffusion.py --model stabilityai/stable-diffusion-2-1
-```
-
-<details>
-  <summary>Output:</summary>
-  
-```bash
-2023-12-03 16:21:17,858 - INFO - Loaded feature_extractor as CLIPImageProcessor from `feature_extractor` subfolder of stabilityai/stable-diffusion-2-1.
-2023-12-03 16:21:17,864 - INFO - Loaded scheduler as DDIMScheduler from `scheduler` subfolder of stabilityai/stable-diffusion-2-1.
-Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00,  5.35it/s]
-2023-12-03 16:21:49,770 - INFO - Running baseline...
-100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:14<00:00,  3.42it/s]
-2023-12-03 16:22:04,551 - INFO - Baseline: 14.78 seconds
-2023-12-03 16:22:04,551 - INFO - Running DeepCache...
-100%|████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:08<00:00,  6.10it/s]
-2023-12-03 16:22:12,911 - INFO - DeepCache: 8.36 seconds
-2023-12-03 16:22:13,417 - INFO - Saved to output.png. Done!
-```
-
-</details>
-
-Currently, our code supports the models that can be loaded by [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/v0.24.0/en/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline). You can specify the model name by the argument `--model`, which by default, is `runwayml/stable-diffusion-v1-5`. 
-
-### Stable Video Diffusion 
-```bash
-python stable_video_diffusion.py
-```
-
-<details>
-  <summary>Output:</summary>
-  
-```bash
-Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00,  8.36it/s]
-2023-12-21 04:56:47,329 - INFO - Running baseline...
-100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [01:27<00:00,  3.49s/it]
-2023-12-21 04:58:26,121 - INFO - Origin: 98.66 seconds
-Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 10.59it/s]
-2023-12-21 04:58:27,202 - INFO - Running DeepCache...
-100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [00:49<00:00,  1.96s/it]
-2023-12-21 04:59:26,607 - INFO - DeepCache: 59.31 seconds
-```
-
-</details>
-
-### DDPM and LDM
-Please check [here](https://github.com/horseee/DeepCache/tree/master/experiments) for the experimental code of DDPM and LDM.
-
-
-## Visualization
-
-Images in the upper line are the baselines and the images in the lower line are accelerated by DeepCache. 
-
-### Stable Diffusion v1.5 (2.15x Acceleration)
-<div align="center">
-  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/sd_examples_high_res.png" width="100%">
-</div>
-
-### LDM-4-G for ImageNet (6.96x Acceleration)
-<div align="center">
-  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/imagenet.png" width="100%">
-</div>
-
-### DDPM for LSUN Church & Bedroom (1.48x Acceleration)
-<div align="center">
-  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/bedroom.png" width="100%">
-  <img src="https://github.com/horseee/Diffusion_DeepCache/blob/master/static/images/church.png" width="100%">
-</div>
-
-## Quantitative Results
-
-### ImageNet
-<div align="center">
-<img width="80%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/151d7639-2501-45cf-8de5-2af3bb5a354b">
-</div>
-
-
-### Stable Diffusion v1.5
-
-<div align="center">
-<img width="80%" alt="image" src="https://github.com/horseee/DeepCache/assets/18592211/e9bd7a8e-07c8-4296-95a2-12d008995807">
-</div>
-
-More results can be found in our [paper](https://arxiv.org/abs/2312.00858)
-
-## Other Implementations and Plugins
-We sincerely thank the authors listed below who implemented DeepCache in plugins or other contexts. 
-
-* Diffusers: https://huggingface.co/docs/diffusers/main/en/optimization/deepcache
-* OneDiff Integration: https://github.com/siliconflow/onediff?tab=readme-ov-file#easy-to-use by @[SiliconFlow](https://github.com/siliconflow). OneDiff also has implementations for DeepCache on SVD, check [this](https://github.com/siliconflow/onediff/pull/438) for details.
-* Comfyui: https://gist.github.com/laksjdjf/435c512bc19636e9c9af4ee7bea9eb86 by @[laksjdjf](https://github.com/laksjdjf)
-* Colab & Gradio: https://github.com/camenduru/DeepCache-colab by @[camenduru](https://github.com/camenduru/DeepCache-colab)
-* WebUI: https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14210 by @[aria1th](https://github.com/aria1th)
-* SD.Next: https://github.com/vladmandic/automatic by @[vladmandic](https://github.com/vladmandic)
-* MLX_deepcache for running SD on M1 Mac: https://github.com/Barry0121/mlx_deepcache by @[Barry Xue](https://github.com/Barry0121)
-
-We warmly welcome contributions from everyone. Please feel free to reach out to us.
-
-
-## Bibtex
-```
-@inproceedings{ma2023deepcache,
-  title={DeepCache: Accelerating Diffusion Models for Free},
-  author={Ma, Xinyin and Fang, Gongfan and Wang, Xinchao},
-  booktitle={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
-  year={2024}
-}
-```
-
-
-
-
-
diff --git a/DeepCache.egg-info/SOURCES.txt b/DeepCache.egg-info/SOURCES.txt
deleted file mode 100644
index 6fb784b..0000000
--- a/DeepCache.egg-info/SOURCES.txt
+++ /dev/null
@@ -1,29 +0,0 @@
-LICENSE
-README.md
-setup.py
-DeepCache/__init__.py
-DeepCache/flops.py
-DeepCache.egg-info/PKG-INFO
-DeepCache.egg-info/SOURCES.txt
-DeepCache.egg-info/dependency_links.txt
-DeepCache.egg-info/requires.txt
-DeepCache.egg-info/top_level.txt
-DeepCache/extension/__init__.py
-DeepCache/extension/deepcache.py
-DeepCache/sd/__init__.py
-DeepCache/sd/pipeline_stable_diffusion.py
-DeepCache/sd/pipeline_text_to_video_zero.py
-DeepCache/sd/pipeline_utils.py
-DeepCache/sd/unet_2d_blocks.py
-DeepCache/sd/unet_2d_condition.py
-DeepCache/sdxl/__init__.py
-DeepCache/sdxl/pipeline_stable_diffusion_xl.py
-DeepCache/sdxl/pipeline_stable_diffusion_xl_img2img.py
-DeepCache/sdxl/pipeline_utils.py
-DeepCache/sdxl/unet_2d_blocks.py
-DeepCache/sdxl/unet_2d_condition.py
-DeepCache/svd/__init__.py
-DeepCache/svd/pipeline_stable_video_diffusion.py
-DeepCache/svd/pipeline_utils.py
-DeepCache/svd/unet_3d_blocks.py
-DeepCache/svd/unet_spatio_temporal_condition.py
\ No newline at end of file
diff --git a/DeepCache.egg-info/dependency_links.txt b/DeepCache.egg-info/dependency_links.txt
deleted file mode 100644
index 8b13789..0000000
--- a/DeepCache.egg-info/dependency_links.txt
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/DeepCache.egg-info/requires.txt b/DeepCache.egg-info/requires.txt
deleted file mode 100644
index de80785..0000000
--- a/DeepCache.egg-info/requires.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-torch
-diffusers
-transformers
diff --git a/DeepCache.egg-info/top_level.txt b/DeepCache.egg-info/top_level.txt
deleted file mode 100644
index 924077e..0000000
--- a/DeepCache.egg-info/top_level.txt
+++ /dev/null
@@ -1 +0,0 @@
-DeepCache

From 016f94bb45972333eab6813d10ce3e3620c9cb43 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 15:17:41 -0700
Subject: [PATCH 3/8] ignore egginfo

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index f81c83c..1dcfb94 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 __pycache__
 output.png
 
+*.egg-info
 
 .idea
\ No newline at end of file

From 84f6cd0170bac99218b68915b474842e41b3324b Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 15:30:48 -0700
Subject: [PATCH 4/8] Privatize internal methods

---
 DeepCache/extension/deepcache.py | 42 ++++++++++++++++----------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index bf6ad43..4d745bf 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -6,8 +6,8 @@ def __init__(self, pipe=None):
 
     def enable(self):
         if not self._enabled:
-            self.reset_states()
-            self.wrap_modules()
+            self._reset_states()
+            self._wrap_modules()
             self._enabled = True
             print("Enabling Deepcache")
         else:
@@ -15,8 +15,8 @@ def enable(self):
 
     def disable(self):
         if self._enabled:
-            self.unwrap_modules()
-            self.reset_states()
+            self._unwrap_modules()
+            self._reset_states()
             print("Disabling Deepcache")
         else:
             print("DeepCache is already disabled.")
@@ -31,7 +31,7 @@ def set_params(self,cache_interval=1, cache_branch_id=0, skip_mode='uniform'):
             'skip_mode': skip_mode
         }
 
-    def is_skip_step(self, block_i, layer_i, blocktype = "down"):
+    def _is_skip_step(self, block_i, layer_i, blocktype ="down"):
         self.start_timestep = self.cur_timestep if self.start_timestep is None else self.start_timestep # For some pipeline that the first timestep != 0
         cache_interval, cache_layer_id, cache_block_id, skip_mode = \
             self.params['cache_interval'], self.params['cache_layer_id'], self.params['cache_block_id'], self.params['skip_mode']
@@ -45,7 +45,7 @@ def is_skip_step(self, block_i, layer_i, blocktype = "down"):
     def is_enter_position(self, block_i, layer_i):
         return block_i == self.params['cache_block_id'] and layer_i == self.params['cache_layer_id']
 
-    def wrap_unet_forward(self):
+    def _wrap_unet_forward(self):
         self.function_dict['unet_forward'] = self.pipe.unet.forward
         def wrapped_forward(*args, **kwargs):
             self.cur_timestep = list(self.pipe.scheduler.timesteps).index(args[1].item())
@@ -53,44 +53,44 @@ def wrapped_forward(*args, **kwargs):
             return result
         self.pipe.unet.forward = wrapped_forward
 
-    def wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype = "down"):
+    def _wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype ="down"):
         self.function_dict[
             (blocktype, block_name, block_i, layer_i)
         ] = block.forward
         def wrapped_forward(*args, **kwargs):
-            skip = self.is_skip_step(block_i, layer_i, blocktype)
+            skip = self._is_skip_step(block_i, layer_i, blocktype)
             result = self.cached_output[(blocktype, block_name, block_i, layer_i)] if skip else self.function_dict[(blocktype, block_name,  block_i, layer_i)](*args, **kwargs)
             if not skip: self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
             return result
         block.forward = wrapped_forward
 
-    def wrap_modules(self):
+    def _wrap_modules(self):
         # 1. wrap unet forward
-        self.wrap_unet_forward()
+        self._wrap_unet_forward()
         # 2. wrap downblock forward
         for block_i, block in enumerate(self.pipe.unet.down_blocks):
             for (layer_i, attention) in enumerate(getattr(block, "attentions", [])):
-                self.wrap_block_forward(attention, "attentions", block_i, layer_i)
+                self._wrap_block_forward(attention, "attentions", block_i, layer_i)
             for (layer_i, resnet) in enumerate(getattr(block, "resnets", [])):
-                self.wrap_block_forward(resnet, "resnet", block_i, layer_i)
+                self._wrap_block_forward(resnet, "resnet", block_i, layer_i)
             for downsampler in getattr(block, "downsamplers", []) if block.downsamplers else []:
-                self.wrap_block_forward(downsampler, "downsampler", block_i, len(getattr(block, "resnets", [])))
-            self.wrap_block_forward(block, "block", block_i, 0, blocktype = "down")
+                self._wrap_block_forward(downsampler, "downsampler", block_i, len(getattr(block, "resnets", [])))
+            self._wrap_block_forward(block, "block", block_i, 0, blocktype ="down")
         # 3. wrap midblock forward
-        self.wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype = "mid")
+        self._wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype ="mid")
         # 4. wrap upblock forward
         block_num = len(self.pipe.unet.up_blocks)
         for block_i, block in enumerate(self.pipe.unet.up_blocks):
             layer_num = len(getattr(block, "resnets", []))
             for (layer_i, attention) in enumerate(getattr(block, "attentions", [])):
-                self.wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype = "up")
+                self._wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
             for (layer_i, resnet) in enumerate(getattr(block, "resnets", [])):
-                self.wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype = "up")
+                self._wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
             for upsampler in getattr(block, "upsamplers", []) if block.upsamplers else []:
-                self.wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype = "up")
-            self.wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype = "up")
+                self._wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype ="up")
+            self._wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype ="up")
 
-    def unwrap_modules(self):
+    def _unwrap_modules(self):
         # 1. unet forward
         self.pipe.unet.forward = self.function_dict['unet_forward']
         # 2. downblock forward
@@ -116,7 +116,7 @@ def unwrap_modules(self):
                 upsampler.forward = self.function_dict[("up", "upsampler", block_num - block_i - 1, 0)]
             block.forward = self.function_dict[("up", "block", block_num - block_i - 1, 0)]
 
-    def reset_states(self):
+    def _reset_states(self):
         self.cur_timestep = 0
         self.function_dict = {}
         self.cached_output = {}

From 52c474e02cb2ca9026dd3000f18cf944420ad828 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 16:23:31 -0700
Subject: [PATCH 5/8] Code fix

---
 DeepCache/extension/deepcache.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index 4d745bf..8fa9b37 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -36,10 +36,12 @@ def _is_skip_step(self, block_i, layer_i, blocktype ="down"):
         cache_interval, cache_layer_id, cache_block_id, skip_mode = \
             self.params['cache_interval'], self.params['cache_layer_id'], self.params['cache_block_id'], self.params['skip_mode']
         if skip_mode == 'uniform':
-            if (self.cur_timestep-self.start_timestep) % cache_interval == 0: return False
+            if (self.cur_timestep-self.start_timestep) % cache_interval == 0:
+                return False
         if block_i > cache_block_id or blocktype == 'mid':
             return True
-        if block_i < cache_block_id: return False
+        if block_i < cache_block_id:
+            return False
         return layer_i >= cache_layer_id if blocktype == 'down' else layer_i > cache_layer_id
         
     def is_enter_position(self, block_i, layer_i):
@@ -60,7 +62,8 @@ def _wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype ="d
         def wrapped_forward(*args, **kwargs):
             skip = self._is_skip_step(block_i, layer_i, blocktype)
             result = self.cached_output[(blocktype, block_name, block_i, layer_i)] if skip else self.function_dict[(blocktype, block_name,  block_i, layer_i)](*args, **kwargs)
-            if not skip: self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
+            if not skip:
+                self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
             return result
         block.forward = wrapped_forward
 
@@ -121,3 +124,17 @@ def _reset_states(self):
         self.function_dict = {}
         self.cached_output = {}
         self.start_timestep = None
+
+
+if __name__ == "__main__":
+
+    from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+    import torch
+
+
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        'stable-diffusion-v1-5/stable-diffusion-v1-5',
+        torch_dtype=torch.float16
+    )
+    print(pipe.unet)
\ No newline at end of file

From bdccd2b66cee8e54dc12c83687af931f21638165 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 16:23:53 -0700
Subject: [PATCH 6/8] update

---
 DeepCache/extension/deepcache.py | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index 8fa9b37..abbb10d 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -124,17 +124,3 @@ def _reset_states(self):
         self.function_dict = {}
         self.cached_output = {}
         self.start_timestep = None
-
-
-if __name__ == "__main__":
-
-    from diffusers import StableDiffusionPipeline, UNet2DConditionModel
-    import torch
-
-
-
-    pipe = StableDiffusionPipeline.from_pretrained(
-        'stable-diffusion-v1-5/stable-diffusion-v1-5',
-        torch_dtype=torch.float16
-    )
-    print(pipe.unet)
\ No newline at end of file

From c170b71312ec68a4bfb6f0ebb3fa128ffe11a6fa Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 16:27:33 -0700
Subject: [PATCH 7/8] revert

---
 DeepCache/extension/deepcache.py | 43 ++++++++++++++++----------------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index abbb10d..fd99325 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -6,8 +6,8 @@ def __init__(self, pipe=None):
 
     def enable(self):
         if not self._enabled:
-            self._reset_states()
-            self._wrap_modules()
+            self.reset_states()
+            self.wrap_modules()
             self._enabled = True
             print("Enabling Deepcache")
         else:
@@ -15,8 +15,9 @@ def enable(self):
 
     def disable(self):
         if self._enabled:
-            self._unwrap_modules()
-            self._reset_states()
+            self.unwrap_modules()
+            self.reset_states()
+            self._enabled = False
             print("Disabling Deepcache")
         else:
             print("DeepCache is already disabled.")
@@ -31,7 +32,7 @@ def set_params(self,cache_interval=1, cache_branch_id=0, skip_mode='uniform'):
             'skip_mode': skip_mode
         }
 
-    def _is_skip_step(self, block_i, layer_i, blocktype ="down"):
+    def is_skip_step(self, block_i, layer_i, blocktype ="down"):
         self.start_timestep = self.cur_timestep if self.start_timestep is None else self.start_timestep # For some pipeline that the first timestep != 0
         cache_interval, cache_layer_id, cache_block_id, skip_mode = \
             self.params['cache_interval'], self.params['cache_layer_id'], self.params['cache_block_id'], self.params['skip_mode']
@@ -47,7 +48,7 @@ def _is_skip_step(self, block_i, layer_i, blocktype ="down"):
     def is_enter_position(self, block_i, layer_i):
         return block_i == self.params['cache_block_id'] and layer_i == self.params['cache_layer_id']
 
-    def _wrap_unet_forward(self):
+    def wrap_unet_forward(self):
         self.function_dict['unet_forward'] = self.pipe.unet.forward
         def wrapped_forward(*args, **kwargs):
             self.cur_timestep = list(self.pipe.scheduler.timesteps).index(args[1].item())
@@ -55,45 +56,45 @@ def wrapped_forward(*args, **kwargs):
             return result
         self.pipe.unet.forward = wrapped_forward
 
-    def _wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype ="down"):
+    def wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype ="down"):
         self.function_dict[
             (blocktype, block_name, block_i, layer_i)
         ] = block.forward
         def wrapped_forward(*args, **kwargs):
-            skip = self._is_skip_step(block_i, layer_i, blocktype)
+            skip = self.is_skip_step(block_i, layer_i, blocktype)
             result = self.cached_output[(blocktype, block_name, block_i, layer_i)] if skip else self.function_dict[(blocktype, block_name,  block_i, layer_i)](*args, **kwargs)
             if not skip:
                 self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
             return result
         block.forward = wrapped_forward
 
-    def _wrap_modules(self):
+    def wrap_modules(self):
         # 1. wrap unet forward
-        self._wrap_unet_forward()
+        self.wrap_unet_forward()
         # 2. wrap downblock forward
         for block_i, block in enumerate(self.pipe.unet.down_blocks):
             for (layer_i, attention) in enumerate(getattr(block, "attentions", [])):
-                self._wrap_block_forward(attention, "attentions", block_i, layer_i)
+                self.wrap_block_forward(attention, "attentions", block_i, layer_i)
             for (layer_i, resnet) in enumerate(getattr(block, "resnets", [])):
-                self._wrap_block_forward(resnet, "resnet", block_i, layer_i)
+                self.wrap_block_forward(resnet, "resnet", block_i, layer_i)
             for downsampler in getattr(block, "downsamplers", []) if block.downsamplers else []:
-                self._wrap_block_forward(downsampler, "downsampler", block_i, len(getattr(block, "resnets", [])))
-            self._wrap_block_forward(block, "block", block_i, 0, blocktype ="down")
+                self.wrap_block_forward(downsampler, "downsampler", block_i, len(getattr(block, "resnets", [])))
+            self.wrap_block_forward(block, "block", block_i, 0, blocktype ="down")
         # 3. wrap midblock forward
-        self._wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype ="mid")
+        self.wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype ="mid")
         # 4. wrap upblock forward
         block_num = len(self.pipe.unet.up_blocks)
         for block_i, block in enumerate(self.pipe.unet.up_blocks):
             layer_num = len(getattr(block, "resnets", []))
             for (layer_i, attention) in enumerate(getattr(block, "attentions", [])):
-                self._wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
+                self.wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
             for (layer_i, resnet) in enumerate(getattr(block, "resnets", [])):
-                self._wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
+                self.wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
             for upsampler in getattr(block, "upsamplers", []) if block.upsamplers else []:
-                self._wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype ="up")
-            self._wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype ="up")
+                self.wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype ="up")
+            self.wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype ="up")
 
-    def _unwrap_modules(self):
+    def unwrap_modules(self):
         # 1. unet forward
         self.pipe.unet.forward = self.function_dict['unet_forward']
         # 2. downblock forward
@@ -119,7 +120,7 @@ def _unwrap_modules(self):
                 upsampler.forward = self.function_dict[("up", "upsampler", block_num - block_i - 1, 0)]
             block.forward = self.function_dict[("up", "block", block_num - block_i - 1, 0)]
 
-    def _reset_states(self):
+    def reset_states(self):
         self.cur_timestep = 0
         self.function_dict = {}
         self.cached_output = {}

From e0b25c570cdc7a2036ca9477e696a534c9dfb8f0 Mon Sep 17 00:00:00 2001
From: Ruixi Fan <ruixi@typeface.ai>
Date: Fri, 13 Sep 2024 16:29:07 -0700
Subject: [PATCH 8/8] Update state

---
 DeepCache/extension/deepcache.py | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/DeepCache/extension/deepcache.py b/DeepCache/extension/deepcache.py
index fd99325..b625ac5 100644
--- a/DeepCache/extension/deepcache.py
+++ b/DeepCache/extension/deepcache.py
@@ -32,17 +32,15 @@ def set_params(self,cache_interval=1, cache_branch_id=0, skip_mode='uniform'):
             'skip_mode': skip_mode
         }
 
-    def is_skip_step(self, block_i, layer_i, blocktype ="down"):
+    def is_skip_step(self, block_i, layer_i, blocktype = "down"):
         self.start_timestep = self.cur_timestep if self.start_timestep is None else self.start_timestep # For some pipeline that the first timestep != 0
         cache_interval, cache_layer_id, cache_block_id, skip_mode = \
             self.params['cache_interval'], self.params['cache_layer_id'], self.params['cache_block_id'], self.params['skip_mode']
         if skip_mode == 'uniform':
-            if (self.cur_timestep-self.start_timestep) % cache_interval == 0:
-                return False
+            if (self.cur_timestep-self.start_timestep) % cache_interval == 0: return False
         if block_i > cache_block_id or blocktype == 'mid':
             return True
-        if block_i < cache_block_id:
-            return False
+        if block_i < cache_block_id: return False
         return layer_i >= cache_layer_id if blocktype == 'down' else layer_i > cache_layer_id
         
     def is_enter_position(self, block_i, layer_i):
@@ -56,15 +54,14 @@ def wrapped_forward(*args, **kwargs):
             return result
         self.pipe.unet.forward = wrapped_forward
 
-    def wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype ="down"):
+    def wrap_block_forward(self, block, block_name, block_i, layer_i, blocktype = "down"):
         self.function_dict[
             (blocktype, block_name, block_i, layer_i)
         ] = block.forward
         def wrapped_forward(*args, **kwargs):
             skip = self.is_skip_step(block_i, layer_i, blocktype)
             result = self.cached_output[(blocktype, block_name, block_i, layer_i)] if skip else self.function_dict[(blocktype, block_name,  block_i, layer_i)](*args, **kwargs)
-            if not skip:
-                self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
+            if not skip: self.cached_output[(blocktype, block_name, block_i, layer_i)] = result
             return result
         block.forward = wrapped_forward
 
@@ -79,20 +76,20 @@ def wrap_modules(self):
                 self.wrap_block_forward(resnet, "resnet", block_i, layer_i)
             for downsampler in getattr(block, "downsamplers", []) if block.downsamplers else []:
                 self.wrap_block_forward(downsampler, "downsampler", block_i, len(getattr(block, "resnets", [])))
-            self.wrap_block_forward(block, "block", block_i, 0, blocktype ="down")
+            self.wrap_block_forward(block, "block", block_i, 0, blocktype = "down")
         # 3. wrap midblock forward
-        self.wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype ="mid")
+        self.wrap_block_forward(self.pipe.unet.mid_block, "mid_block", 0, 0, blocktype = "mid")
         # 4. wrap upblock forward
         block_num = len(self.pipe.unet.up_blocks)
         for block_i, block in enumerate(self.pipe.unet.up_blocks):
             layer_num = len(getattr(block, "resnets", []))
             for (layer_i, attention) in enumerate(getattr(block, "attentions", [])):
-                self.wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
+                self.wrap_block_forward(attention, "attentions", block_num - block_i - 1, layer_num - layer_i - 1, blocktype = "up")
             for (layer_i, resnet) in enumerate(getattr(block, "resnets", [])):
-                self.wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype ="up")
+                self.wrap_block_forward(resnet, "resnet", block_num - block_i - 1, layer_num - layer_i - 1, blocktype = "up")
             for upsampler in getattr(block, "upsamplers", []) if block.upsamplers else []:
-                self.wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype ="up")
-            self.wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype ="up")
+                self.wrap_block_forward(upsampler, "upsampler", block_num - block_i - 1, 0, blocktype = "up")
+            self.wrap_block_forward(block, "block", block_num - block_i - 1, 0, blocktype = "up")
 
     def unwrap_modules(self):
         # 1. unet forward