37 lines
1.1 KiB
Python
37 lines
1.1 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.0, "max": 1.0, "step": 0.05},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_ras"
|
|
CATEGORY = "ras"
|
|
|
|
def apply_ras(self, model: ModelPatcher, sample_ratio: 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)
|
|
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"]
|