From 4705c8908212ce7dd940f626cbded70f6dfe4abc Mon Sep 17 00:00:00 2001 From: thecooltechguy Date: Wed, 13 Dec 2023 19:53:34 +0000 Subject: [PATCH] optimized model downloads, still need to optimize more --- __init__.py | 1 + nodes.py | 6 ++--- prestartup_script.py | 55 ++++++++++++++++++++++++++++++++++++-------- 3 files changed, 48 insertions(+), 14 deletions(-) diff --git a/__init__.py b/__init__.py index 9f0a732..02ded1e 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,7 @@ sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "libs" os.makedirs(os.path.join(folder_paths.models_dir, "MagicAnimate"), exist_ok=True) folder_paths.add_model_folder_path("magic_animate", os.path.join(folder_paths.models_dir, "MagicAnimate")) +folder_paths.folder_names_and_paths['magic_animate'] = (folder_paths.folder_names_and_paths['magic_animate'][0], folder_paths.supported_pt_extensions) # | {'.json'}) magic_animate_checkpoints = folder_paths.get_filename_list("magic_animate") diff --git a/nodes.py b/nodes.py index eb0512d..96c1029 100644 --- a/nodes.py +++ b/nodes.py @@ -19,6 +19,7 @@ from magicanimate.utils.util import save_videos_grid from magicanimate.utils.dist_tools import distributed_init from accelerate.utils import set_seed from collections import OrderedDict +from PIL import Image class MagicAnimateModelLoader: def __init__(self): @@ -89,10 +90,7 @@ class MagicAnimateModelLoader: reference_control_writer = ReferenceAttentionControl(appearance_encoder, do_classifier_free_guidance=True, mode='write', fusion_blocks=config.fusion_blocks) reference_control_reader = ReferenceAttentionControl(unet, do_classifier_free_guidance=True, mode='read', fusion_blocks=config.fusion_blocks) - if config.pretrained_vae_path is not None: - vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path) - else: - vae = AutoencoderKL.from_pretrained(config.pretrained_model_path, subfolder="vae") + vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path) ### Load controlnet controlnet = ControlNetModel.from_pretrained(config.pretrained_controlnet_path) diff --git a/prestartup_script.py b/prestartup_script.py index fa25d6a..ae126dd 100644 --- a/prestartup_script.py +++ b/prestartup_script.py @@ -5,23 +5,58 @@ import sys magic_animate_models_dir = os.path.join(folder_paths.models_dir, "MagicAnimate") os.makedirs(magic_animate_models_dir, exist_ok=True) -if not os.path.exists(os.path.join(magic_animate_models_dir, "appearance_encoder")) or not os.path.exists(os.path.join(magic_animate_models_dir, "densepose_controlnet")) or not os.path.exists(os.path.join(magic_animate_models_dir, "temporal_attention")): - print("Downloading Magic Animate models...") - assert os.system("git lfs install") == 0, "ERROR: Git LFS is not installed. Please install it and restart ComfyUI." - assert os.system(f"cd {folder_paths.models_dir} && git clone https://huggingface.co/zcxu-eric/MagicAnimate") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." +appearance_encoder_path = os.path.join(magic_animate_models_dir, "appearance_encoder") +densepose_controlnet_path = os.path.join(magic_animate_models_dir, "densepose_controlnet") +temporal_attention_path = os.path.join(magic_animate_models_dir, "temporal_attention") -if not os.path.exists(os.path.join(magic_animate_models_dir, "stable-diffusion-v1-5")): - print("Downloading stable-diffusion-v1-5 model from Huggingface...") - assert os.system(f"cd {magic_animate_models_dir} && git clone https://huggingface.co/runwayml/stable-diffusion-v1-5") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model from Huggingface. Please check your internet connection and restart ComfyUI." +if not os.path.exists(appearance_encoder_path): # not os.path.exists(temporal_attention_path): + print("Downloading Magic Animate's appearance encoder...") + os.makedirs(appearance_encoder_path, exist_ok=True) + assert os.system(f"cd {appearance_encoder_path} && wget -q --show-progress --content-disposition 'https://huggingface.co/zcxu-eric/MagicAnimate/resolve/main/appearance_encoder/diffusion_pytorch_model.safetensors?download=true'") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {appearance_encoder_path} && wget -q --show-progress --content-disposition 'https://huggingface.co/zcxu-eric/MagicAnimate/resolve/main/appearance_encoder/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." + +if not os.path.exists(densepose_controlnet_path): + print("Downloading Magic Animate's densepose controlnet...") + os.makedirs(densepose_controlnet_path, exist_ok=True) + assert os.system(f"cd {densepose_controlnet_path} && wget -q --show-progress --content-disposition 'https://huggingface.co/zcxu-eric/MagicAnimate/resolve/main/densepose_controlnet/diffusion_pytorch_model.safetensors?download=true'") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {densepose_controlnet_path} && wget -q --show-progress --content-disposition 'https://huggingface.co/zcxu-eric/MagicAnimate/resolve/main/densepose_controlnet/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." + +if not os.path.exists(temporal_attention_path): + print("Downloading Magic Animate's temporal attention...") + os.makedirs(temporal_attention_path, exist_ok=True) + assert os.system(f"cd {temporal_attention_path} && wget -q --show-progress --content-disposition 'https://huggingface.co/zcxu-eric/MagicAnimate/resolve/main/temporal_attention/temporal_attention.ckpt?download=true'") == 0, "ERROR: Failed to download Magic Animate models. Please check your internet connection and restart ComfyUI." + +sd_15_path = os.path.join(magic_animate_models_dir, "stable-diffusion-v1-5") +if not os.path.exists(sd_15_path): + print("Downloading stable-diffusion-v1-5 checkpoints from Huggingface...") + # assert os.system(f"cd {magic_animate_models_dir} && git clone https://huggingface.co/runwayml/stable-diffusion-v1-5") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model from Huggingface. Please check your internet connection and restart ComfyUI." + os.makedirs(sd_15_path, exist_ok=True) + os.makedirs(os.path.join(sd_15_path, "text_encoder"), exist_ok=True) + os.makedirs(os.path.join(sd_15_path, "tokenizer"), exist_ok=True) + os.makedirs(os.path.join(sd_15_path, "unet"), exist_ok=True) + + assert os.system(f"cd {os.path.join(sd_15_path, 'text_encoder')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/text_encoder/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {os.path.join(sd_15_path, 'text_encoder')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors?download=true'") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." -if not os.path.exists(os.path.join(magic_animate_models_dir, "sd-vae-ft-mse")): + assert os.system(f"cd {os.path.join(sd_15_path, 'tokenizer')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/tokenizer/merges.txt?download=true' -O merges.txt") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {os.path.join(sd_15_path, 'tokenizer')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/tokenizer/special_tokens_map.json?download=true' -O special_tokens_map.json") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {os.path.join(sd_15_path, 'tokenizer')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/tokenizer/tokenizer_config.json?download=true' -O tokenizer_config.json") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {os.path.join(sd_15_path, 'tokenizer')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/tokenizer/vocab.json?download=true' -O vocab.json") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + + assert os.system(f"cd {os.path.join(sd_15_path, 'unet')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/unet/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {os.path.join(sd_15_path, 'unet')} && wget -q --show-progress --content-disposition 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/unet/diffusion_pytorch_model.bin?download=true'") == 0, "ERROR: Failed to download stable-diffusion-v1-5 model files from Huggingface. Please check your internet connection and restart ComfyUI." + +sd_vae_ft_mse_path = os.path.join(magic_animate_models_dir, "sd-vae-ft-mse") +if not os.path.exists(sd_vae_ft_mse_path): print("Downloading sd-vae-ft-mse model from Huggingface...") - assert os.system(f"cd {magic_animate_models_dir} && git clone https://huggingface.co/stabilityai/sd-vae-ft-mse") == 0, "ERROR: Failed to download sd-vae-ft-mse model from Huggingface. Please check your internet connection and restart ComfyUI." + os.makedirs(sd_vae_ft_mse_path, exist_ok=True) + assert os.system(f"cd {sd_vae_ft_mse_path} && wget --content-disposition -q --show-progress 'https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download sd-vae-ft-mse model from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {sd_vae_ft_mse_path} && wget --content-disposition -q --show-progress 'https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.safetensors?download=true'") == 0, "ERROR: Failed to download sd-vae-ft-mse model from Huggingface. Please check your internet connection and restart ComfyUI." control_v11p_sd15_openpose_path = os.path.join(magic_animate_models_dir, "control_v11p_sd15_openpose") if not os.path.exists(control_v11p_sd15_openpose_path): print("Downloading control_v11p_sd15_openpose model from Huggingface...") os.makedirs(control_v11p_sd15_openpose_path, exist_ok=True) assert os.system(f"cd {control_v11p_sd15_openpose_path} && wget --content-disposition -q --show-progress 'https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/resolve/main/diffusion_pytorch_model.safetensors?download=true'") == 0, "ERROR: Failed to download control_v11p_sd15_openpose model from Huggingface. Please check your internet connection and restart ComfyUI." - assert os.system(f"cd {control_v11p_sd15_openpose_path} && wget --content-disposition -q --show-progress 'https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/resolve/main/config.json?download=true'") == 0, "ERROR: Failed to download control_v11p_sd15_openpose model from Huggingface. Please check your internet connection and restart ComfyUI." + assert os.system(f"cd {control_v11p_sd15_openpose_path} && wget --content-disposition -q --show-progress 'https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/resolve/main/config.json?download=true' -O config.json") == 0, "ERROR: Failed to download control_v11p_sd15_openpose model from Huggingface. Please check your internet connection and restart ComfyUI." \ No newline at end of file