diff --git a/nodes.py b/nodes.py index 1bb6e14..71b0769 100644 --- a/nodes.py +++ b/nodes.py @@ -11,6 +11,13 @@ import comfy.utils from contextlib import nullcontext from .lvdm.models.samplers.ddim import DDIMSampler +from contextlib import nullcontext +try: + from accelerate import init_empty_weights + is_accelerate_available = True +except: + pass + def split_and_trim(input_string): # Split the string into an array using '|' as a separator array = input_string.split('|') @@ -63,6 +70,7 @@ class DownloadAndLoadDynamiCrafterModel: CATEGORY = "DynamiCrafterWrapper" def loadmodel(self, dtype, model, fp8_unet=False): + device = mm.get_torch_device() mm.soft_empty_cache() custom_config = { 'dtype': dtype, @@ -103,21 +111,25 @@ class DownloadAndLoadDynamiCrafterModel: model_config = config.pop("model", OmegaConf.create()) model_config['params']['unet_config']['params']['use_checkpoint']=False - self.model = instantiate_from_config(model_config) - self.model = load_model_checkpoint(self.model, model_path) - self.model.eval() + if dtype == "auto": try: if mm.should_use_fp16(): - self.model.to(convert_dtype('fp16')) + precision = (convert_dtype('fp16')) elif mm.should_use_bf16(): - self.model.to(convert_dtype('bf16')) + precision = (convert_dtype('bf16')) else: - self.model.to(convert_dtype('fp32')) + precision = (convert_dtype('fp32')) except: raise AttributeError("ComfyUI version too old, can't autodetect properly. Set your dtype manually.") else: - self.model.to(convert_dtype(dtype)) + precision = (convert_dtype(dtype)) + + with (init_empty_weights() if is_accelerate_available else nullcontext()): + self.model = instantiate_from_config(model_config) + self.model = load_model_checkpoint(self.model, model_path, precision, device) + self.model.to(precision).to(device).eval() + if fp8_unet: self.model.model.diffusion_model = self.model.model.diffusion_model.to(torch.float8_e4m3fn) print(f"Model using dtype: {self.model.dtype}") @@ -596,9 +608,13 @@ class ToonCrafterInterpolation: videos2 = repeat(videos2, 'b c t h w -> b c (repeat t) h w', repeat=frames//2) videos = torch.cat([videos, videos2], dim=2) - z, hs = get_latent_z_with_hidden_states(self.model, videos) - hs = [t.to("cpu") for t in hs] - hidden_states.append(hs) + try: + z, hs = get_latent_z_with_hidden_states(self.model, videos) + hs = [t.to("cpu") for t in hs] + hidden_states.append(hs) + except: + z = get_latent_z(self.model, videos) + hidden_states = None img_tensor_repeat = torch.zeros_like(z) img_tensor_repeat[:,:,:1,:,:] = z[:,:,:1,:,:] @@ -758,10 +774,13 @@ class ToonCrafterDecode: with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=model.first_stage_model.dtype) if autocast_condition else nullcontext(): if mm.XFORMERS_IS_AVAILABLE: print(f"Decoding frames {iteration_counter * 16} - {16 + iteration_counter * 16} out of {num_samples} using xformers") - hs_ = hs[iteration_counter] - hs_ = [t.to(model.first_stage_model.device) for t in hs_] - additional_decode_kwargs = {'ref_context': hs_} - decoded_images = model.decode_first_stage(batch_samples, **additional_decode_kwargs) #b c t h w + if hs is not None: + hs_ = hs[iteration_counter] + hs_ = [t.to(model.first_stage_model.device) for t in hs_] + additional_decode_kwargs = {'ref_context': hs_} + decoded_images = model.decode_first_stage(batch_samples, **additional_decode_kwargs) #b c t h w + else: + decoded_images = model.decode_first_stage(batch_samples) #b c t h w else: raise Exception("XFormers not available, it is required for ToonCrafter decoder. Alternatively you can use a standard VAE Decode -node instead, but this has a negative effect on the image quality though.") diff --git a/requirements.txt b/requirements.txt index 1d39178..4111e25 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ omegaconf>=2.1.1 pytorch_lightning>=2.2.1 tqdm>=4.65.0 transformers>=4.25.1 -timm \ No newline at end of file +timm +accelerate \ No newline at end of file diff --git a/scripts/evaluation/funcs.py b/scripts/evaluation/funcs.py index bf61bc5..a1b03b8 100644 --- a/scripts/evaluation/funcs.py +++ b/scripts/evaluation/funcs.py @@ -5,51 +5,39 @@ import torch from einops import rearrange from safetensors.torch import load_file -def load_model_checkpoint(model, ckpt): - def load_checkpoint(model, ckpt, full_strict): - if "safetensors" in ckpt: - try: - state_dict = load_file(ckpt) - except: - state_dict = torch.load(ckpt, map_location="cpu") - else: - state_dict = torch.load(ckpt, map_location="cpu") - if "state_dict" in list(state_dict.keys()): - state_dict = state_dict["state_dict"] +from contextlib import nullcontext +try: + from accelerate import init_empty_weights + from accelerate.utils import set_module_tensor_to_device + is_accelerate_available = True +except: + pass - filtered_state_dict = { - k: v - for k, v in state_dict.items() - if not (k.startswith("cond_stage_model") or k.startswith("embedder")) - #if not (k.startswith("cond_stage_model")) - } # Filter out keys starting with "cond_stage_model" and "embedder" +def load_model_checkpoint(model, file_path, dtype, device): + if "safetensors" in file_path: try: - model.load_state_dict(filtered_state_dict, strict=full_strict) + state_dict = load_file(file_path) except: - ## rename the keys for 256x256 model - new_pl_sd = OrderedDict() - for k,v in state_dict.items(): - new_pl_sd[k] = v + state_dict = torch.load(file_path, map_location="cpu") + else: + state_dict = torch.load(file_path, map_location="cpu") + if "state_dict" in list(state_dict.keys()): + state_dict = state_dict["state_dict"] - for k in list(new_pl_sd.keys()): - if "framestride_embed" in k: - new_key = k.replace("framestride_embed", "fps_embedding") - new_pl_sd[new_key] = new_pl_sd[k] - del new_pl_sd[k] - model.load_state_dict(new_pl_sd, strict=full_strict) - # else: - # ## deepspeed - # new_pl_sd = OrderedDict() - # for key in state_dict['module'].keys(): - # new_pl_sd[key[16:]]=state_dict['module'][key] - # model.load_state_dict(new_pl_sd, strict=full_strict) + filtered_state_dict = { + k: v + for k, v in state_dict.items() + if not (k.startswith("cond_stage_model") or k.startswith("embedder")) + #if not (k.startswith("cond_stage_model")) + } # Filter out keys starting with "cond_stage_model" and "embedder" + if is_accelerate_available: + for key in filtered_state_dict: + set_module_tensor_to_device(model, key, dtype=dtype, device=device, value=filtered_state_dict[key]) + else: + model.load_state_dict(filtered_state_dict, strict=True) - return model - load_checkpoint(model, ckpt, full_strict=False) - print('>>> model checkpoint loaded.') return model - def load_prompts(prompt_file): f = open(prompt_file, 'r') prompt_list = []