mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2025-12-11 06:54:26 +08:00
convert nodes_freelunch.py to the V3 schema (#10904)
This commit is contained in:
parent
ae676ed105
commit
117bf3f2bd
@ -2,6 +2,8 @@
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
import logging
|
import logging
|
||||||
|
from typing_extensions import override
|
||||||
|
from comfy_api.latest import ComfyExtension, IO
|
||||||
|
|
||||||
def Fourier_filter(x, threshold, scale):
|
def Fourier_filter(x, threshold, scale):
|
||||||
# FFT
|
# FFT
|
||||||
@ -22,21 +24,26 @@ def Fourier_filter(x, threshold, scale):
|
|||||||
return x_filtered.to(x.dtype)
|
return x_filtered.to(x.dtype)
|
||||||
|
|
||||||
|
|
||||||
class FreeU:
|
class FreeU(IO.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def define_schema(cls):
|
||||||
return {"required": { "model": ("MODEL",),
|
return IO.Schema(
|
||||||
"b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.01}),
|
node_id="FreeU",
|
||||||
"b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.01}),
|
category="model_patches/unet",
|
||||||
"s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}),
|
inputs=[
|
||||||
"s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}),
|
IO.Model.Input("model"),
|
||||||
}}
|
IO.Float.Input("b1", default=1.1, min=0.0, max=10.0, step=0.01),
|
||||||
RETURN_TYPES = ("MODEL",)
|
IO.Float.Input("b2", default=1.2, min=0.0, max=10.0, step=0.01),
|
||||||
FUNCTION = "patch"
|
IO.Float.Input("s1", default=0.9, min=0.0, max=10.0, step=0.01),
|
||||||
|
IO.Float.Input("s2", default=0.2, min=0.0, max=10.0, step=0.01),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
IO.Model.Output(),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
CATEGORY = "model_patches/unet"
|
@classmethod
|
||||||
|
def execute(cls, model, b1, b2, s1, s2) -> IO.NodeOutput:
|
||||||
def patch(self, model, b1, b2, s1, s2):
|
|
||||||
model_channels = model.model.model_config.unet_config["model_channels"]
|
model_channels = model.model.model_config.unet_config["model_channels"]
|
||||||
scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
|
scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
|
||||||
on_cpu_devices = {}
|
on_cpu_devices = {}
|
||||||
@ -59,23 +66,31 @@ class FreeU:
|
|||||||
|
|
||||||
m = model.clone()
|
m = model.clone()
|
||||||
m.set_model_output_block_patch(output_block_patch)
|
m.set_model_output_block_patch(output_block_patch)
|
||||||
return (m, )
|
return IO.NodeOutput(m)
|
||||||
|
|
||||||
class FreeU_V2:
|
patch = execute # TODO: remove
|
||||||
|
|
||||||
|
|
||||||
|
class FreeU_V2(IO.ComfyNode):
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def define_schema(cls):
|
||||||
return {"required": { "model": ("MODEL",),
|
return IO.Schema(
|
||||||
"b1": ("FLOAT", {"default": 1.3, "min": 0.0, "max": 10.0, "step": 0.01}),
|
node_id="FreeU_V2",
|
||||||
"b2": ("FLOAT", {"default": 1.4, "min": 0.0, "max": 10.0, "step": 0.01}),
|
category="model_patches/unet",
|
||||||
"s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}),
|
inputs=[
|
||||||
"s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}),
|
IO.Model.Input("model"),
|
||||||
}}
|
IO.Float.Input("b1", default=1.3, min=0.0, max=10.0, step=0.01),
|
||||||
RETURN_TYPES = ("MODEL",)
|
IO.Float.Input("b2", default=1.4, min=0.0, max=10.0, step=0.01),
|
||||||
FUNCTION = "patch"
|
IO.Float.Input("s1", default=0.9, min=0.0, max=10.0, step=0.01),
|
||||||
|
IO.Float.Input("s2", default=0.2, min=0.0, max=10.0, step=0.01),
|
||||||
|
],
|
||||||
|
outputs=[
|
||||||
|
IO.Model.Output(),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
CATEGORY = "model_patches/unet"
|
@classmethod
|
||||||
|
def execute(cls, model, b1, b2, s1, s2) -> IO.NodeOutput:
|
||||||
def patch(self, model, b1, b2, s1, s2):
|
|
||||||
model_channels = model.model.model_config.unet_config["model_channels"]
|
model_channels = model.model.model_config.unet_config["model_channels"]
|
||||||
scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
|
scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
|
||||||
on_cpu_devices = {}
|
on_cpu_devices = {}
|
||||||
@ -105,9 +120,19 @@ class FreeU_V2:
|
|||||||
|
|
||||||
m = model.clone()
|
m = model.clone()
|
||||||
m.set_model_output_block_patch(output_block_patch)
|
m.set_model_output_block_patch(output_block_patch)
|
||||||
return (m, )
|
return IO.NodeOutput(m)
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
patch = execute # TODO: remove
|
||||||
"FreeU": FreeU,
|
|
||||||
"FreeU_V2": FreeU_V2,
|
|
||||||
}
|
class FreelunchExtension(ComfyExtension):
|
||||||
|
@override
|
||||||
|
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||||
|
return [
|
||||||
|
FreeU,
|
||||||
|
FreeU_V2,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
async def comfy_entrypoint() -> FreelunchExtension:
|
||||||
|
return FreelunchExtension()
|
||||||
|
|||||||
@ -53,11 +53,6 @@ class PatchModelAddDownscale(io.ComfyNode):
|
|||||||
return io.NodeOutput(m)
|
return io.NodeOutput(m)
|
||||||
|
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
||||||
# Sampling
|
|
||||||
"PatchModelAddDownscale": "",
|
|
||||||
}
|
|
||||||
|
|
||||||
class ModelDownscaleExtension(ComfyExtension):
|
class ModelDownscaleExtension(ComfyExtension):
|
||||||
@override
|
@override
|
||||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user