nodes/kijai-ComfyUI-Marigold
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

smarter model loading

Browse Source
This commit is contained in:
kijai
2024-04-08 11:13:29 +03:00
parent 20548ee45e
commit 80b3ee8948
1 changed files with 18 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+18 -17
View File
@@ -9,6 +9,7 @@ from .marigold.util.image_util import chw2hwc, colorize_depth_maps, resize_max_r
import comfy.utils import comfy.utils
import comfy.model_management import comfy.model_management
import folder_paths
def colorizedepth(depth_map, colorize_method): def colorizedepth(depth_map, colorize_method):
depth_map = depth_map.cpu().numpy() depth_map = depth_map.cpu().numpy()
@@ -101,19 +102,20 @@ class MarigoldDepthEstimation:
image = image.permute(0, 3, 1, 2).to(device).to(dtype=precision) image = image.permute(0, 3, 1, 2).to(device).to(dtype=precision)
if normalize: if normalize:
image = image * 2.0 - 1.0 image = image * 2.0 - 1.0
diffusers_model_path = os.path.join(folder_paths.models_dir,'diffusers')
#load the diffusers model #load the diffusers model
if model == "Marigold": if model == "Marigold":
folders_to_check = [ folders_to_check = [
"checkpoints/Marigold_v1_merged", os.path.join(script_directory,"checkpoints","Marigold_v1_merged",),
"checkpoints/Marigold", os.path.join(script_directory,"checkpoints","Marigold",),
"../../models/diffusers/Marigold_v1_merged", os.path.join(diffusers_model_path,"Marigold_v1_merged"),
"../../models/diffusers/Marigold", os.path.join(diffusers_model_path,"Marigold")
] ]
elif model == "marigold-lcm-v1-0": elif model == "marigold-lcm-v1-0":
folders_to_check = [ folders_to_check = [
"../../models/diffusers/marigold-lcm-v1-0", os.path.join(diffusers_model_path,"marigold-lcm-v1-0"),
"checkpoints/marigold-lcm-v1-0", os.path.join(diffusers_model_path,"checkpoints","marigold-lcm-v1-0")
] ]
self.custom_config = { self.custom_config = {
"model": model, "model": model,
@@ -125,27 +127,26 @@ class MarigoldDepthEstimation:
# Load the model only if it hasn't been loaded before # Load the model only if it hasn't been loaded before
checkpoint_path = None checkpoint_path = None
for folder in folders_to_check: for folder in folders_to_check:
potential_path = os.path.join(script_directory, folder) if os.path.exists(folder):
if os.path.exists(potential_path): checkpoint_path = folder
checkpoint_path = potential_path
break break
to_ignore = ["*.bin", "*fp16*"]
if checkpoint_path is None: if checkpoint_path is None:
if model == "Marigold": if model == "Marigold":
try: try:
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
checkpoint_path = os.path.join(script_directory, "../../models/diffusers/Marigold") checkpoint_path = os.path.join(diffusers_model_path, "Marigold")
snapshot_download(repo_id="Bingxin/Marigold", ignore_patterns=["*.bin"], local_dir=checkpoint_path, local_dir_use_symlinks=False) snapshot_download(repo_id="Bingxin/Marigold", ignore_patterns=to_ignore, local_dir=checkpoint_path, local_dir_use_symlinks=False)
except: except:
raise FileNotFoundError("No checkpoint directory found.") raise FileNotFoundError(f"No checkpoint directory found at {checkpoint_path}")
if model == "marigold-lcm-v1-0": if model == "marigold-lcm-v1-0":
try: try:
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
checkpoint_path = os.path.join(script_directory, "../../models/diffusers/marigold-lcm-v1-0") checkpoint_path = os.path.join(diffusers_model_path, "marigold-lcm-v1-0")
snapshot_download(repo_id="prs-eth/marigold-lcm-v1-0", ignore_patterns=["*.bin"], local_dir=checkpoint_path, local_dir_use_symlinks=False) snapshot_download(repo_id="prs-eth/marigold-lcm-v1-0", ignore_patterns=to_ignore, local_dir=checkpoint_path, local_dir_use_symlinks=False)
except: except:
raise FileNotFoundError("No checkpoint directory found.") raise FileNotFoundError(f"No checkpoint directory found at {checkpoint_path}")
self.marigold_pipeline = MarigoldPipeline.from_pretrained(checkpoint_path, enable_xformers=False, empty_text_embed=empty_text_embed, noise_scheduler_type=scheduler) self.marigold_pipeline = MarigoldPipeline.from_pretrained(checkpoint_path, enable_xformers=False, empty_text_embed=empty_text_embed, noise_scheduler_type=scheduler)
self.marigold_pipeline = self.marigold_pipeline.to(device).half() if use_fp16 else self.marigold_pipeline.to(device) self.marigold_pipeline = self.marigold_pipeline.to(device).half() if use_fp16 else self.marigold_pipeline.to(device)