diff --git a/multitalk/nodes.py b/multitalk/nodes.py index 0ecea7d..21caa32 100644 --- a/multitalk/nodes.py +++ b/multitalk/nodes.py @@ -22,14 +22,8 @@ class MultiTalkModelLoader: def loadmodel(self, model, base_precision=None): from .multitalk import AudioProjModel - offload_device = mm.unet_offload_device() model_path = folder_paths.get_full_path_or_raise("diffusion_models", model) - if model_path.endswith(".gguf"): - from diffusers.models.model_loading_utils import load_gguf_checkpoint - sd = load_gguf_checkpoint(model_path) - else: - sd = load_torch_file(model_path, device=offload_device, safe_load=True) audio_window=5 intermediate_dim=512 @@ -50,8 +44,7 @@ class MultiTalkModelLoader: multitalk = { "proj_model": multitalk_proj_model, - "sd": sd, - "is_gguf": model_path.endswith(".gguf"), + "model_path": model_path, "model_type": "InfiniteTalk" if "infinite" in model.lower() else "MultiTalk", } diff --git a/nodes.py b/nodes.py index e361521..1c44246 100644 --- a/nodes.py +++ b/nodes.py @@ -1701,20 +1701,21 @@ class WanVideoSampler: is_5b = transformer.out_dim == 48 vae_upscale_factor = 16 if is_5b else 8 + # Load weights if transformer.patched_linear and gguf_reader is None: load_weights(patcher.model.diffusion_model, patcher.model["sd"], weight_dtype, base_dtype=dtype, transformer_load_device=device, block_swap_args=block_swap_args) - if gguf_reader is not None: + if gguf_reader is not None: #handle GGUF load_weights(transformer, patcher.model["sd"], base_dtype=dtype, transformer_load_device=device, patcher=patcher, gguf=True, reader=gguf_reader, block_swap_args=block_swap_args) set_lora_params_gguf(transformer, patcher.patches) transformer.patched_linear = True - elif len(patcher.patches) != 0 and transformer.patched_linear: + elif len(patcher.patches) != 0 and transformer.patched_linear: #handle patched linear layers (unmerged loras, fp8 scaled) log.info(f"Using {len(patcher.patches)} LoRA weight patches for WanVideo model") if not merge_loras and fp8_matmul: raise NotImplementedError("FP8 matmul with unmerged LoRAs is not supported") set_lora_params(transformer, patcher.patches) else: - remove_lora_from_module(transformer) + remove_lora_from_module(transformer) #clear possible unmerged lora weights transformer.lora_scheduling_enabled = transformer_options.get("lora_scheduling_enabled", False) @@ -2386,16 +2387,18 @@ class WanVideoSampler: import copy sample_scheduler_flipped = copy.deepcopy(sample_scheduler) - #rope + # Rotary positional embeddings (RoPE) + + # RoPE base freq scaling as used with CineScale ntk_alphas = [1.0, 1.0, 1.0] if isinstance(rope_function, dict): ntk_alphas = rope_function["ntk_scale_f"], rope_function["ntk_scale_h"], rope_function["ntk_scale_w"] rope_function = rope_function["rope_function"] - + freqs = None transformer.rope_embedder.k = None transformer.rope_embedder.num_frames = None - if "default" in rope_function or bidirectional_sampling: + if "default" in rope_function or bidirectional_sampling: # original RoPE d = transformer.dim // transformer.num_heads freqs = torch.cat([ rope_params(1024, d - 4 * (d // 6), L_test=latent_video_length, k=riflex_freq_index), @@ -2403,7 +2406,7 @@ class WanVideoSampler: rope_params(1024, 2 * (d // 6)) ], dim=1) - elif "comfy" in rope_function: + elif "comfy" in rope_function: # comfy's rope transformer.rope_embedder.k = riflex_freq_index transformer.rope_embedder.num_frames = latent_video_length @@ -2565,37 +2568,37 @@ class WanVideoSampler: base_params = { - 'seq_len': seq_len, - 'device': device, - 'freqs': freqs, - 't': timestep, - 'current_step': idx, - 'last_step': len(timesteps) - 1 == idx, - 'control_lora_enabled': control_lora_enabled, - 'enhance_enabled': enhance_enabled, - 'camera_embed': camera_embed, - 'unianim_data': unianim_data, - 'fun_ref': fun_ref_input if fun_ref_image is not None else None, - 'fun_camera': control_camera_input if control_camera_latents is not None else None, - 'audio_proj': audio_proj if fantasytalking_embeds is not None else None, - 'audio_scale': audio_scale, - "pcd_data": pcd_data_input, - "controlnet": controlnet, - "add_cond": add_cond_input, - "nag_params": text_embeds.get("nag_params", {}), - "nag_context": text_embeds.get("nag_prompt_embeds", None), - "multitalk_audio": multitalk_audio_input if multitalk_audio_embedding is not None else None, - "ref_target_masks": ref_target_masks if multitalk_audio_embedding is not None else None, - "inner_t": [shot_len] if shot_len else None, - "standin_input": standin_input, - "fantasy_portrait_input": fantasy_portrait_input, - "phantom_ref": phantom_ref, - "reverse_time": reverse_time, - "ntk_alphas": ntk_alphas, - "mtv_motion_tokens": mtv_motion_tokens if mtv_input is not None else None, - "mtv_motion_rotary_emb": mtv_motion_rotary_emb if mtv_input is not None else None, - "mtv_strength": mtv_strength[idx] if mtv_input is not None else 1.0, - "mtv_freqs": mtv_freqs if mtv_input is not None else None, + 'seq_len': seq_len, # sequence length + 'device': device, # main device + 'freqs': freqs, # rope freqs + 't': timestep, # current timestep + 'current_step': idx, # current step + 'last_step': len(timesteps) - 1 == idx, # is last step + 'control_lora_enabled': control_lora_enabled, # control lora toggle for patch embed selection + 'enhance_enabled': enhance_enabled, # enhance-a-video toggle + 'camera_embed': camera_embed, # recammaster embedding + 'unianim_data': unianim_data, # unianimate input + 'fun_ref': fun_ref_input if fun_ref_image is not None else None, # Fun model reference latent + 'fun_camera': control_camera_input if control_camera_latents is not None else None, # Fun model camera embed + 'audio_proj': audio_proj if fantasytalking_embeds is not None else None, # FantasyTalking audio projection + 'audio_scale': audio_scale, # FantasyTalking audio scale + "pcd_data": pcd_data_input, # Uni3C input + "controlnet": controlnet, # TheDenk's controlnet input + "add_cond": add_cond_input, # additional conditioning input + "nag_params": text_embeds.get("nag_params", {}), # normalized attention guidance + "nag_context": text_embeds.get("nag_prompt_embeds", None), # normalized attention guidance context + "multitalk_audio": multitalk_audio_input if multitalk_audio_embedding is not None else None, # Multi/InfiniteTalk audio input + "ref_target_masks": ref_target_masks if multitalk_audio_embedding is not None else None, # Multi/InfiniteTalk reference target masks + "inner_t": [shot_len] if shot_len else None, # inner timestep for EchoShot + "standin_input": standin_input, # Stand-in reference input + "fantasy_portrait_input": fantasy_portrait_input, # Fantasy portrait input + "phantom_ref": phantom_ref, # Phantom reference input + "reverse_time": reverse_time, # Reverse RoPE toggle + "ntk_alphas": ntk_alphas, # RoPE freq scaling values + "mtv_motion_tokens": mtv_motion_tokens if mtv_input is not None else None, # MTV-Crafter motion tokens + "mtv_motion_rotary_emb": mtv_motion_rotary_emb if mtv_input is not None else None, # MTV-Crafter RoPE + "mtv_strength": mtv_strength[idx] if mtv_input is not None else 1.0, # MTV-Crafter scaling + "mtv_freqs": mtv_freqs if mtv_input is not None else None, # MTV-Crafter extra RoPE freqs } batch_size = 1 diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 6c90875..a5f6567 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -741,6 +741,13 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None, log.info("Using GGUF to load and assign model weights to device...") # Prepare sd from GGUF readers + + # UniAnimate embedding weight workaround + unianimate_sd = {} + for key in sd.keys(): + if "dwpose_embedding" in key or "randomref_embedding_pose" in key: + unianimate_sd[key] = sd[key] + sd = {} all_tensors = [] for r in reader: @@ -769,6 +776,8 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None, is_gguf_quant = tensor.tensor_type not in [GGMLQuantizationType.F32, GGMLQuantizationType.F16] weights = torch.from_numpy(tensor.data.copy()).to(load_device) sd[tensor.name] = GGUFParameter(weights, quant_type=tensor.tensor_type) if is_gguf_quant else weights + sd.update(unianimate_sd) + del unianimate_sd if not getattr(transformer, "gguf_patched", False): transformer = _replace_with_gguf_linear( @@ -795,7 +804,7 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None, except Exception: block_idx = None - if "loras" in name or "dwpose" in name or "randomref" in name: + if "loras" in name: continue # GGUF: skip GGUFParameter params @@ -820,7 +829,6 @@ def load_weights(transformer, sd=None, weight_dtype=None, base_dtype=None, elif vace_block_idx is not None: if vace_block_idx >= len(transformer.vace_blocks) - block_swap_args.get("vace_blocks_to_swap", 0): load_device = offload_device - # Set tensor to device set_module_tensor_to_device(transformer, name, device=load_device, dtype=dtype_to_use, value=sd[name.replace("_orig_mod.", "")]) cnt += 1 @@ -869,6 +877,7 @@ def patch_stand_in_lora(transformer, lora_sd, transformer_load_device, base_dtyp param.data.copy_(lora_sd["diffusion_model." + name].to(param.device, dtype=param.dtype)) def add_lora_weights(patcher, lora, base_dtype, merge_loras=False): + unianimate_sd = None #spacepxl's control LoRA patch for l in lora: log.info(f"Loading LoRA: {l['name']} with strength: {l['strength']}") @@ -885,7 +894,7 @@ def add_lora_weights(patcher, lora, base_dtype, merge_loras=False): if "dwpose_embedding.0.weight" in lora_sd: #unianimate from .unianimate.nodes import update_transformer log.info("Unianimate LoRA detected, patching model...") - patcher.model.diffusion_model = update_transformer(patcher.model.diffusion_model, lora_sd) + patcher.model.diffusion_model, unianimate_sd = update_transformer(patcher.model.diffusion_model, lora_sd) lora_sd = standardize_lora_key_format(lora_sd) @@ -906,7 +915,7 @@ def add_lora_weights(patcher, lora, base_dtype, merge_loras=False): patcher, _ = load_lora_for_models(patcher, None, lora_sd, lora_strength, 0) del lora_sd - return patcher, control_lora + return patcher, control_lora, unianimate_sd #region Model loading class WanVideoModelLoader: @@ -1033,15 +1042,18 @@ class WanVideoModelLoader: if "vace_blocks.0.after_proj.weight" in sd and not "patch_embedding.weight" in sd: raise ValueError("You are attempting to load a VACE module as a WanVideo model, instead you should use the vace_model input and matching T2V base model") + # currently this can be VAE or MTV-Crafter weights if extra_model is not None: if gguf: if not extra_model["path"].endswith(".gguf"): raise ValueError("With GGUF main model the extra model must also be a GGUF quantized, if the main model already has extra included, you can disconnect the extra module loader") extra_sd, extra_reader = load_gguf(extra_model["path"]) gguf_reader.append(extra_reader) + del extra_reader else: extra_sd = load_torch_file(extra_model["path"], device=transformer_load_device, safe_load=True) sd.update(extra_sd) + del extra_sd first_key = next(iter(sd)) if first_key.startswith("model.diffusion_model."): @@ -1206,40 +1218,54 @@ class WanVideoModelLoader: log.info("FantasyPortrait model detected, patching model...") context_dim = fantasyportrait_model["sd"]["ip_adapter.blocks.0.cross_attn.ip_adapter_single_stream_k_proj.weight"].shape[1] - for block in transformer.blocks: - block.cross_attn.ip_adapter_single_stream_k_proj = nn.Linear(context_dim, dim, bias=False) - block.cross_attn.ip_adapter_single_stream_v_proj = nn.Linear(context_dim, dim, bias=False) + with init_empty_weights(): + for block in transformer.blocks: + block.cross_attn.ip_adapter_single_stream_k_proj = nn.Linear(context_dim, dim, bias=False) + block.cross_attn.ip_adapter_single_stream_v_proj = nn.Linear(context_dim, dim, bias=False) ip_adapter_sd = {} for k, v in fantasyportrait_model["sd"].items(): if k.startswith("ip_adapter."): ip_adapter_sd[k.replace("ip_adapter.", "")] = v sd.update(ip_adapter_sd) + del ip_adapter_sd if multitalk_model is not None: - if multitalk_model["is_gguf"] and not gguf: - raise ValueError("Multitalk/InfiniteTalk model is a GGUF model, main model also has to be a GGUF model.") multitalk_model_type = multitalk_model.get("model_type", "MultiTalk") + log.info(f"{multitalk_model_type} detected, patching model...") + + multitalk_model_path = multitalk_model["model_path"] + if multitalk_model_path.endswith(".gguf") and not gguf: + raise ValueError("Multitalk/InfiniteTalk model is a GGUF model, main model also has to be a GGUF model.") + # init audio module from .multitalk.multitalk import SingleStreamMultiAttention from .wanvideo.modules.model import WanLayerNorm - norm_input_visual = True #dunno what this is for block in transformer.blocks: - block.audio_cross_attn = SingleStreamMultiAttention( - dim=dim, - encoder_hidden_states_dim=768, - num_heads=num_heads, + with init_empty_weights(): + block.norm_x = WanLayerNorm(dim, transformer.eps, elementwise_affine=True) + block.audio_cross_attn = SingleStreamMultiAttention( + dim=dim, + encoder_hidden_states_dim=768, + num_heads=num_heads, qkv_bias=True, class_range=24, class_interval=4, attention_mode=attention_mode, ) - block.norm_x = WanLayerNorm(dim, transformer.eps, elementwise_affine=True) if norm_input_visual else nn.Identity() - log.info(f"{multitalk_model_type} detected, patching model...") transformer.audio_proj = multitalk_model["proj_model"] transformer.multitalk_model_type = multitalk_model_type - sd.update(multitalk_model["sd"]) - + + extra_model_path = multitalk_model["model_path"] + if gguf: + extra_sd, extra_reader = load_gguf(extra_model_path) + gguf_reader.append(extra_reader) + del extra_reader + else: + extra_sd = load_torch_file(extra_model_path, device=transformer_load_device, safe_load=True) + sd.update(extra_sd) + del extra_sd + # Additional cond latents if "add_conv_in.weight" in sd: def zero_module(module): @@ -1266,178 +1292,54 @@ class WanVideoModelLoader: patcher.model.is_patched = False scale_weights = {} + if "fp8" in quantization: + for k, v in sd.items(): + if k.endswith(".scale_weight"): + scale_weights[k] = v.to(base_dtype) + + if "fp8_e4m3fn" in quantization: + weight_dtype = torch.float8_e4m3fn + elif "fp8_e5m2" in quantization: + weight_dtype = torch.float8_e5m2 + else: + weight_dtype = base_dtype + + params_to_keep = {"norm", "bias", "time_in", "patch_embedding", "time_", "img_emb", "modulation", "text_embedding", "adapter", "add", "ref_conv", "audio_proj"} + + control_lora = False + + if not merge_loras and control_lora: + log.warning("Control-LoRA patching is only supported with merge_loras=True") + + if lora is not None: + patcher, control_lora, unianimate_sd = add_lora_weights(patcher, lora, base_dtype, merge_loras=merge_loras) + if unianimate_sd is not None: + log.info("Merging UniAnimate weights to the model...") + sd.update(unianimate_sd) + del unianimate_sd + if not gguf: - if "fp8" in quantization: - for k, v in sd.items(): - if k.endswith(".scale_weight"): - scale_weights[k] = v - if not merge_loras: + if merge_loras and lora is not None: + if not lora_low_mem_load: + load_weights(transformer, sd, weight_dtype, base_dtype, transformer_load_device) + + if control_lora: + patch_control_lora(patcher.model.diffusion_model, device) + patcher.model.is_patched = True + + log.info("Merging LoRA to the model...") + patcher = apply_lora( + patcher, device, transformer_load_device, params_to_keep=params_to_keep, dtype=weight_dtype, base_dtype=base_dtype, state_dict=sd, + low_mem_load=lora_low_mem_load, control_lora=control_lora, scale_weights=scale_weights,) + if not control_lora: + scale_weights.clear() + patcher.patches.clear() + transformer.patched_linear = False + else: from .custom_linear import _replace_linear transformer = _replace_linear(transformer, base_dtype, sd, scale_weights=scale_weights) - - if "fp8_e4m3fn" in quantization: - dtype = torch.float8_e4m3fn - elif "fp8_e5m2" in quantization: - dtype = torch.float8_e5m2 - else: - dtype = base_dtype - params_to_keep = {"norm", "bias", "time_in", "patch_embedding", "time_", "img_emb", "modulation", "text_embedding", "adapter", "add", "ref_conv", "audio_proj"} - if not lora_low_mem_load: - log.info("Using accelerate to load and assign model weights to device...") - param_count = sum(1 for _ in transformer.named_parameters()) - pbar = ProgressBar(param_count) - cnt = 0 - for name, param in tqdm(transformer.named_parameters(), - desc=f"Loading transformer parameters to {transformer_load_device}", - total=param_count, - leave=True): - dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype - dtype_to_use = dtype if sd[name].dtype == dtype else dtype_to_use - if "modulation" in name or "norm" in name or "bias" in name: - dtype_to_use = base_dtype - if "patch_embedding" in name: - dtype_to_use = torch.float32 - set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name]) - cnt += 1 - if cnt % 100 == 0: - pbar.update(100) + transformer.patched_linear = True - #for name, param in transformer.named_parameters(): - # print(name, param.dtype, param.device, param.shape) - pbar.update_absolute(param_count) - pbar.update_absolute(0) - - comfy_model.diffusion_model = transformer - comfy_model.load_device = transformer_load_device - - patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, offload_device) - patcher.model.is_patched = False - - control_lora = False - if lora is not None: - for l in lora: - log.info(f"Loading LoRA: {l['name']} with strength: {l['strength']}") - lora_path = l["path"] - lora_strength = l["strength"] - if isinstance(lora_strength, list): - if merge_loras: - raise ValueError("LoRA strength should be a single value when merge_loras=True") - transformer.lora_scheduling_enabled = True - if lora_strength == 0: - log.warning(f"LoRA {lora_path} has strength 0, skipping...") - continue - lora_sd = load_torch_file(lora_path, safe_load=True) - if "dwpose_embedding.0.weight" in lora_sd: #unianimate - from .unianimate.nodes import update_transformer - log.info("Unianimate LoRA detected, patching model...") - transformer = update_transformer(transformer, lora_sd) - - lora_sd = standardize_lora_key_format(lora_sd) - - if l["blocks"]: - lora_sd = filter_state_dict_by_blocks(lora_sd, l["blocks"], l.get("layer_filter", [])) - - # Filter out any LoRA keys containing 'img' if the base model state_dict has no 'img' keys - if not any('img' in k for k in sd.keys()): - lora_sd = {k: v for k, v in lora_sd.items() if 'img' not in k} - - #spacepxl's control LoRA patch - # for key in lora_sd.keys(): - # print(key) - - if "diffusion_model.patch_embedding.lora_A.weight" in lora_sd: - log.info("Control-LoRA detected, patching model...") - if not merge_loras: - log.warning("Control-LoRA patching is only supported with merge_loras=True, setting it to True") - merge_loras = True - control_lora = True - - in_cls = transformer.patch_embedding.__class__ # nn.Conv3d - old_in_dim = transformer.in_dim # 16 - new_in_dim = lora_sd["diffusion_model.patch_embedding.lora_A.weight"].shape[1] - assert new_in_dim == 32 - - new_in = in_cls( - new_in_dim, - transformer.patch_embedding.out_channels, - transformer.patch_embedding.kernel_size, - transformer.patch_embedding.stride, - transformer.patch_embedding.padding, - ).to(device=device, dtype=torch.float32) - - new_in.weight.zero_() - new_in.bias.zero_() - - new_in.weight[:, :old_in_dim].copy_(transformer.patch_embedding.weight) - new_in.bias.copy_(transformer.patch_embedding.bias) - - transformer.patch_embedding = new_in - transformer.expanded_patch_embedding = new_in - - if "diffusion_model.blocks.0.self_attn.q_loras.down.weight" in lora_sd: - log.info("Stand-In LoRA detected") - for block in transformer.blocks: - block.self_attn.q_loras = LoRALinearLayer(dim, dim, rank=128, device=transformer_load_device, dtype=base_dtype, strength=lora_strength) - block.self_attn.k_loras = LoRALinearLayer(dim, dim, rank=128, device=transformer_load_device, dtype=base_dtype, strength=lora_strength) - block.self_attn.v_loras = LoRALinearLayer(dim, dim, rank=128, device=transformer_load_device, dtype=base_dtype, strength=lora_strength) - for lora in [block.self_attn.q_loras, block.self_attn.k_loras, block.self_attn.v_loras]: - for param in lora.parameters(): - param.requires_grad = False - for name, param in transformer.named_parameters(): - if "lora" in name: - param.data.copy_(lora_sd["diffusion_model." + name].to(param.device, dtype=param.dtype)) - else: - patcher, _ = load_lora_for_models(patcher, None, lora_sd, lora_strength, 0) - - del lora_sd - - if not gguf and merge_loras: - log.info("Patching LoRA to the model...") - patcher = apply_lora( - patcher, device, transformer_load_device, - params_to_keep=params_to_keep, dtype=dtype, base_dtype=base_dtype, state_dict=sd, - low_mem_load=lora_low_mem_load, control_lora=control_lora, scale_weights=scale_weights) - scale_weights.clear() - patcher.patches.clear() - - if gguf: - #from diffusers.quantizers.gguf.utils import _replace_with_gguf_linear, GGUFParameter - from .gguf.gguf import _replace_with_gguf_linear, GGUFParameter - log.info("Using GGUF to load and assign model weights to device...") - param_count = sum(1 for _ in transformer.named_parameters()) - - out_features = sd["blocks.0.self_attn.k.weight"].shape[1] - - patcher.model.diffusion_model = _replace_with_gguf_linear(patcher.model.diffusion_model, base_dtype, sd, patches=patcher.patches) - pbar = ProgressBar(param_count) - cnt = 0 - for name, param in tqdm(patcher.model.diffusion_model.named_parameters(), - desc=f"Loading transformer parameters to {transformer_load_device}", - total=param_count, - leave=True): - if "loras" in name: - continue - #print(name, param.dtype, param.device, param.shape) - if isinstance(param, GGUFParameter): - dtype_to_use = torch.uint8 - elif "patch_embedding" in name: - dtype_to_use = torch.float32 - else: - dtype_to_use = base_dtype - set_module_tensor_to_device(patcher.model.diffusion_model, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name]) - cnt += 1 - if cnt % 100 == 0: - pbar.update(100) - - #for name, param in transformer.named_parameters(): - # print(name, param.dtype, param.device, param.shape) - #patcher.load(device, full_load=True) - pbar.update_absolute(param_count) - - patcher.model.is_patched = True - - patch_linear = (True if "scaled" in quantization or (lora is not None and not merge_loras) else False) - if "fast" in quantization: if lora is not None and not merge_loras: raise NotImplementedError("fp8_fast is not supported with unmerged LoRAs") @@ -1577,8 +1479,6 @@ class WanVideoVAELoader: def loadmodel(self, model_name, precision, compile_args=None): dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] - #with open(os.path.join(script_directory, 'configs', 'hy_vae_config.json')) as f: - # vae_config = json.load(f) model_path = folder_paths.get_full_path("vae", model_name) vae_sd = load_torch_file(model_path, safe_load=True) @@ -1592,8 +1492,9 @@ class WanVideoVAELoader: vae = WanVideoVAE38(dtype=dtype) vae.load_state_dict(vae_sd) + del vae_sd vae.eval() - vae.to(device = offload_device, dtype = dtype) + vae.to(device=offload_device, dtype=dtype) if compile_args is not None: vae.model.decoder = torch.compile(vae.model.decoder, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"])