61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
from .ras import RASConfig, RASManager
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
|
|
class RegionalAdaptiveSampling:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"sample_ratio": (
|
|
"FLOAT",
|
|
{"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},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_ras"
|
|
CATEGORY = "ras"
|
|
|
|
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,
|
|
warmup_steps=warmup_steps,
|
|
hydrate_every=hydrate_every,
|
|
starvation_scale=starvation_scale,
|
|
)
|
|
manager = RASManager(config)
|
|
manager.wrap_model(model)
|
|
return (model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"RegionalAdaptiveSampling": RegionalAdaptiveSampling}
|
|
NODE_DISPLAY_NAME_MAPPING = {"RegionalAdaptiveSampling": "Regional Adaptive Sampling"}
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPING"]
|