add optional arg and copyright header

Author: ziyu-guo

README.md

@@ -12,6 +12,10 @@
 **Figure 1. Accelerating Diffusion Transformer inference across multiple modalities with 50 DDIM Steps on DiT-XL-256x256, 100 DPM-Solver++(3M) SDE steps for a 10s audio sample (spectrogram shown) on Stable Audio Open, 30 Rectified Flow steps on Open-Sora 480p 2s videos**
 
 
+# Updates
+SmoothCache now supports generating cache schedues using a zero-intrusion external helper. See [run_calibration.py](./examples/run_calibration.py) to find out how it generates a schedule compatible with [HuggingFace Diffusers DiTPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/dit/pipeline_dit.py), without requiring any changes to Diffusers implementation!
+
+
 # Introduction
 We introduce **SmoothCache**, a straightforward acceleration technique for DiT architecture models, that's both **training-free, flexible and performant**. By leveraging layer-wise representation error, our method identifies redundancies in the diffusion process, generates a static caching scheme to reuse output featuremaps and therefore reduces the need for computationally expensive operations. This solution works across different models and modalities, can be easily dropped into existing Diffusion Transformer pipelines, can be stacked on different solvers, and requires no additional training or datasets. **SmoothCache** consistently outperforms various solvers designed to accelerate the diffusion process, while matching or surpassing the performance of existing modality-specific caching techniques.
 
@@ -26,7 +30,7 @@ We introduce **SmoothCache**, a straightforward acceleration technique for DiT a
 pip install SmoothCache
 ```
 
-### Usage
+### Usage - Inference
 
 Inspired by [DeepCache](https://raw.githubusercontent.com/horseee/DeepCache), we have implemented drop-in SmoothCache helper classes that easily applies to [Huggingface Diffuser DiTPipeline](https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/dit), and [original DiT implementations](https://github.com/facebookresearch/DiT).
 
@@ -138,6 +142,15 @@ cache_helper.disable()
 save_image(samples, "sample.png", nrow=4, normalize=True, value_range=(-1, 1))
 ```
 
+### Usage - Cache Schedule Generation
+See [run_calibration.py](./examples/run_calibration.py), which generates schedule for the self-attention module ([attn1](https://github.com/huggingface/diffusers/blob/37a5f1b3b69ed284086fb31fb1b49668cba6c365/src/diffusers/models/attention.py#L380)) 
+from Diffusers [BasicTransformerBlock](https://github.com/huggingface/diffusers/blob/37a5f1b3b69ed284086fb31fb1b49668cba6c365/src/diffusers/models/attention.py#L261C7-L261C28) block. 
+
+Note that only self-attention, and not cross-attention, is enabled in the stock config of Diffusers [DiT module](https://github.com/huggingface/diffusers/blob/37a5f1b3b69ed284086fb31fb1b49668cba6c365/src/diffusers/models/transformers/dit_transformer_2d.py#L72-L73). We leave this behavior
+as-is for the purpose of minimal intrusion. 
+
+We welcome all contributions aimed at expending SmoothCache's model coverage and module coverage. 
+
 ## Visualization
 
 (WIP)

SmoothCache/calibration/calibration_helper.py

@@ -1,4 +1,17 @@
-# calibration_helper.py
+# Copyright 2022 Roblox Corporation
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import json
 import re
 import statistics

SmoothCache/calibration/diffuser_calibration_helper.py

@@ -1,6 +1,18 @@
-# SmoothCache/calibration/diffuser_calibration_helper.py
+# Copyright 2022 Roblox Corporation
 
-from typing import List, Union, Type
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional
 import torch.nn as nn
 from .calibration_helper import CalibrationHelper
 
@@ -16,7 +28,8 @@ def __init__(
         calibration_lookahead: int = 3,
         calibration_threshold: float = 0.0,
         schedule_length: int = 50,
-        log_file: str = "calibration_schedule.json"
+        log_file: str = "calibration_schedule.json",
+        components_to_wrap: Optional[List[str]] = None
     ):
         """
         Diffuser-specific CalibrationHelper derived from CalibrationHelper.
@@ -25,18 +38,20 @@ def __init__(
             model (nn.Module): The model to wrap (e.g., pipe.transformer).
             calibration_lookahead (int): Steps to look back for error calculation.
             calibration_threshold (float): Cutoff L1 error value to enable caching.
-            schedule_length (int): Length of the generated schedule, 1:1 mapped to pipeline timesteps
+            schedule_length (int): Length of the generated schedule, 1:1 mapped to pipeline timesteps.
             log_file (str): Path to save the generated schedule JSON.
-
+            components_to_wrap (List[str], optional): List of component names to wrap.
+                Defaults to ['attn1'].
+        
         Raises:
             ImportError: If diffusers' BasicTransformerBlock is unavailable.
         """
         if BasicTransformerBlock is None:
             raise ImportError("Diffusers library not installed or BasicTransformerBlock not found.")
 
         block_classes = [BasicTransformerBlock]
-        components_to_wrap = ['attn1']  # Wrap 'attn1' component
-
+        if components_to_wrap is None:
+            components_to_wrap = ['attn1']
         super().__init__(
             model=model,
             block_classes=block_classes,

SmoothCache/diffuser_cache_helper.py

@@ -15,6 +15,7 @@
 """Helper Class for Diffusion Transformer Implemented at 
    https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/dit"""
 
+from typing import List, Optional
 from .smooth_cache_helper import SmoothCacheHelper
 
 try:

examples/run_calibration.py

@@ -1,4 +1,17 @@
-# example_calibration_run.py
+# Copyright 2022 Roblox Corporation
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 from diffusers import DiTPipeline, DPMSolverMultistepScheduler
 from SmoothCache import DiffuserCalibrationHelper