File size: 1,548 Bytes
82ea528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
log = logging.getLogger(__name__)
import comfy.model_management as mm
from ..sampling.sampler import run_sampler
from ..utils.sampling_utils import prepare_conds
class MochiWrapperSamplerCustomNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MOCHIMODEL",),
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 30.0, "step": 0.01}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"sigmas": ("SIGMAS", {"tooltip": "Override sigma schedule and steps"}),
"latents": ("LATENT", ),
"sampler": ("SAMPLER", ),
"add_noise": ("BOOLEAN", ),
}
}
RETURN_TYPES = ("LATENT",)
RETURN_NAMES = ("samples",)
FUNCTION = "process"
CATEGORY = "MochiEdit/Wrapper"
def process(self, model, positive, negative, cfg, seed, sigmas, latents, sampler, add_noise):
mm.soft_empty_cache()
sigmas = sigmas.tolist()
latents = latents['samples']
positive, negative = prepare_conds(positive, negative)
latents = run_sampler(model, latents, positive, negative, sigmas, cfg, sampler.sampler_function, add_noise, seed)
mm.soft_empty_cache()
return ({"samples": latents},)
|