mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-08-02 23:14:49 +08:00
Make deep shrink behave like it should.
This commit is contained in:
@@ -8,13 +8,14 @@ class PatchModelAddDownscale:
|
||||
"downscale_factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 9.0, "step": 0.001}),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"end_percent": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"downscale_after_skip": ("BOOLEAN", {"default": True}),
|
||||
}}
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
|
||||
CATEGORY = "_for_testing"
|
||||
|
||||
def patch(self, model, block_number, downscale_factor, start_percent, end_percent):
|
||||
def patch(self, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip):
|
||||
sigma_start = model.model.model_sampling.percent_to_sigma(start_percent).item()
|
||||
sigma_end = model.model.model_sampling.percent_to_sigma(end_percent).item()
|
||||
|
||||
@@ -31,7 +32,10 @@ class PatchModelAddDownscale:
|
||||
return h, hsp
|
||||
|
||||
m = model.clone()
|
||||
m.set_model_input_block_patch(input_block_patch)
|
||||
if downscale_after_skip:
|
||||
m.set_model_input_block_patch_after_skip(input_block_patch)
|
||||
else:
|
||||
m.set_model_input_block_patch(input_block_patch)
|
||||
m.set_model_output_block_patch(output_block_patch)
|
||||
return (m, )
|
||||
|
||||
|
Reference in New Issue
Block a user