diff --git a/__init__.py b/__init__.py index b0dbf9b..643a1b5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,2 +1,22 @@ +import sys +import pkg_resources +import subprocess + +def is_module_installed(module_name): + installed_packages = [d.key for d in pkg_resources.working_set] + return module_name in installed_packages + +def install_module(module_name): + subprocess.check_call([sys.executable, "-m", "pip", "install", module_name]) + +module_name = "diffusers" +if not is_module_installed(module_name): + print(f"### ComfyUI-LCM: {module_name} is not installed. Installing now...") + install_module(module_name) + print(f"### ComfyUI-LCM: {module_name} has been installed.") +else: + print(f"### ComfyUI-LCM: {module_name} is already installed.") + + from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py index e8b5240..2b5a9fc 100644 --- a/nodes.py +++ b/nodes.py @@ -24,9 +24,9 @@ class LCMSampler: @classmethod def INPUT_TYPES(s): return {"required": - {"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + { "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "steps": ("INT", {"default": 4, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), "size": ("INT", {"default": 512, "min": 512, "max": 768}), "num_images": ("INT", {"default": 1, "min": 1, "max": 64}), @@ -34,16 +34,16 @@ class LCMSampler: } } - RETURN_TYPES = ("LATENT",) + RETURN_TYPES = ("IMAGE",) FUNCTION = "sample" CATEGORY = "sampling" - def sample(self, ckpt_name, seed, steps, cfg, positive_prompt, size, num_images): + def sample(self, seed, steps, cfg, positive_prompt, size, num_images): if self.pipe is None: self.pipe = LatentConsistencyModelPipeline.from_pretrained( - "SimianLuo/LCM_Dreamshaper_v7",#folder_paths.get_annotated_filepath(ckpt_name, "checkpoints"), - #local_files_only=True, + pretrained_model_name_or_path="SimianLuo/LCM_Dreamshaper_v7", + local_files_only=True, scheduler=self.scheduler ) self.pipe.to(get_torch_device()) @@ -58,13 +58,14 @@ class LCMSampler: guidance_scale=cfg, num_inference_steps=steps, num_images_per_prompt=num_images, - lcm_origin_steps=50, - output_type="latent", + lcm_origin_steps=4, + output_type="np", ).images print("LCM inference time: ", time.time() - start_time, "seconds") + images_tensor = torch.from_numpy(result) - return ({"samples":result},) + return (images_tensor,) NODE_CLASS_MAPPINGS = { "LCMSampler": LCMSampler