add readme

README.md

@@ -1 +1,13 @@
 # RegionalAdaptiveSampling
+
+[Regional Adaptive Sampling](https://github.com/microsoft/RAS) is a new technique for accelerating the inference of diffusion transformers. 
+It essentially works as a KV Cache inside the model, picking regions that are likely to be updated by each diffusion step and passing in only those tokens.
+
+This implementation is simple to use, and compatible with Flux (dev & schnell) at HunYuanVideo. I may add support for other models in the future.
+
+## Usage
+Apply the `Regional Adaptive Sampling` node to the desired model. It has the following parameters:  
+- **sample_ratio**: The percent of tokens to keep in the model on a RAS pass. Anything below 0.3 is usually very bad quality.
+- **warmup_steps**: The number of steps to do without RAS at the beginning. Setting higher will decrease the speedup, and setting it lower will degrade the composition.
+- **hydrate_every**: Every `hydrate_every` steps, we do a full run through the model with all tokens, to refresh the stale cache. Set to 0 to disable and do full RAS.
+- **starvation_scale**: Controls how the model decides which part of the image to focus on. Increasing it will probably shift quality from the main subject to the background. The default of 0.1 is what's used in the paper, and I haven't tried anything else.

__init__.py

@@ -10,7 +10,19 @@ class RegionalAdaptiveSampling:
                 "model": ("MODEL",),
                 "sample_ratio": (
                     "FLOAT",
-                    {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05},
+                    {"default": 0.5, "min": 0.05, "max": 1.0, "step": 0.05},
+                ),
+                "warmup_steps": (
+                    "INT",
+                    {"default": 4, "min": 0, "max": 100},
+                ),
+                "hydrate_every": (
+                    "INT",
+                    {"default": 4, "min": 0, "max": 100},
+                ),
+                "starvation_scale": (
+                    "FLOAT",
+                    {"default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01},
                 ),
             }
         }
@@ -19,13 +31,25 @@ class RegionalAdaptiveSampling:
     FUNCTION = "apply_ras"
     CATEGORY = "ras"
 
-    def apply_ras(self, model: ModelPatcher, sample_ratio: float):
+    def apply_ras(
+        self,
+        model: ModelPatcher,
+        sample_ratio: float,
+        warmup_steps: int,
+        hydrate_every: int,
+        starvation_scale: float,
+    ):
         model = model.clone()
         # unpatch the model
         # this makes sure that we're wrapping the model "in a pure state"
         # the model will repatch itself later
         model.unpatch_model()
-        config = RASConfig(sample_ratio=sample_ratio)
+        config = RASConfig(
+            sample_ratio=sample_ratio,
+            warmup_steps=warmup_steps,
+            hydrate_every=hydrate_every,
+            starvation_scale=starvation_scale,
+        )
         manager = RASManager(config)
         manager.wrap_model(model)
         return (model,)

ras.py

@@ -33,7 +33,7 @@ def take_attributes_from(source, target, keys):
 
 @dataclass
 class RASConfig:
-    start_step: int = 4
+    warmup_steps: int = 4
     hydrate_every: int = 5
     sample_ratio: float = 0.5
     starvation_scale: float = 0.1
@@ -131,10 +131,13 @@ class RASManager:
         return int(i.item())
 
     def skip_ratio(self, timestep: int) -> float:
-        if timestep < self.config.start_step or (
-            timestep % self.config.hydrate_every == 0
-        ):
-            return 1
+        if timestep < self.config.warmup_steps:
+            return 0
+        if self.config.hydrate_every:
+            if (
+                1 + timestep - self.config.warmup_steps
+            ) % self.config.hydrate_every == 0:
+                return 0
         return 1.0 - self.config.sample_ratio
 
     def select_indices(self, diff: Tensor, timestep: int):
@@ -172,7 +175,7 @@ class RASManager:
         metric *= torch.exp(self.config.starvation_scale * self.drop_count)
         indices = torch.sort(metric, dim=-1, descending=False).indices
         skip_ratio = self.skip_ratio(timestep)
-        if skip_ratio >= 0.99:
+        if skip_ratio <= 0.01:
             # we're not dropping anything -- remove the live_indices
             # we use the value None to indicate a full hydrate
             self.live_img_indices = None