From c599c7e53c25ad4d3431366cd00035a0fb5cc930 Mon Sep 17 00:00:00 2001
From: Raphael Walker <slickytail.mc@gmail.com>
Date: Fri, 28 Feb 2025 11:34:04 +0100
Subject: [PATCH] add readme

---
 README.md   | 12 ++++++++++++
 __init__.py | 30 +++++++++++++++++++++++++++---
 ras.py      | 15 +++++++++------
 3 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 6539206..40ac53d 100644
--- a/README.md
+++ b/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.
diff --git a/__init__.py b/__init__.py
index 2710758..8b6f5b8 100644
--- a/__init__.py
+++ b/__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,)
diff --git a/ras.py b/ras.py
index db479ae..8eefd68 100644
--- a/ras.py
+++ b/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