68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
from nodes import KSampler, KSamplerAdvanced
|
|
|
|
from .patch import patch_all, unpatch_all
|
|
from .efficient import (
|
|
NODE_CLASS_MAPPINGS as EFFICIENCY_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as EFFICIENCY_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
|
|
|
|
class KSamplerWithSharpness(KSampler):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = KSampler.INPUT_TYPES()
|
|
inputs["optional"] = {
|
|
"sharpness": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
|
)
|
|
}
|
|
|
|
return inputs
|
|
|
|
CATEGORY = "ArtVenture/Sampling"
|
|
|
|
def sample(self, *args, sharpness=2.0, **kwargs):
|
|
patch.sharpness = sharpness
|
|
patch_all()
|
|
results = super().sample(*args, **kwargs)
|
|
unpatch_all()
|
|
return results
|
|
|
|
|
|
class KSamplerAdvancedWithSharpness(KSamplerAdvanced):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = KSamplerAdvanced.INPUT_TYPES()
|
|
inputs["optional"] = {
|
|
"sharpness": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
|
)
|
|
}
|
|
|
|
return inputs
|
|
|
|
CATEGORY = "ArtVenture/Sampling"
|
|
|
|
def sample(self, *args, sharpness=2.0, **kwargs):
|
|
patch.sharpness = sharpness
|
|
patch_all()
|
|
results = super().sample(*args, **kwargs)
|
|
unpatch_all()
|
|
return results
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Fooocus_KSampler": KSamplerWithSharpness,
|
|
"Fooocus_KSamplerAdvanced": KSamplerAdvancedWithSharpness,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Fooocus_KSampler": "KSampler Fooocus",
|
|
"Fooocus_KSamplerAdvanced": "KSampler Adv. Fooocus",
|
|
}
|
|
|
|
NODE_CLASS_MAPPINGS.update(EFFICIENCY_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(EFFICIENCY_NODE_DISPLAY_NAME_MAPPINGS)
|